From 7e8fd9e5906a0c69a2686d450596f5494a6855d1 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Wed, 1 Jul 2026 02:32:09 -1000 Subject: [PATCH 01/16] Fix wait_agent/research supervisor coordination and TUI sub-agent monitor. Accumulate child streaming output for parent polling, add auto-cursor and cancel-aware waits, inherit parent ACP connections for child turns, and improve sub-agent list/transcript UX in the TUI. --- crates/core/src/skills.rs | 1 - crates/core/src/tools/handlers/agent.rs | 136 +++++--- crates/protocol/src/acp_ts.rs | 6 +- crates/protocol/src/agent.rs | 155 ++++++++- crates/server/src/runtime.rs | 3 + crates/server/src/runtime/agents.rs | 27 +- .../server/src/runtime/agents/coordinator.rs | 63 +++- crates/server/src/runtime/agents/lifecycle.rs | 10 +- .../src/runtime/handlers/acp/session.rs | 1 + crates/server/src/runtime/handlers/turn.rs | 10 +- crates/server/src/runtime/research_events.rs | 26 +- crates/server/src/runtime/research_parsing.rs | 17 + crates/server/src/subagent.rs | 134 +++++++- crates/server/tests/deep_research_e2e.rs | 2 +- crates/server/tests/subagent_lifecycle.rs | 38 +- .../tests/support/subagent_lifecycle.rs | 12 +- crates/tui/src/chatwidget/input.rs | 3 + .../tui/src/chatwidget/subagent_live_list.rs | 3 +- crates/tui/src/chatwidget/subagent_monitor.rs | 325 ++++++++++++++++-- crates/tui/src/chatwidget_tests.rs | 44 ++- crates/tui/src/events.rs | 4 + crates/tui/src/worker.rs | 11 +- crates/tui/src/worker/acp_events.rs | 24 +- .../L2-DES-AGENT-003-subagent-architecture.md | 27 +- 24 files changed, 896 insertions(+), 186 deletions(-) diff --git a/crates/core/src/skills.rs b/crates/core/src/skills.rs index 1dab9b11..e861e0c4 100644 --- a/crates/core/src/skills.rs +++ b/crates/core/src/skills.rs @@ -1,6 +1,5 @@ //! Core-facing skill catalog wrapper backed by `devo-skills`. -use devo_util_paths; use std::path::Path; use std::path::PathBuf; diff --git a/crates/core/src/tools/handlers/agent.rs b/crates/core/src/tools/handlers/agent.rs index 74feaa39..55901793 100644 --- a/crates/core/src/tools/handlers/agent.rs +++ b/crates/core/src/tools/handlers/agent.rs @@ -2,19 +2,32 @@ use std::collections::BTreeMap; use std::sync::Arc; use async_trait::async_trait; -use devo_protocol::{ - AgentContextMode, AgentListParams, AgentMessageParams, AgentToolPolicy, CloseAgentParams, - SessionId, SpawnAgentParams, WaitAgentParams, -}; - -use crate::contracts::{ - ToolCallError, ToolContext, ToolProgress, ToolProgressSender, ToolResult, ToolResultContent, -}; +use devo_protocol::AgentContextMode; +use devo_protocol::AgentListParams; +use devo_protocol::AgentMessageParams; +use devo_protocol::AgentToolPolicy; +use devo_protocol::CloseAgentParams; +use devo_protocol::ParentAgentInfo; +use devo_protocol::ParentAgentListResult; +use devo_protocol::ParentSpawnAgentResult; +use devo_protocol::SessionId; +use devo_protocol::SpawnAgentParams; +use devo_protocol::WaitAgentParams; + +use crate::contracts::ToolCallError; +use crate::contracts::ToolContext; +use crate::contracts::ToolProgress; +use crate::contracts::ToolProgressSender; +use crate::contracts::ToolResult; +use crate::contracts::ToolResultContent; use crate::json_schema::JsonSchema; use crate::registry::ToolExposure; use crate::registry::ToolRegistryBuilder; use crate::tool_handler::ToolHandler; -use crate::tool_spec::{ToolExecutionMode, ToolOutputMode, ToolPreparationFeedback, ToolSpec}; +use crate::tool_spec::ToolExecutionMode; +use crate::tool_spec::ToolOutputMode; +use crate::tool_spec::ToolPreparationFeedback; +use crate::tool_spec::ToolSpec; #[derive(Clone, Copy)] enum AgentToolKind { @@ -111,7 +124,7 @@ impl ToolHandler for AgentToolHandler { ephemeral: false, }) .await?; - json_result(result, "agent spawned") + json_result(ParentSpawnAgentResult::from(result), "agent spawned") } AgentToolKind::SendMessage => { let input: AgentMessageInput = parse_input(input)?; @@ -136,7 +149,7 @@ impl ToolHandler for AgentToolHandler { session_id, target: input.target, after_sequence: input.after_sequence, - timeout_ms: input.timeout_ms, + timeout_secs: input.timeout_secs, }; let result = tokio::select! { result = coordinator.wait_agent(params) => result?, @@ -152,7 +165,12 @@ impl ToolHandler for AgentToolHandler { path_prefix: input.path_prefix, }) .await?; - json_result(serde_json::json!({ "agents": agents }), "agents listed") + json_result( + ParentAgentListResult { + agents: agents.into_iter().map(ParentAgentInfo::from).collect(), + }, + "agents listed", + ) } AgentToolKind::Close => { let input: CloseAgentInput = parse_input(input)?; @@ -188,7 +206,7 @@ struct WaitAgentInput { #[serde(default)] after_sequence: Option, #[serde(default)] - timeout_ms: Option, + timeout_secs: Option, } #[derive(serde::Deserialize)] @@ -242,19 +260,19 @@ fn spec(name: &str, description: &str, schema: JsonSchema) -> ToolSpec { fn spawn_spec() -> ToolSpec { spec( "spawn_agent", - "Launch a new Devo agent to handle complex, multi-step tasks autonomously.\n\nUse this from a parent session to parallelize independent research or implementation work, then use wait_agent to collect child output. Launch multiple agents concurrently whenever possible when the work is independent.\n\nWriting the prompt:\n- Brief the agent like a smart colleague who just walked into the room: it has not seen this conversation, does not know what you have tried, and does not understand why the task matters unless you tell it.\n- Explain what you are trying to accomplish and why.\n- Describe what you have already learned or ruled out.\n- Give enough context about the surrounding problem that the agent can make judgment calls rather than just following narrow instructions.\n- If you need a short response, say so, for example \"report in under 200 words\".\n- Lookups: hand over the exact command. Investigations: hand over the question; prescribed steps become dead weight when the premise is wrong.\n- Terse command-style prompts produce shallow, generic work.\n- Never delegate understanding. Do not write \"based on your findings, fix the bug\" or \"based on the research, implement it.\" Those phrases push synthesis onto the agent instead of doing it yourself. Write prompts that prove you understood: include file paths, line numbers, and what specifically to change.\n\nWhen not to use an agent:\n- If you want to read a specific file path, use the read tool or file search instead.\n- If you are searching for a specific class, symbol, or string, use grep/code search instead when available.\n- If you are searching within a specific file or a small set of files, read those files directly.\n- Do not use an agent for tasks unrelated to available agent descriptions or the current request.\n\nThe agent's result is not visible to the user. To show the user the result, send a concise summary after wait_agent returns.", + "Launch a new child agent for complex multi-step work. Agent coordination tools (spawn_agent, send_message, wait_agent, list_agents, close_agent) are parent-only.\n\nTypical flow: spawn_agent -> wait_agent until status completes -> optionally send_message for a follow-up turn -> wait_agent again. Use list_agents for status without child text; close_agent to stop a child. Parallelize independent work by spawning multiple children whenever possible.\n\nChild output is only visible through wait_agent. Never infer or summarize child findings before wait_agent returns.\n\nWriting the prompt:\n- Brief the agent like a colleague who just arrived: no shared conversation unless fork_turns provides history.\n- State goal, why it matters, what you already ruled out, and the expected deliverable.\n- Lookups: give exact commands. Investigations: give the question, not a brittle script.\n- Never delegate understanding with phrases like \"based on your findings, fix it.\" Include concrete paths, symbols, and constraints.\n\nWhen not to use:\n- Reading a known file path -> read tool.\n- Searching a symbol or string -> grep/code search.\n- Small scoped file reads -> read directly.\n\nThe user does not see child output directly. Summarize for the user after wait_agent returns.\n\nExample: spawn_agent({message:\"Survey crates/server for wait_agent usage and summarize call sites.\"})", JsonSchema::object( BTreeMap::from([ ( "message".to_string(), JsonSchema::string(Some( - "Initial task message for the child agent. Include the goal, scope, files or subsystems to inspect, context needed for judgment calls, and the expected result.", + "Initial child task. Include goal, scope, relevant paths, and expected result format.", )), ), ( "fork_turns".to_string(), JsonSchema::string(Some( - "Optional history fork mode. In coding-agent sessions, use \"all\" (default) when the child needs stable completed parent history; it excludes the active parent turn. Use \"none\" for a clean child context containing only the task message. In DeepResearch sessions, this is always forced to \"none\". Do not assume the child sees your active turn unless you include needed context in message.", + "\"all\" (coding-agent default): stable completed parent history, excludes the active parent turn. \"none\": only the task message. DeepResearch sessions always force \"none\".", )), ), ]), @@ -267,7 +285,7 @@ fn spawn_spec() -> ToolSpec { fn send_message_spec() -> ToolSpec { spec( "send_message", - "Parent-only tool that sends additional user input to an existing child agent. If the child is idle, the message starts a child turn; if active, it queues for the next turn. Use this to continue a previously spawned agent with the agent's full context preserved. Each newly spawned agent starts from its configured fork context, so provide a complete task description when spawning rather than relying on hidden assumptions.", + "Send more input to an existing child agent. Idle children start a new turn immediately; active children queue the message for the next turn.\n\nWhen to use:\n- Follow up after a completed turn on the same child.\n- Correct or narrow the task without spawning a duplicate worker.\n\nWhen not to use:\n- Collecting output -> wait_agent.\n- Checking if still running -> list_agents.\n- Stopping a child -> close_agent.\n\nMulti-turn rule: each send_message starts a new child turn. Prior wait_agent results belong to the previous turn only. After send_message, call wait_agent again and wait for a fresh status event before treating output as final.\n\nExample: send_message({target:\"brave-apple\", message:\"Also check error paths in coordinator.rs.\"})", message_schema(), ) } @@ -275,26 +293,24 @@ fn send_message_spec() -> ToolSpec { fn wait_agent_spec() -> ToolSpec { spec( "wait_agent", - "Parent-only tool that polls child assistant output and terminal status events. Use after spawn_agent or send_message to collect incremental child results.\n\nDo not peek at generated transcript files or tail child output unless the user explicitly asks for a progress check. Reading a transcript mid-flight pulls the child agent's tool noise into your context and defeats the point of delegation.\n\nDo not race. After launching a child agent, you know nothing about what it found. Never fabricate or predict child results in any format: not as prose, summary, or structured output. If the user asks a follow-up before output lands, tell them the child agent is still running and give status, not a guess.", + "Collect child assistant output and turn-completion status. Each assistant_message event is the full accumulated text for that turn, not token deltas.\n\nDecision tree:\n1. After spawn_agent or send_message -> wait_agent with a longer timeout_secs (e.g. 60-120) until a status event (completed/failed/interrupted/closed).\n2. If timed_out with no events -> list_agents; if still running, wait_agent again with a short timeout_secs (e.g. 2-5).\n3. If completed with assistant_message -> use the output; send_message only if more work is needed.\n4. If off-track, stuck, or user wants to stop -> close_agent.\n\nDo not read child transcript files mid-flight. Do not fabricate child results. If the user asks early, report list_agents status instead of guessing.\n\nExample: wait_agent({\"target\":\"brave-apple\",\"timeout_secs\":90}) -> on timed_out with no events, list_agents({}) then wait_agent({\"target\":\"brave-apple\",\"timeout_secs\":3})", JsonSchema::object( BTreeMap::from([ ( "target".to_string(), JsonSchema::string(Some( - "Optional child agent path, generated nickname, or session id. Omit to poll all direct children.", + "Child agent_nickname or agent_path from spawn_agent, list_agents, or prior wait_agent output. Omit to poll all direct children.", )), ), ( "after_sequence".to_string(), JsonSchema::integer(Some( - "Only return output events with sequence greater than this value. To poll incrementally, pass the largest sequence value from the previous events list.", + "Return events with sequence greater than this value. Omit on first poll to use the runtime cursor.", )), ), ( - "timeout_ms".to_string(), - JsonSchema::integer(Some( - "Optional wait timeout in milliseconds. If no newer output exists, wait up to this timeout before returning timed_out.", - )), + "timeout_secs".to_string(), + JsonSchema::integer(Some("Wait up to this many seconds (default 5, max 120).")), ), ]), None, @@ -306,11 +322,11 @@ fn wait_agent_spec() -> ToolSpec { fn list_agents_spec() -> ToolSpec { spec( "list_agents", - "Parent-only tool that lists child agents for the current session, including generated path, nickname, status, and last task message.", + "Lightweight status snapshot for child agents. Does not return assistant text and does not block.\n\nWhen to use:\n- Right after spawn_agent to confirm registration and copy agent_nickname/agent_path.\n- After wait_agent timed_out with no events to see if the child is still running.\n- Before send_message or close_agent when multiple children exist.\n- When the user asks for progress without needing findings yet.\n\nWhen not to use:\n- Collecting child findings -> wait_agent.\n- Stopping a child -> close_agent.\n\nStatus values: running, completed, failed, interrupted, closed, spawning. running with an empty wait_agent poll usually means the child is still working.\n\nExample: list_agents({})", JsonSchema::object( BTreeMap::from([( "path_prefix".to_string(), - JsonSchema::string(Some("Optional generated child path or path prefix filter.")), + JsonSchema::string(Some("Optional agent_path prefix filter.")), )]), None, Some(false), @@ -321,12 +337,12 @@ fn list_agents_spec() -> ToolSpec { fn close_agent_spec() -> ToolSpec { spec( "close_agent", - "Parent-only tool that closes an existing child agent, interrupts active work if needed, and records a terminal output event.", + "Stop a child agent, interrupt active work if needed, and record a terminal status event.\n\nWhen to use:\n- Child is off-track or producing useless work.\n- Child stays running with no useful progress after list_agents + short wait_agent polls.\n- User asks to cancel or you no longer need the worker.\n- You spawned the wrong worker and will not use its output.\n\nWhen not to use:\n- Collecting output from a healthy completed child -> wait_agent first, then close if cleanup is needed.\n- Checking status -> list_agents.\n- Sending corrections -> send_message.\n\nExample: close_agent({target:\"brave-apple\"})", JsonSchema::object( BTreeMap::from([( "target".to_string(), JsonSchema::string(Some( - "Target child agent path, generated nickname, or session id.", + "Child agent_nickname or agent_path from spawn_agent or list_agents.", )), )]), Some(vec!["target".to_string()]), @@ -341,14 +357,12 @@ fn message_schema() -> JsonSchema { ( "target".to_string(), JsonSchema::string(Some( - "Target child agent path, generated nickname, or session id.", + "Child agent_nickname or agent_path from spawn_agent or list_agents.", )), ), ( "message".to_string(), - JsonSchema::string(Some( - "Additional task input to deliver to the child as user text.", - )), + JsonSchema::string(Some("Follow-up user message for the child agent.")), ), ]), Some(vec!["target".to_string(), "message".to_string()]), @@ -543,36 +557,54 @@ mod tests { .as_str() .expect("fork_turns description"); + assert!(spawn.description.contains("parent-only")); assert!(spawn.description.contains("wait_agent")); - assert!(fork_description.contains("\"all\" (default)")); + assert!( + spawn + .description + .contains("only visible through wait_agent") + ); + assert!(fork_description.contains("\"all\" (coding-agent default)")); assert!(fork_description.contains("stable completed parent history")); assert!(fork_description.contains("excludes the active parent turn")); assert!(fork_description.contains("DeepResearch sessions")); - assert!(fork_description.contains("always forced to \"none\"")); + assert!(fork_description.contains("always force \"none\"")); assert!(fork_description.contains("\"none\"")); - assert!(send_message_spec().description.contains("Parent-only")); - assert!( - send_message_spec() - .description - .contains("queues for the next turn") - ); + + let send = send_message_spec(); + assert!(!send.description.contains("Parent-only")); assert!( - wait_agent_spec() - .description - .contains("polls child assistant output") + send.description + .contains("queue the message for the next turn") ); - let wait_schema = wait_agent_spec().input_schema.to_json_value(); + assert!(send.description.contains("Multi-turn rule")); + assert!(send.description.contains("wait_agent again")); + + let wait = wait_agent_spec(); + assert!(!wait.description.contains("Parent-only")); + assert!(wait.description.contains("Decision tree")); + assert!(wait.description.contains(r#"wait_agent({"target""#)); + assert!(wait.description.contains("assistant_message")); + assert!(wait.description.contains("not token deltas")); + assert!(wait.description.contains("list_agents")); + assert!(wait.description.contains("close_agent")); + let wait_schema = wait.input_schema.to_json_value(); let after_sequence_description = wait_schema["properties"]["after_sequence"]["description"] .as_str() .expect("after_sequence description"); - assert!(after_sequence_description.contains("largest sequence value")); - assert!(!after_sequence_description.contains("previous next_sequence")); - assert!(list_agents_spec().description.contains("generated path")); - assert!( - close_agent_spec() - .description - .contains("terminal output event") - ); + assert!(after_sequence_description.contains("runtime cursor")); + assert!(wait_schema["properties"].get("timeout_secs").is_some()); + + let list = list_agents_spec(); + assert!(!list.description.contains("Parent-only")); + assert!(list.description.contains("timed_out with no events")); + assert!(list.description.contains("Does not return assistant text")); + assert!(list.description.contains("agent_nickname")); + + let close = close_agent_spec(); + assert!(!close.description.contains("Parent-only")); + assert!(close.description.contains("off-track")); + assert!(close.description.contains("terminal status event")); } #[tokio::test] diff --git a/crates/protocol/src/acp_ts.rs b/crates/protocol/src/acp_ts.rs index 17768bbf..dc8dc2cd 100644 --- a/crates/protocol/src/acp_ts.rs +++ b/crates/protocol/src/acp_ts.rs @@ -333,7 +333,11 @@ pub fn generate_protocol_typescript() -> String { push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); - push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); + push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); push_decl::(&cfg, &mut output); diff --git a/crates/protocol/src/agent.rs b/crates/protocol/src/agent.rs index bf7ee646..083959b2 100644 --- a/crates/protocol/src/agent.rs +++ b/crates/protocol/src/agent.rs @@ -25,6 +25,21 @@ pub enum AgentContextMode { DeepResearch, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] +#[serde(rename_all = "snake_case")] +pub enum AgentOutputEventKind { + #[serde(alias = "assistant_delta")] + AssistantDelta, + AssistantMessage, + Status, +} + +impl AgentOutputEventKind { + pub fn is_assistant_text(self) -> bool { + matches!(self, Self::AssistantDelta | Self::AssistantMessage) + } +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] pub struct SpawnAgentParams { pub session_id: SessionId, @@ -49,6 +64,24 @@ pub struct SpawnAgentResult { pub status: String, } +/// Model-facing spawn result: address children by path or nickname, not session ids. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] +pub struct ParentSpawnAgentResult { + pub agent_path: String, + pub agent_nickname: String, + pub status: String, +} + +impl From for ParentSpawnAgentResult { + fn from(result: SpawnAgentResult) -> Self { + Self { + agent_path: result.agent_path, + agent_nickname: result.agent_nickname, + status: result.status, + } + } +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] pub struct AgentMessageParams { pub session_id: SessionId, @@ -69,7 +102,16 @@ pub struct WaitAgentParams { #[serde(default, skip_serializing_if = "Option::is_none")] pub after_sequence: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - pub timeout_ms: Option, + pub timeout_secs: Option, +} + +pub const DEFAULT_WAIT_AGENT_TIMEOUT_SECS: u64 = 5; +pub const MAX_WAIT_AGENT_TIMEOUT_SECS: u64 = 120; + +pub fn resolve_wait_agent_timeout(timeout_secs: Option) -> u64 { + timeout_secs + .unwrap_or(DEFAULT_WAIT_AGENT_TIMEOUT_SECS) + .min(MAX_WAIT_AGENT_TIMEOUT_SECS) } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] @@ -86,19 +128,20 @@ pub struct AgentMailboxMessage { #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] pub struct WaitAgentResult { - pub events: Vec, + pub events: Vec, pub next_sequence: u64, pub timed_out: bool, } -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] +/// Buffered child output stored server-side (includes routing metadata). +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct AgentOutputEvent { pub sequence: u64, pub child_session_id: SessionId, pub agent_path: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub turn_id: Option, - pub kind: String, + pub kind: AgentOutputEventKind, #[serde(default, skip_serializing_if = "Option::is_none")] pub text: Option, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -106,6 +149,38 @@ pub struct AgentOutputEvent { pub created_at: DateTime, } +/// Model-facing wait output: address children by path or nickname; omit internal session and turn ids. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] +pub struct ParentAgentOutputEvent { + pub sequence: u64, + pub agent_path: String, + pub agent_nickname: String, + pub kind: AgentOutputEventKind, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub text: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub status: Option, +} + +impl From for ParentAgentOutputEvent { + fn from(event: AgentOutputEvent) -> Self { + let agent_nickname = event + .agent_path + .rsplit('/') + .next() + .unwrap_or(&event.agent_path) + .to_string(); + Self { + sequence: event.sequence, + agent_path: event.agent_path, + agent_nickname, + kind: event.kind, + text: event.text, + status: event.status, + } + } +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] pub struct AgentInfo { pub session_id: SessionId, @@ -119,6 +194,29 @@ pub struct AgentInfo { pub last_task_message: Option, } +/// Model-facing list entry: address children by path or nickname. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] +pub struct ParentAgentInfo { + pub agent_path: String, + pub agent_nickname: String, + pub agent_role: String, + pub status: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub last_task_message: Option, +} + +impl From for ParentAgentInfo { + fn from(info: AgentInfo) -> Self { + Self { + agent_path: info.agent_path, + agent_nickname: info.agent_nickname, + agent_role: info.agent_role, + status: info.status, + last_task_message: info.last_task_message, + } + } +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] pub struct AgentListParams { pub session_id: SessionId, @@ -131,6 +229,11 @@ pub struct AgentListResult { pub agents: Vec, } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] +pub struct ParentAgentListResult { + pub agents: Vec, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] pub struct AgentStatusParams { pub session_id: SessionId, @@ -149,6 +252,14 @@ pub struct CloseAgentResult { pub status: String, } +pub fn wait_agent_cursor_key(target: Option<&str>) -> String { + target + .map(str::trim) + .filter(|target| !target.is_empty()) + .unwrap_or_default() + .to_string() +} + #[cfg(test)] mod tests { use pretty_assertions::assert_eq; @@ -176,15 +287,13 @@ mod tests { status: "running".to_string(), }, "wait": WaitAgentResult { - events: vec![AgentOutputEvent { + events: vec![ParentAgentOutputEvent { sequence: 1, - child_session_id, agent_path: "root/review".to_string(), - turn_id: Some(TurnId::new()), - kind: "assistant_delta".to_string(), + agent_nickname: "review".to_string(), + kind: AgentOutputEventKind::AssistantMessage, text: Some("done".to_string()), status: None, - created_at: Utc::now(), }], next_sequence: 2, timed_out: false, @@ -196,4 +305,32 @@ mod tests { assert_eq!(restored, payloads); } + + #[test] + fn resolve_wait_agent_timeout_uses_default_and_clamps() { + assert_eq!(resolve_wait_agent_timeout(None), 5); + assert_eq!(resolve_wait_agent_timeout(Some(30)), 30); + assert_eq!(resolve_wait_agent_timeout(Some(999)), 120); + } + + #[test] + fn parent_visible_output_omits_internal_ids_and_includes_nickname() { + let child_session_id = SessionId::new(); + let event = AgentOutputEvent { + sequence: 1, + child_session_id, + agent_path: "root/review".to_string(), + turn_id: Some(TurnId::new()), + kind: AgentOutputEventKind::AssistantMessage, + text: Some("done".to_string()), + status: None, + created_at: Utc::now(), + }; + let parent_visible = ParentAgentOutputEvent::from(event); + assert_eq!(parent_visible.agent_nickname, "review"); + let json = serde_json::to_value(parent_visible).expect("serialize parent output"); + assert!(json.get("child_session_id").is_none()); + assert!(json.get("turn_id").is_none()); + assert!(json.get("created_at").is_none()); + } } diff --git a/crates/server/src/runtime.rs b/crates/server/src/runtime.rs index 0924bed6..35d8cc2c 100644 --- a/crates/server/src/runtime.rs +++ b/crates/server/src/runtime.rs @@ -222,6 +222,8 @@ pub struct ServerRuntime { agent_mailboxes: Mutex>, /// Per-parent child-output buffers used by wait_agent polling. agent_output_buffers: Mutex>, + /// Per-parent `wait_agent` sequence cursors keyed by optional target string. + agent_wait_cursors: Mutex>>, /// Child agents owned by an active `/research` pipeline. research_child_agents: Mutex>>, /// Latest subagent turn usage grouped under the parent turn that requested the work. @@ -366,6 +368,7 @@ impl ServerRuntime { agent_registries: Mutex::new(HashMap::new()), agent_mailboxes: Mutex::new(HashMap::new()), agent_output_buffers: Mutex::new(HashMap::new()), + agent_wait_cursors: Mutex::new(HashMap::new()), research_child_agents: Mutex::new(HashMap::new()), subagent_usage: Mutex::new(subagent_usage::SubagentUsageState::default()), reference_searches: Mutex::new(HashMap::new()), diff --git a/crates/server/src/runtime/agents.rs b/crates/server/src/runtime/agents.rs index 00623611..8ef61df4 100644 --- a/crates/server/src/runtime/agents.rs +++ b/crates/server/src/runtime/agents.rs @@ -10,8 +10,6 @@ mod coordinator; mod handlers; mod lifecycle; -const DEFAULT_WAIT_AGENT_TIMEOUT: Duration = Duration::from_secs(300); -const MAX_WAIT_AGENT_TIMEOUT: Duration = Duration::from_secs(900); const AGENT_NAME_ADJECTIVES: &[&str] = &[ "brave", "clever", "silent", "happy", "gentle", "swift", "bright", "lazy", "wild", "calm", "fuzzy", "tiny", "bold", "lucky", "mighty", @@ -275,6 +273,15 @@ impl ServerRuntime { }, ) .await; + if let Some(parent_turn_id) = parent_active_turn_id { + self.record_subagent_status_event( + parent_session_id, + child_session_id, + SubagentStatus::Spawning, + parent_turn_id, + ) + .await; + } self.register_subagent_usage_owner( parent_session_id, child_session_id, @@ -555,6 +562,16 @@ impl ServerRuntime { let turn_for_task = turn.clone(); let turn_config_for_task = turn_config.clone(); let cancel_token = CancellationToken::new(); + let parent_session_id = { + let session = session_arc.lock().await; + session.summary.parent_session_id + }; + if let Some(parent_session_id) = parent_session_id { + let mut connections = self.active_turn_connections.lock().await; + if let Some(connection_id) = connections.get(&parent_session_id).copied() { + connections.insert(session_id, connection_id); + } + } self.active_turn_cancellations .lock() .await @@ -679,7 +696,7 @@ impl ServerRuntime { Ok(()) } - async fn resolve_child_agent( + pub(in crate::runtime) async fn resolve_child_agent( &self, parent_session_id: SessionId, target: &str, @@ -889,7 +906,7 @@ impl ServerRuntime { child_session_id, agent_path, turn_id: payload.context.turn_id, - kind: "assistant_delta".to_string(), + kind: devo_protocol::AgentOutputEventKind::AssistantMessage, text: Some(payload.delta.clone()), status: None, created_at: Utc::now(), @@ -937,7 +954,7 @@ impl ServerRuntime { child_session_id, agent_path, turn_id: Some(turn_id), - kind: "status".to_string(), + kind: devo_protocol::AgentOutputEventKind::Status, text, status: Some(status.as_str().to_string()), created_at: Utc::now(), diff --git a/crates/server/src/runtime/agents/coordinator.rs b/crates/server/src/runtime/agents/coordinator.rs index e67aba4a..e85c3be1 100644 --- a/crates/server/src/runtime/agents/coordinator.rs +++ b/crates/server/src/runtime/agents/coordinator.rs @@ -1,6 +1,36 @@ use super::*; impl ServerRuntime { + fn wait_agent_cursor_key(target: Option<&str>) -> String { + devo_protocol::wait_agent_cursor_key(target) + } + + async fn wait_agent_cursor(&self, parent_session_id: SessionId, target_key: &str) -> u64 { + self.agent_wait_cursors + .lock() + .await + .get(&parent_session_id) + .and_then(|cursors| cursors.get(target_key).copied()) + .unwrap_or_default() + } + + async fn update_wait_agent_cursor( + &self, + parent_session_id: SessionId, + target_key: &str, + consumed_sequence: u64, + ) { + if consumed_sequence == 0 { + return; + } + self.agent_wait_cursors + .lock() + .await + .entry(parent_session_id) + .or_default() + .insert(target_key.to_string(), consumed_sequence); + } + async fn send_message_inner( self: &Arc, params: devo_protocol::AgentMessageParams, @@ -17,24 +47,43 @@ impl ServerRuntime { &self, params: devo_protocol::WaitAgentParams, ) -> Result { - let timeout = params - .timeout_ms - .map(Duration::from_millis) - .unwrap_or(DEFAULT_WAIT_AGENT_TIMEOUT) - .min(MAX_WAIT_AGENT_TIMEOUT); + let timeout = Duration::from_secs(devo_protocol::resolve_wait_agent_timeout( + params.timeout_secs, + )); let target_session_ids = self .resolve_wait_agent_targets(params.session_id, params.target.as_deref()) .await?; + let cursor_key = Self::wait_agent_cursor_key(params.target.as_deref()); + let effective_after_sequence = match params.after_sequence { + Some(after_sequence) => after_sequence, + None => self.wait_agent_cursor(params.session_id, &cursor_key).await, + }; let output_buffer = self.output_buffer(params.session_id).await; + let cancel = self + .active_turn_cancellations + .lock() + .await + .get(¶ms.session_id) + .cloned(); let (events, next_sequence, timed_out) = output_buffer .wait_after( - params.after_sequence.unwrap_or_default(), + effective_after_sequence, &target_session_ids, timeout, + cancel, ) .await; + if let Some(consumed_sequence) = events.iter().map(|event| event.sequence).max() + && params.after_sequence.is_none() + { + self.update_wait_agent_cursor(params.session_id, &cursor_key, consumed_sequence) + .await; + } Ok(devo_protocol::WaitAgentResult { - events, + events: events + .into_iter() + .map(devo_protocol::ParentAgentOutputEvent::from) + .collect(), next_sequence, timed_out, }) diff --git a/crates/server/src/runtime/agents/lifecycle.rs b/crates/server/src/runtime/agents/lifecycle.rs index ba4131ee..053c3e53 100644 --- a/crates/server/src/runtime/agents/lifecycle.rs +++ b/crates/server/src/runtime/agents/lifecycle.rs @@ -263,20 +263,18 @@ mod tests { session_id: parent_session_id, target: None, after_sequence: None, - timeout_ms: Some(1), + timeout_secs: Some(0), }) .await?; assert_eq!( wait_result.events, - vec![devo_protocol::AgentOutputEvent { + vec![devo_protocol::ParentAgentOutputEvent { sequence: 1, - child_session_id, agent_path: "root/review".to_string(), - turn_id: wait_result.events[0].turn_id, - kind: "status".to_string(), + agent_nickname: "review".to_string(), + kind: devo_protocol::AgentOutputEventKind::Status, text: Some("failed to start child agent: rollout append failed".to_string()), status: Some("failed".to_string()), - created_at: wait_result.events[0].created_at, }] ); diff --git a/crates/server/src/runtime/handlers/acp/session.rs b/crates/server/src/runtime/handlers/acp/session.rs index fed2c4d0..24194c2c 100644 --- a/crates/server/src/runtime/handlers/acp/session.rs +++ b/crates/server/src/runtime/handlers/acp/session.rs @@ -518,6 +518,7 @@ impl ServerRuntime { self.goal_stores.lock().await.remove(&session_id); self.agent_mailboxes.lock().await.remove(&session_id); self.agent_output_buffers.lock().await.remove(&session_id); + self.agent_wait_cursors.lock().await.remove(&session_id); { let mut registries = self.agent_registries.lock().await; registries.remove(&session_id); diff --git a/crates/server/src/runtime/handlers/turn.rs b/crates/server/src/runtime/handlers/turn.rs index 20b62323..7fd3f892 100644 --- a/crates/server/src/runtime/handlers/turn.rs +++ b/crates/server/src/runtime/handlers/turn.rs @@ -708,9 +708,13 @@ impl ServerRuntime { if let Some(task) = self.active_tasks.lock().await.remove(¶ms.session_id) { task.abort(); } - Arc::clone(self) - .close_research_child_agents(params.session_id) - .await; + let close_children_runtime = Arc::clone(self); + let close_children_parent = params.session_id; + tokio::spawn(async move { + close_children_runtime + .close_research_child_agents(close_children_parent) + .await; + }); let interrupted_turn = { let mut session = session_arc.lock().await; let Some(mut turn) = session.active_turn.take() else { diff --git a/crates/server/src/runtime/research_events.rs b/crates/server/src/runtime/research_events.rs index b09a5a1c..6ddedaf3 100644 --- a/crates/server/src/runtime/research_events.rs +++ b/crates/server/src/runtime/research_events.rs @@ -8,7 +8,7 @@ use super::research_capture::{ use super::research_parsing::{ is_request_user_input_tool_name, is_spawn_agent_tool_name, request_user_input_exchanges_from_response, request_user_input_questions_from_input, - spawn_agent_child_session_id, tool_content_to_json, + spawn_agent_child_session_id, spawn_agent_child_target, tool_content_to_json, }; use super::research_stages::ResearchStageKind; use super::research_stages::StreamedResearchArtifact; @@ -200,13 +200,23 @@ impl ServerRuntime { summary, } => { let output = tool_content_to_json(content); - if is_spawn_agent_tool_name(&tool_name) - && !is_error - && let Some(child_session_id) = spawn_agent_child_session_id(&output) - { - self.remember_research_child_agent(session_id, child_session_id) - .await; - capture.spawned_worker_count += 1; + if is_spawn_agent_tool_name(&tool_name) && !is_error { + let child_session_id = + if let Some(child_session_id) = spawn_agent_child_session_id(&output) { + Some(child_session_id) + } else if let Some(target) = spawn_agent_child_target(&output) { + self.resolve_child_agent(session_id, &target) + .await + .ok() + .map(|metadata| metadata.session_id) + } else { + None + }; + if let Some(child_session_id) = child_session_id { + self.remember_research_child_agent(session_id, child_session_id) + .await; + capture.spawned_worker_count += 1; + } } if let Some(pending) = capture.pending_tools.remove(&tool_use_id) { self.complete_item( diff --git a/crates/server/src/runtime/research_parsing.rs b/crates/server/src/runtime/research_parsing.rs index 7cc9f65d..621a7d21 100644 --- a/crates/server/src/runtime/research_parsing.rs +++ b/crates/server/src/runtime/research_parsing.rs @@ -93,6 +93,23 @@ pub(super) fn spawn_agent_child_session_id(output: &serde_json::Value) -> Option .map(|result| result.child_session_id) } +pub(super) fn spawn_agent_child_target(output: &serde_json::Value) -> Option { + output + .get("agent_path") + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|target| !target.is_empty()) + .map(str::to_string) + .or_else(|| { + output + .get("agent_nickname") + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|target| !target.is_empty()) + .map(str::to_string) + }) +} + pub(super) fn tool_content_to_json(content: ToolContent) -> serde_json::Value { match content { ToolContent::Text(text) => serde_json::Value::String(text), diff --git a/crates/server/src/subagent.rs b/crates/server/src/subagent.rs index d9aa514f..3909fae1 100644 --- a/crates/server/src/subagent.rs +++ b/crates/server/src/subagent.rs @@ -14,14 +14,27 @@ use chrono::{DateTime, Utc}; use devo_protocol::AgentInfo; use devo_protocol::AgentMailboxMessage; use devo_protocol::AgentOutputEvent; +use devo_protocol::AgentOutputEventKind; use devo_protocol::SessionId; use serde::{Deserialize, Serialize}; use tokio::sync::Mutex; use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; // ── Agent Registry ────────────────────────────────────────────────── /// Per-root-session registry of all spawned subagents. +/// +/// The registry models a parent/child tree: any registered agent may in principle +/// spawn further children, and [`parent_to_children`](Self::parent_to_children) / +/// [`child_to_parent`](Self::child_to_parent) track that hierarchy for arbitrary depth. +/// +/// In current product usage, only the root session agent is expected to coordinate +/// subagents; child sessions do not receive agent coordination tools and must not +/// spawn nested subagents. That restriction is enforced at runtime (tool policy and +/// prompts), not by flattening this data structure. The tree-shaped registry is kept +/// intentionally so deeper nesting can be enabled later without redesigning storage +/// or lookup. #[derive(Debug, Clone, Default)] pub struct AgentRegistry { pub agents: HashMap, @@ -295,6 +308,21 @@ impl SubagentOutputBuffer { pub async fn push(&self, mut event: AgentOutputEvent) -> AgentOutputEvent { let mut inner = self.inner.lock().await; + if event.kind.is_assistant_text() { + event.kind = AgentOutputEventKind::AssistantMessage; + if let Some(last) = inner.events.back_mut() + && last.kind == AgentOutputEventKind::AssistantMessage + && last.child_session_id == event.child_session_id + && last.turn_id == event.turn_id + && last.status.is_none() + { + if let Some(delta) = event.text.take() { + last.text.get_or_insert_with(String::new).push_str(&delta); + } + last.created_at = event.created_at; + return last.clone(); + } + } event.sequence = inner.next_sequence; inner.next_sequence = inner.next_sequence.saturating_add(1); inner.events.push_back(event.clone()); @@ -308,10 +336,20 @@ impl SubagentOutputBuffer { after_sequence: u64, target_session_ids: &[SessionId], timeout: Duration, + cancel: Option, ) -> (Vec, u64, bool) { + if target_session_ids.is_empty() { + let inner = self.inner.lock().await; + return (Vec::new(), inner.next_sequence, true); + } let target_session_ids = target_session_ids.iter().copied().collect::>(); let start = Instant::now(); loop { + if cancel.as_ref().is_some_and(CancellationToken::is_cancelled) { + let (_, next_sequence) = + self.events_after(after_sequence, &target_session_ids).await; + return (Vec::new(), next_sequence, true); + } let notified = self.notify.notified(); let (events, next_sequence) = self.events_after(after_sequence, &target_session_ids).await; @@ -322,13 +360,27 @@ impl SubagentOutputBuffer { if elapsed >= timeout { return (Vec::new(), next_sequence, true); } - if tokio::time::timeout(timeout.saturating_sub(elapsed), notified) - .await - .is_err() - { - let (_, next_sequence) = - self.events_after(after_sequence, &target_session_ids).await; - return (Vec::new(), next_sequence, true); + let remaining = timeout.saturating_sub(elapsed); + let sleep = tokio::time::sleep(remaining); + tokio::pin!(sleep); + tokio::select! { + _ = &mut sleep => { + let (_, next_sequence) = + self.events_after(after_sequence, &target_session_ids).await; + return (Vec::new(), next_sequence, true); + } + _ = notified => {} + _ = async { + if let Some(token) = &cancel { + token.cancelled().await; + } else { + std::future::pending::<()>().await; + } + } => { + let (_, next_sequence) = + self.events_after(after_sequence, &target_session_ids).await; + return (Vec::new(), next_sequence, true); + } } } } @@ -373,6 +425,74 @@ pub enum SubagentError { #[cfg(test)] mod tests { use super::*; + use pretty_assertions::assert_eq; + + #[tokio::test] + async fn output_buffer_accumulates_assistant_deltas_per_turn() { + let buffer = SubagentOutputBuffer::new(); + let child = SessionId::new(); + let turn_id = devo_protocol::TurnId::new(); + let base = || AgentOutputEvent { + sequence: 0, + child_session_id: child, + agent_path: "root/worker".into(), + turn_id: Some(turn_id), + kind: AgentOutputEventKind::AssistantDelta, + text: None, + status: None, + created_at: Utc::now(), + }; + + let first = buffer + .push(AgentOutputEvent { + text: Some("alpha ".into()), + ..base() + }) + .await; + let second = buffer + .push(AgentOutputEvent { + text: Some("beta".into()), + ..base() + }) + .await; + assert_eq!(first.sequence, 1); + assert_eq!(second.sequence, 1); + assert_eq!(second.text.as_deref(), Some("alpha beta")); + assert_eq!(second.kind, AgentOutputEventKind::AssistantMessage); + + let (events, next_sequence, timed_out) = buffer + .wait_after(0, &[child], Duration::from_millis(1), None) + .await; + assert!(!timed_out); + assert_eq!(next_sequence, 2); + assert_eq!(events.len(), 1); + assert_eq!(events[0].text.as_deref(), Some("alpha beta")); + } + + #[tokio::test] + async fn wait_after_empty_targets_returns_immediately() { + let buffer = SubagentOutputBuffer::new(); + let child = SessionId::new(); + buffer + .push(AgentOutputEvent { + sequence: 0, + child_session_id: child, + agent_path: "root/worker".into(), + turn_id: None, + kind: AgentOutputEventKind::Status, + text: None, + status: Some("running".into()), + created_at: Utc::now(), + }) + .await; + + let (events, next_sequence, timed_out) = buffer + .wait_after(0, &[], Duration::from_secs(60), None) + .await; + assert!(timed_out); + assert!(events.is_empty()); + assert_eq!(next_sequence, 2); + } #[test] fn agent_registry_register_and_lookup() { diff --git a/crates/server/tests/deep_research_e2e.rs b/crates/server/tests/deep_research_e2e.rs index 8bdd171c..3479f8c0 100644 --- a/crates/server/tests/deep_research_e2e.rs +++ b/crates/server/tests/deep_research_e2e.rs @@ -1578,7 +1578,7 @@ fn supervisor_stream_events(request: &ModelRequest) -> Vec> .map(str::to_string) }) .unwrap_or_default(); - let mut input = serde_json::json!({ "timeout_ms": 900000 }); + let mut input = serde_json::json!({ "timeout_secs": 120 }); if !target.is_empty() { input["target"] = serde_json::Value::String(target); } diff --git a/crates/server/tests/subagent_lifecycle.rs b/crates/server/tests/subagent_lifecycle.rs index 72069725..c7132c43 100644 --- a/crates/server/tests/subagent_lifecycle.rs +++ b/crates/server/tests/subagent_lifecycle.rs @@ -5,9 +5,10 @@ use anyhow::Context; use anyhow::Result; use devo_protocol::AgentInfo; use devo_protocol::AgentMessageResult; -use devo_protocol::AgentOutputEvent; +use devo_protocol::AgentOutputEventKind; use devo_protocol::ErrorResponse; use devo_protocol::ModelRequest; +use devo_protocol::ParentAgentOutputEvent; use devo_protocol::ProtocolErrorCode; use pretty_assertions::assert_eq; use tempfile::TempDir; @@ -333,32 +334,28 @@ async fn wait_agent_reports_child_output_and_terminal_status() -> Result<()> { spawn_result.child_session_id, ) .await?; - let wait_result = request_agent_wait(&runtime, connection_id, parent_session_id, 1_000).await?; + let wait_result = request_agent_wait(&runtime, connection_id, parent_session_id, 1).await?; assert_eq!(wait_result.timed_out, false); assert_eq!(wait_result.next_sequence, 3); assert_eq!( wait_result.events, vec![ - AgentOutputEvent { + ParentAgentOutputEvent { sequence: 1, - child_session_id: spawn_result.child_session_id, agent_path: spawn_result.agent_path.clone(), - turn_id: wait_result.events[0].turn_id, - kind: "assistant_delta".to_string(), + agent_nickname: spawn_result.agent_nickname.clone(), + kind: AgentOutputEventKind::AssistantMessage, text: Some("child finished review".to_string()), status: None, - created_at: wait_result.events[0].created_at, }, - AgentOutputEvent { + ParentAgentOutputEvent { sequence: 2, - child_session_id: spawn_result.child_session_id, agent_path: spawn_result.agent_path.clone(), - turn_id: wait_result.events[1].turn_id, - kind: "status".to_string(), + agent_nickname: spawn_result.agent_nickname.clone(), + kind: AgentOutputEventKind::Status, text: None, status: Some("completed".to_string()), - created_at: wait_result.events[1].created_at, }, ] ); @@ -399,7 +396,7 @@ async fn wait_agent_polls_incremental_child_output() -> Result<()> { parent_session_id, Some(spawn_result.child_session_id), Some(0), - 1_000, + 1, ) .await?; assert_eq!( @@ -408,7 +405,11 @@ async fn wait_agent_polls_incremental_child_output() -> Result<()> { .iter() .filter_map(|event| event.text.as_deref()) .collect::>(), - vec!["alpha ", "beta"] + vec!["alpha beta"] + ); + assert_eq!( + first_poll.events[0].kind, + AgentOutputEventKind::AssistantMessage ); let second_poll = request_agent_wait_with( @@ -426,10 +427,9 @@ async fn wait_agent_polls_incremental_child_output() -> Result<()> { .iter() .map(|event| event.sequence) .collect::>(), - vec![2, 3] + vec![2] ); - assert_eq!(second_poll.events[0].text.as_deref(), Some("beta")); - assert_eq!(second_poll.events[1].status.as_deref(), Some("completed")); + assert_eq!(second_poll.events[0].status.as_deref(), Some("completed")); Ok(()) } @@ -440,7 +440,7 @@ async fn wait_agent_preserves_full_child_report_for_parent_model() -> Result<()> let full_report = format!("{}END_OF_LONG_SURVEY_REPORT", "survey finding ".repeat(900)); let provider = Arc::new(ScriptedProvider::new([ ScriptedProvider::completed(&full_report), - ScriptedProvider::wait_agent_tool_call(1_000), + ScriptedProvider::wait_agent_tool_call(120), ScriptedProvider::completed("parent consumed child report"), ])); let runtime = build_runtime(data_root.path(), Arc::clone(&provider) as _)?; @@ -656,7 +656,7 @@ async fn close_agent_records_closed_output_event_once() -> Result<()> { ) .await?; assert_eq!(close_result.status, "closed"); - let wait_result = request_agent_wait(&runtime, connection_id, parent_session_id, 1_000).await?; + let wait_result = request_agent_wait(&runtime, connection_id, parent_session_id, 1).await?; assert_eq!( wait_result .events diff --git a/crates/server/tests/support/subagent_lifecycle.rs b/crates/server/tests/support/subagent_lifecycle.rs index bb7f9a9a..703af1f8 100644 --- a/crates/server/tests/support/subagent_lifecycle.rs +++ b/crates/server/tests/support/subagent_lifecycle.rs @@ -121,8 +121,8 @@ impl ScriptedProvider { ]) } - pub fn wait_agent_tool_call(timeout_ms: u64) -> StreamScript { - let input = serde_json::json!({ "timeout_ms": timeout_ms }); + pub fn wait_agent_tool_call(timeout_secs: u64) -> StreamScript { + let input = serde_json::json!({ "timeout_secs": timeout_secs }); let tool_call_id = "wait-agent-call".to_string(); StreamScript::Events(vec![ StreamEvent::ToolCallStart { @@ -407,7 +407,7 @@ pub async fn request_agent_wait( runtime: &Arc, connection_id: u64, session_id: devo_protocol::SessionId, - timeout_ms: u64, + timeout_secs: u64, ) -> Result { request_agent_wait_with( runtime, @@ -415,7 +415,7 @@ pub async fn request_agent_wait( session_id, None::, None, - timeout_ms, + timeout_secs, ) .await } @@ -426,7 +426,7 @@ pub async fn request_agent_wait_with( session_id: devo_protocol::SessionId, target: Option, after_sequence: Option, - timeout_ms: u64, + timeout_secs: u64, ) -> Result { let target = target .map(serde_json::to_value) @@ -442,7 +442,7 @@ pub async fn request_agent_wait_with( "session_id": session_id, "target": target, "after_sequence": after_sequence, - "timeout_ms": timeout_ms + "timeout_secs": timeout_secs } }), ) diff --git a/crates/tui/src/chatwidget/input.rs b/crates/tui/src/chatwidget/input.rs index e7ea8a0d..6b5dc358 100644 --- a/crates/tui/src/chatwidget/input.rs +++ b/crates/tui/src/chatwidget/input.rs @@ -5,6 +5,8 @@ use std::path::Path; +use std::time::Instant; + use crossterm::event::KeyCode; use crossterm::event::KeyEvent; use crossterm::event::KeyEventKind; @@ -213,6 +215,7 @@ impl ChatWidget { pub(crate) fn pre_draw_tick(&mut self) { self.advance_startup_header_animation(); self.run_stream_commit_tick(); + self.tick_subagent_monitor(Instant::now()); self.bottom_pane.pre_draw_tick(); } diff --git a/crates/tui/src/chatwidget/subagent_live_list.rs b/crates/tui/src/chatwidget/subagent_live_list.rs index 5b545bad..5ba171a8 100644 --- a/crates/tui/src/chatwidget/subagent_live_list.rs +++ b/crates/tui/src/chatwidget/subagent_live_list.rs @@ -178,6 +178,7 @@ fn take_suffix_by_width(text: &str, max_width: usize) -> String { fn status_marker_style(status: &str) -> Style { match status.to_ascii_lowercase().as_str() { "idle" => Style::default().fg(COMPLETED_COLOR).bold(), + "working" | "running" | "active_turn" => Style::default().fg(RUNNING_COLOR).bold(), "waiting_client" => Style::default().fg(REASONING_ACCENT_COLOR).bold(), _ => Style::default().fg(RUNNING_COLOR).bold(), } @@ -185,7 +186,7 @@ fn status_marker_style(status: &str) -> Style { fn status_text_style(status: &str) -> Style { match status.to_ascii_lowercase().as_str() { - "running" | "active_turn" => Style::default().fg(RUNNING_COLOR).bold(), + "running" | "active_turn" | "working" => Style::default().fg(RUNNING_COLOR).bold(), "idle" => Style::default().fg(COMPLETED_COLOR).bold(), "waiting_client" => Style::default().fg(REASONING_ACCENT_COLOR).bold(), _ => Style::default().fg(MUTED_COLOR), diff --git a/crates/tui/src/chatwidget/subagent_monitor.rs b/crates/tui/src/chatwidget/subagent_monitor.rs index ddf6c8b4..f629ec75 100644 --- a/crates/tui/src/chatwidget/subagent_monitor.rs +++ b/crates/tui/src/chatwidget/subagent_monitor.rs @@ -5,6 +5,8 @@ //! the selected child through the normal transcript pager. use std::collections::HashMap; +use std::time::Duration; +use std::time::Instant; use crossterm::event::KeyCode; use crossterm::event::KeyEvent; @@ -31,13 +33,18 @@ use super::ActiveCellTranscriptKey; use super::ChatWidget; use super::DotStatus; use super::TranscriptOverlayCell; +use super::UserMessage; use super::subagent_live_list; use super::subagent_live_list::SubagentLiveListRow; use super::subagent_live_list::SubagentLiveListRowKey; +const SUBAGENT_INACTIVE_GRACE: Duration = Duration::from_secs(12); +const SUBAGENT_CTRL_X_REVEAL: Duration = Duration::from_secs(15); + #[derive(Debug, Default)] pub(super) struct SubagentMonitorState { live_list_focused: bool, + list_reveal_until: Option, agents: Vec, selected: Option, user_selected: bool, @@ -54,6 +61,7 @@ struct SubagentSessionView { active_turn: Option, latest_preview: String, has_runtime_update: bool, + last_activity_at: Option, revision: u64, } @@ -73,6 +81,7 @@ struct MonitorToolItem { #[derive(Clone, Copy, Debug)] enum MonitorTranscriptKind { + User, Assistant, Reasoning, Tool, @@ -94,19 +103,53 @@ impl ChatWidget { } pub(super) fn focus_subagent_live_list(&mut self) { - if !self.has_live_subagents() { + if self.subagent_monitor.agents.is_empty() { self.subagent_monitor.live_list_focused = false; self.set_status_message("No active sub-agents"); self.frame_requester.schedule_frame(); return; } + let now = Instant::now(); + if !self.has_visible_subagent_list(now) { + self.subagent_monitor.list_reveal_until = Some(now + SUBAGENT_CTRL_X_REVEAL); + self.frame_requester + .schedule_frame_in(SUBAGENT_CTRL_X_REVEAL); + } + self.subagent_monitor.live_list_focused = true; - self.ensure_live_subagent_selected(); + self.ensure_visible_subagent_selected(now); self.set_status_message("Select sub-agent"); self.frame_requester.schedule_frame(); } + pub(super) fn tick_subagent_monitor(&mut self, now: Instant) { + let reveal_expired = self + .subagent_monitor + .list_reveal_until + .is_some_and(|until| until <= now); + if reveal_expired { + self.subagent_monitor.list_reveal_until = None; + self.sync_subagent_hint(now); + self.frame_requester.schedule_frame(); + return; + } + + let had_visible = self.has_visible_subagent_list(now); + self.sync_subagent_hint(now); + if had_visible + && !self.has_visible_subagent_list(now) + && self.subagent_monitor.agents.iter().any(|agent| { + self.subagent_monitor + .sessions + .contains_key(&agent.session_id) + }) + { + self.frame_requester + .schedule_frame_in(SUBAGENT_INACTIVE_GRACE); + } + } + pub(super) fn handle_subagent_live_list_key_event(&mut self, key: KeyEvent) { if !matches!(key.kind, KeyEventKind::Press | KeyEventKind::Repeat) { return; @@ -184,7 +227,7 @@ impl ChatWidget { "subagent discovered by chat widget" ); self.upsert_subagent(agent); - self.sync_subagent_hint(); + self.sync_subagent_hint(Instant::now()); self.frame_requester.schedule_frame(); } @@ -226,7 +269,7 @@ impl ChatWidget { view.agent = Some(updated_agent); } - self.sync_subagent_hint(); + self.sync_subagent_hint(Instant::now()); let view = self.subagent_monitor.sessions.get(&session_id); let latest_preview = view.map(|view| view.latest_preview.as_str()).unwrap_or(""); let has_runtime_update = view.is_some_and(|view| view.has_runtime_update); @@ -313,6 +356,18 @@ impl ChatWidget { .then(|| self.subagent_live_tail_lines(view, width)) } + #[cfg(test)] + pub(crate) fn expire_subagent_inactivity_for_test(&mut self) { + let expired = Instant::now() + .checked_sub(SUBAGENT_INACTIVE_GRACE + Duration::from_secs(1)) + .unwrap_or_else(Instant::now); + for view in self.subagent_monitor.sessions.values_mut() { + view.last_activity_at = Some(expired); + } + self.subagent_monitor.list_reveal_until = None; + self.sync_subagent_hint(Instant::now()); + } + #[cfg(test)] pub(crate) fn is_subagent_monitor_open_for_test(&self) -> bool { self.subagent_monitor.live_list_focused @@ -347,23 +402,35 @@ impl ChatWidget { .entry(session_id) .or_default(); view.status = agent.status.clone(); - view.agent = Some(agent); + view.agent = Some(agent.clone()); + if is_active_subagent_status(&normalize_subagent_display_status(&agent.status)) { + view.touch_activity(); + if !is_terminal_status(&view.status) { + view.status = "working".to_string(); + } + } + if let Some(message) = agent.last_task_message.as_deref() { + view.seed_initial_task_message(message); + } } - fn sync_subagent_hint(&mut self) { - let has_live = self.has_live_subagents(); - if !has_live { + fn sync_subagent_hint(&mut self, now: Instant) { + let has_active = self.has_active_subagents(); + let has_visible = self.has_visible_subagent_list(now); + if !has_visible { self.subagent_monitor.live_list_focused = false; self.subagent_monitor.selected = None; self.subagent_monitor.user_selected = false; } else { - self.ensure_live_subagent_selected(); + self.ensure_visible_subagent_selected(now); } - self.bottom_pane.set_subagent_hint_visible(has_live); + self.bottom_pane.set_subagent_hint_visible(has_active); tracing::debug!( target: "devo_tui::subagent", - has_live, - live_count = self.live_subagent_ids().len(), + has_active, + has_visible, + live_count = self.active_subagent_ids().len(), + visible_count = self.visible_subagent_ids(now).len(), selected = ?self.subagent_monitor.selected, focused = self.subagent_monitor.live_list_focused, user_selected = self.subagent_monitor.user_selected, @@ -371,40 +438,86 @@ impl ChatWidget { ); } - fn has_live_subagents(&self) -> bool { + fn has_active_subagents(&self) -> bool { + self.subagent_monitor.agents.iter().any(|agent| { + is_active_subagent_status(&self.subagent_display_status(agent)) + || self + .subagent_monitor + .sessions + .get(&agent.session_id) + .is_some_and(SubagentSessionView::has_live_tail) + }) + } + + fn has_visible_subagent_list(&self, now: Instant) -> bool { + if self + .subagent_monitor + .list_reveal_until + .is_some_and(|until| until > now) + { + return !self.subagent_monitor.agents.is_empty(); + } self.subagent_monitor .agents .iter() - .any(|agent| is_live_status(&self.subagent_status_for_agent(agent))) + .any(|agent| self.is_agent_visible_in_list(agent, now)) + } + + fn is_agent_visible_in_list(&self, agent: &SubagentMonitorAgent, now: Instant) -> bool { + if is_active_subagent_status(&self.subagent_display_status(agent)) { + return true; + } + let Some(view) = self.subagent_monitor.sessions.get(&agent.session_id) else { + return false; + }; + if view.has_live_tail() { + return true; + } + view.last_activity_at + .is_some_and(|at| now.duration_since(at) < SUBAGENT_INACTIVE_GRACE) + } + + fn has_live_subagents(&self) -> bool { + self.has_active_subagents() } fn selected_live_subagent(&self) -> Option { let selected = self.subagent_monitor.selected?; + let now = Instant::now(); self.subagent_agent(selected) - .filter(|agent| is_live_status(&self.subagent_status_for_agent(agent))) + .filter(|agent| self.is_agent_visible_in_list(agent, now)) .map(|agent| agent.session_id) } - fn ensure_live_subagent_selected(&mut self) { + fn ensure_visible_subagent_selected(&mut self, now: Instant) { if self.selected_live_subagent().is_some() && (self.subagent_monitor.live_list_focused || self.subagent_monitor.user_selected) { return; } - self.subagent_monitor.selected = self.live_subagent_ids().last().copied(); + self.subagent_monitor.selected = self + .visible_subagent_ids(now) + .last() + .copied() + .or_else(|| self.active_subagent_ids().last().copied()); self.subagent_monitor.user_selected = false; } + fn ensure_live_subagent_selected(&mut self) { + self.ensure_visible_subagent_selected(Instant::now()); + } + fn select_relative_live_subagent(&mut self, delta: isize) { - let live_ids = self.live_subagent_ids(); - if live_ids.is_empty() { + let now = Instant::now(); + let visible_ids = self.visible_subagent_ids(now); + if visible_ids.is_empty() { return; } let current = self .subagent_monitor .selected .and_then(|selected| { - live_ids + visible_ids .iter() .position(|session_id| *session_id == selected) }) @@ -414,29 +527,73 @@ impl ChatWidget { } else { current .saturating_add(delta as usize) - .min(live_ids.len().saturating_sub(1)) + .min(visible_ids.len().saturating_sub(1)) }; - self.subagent_monitor.selected = Some(live_ids[next]); + self.subagent_monitor.selected = Some(visible_ids[next]); self.subagent_monitor.user_selected = true; self.frame_requester.schedule_frame(); } - fn live_subagent_ids(&self) -> Vec { + fn active_subagent_ids(&self) -> Vec { self.subagent_monitor .agents .iter() - .filter(|agent| is_live_status(&self.subagent_status_for_agent(agent))) + .filter(|agent| { + is_active_subagent_status(&self.subagent_display_status(agent)) + || self + .subagent_monitor + .sessions + .get(&agent.session_id) + .is_some_and(SubagentSessionView::has_live_tail) + }) .map(|agent| agent.session_id) .collect() } - fn subagent_live_list_rows(&self) -> Vec { - let mut rows = self + fn visible_subagent_ids(&self, now: Instant) -> Vec { + let mut indexed = self .subagent_monitor .agents .iter() - .filter(|agent| is_live_status(&self.subagent_status_for_agent(agent))) - .map(|agent| { + .enumerate() + .filter(|(_, agent)| self.is_agent_visible_in_list(agent, now)) + .collect::>(); + indexed.sort_by(|(left_index, left), (right_index, right)| { + let left_active = is_active_subagent_status(&self.subagent_display_status(left)) + || self + .subagent_monitor + .sessions + .get(&left.session_id) + .is_some_and(SubagentSessionView::has_live_tail); + let right_active = is_active_subagent_status(&self.subagent_display_status(right)) + || self + .subagent_monitor + .sessions + .get(&right.session_id) + .is_some_and(SubagentSessionView::has_live_tail); + match (left_active, right_active) { + (true, false) => std::cmp::Ordering::Less, + (false, true) => std::cmp::Ordering::Greater, + _ => left_index.cmp(right_index), + } + }); + indexed + .into_iter() + .map(|(_, agent)| agent.session_id) + .collect() + } + + fn live_subagent_ids(&self) -> Vec { + self.active_subagent_ids() + } + + fn subagent_live_list_rows(&self) -> Vec { + let now = Instant::now(); + let mut rows = self + .visible_subagent_ids(now) + .into_iter() + .filter_map(|session_id| { + let agent = self.subagent_agent(session_id)?; let view = self.subagent_monitor.sessions.get(&agent.session_id); let preview = if let Some(view) = view && view.has_runtime_update @@ -450,12 +607,12 @@ impl ChatWidget { .and_then(tail_preview) .unwrap_or_else(|| "Waiting for updates".to_string()) }; - SubagentLiveListRow { + Some(SubagentLiveListRow { key: SubagentLiveListRowKey::Session(agent.session_id), name: agent.nickname.clone(), - status: self.subagent_status_for_agent(agent), + status: self.subagent_display_status(agent), preview, - } + }) }) .collect::>(); rows.extend( @@ -487,12 +644,24 @@ impl ChatWidget { .unwrap_or_else(|| agent.status.clone()) } + fn subagent_display_status(&self, agent: &SubagentMonitorAgent) -> String { + let stored = self.subagent_status_for_agent(agent); + let Some(view) = self.subagent_monitor.sessions.get(&agent.session_id) else { + return normalize_subagent_display_status(&stored); + }; + if view.has_live_tail() || view.active_turn.is_some() { + return "working".to_string(); + } + normalize_subagent_display_status(&stored) + } + fn subagent_transcript_item_cell( &self, item: &MonitorTranscriptItem, width: u16, ) -> TranscriptOverlayCell { let lines = match item.kind { + MonitorTranscriptKind::User => Vec::new(), MonitorTranscriptKind::Assistant => history_cell::AgentMarkdownCell::new( item.body.clone(), &self.session.cwd, @@ -542,7 +711,8 @@ impl ChatWidget { TranscriptOverlayCell { lines, is_stream_continuation: false, - user_message: None, + user_message: matches!(item.kind, MonitorTranscriptKind::User) + .then(|| UserMessage::from(item.body.clone())), is_selected_user: false, } } @@ -617,7 +787,7 @@ impl SubagentSessionView { session_id: _, turn_id, } => { - self.status = "running".to_string(); + self.mark_working(); self.active_turn = Some(turn_id); self.set_latest_preview("Started turn"); } @@ -626,6 +796,7 @@ impl SubagentSessionView { item_id, kind, } => { + self.mark_working(); self.active_text.insert( text_key(Some(item_id), kind), MonitorTextItem { @@ -642,6 +813,7 @@ impl SubagentSessionView { kind, delta, } => { + self.mark_working(); let latest_preview = { let latest = self .active_text @@ -683,6 +855,7 @@ impl SubagentSessionView { tool_use_id, summary, } => { + self.mark_working(); let latest_preview = format!("Running {summary}"); self.active_tools .entry(tool_use_id) @@ -699,6 +872,7 @@ impl SubagentSessionView { tool_use_id, delta, } => { + self.mark_working(); let latest_preview = { let tool = self .active_tools @@ -771,6 +945,7 @@ impl SubagentSessionView { session_id: _, status, } => { + self.touch_activity(); self.status = status.clone(); self.active_turn = None; self.flush_active_items(); @@ -786,6 +961,7 @@ impl SubagentSessionView { session_id: _, message, } => { + self.touch_activity(); self.status = "failed".to_string(); self.active_turn = None; self.flush_active_items(); @@ -797,18 +973,75 @@ impl SubagentSessionView { is_error: true, }); } + SubagentMonitorEvent::TaskMessage { + session_id: _, + message, + } => { + self.mark_working(); + self.append_task_message(message); + } SubagentMonitorEvent::SessionStatusChanged { session_id: _, status, } => { if !is_terminal_status(&self.status) { - self.status = format!("{status:?}").to_lowercase(); + self.status = normalize_runtime_status_for_subagent(status); self.set_latest_preview(format!("Status {}", self.status)); } } } } + fn touch_activity(&mut self) { + self.last_activity_at = Some(Instant::now()); + } + + fn mark_working(&mut self) { + self.touch_activity(); + if !is_terminal_status(&self.status) { + self.status = "working".to_string(); + } + } + + fn seed_initial_task_message(&mut self, message: &str) { + if message.trim().is_empty() + || self + .transcript + .iter() + .any(|item| matches!(item.kind, MonitorTranscriptKind::User)) + { + return; + } + self.transcript.insert( + 0, + MonitorTranscriptItem { + kind: MonitorTranscriptKind::User, + title: String::new(), + body: message.to_string(), + is_error: false, + }, + ); + self.touch_activity(); + } + + fn append_task_message(&mut self, message: String) { + if message.trim().is_empty() { + return; + } + if self.transcript.last().is_some_and(|item| { + matches!(item.kind, MonitorTranscriptKind::User) && item.body == message + }) { + return; + } + self.transcript.push(MonitorTranscriptItem { + kind: MonitorTranscriptKind::User, + title: String::new(), + body: message, + is_error: false, + }); + self.touch_activity(); + } + fn has_live_tail(&self) -> bool { !self.active_text.is_empty() || !self.active_tools.is_empty() } @@ -865,6 +1098,7 @@ impl SubagentMonitorEvent { | Self::PlanUpdated { session_id, .. } | Self::TurnFinished { session_id, .. } | Self::TurnFailed { session_id, .. } + | Self::TaskMessage { session_id, .. } | Self::SessionStatusChanged { session_id, .. } => *session_id, } } @@ -883,12 +1117,31 @@ fn subagent_monitor_event_kind(event: &SubagentMonitorEvent) -> &'static str { SubagentMonitorEvent::PlanUpdated { .. } => "plan_updated", SubagentMonitorEvent::TurnFinished { .. } => "turn_finished", SubagentMonitorEvent::TurnFailed { .. } => "turn_failed", + SubagentMonitorEvent::TaskMessage { .. } => "task_message", SubagentMonitorEvent::SessionStatusChanged { .. } => "session_status_changed", } } -fn is_live_status(status: &str) -> bool { - !is_terminal_status(status) +fn is_active_subagent_status(status: &str) -> bool { + matches!( + status.to_ascii_lowercase().as_str(), + "working" | "running" | "spawning" | "activeturn" + ) +} + +fn normalize_subagent_display_status(status: &str) -> String { + match status.to_ascii_lowercase().as_str() { + "running" | "spawning" | "activeturn" => "working".to_string(), + other => other.to_string(), + } +} + +fn normalize_runtime_status_for_subagent(status: devo_protocol::SessionRuntimeStatus) -> String { + match status { + devo_protocol::SessionRuntimeStatus::ActiveTurn => "working".to_string(), + devo_protocol::SessionRuntimeStatus::Idle => "idle".to_string(), + other => format!("{other:?}").to_lowercase(), + } } fn is_terminal_status(status: &str) -> bool { diff --git a/crates/tui/src/chatwidget_tests.rs b/crates/tui/src/chatwidget_tests.rs index 2dbf773c..556eb24d 100644 --- a/crates/tui/src/chatwidget_tests.rs +++ b/crates/tui/src/chatwidget_tests.rs @@ -5581,14 +5581,14 @@ fn subagent_discovery_shows_inline_live_list_without_focusing_it() { assert_eq!(widget.selected_subagent_for_test(), Some(child)); let rows = rendered_rows(&widget, 160, 18).join("\n"); assert!(rows.contains("ctrl + x agents"), "rows:\n{rows}"); - assert!(rows.contains("reviewer: running"), "rows:\n{rows}"); + assert!(rows.contains("reviewer: working"), "rows:\n{rows}"); assert!(rows.contains("checking files"), "rows:\n{rows}"); let rendered = rendered_rows(&widget, 160, 18); let live_prefix = " ".repeat(usize::from(LIVE_PREFIX_COLS)); assert!( rendered .iter() - .any(|row| row.starts_with(&format!("{live_prefix}● reviewer: running"))), + .any(|row| row.starts_with(&format!("{live_prefix}● reviewer: working"))), "live-list title should use the shared live prefix only:\n{}", rendered.join("\n") ); @@ -5679,7 +5679,7 @@ fn request_user_input_and_subagent_live_list_are_visible_together() { }); let rows = rendered_rows(&widget, 120, 28).join("\n"); - assert!(rows.contains("researcher: running"), "rows:\n{rows}"); + assert!(rows.contains("researcher: working"), "rows:\n{rows}"); assert!(rows.contains("Working"), "rows:\n{rows}"); assert!( rows.contains("Which scope should research use?"), @@ -5711,8 +5711,8 @@ fn ctrl_x_focuses_inline_live_list_and_ctrl_x_esc_or_q_exits() { assert_eq!(widget.selected_subagent_for_test(), Some(second)); let rows = rendered_rows(&widget, 160, 18).join("\n"); assert!(!rows.contains("Sub-agents"), "rows:\n{rows}"); - assert!(rows.contains("first: running"), "rows:\n{rows}"); - assert!(rows.contains("second: running"), "rows:\n{rows}"); + assert!(rows.contains("first: working"), "rows:\n{rows}"); + assert!(rows.contains("second: working"), "rows:\n{rows}"); assert!(rows.contains("run first"), "rows:\n{rows}"); assert!(!rows.contains("root/second"), "rows:\n{rows}"); @@ -5807,6 +5807,10 @@ fn terminal_subagent_status_hides_ctrl_x_hint_when_no_live_children_remain() { assert!(!widget.has_live_subagents_for_test()); assert!(!widget.is_subagent_monitor_open_for_test()); let rows = rendered_rows(&widget, 100, 18).join("\n"); + assert!(rows.contains("builder: completed"), "rows:\n{rows}"); + widget.expire_subagent_inactivity_for_test(); + let rows = rendered_rows(&widget, 100, 18).join("\n"); + assert!(!rows.contains("builder: completed"), "rows:\n{rows}"); assert!(!rows.contains("ctrl + x agents"), "rows:\n{rows}"); } @@ -5833,6 +5837,9 @@ fn terminal_cancelled_subagent_disappears_from_live_list() { assert!(!widget.has_live_subagents_for_test()); let rows = rendered_rows(&widget, 100, 18).join("\n"); + assert!(rows.contains("builder: cancelled"), "rows:\n{rows}"); + widget.expire_subagent_inactivity_for_test(); + let rows = rendered_rows(&widget, 100, 18).join("\n"); assert!(!rows.contains("builder: cancelled"), "rows:\n{rows}"); assert!(!rows.contains("ctrl + x agents"), "rows:\n{rows}"); } @@ -6084,6 +6091,33 @@ fn subagent_live_list_preview_updates_from_completed_reasoning_tool_and_plan_tai assert!(!rows.contains("old plan note"), "rows:\n{rows}"); } +#[test] +fn subagent_transcript_overlay_includes_spawn_task_message() { + let model = Model { + slug: "test-model".to_string(), + display_name: "Test Model".to_string(), + ..Model::default() + }; + let (mut widget, _app_event_rx) = widget_with_model(model, PathBuf::from(".")); + let parent = SessionId::new(); + let child = SessionId::new(); + + widget.handle_worker_event(crate::events::WorkerEvent::SubagentDiscovered { + agent: monitor_agent(child, parent, "builder"), + }); + + let cells = widget + .subagent_transcript_overlay_cells(child, 80) + .expect("overlay cells"); + assert!( + cells + .first() + .and_then(|cell| cell.user_message.as_ref()) + .is_some_and(|message| { message.text.contains("run builder") }), + "expected spawn task message at top of subagent overlay" + ); +} + #[test] fn subagent_live_list_enter_emits_overlay_request_for_selected_child() { let model = Model { diff --git a/crates/tui/src/events.rs b/crates/tui/src/events.rs index 9ca06b7f..99462e4b 100644 --- a/crates/tui/src/events.rs +++ b/crates/tui/src/events.rs @@ -121,6 +121,10 @@ pub(crate) enum SubagentMonitorEvent { session_id: SessionId, message: String, }, + TaskMessage { + session_id: SessionId, + message: String, + }, SessionStatusChanged { session_id: SessionId, status: SessionRuntimeStatus, diff --git a/crates/tui/src/worker.rs b/crates/tui/src/worker.rs index 86d0c96b..14bd2f05 100644 --- a/crates/tui/src/worker.rs +++ b/crates/tui/src/worker.rs @@ -104,6 +104,7 @@ use acp_events::acp_terminal_output_event_with_session; use acp_events::parse_acp_session_notification; use acp_events::session_metadata_from_acp_update; use acp_events::spawn_agent_result_from_acp_update; +use acp_events::spawn_task_message_from_acp_update; use acp_events::subagent_monitor_events_from_acp_session_notification_with_terminal_state; #[cfg(test)] use acp_events::worker_events_from_acp_notification; @@ -157,7 +158,8 @@ async fn maybe_discover_spawned_subagent_from_acp_update( }; let child_session_id = spawn_result.child_session_id; if child_agent_sessions.contains(&child_session_id) { - return; + // Child may already be registered from session_info_update; still hydrate + // status and last_task_message from agent/list when spawn completes. } let listed_agent = match client @@ -190,11 +192,10 @@ async fn maybe_discover_spawned_subagent_from_acp_update( nickname: spawn_result.agent_nickname, role: "default".to_string(), status: spawn_result.status, - last_task_message: None, + last_task_message: spawn_task_message_from_acp_update(update), }); - if child_agent_sessions.insert(agent.session_id) { - let _ = event_tx.send(WorkerEvent::SubagentDiscovered { agent }); - } + child_agent_sessions.insert(agent.session_id); + let _ = event_tx.send(WorkerEvent::SubagentDiscovered { agent }); } /// Immutable runtime configuration used to construct the background server client worker. diff --git a/crates/tui/src/worker/acp_events.rs b/crates/tui/src/worker/acp_events.rs index bc65de43..22b7dee6 100644 --- a/crates/tui/src/worker/acp_events.rs +++ b/crates/tui/src/worker/acp_events.rs @@ -319,6 +319,18 @@ pub(super) fn session_metadata_from_acp_update( serde_json::from_value(meta.get(DEVO_SESSION_META)?.clone()).ok() } +pub(super) fn spawn_task_message_from_acp_update(update: &AcpSessionUpdate) -> Option { + let raw_input = match update { + AcpSessionUpdate::ToolCall { raw_input, .. } => raw_input.as_ref(), + AcpSessionUpdate::ToolCallUpdate { raw_input, .. } => raw_input.as_ref(), + _ => None, + }?; + raw_input + .get("message") + .and_then(serde_json::Value::as_str) + .map(str::to_string) +} + pub(super) fn spawn_agent_result_from_acp_update( update: &AcpSessionUpdate, ) -> Option { @@ -450,8 +462,16 @@ pub(super) fn subagent_monitor_events_from_acp_session_notification_with_termina owner_session_id: Some(session_id), }, ), - AcpSessionUpdate::UserMessageChunk { .. } - | AcpSessionUpdate::SessionInfoUpdate { .. } + AcpSessionUpdate::UserMessageChunk { content, .. } => acp_content_display_text(&content) + .into_iter() + .map(|message| WorkerEvent::SubagentMonitor { + event: SubagentMonitorEvent::TaskMessage { + session_id, + message, + }, + }) + .collect(), + AcpSessionUpdate::SessionInfoUpdate { .. } | AcpSessionUpdate::AvailableCommandsUpdate { .. } | AcpSessionUpdate::CurrentModeUpdate { .. } | AcpSessionUpdate::ConfigOptionUpdate { .. } diff --git a/specs/L2/agent/L2-DES-AGENT-003-subagent-architecture.md b/specs/L2/agent/L2-DES-AGENT-003-subagent-architecture.md index 1bd9e688..d8a56e32 100644 --- a/specs/L2/agent/L2-DES-AGENT-003-subagent-architecture.md +++ b/specs/L2/agent/L2-DES-AGENT-003-subagent-architecture.md @@ -65,7 +65,7 @@ Parent agents need to send additional input to child agents without treating the The parent must be able to monitor child progress and completion without receiving child-authored mailbox messages. -**Decision**: Each parent session has a sequence-numbered output buffer for direct child assistant text and terminal status events. Child assistant deltas are appended as they stream. Child terminal status changes are appended as status events. The `wait_agent` tool polls this buffer with an optional target and sequence cursor. +**Decision**: Each parent session has a sequence-numbered output buffer for direct child assistant text and terminal status events. Child assistant streaming deltas for the same turn are accumulated into a single `assistant_message` event before the parent polls them. Child terminal status changes are appended as status events. The `wait_agent` tool polls this buffer with an optional target and sequence cursor. When `after_sequence` is omitted, the runtime fills it from a per-parent, per-target cursor so repeated polls do not re-deliver consumed events. ### DD-6: Subagents inherit permission and safety boundaries, never bypass them @@ -264,21 +264,24 @@ Sender Agent Mailbox Recipient Agent The parent does not receive child-authored mailbox messages. Instead, each parent session has an output buffer: ```rust -struct AgentOutputEvent { +enum AgentOutputEventKind { + AssistantMessage, // full accumulated assistant text for one child turn + Status, // terminal turn status (completed, failed, ...) +} + +struct ParentAgentOutputEvent { sequence: u64, - child_session_id: SessionId, agent_path: String, - turn_id: Option, - kind: String, // "assistant_delta" or "status" + agent_nickname: String, + kind: AgentOutputEventKind, text: Option, status: Option, - created_at: DateTime, } ``` -`wait_agent` reads events after an optional `after_sequence` cursor. If matching events already exist, it returns immediately. Otherwise it waits with a deadline and returns either new events or `timed_out = true`. +`wait_agent` reads events after an optional `after_sequence` cursor. When `after_sequence` is omitted, the runtime substitutes the stored cursor for the target key. If matching events already exist, it returns immediately. Otherwise it waits with a deadline and returns either new events or `timed_out = true`. -Timeout bounds are configurable per session (`min_wait_timeout_ms`, `max_wait_timeout_ms`, `default_wait_timeout_ms`). +Default and maximum wait timeouts are server constants (`default_wait_timeout_secs = 5`, `max_wait_timeout_secs = 120`). Callers pass `timeout_secs`. ### Agent Status Lifecycle @@ -345,11 +348,11 @@ Polls child assistant output and terminal status events, optionally waiting for | Parameter | Required | Type | Description | |-----------|----------|------|-------------| -| `target` | No | string | Optional child agent path or session id | -| `after_sequence` | No | integer | Only return events after this parent-buffer sequence | -| `timeout_ms` | No | integer | Wait timeout in milliseconds (clamped to `[min_wait_timeout_ms, max_wait_timeout_ms]`) | +| `target` | No | string | Optional child agent path or nickname | +| `after_sequence` | No | integer | Only return events after this parent-buffer sequence; omitted values use the runtime per-target cursor | +| `timeout_secs` | No | integer | Wait timeout in seconds (default 5, max 120). Returns immediately when matching events already exist. | -**Output**: `{ "events": AgentOutputEvent[], "next_sequence": integer, "timed_out": bool }`. +**Output**: `{ "events": ParentAgentOutputEvent[], "next_sequence": integer, "timed_out": bool }`. Model-facing events omit internal session and turn ids; address children by `agent_path` or nickname. **Behavior**: If matching output events after `after_sequence` already exist, returns immediately. Otherwise waits until a matching event arrives or the timeout expires. From 36105d5a431b118202d3ea095c46ee1dd971e18f Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Wed, 1 Jul 2026 04:26:29 -1000 Subject: [PATCH 02/16] chore: try to fix /research supervisor stuck but still in working --- crates/core/src/tools/handlers/webfetch.rs | 75 ++++---- crates/protocol/src/acp_event_to_update.rs | 37 +++- crates/server/src/runtime/agents.rs | 6 +- crates/server/src/runtime/agents/lifecycle.rs | 3 +- crates/server/src/runtime/items.rs | 6 +- crates/server/src/runtime/research.rs | 3 +- crates/server/src/runtime/subagent_usage.rs | 9 +- crates/server/src/runtime/turn_exec.rs | 59 +----- crates/server/src/transport.rs | 90 ++++----- crates/tui/src/chatwidget/text_stream.rs | 10 +- crates/tui/src/chatwidget/transcript_view.rs | 175 +++++++++++------- 11 files changed, 254 insertions(+), 219 deletions(-) diff --git a/crates/core/src/tools/handlers/webfetch.rs b/crates/core/src/tools/handlers/webfetch.rs index 5a7e26f7..de68cd71 100644 --- a/crates/core/src/tools/handlers/webfetch.rs +++ b/crates/core/src/tools/handlers/webfetch.rs @@ -122,10 +122,44 @@ impl ToolHandler for WebFetchHandler { .header(reqwest::header::ACCEPT, accept) .header(reqwest::header::ACCEPT_LANGUAGE, "en-US,en;q=0.9"); - let response = timeout(Duration::from_millis(timeout_ms), request.send()).await; - let response = match response { - Ok(result) => result - .map_err(|e| ToolCallError::ExecutionFailed(format!("Request failed: {e}")))?, + let fetch_result = timeout(Duration::from_millis(timeout_ms), async { + let response = request + .send() + .await + .map_err(|e| ToolCallError::ExecutionFailed(format!("Request failed: {e}")))?; + if !response.status().is_success() { + let msg = format!("Request failed with status code: {}", response.status()); + return Err(ToolCallError::ExecutionFailed(msg)); + } + if response + .content_length() + .is_some_and(|len| len as usize > MAX_RESPONSE_SIZE) + { + return Err(ToolCallError::ExecutionFailed("response too large".into())); + } + let content_type = response + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .unwrap_or("") + .to_string(); + let bytes = response.bytes().await.map_err(|e| { + ToolCallError::ExecutionFailed(format!("Failed to read response: {e}")) + })?; + Ok((bytes, content_type)) + }) + .await; + + let (bytes, content_type) = match fetch_result { + Ok(Ok(payload)) => payload, + Ok(Err(error)) => { + let message = error.to_string(); + return Ok(ToolResult::error( + ToolResultContent::Text(message.clone()), + "HTTP error", + error, + )); + } Err(_) => { return Ok(ToolResult::error( ToolResultContent::Text("Request timed out".into()), @@ -135,19 +169,7 @@ impl ToolHandler for WebFetchHandler { } }; - if !response.status().is_success() { - let msg = format!("Request failed with status code: {}", response.status()); - return Ok(ToolResult::error( - ToolResultContent::Text(msg.clone()), - "HTTP error", - ToolCallError::ExecutionFailed(msg), - )); - } - - if response - .content_length() - .is_some_and(|len| len as usize > MAX_RESPONSE_SIZE) - { + if bytes.len() > MAX_RESPONSE_SIZE { return Ok(ToolResult::error( ToolResultContent::Text("Response too large (exceeds 5MB limit)".into()), "Response too large", @@ -155,12 +177,6 @@ impl ToolHandler for WebFetchHandler { )); } - let content_type = response - .headers() - .get(reqwest::header::CONTENT_TYPE) - .and_then(|value| value.to_str().ok()) - .unwrap_or("") - .to_string(); let mime = content_type .split(';') .next() @@ -169,19 +185,6 @@ impl ToolHandler for WebFetchHandler { .to_lowercase(); let title = format!("{url} ({content_type})"); - let bytes = response - .bytes() - .await - .map_err(|e| ToolCallError::ExecutionFailed(format!("Failed to read response: {e}")))?; - - if bytes.len() > MAX_RESPONSE_SIZE { - return Ok(ToolResult::error( - ToolResultContent::Text("Response too large (exceeds 5MB limit)".into()), - "Response too large", - ToolCallError::ExecutionFailed("response too large".into()), - )); - } - if is_image_mime(&mime) { return Ok(ToolResult::success( ToolResultContent::Mixed { diff --git a/crates/protocol/src/acp_event_to_update.rs b/crates/protocol/src/acp_event_to_update.rs index 6340cbdc..54cf4aec 100644 --- a/crates/protocol/src/acp_event_to_update.rs +++ b/crates/protocol/src/acp_event_to_update.rs @@ -358,12 +358,40 @@ fn acp_plan_entry_from_turn_plan_step(step: &TurnPlanStepPayload) -> AcpPlanEntr } fn tool_title(tool_name: &str, parameters: &serde_json::Value) -> String { - parameters + if let Some(command) = parameters .get("command") .or_else(|| parameters.get("cmd")) .and_then(serde_json::Value::as_str) - .map(str::to_string) - .unwrap_or_else(|| tool_name.to_string()) + { + return command.to_string(); + } + if matches!( + tool_name, + "webfetch" | "web_fetch" | "web-fetch" | "fetch_url" | "fetch-url" + ) && let Some(url) = parameters + .get("url") + .and_then(serde_json::Value::as_str) + .filter(|url| !url.is_empty()) + { + return format!("web_fetch: {url}"); + } + if matches!(tool_name, "web_search" | "websearch" | "web-search") + && let Some(query) = parameters + .get("query") + .and_then(serde_json::Value::as_str) + .filter(|query| !query.is_empty()) + { + return format!("web_search: {query}"); + } + if tool_name == "spawn_agent" + && let Some(message) = parameters + .get("message") + .and_then(serde_json::Value::as_str) + .filter(|message| !message.is_empty()) + { + return format!("spawn_agent: {message}"); + } + tool_name.to_string() } fn tool_kind_from_name(tool_name: &str) -> AcpToolKind { @@ -371,7 +399,8 @@ fn tool_kind_from_name(tool_name: &str) -> AcpToolKind { "read" | "grep" | "glob" | "lsp" => AcpToolKind::Read, "apply_patch" | "edit" | "write" => AcpToolKind::Edit, "bash" | "shell_command" | "exec_command" => AcpToolKind::Execute, - "web_search" | "websearch" | "web_fetch" | "websearch_query" => AcpToolKind::Fetch, + "web_search" | "websearch" | "web_fetch" | "webfetch" | "web-fetch" | "fetch_url" + | "fetch-url" | "websearch_query" => AcpToolKind::Fetch, "agent" => AcpToolKind::Think, _ => AcpToolKind::Other, } diff --git a/crates/server/src/runtime/agents.rs b/crates/server/src/runtime/agents.rs index 8ef61df4..9da543e5 100644 --- a/crates/server/src/runtime/agents.rs +++ b/crates/server/src/runtime/agents.rs @@ -93,6 +93,8 @@ impl ServerRuntime { parent_tool_registry, runtime_context, ) = parent_snapshot; + let parent_usage_turn_id = + parent_active_turn_id.or_else(|| parent_latest_turn.as_ref().map(|turn| turn.turn_id)); let nickname = self .generate_unique_agent_name(parent_session_id, child_session_id) @@ -273,7 +275,7 @@ impl ServerRuntime { }, ) .await; - if let Some(parent_turn_id) = parent_active_turn_id { + if let Some(parent_turn_id) = parent_usage_turn_id { self.record_subagent_status_event( parent_session_id, child_session_id, @@ -285,7 +287,7 @@ impl ServerRuntime { self.register_subagent_usage_owner( parent_session_id, child_session_id, - parent_active_turn_id, + parent_usage_turn_id, ) .await; if !summary.ephemeral diff --git a/crates/server/src/runtime/agents/lifecycle.rs b/crates/server/src/runtime/agents/lifecycle.rs index 053c3e53..b24f0984 100644 --- a/crates/server/src/runtime/agents/lifecycle.rs +++ b/crates/server/src/runtime/agents/lifecycle.rs @@ -37,7 +37,8 @@ impl ServerRuntime { ) { self.set_agent_status(parent_session_id, child_session_id, SubagentStatus::Failed) .await; - if let Some(session_arc) = self.sessions.lock().await.get(&child_session_id).cloned() { + let session_arc = self.sessions.lock().await.get(&child_session_id).cloned(); + if let Some(session_arc) = session_arc { let mut session = session_arc.lock().await; session.summary.status = SessionRuntimeStatus::Idle; session.summary.updated_at = Utc::now(); diff --git a/crates/server/src/runtime/items.rs b/crates/server/src/runtime/items.rs index 499bb61e..14a40f2d 100644 --- a/crates/server/src/runtime/items.rs +++ b/crates/server/src/runtime/items.rs @@ -355,7 +355,8 @@ impl ServerRuntime { turn_status: Option, worklog: Option, ) { - if let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() { + let session_arc = self.sessions.lock().await.get(&session_id).cloned(); + if let Some(session_arc) = session_arc { let record = { let mut session = session_arc.lock().await; if let Some(history_item) = history_item_from_turn_item(&turn_item) { @@ -402,7 +403,8 @@ impl ServerRuntime { } pub(super) async fn allocate_item_sequence(&self, session_id: SessionId) -> u64 { - if let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() { + let session_arc = self.sessions.lock().await.get(&session_id).cloned(); + if let Some(session_arc) = session_arc { let mut session = session_arc.lock().await; let item_seq = session.next_item_seq; session.loaded_item_count += 1; diff --git a/crates/server/src/runtime/research.rs b/crates/server/src/runtime/research.rs index 04cf26de..4e8adb1c 100644 --- a/crates/server/src/runtime/research.rs +++ b/crates/server/src/runtime/research.rs @@ -685,7 +685,8 @@ impl ServerRuntime { })) .await; let response = rx.await?; - if let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() { + let session_arc = self.sessions.lock().await.get(&session_id).cloned(); + if let Some(session_arc) = session_arc { let mut session = session_arc.lock().await; session.summary.status = SessionRuntimeStatus::ActiveTurn; } diff --git a/crates/server/src/runtime/subagent_usage.rs b/crates/server/src/runtime/subagent_usage.rs index 01c0f4ca..f00e48ce 100644 --- a/crates/server/src/runtime/subagent_usage.rs +++ b/crates/server/src/runtime/subagent_usage.rs @@ -369,13 +369,16 @@ impl ServerRuntime { } async fn apply_parent_usage_snapshot(&self, snapshot: ParentUsageSnapshot) { - if let Some(session_arc) = self + // Bind the session Arc before locking it so the global `sessions` guard is + // released first. Holding `sessions` across the per-session `.await` would + // pin the process-wide lock and can deadlock every other request. + let session_arc = self .sessions .lock() .await .get(&snapshot.session_id) - .cloned() - { + .cloned(); + if let Some(session_arc) = session_arc { let mut session = session_arc.lock().await; snapshot.apply_to_session(&mut session); } diff --git a/crates/server/src/runtime/turn_exec.rs b/crates/server/src/runtime/turn_exec.rs index 0b44a10b..c0341d7e 100644 --- a/crates/server/src/runtime/turn_exec.rs +++ b/crates/server/src/runtime/turn_exec.rs @@ -856,7 +856,8 @@ impl ServerRuntime { ) { self.capture_turn_workspace_baseline(session_id, turn.turn_id, cwd.clone()) .await; - if let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() { + let session_arc = self.sessions.lock().await.get(&session_id).cloned(); + if let Some(session_arc) = session_arc { session_arc.lock().await.turn_approval_cache = crate::execution::ApprovalGrantCache::default(); } @@ -1088,7 +1089,8 @@ impl ServerRuntime { self.capture_turn_workspace_baseline(session_id, turn.turn_id, cwd) .await; } - if let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() { + let session_arc = self.sessions.lock().await.get(&session_id).cloned(); + if let Some(session_arc) = session_arc { session_arc.lock().await.turn_approval_cache = crate::execution::ApprovalGrantCache::default(); } @@ -1149,7 +1151,6 @@ impl ServerRuntime { let mut proposed_plan_leading_normal = String::new(); let mut latest_usage: Option = None; let mut stop_reason: Option = None; - let mut usage_base: Option<(usize, usize, usize, usize)> = None; while let Some(event) = event_rx.recv().await { let assistant_token_text = query_event_trace_token_preview(&event); if let Some(assistant_token_text) = assistant_token_text.as_deref() { @@ -1778,52 +1779,7 @@ impl ServerRuntime { QueryEvent::UsageDelta { usage } | QueryEvent::Usage { usage } => { let turn_usage = TurnUsage::from_usage(&usage); latest_usage = Some(turn_usage.clone()); - let display_total_tokens = usage.display_total_tokens(); - let cache_read_input_tokens = usage.cache_read_input_tokens.unwrap_or(0); - - let base = if let Some(base) = usage_base { - base - } else { - let base = { - let session = event_session_arc.lock().await; - ( - session.summary.total_input_tokens, - session.summary.total_output_tokens, - session.summary.total_tokens, - session.summary.total_cache_read_tokens, - ) - }; - usage_base = Some(base); - base - }; - let total_input_tokens = base.0 + usage.input_tokens; - let total_output_tokens = base.1 + usage.output_tokens; - let total_tokens = base.2 + display_total_tokens; - let total_cache_read_tokens = base.3 + cache_read_input_tokens; if usage_parent_session_id.is_some() { - { - let mut session = event_session_arc.lock().await; - session.summary.total_input_tokens = total_input_tokens; - session.summary.total_output_tokens = total_output_tokens; - session.summary.total_tokens = total_tokens; - session.summary.total_cache_read_tokens = total_cache_read_tokens; - session.summary.last_query_total_tokens = display_total_tokens; - } - let _ = runtime - .broadcast_event(ServerEvent::TurnUsageUpdated( - TurnUsageUpdatedPayload { - session_id, - turn_id: turn_for_events.turn_id, - usage: turn_usage.clone(), - total_input_tokens, - total_output_tokens, - total_tokens, - total_cache_read_tokens, - last_query_input_tokens: usage.input_tokens, - context_window: usage_context_window, - }, - )) - .await; let _ = runtime .publish_subagent_turn_usage( session_id, @@ -2096,7 +2052,12 @@ impl ServerRuntime { .as_ref() .and_then(|summary| summary.latest_usage.clone()); let terminal_stop_reason = event_summary.and_then(|summary| summary.stop_reason); - if usage_parent_session_id.is_none() + if usage_parent_session_id.is_some() + && let Some(usage) = latest_usage.clone() + { + self.publish_subagent_turn_usage(session_id, turn.turn_id, usage) + .await; + } else if usage_parent_session_id.is_none() && let Some(snapshot) = self.parent_usage_snapshot(session_id, turn.turn_id).await { latest_usage = Some(snapshot.turn_usage.to_turn_usage()); diff --git a/crates/server/src/transport.rs b/crates/server/src/transport.rs index cf51c374..ce39c0ce 100644 --- a/crates/server/src/transport.rs +++ b/crates/server/src/transport.rs @@ -322,20 +322,13 @@ async fn run_stdio(runtime: Arc) -> Result<()> { continue; } let value: serde_json::Value = serde_json::from_str(&line)?; - if let Some(response) = runtime - .handle_incoming_with_actions(connection_id, value) - .await - && !send_incoming_response( - &runtime, - &sender_clone, - response, - connection_id, - "stdio_notifications", - ) - .await - { - break; - } + spawn_incoming_message_handler( + Arc::clone(&runtime), + connection_id, + sender_clone.clone(), + value, + "stdio_notifications", + ); } runtime.unregister_connection(connection_id).await; @@ -451,20 +444,13 @@ async fn handle_internal_proxy_connection( match frame { Message::Text(text) => { let value: serde_json::Value = serde_json::from_str(&text)?; - if let Some(response) = runtime - .handle_incoming_with_actions(connection_id, value) - .await - && !send_incoming_response( - &runtime, - &sender_clone, - response, - connection_id, - "internal_proxy_notifications", - ) - .await - { - break; - } + spawn_incoming_message_handler( + Arc::clone(&runtime), + connection_id, + sender_clone.clone(), + value, + "internal_proxy_notifications", + ); } Message::Close(_) => break, _ => {} @@ -570,20 +556,13 @@ async fn handle_websocket_connection( match frame { Message::Text(text) => { let value: serde_json::Value = serde_json::from_str(&text)?; - if let Some(response) = runtime - .handle_incoming_with_actions(connection_id, value) - .await - && !send_incoming_response( - &runtime, - &sender_clone, - response, - connection_id, - "websocket_notifications", - ) - .await - { - break; - } + spawn_incoming_message_handler( + Arc::clone(&runtime), + connection_id, + sender_clone.clone(), + value, + "websocket_notifications", + ); } Message::Close(_) => break, _ => {} @@ -596,6 +575,33 @@ async fn handle_websocket_connection( Ok(()) } +/// Dispatch a single decoded client message without blocking the connection's +/// read loop. +/// +/// The reader task MUST stay responsive so that it can always deliver +/// `turn/interrupt` requests and, critically, client responses to +/// server-initiated requests (approvals, `fs/read_text_file`, user input). +/// Handling a message inline would let a blocked handler or a backpressured +/// response send freeze the reader, which deadlocks the whole connection (the +/// awaited client reply can never be read). Spawning per message keeps the +/// reader draining while individual handlers make progress concurrently. +fn spawn_incoming_message_handler( + runtime: Arc, + connection_id: u64, + sender: mpsc::Sender, + value: serde_json::Value, + queue: &'static str, +) { + tokio::spawn(async move { + if let Some(response) = runtime + .handle_incoming_with_actions(connection_id, value) + .await + { + send_incoming_response(&runtime, &sender, response, connection_id, queue).await; + } + }); +} + async fn send_incoming_response( runtime: &Arc, sender: &mpsc::Sender, diff --git a/crates/tui/src/chatwidget/text_stream.rs b/crates/tui/src/chatwidget/text_stream.rs index 861da952..ff8ef17b 100644 --- a/crates/tui/src/chatwidget/text_stream.rs +++ b/crates/tui/src/chatwidget/text_stream.rs @@ -391,14 +391,8 @@ impl ChatWidget { } fn active_text_item_insert_index(&self, kind: TextItemKind) -> usize { - match kind { - TextItemKind::Reasoning | TextItemKind::ResearchArtifact => self - .active_text_items - .iter() - .position(|item| item.kind == TextItemKind::Assistant) - .unwrap_or(self.active_text_items.len()), - TextItemKind::Assistant => self.active_text_items.len(), - } + let _ = kind; + self.active_text_items.len() } fn commit_completed_text_items(&mut self) { diff --git a/crates/tui/src/chatwidget/transcript_view.rs b/crates/tui/src/chatwidget/transcript_view.rs index 242888d7..a735caea 100644 --- a/crates/tui/src/chatwidget/transcript_view.rs +++ b/crates/tui/src/chatwidget/transcript_view.rs @@ -32,6 +32,11 @@ pub(crate) struct TranscriptOverlayCell { pub(crate) is_selected_user: bool, } +enum LiveViewportLineMode { + Display, + Transcript, +} + impl ChatWidget { pub(crate) fn active_cell_transcript_key(&self) -> Option { let active_cell = self.active_cell.as_ref()?; @@ -154,50 +159,19 @@ impl ChatWidget { } pub(super) fn active_viewport_lines(&self, width: u16) -> Vec> { - let mut lines = Vec::new(); - if let Some(cell) = &self.active_cell { - Self::extend_lines_with_separator(&mut lines, cell.display_lines(width)); - } - for item in &self.active_text_items { - if let Some(cell) = &item.cell { - Self::extend_lines_with_separator(&mut lines, cell.display_lines(width)); - } - } - // Pending tool calls are shown with a pending (cyan) dot until their results arrive. - for pending in &self.pending_tool_calls { - let pending_lines = if let Some(start_time) = pending.start_time { - let mut lines = vec![Line::from(vec![ - crate::exec_cell::spinner(Some(start_time), true), - " ".into(), - Span::styled(pending.title.clone(), Self::tool_text_style()), - ])]; - lines.extend(pending.lines.clone()); - lines - } else { - pending.lines.clone() - }; - Self::extend_lines_with_separator( - &mut lines, - history_cell::AgentMessageCell::new_with_prefix( - pending_lines, - Self::pending_dot_prefix(), - " ", - false, - ) - .display_lines(width), - ); - } - Self::trim_trailing_blank_lines(&mut lines); - lines + self.live_viewport_lines(width, LiveViewportLineMode::Display) } - fn live_transcript_lines(&self, width: u16) -> Vec> { + fn live_viewport_lines(&self, width: u16, mode: LiveViewportLineMode) -> Vec> { let mut lines = Vec::new(); + let cell_lines = |cell: &dyn history_cell::HistoryCell| match mode { + LiveViewportLineMode::Display => cell.display_lines(width), + LiveViewportLineMode::Transcript => cell.transcript_lines(width), + }; if let Some(cell) = &self.active_cell { - Self::extend_lines_with_separator(&mut lines, cell.transcript_lines(width)); + Self::extend_lines_with_separator(&mut lines, cell_lines(cell.as_ref())); } - // Build a merged list of (seq, LiveItem) sorted by seq #[allow(clippy::large_enum_variant)] enum LiveItem { Text(usize), @@ -221,65 +195,124 @@ impl ChatWidget { match item { LiveItem::Text(idx) => { if let Some(cell) = &self.active_text_items[idx].cell { - Self::extend_lines_with_separator(&mut lines, cell.transcript_lines(width)); + Self::extend_lines_with_separator(&mut lines, cell_lines(cell.as_ref())); } } LiveItem::Tool(tool_use_id) => { if let Some(tool_call) = self.active_tool_calls.get(&tool_use_id) { - let transcript_lines = match (&tool_call.tool_name, &tool_call.input) { - (Some(tool_name), Some(input)) => ToolIoCell::from_text_output( - ToolIoCellOptions { - title_line: Some(Self::running_tool_line(&tool_call.title)), - dot_prefix: Self::pending_dot_prefix(), - subsequent_prefix: " ".into(), - output_style: Self::tool_text_style(), - show_empty_ellipsis: false, - }, - tool_name.clone(), - input.clone(), - tool_call.output.clone(), - ) - .transcript_lines(width), - _ => history_cell::AgentMessageCell::new_with_prefix( - tool_call.lines.clone(), - Self::pending_dot_prefix(), - " ", - false, - ) - .transcript_lines(width), + let tool_lines = match mode { + LiveViewportLineMode::Display => { + Self::live_tool_display_lines(width, tool_call) + } + LiveViewportLineMode::Transcript => { + Self::live_tool_transcript_lines(width, tool_call) + } }; - Self::extend_lines_with_separator(&mut lines, transcript_lines); + Self::extend_lines_with_separator(&mut lines, tool_lines); } } } } for pending in &self.pending_tool_calls { let pending_lines = if let Some(start_time) = pending.start_time { - let mut lines = vec![Line::from(vec![ + let mut pending_lines = vec![Line::from(vec![ crate::exec_cell::spinner(Some(start_time), true), " ".into(), Span::styled(pending.title.clone(), Self::tool_text_style()), ])]; - lines.extend(pending.lines.clone()); - lines + pending_lines.extend(pending.lines.clone()); + pending_lines } else { pending.lines.clone() }; Self::extend_lines_with_separator( &mut lines, - history_cell::AgentMessageCell::new_with_prefix( - pending_lines, - Self::pending_dot_prefix(), - " ", - false, - ) - .transcript_lines(width), + match mode { + LiveViewportLineMode::Display => { + history_cell::AgentMessageCell::new_with_prefix( + pending_lines, + Self::pending_dot_prefix(), + " ", + false, + ) + .display_lines(width) + } + LiveViewportLineMode::Transcript => { + history_cell::AgentMessageCell::new_with_prefix( + pending_lines, + Self::pending_dot_prefix(), + " ", + false, + ) + .transcript_lines(width) + } + }, ); } Self::trim_trailing_blank_lines(&mut lines); lines } + fn live_tool_display_lines( + width: u16, + tool_call: &super::ActiveToolCall, + ) -> Vec> { + match (&tool_call.tool_name, &tool_call.input) { + (Some(tool_name), Some(input)) => ToolIoCell::from_text_output( + ToolIoCellOptions { + title_line: Some(Self::running_tool_line(&tool_call.title)), + dot_prefix: Self::pending_dot_prefix(), + subsequent_prefix: " ".into(), + output_style: Self::tool_text_style(), + show_empty_ellipsis: false, + }, + tool_name.clone(), + input.clone(), + tool_call.output.clone(), + ) + .display_lines(width), + _ => history_cell::AgentMessageCell::new_with_prefix( + tool_call.lines.clone(), + Self::pending_dot_prefix(), + " ", + false, + ) + .display_lines(width), + } + } + + fn live_tool_transcript_lines( + width: u16, + tool_call: &super::ActiveToolCall, + ) -> Vec> { + match (&tool_call.tool_name, &tool_call.input) { + (Some(tool_name), Some(input)) => ToolIoCell::from_text_output( + ToolIoCellOptions { + title_line: Some(Self::running_tool_line(&tool_call.title)), + dot_prefix: Self::pending_dot_prefix(), + subsequent_prefix: " ".into(), + output_style: Self::tool_text_style(), + show_empty_ellipsis: false, + }, + tool_name.clone(), + input.clone(), + tool_call.output.clone(), + ) + .transcript_lines(width), + _ => history_cell::AgentMessageCell::new_with_prefix( + tool_call.lines.clone(), + Self::pending_dot_prefix(), + " ", + false, + ) + .transcript_lines(width), + } + } + + fn live_transcript_lines(&self, width: u16) -> Vec> { + self.live_viewport_lines(width, LiveViewportLineMode::Transcript) + } + fn extend_lines_with_separator(target: &mut Vec>, mut next: Vec>) { if next.is_empty() { return; From ceac1eb77cef8edc01ee1338cf9e19b170090626 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Thu, 2 Jul 2026 04:21:09 -1000 Subject: [PATCH 03/16] fix: resolve session actor deadlocks and turn cancellation hangs Turns now execute inline on each session's actor task instead of a separately spawned task, so interrupting/closing a session could no longer be stopped via JoinHandle::abort() and would hang forever waiting on a provider response. Race the model query against the turn's cancellation token instead, and stop removing that token from the shared map when cancelling it (only finalize should remove it), since the removal raced with the query fetching its own clone and could silently hand it a fresh, disconnected token. Also fixes several self-deadlocks introduced by the actor refactor, where turn-completion code tried to read the currently-executing actor's own state through its mailbox (which is not polled again until the in-flight turn finishes): subagent stop hooks, goal accounting's plan-mode check, and duplicate "closed" notifications when a child agent's turn is interrupted mid-close. Additionally adds request/stream idle timeouts to the provider HTTP client so a stalled upstream connection can no longer wedge a turn indefinitely, and folds turn_exec.rs into a turn_exec/ module split for maintainability. --- .cargo/config.toml | 3 + Cargo.lock | 588 +--- .../desktop/src/main/acp-stdio-client.test.ts | 6 +- apps/desktop/src/main/acp-stdio-client.ts | 1 - apps/desktop/src/main/devo-manager.ts | 3 +- crates/arg0/src/lib.rs | 98 +- crates/cli/src/main.rs | 43 +- crates/client/src/client_core.rs | 11 +- crates/client/src/stdio.rs | 37 +- crates/client/src/websocket.rs | 2 +- crates/core/src/session.rs | 27 + crates/core/src/tools/handlers/webfetch.rs | 261 +- crates/network-proxy/src/lib.rs | 34 +- crates/provider/src/anthropic/messages.rs | 34 +- crates/provider/src/http.rs | 21 +- crates/provider/src/lib.rs | 1 + .../provider/src/openai/chat_completions.rs | 34 +- .../src/openai/chat_completions/stream.rs | 17 +- crates/provider/src/openai/responses.rs | 39 +- crates/provider/src/timeout.rs | 89 + crates/server/Cargo.toml | 11 - crates/server/src/bootstrap.rs | 167 +- crates/server/src/execution.rs | 14 +- crates/server/src/lib.rs | 2 - crates/server/src/persistence.rs | 6 +- crates/server/src/runtime.rs | 74 +- crates/server/src/runtime/agents.rs | 494 +-- crates/server/src/runtime/agents/lifecycle.rs | 82 +- crates/server/src/runtime/approval.rs | 265 +- crates/server/src/runtime/command_exec.rs | 8 +- crates/server/src/runtime/connection.rs | 164 +- crates/server/src/runtime/goal_accounting.rs | 9 +- .../server/src/runtime/goal_continuation.rs | 355 +- crates/server/src/runtime/goal_handlers.rs | 14 +- .../server/src/runtime/handlers/acp/prompt.rs | 51 +- .../src/runtime/handlers/acp/session.rs | 61 +- .../runtime/handlers/acp/session_support.rs | 79 +- .../runtime/handlers/acp_config_options.rs | 368 +- .../server/src/runtime/handlers/compaction.rs | 57 +- .../src/runtime/handlers/message_edit.rs | 134 +- crates/server/src/runtime/handlers/session.rs | 300 +- crates/server/src/runtime/handlers/turn.rs | 818 +++-- .../src/runtime/handlers/workspace_changes.rs | 24 +- crates/server/src/runtime/hooks.rs | 82 +- crates/server/src/runtime/items.rs | 286 +- crates/server/src/runtime/lifecycle.rs | 165 +- crates/server/src/runtime/research.rs | 453 ++- crates/server/src/runtime/research_capture.rs | 1 + crates/server/src/runtime/research_events.rs | 26 +- crates/server/src/runtime/research_stages.rs | 14 +- .../src/runtime/research_tool_runtime.rs | 25 +- .../src/runtime/session_actor/commands.rs | 198 ++ .../src/runtime/session_actor/handle.rs | 616 ++++ .../server/src/runtime/session_actor/loop_.rs | 590 ++++ .../server/src/runtime/session_actor/mod.rs | 11 + .../src/runtime/session_actor/registry.rs | 199 ++ .../src/runtime/session_actor/snapshots.rs | 110 + .../server/src/runtime/session_actor/state.rs | 204 ++ .../server/src/runtime/session_actor/turn.rs | 107 + .../src/runtime/session_actor/turn_inline.rs | 71 + .../server/src/runtime/session_interactive.rs | 147 + crates/server/src/runtime/subagent_usage.rs | 70 +- crates/server/src/runtime/turn_exec.rs | 2985 ----------------- .../src/runtime/turn_exec/event_stream.rs | 664 ++++ .../server/src/runtime/turn_exec/finalize.rs | 320 ++ .../server/src/runtime/turn_exec/followup.rs | 192 ++ .../src/runtime/turn_exec/item_stream.rs | 244 ++ crates/server/src/runtime/turn_exec/mod.rs | 70 + crates/server/src/runtime/turn_exec/query.rs | 197 ++ crates/server/src/runtime/turn_exec/shell.rs | 211 ++ crates/server/src/runtime/turn_exec/tests.rs | 371 ++ .../src/runtime/turn_exec/tool_display.rs | 404 +++ .../src/runtime/turn_exec/tool_results.rs | 430 +++ crates/server/src/runtime/turn_exec/trace.rs | 96 + crates/server/src/runtime/turn_exec/types.rs | 85 + crates/server/src/runtime/turn_reservation.rs | 76 +- crates/server/src/runtime/user_input.rs | 46 +- .../server/src/runtime/workspace_baseline.rs | 16 +- crates/server/src/server_tray.rs | 673 ---- crates/server/tests/deep_research_e2e.rs | 284 +- crates/server/tests/subagent_lifecycle.rs | 52 +- .../tests/support/subagent_lifecycle.rs | 51 + .../tui/src/chatwidget/subagent_live_list.rs | 5 +- crates/tui/src/chatwidget/subagent_monitor.rs | 121 +- crates/tui/src/chatwidget/text_stream.rs | 83 +- crates/tui/src/chatwidget/transcript_view.rs | 53 +- crates/tui/src/chatwidget_tests.rs | 71 +- crates/tui/src/custom_terminal.rs | 2 +- crates/tui/src/interactive.rs | 25 +- crates/tui/src/streaming/controller.rs | 4 + crates/tui/src/worker.rs | 296 +- crates/tui/src/worker/acp_events.rs | 91 + 92 files changed, 9414 insertions(+), 7085 deletions(-) create mode 100644 crates/provider/src/timeout.rs create mode 100644 crates/server/src/runtime/session_actor/commands.rs create mode 100644 crates/server/src/runtime/session_actor/handle.rs create mode 100644 crates/server/src/runtime/session_actor/loop_.rs create mode 100644 crates/server/src/runtime/session_actor/mod.rs create mode 100644 crates/server/src/runtime/session_actor/registry.rs create mode 100644 crates/server/src/runtime/session_actor/snapshots.rs create mode 100644 crates/server/src/runtime/session_actor/state.rs create mode 100644 crates/server/src/runtime/session_actor/turn.rs create mode 100644 crates/server/src/runtime/session_actor/turn_inline.rs create mode 100644 crates/server/src/runtime/session_interactive.rs delete mode 100644 crates/server/src/runtime/turn_exec.rs create mode 100644 crates/server/src/runtime/turn_exec/event_stream.rs create mode 100644 crates/server/src/runtime/turn_exec/finalize.rs create mode 100644 crates/server/src/runtime/turn_exec/followup.rs create mode 100644 crates/server/src/runtime/turn_exec/item_stream.rs create mode 100644 crates/server/src/runtime/turn_exec/mod.rs create mode 100644 crates/server/src/runtime/turn_exec/query.rs create mode 100644 crates/server/src/runtime/turn_exec/shell.rs create mode 100644 crates/server/src/runtime/turn_exec/tests.rs create mode 100644 crates/server/src/runtime/turn_exec/tool_display.rs create mode 100644 crates/server/src/runtime/turn_exec/tool_results.rs create mode 100644 crates/server/src/runtime/turn_exec/trace.rs create mode 100644 crates/server/src/runtime/turn_exec/types.rs delete mode 100644 crates/server/src/server_tray.rs diff --git a/.cargo/config.toml b/.cargo/config.toml index 36945c07..b9b1e276 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -1,3 +1,6 @@ +[env] +RUST_MIN_STACK = "16777216" + [target.x86_64-unknown-linux-musl] rustflags = ["-C", "target-feature=+crt-static", "-C", "strip=symbols"] diff --git a/Cargo.lock b/Cargo.lock index f5f7a35a..4b405433 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -96,31 +96,6 @@ version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" -[[package]] -name = "android-activity" -version = "0.6.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f2a1bb052857d5dd49572219344a7332b31b76405648eabac5bc68978251bcd" -dependencies = [ - "android-properties", - "bitflags 2.13.0", - "cc", - "jni", - "libc", - "log", - "ndk", - "ndk-context", - "ndk-sys", - "num_enum", - "thiserror 2.0.18", -] - -[[package]] -name = "android-properties" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fc7eb209b1518d6bb87b283c20095f5228ecda460da70b44f0802523dea6da04" - [[package]] name = "android_system_properties" version = "0.1.5" @@ -234,11 +209,11 @@ dependencies = [ "clipboard-win", "image", "log", - "objc2 0.6.4", - "objc2-app-kit 0.3.2", + "objc2", + "objc2-app-kit", "objc2-core-foundation", "objc2-core-graphics", - "objc2-foundation 0.3.2", + "objc2-foundation", "parking_lot", "percent-encoding", "windows-sys 0.60.2", @@ -552,9 +527,6 @@ name = "bitflags" version = "2.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" -dependencies = [ - "serde_core", -] [[package]] name = "blake3" @@ -597,24 +569,6 @@ dependencies = [ "generic-array", ] -[[package]] -name = "block2" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c132eebf10f5cad5289222520a4a058514204aed6d791f1cf4fe8088b82d15f" -dependencies = [ - "objc2 0.5.2", -] - -[[package]] -name = "block2" -version = "0.6.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cdeb9d870516001442e364c5220d3574d2da8dc765554b4a617230d33fa58ef5" -dependencies = [ - "objc2 0.6.4", -] - [[package]] name = "blocking" version = "1.6.2" @@ -760,20 +714,6 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ade8366b8bd5ba243f0a58f036cc0ca8a2f069cff1a2351ef1cac6b083e16fc0" -[[package]] -name = "calloop" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b99da2f8558ca23c71f4fd15dc57c906239752dd27ff3c00a1d56b685b7cbfec" -dependencies = [ - "bitflags 2.13.0", - "log", - "polling", - "rustix 0.38.44", - "slab", - "thiserror 1.0.69", -] - [[package]] name = "cassowary" version = "0.3.0" @@ -1092,30 +1032,6 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" -[[package]] -name = "core-graphics" -version = "0.23.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c07782be35f9e1140080c6b96f0d44b739e2278479f64e02fdab4e32dfd8b081" -dependencies = [ - "bitflags 1.3.2", - "core-foundation 0.9.4", - "core-graphics-types", - "foreign-types", - "libc", -] - -[[package]] -name = "core-graphics-types" -version = "0.1.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45390e6114f68f718cc7a830514a96f903cccd70d02a8f6d9f643ac4ba45afaf" -dependencies = [ - "bitflags 1.3.2", - "core-foundation 0.9.4", - "libc", -] - [[package]] name = "countio" version = "0.3.0" @@ -1283,12 +1199,6 @@ version = "0.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52560adf09603e58c9a7ee1fe1dcb95a16927b17c127f0ac02d6e768a0e25bc1" -[[package]] -name = "cursor-icon" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f27ae1dd37df86211c42e150270f82743308803d90a6f6e6651cd730d5e1732f" - [[package]] name = "darling" version = "0.20.11" @@ -1806,8 +1716,6 @@ dependencies = [ "diffy", "fs2", "futures", - "image", - "objc2-foundation 0.3.2", "pretty_assertions", "rusqlite", "serde", @@ -1822,10 +1730,7 @@ dependencies = [ "toml", "tracing", "tracing-subscriber", - "tray-icon", "uuid", - "windows-sys 0.61.2", - "winit", ] [[package]] @@ -2067,12 +1972,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "dispatch" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bd0c93bb4b0c6d9b77f4435b0ae98c24d17f1c45b2ff844c6151a07256ca923b" - [[package]] name = "dispatch2" version = "0.3.1" @@ -2080,7 +1979,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e0e367e4e7da84520dedcac1901e4da967309406d1e51017ae1abfb97adbd38" dependencies = [ "bitflags 2.13.0", - "objc2 0.6.4", + "objc2", ] [[package]] @@ -2094,15 +1993,6 @@ dependencies = [ "syn", ] -[[package]] -name = "dlib" -version = "0.5.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ab8ecd87370524b461f8557c119c405552c396ed91fc0a8eec68679eab26f94a" -dependencies = [ - "libloading", -] - [[package]] name = "dotenvy" version = "0.15.7" @@ -2115,12 +2005,6 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "75b325c5dbd37f80359721ad39aca5a29fb04c89279657cffdda8736d0c0b9d2" -[[package]] -name = "dpi" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d8b14ccef22fc6f5a8f4d7d768562a182c04ce9a3b3157b91390b52ddfdf1a76" - [[package]] name = "dtor" version = "0.1.1" @@ -2372,33 +2256,6 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" -[[package]] -name = "foreign-types" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d737d9aa519fb7b749cbc3b962edcf310a8dd1f4b67c91c4f83975dbdd17d965" -dependencies = [ - "foreign-types-macros", - "foreign-types-shared", -] - -[[package]] -name = "foreign-types-macros" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a5c6c585bc94aaf2c7b51dd4c2ba22680844aba4c687be581871a6f518c5742" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "foreign-types-shared" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b" - [[package]] name = "form_urlencoded" version = "1.2.2" @@ -3339,7 +3196,7 @@ dependencies = [ "cfg-if 1.0.4", "combine", "jni-macros", - "jni-sys 0.4.1", + "jni-sys", "log", "simd_cesu8", "thiserror 2.0.18", @@ -3360,15 +3217,6 @@ dependencies = [ "syn", ] -[[package]] -name = "jni-sys" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41a652e1f9b6e0275df1f15b32661cf0d4b78d4d87ddec5e0c3c20f097433258" -dependencies = [ - "jni-sys 0.4.1", -] - [[package]] name = "jni-sys" version = "0.4.1" @@ -3420,17 +3268,6 @@ dependencies = [ "serde_json", ] -[[package]] -name = "keyboard-types" -version = "0.7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b750dcadc39a09dbadd74e118f6dd6598df77fa01df0cfcdc52c28dece74528a" -dependencies = [ - "bitflags 2.13.0", - "serde", - "unicode-segmentation", -] - [[package]] name = "keyring" version = "3.6.3" @@ -3513,26 +3350,13 @@ dependencies = [ "pkg-config", ] -[[package]] -name = "libloading" -version = "0.8.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" -dependencies = [ - "cfg-if 1.0.4", - "windows-link", -] - [[package]] name = "libredox" version = "0.1.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7ddbf48fd451246b1f8c2610bd3b4ac0cc6e149d89832867093ab69a17194f08" dependencies = [ - "bitflags 2.13.0", "libc", - "plain", - "redox_syscall 0.7.5", ] [[package]] @@ -3797,25 +3621,6 @@ dependencies = [ "pxfm", ] -[[package]] -name = "muda" -version = "0.19.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1dd04e60bc0b07438a6771710ee1698f98f6ebbc7f89b61264af1563b8aeb878" -dependencies = [ - "crossbeam-channel", - "dpi", - "keyboard-types", - "objc2 0.6.4", - "objc2-app-kit 0.3.2", - "objc2-core-foundation", - "objc2-foundation 0.3.2", - "once_cell", - "png", - "thiserror 2.0.18", - "windows-sys 0.61.2", -] - [[package]] name = "ndarray" version = "0.16.1" @@ -3831,35 +3636,12 @@ dependencies = [ "rawpointer", ] -[[package]] -name = "ndk" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3f42e7bbe13d351b6bead8286a43aac9534b82bd3cc43e47037f012ebfd62d4" -dependencies = [ - "bitflags 2.13.0", - "jni-sys 0.3.1", - "log", - "ndk-sys", - "num_enum", - "thiserror 1.0.69", -] - [[package]] name = "ndk-context" version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b" -[[package]] -name = "ndk-sys" -version = "0.6.0+11769913" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ee6cda3051665f1fb8d9e08fc35c96d5a244fb1be711a03b71118828afc9a873" -dependencies = [ - "jni-sys 0.3.1", -] - [[package]] name = "nix" version = "0.28.0" @@ -4081,28 +3863,6 @@ dependencies = [ "libc", ] -[[package]] -name = "num_enum" -version = "0.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d0bca838442ec211fa11de3a8b0e0e8f3a4522575b5c4c06ed722e005036f26" -dependencies = [ - "num_enum_derive", - "rustversion", -] - -[[package]] -name = "num_enum_derive" -version = "0.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "680998035259dcfcafe653688bf2aa6d3e2dc05e98be6ab46afb089dc84f1df8" -dependencies = [ - "proc-macro-crate", - "proc-macro2", - "quote", - "syn", -] - [[package]] name = "number_prefix" version = "0.4.0" @@ -4129,22 +3889,6 @@ dependencies = [ "url", ] -[[package]] -name = "objc-sys" -version = "0.3.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cdb91bdd390c7ce1a8607f35f3ca7151b65afc0ff5ff3b34fa350f7d7c7e4310" - -[[package]] -name = "objc2" -version = "0.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46a785d4eeff09c14c487497c162e92766fbb3e4059a71840cecc03d9a50b804" -dependencies = [ - "objc-sys", - "objc2-encode", -] - [[package]] name = "objc2" version = "0.6.4" @@ -4154,22 +3898,6 @@ dependencies = [ "objc2-encode", ] -[[package]] -name = "objc2-app-kit" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e4e89ad9e3d7d297152b17d39ed92cd50ca8063a89a9fa569046d41568891eff" -dependencies = [ - "bitflags 2.13.0", - "block2 0.5.1", - "libc", - "objc2 0.5.2", - "objc2-core-data", - "objc2-core-image", - "objc2-foundation 0.2.2", - "objc2-quartz-core", -] - [[package]] name = "objc2-app-kit" version = "0.3.2" @@ -4177,46 +3905,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d49e936b501e5c5bf01fda3a9452ff86dc3ea98ad5f283e1455153142d97518c" dependencies = [ "bitflags 2.13.0", - "objc2 0.6.4", - "objc2-core-foundation", + "objc2", "objc2-core-graphics", - "objc2-foundation 0.3.2", -] - -[[package]] -name = "objc2-cloud-kit" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74dd3b56391c7a0596a295029734d3c1c5e7e510a4cb30245f8221ccea96b009" -dependencies = [ - "bitflags 2.13.0", - "block2 0.5.1", - "objc2 0.5.2", - "objc2-core-location", - "objc2-foundation 0.2.2", -] - -[[package]] -name = "objc2-contacts" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5ff520e9c33812fd374d8deecef01d4a840e7b41862d849513de77e44aa4889" -dependencies = [ - "block2 0.5.1", - "objc2 0.5.2", - "objc2-foundation 0.2.2", -] - -[[package]] -name = "objc2-core-data" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "617fbf49e071c178c0b24c080767db52958f716d9eabdf0890523aeae54773ef" -dependencies = [ - "bitflags 2.13.0", - "block2 0.5.1", - "objc2 0.5.2", - "objc2-foundation 0.2.2", + "objc2-foundation", ] [[package]] @@ -4227,7 +3918,7 @@ checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" dependencies = [ "bitflags 2.13.0", "dispatch2", - "objc2 0.6.4", + "objc2", ] [[package]] @@ -4238,54 +3929,17 @@ checksum = "e022c9d066895efa1345f8e33e584b9f958da2fd4cd116792e15e07e4720a807" dependencies = [ "bitflags 2.13.0", "dispatch2", - "objc2 0.6.4", + "objc2", "objc2-core-foundation", "objc2-io-surface", ] -[[package]] -name = "objc2-core-image" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "55260963a527c99f1819c4f8e3b47fe04f9650694ef348ffd2227e8196d34c80" -dependencies = [ - "block2 0.5.1", - "objc2 0.5.2", - "objc2-foundation 0.2.2", - "objc2-metal", -] - -[[package]] -name = "objc2-core-location" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "000cfee34e683244f284252ee206a27953279d370e309649dc3ee317b37e5781" -dependencies = [ - "block2 0.5.1", - "objc2 0.5.2", - "objc2-contacts", - "objc2-foundation 0.2.2", -] - [[package]] name = "objc2-encode" version = "4.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ef25abbcd74fb2609453eb695bd2f860d389e457f67dc17cafc8b8cbc89d0c33" -[[package]] -name = "objc2-foundation" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ee638a5da3799329310ad4cfa62fbf045d5f56e3ef5ba4149e7452dcf89d5a8" -dependencies = [ - "bitflags 2.13.0", - "block2 0.5.1", - "dispatch", - "libc", - "objc2 0.5.2", -] - [[package]] name = "objc2-foundation" version = "0.3.2" @@ -4293,8 +3947,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272" dependencies = [ "bitflags 2.13.0", - "block2 0.6.2", - "objc2 0.6.4", + "objc2", "objc2-core-foundation", ] @@ -4315,57 +3968,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "180788110936d59bab6bd83b6060ffdfffb3b922ba1396b312ae795e1de9d81d" dependencies = [ "bitflags 2.13.0", - "objc2 0.6.4", + "objc2", "objc2-core-foundation", ] -[[package]] -name = "objc2-link-presentation" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1a1ae721c5e35be65f01a03b6d2ac13a54cb4fa70d8a5da293d7b0020261398" -dependencies = [ - "block2 0.5.1", - "objc2 0.5.2", - "objc2-app-kit 0.2.2", - "objc2-foundation 0.2.2", -] - -[[package]] -name = "objc2-metal" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dd0cba1276f6023976a406a14ffa85e1fdd19df6b0f737b063b95f6c8c7aadd6" -dependencies = [ - "bitflags 2.13.0", - "block2 0.5.1", - "objc2 0.5.2", - "objc2-foundation 0.2.2", -] - -[[package]] -name = "objc2-quartz-core" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e42bee7bff906b14b167da2bac5efe6b6a07e6f7c0a21a7308d40c960242dc7a" -dependencies = [ - "bitflags 2.13.0", - "block2 0.5.1", - "objc2 0.5.2", - "objc2-foundation 0.2.2", - "objc2-metal", -] - -[[package]] -name = "objc2-symbols" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0a684efe3dec1b305badae1a28f6555f6ddd3bb2c2267896782858d5a78404dc" -dependencies = [ - "objc2 0.5.2", - "objc2-foundation 0.2.2", -] - [[package]] name = "objc2-system-configuration" version = "0.3.2" @@ -4375,51 +3981,6 @@ dependencies = [ "objc2-core-foundation", ] -[[package]] -name = "objc2-ui-kit" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8bb46798b20cd6b91cbd113524c490f1686f4c4e8f49502431415f3512e2b6f" -dependencies = [ - "bitflags 2.13.0", - "block2 0.5.1", - "objc2 0.5.2", - "objc2-cloud-kit", - "objc2-core-data", - "objc2-core-image", - "objc2-core-location", - "objc2-foundation 0.2.2", - "objc2-link-presentation", - "objc2-quartz-core", - "objc2-symbols", - "objc2-uniform-type-identifiers", - "objc2-user-notifications", -] - -[[package]] -name = "objc2-uniform-type-identifiers" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44fa5f9748dbfe1ca6c0b79ad20725a11eca7c2218bceb4b005cb1be26273bfe" -dependencies = [ - "block2 0.5.1", - "objc2 0.5.2", - "objc2-foundation 0.2.2", -] - -[[package]] -name = "objc2-user-notifications" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76cfcbf642358e8689af64cee815d139339f3ed8ad05103ed5eaf73db8d84cb3" -dependencies = [ - "bitflags 2.13.0", - "block2 0.5.1", - "objc2 0.5.2", - "objc2-core-location", - "objc2-foundation 0.2.2", -] - [[package]] name = "once_cell" version = "1.21.4" @@ -4472,16 +4033,6 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" -[[package]] -name = "orbclient" -version = "0.3.55" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5df339f526ea9a60e371768d50efc2f2508c7203290731565d1f7a6f71d21747" -dependencies = [ - "libc", - "libredox", -] - [[package]] name = "ordered-stream" version = "0.2.0" @@ -4535,7 +4086,7 @@ checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" dependencies = [ "cfg-if 1.0.4", "libc", - "redox_syscall 0.5.18", + "redox_syscall", "smallvec", "windows-link", ] @@ -4618,12 +4169,6 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" -[[package]] -name = "plain" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4596b6d070b27117e987119b4dac604f3c58cfb0b191112e24771b2faeac1a6" - [[package]] name = "plist" version = "1.8.0" @@ -5053,15 +4598,6 @@ dependencies = [ "libc", ] -[[package]] -name = "redox_syscall" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa" -dependencies = [ - "bitflags 1.3.2", -] - [[package]] name = "redox_syscall" version = "0.5.18" @@ -5071,15 +4607,6 @@ dependencies = [ "bitflags 2.13.0", ] -[[package]] -name = "redox_syscall" -version = "0.7.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4666a1a60d8412eab19d94f6d13dcc9cea0a5ef4fdf6a5db306537413c661b1b" -dependencies = [ - "bitflags 2.13.0", -] - [[package]] name = "redox_users" version = "0.5.2" @@ -6001,15 +5528,6 @@ dependencies = [ "serde", ] -[[package]] -name = "smol_str" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dd538fb6910ac1099850255cf94a94df6551fbdd602454387d0adb2d1ca6dead" -dependencies = [ - "serde", -] - [[package]] name = "smol_str" version = "0.3.2" @@ -6758,26 +6276,6 @@ dependencies = [ "tracing-serde", ] -[[package]] -name = "tray-icon" -version = "0.24.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "65ba1e5f6b9ef9fd87e21b9c6f351554dbd717960089168fcfdef854686961dc" -dependencies = [ - "crossbeam-channel", - "dirs", - "muda", - "objc2 0.6.4", - "objc2-app-kit 0.3.2", - "objc2-core-foundation", - "objc2-core-graphics", - "objc2-foundation 0.3.2", - "once_cell", - "png", - "thiserror 2.0.18", - "windows-sys 0.61.2", -] - [[package]] name = "tree-sitter" version = "0.25.10" @@ -7592,8 +7090,8 @@ dependencies = [ "jni", "log", "ndk-context", - "objc2 0.6.4", - "objc2-foundation 0.3.2", + "objc2", + "objc2-foundation", "url", "web-sys", ] @@ -8033,45 +7531,6 @@ version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" -[[package]] -name = "winit" -version = "0.30.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6755fa58a9f8350bd1e472d4c3fcc25f824ec358933bba33306d0b63df5978d" -dependencies = [ - "android-activity", - "atomic-waker", - "bitflags 2.13.0", - "block2 0.5.1", - "calloop", - "cfg_aliases 0.2.1", - "concurrent-queue", - "core-foundation 0.9.4", - "core-graphics", - "cursor-icon", - "dpi", - "js-sys", - "libc", - "ndk", - "objc2 0.5.2", - "objc2-app-kit 0.2.2", - "objc2-foundation 0.2.2", - "objc2-ui-kit", - "orbclient", - "pin-project", - "redox_syscall 0.4.1", - "rustix 0.38.44", - "smol_str 0.2.2", - "tracing", - "unicode-segmentation", - "wasm-bindgen", - "wasm-bindgen-futures", - "web-sys", - "web-time", - "windows-sys 0.52.0", - "xkbcommon-dl", -] - [[package]] name = "winnow" version = "0.7.15" @@ -8391,25 +7850,6 @@ dependencies = [ "winapi", ] -[[package]] -name = "xkbcommon-dl" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d039de8032a9a8856a6be89cea3e5d12fdd82306ab7c94d74e6deab2460651c5" -dependencies = [ - "bitflags 2.13.0", - "dlib", - "log", - "once_cell", - "xkeysym", -] - -[[package]] -name = "xkeysym" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9cc00251562a284751c9973bace760d86c0276c471b4be569fe6b068ee97a56" - [[package]] name = "yaml-rust" version = "0.4.5" diff --git a/apps/desktop/src/main/acp-stdio-client.test.ts b/apps/desktop/src/main/acp-stdio-client.test.ts index a3e4895d..4b7b85cb 100644 --- a/apps/desktop/src/main/acp-stdio-client.test.ts +++ b/apps/desktop/src/main/acp-stdio-client.test.ts @@ -46,12 +46,12 @@ describe("StdioAcpClient", () => { const env = buildServerProcessEnv({ baseEnv: { PATH: "/usr/bin", KEEP: "base" }, homeDir: "/Users/tester", - optionsEnv: { DEVO_SUPPRESS_SERVER_TRAY: "1", PATH: "/custom/bin" }, + optionsEnv: { CUSTOM_FLAG: "1", PATH: "/custom/bin" }, pathSeparator: ":", }) expect(env).toMatchObject({ - DEVO_SUPPRESS_SERVER_TRAY: "1", + CUSTOM_FLAG: "1", KEEP: "base", PATH: "/Users/tester/.devo/bin:/custom/bin", }) @@ -61,13 +61,11 @@ describe("StdioAcpClient", () => { const env = buildServerProcessEnv({ baseEnv: { PATH: "/usr/bin", KEEP: "base" }, homeDir: "/Users/tester", - optionsEnv: { DEVO_SUPPRESS_SERVER_TRAY: "1" }, pathSeparator: ":", runtimeBinDir: "/Applications/Devo.app/Contents/Resources/runtime/bin", }) expect(env).toMatchObject({ - DEVO_SUPPRESS_SERVER_TRAY: "1", KEEP: "base", PATH: "/Applications/Devo.app/Contents/Resources/runtime/bin:/usr/bin", }) diff --git a/apps/desktop/src/main/acp-stdio-client.ts b/apps/desktop/src/main/acp-stdio-client.ts index eddbd1cf..718ca751 100644 --- a/apps/desktop/src/main/acp-stdio-client.ts +++ b/apps/desktop/src/main/acp-stdio-client.ts @@ -9,7 +9,6 @@ import type { AcpTrafficLogRecord, AcpTrafficLogger } from "./acp-traffic-log" import { createLogger } from "./logger" const log = createLogger("acp-stdio-client") -export const SUPPRESS_SERVER_TRAY_ENV = "DEVO_SUPPRESS_SERVER_TRAY" export const DESKTOP_NETWORK_PROXY_MODE_ENV = "DEVO_DESKTOP_NETWORK_PROXY_MODE" export const DESKTOP_NETWORK_PROXY_URL_ENV = "DEVO_DESKTOP_NETWORK_PROXY_URL" export const DESKTOP_NETWORK_NO_PROXY_ENV = "DEVO_DESKTOP_NETWORK_NO_PROXY" diff --git a/apps/desktop/src/main/devo-manager.ts b/apps/desktop/src/main/devo-manager.ts index 2e666db5..6bb8d59f 100644 --- a/apps/desktop/src/main/devo-manager.ts +++ b/apps/desktop/src/main/devo-manager.ts @@ -6,7 +6,7 @@ import { type AcpTrafficLogger, type AcpTrafficLogState, } from "./acp-traffic-log" -import { StdioAcpClient, SUPPRESS_SERVER_TRAY_ENV } from "./acp-stdio-client" +import { StdioAcpClient } from "./acp-stdio-client" import { resolveDevoProgram } from "./devo-program" import { createLogger } from "./logger" import { startNotificationWatcher, stopNotificationWatcher } from "./notification-watcher" @@ -143,7 +143,6 @@ function getOrCreateClient(): StdioAcpClient { }) stdioClient = new StdioAcpClient({ program, - env: { [SUPPRESS_SERVER_TRAY_ENV]: "1" }, networkProxy: getSettings().servers.networkProxy, trafficLogger: getAcpTrafficLogger(), }) diff --git a/crates/arg0/src/lib.rs b/crates/arg0/src/lib.rs index 6260ed51..3c03c993 100644 --- a/crates/arg0/src/lib.rs +++ b/crates/arg0/src/lib.rs @@ -29,8 +29,6 @@ use anyhow::Result; /// Alias name for the server sub‑binary. const SERVER_ALIAS: &str = "devo-server"; const ALIAS_SENTINEL_PREFIX: &str = "--devo-alias="; -/// Internal env var used by Devo Desktop when it owns the visible tray icon. -pub const SUPPRESS_SERVER_TRAY_ENV: &str = "DEVO_SUPPRESS_SERVER_TRAY"; /// Directory (under DEVO_HOME/tmp/arg0) where alias entries are created. const ALIAS_TEMP_ROOT: &str = "arg0"; @@ -190,39 +188,6 @@ where filtered } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum ServerAliasDispatchMode { - AsyncRuntime, - MacosMainThreadTray, -} - -/// Returns true when a child server process should avoid creating its own tray. -pub fn suppress_server_tray_from_env() -> bool { - let value = std::env::var(SUPPRESS_SERVER_TRAY_ENV).ok(); - suppress_server_tray_value(value.as_deref()) -} - -/// Interprets [`SUPPRESS_SERVER_TRAY_ENV`] without reading process env. -pub fn suppress_server_tray_value(value: Option<&str>) -> bool { - value == Some("1") -} - -fn server_alias_dispatch_mode() -> ServerAliasDispatchMode { - server_alias_dispatch_mode_for(suppress_server_tray_from_env()) -} - -fn server_alias_dispatch_mode_for(suppress_server_tray: bool) -> ServerAliasDispatchMode { - if suppress_server_tray { - return ServerAliasDispatchMode::AsyncRuntime; - } - - if cfg!(target_os = "macos") { - ServerAliasDispatchMode::MacosMainThreadTray - } else { - ServerAliasDispatchMode::AsyncRuntime - } -} - /// Build a multi‑thread Tokio runtime. fn build_runtime() -> Result { let mut builder = tokio::runtime::Builder::new_multi_thread(); @@ -232,32 +197,9 @@ fn build_runtime() -> Result { } fn run_server_alias_dispatch() -> Result<()> { - match server_alias_dispatch_mode() { - ServerAliasDispatchMode::AsyncRuntime => { - let runtime = build_runtime()?; - runtime.block_on(run_server_dispatch()); - Ok(()) - } - ServerAliasDispatchMode::MacosMainThreadTray => { - run_macos_server_dispatch(); - Ok(()) - } - } -} - -#[cfg(target_os = "macos")] -fn run_macos_server_dispatch() { - let args = parse_server_dispatch_args(); - let _logging = install_server_logging_for_dispatch(); - if let Err(err) = devo_server::run_server_process_with_macos_tray(args) { - eprintln!("server error: {err}"); - std::process::exit(1); - } -} - -#[cfg(not(target_os = "macos"))] -fn run_macos_server_dispatch() { - unreachable!("macOS server dispatch is only selected on macOS"); + let runtime = build_runtime()?; + runtime.block_on(run_server_dispatch()); + Ok(()) } fn parse_server_dispatch_args() -> devo_server::ServerProcessArgs { @@ -301,7 +243,7 @@ fn install_server_logging_for_dispatch() -> Option { async fn run_server_dispatch() { let args = parse_server_dispatch_args(); let _logging = install_server_logging_for_dispatch(); - if let Err(err) = devo_server::run_server_process(args).await { + if let Err(err) = devo_server::run_server_process(args, devo_server::ServerProcessRunOptions::default()).await { eprintln!("server error: {err}"); std::process::exit(1); } @@ -597,38 +539,6 @@ mod tests { ); } - #[test] - fn server_alias_dispatch_mode_matches_platform_threading_contract() { - let expected = if cfg!(target_os = "macos") { - super::ServerAliasDispatchMode::MacosMainThreadTray - } else { - super::ServerAliasDispatchMode::AsyncRuntime - }; - - assert_eq!(super::server_alias_dispatch_mode(), expected); - } - - #[test] - fn server_alias_dispatch_mode_suppresses_macos_tray_when_requested() { - assert_eq!( - super::server_alias_dispatch_mode_for(/*suppress_server_tray*/ true), - super::ServerAliasDispatchMode::AsyncRuntime - ); - } - - #[test] - fn suppress_server_tray_value_only_accepts_one() { - assert_eq!( - [ - super::suppress_server_tray_value(None), - super::suppress_server_tray_value(Some("")), - super::suppress_server_tray_value(Some("true")), - super::suppress_server_tray_value(Some("1")), - ], - [false, false, false, true] - ); - } - #[test] fn lock_unavailable_detection_handles_would_block() { assert_eq!( diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index c301d458..a75553ef 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -21,10 +21,9 @@ use devo_core::UpdateCheckOutcome; use devo_core::UpdateChecker; use devo_core::format_update_notification; use devo_server::ServerProcessArgs; +use devo_server::ServerProcessRunOptions; use devo_server::ServerTransportMode; use devo_server::run_server_process; -#[cfg(target_os = "macos")] -use devo_server::run_server_process_with_macos_tray; use devo_util_paths::find_devo_home; use tracing_subscriber::filter::LevelFilter; @@ -226,7 +225,7 @@ async fn run_cli() -> Result<()> { }) => { let args = server_process_args_from_cli(&cli).expect("server command args"); let _logging = install_server_logging(&cli)?; - run_server_process(args).await + run_server_process(args, ServerProcessRunOptions::default()).await } None => { maybe_print_startup_update(&cli).await; @@ -253,37 +252,10 @@ async fn run_cli() -> Result<()> { } } -#[cfg(target_os = "macos")] -fn direct_server_early_dispatch() -> devo_arg0::EarlyDispatch { - let cli = match Cli::try_parse() { - Ok(cli) => cli, - Err(_) => return devo_arg0::EarlyDispatch::Continue, - }; - if !should_direct_dispatch_macos_server(&cli, devo_arg0::suppress_server_tray_from_env()) { - return devo_arg0::EarlyDispatch::Continue; - } - let Some(args) = server_process_args_from_cli(&cli) else { - return devo_arg0::EarlyDispatch::Continue; - }; - devo_arg0::EarlyDispatch::Handled(run_macos_server_process_from_cli(&cli, args)) -} - -#[cfg(not(target_os = "macos"))] fn direct_server_early_dispatch() -> devo_arg0::EarlyDispatch { devo_arg0::EarlyDispatch::Continue } -#[cfg(target_os = "macos")] -fn run_macos_server_process_from_cli(cli: &Cli, args: ServerProcessArgs) -> Result<()> { - let _logging = install_server_logging(cli)?; - run_server_process_with_macos_tray(args) -} - -#[cfg(any(target_os = "macos", test))] -fn should_direct_dispatch_macos_server(cli: &Cli, suppress_server_tray: bool) -> bool { - !suppress_server_tray && server_process_args_from_cli(cli).is_some() -} - fn server_process_args_from_cli(cli: &Cli) -> Option { match &cli.command { Some(Command::Server { @@ -654,17 +626,6 @@ mod tests { ); } - #[test] - fn direct_server_dispatch_suppresses_macos_tray_when_requested() { - let cli = - Cli::try_parse_from(["devo", "server", "--transport", "stdio"]).expect("parse server"); - - assert_eq!( - super::should_direct_dispatch_macos_server(&cli, /*suppress_server_tray*/ true), - false - ); - } - #[test] fn server_process_args_from_cli_skips_non_server_command() { let cli = Cli::try_parse_from(["devo", "doctor"]).expect("parse doctor"); diff --git a/crates/client/src/client_core.rs b/crates/client/src/client_core.rs index 762f9b5d..7adbb003 100644 --- a/crates/client/src/client_core.rs +++ b/crates/client/src/client_core.rs @@ -331,14 +331,19 @@ impl ServerClientCore { Ok(success.result) } - pub(crate) async fn turn_start(&mut self, params: TurnStartParams) -> Result<()> { + pub(crate) async fn turn_start(&mut self, params: TurnStartParams) -> Result { match self .request_devo::<_, TurnStartResult>("turn/start", params.clone()) .await { - Ok(_) => Ok(()), + Ok(result) => Ok(result), Err(error) if is_method_not_found_error(&error) => { - self.turn_start_acp_prompt_detached(params).await + self.turn_start_acp_prompt_detached(params).await?; + Ok(TurnStartResult::Started { + turn_id: TurnId::new(), + status: TurnStatus::Running, + accepted_at: Utc::now(), + }) } Err(error) => Err(error), } diff --git a/crates/client/src/stdio.rs b/crates/client/src/stdio.rs index 3f02c866..b64f8dce 100644 --- a/crates/client/src/stdio.rs +++ b/crates/client/src/stdio.rs @@ -130,10 +130,12 @@ use devo_protocol::SkillSetEnabledParams; use devo_protocol::SkillSetEnabledResult; use devo_protocol::SpawnAgentParams; use devo_protocol::SpawnAgentResult; +use devo_protocol::TurnId; use devo_protocol::TurnInterruptParams; use devo_protocol::TurnInterruptResult; use devo_protocol::TurnStartParams; use devo_protocol::TurnStartResult; +use devo_protocol::TurnStatus; use devo_protocol::TurnSteerParams; use devo_protocol::TurnSteerResult; use devo_protocol::devo_extension_inner_method; @@ -166,7 +168,6 @@ type PendingResponses = Arc, - pub client_capabilities: devo_protocol::AcpClientCapabilities, } #[derive(Debug, Clone)] @@ -182,7 +183,6 @@ pub struct StdioServerClient { acp_pending_permissions: AcpPendingPermissions, acp_terminals: AcpTerminalManager, acp_agent_capabilities: Option, - client_capabilities: devo_protocol::AcpClientCapabilities, next_request_id: AtomicU64, notifications_rx: mpsc::UnboundedReceiver, notifications_tx: mpsc::UnboundedSender, @@ -234,22 +234,24 @@ impl StdioServerClient { acp_pending_permissions, acp_terminals, acp_agent_capabilities: None, - client_capabilities: config.client_capabilities, next_request_id: AtomicU64::new(1), notifications_rx, notifications_tx, }) } - pub async fn initialize(&mut self) -> Result { + pub async fn initialize( + &mut self, + client_capabilities: &devo_protocol::AcpClientCapabilities, + ) -> Result { tracing::info!("initializing stdio server client"); let result: AcpInitializeResult = timeout( - Duration::from_secs(10), + Duration::from_secs(3), self.request( ACP_INITIALIZE_METHOD, AcpInitializeParams { protocol_version: 1, - client_capabilities: self.client_capabilities.clone(), + client_capabilities: client_capabilities.clone(), client_info: Some( AcpImplementation::new("devo", env!("CARGO_PKG_VERSION")) .with_title("Devo"), @@ -581,14 +583,19 @@ impl StdioServerClient { self.request_devo("command/exec/terminate", params).await } - pub async fn turn_start(&mut self, params: TurnStartParams) -> Result<()> { + pub async fn turn_start(&mut self, params: TurnStartParams) -> Result { match self .request_devo::<_, TurnStartResult>("turn/start", params.clone()) .await { - Ok(_) => Ok(()), + Ok(result) => Ok(result), Err(error) if is_method_not_found_error(&error) => { - self.turn_start_acp_prompt_detached(params).await + self.turn_start_acp_prompt_detached(params).await?; + Ok(TurnStartResult::Started { + turn_id: TurnId::new(), + status: TurnStatus::Running, + accepted_at: Utc::now(), + }) } Err(error) => Err(error), } @@ -1367,15 +1374,16 @@ mod tests { acp_pending_permissions, acp_terminals, acp_agent_capabilities: None, - client_capabilities: client_capabilities.clone(), next_request_id: AtomicU64::new(1), notifications_rx, notifications_tx, }; let mut stdout_lines = BufReader::new(stdout).lines(); + let expected_capabilities = + serde_json::to_value(&client_capabilities).expect("serialize client capabilities"); let initialize = tokio::spawn(async move { - let result = client.initialize().await; + let result = client.initialize(&client_capabilities).await; (result, client) }); @@ -1384,7 +1392,7 @@ mod tests { assert_eq!(request["params"]["protocolVersion"], serde_json::json!(1)); assert_eq!( request["params"]["clientCapabilities"], - serde_json::to_value(&client_capabilities).expect("serialize client capabilities") + expected_capabilities ); pending .lock() @@ -1421,7 +1429,6 @@ mod tests { acp_pending_permissions, acp_terminals, acp_agent_capabilities: None, - client_capabilities: Default::default(), next_request_id: AtomicU64::new(1), notifications_rx, notifications_tx, @@ -1518,7 +1525,6 @@ mod tests { }, ..Default::default() }), - client_capabilities: Default::default(), next_request_id: AtomicU64::new(1), notifications_rx, notifications_tx, @@ -1626,7 +1632,6 @@ mod tests { }, ..Default::default() }), - client_capabilities: Default::default(), next_request_id: AtomicU64::new(1), notifications_rx, notifications_tx, @@ -1756,7 +1761,6 @@ mod tests { acp_pending_permissions, acp_terminals, acp_agent_capabilities: None, - client_capabilities: Default::default(), next_request_id: AtomicU64::new(1), notifications_rx, notifications_tx, @@ -1839,7 +1843,6 @@ mod tests { acp_pending_permissions, acp_terminals, acp_agent_capabilities: None, - client_capabilities: Default::default(), next_request_id: AtomicU64::new(1), notifications_rx, notifications_tx, diff --git a/crates/client/src/websocket.rs b/crates/client/src/websocket.rs index a9d67d27..a718a2ca 100644 --- a/crates/client/src/websocket.rs +++ b/crates/client/src/websocket.rs @@ -293,7 +293,7 @@ impl WebSocketServerClient { .await } - pub async fn turn_start(&mut self, params: TurnStartParams) -> Result<()> { + pub async fn turn_start(&mut self, params: TurnStartParams) -> Result { self.core.turn_start(params).await } diff --git a/crates/core/src/session.rs b/crates/core/src/session.rs index 1a9956b0..957c6a44 100644 --- a/crates/core/src/session.rs +++ b/crates/core/src/session.rs @@ -291,6 +291,33 @@ impl SessionState { } } + /// Clones session state for read-only export without active turn bookkeeping. + pub fn snapshot_for_export(&self) -> Self { + Self { + id: self.id.clone(), + config: self.config.clone(), + messages: self.messages.clone(), + prompt_messages: self.prompt_messages.clone(), + session_context: self.session_context.clone(), + latest_turn_context: self.latest_turn_context.clone(), + active_goal: self.active_goal.clone(), + collaboration_mode: self.collaboration_mode, + cwd: self.cwd.clone(), + turn_count: self.turn_count, + total_input_tokens: self.total_input_tokens, + total_output_tokens: self.total_output_tokens, + total_tokens: self.total_tokens, + total_cache_creation_tokens: self.total_cache_creation_tokens, + total_cache_read_tokens: self.total_cache_read_tokens, + prompt_token_estimate: self.prompt_token_estimate, + last_input_tokens: self.last_input_tokens, + last_turn_tokens: self.last_turn_tokens, + pending_turn_queue: Arc::clone(&self.pending_turn_queue), + btw_input_queue: Arc::clone(&self.btw_input_queue), + turn_state: None, + } + } + pub fn push_message(&mut self, msg: Message) { if let Some(prompt_messages) = self.prompt_messages.as_mut() { self.messages.push(msg.clone()); diff --git a/crates/core/src/tools/handlers/webfetch.rs b/crates/core/src/tools/handlers/webfetch.rs index de68cd71..f75786f7 100644 --- a/crates/core/src/tools/handlers/webfetch.rs +++ b/crates/core/src/tools/handlers/webfetch.rs @@ -1,17 +1,23 @@ use async_trait::async_trait; use base64::Engine; -use tokio::time::{Duration, timeout}; +use tokio::time::Duration; +use tokio::time::timeout; -use crate::contracts::{ - ToolCallError, ToolContext, ToolProgressSender, ToolResult, ToolResultContent, -}; +use crate::contracts::ToolCallError; +use crate::contracts::ToolContext; +use crate::contracts::ToolProgressSender; +use crate::contracts::ToolResult; +use crate::contracts::ToolResultContent; use crate::json_schema::JsonSchema; use crate::tool_handler::ToolHandler; -use crate::tool_spec::{ToolCapabilityTag, ToolExecutionMode, ToolOutputMode, ToolSpec}; +use crate::tool_spec::ToolCapabilityTag; +use crate::tool_spec::ToolExecutionMode; +use crate::tool_spec::ToolOutputMode; +use crate::tool_spec::ToolSpec; const MAX_RESPONSE_SIZE: usize = 5 * 1024 * 1024; -const DEFAULT_TIMEOUT_MS: u64 = 30_000; -const MAX_TIMEOUT_MS: u64 = 120_000; +const DEFAULT_TIMEOUT_SECS: u64 = 30; +const MAX_TIMEOUT_SECS: u64 = 120; const WEBFETCH_DESCRIPTION: &str = include_str!("../webfetch.txt"); pub struct WebFetchHandler { @@ -56,7 +62,7 @@ impl WebFetchHandler { supports_parallel: true, preparation_feedback: crate::tool_spec::ToolPreparationFeedback::None, display_name: None, - supports_cancellation: None, + supports_cancellation: Some(true), supports_streaming: None, }, } @@ -85,11 +91,9 @@ impl ToolHandler for WebFetchHandler { } let format = input["format"].as_str().unwrap_or("markdown"); - let timeout_ms = input["timeout"] - .as_u64() - .unwrap_or(DEFAULT_TIMEOUT_MS / 1000) - .saturating_mul(1000) - .min(MAX_TIMEOUT_MS); + let timeout_secs = parse_timeout_secs(&input); + let timeout_ms = timeout_secs.saturating_mul(1_000); + let timeout_duration = Duration::from_millis(timeout_ms); let accept = match format { "markdown" => { @@ -107,7 +111,9 @@ impl ToolHandler for WebFetchHandler { no_proxy: ctx.network_no_proxy.clone(), }; let client = devo_network_proxy::apply_proxy_config( - reqwest::Client::builder(), + reqwest::Client::builder() + .connect_timeout(devo_provider::timeout::connect_timeout()) + .timeout(timeout_duration), &proxy_config, ) .map_err(|e| { @@ -117,16 +123,17 @@ impl ToolHandler for WebFetchHandler { .build() .map_err(|e| ToolCallError::ExecutionFailed(format!("Failed to create HTTP client: {e}")))?; - let request = client - .get(url) - .header(reqwest::header::ACCEPT, accept) - .header(reqwest::header::ACCEPT_LANGUAGE, "en-US,en;q=0.9"); - - let fetch_result = timeout(Duration::from_millis(timeout_ms), async { - let response = request + let cancel_token = ctx.cancel_token.clone(); + let url = url.to_string(); + let format = format.to_string(); + let operation = async move { + let response = client + .get(&url) + .header(reqwest::header::ACCEPT, accept) + .header(reqwest::header::ACCEPT_LANGUAGE, "en-US,en;q=0.9") .send() .await - .map_err(|e| ToolCallError::ExecutionFailed(format!("Request failed: {e}")))?; + .map_err(|error| map_webfetch_request_error(error, timeout_secs))?; if !response.status().is_success() { let msg = format!("Request failed with status code: {}", response.status()); return Err(ToolCallError::ExecutionFailed(msg)); @@ -143,89 +150,148 @@ impl ToolHandler for WebFetchHandler { .and_then(|value| value.to_str().ok()) .unwrap_or("") .to_string(); - let bytes = response.bytes().await.map_err(|e| { - ToolCallError::ExecutionFailed(format!("Failed to read response: {e}")) - })?; - Ok((bytes, content_type)) + let bytes = response + .bytes() + .await + .map_err(|error| map_webfetch_request_error(error, timeout_secs))?; + if bytes.len() > MAX_RESPONSE_SIZE { + return Err(ToolCallError::ExecutionFailed("response too large".into())); + } + + let url = url.clone(); + let format = format.clone(); + let bytes = bytes.to_vec(); + tokio::task::spawn_blocking(move || { + build_webfetch_result(&url, bytes, &content_type, &format) + }) + .await + .map_err(|error| { + ToolCallError::ExecutionFailed(format!("response processing failed: {error}")) + })? + }; + + let fetch_result = timeout(timeout_duration, async { + tokio::select! { + result = operation => result, + () = cancel_token.cancelled() => Err(ToolCallError::Cancelled), + } }) .await; - let (bytes, content_type) = match fetch_result { - Ok(Ok(payload)) => payload, + match fetch_result { + Ok(Ok(result)) => Ok(result), Ok(Err(error)) => { + if matches!(error, ToolCallError::Cancelled) { + return Ok(ToolResult::error( + ToolResultContent::Text("Request cancelled".into()), + "Cancelled", + error, + )); + } + if matches!(error, ToolCallError::TimedOut(_)) { + return Ok(ToolResult::error( + ToolResultContent::Text("Request timed out".into()), + "Timeout", + error, + )); + } let message = error.to_string(); - return Ok(ToolResult::error( + Ok(ToolResult::error( ToolResultContent::Text(message.clone()), "HTTP error", error, - )); - } - Err(_) => { - return Ok(ToolResult::error( - ToolResultContent::Text("Request timed out".into()), - "Timeout", - ToolCallError::TimedOut(timeout_ms / 1000), - )); + )) } - }; - - if bytes.len() > MAX_RESPONSE_SIZE { - return Ok(ToolResult::error( - ToolResultContent::Text("Response too large (exceeds 5MB limit)".into()), - "Response too large", - ToolCallError::ExecutionFailed("response too large".into()), - )); + Err(_) => Ok(ToolResult::error( + ToolResultContent::Text("Request timed out".into()), + "Timeout", + ToolCallError::TimedOut(timeout_secs), + )), } + } +} - let mime = content_type - .split(';') - .next() - .unwrap_or("") - .trim() - .to_lowercase(); - let title = format!("{url} ({content_type})"); +fn build_webfetch_result( + url: &str, + bytes: Vec, + content_type: &str, + format: &str, +) -> Result { + let mime = content_type + .split(';') + .next() + .unwrap_or("") + .trim() + .to_lowercase(); + let title = format!("{url} ({content_type})"); - if is_image_mime(&mime) { - return Ok(ToolResult::success( - ToolResultContent::Mixed { - text: Some("Image fetched successfully".to_string()), - json: Some(serde_json::json!({ - "title": title, - "mime": mime, - "image_base64": base64::engine::general_purpose::STANDARD.encode(bytes), - })), - }, - "Image fetched", - )); - } + if is_image_mime(&mime) { + return Ok(ToolResult::success( + ToolResultContent::Mixed { + text: Some("Image fetched successfully".to_string()), + json: Some(serde_json::json!({ + "title": title, + "mime": mime, + "image_base64": base64::engine::general_purpose::STANDARD.encode(bytes), + })), + }, + "Image fetched", + )); + } - let content = String::from_utf8_lossy(&bytes).into_owned(); - let output = match format { - "text" => { - if content_type.contains("text/html") { - extract_text_from_html(&content) - } else { - content - } + let content = String::from_utf8_lossy(&bytes).into_owned(); + let output = match format { + "text" => { + if content_type.contains("text/html") { + extract_text_from_html(&content) + } else { + content } - "html" => content, - "markdown" => { - if content_type.contains("text/html") { - convert_html_to_markdown(&content) + } + "html" => content, + "markdown" => { + if content_type.contains("text/html") { + convert_html_to_markdown(&content) + } else { + content + } + } + _ => content, + }; + + Ok(ToolResult::success( + ToolResultContent::Mixed { + text: Some(output), + json: Some(serde_json::json!({ "title": title, "mime": mime })), + }, + "Content fetched", + )) +} + +fn parse_timeout_secs(input: &serde_json::Value) -> u64 { + let raw = &input["timeout"]; + let secs = raw + .as_u64() + .or_else(|| raw.as_i64().and_then(|value| u64::try_from(value).ok())) + .or_else(|| { + raw.as_f64().and_then(|value| { + if value.is_finite() && value > 0.0 { + Some(value.round() as u64) } else { - content + None } - } - _ => content, - }; + }) + }) + .or_else(|| raw.as_str().and_then(|value| value.parse().ok())) + .unwrap_or(DEFAULT_TIMEOUT_SECS); + secs.clamp(1, MAX_TIMEOUT_SECS) +} - Ok(ToolResult::success( - ToolResultContent::Mixed { - text: Some(output), - json: Some(serde_json::json!({ "title": title, "mime": mime })), - }, - "Content fetched", - )) +fn map_webfetch_request_error(error: reqwest::Error, timeout_secs: u64) -> ToolCallError { + if error.is_timeout() { + ToolCallError::TimedOut(timeout_secs) + } else { + ToolCallError::ExecutionFailed(format!("Request failed: {error}")) } } @@ -300,6 +366,25 @@ mod tests { use pretty_assertions::assert_eq; use super::extract_text_from_html; + use super::parse_timeout_secs; + + #[test] + fn parse_timeout_secs_accepts_integers_floats_and_strings() { + assert_eq!(parse_timeout_secs(&serde_json::json!({ "timeout": 5 })), 5); + assert_eq!( + parse_timeout_secs(&serde_json::json!({ "timeout": 5.0 })), + 5 + ); + assert_eq!( + parse_timeout_secs(&serde_json::json!({ "timeout": "10" })), + 10 + ); + assert_eq!(parse_timeout_secs(&serde_json::json!({})), 30); + assert_eq!( + parse_timeout_secs(&serde_json::json!({ "timeout": 999 })), + 120 + ); + } #[test] fn extract_text_from_html_skips_case_insensitive_script_blocks() { diff --git a/crates/network-proxy/src/lib.rs b/crates/network-proxy/src/lib.rs index 7f92609a..faa701a3 100644 --- a/crates/network-proxy/src/lib.rs +++ b/crates/network-proxy/src/lib.rs @@ -100,7 +100,8 @@ where .no_proxy(no_proxy.clone()); builder = builder.proxy(proxy); } - if let Some(proxy_url) = proxies.https.as_deref() { + let https_proxy_url = proxies.https.as_deref().or(proxies.http.as_deref()); + if let Some(proxy_url) = https_proxy_url { let proxy = reqwest::Proxy::https(proxy_url) .with_context(|| format!("invalid https_proxy URL `{proxy_url}`"))? .no_proxy(no_proxy.clone()); @@ -265,6 +266,37 @@ mod tests { ); } + #[tokio::test] + async fn http_proxy_env_applies_to_https_requests() { + let target_url = spawn_response_server("proxied").await; + let proxy_url = spawn_response_server("proxy").await; + let env = move |key: &str| -> std::result::Result { + match key { + "http_proxy" => Ok(proxy_url.clone()), + _ => Err(std::env::VarError::NotPresent), + } + }; + + let client = apply_proxy_with_env( + reqwest::Client::builder().timeout(Duration::from_secs(5)), + None, + env, + ) + .expect("proxy config") + .build() + .expect("client"); + let body = client + .get(target_url) + .send() + .await + .expect("proxied response") + .text() + .await + .expect("response body"); + + assert_eq!(body, "proxy"); + } + #[tokio::test] async fn no_proxy_bypasses_matching_env_proxy() { let target_url = spawn_response_server("direct").await; diff --git a/crates/provider/src/anthropic/messages.rs b/crates/provider/src/anthropic/messages.rs index 37031369..eca3ec41 100644 --- a/crates/provider/src/anthropic/messages.rs +++ b/crates/provider/src/anthropic/messages.rs @@ -21,7 +21,6 @@ use devo_protocol::StreamEvent; use devo_protocol::Usage; use devo_protocol::normalize_tool_result_messages; use futures::Stream; -use futures::StreamExt; use reqwest::Client; use reqwest::header::ACCEPT_ENCODING; use reqwest::header::CACHE_CONTROL; @@ -46,11 +45,13 @@ use crate::dsml::DsmlToolCallHealer; use crate::hosted_tools::append_anthropic_hosted_tools; use crate::http::invalid_status_error; use crate::merge_extra_body; +use crate::timeout; /// /// Anthropic provider backed by the official HTTP API. pub struct AnthropicProvider { client: Client, + streaming_client: Client, base_url: String, api_key: Option, http_options: ProviderHttpOptions, @@ -61,7 +62,10 @@ impl AnthropicProvider { let http_options = ProviderHttpOptions::default(); Self { client: http_options - .build_client(None) + .build_request_client() + .unwrap_or_else(|_| Client::new()), + streaming_client: http_options + .build_streaming_client() .unwrap_or_else(|_| Client::new()), base_url: base_url.into(), api_key: None, @@ -75,7 +79,8 @@ impl AnthropicProvider { } pub fn with_http_options(mut self, http_options: ProviderHttpOptions) -> Result { - self.client = http_options.build_client(None)?; + self.client = http_options.build_request_client()?; + self.streaming_client = http_options.build_streaming_client()?; self.http_options = http_options; Ok(self) } @@ -84,9 +89,8 @@ impl AnthropicProvider { format!("{}/v1/messages", self.base_url.trim_end_matches('/')) } - fn request_builder(&self, body: &Value) -> reqwest::RequestBuilder { - let builder = self - .client + fn post_builder(&self, client: &Client, body: &Value) -> reqwest::RequestBuilder { + let builder = client .post(self.endpoint()) .header("anthropic-version", "2023-06-01") .header(CONTENT_TYPE, HeaderValue::from_static("application/json")); @@ -101,8 +105,12 @@ impl AnthropicProvider { self.http_options.apply_custom_headers(builder).json(body) } + fn request_builder(&self, body: &Value) -> reqwest::RequestBuilder { + self.post_builder(&self.client, body) + } + fn streaming_request_builder(&self, body: &Value) -> reqwest::RequestBuilder { - self.request_builder(body) + self.post_builder(&self.streaming_client, body) .header(CACHE_CONTROL, HeaderValue::from_static("no-cache")) .header(ACCEPT_ENCODING, HeaderValue::from_static("identity")) } @@ -377,7 +385,17 @@ impl ModelProviderSDK for AnthropicProvider { let mut dsml_text_filters = BTreeMap::new(); futures::pin_mut!(event_source); - while let Some(event) = event_source.next().await { + loop { + let event = match timeout::next_eventsource_event(&mut event_source).await { + Ok(Some(event)) => event, + Ok(None) => break, + Err(idle) => { + Err(anyhow::anyhow!( + "anthropic stream idle timeout for model {}: {idle}", + request.model + ))? + } + }; let event = match event { Ok(event) => event, Err(reqwest_eventsource::Error::InvalidStatusCode(status, response)) => { diff --git a/crates/provider/src/http.rs b/crates/provider/src/http.rs index 773bcc68..38382fc0 100644 --- a/crates/provider/src/http.rs +++ b/crates/provider/src/http.rs @@ -1,5 +1,3 @@ -use std::time::Duration; - use anyhow::Context; use anyhow::Result; use devo_network_proxy::NetworkProxyConfig; @@ -46,16 +44,25 @@ impl ProviderHttpOptions { self.network_proxy.proxy_url.as_deref() } - pub(crate) fn build_client(&self, timeout: Option) -> Result { - let mut builder = Client::builder(); - if let Some(timeout) = timeout { - builder = builder.timeout(timeout); - } + /// HTTP client for non-streaming requests with total request timeout. + pub(crate) fn build_request_client(&self) -> Result { + let builder = Client::builder() + .connect_timeout(crate::timeout::connect_timeout()) + .timeout(crate::timeout::request_timeout()); devo_network_proxy::apply_proxy_config(builder, &self.network_proxy)? .build() .context("failed to build provider HTTP client") } + /// HTTP client for SSE streaming. Duration is bounded by per-chunk idle + /// timeout in the stream layer, not a single wall-clock request timeout. + pub(crate) fn build_streaming_client(&self) -> Result { + let builder = Client::builder().connect_timeout(crate::timeout::connect_timeout()); + devo_network_proxy::apply_proxy_config(builder, &self.network_proxy)? + .build() + .context("failed to build provider streaming HTTP client") + } + pub(crate) fn apply_custom_headers(&self, builder: RequestBuilder) -> RequestBuilder { if self.custom_headers.is_empty() { builder diff --git a/crates/provider/src/lib.rs b/crates/provider/src/lib.rs index 93b093b2..08b09417 100644 --- a/crates/provider/src/lib.rs +++ b/crates/provider/src/lib.rs @@ -14,6 +14,7 @@ mod provider; mod request; pub mod router; mod text_normalization; +pub mod timeout; pub use http::ProviderHttpOptions; pub use provider::*; diff --git a/crates/provider/src/openai/chat_completions.rs b/crates/provider/src/openai/chat_completions.rs index 39c0aa86..f28cc6ef 100644 --- a/crates/provider/src/openai/chat_completions.rs +++ b/crates/provider/src/openai/chat_completions.rs @@ -45,6 +45,7 @@ use crate::text_normalization::split_tagged_text; /// Works with OpenAI chat-completion servers by changing the base URL. pub struct OpenAIProvider { client: Client, + streaming_client: Client, base_url: String, api_key: Option, http_options: ProviderHttpOptions, @@ -52,17 +53,13 @@ pub struct OpenAIProvider { impl OpenAIProvider { pub fn new(base_url: impl Into) -> Self { - // TODO: Should we pertain the env `DEV_REQUEST_TIMEOUT` here ? - // maybe move to config toml is a good choice. - let timeout_secs = std::env::var("DEVO_REQUEST_TIMEOUT") - .ok() - .and_then(|v| v.parse::().ok()) - .unwrap_or(300); - let timeout = std::time::Duration::from_secs(timeout_secs); let http_options = ProviderHttpOptions::default(); Self { client: http_options - .build_client(Some(timeout)) + .build_request_client() + .unwrap_or_else(|_| Client::new()), + streaming_client: http_options + .build_streaming_client() .unwrap_or_else(|_| Client::new()), base_url: base_url.into(), api_key: None, @@ -76,12 +73,8 @@ impl OpenAIProvider { } pub fn with_http_options(mut self, http_options: ProviderHttpOptions) -> Result { - let timeout_secs = std::env::var("DEVO_REQUEST_TIMEOUT") - .ok() - .and_then(|v| v.parse::().ok()) - .unwrap_or(300); - let timeout = std::time::Duration::from_secs(timeout_secs); - self.client = http_options.build_client(Some(timeout))?; + self.client = http_options.build_request_client()?; + self.streaming_client = http_options.build_streaming_client()?; self.http_options = http_options; Ok(self) } @@ -90,9 +83,8 @@ impl OpenAIProvider { format!("{}/chat/completions", self.base_url.trim_end_matches('/')) } - fn request_builder(&self, body: &Value) -> reqwest::RequestBuilder { - let builder = self - .client + fn post_builder(&self, client: &Client, body: &Value) -> reqwest::RequestBuilder { + let builder = client .post(self.endpoint()) .header(CONTENT_TYPE, "application/json"); let builder = if let Some(api_key) = &self.api_key { @@ -102,6 +94,14 @@ impl OpenAIProvider { }; self.http_options.apply_custom_headers(builder).json(body) } + + fn request_builder(&self, body: &Value) -> reqwest::RequestBuilder { + self.post_builder(&self.client, body) + } + + pub(super) fn streaming_request_builder(&self, body: &Value) -> reqwest::RequestBuilder { + self.post_builder(&self.streaming_client, body) + } } #[derive(Debug, Clone, Deserialize, Serialize)] diff --git a/crates/provider/src/openai/chat_completions/stream.rs b/crates/provider/src/openai/chat_completions/stream.rs index 4ac92834..7975c9d1 100644 --- a/crates/provider/src/openai/chat_completions/stream.rs +++ b/crates/provider/src/openai/chat_completions/stream.rs @@ -7,7 +7,7 @@ use devo_protocol::{ ModelRequest, ModelResponse, ResponseContent, ResponseExtra, ResponseMetadata, StopReason, StreamEvent, Usage, }; -use futures::{Stream, StreamExt}; +use futures::Stream; use reqwest_eventsource::{Event, EventSource}; use serde::Deserialize; use serde_json::Value; @@ -23,6 +23,7 @@ use crate::dsml::DsmlTextStreamFilter; use crate::dsml::DsmlToolCallHealer; use crate::http::invalid_status_error; use crate::text_normalization::{TaggedTextFragment, TaggedTextParser}; +use crate::timeout; /// /// Represents a streamed chunk of a chat completion response returned by the model, based on the provided input. @@ -79,13 +80,23 @@ pub(super) async fn completion_stream( "sending openai streaming request" ); - let event_source = EventSource::new(provider.request_builder(&body)) + let event_source = EventSource::new(provider.streaming_request_builder(&body)) .context("failed to create openai event source")?; let stream = async_stream::try_stream! { let mut state = ChatCompletionStreamState::for_request(&request); futures::pin_mut!(event_source); - while let Some(event) = event_source.next().await { + loop { + let event = match timeout::next_eventsource_event(&mut event_source).await { + Ok(Some(event)) => event, + Ok(None) => break, + Err(idle) => { + Err(anyhow::anyhow!( + "openai stream idle timeout for model {}: {idle}", + request.model + ))? + } + }; let event = match event { Ok(event) => event, Err(reqwest_eventsource::Error::InvalidStatusCode(status, response)) => { diff --git a/crates/provider/src/openai/responses.rs b/crates/provider/src/openai/responses.rs index 5a80d69a..e1c29d1e 100644 --- a/crates/provider/src/openai/responses.rs +++ b/crates/provider/src/openai/responses.rs @@ -6,7 +6,7 @@ use devo_protocol::{ ModelRequest, ModelResponse, RequestContent, ResponseContent, ResponseExtra, ResponseMetadata, StopReason, StreamEvent, Usage, }; -use futures::{Stream, StreamExt}; +use futures::Stream; use reqwest::Client; use reqwest::header::{AUTHORIZATION, CONTENT_TYPE}; use reqwest_eventsource::{Event, EventSource}; @@ -16,6 +16,7 @@ use tracing::debug; use crate::hosted_tools::append_openai_responses_hosted_tools; use crate::http::invalid_status_error; use crate::text_normalization::{TaggedTextFragment, TaggedTextParser, split_tagged_text}; +use crate::timeout; use crate::{ModelProviderSDK, ProviderHttpOptions, merge_extra_body}; use super::capabilities::{OpenAITransport, resolve_request_profile}; @@ -30,6 +31,7 @@ use super::{ /// chat-completions adapter so the transport can evolve independently. pub struct OpenAIResponsesProvider { client: Client, + streaming_client: Client, base_url: String, api_key: Option, http_options: ProviderHttpOptions, @@ -40,7 +42,10 @@ impl OpenAIResponsesProvider { let http_options = ProviderHttpOptions::default(); Self { client: http_options - .build_client(None) + .build_request_client() + .unwrap_or_else(|_| Client::new()), + streaming_client: http_options + .build_streaming_client() .unwrap_or_else(|_| Client::new()), base_url: base_url.into(), api_key: None, @@ -54,7 +59,8 @@ impl OpenAIResponsesProvider { } pub fn with_http_options(mut self, http_options: ProviderHttpOptions) -> Result { - self.client = http_options.build_client(None)?; + self.client = http_options.build_request_client()?; + self.streaming_client = http_options.build_streaming_client()?; self.http_options = http_options; Ok(self) } @@ -63,9 +69,8 @@ impl OpenAIResponsesProvider { format!("{}/responses", self.base_url.trim_end_matches('/')) } - fn request_builder(&self, body: &Value) -> reqwest::RequestBuilder { - let builder = self - .client + fn post_builder(&self, client: &Client, body: &Value) -> reqwest::RequestBuilder { + let builder = client .post(self.endpoint()) .header(CONTENT_TYPE, "application/json"); let builder = if let Some(api_key) = &self.api_key { @@ -75,6 +80,14 @@ impl OpenAIResponsesProvider { }; self.http_options.apply_custom_headers(builder).json(body) } + + fn request_builder(&self, body: &Value) -> reqwest::RequestBuilder { + self.post_builder(&self.client, body) + } + + fn streaming_request_builder(&self, body: &Value) -> reqwest::RequestBuilder { + self.post_builder(&self.streaming_client, body) + } } /// Builds the exact OpenAI Responses request body used by this provider. @@ -515,7 +528,7 @@ impl ModelProviderSDK for OpenAIResponsesProvider { "sending openai responses streaming request" ); - let event_source = EventSource::new(self.request_builder(&body)) + let event_source = EventSource::new(self.streaming_request_builder(&body)) .context("failed to create openai responses event source")?; let stream = async_stream::try_stream! { let mut text_buf = String::new(); @@ -530,7 +543,17 @@ impl ModelProviderSDK for OpenAIResponsesProvider { let mut text_started = false; futures::pin_mut!(event_source); - while let Some(event) = event_source.next().await { + loop { + let event = match timeout::next_eventsource_event(&mut event_source).await { + Ok(Some(event)) => event, + Ok(None) => break, + Err(idle) => { + Err(anyhow::anyhow!( + "openai responses stream idle timeout for model {}: {idle}", + request.model + ))? + } + }; let event = match event { Ok(event) => event, Err(reqwest_eventsource::Error::InvalidStatusCode(status, response)) => { diff --git a/crates/provider/src/timeout.rs b/crates/provider/src/timeout.rs new file mode 100644 index 00000000..8bf27371 --- /dev/null +++ b/crates/provider/src/timeout.rs @@ -0,0 +1,89 @@ +//! Provider HTTP and SSE stream timeout configuration. +//! +//! Non-streaming requests use a total [`request_timeout`]. Streaming responses +//! use a per-chunk idle [`stream_idle_timeout`] that resets whenever a new SSE +//! event arrives, so long generations are not cut off by a fixed wall-clock cap. + +use std::time::Duration; + +use futures::StreamExt; +use reqwest_eventsource::{Event, EventSource}; + +/// Total wall-clock timeout for non-streaming provider HTTP requests. +pub const REQUEST_TIMEOUT_SECS: u64 = 120; + +/// TCP/TLS connect timeout for provider HTTP clients. +pub const CONNECT_TIMEOUT_SECS: u64 = 30; + +/// Maximum idle time between consecutive SSE events during streaming. +pub const STREAM_IDLE_TIMEOUT_SECS: u64 = 120; + +/// Total timeout for non-streaming provider HTTP requests. +#[inline] +pub fn request_timeout() -> Duration { + Duration::from_secs(REQUEST_TIMEOUT_SECS) +} + +/// TCP/TLS connect timeout for provider HTTP clients. +#[inline] +pub fn connect_timeout() -> Duration { + Duration::from_secs(CONNECT_TIMEOUT_SECS) +} + +/// Idle timeout between consecutive SSE events during streaming. +#[inline] +pub fn stream_idle_timeout() -> Duration { + Duration::from_secs(STREAM_IDLE_TIMEOUT_SECS) +} + +/// Waits for the next SSE event, failing when no data arrives within +/// [`stream_idle_timeout`]. +pub async fn next_eventsource_event( + event_source: &mut EventSource, +) -> Result>, StreamIdleTimeoutError> { + match tokio::time::timeout(stream_idle_timeout(), event_source.next()).await { + Ok(event) => Ok(event), + Err(_) => Err(StreamIdleTimeoutError { + idle_timeout: stream_idle_timeout(), + }), + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct StreamIdleTimeoutError { + pub idle_timeout: Duration, +} + +impl std::fmt::Display for StreamIdleTimeoutError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + formatter, + "provider stream idle timeout after {}s without receiving data", + self.idle_timeout.as_secs() + ) + } +} + +impl std::error::Error for StreamIdleTimeoutError {} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn request_timeout_is_two_minutes() { + assert_eq!(request_timeout(), Duration::from_secs(120)); + } + + #[test] + fn connect_timeout_is_thirty_seconds() { + assert_eq!(connect_timeout(), Duration::from_secs(30)); + } + + #[test] + fn stream_idle_timeout_is_two_minutes() { + assert_eq!(stream_idle_timeout(), Duration::from_secs(120)); + } +} diff --git a/crates/server/Cargo.toml b/crates/server/Cargo.toml index 10231ef5..91ea2059 100644 --- a/crates/server/Cargo.toml +++ b/crates/server/Cargo.toml @@ -42,17 +42,6 @@ tracing = { workspace = true } tracing-subscriber = { workspace = true } uuid = { workspace = true } -[target.'cfg(any(windows, target_os = "macos"))'.dependencies] -image = { workspace = true, features = ["png"] } -tray-icon = { workspace = true } - -[target.'cfg(windows)'.dependencies] -windows-sys = { workspace = true } - -[target.'cfg(target_os = "macos")'.dependencies] -objc2-foundation = { workspace = true } -winit = { workspace = true } - [dev-dependencies] async-trait = { workspace = true } pretty_assertions = "1" diff --git a/crates/server/src/bootstrap.rs b/crates/server/src/bootstrap.rs index 616c0bb0..438b4951 100644 --- a/crates/server/src/bootstrap.rs +++ b/crates/server/src/bootstrap.rs @@ -1,5 +1,4 @@ use std::sync::Arc; -use std::sync::mpsc::Sender; use anyhow::Result; use clap::Parser; @@ -22,8 +21,6 @@ use crate::execution::ServerRuntimeDependencies; use crate::load_server_provider; use crate::resolve_listen_targets; use crate::run_listeners_with_internal_proxy; -#[cfg(windows)] -use crate::server_tray::ServerTray; use crate::singleton::ServerControlAction; use crate::singleton::SingletonRole; use crate::singleton::acquire_singleton_role; @@ -33,9 +30,6 @@ use crate::transport::DEFAULT_WEBSOCKET_BIND_ADDRESS; use crate::transport::InternalProxyControl; use crate::transport::InternalProxyEndpoint; -/// Matches `devo_arg0::SUPPRESS_SERVER_TRAY_ENV` without adding a dependency cycle. -const SUPPRESS_SERVER_TRAY_ENV: &str = "DEVO_SUPPRESS_SERVER_TRAY"; - #[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] pub enum ServerTransportMode { Config, @@ -79,56 +73,16 @@ impl ServerProcessArgs { } } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum ServerTrayStartup { - Enabled, - Disabled, -} - -pub(crate) struct ServerProcessRunOptions { - pub(crate) external_shutdown: Option, - pub(crate) tray_startup: ServerTrayStartup, - pub(crate) real_server_started: Option>, -} - -impl Default for ServerProcessRunOptions { - fn default() -> Self { - Self { - external_shutdown: None, - tray_startup: server_tray_startup_from_env(), - real_server_started: None, - } - } -} - -fn server_tray_startup_from_env() -> ServerTrayStartup { - let value = std::env::var(SUPPRESS_SERVER_TRAY_ENV).ok(); - server_tray_startup_for(value.as_deref() == Some("1")) -} - -fn server_tray_startup_for(suppress_server_tray: bool) -> ServerTrayStartup { - if suppress_server_tray { - ServerTrayStartup::Disabled - } else { - ServerTrayStartup::Enabled - } +#[derive(Default)] +pub struct ServerProcessRunOptions { + pub external_shutdown: Option, } /// Starts the transport-facing server runtime using the resolved application /// configuration and listener set. -pub async fn run_server_process(args: ServerProcessArgs) -> Result<()> { - run_server_process_inner(args, ServerProcessRunOptions::default()).await -} - -/// Starts the macOS server process with AppKit owned by the process main thread. -#[cfg(target_os = "macos")] -pub fn run_server_process_with_macos_tray(args: ServerProcessArgs) -> Result<()> { - crate::server_tray::run_server_process_with_macos_tray(args) -} - -pub(crate) async fn run_server_process_inner( +pub async fn run_server_process( args: ServerProcessArgs, - mut options: ServerProcessRunOptions, + options: ServerProcessRunOptions, ) -> Result<()> { let resolver = FileSystemConfigPathResolver::from_env()?; let action = args.action()?; @@ -265,80 +219,30 @@ pub(crate) async fn run_server_process_inner( tracing::info!("starting persisted session restore"); runtime.load_persisted_sessions().await?; tracing::info!("persisted session restore completed"); - tracing::info!("server bootstrap completed; starting listeners"); - if let Some(real_server_started) = options.real_server_started.take() { - let _ = real_server_started.send(()); - } + let shutdown_signal = tokio_util::sync::CancellationToken::new(); let internal_proxy_control = InternalProxyControl::new(shutdown_signal.clone()); let external_shutdown = options.external_shutdown.clone(); - #[cfg(not(windows))] - match options.tray_startup { - ServerTrayStartup::Enabled | ServerTrayStartup::Disabled => {} - } - - #[cfg(windows)] - let mut server_tray = match options.tray_startup { - ServerTrayStartup::Enabled => match ServerTray::start() { - Ok(tray) => Some(tray), - Err(error) => { - tracing::warn!(%error, "failed to start server tray icon"); - None - } - }, - ServerTrayStartup::Disabled => None, - }; - #[cfg(windows)] - { - tokio::select! { - result = run_listeners_with_internal_proxy( - runtime.clone(), - &effective_listen, - internal_proxy, - singleton_metadata.token.clone(), - internal_proxy_control, - ) => { - result?; - } - result = tokio::signal::ctrl_c() => { - result?; - tracing::info!("server shutdown requested"); - } - _ = wait_for_server_tray_shutdown(&mut server_tray) => { - tracing::info!("server shutdown requested from tray icon"); - } - _ = wait_for_external_shutdown(external_shutdown.as_ref()) => { - tracing::info!("server shutdown requested from external process controller"); - } - _ = shutdown_signal.cancelled() => { - tracing::info!("server shutdown requested from singleton control"); - } + tokio::select! { + result = run_listeners_with_internal_proxy( + runtime.clone(), + &effective_listen, + internal_proxy, + singleton_metadata.token.clone(), + internal_proxy_control, + ) => { + result?; } - } - - #[cfg(not(windows))] - { - tokio::select! { - result = run_listeners_with_internal_proxy( - runtime.clone(), - &effective_listen, - internal_proxy, - singleton_metadata.token.clone(), - internal_proxy_control, - ) => { - result?; - } - result = tokio::signal::ctrl_c() => { - result?; - tracing::info!("server shutdown requested"); - } - _ = wait_for_external_shutdown(external_shutdown.as_ref()) => { - tracing::info!("server shutdown requested from external process controller"); - } - _ = shutdown_signal.cancelled() => { - tracing::info!("server shutdown requested from singleton control"); - } + result = tokio::signal::ctrl_c() => { + result?; + tracing::info!("server shutdown requested"); + } + _ = wait_for_external_shutdown(external_shutdown.as_ref()) => { + tracing::info!("server shutdown requested from external process controller"); + } + _ = shutdown_signal.cancelled() => { + tracing::info!("server shutdown requested from singleton control"); } } @@ -356,15 +260,6 @@ fn print_existing_server_status(metadata: &crate::singleton::ServerLockMetadata, println!("started_at: {}", metadata.started_at); } -#[cfg(windows)] -async fn wait_for_server_tray_shutdown(server_tray: &mut Option) { - if let Some(tray) = server_tray { - tray.shutdown_requested().await; - } else { - std::future::pending::<()>().await; - } -} - async fn wait_for_external_shutdown( external_shutdown: Option<&tokio_util::sync::CancellationToken>, ) { @@ -394,20 +289,6 @@ mod tests { ); } - #[test] - fn server_tray_startup_honors_desktop_suppression() { - assert_eq!( - [ - super::server_tray_startup_for(/*suppress_server_tray*/ false), - super::server_tray_startup_for(/*suppress_server_tray*/ true), - ], - [ - super::ServerTrayStartup::Enabled, - super::ServerTrayStartup::Disabled, - ] - ); - } - #[test] fn server_process_args_accept_stdio_transport_override() { let args = ServerProcessArgs::parse_from(["devo-server", "--transport", "stdio"]); diff --git a/crates/server/src/execution.rs b/crates/server/src/execution.rs index c0103937..1ebf74c0 100644 --- a/crates/server/src/execution.rs +++ b/crates/server/src/execution.rs @@ -1,4 +1,3 @@ -use std::collections::HashMap; use std::collections::HashSet; use std::collections::VecDeque; use std::path::Path; @@ -60,7 +59,7 @@ pub(crate) struct PendingUserInput { pub(crate) tx: oneshot::Sender, } -#[derive(Default)] +#[derive(Clone, Default)] pub(crate) struct ApprovalGrantCache { pub(crate) tools: HashSet, pub(crate) hosts: HashSet, @@ -227,10 +226,6 @@ pub(crate) struct RuntimeSession { pub(crate) next_item_seq: u64, /// First user input captured from the session's first turn, used for title generation. pub(crate) first_user_input: Option, - /// Active approval requests waiting for client decisions. - pub(crate) pending_approvals: HashMap, - /// Active request_user_input calls waiting for client answers. - pub(crate) pending_user_inputs: HashMap, /// Session-specific tool registry, used when the session was created with /// request-scoped tool sources such as ACP MCP servers. pub(crate) tool_registry: Option>, @@ -240,13 +235,6 @@ pub(crate) struct RuntimeSession { pub(crate) turn_approval_cache: ApprovalGrantCache, } -impl RuntimeSession { - /// Wraps a new runtime session in an async mutex for storage in the session map. - pub(crate) fn shared(self) -> Arc> { - Arc::new(Mutex::new(self)) - } -} - #[cfg(test)] mod tests { use std::pin::Pin; diff --git a/crates/server/src/lib.rs b/crates/server/src/lib.rs index a1434123..12fb4014 100644 --- a/crates/server/src/lib.rs +++ b/crates/server/src/lib.rs @@ -13,8 +13,6 @@ mod projection; mod protocol; mod provider_config; mod runtime; -#[cfg(any(windows, target_os = "macos"))] -mod server_tray; mod session; mod session_context; mod singleton; diff --git a/crates/server/src/persistence.rs b/crates/server/src/persistence.rs index b05c308a..45e9bc65 100644 --- a/crates/server/src/persistence.rs +++ b/crates/server/src/persistence.rs @@ -341,12 +341,12 @@ impl RolloutStore { pub(crate) async fn load_sessions( &self, deps: &ServerRuntimeDependencies, - ) -> Result>>> { + ) -> Result> { let mut sessions = HashMap::new(); for rollout_path in self.rollout_paths()? { match self.load_session_from_rollout(&rollout_path, deps).await { Ok(recovered) => { - sessions.insert(recovered.summary.session_id, recovered.shared()); + sessions.insert(recovered.summary.session_id, recovered); } Err(error) => { tracing::warn!( @@ -916,8 +916,6 @@ impl ReplayState { deferred_reasoning: None, next_item_seq: self.next_item_seq.max(1), first_user_input: None, - pending_approvals: std::collections::HashMap::new(), - pending_user_inputs: std::collections::HashMap::new(), tool_registry: None, session_approval_cache: crate::execution::ApprovalGrantCache::default(), turn_approval_cache: crate::execution::ApprovalGrantCache::default(), diff --git a/crates/server/src/runtime.rs b/crates/server/src/runtime.rs index 35d8cc2c..98cd11b6 100644 --- a/crates/server/src/runtime.rs +++ b/crates/server/src/runtime.rs @@ -49,8 +49,6 @@ use devo_core::tools::PermissionChecker; use devo_core::tools::ToolAgentScope; use devo_core::tools::ToolCall; use devo_core::tools::ToolCallError; -use devo_core::tools::ToolCallResult; -use devo_core::tools::ToolContent; use devo_core::tools::ToolExecutionOptions; use devo_core::tools::ToolPermissionRequest; use devo_core::tools::ToolRegistry; @@ -62,13 +60,11 @@ use devo_protocol::{ WorkspaceDiffDetail, }; use devo_safety::PermissionMode; -use devo_util_shell_command::parse_command::parse_command; use crate::ApprovalDecisionValue; use crate::ApprovalScopeValue; use crate::ClientMethod; use crate::ClientTransportKind; -use crate::CommandExecutionPayload; use crate::ConnectionState; use crate::ErrorResponse; use crate::EventContext; @@ -128,7 +124,6 @@ use crate::approval_reviewer::ReviewerDecision; use crate::approval_reviewer::build_approval_review_request; use crate::approval_reviewer::parse_reviewer_decision; use crate::db::QueueType; -use crate::db::SessionStats; use crate::execution::PendingApproval; use crate::execution::PendingUserInput; use crate::execution::RuntimeSession; @@ -181,6 +176,8 @@ mod research_stages; mod research_streaming; mod research_tool_runtime; mod research_tools; +mod session_actor; +mod session_interactive; mod skills; mod subagent_usage; mod turn_exec; @@ -196,18 +193,33 @@ pub(crate) use connection::SubscriptionFilter; pub(crate) use items::render_input_items; pub(crate) use research_tools::extract_written_file_path; pub(crate) use research_tools::is_write_tool_name; +use session_actor::SessionHandle; +use session_interactive::SessionInteractiveLanes; use turn_exec::ExecuteTurnRequest; +pub(crate) use session_actor::SessionActorState; + pub struct ServerRuntime { metadata: InitializeResult, deps: ServerRuntimeDependencies, rollout_store: RolloutStore, goal_durable_store: GoalDurableStore, - /// Thread safe hashmap as sessions container, there are allowed multiple sessions. - sessions: Mutex>>>, + /// Per-session actor handles; map lock must not be held across await. + sessions: Mutex>, + /// Interactive approval and user-input waits outside session actors. + session_interactive: SessionInteractiveLanes, + /// Spawn snapshots for in-flight parent turns keyed by session and turn id. + active_spawn_snapshots: + Mutex>>>, + /// Stream state shared with the turn event task while a session actor turn runs. + active_stream_states: Mutex< + HashMap>>, + >, connections: Arc>>, active_tasks: Mutex>, active_turn_cancellations: Mutex>, + /// Active turn ids tracked at runtime level so cancel/interrupt can avoid session-actor mailbox round-trips while a turn is blocked in permission wait. + active_turn_ids: Mutex>, active_turn_connections: Mutex>, terminal_turn_statuses: Mutex>, acp_prompt_waiters: Mutex>>>, @@ -235,10 +247,12 @@ pub struct ServerRuntime { command_exec_manager: command_exec::CommandExecManager, /// Turn-scoped workspace baselines captured at actual execution start. active_workspace_baselines: Mutex>, + /// Weak back-reference used when session actors need the owning runtime `Arc`. + self_weak: std::sync::Weak, } #[derive(Debug, Clone, PartialEq, Eq)] -enum TurnInputMode { +pub(crate) enum TurnInputMode { VisibleUserMessage, HiddenGoalContinuation { goal: devo_protocol::ThreadGoal }, } @@ -291,33 +305,10 @@ fn requested_model_selection<'a>( .or_else(|| session_model_selection(session)) } -fn apply_turn_config_to_session_summary(summary: &mut SessionMetadata, turn_config: &TurnConfig) { - summary.model = Some(turn_config.model.slug.clone()); - summary.model_binding_id = turn_config.model_binding_id.clone(); - summary.reasoning_effort_selection = turn_config.reasoning_effort_selection.clone(); -} - -fn string_field_from_pending_metadata( - metadata: Option<&serde_json::Value>, - key: &str, -) -> Option { - metadata? - .get(key)? - .as_str() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string) -} - -fn model_selection_from_pending_metadata(metadata: Option<&serde_json::Value>) -> Option { - string_field_from_pending_metadata(metadata, "model_binding_id") - .or_else(|| string_field_from_pending_metadata(metadata, "model")) -} - const SUBAGENT_USAGE_PARENT_SESSION_ID_METADATA: &str = "devo_subagent_usage_parent_session_id"; const SUBAGENT_USAGE_PARENT_TURN_ID_METADATA: &str = "devo_subagent_usage_parent_turn_id"; -fn subagent_usage_owner_pending_metadata( +pub(super) fn subagent_usage_owner_pending_metadata( parent_session_id: SessionId, parent_turn_id: Option, ) -> serde_json::Value { @@ -327,23 +318,11 @@ fn subagent_usage_owner_pending_metadata( }) } -fn subagent_usage_owner_from_pending_metadata( - metadata: Option<&serde_json::Value>, -) -> Option<(SessionId, Option)> { - let parent_session_id = - string_field_from_pending_metadata(metadata, SUBAGENT_USAGE_PARENT_SESSION_ID_METADATA) - .and_then(|value| SessionId::try_from(value).ok())?; - let parent_turn_id = - string_field_from_pending_metadata(metadata, SUBAGENT_USAGE_PARENT_TURN_ID_METADATA) - .and_then(|value| TurnId::try_from(value).ok()); - Some((parent_session_id, parent_turn_id)) -} - impl ServerRuntime { pub fn new(server_home: PathBuf, deps: ServerRuntimeDependencies) -> Arc { let rollout_store = RolloutStore::new(server_home.clone()); let goal_durable_store = GoalDurableStore::new(server_home.clone()); - Arc::new(Self { + Arc::new_cyclic(|self_weak| Self { metadata: InitializeResult { server_name: "devo-server".into(), server_version: env!("CARGO_PKG_VERSION").into(), @@ -355,9 +334,13 @@ impl ServerRuntime { rollout_store, goal_durable_store, sessions: Mutex::new(HashMap::new()), + session_interactive: SessionInteractiveLanes::default(), + active_spawn_snapshots: Mutex::new(HashMap::new()), + active_stream_states: Mutex::new(HashMap::new()), connections: Arc::new(Mutex::new(HashMap::new())), active_tasks: Mutex::new(HashMap::new()), active_turn_cancellations: Mutex::new(HashMap::new()), + active_turn_ids: Mutex::new(HashMap::new()), active_turn_connections: Mutex::new(HashMap::new()), terminal_turn_statuses: Mutex::new(VecDeque::new()), acp_prompt_waiters: Mutex::new(HashMap::new()), @@ -374,6 +357,7 @@ impl ServerRuntime { reference_searches: Mutex::new(HashMap::new()), command_exec_manager: command_exec::CommandExecManager::new(), active_workspace_baselines: Mutex::new(HashMap::new()), + self_weak: self_weak.clone(), }) } } diff --git a/crates/server/src/runtime/agents.rs b/crates/server/src/runtime/agents.rs index 9da543e5..08488174 100644 --- a/crates/server/src/runtime/agents.rs +++ b/crates/server/src/runtime/agents.rs @@ -60,39 +60,33 @@ impl ServerRuntime { )); } - let parent_arc = self.session_arc(parent_session_id).await?; - let parent_snapshot = { - let parent = parent_arc.lock().await; - let stable_items = if fork_turns == "all" { - let active_turn_id = parent.active_turn.as_ref().map(|turn| turn.turn_id); - parent - .persisted_turn_items - .iter() - .filter(|item| active_turn_id.is_none_or(|turn_id| item.turn_id != turn_id)) - .cloned() - .collect::>() - } else { - Vec::new() - }; - ( - parent.summary.clone(), - parent.config.clone(), - stable_items, - parent.latest_turn.clone(), - parent.active_turn.as_ref().map(|turn| turn.turn_id), - parent.tool_registry.clone(), - Arc::clone(&parent.runtime_context), - ) + let parent_handle = self.session(parent_session_id).await.ok_or_else(|| { + ToolCallError::InvalidInput(format!("session not found: {parent_session_id}")) + })?; + let parent_snapshot = if let Some(snapshot) = self + .active_spawn_snapshot_for_session(parent_session_id) + .await + { + snapshot + } else { + parent_handle.spawn_snapshot().await.ok_or_else(|| { + ToolCallError::InvalidInput(format!( + "failed to snapshot parent session: {parent_session_id}" + )) + })? }; - let ( - parent_summary, - parent_config, - stable_items, - parent_latest_turn, - parent_active_turn_id, - parent_tool_registry, - runtime_context, - ) = parent_snapshot; + let stable_items = if fork_turns == "all" { + parent_snapshot.stable_items + } else { + Vec::new() + }; + + let parent_summary = parent_snapshot.parent_summary; + let parent_config = parent_snapshot.parent_config; + let parent_latest_turn = parent_snapshot.parent_latest_turn; + let parent_active_turn_id = parent_snapshot.parent_active_turn_id; + let parent_tool_registry = parent_snapshot.parent_tool_registry; + let runtime_context = parent_snapshot.runtime_context; let parent_usage_turn_id = parent_active_turn_id.or_else(|| parent_latest_turn.as_ref().map(|turn| turn.turn_id)); @@ -233,16 +227,12 @@ impl ServerRuntime { deferred_reasoning: None, next_item_seq: 1, first_user_input: Some(params.message.clone()), - pending_approvals: HashMap::new(), - pending_user_inputs: std::collections::HashMap::new(), tool_registry: parent_tool_registry, session_approval_cache: crate::execution::ApprovalGrantCache::default(), turn_approval_cache: crate::execution::ApprovalGrantCache::default(), }; - self.sessions - .lock() - .await - .insert(child_session_id, child_session.shared()); + let child_state = SessionActorState::from_runtime_session(child_session); + self.insert_session_actor(child_state).await; self.agent_mailboxes .lock() .await @@ -302,6 +292,12 @@ impl ServerRuntime { let start_runtime = Arc::clone(self); let start_summary = summary; let start_message = params.message.clone(); + tracing::debug!( + parent_session_id = %parent_session_id, + child_session_id = %child_session_id, + agent_path = %agent_path, + "subagent startup task spawned" + ); tokio::spawn(async move { start_runtime .broadcast_event(ServerEvent::SessionStarted(SessionEventPayload { @@ -377,18 +373,6 @@ impl ServerRuntime { }) } - async fn session_arc( - &self, - session_id: SessionId, - ) -> Result>, ToolCallError> { - self.sessions - .lock() - .await - .get(&session_id) - .cloned() - .ok_or_else(|| ToolCallError::InvalidInput(format!("session not found: {session_id}"))) - } - async fn mailbox(&self, session_id: SessionId) -> SubagentMailbox { self.agent_mailboxes .lock() @@ -444,47 +428,47 @@ impl ServerRuntime { input_text: String, queued_metadata: Option, ) -> Result { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { - return Err(ToolCallError::InvalidInput(format!( - "session not found: {session_id}" - ))); - }; - let queued_active_turn = { - let session = session_arc.lock().await; - if session.max_turns.is_some_and(|max_turns| { - max_turns == 0 - || session - .active_turn - .as_ref() - .is_some_and(|turn| turn.sequence >= max_turns) - || session - .latest_turn - .as_ref() - .is_some_and(|turn| turn.sequence >= max_turns) - }) { - return Err(ToolCallError::InvalidInput( - "agent maximum turn count reached".to_string(), - )); - } - session.active_turn.as_ref().map(|turn| { - ( - turn.clone(), - Arc::clone(&session.pending_turn_queue), - session.summary.ephemeral, - ) - }) - }; - if let Some((active_turn, pending_turn_queue, is_ephemeral)) = queued_active_turn { - let item = devo_core::PendingInputItem::new( - devo_core::PendingInputKind::UserText { text: input_text }, + let session_handle = self.session(session_id).await.ok_or_else(|| { + ToolCallError::InvalidInput(format!("session not found: {session_id}")) + })?; + + let reservation = session_handle + .turn_reservation_snapshot() + .await + .ok_or_else(|| { + ToolCallError::InvalidInput(format!( + "failed to snapshot session reservation: {session_id}" + )) + })?; + + if reservation.max_turns.is_some_and(|max_turns| { + max_turns == 0 + || reservation + .active_turn + .as_ref() + .is_some_and(|turn| turn.sequence >= max_turns) + || reservation + .latest_turn + .as_ref() + .is_some_and(|turn| turn.sequence >= max_turns) + }) { + return Err(ToolCallError::InvalidInput( + "agent maximum turn count reached".to_string(), + )); + } + + if let Some(active_turn) = reservation.active_turn.clone() { + let item = devo_protocol::PendingInputItem::new( + devo_protocol::PendingInputKind::UserText { text: input_text }, queued_metadata, Utc::now(), ); - pending_turn_queue + reservation + .pending_turn_queue .lock() .expect("pending turn queue mutex should not be poisoned") .push_back(item.clone()); - if !is_ephemeral + if !reservation.ephemeral && let Err(error) = self .deps .db @@ -500,54 +484,50 @@ impl ServerRuntime { return Ok(active_turn); } - let (turn_config, resolved_request) = { - let session = session_arc.lock().await; - let turn_config = session.runtime_context.resolve_turn_config( - session_model_selection(&session.summary), - session.summary.reasoning_effort_selection.clone(), - ); - let resolved_request = turn_config.model.resolve_reasoning_effort_selection( - turn_config.reasoning_effort_selection.as_deref(), - ); - (turn_config, resolved_request) - }; + let turn_config = reservation.runtime_context.resolve_turn_config( + session_model_selection(&reservation.summary), + reservation.summary.reasoning_effort_selection.clone(), + ); + let resolved_request = turn_config + .model + .resolve_reasoning_effort_selection(turn_config.reasoning_effort_selection.as_deref()); + let request_model = turn_config.provider_request_model(&resolved_request.request_model); let now = Utc::now(); - let turn = { - let mut session = session_arc.lock().await; - let turn = TurnMetadata { - turn_id: TurnId::new(), - session_id, - sequence: session - .latest_turn - .as_ref() - .map_or(1, |turn| turn.sequence + 1), - status: TurnStatus::Running, - kind: devo_core::TurnKind::Regular, - model: turn_config.model.slug.clone(), - model_binding_id: turn_config.model_binding_id.clone(), - reasoning_effort_selection: turn_config.reasoning_effort_selection.clone(), - reasoning_effort: resolved_request.effective_reasoning_effort, - request_model, - request_thinking: resolved_request.request_thinking.clone(), - started_at: now, - completed_at: None, - usage: None, - stop_reason: None, - failure_reason: None, - }; - session.summary.status = SessionRuntimeStatus::ActiveTurn; - session.summary.updated_at = now; - session.summary.last_activity_at = now; - apply_turn_config_to_session_summary(&mut session.summary, &turn_config); - session.active_turn = Some(turn.clone()); - turn + let sequence = reservation + .latest_turn + .as_ref() + .map_or(1, |turn| turn.sequence + 1); + let turn = TurnMetadata { + turn_id: TurnId::new(), + session_id, + sequence, + status: TurnStatus::Running, + kind: devo_core::TurnKind::Regular, + model: turn_config.model.slug.clone(), + model_binding_id: turn_config.model_binding_id.clone(), + reasoning_effort_selection: turn_config.reasoning_effort_selection.clone(), + reasoning_effort: resolved_request.effective_reasoning_effort, + request_model, + request_thinking: resolved_request.request_thinking.clone(), + started_at: now, + completed_at: None, + usage: None, + stop_reason: None, + failure_reason: None, }; + + session_handle + .begin_active_turn(turn.clone(), turn_config.clone()) + .await; + if let Err(error) = self.append_turn_start(session_id, &turn).await { - self.clear_active_turn_reservation(&session_arc, turn.turn_id) + let _ = session_handle + .clear_active_turn_if_matches(turn.turn_id) .await; return Err(error); } + self.broadcast_event(ServerEvent::SessionStatusChanged( SessionStatusChangedPayload { session_id, @@ -560,15 +540,12 @@ impl ServerRuntime { turn: turn.clone(), })) .await; + let runtime = Arc::clone(self); let turn_for_task = turn.clone(); let turn_config_for_task = turn_config.clone(); let cancel_token = CancellationToken::new(); - let parent_session_id = { - let session = session_arc.lock().await; - session.summary.parent_session_id - }; - if let Some(parent_session_id) = parent_session_id { + if let Some(parent_session_id) = reservation.parent_session_id { let mut connections = self.active_turn_connections.lock().await; if let Some(connection_id) = connections.get(&parent_session_id).copied() { connections.insert(session_id, connection_id); @@ -578,6 +555,10 @@ impl ServerRuntime { .lock() .await .insert(session_id, cancel_token); + self.active_turn_ids + .lock() + .await + .insert(session_id, turn.turn_id); let task = tokio::spawn(async move { runtime .execute_turn(ExecuteTurnRequest { @@ -604,24 +585,22 @@ impl ServerRuntime { session_id: SessionId, turn: &TurnMetadata, ) -> Result<(), ToolCallError> { - let session_arc = self - .sessions - .lock() + let session_handle = self.session(session_id).await.ok_or_else(|| { + ToolCallError::InvalidInput(format!("session not found: {session_id}")) + })?; + let persistence_snapshot = session_handle + .turn_persistence_snapshot() .await - .get(&session_id) - .cloned() .ok_or_else(|| { - ToolCallError::InvalidInput(format!("session not found: {session_id}")) + ToolCallError::InvalidInput(format!( + "failed to snapshot turn persistence: {session_id}" + )) })?; - let (record, session_context, turn_context) = { - let session = session_arc.lock().await; - let core_session = session.core_session.lock().await; - ( - session.record.clone(), - core_session.session_context.clone(), - core_session.latest_turn_context.clone(), - ) - }; + let (record, session_context, turn_context) = ( + persistence_snapshot.record.clone(), + persistence_snapshot.session_context.clone(), + persistence_snapshot.latest_turn_context.clone(), + ); if let Some(record) = record { self.rollout_store .append_turn( @@ -770,12 +749,13 @@ impl ServerRuntime { } async fn session_agent_path(&self, session_id: SessionId) -> String { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { + let Some(session_handle) = self.session(session_id).await else { return "root".to_string(); }; - let session = session_arc.lock().await; - session - .summary + let Some(summary) = session_handle.summary().await else { + return "root".to_string(); + }; + summary .agent_path .clone() .unwrap_or_else(|| "root".to_string()) @@ -786,11 +766,73 @@ impl ServerRuntime { child_session_id: SessionId, turn: &TurnMetadata, ) { - let Some((parent_session_id, _agent_path)) = - self.child_parent_and_path(child_session_id).await + let Some((parent_session_id, status)) = self + .resolve_terminal_subagent_status(child_session_id, turn) + .await + else { + return; + }; + let detail = self + .subagent_terminal_status_detail(child_session_id, turn.turn_id, status) + .await; + self.finish_subagent_turn_completion( + parent_session_id, + child_session_id, + turn, + status, + detail, + ) + .await; + if subagent_stop_hook_applies(status) { + self.run_subagent_stop_hook(child_session_id).await; + } + } + + /// Same as `handle_subagent_turn_completed`, but reads any data owned by + /// the currently-executing session actor directly from `state` instead of + /// round-tripping through the session actor mailbox. + /// + /// Must be used when `child_session_id` is the session actor currently + /// executing this code: that actor's mailbox is not being polled until + /// the in-flight turn finishes, so a mailbox round-trip here would + /// deadlock forever waiting on itself. + pub(super) async fn handle_subagent_turn_completed_for_actor_state( + &self, + state: &SessionActorState, + child_session_id: SessionId, + turn: &TurnMetadata, + ) { + let Some((parent_session_id, status)) = self + .resolve_terminal_subagent_status(child_session_id, turn) + .await else { return; }; + let detail = subagent_terminal_status_detail_from_stable_items( + &state.persisted_turn_items, + turn.turn_id, + status, + ); + self.finish_subagent_turn_completion( + parent_session_id, + child_session_id, + turn, + status, + detail, + ) + .await; + if subagent_stop_hook_applies(status) { + self.run_subagent_stop_hook_for_actor_state(state, child_session_id) + .await; + } + } + + async fn resolve_terminal_subagent_status( + &self, + child_session_id: SessionId, + turn: &TurnMetadata, + ) -> Option<(SessionId, SubagentStatus)> { + let (parent_session_id, _agent_path) = self.child_parent_and_path(child_session_id).await?; let status = match turn.status { TurnStatus::Completed => SubagentStatus::Completed, TurnStatus::Interrupted => SubagentStatus::Interrupted, @@ -807,11 +849,19 @@ impl ServerRuntime { } else { status }; + Some((parent_session_id, status)) + } + + async fn finish_subagent_turn_completion( + &self, + parent_session_id: SessionId, + child_session_id: SessionId, + turn: &TurnMetadata, + status: SubagentStatus, + detail: Option, + ) { self.set_agent_status(parent_session_id, child_session_id, status) .await; - let detail = self - .subagent_terminal_status_detail(child_session_id, turn.turn_id, status) - .await; self.record_subagent_status_event_with_text( parent_session_id, child_session_id, @@ -820,16 +870,6 @@ impl ServerRuntime { detail, ) .await; - if matches!( - status, - SubagentStatus::Completed - | SubagentStatus::Failed - | SubagentStatus::Interrupted - | SubagentStatus::Canceled - | SubagentStatus::Closed - ) { - self.run_subagent_stop_hook(child_session_id).await; - } } async fn subagent_terminal_status_detail( @@ -841,35 +881,9 @@ impl ServerRuntime { if status != SubagentStatus::Failed { return None; } - let session_arc = self.sessions.lock().await.get(&child_session_id).cloned()?; - let session = session_arc.lock().await; - session.persisted_turn_items.iter().rev().find_map(|item| { - if item.turn_id != turn_id { - return None; - } - match &item.turn_item { - TurnItem::AgentMessage(TextItem { text }) if !text.trim().is_empty() => { - Some(text.trim().to_string()) - } - TurnItem::UserMessage(_) - | TurnItem::SteerInput(_) - | TurnItem::HookPrompt(_) - | TurnItem::AgentMessage(_) - | TurnItem::Plan(_) - | TurnItem::Reasoning(_) - | TurnItem::ToolCall(_) - | TurnItem::ToolProgress(_) - | TurnItem::ToolResult(_) - | TurnItem::CommandExecution(_) - | TurnItem::ApprovalRequest(_) - | TurnItem::ApprovalDecision(_) - | TurnItem::WebSearch(_) - | TurnItem::ImageGeneration(_) - | TurnItem::ContextCompaction(_) - | TurnItem::ResearchArtifact(_) - | TurnItem::TurnSummary(_) => None, - } - }) + let session_handle = self.sessions.lock().await.get(&child_session_id).cloned()?; + let snapshot = session_handle.spawn_snapshot().await?; + subagent_terminal_status_detail_from_stable_items(&snapshot.stable_items, turn_id, status) } pub(super) async fn child_parent_and_path( @@ -977,6 +991,38 @@ impl ServerRuntime { .is_some_and(|metadata| metadata.close_requested) } + pub(super) async fn interrupt_all_child_agents(self: Arc, parent_session_id: SessionId) { + let child_session_ids = { + let registries = self.agent_registries.lock().await; + registries + .get(&parent_session_id) + .map(|registry| registry.children_of(parent_session_id)) + .unwrap_or_default() + }; + let research_children = self + .research_child_agents + .lock() + .await + .get(&parent_session_id) + .cloned() + .unwrap_or_default(); + Arc::clone(&self) + .close_research_child_agents(parent_session_id) + .await; + for child_session_id in child_session_ids { + if research_children.contains(&child_session_id) { + continue; + } + let _ = self.interrupt_child_runtime_work(child_session_id).await; + self.set_agent_status( + parent_session_id, + child_session_id, + SubagentStatus::Interrupted, + ) + .await; + } + } + async fn close_child_agent( self: &Arc, parent_session_id: SessionId, @@ -1009,6 +1055,16 @@ impl ServerRuntime { } terminal }; + // A running turn's cancellation is now observed by the session actor's + // own turn-completion handling (`handle_subagent_turn_completed_for_actor_state`), + // which independently resolves and records the terminal "closed" status + // once it sees `close_requested`. Track whether a turn was actually in + // flight so we don't also send a duplicate closed notification below. + let had_active_turn = self + .active_turn_cancellations + .lock() + .await + .contains_key(&child_session_id); let interrupted_turn = self.interrupt_child_runtime_work(child_session_id).await; if already_terminal && interrupted_turn.is_none() { let status = self @@ -1033,7 +1089,7 @@ impl ServerRuntime { .await; self.handle_subagent_turn_completed(child_session_id, &turn) .await; - } else { + } else if !had_active_turn { self.send_closed_notification(parent_session_id, child_session_id) .await; } @@ -1109,6 +1165,54 @@ struct AgentRoute { to_agent_path: String, } +fn subagent_stop_hook_applies(status: SubagentStatus) -> bool { + matches!( + status, + SubagentStatus::Completed + | SubagentStatus::Failed + | SubagentStatus::Interrupted + | SubagentStatus::Canceled + | SubagentStatus::Closed + ) +} + +fn subagent_terminal_status_detail_from_stable_items( + stable_items: &[crate::execution::PersistedTurnItem], + turn_id: TurnId, + status: SubagentStatus, +) -> Option { + if status != SubagentStatus::Failed { + return None; + } + stable_items.iter().rev().find_map(|item| { + if item.turn_id != turn_id { + return None; + } + match &item.turn_item { + TurnItem::AgentMessage(TextItem { text }) if !text.trim().is_empty() => { + Some(text.trim().to_string()) + } + TurnItem::UserMessage(_) + | TurnItem::SteerInput(_) + | TurnItem::HookPrompt(_) + | TurnItem::AgentMessage(_) + | TurnItem::Plan(_) + | TurnItem::Reasoning(_) + | TurnItem::ToolCall(_) + | TurnItem::ToolProgress(_) + | TurnItem::ToolResult(_) + | TurnItem::CommandExecution(_) + | TurnItem::ApprovalRequest(_) + | TurnItem::ApprovalDecision(_) + | TurnItem::WebSearch(_) + | TurnItem::ImageGeneration(_) + | TurnItem::ContextCompaction(_) + | TurnItem::ResearchArtifact(_) + | TurnItem::TurnSummary(_) => None, + } + }) +} + fn generated_name_start_index(child_session_id: SessionId, max_count: usize) -> usize { child_session_id .to_string() diff --git a/crates/server/src/runtime/agents/lifecycle.rs b/crates/server/src/runtime/agents/lifecycle.rs index b24f0984..7cd626dc 100644 --- a/crates/server/src/runtime/agents/lifecycle.rs +++ b/crates/server/src/runtime/agents/lifecycle.rs @@ -19,6 +19,28 @@ impl ServerRuntime { let Some(context) = self.hook_context_for_session(child_session_id).await else { return; }; + Self::run_subagent_stop_hook_with_context(context).await; + } + + /// Same as `run_subagent_stop_hook`, but reads hook context directly from + /// `state` instead of round-tripping through the session actor mailbox. + /// + /// Must be used when `child_session_id` is the session actor currently + /// executing this code: that actor's mailbox is not being polled until + /// the in-flight turn finishes, so a mailbox round-trip here would + /// deadlock forever waiting on itself. + pub(super) async fn run_subagent_stop_hook_for_actor_state( + &self, + state: &SessionActorState, + child_session_id: SessionId, + ) { + let Some(context) = Self::hook_context_from_actor_state(state, child_session_id) else { + return; + }; + Self::run_subagent_stop_hook_with_context(context).await; + } + + async fn run_subagent_stop_hook_with_context(context: devo_core::HookRuntimeContext) { context .runner .run(devo_core::HookInput::new( @@ -37,12 +59,8 @@ impl ServerRuntime { ) { self.set_agent_status(parent_session_id, child_session_id, SubagentStatus::Failed) .await; - let session_arc = self.sessions.lock().await.get(&child_session_id).cloned(); - if let Some(session_arc) = session_arc { - let mut session = session_arc.lock().await; - session.summary.status = SessionRuntimeStatus::Idle; - session.summary.updated_at = Utc::now(); - session.summary.last_activity_at = session.summary.updated_at; + if let Some(session_handle) = self.sessions.lock().await.get(&child_session_id).cloned() { + session_handle.set_session_idle(None).await; } self.broadcast_event(ServerEvent::SessionStatusChanged( SessionStatusChangedPayload { @@ -66,28 +84,25 @@ impl ServerRuntime { self: &Arc, child_session_id: SessionId, ) -> Option { + // Cancel via a clone rather than `remove` so a concurrent read of the + // same token (e.g. `run_turn_model_query` fetching it to race against + // the in-flight query) cannot lose the signal by finding the map entry + // already gone and falling back to a fresh, disconnected token. The + // entry itself is cleaned up later by `finalize_executed_turn`. if let Some(cancel_token) = self .active_turn_cancellations .lock() .await - .remove(&child_session_id) + .get(&child_session_id) + .cloned() { cancel_token.cancel(); } if let Some(task) = self.active_tasks.lock().await.remove(&child_session_id) { task.abort(); } - let session_arc = self.sessions.lock().await.get(&child_session_id).cloned()?; - let mut session = session_arc.lock().await; - session.summary.status = SessionRuntimeStatus::Idle; - session.summary.updated_at = Utc::now(); - session.summary.last_activity_at = session.summary.updated_at; - session.active_turn.take().map(|mut turn| { - turn.status = TurnStatus::Interrupted; - turn.completed_at = Some(Utc::now()); - session.latest_turn = Some(turn.clone()); - turn - }) + let session_handle = self.sessions.lock().await.get(&child_session_id).cloned()?; + session_handle.interrupt_active_turn().await? } } @@ -306,21 +321,10 @@ mod tests { let session_id = response.result.session.session_id; let bad_rollout_path = data_root.path().join("rollout-dir"); std::fs::create_dir(&bad_rollout_path)?; - let session_arc = runtime - .sessions - .lock() - .await - .get(&session_id) - .cloned() - .expect("session"); - { - let mut session = session_arc.lock().await; - session - .record - .as_mut() - .expect("durable record") - .rollout_path = bad_rollout_path; - } + let session_handle = runtime.session(session_id).await.expect("session"); + session_handle + .update_record_rollout_path(bad_rollout_path) + .await; let error = runtime .start_runtime_turn( @@ -331,12 +335,16 @@ mod tests { ) .await .expect_err("append failure"); - let session = session_arc.lock().await; + let reservation = session_handle + .turn_reservation_snapshot() + .await + .expect("turn reservation snapshot"); + let summary = session_handle.summary().await.expect("summary"); assert!(matches!(error, ToolCallError::InternalError(_))); - assert_eq!(session.active_turn, None); - assert_eq!(session.summary.status, SessionRuntimeStatus::Idle); - assert_eq!(session.latest_turn, None); + assert_eq!(reservation.active_turn, None); + assert_eq!(summary.status, SessionRuntimeStatus::Idle); + assert_eq!(reservation.latest_turn, None); Ok(()) } diff --git a/crates/server/src/runtime/approval.rs b/crates/server/src/runtime/approval.rs index 062c68b8..82d62344 100644 --- a/crates/server/src/runtime/approval.rs +++ b/crates/server/src/runtime/approval.rs @@ -1,5 +1,7 @@ use super::*; +use crate::runtime::session_interactive::complete_approval_wait; + use std::path::Component; use std::path::Path; @@ -59,7 +61,8 @@ impl ServerRuntime { if self.approval_cache_allows(session_id, &request).await { return Ok(()); } - match policy_decision(&permission_profile, &request) { + let policy = policy_decision(&permission_profile, &request); + match policy { PolicyAuthorization::Allow => Ok(()), PolicyAuthorization::Ask => { if let Some(reason) = self @@ -138,17 +141,22 @@ impl ServerRuntime { request: &ToolPermissionRequest, ) -> AutoReviewOutcome { let (model, runtime_context) = { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { + let Some(session_handle) = self.session(session_id).await else { + return AutoReviewOutcome::AskUser; + }; + let Some(summary) = session_handle.summary().await else { + return AutoReviewOutcome::AskUser; + }; + let Some(runtime_context) = session_handle.runtime_context().await else { return AutoReviewOutcome::AskUser; }; - let session = session_arc.lock().await; + ( - session - .summary + summary .model .clone() - .unwrap_or_else(|| session.runtime_context.default_model.clone()), - Arc::clone(&session.runtime_context), + .unwrap_or_else(|| runtime_context.default_model.clone()), + runtime_context, ) }; let response = match runtime_context @@ -258,12 +266,65 @@ impl ServerRuntime { session_id: SessionId, request: &ToolPermissionRequest, ) -> bool { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { + if self + .session_approval_cache_allows(session_id, request) + .await + { + return true; + } + if let Some(parent_session_id) = self.parent_session_id(session_id).await { + return self + .session_approval_cache_allows(parent_session_id, request) + .await; + } + false + } + + async fn session_approval_cache_allows( + &self, + session_id: SessionId, + request: &ToolPermissionRequest, + ) -> bool { + if let Some(stream) = self.active_stream_state(session_id).await { + let stream = stream.lock().await; + if let Some(inline) = stream.turn_inline.as_ref() { + return cache_allows(&inline.session_approval_cache, request) + || cache_allows(&inline.turn_approval_cache, request); + } + } + let Some(session_handle) = self.session(session_id).await else { + return false; + }; + let Some(cache) = session_handle.approval_cache_snapshot().await else { return false; }; - let session = session_arc.lock().await; - cache_allows(&session.session_approval_cache, request) - || cache_allows(&session.turn_approval_cache, request) + cache_allows(&cache.session_approval_cache, request) + || cache_allows(&cache.turn_approval_cache, request) + } + + async fn parent_session_id(&self, session_id: SessionId) -> Option { + if let Some(stream) = self.active_stream_state(session_id).await { + let stream = stream.lock().await; + if let Some(inline) = stream.turn_inline.as_ref() { + return inline.summary.parent_session_id; + } + } + let session_handle = self.sessions.lock().await.get(&session_id).cloned()?; + session_handle.parent_session_id().await.and_then(|p| p) + } + + /// Sub-agent turns route interactive approvals through the parent session so + /// the active ACP connection and approval cache stay aligned with the UI. + async fn permission_host_session_id(&self, session_id: SessionId) -> SessionId { + let Some(parent_session_id) = self.parent_session_id(session_id).await else { + return session_id; + }; + let connections = self.active_turn_connections.lock().await; + if connections.contains_key(&parent_session_id) { + parent_session_id + } else { + session_id + } } async fn request_tool_approval( @@ -271,53 +332,66 @@ impl ServerRuntime { session_id: SessionId, request: ToolPermissionRequest, ) -> Result<(), String> { - let approval_id = format!("approval-{}", request.tool_call_id); - let (tx, rx) = oneshot::channel(); + let host_session_id = self.permission_host_session_id(session_id).await; let available_scopes = approval_scopes_for_request(&request); - let Some(connection_id) = self - .active_turn_connections - .lock() - .await - .get(&session_id) - .copied() - else { + let connection_id = { + let connections = self.active_turn_connections.lock().await; + connections + .get(&host_session_id) + .or_else(|| connections.get(&session_id)) + .copied() + }; + let Some(connection_id) = connection_id else { return Err("no ACP client connection is available for permission request".to_string()); }; - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { - return Err("session does not exist".to_string()); - }; - { - let mut session = session_arc.lock().await; - session.pending_approvals.insert( - approval_id.clone(), - PendingApproval { - tool_name: request.tool_name.clone(), - path: request.path.clone(), - host: request.host.clone(), - command_prefix: request.command_prefix.clone(), - tx, - }, + if host_session_id != session_id { + tracing::debug!( + child_session_id = %session_id, + parent_session_id = %host_session_id, + tool = %request.tool_name, + "routing sub-agent permission request through parent session" ); } - let request_params = acp_request_permission_params(session_id, &request, &available_scopes); + let approval_id = request.tool_call_id.clone(); + let (tx, rx) = oneshot::channel(); + let pending = PendingApproval { + tool_name: request.tool_name.clone(), + path: request.path.clone(), + host: request.host.clone(), + command_prefix: request.command_prefix.clone(), + tx, + }; + self.session_interactive + .register_pending_approval(host_session_id, approval_id.clone(), pending) + .await; + + let request_params = + acp_request_permission_params(host_session_id, &request, &available_scopes); + let cancel_token = { + let cancellations = self.active_turn_cancellations.lock().await; + cancellations + .get(&host_session_id) + .or_else(|| cancellations.get(&session_id)) + .cloned() + .unwrap_or_else(CancellationToken::new) + }; let response = match self - .send_request_to_connection( + .send_request_to_connection_cancellable( connection_id, devo_protocol::ACP_SESSION_REQUEST_PERMISSION_METHOD, serde_json::to_value(request_params) .expect("serialize ACP permission request params"), + cancel_token, ) .await { Ok(response) => response, Err(error) => { - session_arc - .lock() - .await - .pending_approvals - .remove(&approval_id); + self.session_interactive + .remove_pending_approval(host_session_id, &approval_id) + .await; return Err(format!("permission request failed: {error}")); } }; @@ -325,11 +399,9 @@ impl ServerRuntime { match serde_json::from_value(response) { Ok(response) => response, Err(error) => { - session_arc - .lock() - .await - .pending_approvals - .remove(&approval_id); + self.session_interactive + .remove_pending_approval(host_session_id, &approval_id) + .await; return Err(format!( "invalid session/request_permission response: {error}" )); @@ -338,32 +410,45 @@ impl ServerRuntime { let (decision, scope) = match approval_decision_from_acp_outcome(response.outcome) { Ok(decision) => decision, Err(error) => { - session_arc - .lock() - .await - .pending_approvals - .remove(&approval_id); + self.session_interactive + .remove_pending_approval(host_session_id, &approval_id) + .await; return Err(error); } }; - let pending = { - let mut session = session_arc.lock().await; - let Some(pending) = session.pending_approvals.remove(&approval_id) else { - return Err("approval request was already resolved".to_string()); - }; - if matches!(decision, ApprovalDecisionValue::Approve) { - apply_approval_scope(&mut session, &scope, &pending); - } - pending - }; - let _ = pending.tx.send(decision); - match rx.await { - Ok(ApprovalDecisionValue::Approve) => Ok(()), - Ok(ApprovalDecisionValue::Deny) => Err("rejected by user".to_string()), - Ok(ApprovalDecisionValue::Cancel) => Err("cancelled by user".to_string()), - Err(_) => Err("approval channel closed".to_string()), + if let Some(pending) = self + .session_interactive + .remove_pending_approval(host_session_id, &approval_id) + .await + { + let _ = pending.tx.send(decision.clone()); + if matches!(decision, ApprovalDecisionValue::Approve) + && let Some(session_handle) = self.session(host_session_id).await + { + let (scope_tx, _) = oneshot::channel(); + session_handle + .apply_approval_scope( + scope, + PendingApproval { + tool_name: pending.tool_name, + path: pending.path, + host: pending.host, + command_prefix: pending.command_prefix, + tx: scope_tx, + }, + ) + .await; + } } + + complete_approval_wait(rx) + .await + .and_then(|decision| match decision { + ApprovalDecisionValue::Approve => Ok(()), + ApprovalDecisionValue::Deny => Err("rejected by user".to_string()), + ApprovalDecisionValue::Cancel => Err("cancelled by user".to_string()), + }) } } @@ -521,52 +606,6 @@ fn approval_decision_from_acp_outcome( } } -fn apply_approval_scope( - session: &mut RuntimeSession, - scope: &ApprovalScopeValue, - pending: &PendingApproval, -) { - match scope { - ApprovalScopeValue::Once => {} - ApprovalScopeValue::Turn => { - session - .turn_approval_cache - .tools - .insert(pending.tool_name.clone()); - } - ApprovalScopeValue::Session => { - session - .session_approval_cache - .tools - .insert(pending.tool_name.clone()); - } - ApprovalScopeValue::PathPrefix => { - if let Some(path) = pending.path.clone() { - session.turn_approval_cache.path_prefixes.insert(path); - } - } - ApprovalScopeValue::Host => { - if let Some(host) = pending.host.clone() { - session.turn_approval_cache.hosts.insert(host); - } - } - ApprovalScopeValue::Tool => { - session - .turn_approval_cache - .tools - .insert(pending.tool_name.clone()); - } - ApprovalScopeValue::CommandPrefix => { - if let Some(command_prefix) = pending.command_prefix.clone() { - session - .session_approval_cache - .command_prefixes - .insert(command_prefix); - } - } - } -} - fn cache_allows( cache: &crate::execution::ApprovalGrantCache, request: &ToolPermissionRequest, diff --git a/crates/server/src/runtime/command_exec.rs b/crates/server/src/runtime/command_exec.rs index 21ad5547..0fdf899b 100644 --- a/crates/server/src/runtime/command_exec.rs +++ b/crates/server/src/runtime/command_exec.rs @@ -470,7 +470,13 @@ impl ServerRuntime { format!("session not found: {session_id}"), ) })?; - Ok(session.lock().await.summary.cwd.clone()) + let Some(summary) = session.summary().await else { + return Err(( + ProtocolErrorCode::InternalError, + "failed to read session summary".to_string(), + )); + }; + Ok(summary.cwd) } } diff --git a/crates/server/src/runtime/connection.rs b/crates/server/src/runtime/connection.rs index 8eb514ec..9fe83066 100644 --- a/crates/server/src/runtime/connection.rs +++ b/crates/server/src/runtime/connection.rs @@ -38,6 +38,56 @@ struct PendingConnectionNotification { value: serde_json::Value, } +struct QueuedConnectionNotification { + notification: PendingConnectionNotification, + delivered: Option>, +} + +fn spawn_connection_notification_writer( + mut queue_rx: mpsc::UnboundedReceiver, +) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + while let Some(queued) = queue_rx.recv().await { + let sent = send_connection_notification(queued.notification).await; + if let Some(delivered) = queued.delivered { + let _ = delivered.send(sent); + } + } + }) +} + +fn enqueue_connection_notification( + queue_tx: &mpsc::UnboundedSender, + notification: PendingConnectionNotification, +) -> bool { + queue_tx + .send(QueuedConnectionNotification { + notification, + delivered: None, + }) + .is_ok() +} + +async fn enqueue_connection_notification_and_wait_for_delivery( + queue_tx: &mpsc::UnboundedSender, + notification: PendingConnectionNotification, +) -> bool { + let (delivered_tx, delivered_rx) = oneshot::channel(); + if queue_tx + .send(QueuedConnectionNotification { + notification, + delivered: Some(delivered_tx), + }) + .is_err() + { + return false; + } + match delivered_rx.await { + Ok(sent) => sent, + Err(_) => false, + } +} + #[derive(Clone, Copy)] enum PendingConnectionMessageKind { Notification, @@ -148,6 +198,9 @@ impl ServerRuntime { sender: mpsc::Sender, ) -> u64 { let connection_id = self.next_connection_id.fetch_add(1, Ordering::SeqCst); + let (notification_queue_tx, notification_queue_rx) = mpsc::unbounded_channel(); + let _notification_writer = + spawn_connection_notification_writer(notification_queue_rx); let mut connections = self.connections.lock().await; connections.insert( connection_id, @@ -157,6 +210,7 @@ impl ServerRuntime { acp_authenticated: false, acp_client_capabilities: crate::AcpClientCapabilities::default(), sender, + notification_queue_tx, opt_out_notification_methods: HashSet::new(), subscriptions: Vec::new(), next_event_seq: 1, @@ -625,17 +679,20 @@ impl ServerRuntime { let event_seq = connection.next_seq(); let event = event.with_seq(event_seq); let (method, value) = acp_notification_from_server_event(method, &event); - Some(PendingConnectionNotification { - connection_id, - kind: PendingConnectionMessageKind::Notification, - method, - event_seq, - sender: connection.sender.clone(), - value, - }) + Some(( + connection.notification_queue_tx.clone(), + PendingConnectionNotification { + connection_id, + kind: PendingConnectionMessageKind::Notification, + method, + event_seq, + sender: connection.sender.clone(), + value, + }, + )) }; - if let Some(notification) = notification { - send_connection_notification(notification).await; + if let Some((queue_tx, notification)) = notification { + enqueue_connection_notification(&queue_tx, notification); } } @@ -668,19 +725,22 @@ impl ServerRuntime { let event_seq = connection.next_seq(); let event = event.clone().with_seq(event_seq); let (method, value) = acp_notification_from_server_event(method, &event); - Some(PendingConnectionNotification { - connection_id: *connection_id, - kind: PendingConnectionMessageKind::Notification, - method, - event_seq, - sender: connection.sender.clone(), - value, - }) + Some(( + connection.notification_queue_tx.clone(), + PendingConnectionNotification { + connection_id: *connection_id, + kind: PendingConnectionMessageKind::Notification, + method, + event_seq, + sender: connection.sender.clone(), + value, + }, + )) }) .collect::>() }; - for notification in notifications { - send_connection_notification(notification).await; + for (queue_tx, notification) in notifications { + enqueue_connection_notification(&queue_tx, notification); } } @@ -688,11 +748,16 @@ impl ServerRuntime { let Some(session_id) = session_activity_event_id(event) else { return; }; - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { - return; - }; - let mut session = session_arc.lock().await; - session.summary.last_activity_at = session.summary.last_activity_at.max(chrono::Utc::now()); + if let Some(stream) = self.active_stream_state(session_id).await { + let mut stream = stream.lock().await; + if let Some(inline) = stream.turn_inline.as_mut() { + inline.summary.last_activity_at = inline.summary.last_activity_at.max(Utc::now()); + return; + } + } + if let Some(session_handle) = self.session(session_id).await { + session_handle.touch_last_activity().await; + } } pub(super) async fn send_raw_to_connection( @@ -700,35 +765,39 @@ impl ServerRuntime { connection_id: u64, value: serde_json::Value, ) { - let notification = { + let (queue_tx, notification) = { let connections = self.connections.lock().await; let Some(connection) = connections.get(&connection_id) else { return; }; - PendingConnectionNotification { - connection_id, - kind: PendingConnectionMessageKind::JsonRpcResponse, - method: "".to_string(), - event_seq: 0, - sender: connection.sender.clone(), - value, - } + ( + connection.notification_queue_tx.clone(), + PendingConnectionNotification { + connection_id, + kind: PendingConnectionMessageKind::JsonRpcResponse, + method: "".to_string(), + event_seq: 0, + sender: connection.sender.clone(), + value, + }, + ) }; - send_connection_notification(notification).await; + let _ = enqueue_connection_notification_and_wait_for_delivery(&queue_tx, notification).await; } - pub(super) async fn send_request_to_connection( + pub(super) async fn send_request_to_connection_cancellable( &self, connection_id: u64, method: &str, params: serde_json::Value, + cancel_token: CancellationToken, ) -> Result { self.send_request_to_connection_inner( connection_id, method, params, /*timeout_duration*/ None, - CancellationToken::new(), + cancel_token, ) .await } @@ -759,7 +828,7 @@ impl ServerRuntime { timeout_duration: Option, cancel_token: CancellationToken, ) -> Result { - let (request_id, receiver, notification) = { + let (request_id, receiver, queue_tx, notification) = { let mut connections = self.connections.lock().await; let Some(connection) = connections.get_mut(&connection_id) else { return Err("client connection does not exist".to_string()); @@ -777,6 +846,7 @@ impl ServerRuntime { ( request_id, rx, + connection.notification_queue_tx.clone(), PendingConnectionNotification { connection_id, kind: PendingConnectionMessageKind::ClientRequest, @@ -792,7 +862,7 @@ impl ServerRuntime { connection_id, request_id, ); - if !send_connection_notification(notification).await { + if !enqueue_connection_notification(&queue_tx, notification) { pending_request.remove().await; return Err("client connection closed before request was sent".to_string()); } @@ -825,11 +895,16 @@ impl ServerRuntime { } } None => { - let message = receiver - .await - .map_err(|_| "client connection closed before responding".to_string())??; - pending_request.disarm(); - message + tokio::select! { + _ = cancel_token.cancelled() => { + pending_request.remove().await; + return Err("client request cancelled".to_string()); + } + result = receiver => { + pending_request.disarm(); + result.map_err(|_| "client connection closed before responding".to_string())?? + } + } } }; if let Some(error) = message.get("error") { @@ -1042,6 +1117,7 @@ pub(crate) struct ConnectionRuntime { pub(crate) acp_authenticated: bool, pub(crate) acp_client_capabilities: crate::AcpClientCapabilities, pub(crate) sender: mpsc::Sender, + notification_queue_tx: mpsc::UnboundedSender, pub(crate) opt_out_notification_methods: HashSet, pub(crate) subscriptions: Vec, next_event_seq: u64, diff --git a/crates/server/src/runtime/goal_accounting.rs b/crates/server/src/runtime/goal_accounting.rs index 382c0282..1b5d7553 100644 --- a/crates/server/src/runtime/goal_accounting.rs +++ b/crates/server/src/runtime/goal_accounting.rs @@ -73,14 +73,7 @@ impl ServerRuntime { } async fn session_is_plan_mode(&self, session_id: SessionId) -> bool { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { - return false; - }; - let core_session = { - let session = session_arc.lock().await; - Arc::clone(&session.core_session) - }; - core_session.lock().await.collaboration_mode == devo_protocol::CollaborationMode::Plan + self.session_collaboration_mode(session_id).await == Some(devo_protocol::CollaborationMode::Plan) } } diff --git a/crates/server/src/runtime/goal_continuation.rs b/crates/server/src/runtime/goal_continuation.rs index c014ae33..387e10b1 100644 --- a/crates/server/src/runtime/goal_continuation.rs +++ b/crates/server/src/runtime/goal_continuation.rs @@ -6,6 +6,7 @@ use super::*; use crate::goal::GoalStatus; +use crate::runtime::session_actor::SessionHandle; use futures::future::BoxFuture; struct GoalContinuationCandidate { @@ -19,74 +20,68 @@ impl ServerRuntime { session_id: SessionId, ) -> BoxFuture<'_, ()> { Box::pin(async move { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { + let Some(session_handle) = self.session(session_id).await else { return; }; if self - .pause_goal_continuation_after_failed_turn(session_id, &session_arc) + .pause_goal_continuation_after_failed_turn(session_id, &session_handle) .await { return; } - if !session_allows_goal_continuation(&session_arc).await { + if !self.session_allows_goal_continuation(session_id).await { return; } let Some(candidate) = self.goal_continuation_candidate(session_id).await else { return; }; - if !session_allows_goal_continuation(&session_arc).await { + if !self.session_allows_goal_continuation(session_id).await { return; } - let (turn_config, resolved_request) = { - let session = session_arc.lock().await; - let requested_model = session_model_selection(&session.summary); - let requested_reasoning_effort_selection = - session.summary.reasoning_effort_selection.clone(); - let turn_config = session - .runtime_context - .resolve_turn_config(requested_model, requested_reasoning_effort_selection); - let resolved_request = turn_config.model.resolve_reasoning_effort_selection( - turn_config.reasoning_effort_selection.as_deref(), - ); - (turn_config, resolved_request) + let Some(reservation) = session_handle.turn_reservation_snapshot().await else { + return; }; + let requested_model = session_model_selection(&reservation.summary); + let requested_reasoning_effort_selection = + reservation.summary.reasoning_effort_selection.clone(); + let turn_config = reservation + .runtime_context + .resolve_turn_config(requested_model, requested_reasoning_effort_selection); + let resolved_request = turn_config.model.resolve_reasoning_effort_selection( + turn_config.reasoning_effort_selection.as_deref(), + ); let request_model = turn_config.provider_request_model(&resolved_request.request_model); let now = Utc::now(); - let turn = { - let mut session = session_arc.lock().await; - if !session_has_goal_continuation_capacity_locked(&session) { - return; - } - let turn = TurnMetadata { - turn_id: TurnId::new(), - session_id, - sequence: session - .latest_turn - .as_ref() - .map_or(1, |turn| turn.sequence + 1), - status: TurnStatus::Running, - kind: devo_core::TurnKind::Regular, - model: turn_config.model.slug.clone(), - model_binding_id: turn_config.model_binding_id.clone(), - reasoning_effort_selection: turn_config.reasoning_effort_selection.clone(), - reasoning_effort: resolved_request.effective_reasoning_effort, - request_model, - request_thinking: resolved_request.request_thinking.clone(), - started_at: now, - completed_at: None, - usage: None, - stop_reason: None, - failure_reason: None, - }; - session.summary.status = SessionRuntimeStatus::ActiveTurn; - session.summary.updated_at = now; - session.summary.last_activity_at = now; - apply_turn_config_to_session_summary(&mut session.summary, &turn_config); - session.active_turn = Some(turn.clone()); - turn + let turn = TurnMetadata { + turn_id: TurnId::new(), + session_id, + sequence: reservation + .latest_turn + .as_ref() + .map_or(1, |turn| turn.sequence + 1), + status: TurnStatus::Running, + kind: devo_core::TurnKind::Regular, + model: turn_config.model.slug.clone(), + model_binding_id: turn_config.model_binding_id.clone(), + reasoning_effort_selection: turn_config.reasoning_effort_selection.clone(), + reasoning_effort: resolved_request.effective_reasoning_effort, + request_model, + request_thinking: resolved_request.request_thinking.clone(), + started_at: now, + completed_at: None, + usage: None, + stop_reason: None, + failure_reason: None, }; + if !session_handle + .try_begin_active_turn(turn.clone(), turn_config.clone()) + .await + .unwrap_or(false) + { + return; + } self.active_goal_continuation_turns .lock() .await @@ -99,7 +94,7 @@ impl ServerRuntime { .mark_goal_continuation_turn_started(session_id, &candidate.goal_id, turn.turn_id) .await { - self.clear_goal_continuation_turn_reservation(&session_arc, turn.turn_id) + self.clear_goal_continuation_turn_reservation(&session_handle, turn.turn_id) .await; return; } @@ -114,20 +109,20 @@ impl ServerRuntime { error = %error, "failed to persist goal continuation turn start" ); - self.clear_goal_continuation_turn_reservation(&session_arc, turn.turn_id) + self.clear_goal_continuation_turn_reservation(&session_handle, turn.turn_id) .await; return; } if !goal_continuation_turn_still_current( self, - &session_arc, + &session_handle, session_id, turn.turn_id, &candidate.goal_id, ) .await { - self.clear_goal_continuation_turn_reservation(&session_arc, turn.turn_id) + self.clear_goal_continuation_turn_reservation(&session_handle, turn.turn_id) .await; return; } @@ -141,14 +136,14 @@ impl ServerRuntime { .await; if !goal_continuation_turn_still_current( self, - &session_arc, + &session_handle, session_id, turn.turn_id, &candidate.goal_id, ) .await { - self.clear_goal_continuation_turn_reservation(&session_arc, turn.turn_id) + self.clear_goal_continuation_turn_reservation(&session_handle, turn.turn_id) .await; return; } @@ -159,14 +154,14 @@ impl ServerRuntime { .await; if !goal_continuation_turn_still_current( self, - &session_arc, + &session_handle, session_id, turn.turn_id, &candidate.goal_id, ) .await { - self.clear_goal_continuation_turn_reservation(&session_arc, turn.turn_id) + self.clear_goal_continuation_turn_reservation(&session_handle, turn.turn_id) .await; return; } @@ -186,17 +181,18 @@ impl ServerRuntime { } else { let mut cancellations = self.active_turn_cancellations.lock().await; let mut active_tasks = self.active_tasks.lock().await; - let still_active_turn = { - let session = session_arc.lock().await; - session - .active_turn - .as_ref() - .is_some_and(|active_turn| active_turn.turn_id == turn.turn_id) - }; + let still_active_turn = session_handle + .active_turn_id() + .await + .is_some_and(|active_turn_id| active_turn_id == Some(turn.turn_id)); if !still_active_turn { false } else { cancellations.insert(session_id, cancel_token); + self.active_turn_ids + .lock() + .await + .insert(session_id, turn.turn_id); let task = tokio::spawn(async move { runtime .execute_turn(ExecuteTurnRequest { @@ -219,12 +215,40 @@ impl ServerRuntime { } }; if !task_started { - self.clear_goal_continuation_turn_reservation(&session_arc, turn.turn_id) + self.clear_goal_continuation_turn_reservation(&session_handle, turn.turn_id) .await; } }) } + async fn session_allows_goal_continuation(&self, session_id: SessionId) -> bool { + if self + .session_interactive + .has_pending_interactive(session_id) + .await + { + return false; + } + let Some(session_handle) = self.session(session_id).await else { + return false; + }; + let Some(reservation) = session_handle.turn_reservation_snapshot().await else { + return false; + }; + if reservation.active_turn.is_some() { + return false; + } + if session_handle.collaboration_mode().await == Some(devo_protocol::CollaborationMode::Plan) + { + return false; + } + reservation + .pending_turn_queue + .lock() + .expect("pending turn queue mutex should not be poisoned") + .is_empty() + } + async fn goal_continuation_candidate( &self, session_id: SessionId, @@ -243,15 +267,21 @@ impl ServerRuntime { async fn pause_goal_continuation_after_failed_turn( &self, session_id: SessionId, - session_arc: &Arc>, + session_handle: &SessionHandle, ) -> bool { let (latest_turn, failure_message) = { - let session = session_arc.lock().await; - let latest_turn = session.latest_turn.clone(); - let failure_message = latest_turn.as_ref().and_then(|turn| { - latest_failed_turn_error_message(&session.persisted_turn_items, turn) - }); - (latest_turn, failure_message) + let reservation = session_handle.turn_reservation_snapshot().await; + let snapshot = session_handle.spawn_snapshot().await; + match (reservation, snapshot) { + (Some(reservation), Some(snapshot)) => { + let latest_turn = reservation.latest_turn; + let failure_message = latest_turn.as_ref().and_then(|turn| { + latest_failed_turn_error_message(&snapshot.stable_items, turn) + }); + (latest_turn, failure_message) + } + _ => (None, None), + } }; let mut stores = self.goal_stores.lock().await; @@ -295,18 +325,17 @@ impl ServerRuntime { session_id: SessionId, turn: &TurnMetadata, ) -> anyhow::Result<()> { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { + let Some(session_handle) = self.session(session_id).await else { return Ok(()); }; - let (record, session_context, turn_context) = { - let session = session_arc.lock().await; - let core_session = session.core_session.lock().await; - ( - session.record.clone(), - core_session.session_context.clone(), - core_session.latest_turn_context.clone(), - ) + let Some(persistence) = session_handle.turn_persistence_snapshot().await else { + return Ok(()); }; + let (record, session_context, turn_context) = ( + persistence.record, + persistence.session_context, + persistence.latest_turn_context, + ); if let Some(record) = record { self.rollout_store.append_turn( &record, @@ -384,7 +413,7 @@ impl ServerRuntime { async fn clear_goal_continuation_turn_reservation( &self, - session_arc: &Arc>, + session_handle: &SessionHandle, turn_id: TurnId, ) { let session_id = { @@ -410,31 +439,17 @@ impl ServerRuntime { .lock() .await .remove(&turn_id); - let mut session = session_arc.lock().await; - if session - .active_turn - .as_ref() - .is_some_and(|active| active.turn_id == turn_id) - { - session.active_turn = None; - session.summary.status = SessionRuntimeStatus::Idle; - } + session_handle.clear_active_turn_if_matches(turn_id).await; } async fn complete_deferred_items_for_goal_turn( &self, - session_arc: &Arc>, + session_handle: &SessionHandle, session_id: SessionId, turn_id: TurnId, ) { - let (deferred_assistant, deferred_reasoning) = { - let mut session = session_arc.lock().await; - ( - session.deferred_assistant.take(), - session.deferred_reasoning.take(), - ) - }; - if let Some((item_id, item_seq, text)) = deferred_assistant { + let deferred = session_handle.take_deferred_items().await; + if let Some((item_id, item_seq, text)) = deferred.assistant { self.complete_item( session_id, turn_id, @@ -446,7 +461,7 @@ impl ServerRuntime { ) .await; } - if let Some((item_id, item_seq, text)) = deferred_reasoning { + if let Some((item_id, item_seq, text)) = deferred.reasoning { self.complete_item( session_id, turn_id, @@ -473,20 +488,17 @@ impl ServerRuntime { else { return false; }; - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { + let Some(session_handle) = self.session(session_id).await else { self.goal_continuation_turn_goals .lock() .await .remove(&turn_id); return false; }; - let active_turn_matches = { - let session = session_arc.lock().await; - session - .active_turn - .as_ref() - .is_some_and(|active_turn| active_turn.turn_id == turn_id) - }; + let active_turn_matches = session_handle + .active_turn_id() + .await + .is_some_and(|active_turn_id| active_turn_id == Some(turn_id)); if !active_turn_matches { self.goal_continuation_turn_goals .lock() @@ -494,61 +506,25 @@ impl ServerRuntime { .remove(&turn_id); return false; } - self.complete_deferred_items_for_goal_turn(&session_arc, session_id, turn_id) + self.complete_deferred_items_for_goal_turn(&session_handle, session_id, turn_id) .await; - let interrupted_turn = { - let mut session = session_arc.lock().await; - let Some(active_turn) = session.active_turn.as_ref() else { - return false; - }; - if active_turn.turn_id != turn_id { - return false; - } - let mut turn = session - .active_turn - .take() - .expect("active turn checked above"); - turn.status = TurnStatus::Interrupted; - turn.completed_at = Some(Utc::now()); - session.latest_turn = Some(turn.clone()); - session.summary.status = SessionRuntimeStatus::Idle; - session.summary.updated_at = Utc::now(); - session.summary.last_activity_at = session.summary.updated_at; - let totals = session.core_session.try_lock().ok().map(|core_session| { - ( - core_session.total_input_tokens, - core_session.total_output_tokens, - core_session.total_tokens, - core_session.total_cache_creation_tokens, - core_session.total_cache_read_tokens, - core_session.prompt_token_estimate, - ) - }); - if let Some(( - total_input_tokens, - total_output_tokens, - total_tokens, - total_cache_creation_tokens, - total_cache_read_tokens, - prompt_token_estimate, - )) = totals - { - session.summary.total_input_tokens = total_input_tokens; - session.summary.total_output_tokens = total_output_tokens; - session.summary.total_tokens = total_tokens; - session.summary.total_cache_creation_tokens = total_cache_creation_tokens; - session.summary.total_cache_read_tokens = total_cache_read_tokens; - session.summary.prompt_token_estimate = prompt_token_estimate; - } - turn + let Some(interrupted_turn) = session_handle.interrupt_active_turn().await.flatten() else { + return false; }; + if interrupted_turn.turn_id != turn_id { + return false; + } + // Cancel via a clone rather than `remove`: see the comment in + // `interrupt_child_runtime_work` for why removing here races with + // `run_turn_model_query` fetching the same token. if let Some(cancel_token) = self .active_turn_cancellations .lock() .await - .remove(&session_id) + .get(&session_id) + .cloned() { cancel_token.cancel(); } @@ -557,16 +533,14 @@ impl ServerRuntime { } let (record, session_context, turn_context) = { - let session = session_arc.lock().await; - let core_session_lock = session.core_session.try_lock(); - if let Ok(core_session) = core_session_lock { - ( - session.record.clone(), - core_session.session_context.clone(), - core_session.latest_turn_context.clone(), - ) - } else { - (session.record.clone(), None, None) + let persistence = session_handle.turn_persistence_snapshot().await; + match persistence { + Some(persistence) => ( + persistence.record, + persistence.session_context, + persistence.latest_turn_context, + ), + None => (None, None, None), } }; if let Some(record) = record @@ -615,47 +589,9 @@ impl ServerRuntime { } } -async fn session_allows_goal_continuation(session_arc: &Arc>) -> bool { - let (core_session, pending_turn_queue) = { - let session = session_arc.lock().await; - if !session_has_goal_continuation_capacity_locked(&session) { - return false; - } - ( - Arc::clone(&session.core_session), - Arc::clone(&session.pending_turn_queue), - ) - }; - let plan_mode = { - let core_session = core_session.lock().await; - core_session.collaboration_mode == devo_protocol::CollaborationMode::Plan - }; - if plan_mode { - return false; - } - pending_turn_queue - .lock() - .expect("pending turn queue mutex should not be poisoned") - .is_empty() -} - -fn session_has_goal_continuation_capacity_locked(session: &RuntimeSession) -> bool { - if session.active_turn.is_some() - || !session.pending_approvals.is_empty() - || !session.pending_user_inputs.is_empty() - { - return false; - } - session - .pending_turn_queue - .lock() - .expect("pending turn queue mutex should not be poisoned") - .is_empty() -} - async fn goal_continuation_turn_still_current( runtime: &ServerRuntime, - session_arc: &Arc>, + session_handle: &SessionHandle, session_id: SessionId, turn_id: TurnId, goal_id: &GoalId, @@ -669,13 +605,10 @@ async fn goal_continuation_turn_still_current( if !still_reserved { return false; } - let still_active_turn = { - let session = session_arc.lock().await; - session - .active_turn - .as_ref() - .is_some_and(|active_turn| active_turn.turn_id == turn_id) - }; + let still_active_turn = session_handle + .active_turn_id() + .await + .is_some_and(|active_turn_id| active_turn_id == Some(turn_id)); if !still_active_turn { return false; } diff --git a/crates/server/src/runtime/goal_handlers.rs b/crates/server/src/runtime/goal_handlers.rs index 3fae4bb7..42b288d4 100644 --- a/crates/server/src/runtime/goal_handlers.rs +++ b/crates/server/src/runtime/goal_handlers.rs @@ -530,19 +530,9 @@ impl ServerRuntime { session_id: SessionId, goal: Option, ) { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { + let Some(session_handle) = self.session(session_id).await else { return; }; - let core_session = { - let session = session_arc.lock().await; - Arc::clone(&session.core_session) - }; - if let Ok(mut core_session) = core_session.try_lock() { - if let Some(goal) = goal { - core_session.set_active_goal(goal); - } else { - core_session.clear_active_goal(); - } - } + session_handle.set_active_goal(goal).await; } } diff --git a/crates/server/src/runtime/handlers/acp/prompt.rs b/crates/server/src/runtime/handlers/acp/prompt.rs index 893cc814..cdbb283e 100644 --- a/crates/server/src/runtime/handlers/acp/prompt.rs +++ b/crates/server/src/runtime/handlers/acp/prompt.rs @@ -113,35 +113,40 @@ impl ServerRuntime { return; } }; - let Some(turn_id) = self.active_turn_id(params.session_id).await else { + let Some(turn_id) = self.runtime_active_turn_id(params.session_id).await else { tracing::debug!(session_id = %params.session_id, "session/cancel had no active turn"); return; }; - let _ = self - .handle_turn_interrupt( - serde_json::Value::Null, - serde_json::to_value(TurnInterruptParams { - session_id: params.session_id, - turn_id, - reason: Some("cancelled by ACP client".to_string()), - }) - .expect("serialize turn interrupt params"), - ) - .await; + if let Some(cancel_token) = self + .active_turn_cancellations + .lock() + .await + .get(¶ms.session_id) + .cloned() + { + cancel_token.cancel(); + } + if let Some(task) = self.active_tasks.lock().await.remove(¶ms.session_id) { + task.abort(); + } + let runtime = Arc::clone(self); + tokio::spawn(async move { + let _ = runtime + .handle_turn_interrupt( + serde_json::Value::Null, + serde_json::to_value(TurnInterruptParams { + session_id: params.session_id, + turn_id, + reason: Some("cancelled by ACP client".to_string()), + }) + .expect("serialize turn interrupt params"), + ) + .await; + }); } async fn session_has_active_turn(&self, session_id: SessionId) -> bool { - self.active_turn_id(session_id).await.is_some() - } - - async fn active_turn_id(&self, session_id: SessionId) -> Option { - let session = self.sessions.lock().await.get(&session_id).cloned()?; - session - .lock() - .await - .active_turn - .as_ref() - .map(|turn| turn.turn_id) + self.runtime_active_turn_id(session_id).await.is_some() } pub(crate) async fn wait_for_acp_prompt_stop_reason( diff --git a/crates/server/src/runtime/handlers/acp/session.rs b/crates/server/src/runtime/handlers/acp/session.rs index 24194c2c..d5afc37e 100644 --- a/crates/server/src/runtime/handlers/acp/session.rs +++ b/crates/server/src/runtime/handlers/acp/session.rs @@ -410,16 +410,8 @@ impl ServerRuntime { ) -> Result, String> { let session_ids = self.collect_session_delete_tree(root_session_id).await; for session_id in &session_ids { - if self.sessions.lock().await.contains_key(session_id) { - self.handle_acp_session_cancel( - serde_json::to_value(AcpCancelParams { - session_id: *session_id, - meta: None, - }) - .expect("serialize ACP cancel params"), - ) + self.await_session_turn_interrupt_before_delete(*session_id) .await; - } } let mut deleted_session_ids = Vec::new(); @@ -457,19 +449,19 @@ impl ServerRuntime { } async fn collect_session_delete_tree(&self, root_session_id: SessionId) -> Vec { - let session_entries = { + let session_ids_in_runtime: Vec = { let sessions = self.sessions.lock().await; - sessions - .iter() - .map(|(session_id, session_arc)| (*session_id, Arc::clone(session_arc))) - .collect::>() + sessions.keys().copied().collect() }; let mut parent_by_session = Vec::new(); - for (session_id, session_arc) in session_entries { - let session = session_arc.lock().await; - parent_by_session.push((session_id, session.summary.parent_session_id)); + for session_id in session_ids_in_runtime { + let parent_id = self + .session_parent_id_snapshot(session_id) + .await + .flatten(); + parent_by_session.push((session_id, parent_id)); } - parent_by_session.sort_by_key(|(session_id, _)| session_id.to_string()); + parent_by_session.sort_by_key(|(session_id, _parent_id)| session_id.to_string()); let mut seen = std::collections::HashSet::new(); let mut session_ids = Vec::new(); @@ -488,18 +480,49 @@ impl ServerRuntime { session_ids } + async fn await_session_turn_interrupt_before_delete( + self: &Arc, + session_id: SessionId, + ) { + let Some(turn_id) = self.runtime_active_turn_id(session_id).await else { + return; + }; + let receiver = self.subscribe_terminal_turn_status(turn_id).await; + if self.recent_terminal_turn_status(turn_id).await.is_some() { + return; + } + if let Some(cancel_token) = self + .active_turn_cancellations + .lock() + .await + .get(&session_id) + .cloned() + { + cancel_token.cancel(); + } + if let Some(task) = self.active_tasks.lock().await.remove(&session_id) { + task.abort(); + } + let _ = tokio::time::timeout(std::time::Duration::from_secs(5), receiver).await; + } + async fn clear_deleted_session_runtime_state(&self, session_id: SessionId) { if let Some(task) = self.active_tasks.lock().await.remove(&session_id) { task.abort(); } + // Cancel via a clone rather than `remove`: see the comment in + // `interrupt_child_runtime_work` for why removing here races with + // `run_turn_model_query` fetching the same token. if let Some(cancellation) = self .active_turn_cancellations .lock() .await - .remove(&session_id) + .get(&session_id) + .cloned() { cancellation.cancel(); } + self.active_turn_ids.lock().await.remove(&session_id); self.active_turn_connections .lock() .await diff --git a/crates/server/src/runtime/handlers/acp/session_support.rs b/crates/server/src/runtime/handlers/acp/session_support.rs index 39c25bab..15eb7065 100644 --- a/crates/server/src/runtime/handlers/acp/session_support.rs +++ b/crates/server/src/runtime/handlers/acp/session_support.rs @@ -15,7 +15,16 @@ impl ServerRuntime { "session does not exist".to_string(), )); }; - let stored_cwd = session_arc.lock().await.summary.cwd.clone(); + let stored_cwd = session_arc + .summary() + .await + .ok_or_else(|| { + ( + AcpErrorCode::ServerError, + "session actor unavailable".to_string(), + ) + })? + .cwd; if stored_cwd != cwd { return Err(( AcpErrorCode::InvalidParams, @@ -67,40 +76,42 @@ impl ServerRuntime { let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { return Err("session does not exist".to_string()); }; - let (summary, core_session, profile) = { - let mut session = session_arc.lock().await; - if session.summary.additional_directories == additional_directories { - return Ok(session.summary.clone()); - } - let updated_at = Utc::now(); - session.summary.additional_directories = additional_directories.clone(); - session.summary.updated_at = updated_at; - if let Some(record) = session.record.as_mut() { - record.additional_directories = additional_directories.clone(); - record.updated_at = updated_at; - if let Err(error) = self.rollout_store.append_session_meta(record) { - return Err(format!( - "failed to persist ACP session additional directories: {error}" - )); - } - } + let snapshot = session_arc + .hook_context_snapshot() + .await + .ok_or_else(|| "session actor unavailable".to_string())?; + if snapshot.summary.additional_directories == additional_directories { + return Ok(snapshot.summary); + } - let profile = devo_safety::RuntimePermissionProfile::from_preset( - session.config.permission_profile.preset, - session.summary.cwd.clone(), - ) - .with_additional_roots(additional_directories.clone()); - session.config.permission_profile = profile.clone(); - ( - session.summary.clone(), - Arc::clone(&session.core_session), - profile, - ) - }; - core_session.lock().await.config.permission_profile = profile; + let updated_at = Utc::now(); + let mut updated_summary = snapshot.summary.clone(); + updated_summary.additional_directories = additional_directories.clone(); + updated_summary.updated_at = updated_at; + + // Apply new permission roots so tool authorization uses the updated workspace. + let profile = devo_safety::RuntimePermissionProfile::from_preset( + snapshot.config.permission_profile.preset, + updated_summary.cwd.clone(), + ) + .with_additional_roots(additional_directories.clone()); + if !session_arc.apply_permission_profile(profile).await { + return Err("failed to apply updated permission profile".to_string()); + } + session_arc.update_summary(updated_summary.clone()).await; + + if let Some(mut record) = snapshot.record { + record.additional_directories = additional_directories; + record.updated_at = updated_at; + if let Err(error) = self.rollout_store.append_session_meta(&record) { + return Err(format!( + "failed to persist ACP session additional directories: {error}" + )); + } + } - if !summary.ephemeral - && let Err(error) = self.deps.db.upsert_session(&summary) + if !updated_summary.ephemeral + && let Err(error) = self.deps.db.upsert_session(&updated_summary) { tracing::warn!( session_id = %session_id, @@ -108,7 +119,7 @@ impl ServerRuntime { "failed to persist ACP session additional directories" ); } - Ok(summary) + Ok(updated_summary) } pub(super) async fn send_acp_history_updates( diff --git a/crates/server/src/runtime/handlers/acp_config_options.rs b/crates/server/src/runtime/handlers/acp_config_options.rs index 8f390244..ea7eb3ce 100644 --- a/crates/server/src/runtime/handlers/acp_config_options.rs +++ b/crates/server/src/runtime/handlers/acp_config_options.rs @@ -2,12 +2,11 @@ use super::super::*; use std::collections::BTreeSet; -use chrono::DateTime; - use crate::AcpErrorCode; use crate::AcpSetConfigOptionParams; -use crate::execution::RuntimeSession; +use crate::runtime::session_actor::snapshots::HookContextSnapshot; use crate::session_context::SessionRuntimeContext; +use devo_core::SessionConfig; use devo_core::TurnConfig; use devo_protocol::AcpSessionConfigOption; use devo_protocol::AcpSessionConfigOptionCategory; @@ -30,8 +29,15 @@ impl ServerRuntime { let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { return Err("session does not exist".to_string()); }; - let session = session_arc.lock().await; - Ok(acp_config_options_for_session(&session)) + let snapshot: HookContextSnapshot = session_arc + .hook_context_snapshot() + .await + .ok_or_else(|| "session actor unavailable".to_string())?; + Ok(acp_config_options_for_session( + &snapshot.runtime_context, + &snapshot.summary, + &snapshot.config, + )) } pub(crate) fn acp_model_config_options_for_context( @@ -53,126 +59,192 @@ impl ServerRuntime { )); }; - let (updated_session, config_options) = { - let mut session = session_arc.lock().await; - match params.config_id.as_str() { - ACP_MODEL_CONFIG_ID => { - let model_option = acp_model_config_option_for_session(&session); - let value_is_allowed = match &model_option { - AcpSessionConfigOption::Select { options, .. } => { - select_options_contain_value(options, ¶ms.value) - } - }; - if !value_is_allowed { - return Err(( - AcpErrorCode::InvalidParams, - format!( - "invalid value '{}' for session config option '{}'", - params.value, params.config_id - ), - )); + let snapshot: HookContextSnapshot = + session_arc.hook_context_snapshot().await.ok_or_else(|| { + ( + AcpErrorCode::ServerError, + "session actor unavailable".to_string(), + ) + })?; + + let (updated_session, config_options) = match params.config_id.as_str() { + ACP_MODEL_CONFIG_ID => { + let model_option = acp_model_config_option_for_session( + &snapshot.runtime_context, + &snapshot.summary, + &snapshot.config, + ); + let value_is_allowed = match &model_option { + AcpSessionConfigOption::Select { options, .. } => { + select_options_contain_value(options, ¶ms.value) } + }; + if !value_is_allowed { + return Err(( + AcpErrorCode::InvalidParams, + format!( + "invalid value '{}' for session config option '{}'", + params.value, params.config_id + ), + )); + } - let mut turn_config = session.runtime_context.resolve_turn_config( - Some(params.value.as_str()), - session.summary.reasoning_effort_selection.clone(), - ); - turn_config.reasoning_effort_selection = current_reasoning_effort_value( - &turn_config.model, - session.summary.reasoning_effort_selection.as_deref(), - ); - apply_turn_config_to_session_summary(&mut session.summary, &turn_config); - let updated_at = Utc::now(); - session.summary.updated_at = updated_at; - if let Err(error) = persist_session_config_summary( - &self.rollout_store, - &mut session, - updated_at, - ) { - return Err((AcpErrorCode::ServerError, error)); - } - ( - Some(session.summary.clone()), - acp_config_options_for_session(&session), + let mut turn_config = snapshot.runtime_context.resolve_turn_config( + Some(params.value.as_str()), + snapshot.summary.reasoning_effort_selection.clone(), + ); + turn_config.reasoning_effort_selection = current_reasoning_effort_value( + &turn_config.model, + snapshot.summary.reasoning_effort_selection.as_deref(), + ); + + let updated = session_arc + .update_session_metadata( + Some(turn_config.model.slug.clone()), + turn_config.model_binding_id.clone(), + turn_config.reasoning_effort_selection.clone(), ) + .await + .ok_or_else(|| { + ( + AcpErrorCode::ServerError, + "failed to update session".to_string(), + ) + })?; + + let updated_snapshot = + session_arc.hook_context_snapshot().await.ok_or_else(|| { + ( + AcpErrorCode::ServerError, + "session actor unavailable".to_string(), + ) + })?; + if let Some(record) = updated_snapshot.record.as_ref() + && let Err(error) = self.rollout_store.append_session_meta(record) + { + return Err((AcpErrorCode::ServerError, error.to_string())); } - ACP_REASONING_EFFORT_CONFIG_ID => { - let Some(reasoning_effort_option) = - acp_reasoning_effort_config_option_for_session(&session) - else { - return Err(( - AcpErrorCode::InvalidParams, - format!("unknown session config option '{}'", params.config_id), - )); - }; - let value_is_allowed = match &reasoning_effort_option { - AcpSessionConfigOption::Select { options, .. } => { - select_options_contain_value(options, ¶ms.value) - } - }; - if !value_is_allowed { - return Err(( - AcpErrorCode::InvalidParams, - format!( - "invalid value '{}' for session config option '{}'", - params.value, params.config_id - ), - )); - } - let mut turn_config = session.runtime_context.resolve_turn_config( - session_model_selection(&session.summary), - Some(params.value.clone()), - ); - turn_config.reasoning_effort_selection = Some(params.value.clone()); - apply_turn_config_to_session_summary(&mut session.summary, &turn_config); - let updated_at = Utc::now(); - session.summary.updated_at = updated_at; - if let Err(error) = persist_session_config_summary( - &self.rollout_store, - &mut session, - updated_at, - ) { - return Err((AcpErrorCode::ServerError, error)); + ( + Some(updated), + acp_config_options_for_session( + &updated_snapshot.runtime_context, + &updated_snapshot.summary, + &updated_snapshot.config, + ), + ) + } + ACP_REASONING_EFFORT_CONFIG_ID => { + let Some(reasoning_effort_option) = acp_reasoning_effort_config_option_for_session( + &snapshot.runtime_context, + &snapshot.summary, + ) else { + return Err(( + AcpErrorCode::InvalidParams, + format!("unknown session config option '{}'", params.config_id), + )); + }; + let value_is_allowed = match &reasoning_effort_option { + AcpSessionConfigOption::Select { options, .. } => { + select_options_contain_value(options, ¶ms.value) } - ( - Some(session.summary.clone()), - acp_config_options_for_session(&session), - ) + }; + if !value_is_allowed { + return Err(( + AcpErrorCode::InvalidParams, + format!( + "invalid value '{}' for session config option '{}'", + params.value, params.config_id + ), + )); } - ACP_MODE_CONFIG_ID => { - let Some(preset) = permission_preset_from_value(¶ms.value) else { - return Err(( - AcpErrorCode::InvalidParams, - format!( - "invalid value '{}' for session config option '{}'", - params.value, params.config_id - ), - )); - }; - let profile = safety_profile_from_protocol( - preset, - session.summary.cwd.clone(), - session.summary.additional_directories.clone(), - ); - { - let mut core_session = session.core_session.lock().await; - core_session.config.permission_mode = profile.permission_mode(); - core_session.config.permission_profile = profile.clone(); - } - session.config.permission_mode = profile.permission_mode(); - session.config.permission_profile = profile; - session.session_approval_cache = - crate::execution::ApprovalGrantCache::default(); - session.turn_approval_cache = crate::execution::ApprovalGrantCache::default(); - (None, acp_config_options_for_session(&session)) + + let mut turn_config = snapshot.runtime_context.resolve_turn_config( + session_model_selection(&snapshot.summary), + Some(params.value.clone()), + ); + turn_config.reasoning_effort_selection = Some(params.value.clone()); + + let updated = session_arc + .update_session_metadata( + Some(turn_config.model.slug.clone()), + turn_config.model_binding_id.clone(), + turn_config.reasoning_effort_selection.clone(), + ) + .await + .ok_or_else(|| { + ( + AcpErrorCode::ServerError, + "failed to update session".to_string(), + ) + })?; + + let updated_snapshot = + session_arc.hook_context_snapshot().await.ok_or_else(|| { + ( + AcpErrorCode::ServerError, + "session actor unavailable".to_string(), + ) + })?; + if let Some(record) = updated_snapshot.record.as_ref() + && let Err(error) = self.rollout_store.append_session_meta(record) + { + return Err((AcpErrorCode::ServerError, error.to_string())); } - _ => { + + ( + Some(updated), + acp_config_options_for_session( + &updated_snapshot.runtime_context, + &updated_snapshot.summary, + &updated_snapshot.config, + ), + ) + } + ACP_MODE_CONFIG_ID => { + let Some(preset) = permission_preset_from_value(¶ms.value) else { return Err(( AcpErrorCode::InvalidParams, - format!("unknown session config option '{}'", params.config_id), + format!( + "invalid value '{}' for session config option '{}'", + params.value, params.config_id + ), + )); + }; + let profile = safety_profile_from_protocol( + preset, + snapshot.summary.cwd.clone(), + snapshot.summary.additional_directories.clone(), + ); + + if !session_arc.apply_permission_profile(profile.clone()).await { + return Err(( + AcpErrorCode::ServerError, + "failed to apply permission profile".to_string(), )); } + + let updated_snapshot = + session_arc.hook_context_snapshot().await.ok_or_else(|| { + ( + AcpErrorCode::ServerError, + "session actor unavailable".to_string(), + ) + })?; + ( + None, + acp_config_options_for_session( + &updated_snapshot.runtime_context, + &updated_snapshot.summary, + &updated_snapshot.config, + ), + ) + } + _ => { + return Err(( + AcpErrorCode::InvalidParams, + format!("unknown session config option '{}'", params.config_id), + )); } }; @@ -191,20 +263,25 @@ impl ServerRuntime { } } -fn acp_config_options_for_session(session: &RuntimeSession) -> Vec { - let mut options = acp_model_and_reasoning_options_for_session(session); - options.push(acp_mode_config_option_for_session(session)); +fn acp_config_options_for_session( + runtime_context: &SessionRuntimeContext, + summary: &SessionMetadata, + config: &SessionConfig, +) -> Vec { + let mut options = acp_model_and_reasoning_options_for_session(runtime_context, summary); + options.push(acp_mode_config_option_for_session(config)); options } fn acp_model_and_reasoning_options_for_session( - session: &RuntimeSession, + runtime_context: &SessionRuntimeContext, + summary: &SessionMetadata, ) -> Vec { - let turn_config = session.runtime_context.resolve_turn_config( - session_model_selection(&session.summary), - session.summary.reasoning_effort_selection.clone(), + let turn_config = runtime_context.resolve_turn_config( + session_model_selection(summary), + summary.reasoning_effort_selection.clone(), ); - acp_model_and_reasoning_options_for_context(session.runtime_context.as_ref(), &turn_config) + acp_model_and_reasoning_options_for_context(runtime_context, &turn_config) } fn acp_model_and_reasoning_options_for_context( @@ -223,9 +300,9 @@ fn acp_model_and_reasoning_options_for_context( options } -fn acp_mode_config_option_for_session(session: &RuntimeSession) -> AcpSessionConfigOption { +fn acp_mode_config_option_for_session(config: &SessionConfig) -> AcpSessionConfigOption { let current_value = permission_preset_value(permission_preset_from_safety( - session.config.permission_profile.preset, + config.permission_profile.preset, )) .to_string(); AcpSessionConfigOption::Select { @@ -271,12 +348,16 @@ fn acp_mode_config_option_for_session(session: &RuntimeSession) -> AcpSessionCon } } -fn acp_model_config_option_for_session(session: &RuntimeSession) -> AcpSessionConfigOption { - let turn_config = session.runtime_context.resolve_turn_config( - session_model_selection(&session.summary), - session.summary.reasoning_effort_selection.clone(), +fn acp_model_config_option_for_session( + runtime_context: &SessionRuntimeContext, + summary: &SessionMetadata, + _config: &SessionConfig, +) -> AcpSessionConfigOption { + let turn_config = runtime_context.resolve_turn_config( + session_model_selection(summary), + summary.reasoning_effort_selection.clone(), ); - acp_model_config_option_for_turn_config(session.runtime_context.as_ref(), &turn_config) + acp_model_config_option_for_turn_config(runtime_context, &turn_config) } fn acp_model_config_option_for_turn_config( @@ -355,11 +436,12 @@ fn acp_model_config_option_for_turn_config( } fn acp_reasoning_effort_config_option_for_session( - session: &RuntimeSession, + runtime_context: &SessionRuntimeContext, + summary: &SessionMetadata, ) -> Option { - let turn_config = session.runtime_context.resolve_turn_config( - session_model_selection(&session.summary), - session.summary.reasoning_effort_selection.clone(), + let turn_config = runtime_context.resolve_turn_config( + session_model_selection(summary), + summary.reasoning_effort_selection.clone(), ); acp_reasoning_effort_config_option_for_turn_config(&turn_config) } @@ -434,28 +516,6 @@ fn current_reasoning_effort_value(model: &Model, selection: Option<&str>) -> Opt .or_else(|| option_values.first().cloned()) } -fn persist_session_config_summary( - rollout_store: &crate::persistence::RolloutStore, - session: &mut RuntimeSession, - updated_at: DateTime, -) -> Result<(), String> { - let model = session.summary.model.clone(); - let model_binding_id = session.summary.model_binding_id.clone(); - let reasoning_effort_selection = session.summary.reasoning_effort_selection.clone(); - if let Some(record) = session.record.as_mut() { - record.model = model; - record.model_binding_id = model_binding_id; - record.reasoning_effort_selection = reasoning_effort_selection; - record.updated_at = updated_at; - if let Err(error) = rollout_store.append_session_meta(record) { - return Err(format!( - "failed to persist ACP session config option: {error}" - )); - } - } - Ok(()) -} - pub(crate) fn select_options_contain_value( options: &AcpSessionConfigSelectOptions, value: &str, diff --git a/crates/server/src/runtime/handlers/compaction.rs b/crates/server/src/runtime/handlers/compaction.rs index b23d5c07..8c2527b9 100644 --- a/crates/server/src/runtime/handlers/compaction.rs +++ b/crates/server/src/runtime/handlers/compaction.rs @@ -17,7 +17,7 @@ impl ServerRuntime { } }; - let session_arc = match self.sessions.lock().await.get(¶ms.session_id).cloned() { + let session_handle = match self.session(params.session_id).await { Some(session) => session, None => { return self.error_response( @@ -28,15 +28,18 @@ impl ServerRuntime { } }; - let summary = { - let runtime_session = session_arc.lock().await; - runtime_session.summary.clone() + let Some(summary) = session_handle.summary().await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); }; let runtime = Arc::clone(self); tokio::spawn(async move { if let Err(panic) = - AssertUnwindSafe(runtime.run_session_compaction(params.session_id, session_arc)) + AssertUnwindSafe(runtime.run_session_compaction(params.session_id, session_handle)) .catch_unwind() .await { @@ -59,12 +62,11 @@ impl ServerRuntime { pub(crate) async fn run_session_compaction( self: Arc, session_id: SessionId, - session_arc: Arc>, + session_handle: crate::runtime::session_actor::SessionHandle, ) { tracing::info!(session_id = %session_id, "session compaction task started"); - let started_summary = { - let runtime_session = session_arc.lock().await; - runtime_session.summary.clone() + let Some(started_summary) = session_handle.summary().await else { + return; }; self.broadcast_event(ServerEvent::SessionCompactionStarted(SessionEventPayload { session: started_summary, @@ -81,7 +83,17 @@ impl ServerRuntime { .await; let result = { - let runtime_session = session_arc.lock().await; + let Some(runtime_session) = session_handle.export_runtime_session().await else { + tracing::warn!(session_id = %session_id, "session compaction failed: session unavailable"); + self.broadcast_event(ServerEvent::SessionCompactionFailed( + SessionCompactionFailedPayload { + session_id, + message: "compaction failed: session unavailable".to_string(), + }, + )) + .await; + return; + }; let core_session = runtime_session.core_session.lock().await; let items: Vec = core_session @@ -143,7 +155,10 @@ impl ServerRuntime { match result { Ok(CompactAction::Replaced(compacted_items)) => { - let mut runtime_session = session_arc.lock().await; + let Some(mut runtime_session) = session_handle.export_runtime_session().await + else { + return; + }; let preserved_item_ids = Self::preserved_item_ids_from_compacted( &runtime_session.persisted_turn_items, &compacted_items, @@ -292,7 +307,13 @@ impl ServerRuntime { } } let summary = runtime_session.summary.clone(); - drop(runtime_session); + session_handle + .replace_state( + crate::runtime::session_actor::SessionActorState::from_runtime_session( + runtime_session, + ), + ) + .await; self.run_session_hook( session_id, devo_core::HookEvent::PostCompact, @@ -314,6 +335,13 @@ impl ServerRuntime { } let summary = runtime_session.summary.clone(); + session_handle + .replace_state( + crate::runtime::session_actor::SessionActorState::from_runtime_session( + runtime_session, + ), + ) + .await; tracing::info!(session_id = %session_id, "session compaction completed with replacement"); self.broadcast_event(ServerEvent::SessionCompactionCompleted( SessionEventPayload { session: summary }, @@ -321,8 +349,9 @@ impl ServerRuntime { .await; } Ok(CompactAction::Skipped) => { - let runtime_session = session_arc.lock().await; - let summary = runtime_session.summary.clone(); + let Some(summary) = session_handle.summary().await else { + return; + }; tracing::info!(session_id = %session_id, "session compaction completed without replacement"); self.broadcast_event(ServerEvent::SessionCompactionCompleted( SessionEventPayload { session: summary }, diff --git a/crates/server/src/runtime/handlers/message_edit.rs b/crates/server/src/runtime/handlers/message_edit.rs index 463d5652..352d90ea 100644 --- a/crates/server/src/runtime/handlers/message_edit.rs +++ b/crates/server/src/runtime/handlers/message_edit.rs @@ -44,20 +44,22 @@ impl ServerRuntime { ); }; - let Some(session_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + let Some(session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, "session does not exist", ); }; - let (workspace_root, runtime_context) = { - let session = session_arc.lock().await; - ( - session.summary.cwd.clone(), - Arc::clone(&session.runtime_context), - ) + let Some(hook_context) = session_handle.hook_context_snapshot().await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); }; + let workspace_root = hook_context.summary.cwd.clone(); + let runtime_context = hook_context.runtime_context; let Some(resolved_input) = (match runtime_context .resolve_input_items(&edited_input, Some(workspace_root.as_path())) { @@ -128,7 +130,13 @@ impl ServerRuntime { "message/editPrevious queued-only edits are not implemented", ); } - let session = session_arc.lock().await; + let Some(session) = session_handle.export_runtime_session().await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); + }; if session.active_turn.is_some() { return self.error_response( request_id, @@ -189,6 +197,14 @@ impl ServerRuntime { let requested_reasoning_effort_selection = session.summary.reasoning_effort_selection.clone(); let runtime_context = Arc::clone(&session.runtime_context); + let (session_context, turn_context, collaboration_mode) = { + let core_session = session.core_session.lock().await; + ( + core_session.session_context.clone(), + core_session.latest_turn_context.clone(), + core_session.collaboration_mode, + ) + }; drop(session); let turn_config = runtime_context.resolve_turn_config( @@ -276,15 +292,6 @@ impl ServerRuntime { stop_reason: None, failure_reason: None, }; - let (session_context, turn_context, collaboration_mode) = { - let session = session_arc.lock().await; - let core_session = session.core_session.lock().await; - ( - core_session.session_context.clone(), - core_session.latest_turn_context.clone(), - core_session.collaboration_mode, - ) - }; if let Err(error) = self .rollout_store .append_message_edit_recorded(&record, edit_record.clone()) @@ -391,46 +398,52 @@ impl ServerRuntime { ); } + let Some(mut session) = session_handle.export_runtime_session().await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); + }; + session + .persisted_turn_items + .retain(|item| item.turn_id != target_turn_id); + let mut rebuilt_messages = Vec::new(); + let mut rebuilt_history_items = Vec::new(); + let mut tool_names_by_id = HashMap::new(); + for item in &session.persisted_turn_items { + crate::persistence::apply_turn_item( + &mut rebuilt_messages, + &mut rebuilt_history_items, + &mut tool_names_by_id, + &item.turn_kind, + item.turn_item.clone(), + ); + } + if let Some(history_item) = history_item_from_turn_item(&replacement_item) { + rebuilt_history_items.push(history_item); + } + session.history_items = rebuilt_history_items; + session + .persisted_turn_items + .push(crate::execution::PersistedTurnItem { + turn_id: replacement_turn_id, + turn_kind: replacement_turn.kind.clone(), + item_id: replacement_message_id, + turn_item: replacement_item.clone(), + }); + let branch_turn_count = session + .persisted_turn_items + .iter() + .filter(|item| matches!(item.turn_item, TurnItem::UserMessage(_))) + .count() + .saturating_sub(1); + session.latest_compaction_snapshot = None; + session.summary.status = SessionRuntimeStatus::ActiveTurn; + session.summary.updated_at = now; + session.summary.last_activity_at = now; + session.active_turn = Some(replacement_turn.clone()); { - let mut session = session_arc.lock().await; - session - .persisted_turn_items - .retain(|item| item.turn_id != target_turn_id); - let mut rebuilt_messages = Vec::new(); - let mut rebuilt_history_items = Vec::new(); - let mut tool_names_by_id = HashMap::new(); - for item in &session.persisted_turn_items { - crate::persistence::apply_turn_item( - &mut rebuilt_messages, - &mut rebuilt_history_items, - &mut tool_names_by_id, - &item.turn_kind, - item.turn_item.clone(), - ); - } - if let Some(history_item) = history_item_from_turn_item(&replacement_item) { - rebuilt_history_items.push(history_item); - } - session.history_items = rebuilt_history_items; - session - .persisted_turn_items - .push(crate::execution::PersistedTurnItem { - turn_id: replacement_turn_id, - turn_kind: replacement_turn.kind.clone(), - item_id: replacement_message_id, - turn_item: replacement_item.clone(), - }); - let branch_turn_count = session - .persisted_turn_items - .iter() - .filter(|item| matches!(item.turn_item, TurnItem::UserMessage(_))) - .count() - .saturating_sub(1); - session.latest_compaction_snapshot = None; - session.summary.status = SessionRuntimeStatus::ActiveTurn; - session.summary.updated_at = now; - session.summary.last_activity_at = now; - session.active_turn = Some(replacement_turn.clone()); let mut core_session = session.core_session.lock().await; core_session.messages = rebuilt_messages; core_session.prompt_messages = None; @@ -443,11 +456,20 @@ impl ServerRuntime { } } } + session_handle + .replace_state( + crate::runtime::session_actor::SessionActorState::from_runtime_session(session), + ) + .await; self.active_turn_cancellations .lock() .await .insert(params.session_id, CancellationToken::new()); + self.active_turn_ids + .lock() + .await + .insert(params.session_id, replacement_turn.turn_id); let runtime = Arc::clone(self); let replacement_turn_for_task = replacement_turn.clone(); let turn_config_for_task = turn_config.clone(); diff --git a/crates/server/src/runtime/handlers/session.rs b/crates/server/src/runtime/handlers/session.rs index 079560e8..00dedd9f 100644 --- a/crates/server/src/runtime/handlers/session.rs +++ b/crates/server/src/runtime/handlers/session.rs @@ -106,36 +106,32 @@ impl ServerRuntime { let config = core_session.config.clone(); let pending_turn_queue = Arc::clone(&core_session.pending_turn_queue); let btw_input_queue = Arc::clone(&core_session.btw_input_queue); - self.sessions.lock().await.insert( - session_id, - RuntimeSession { - runtime_context, - record, - summary: summary.clone(), - config, - core_session: Arc::new(Mutex::new(core_session)), - active_turn: None, - latest_turn: None, - loaded_item_count: 0, - history_items: Vec::new(), - persisted_turn_items: Vec::new(), - latest_compaction_snapshot: None, - pending_turn_queue, - btw_input_queue, - agent_tool_policy: Default::default(), - max_turns: None, - deferred_assistant: None, - deferred_reasoning: None, - next_item_seq: 1, - first_user_input: None, - pending_approvals: std::collections::HashMap::new(), - pending_user_inputs: std::collections::HashMap::new(), - tool_registry, - session_approval_cache: crate::execution::ApprovalGrantCache::default(), - turn_approval_cache: crate::execution::ApprovalGrantCache::default(), - } - .shared(), - ); + let actor_state = SessionActorState { + runtime_context, + record, + summary: summary.clone(), + config, + core: core_session, + stream: Arc::new(tokio::sync::Mutex::new( + crate::runtime::session_actor::state::SessionStreamState::default(), + )), + active_turn: None, + latest_turn: None, + loaded_item_count: 0, + history_items: Vec::new(), + persisted_turn_items: Vec::new(), + latest_compaction_snapshot: None, + pending_turn_queue, + btw_input_queue, + agent_tool_policy: Default::default(), + max_turns: None, + next_item_seq: 1, + first_user_input: None, + tool_registry, + session_approval_cache: crate::execution::ApprovalGrantCache::default(), + turn_approval_cache: crate::execution::ApprovalGrantCache::default(), + }; + self.insert_session_actor(actor_state).await; self.subscribe_connection_to_session(connection_id, session_id, None) .await; @@ -181,24 +177,7 @@ impl ServerRuntime { } pub(crate) async fn list_session_summaries(&self) -> Vec { - let sessions = self - .sessions - .lock() - .await - .values() - .cloned() - .collect::>(); - let mut summaries = Vec::with_capacity(sessions.len()); - for session in sessions { - summaries.push(session.lock().await.summary.clone()); - } - summaries.sort_by(|left, right| { - right - .last_activity_at - .cmp(&left.last_activity_at) - .then_with(|| right.updated_at.cmp(&left.updated_at)) - }); - summaries + self.list_session_summaries_from_actors().await } pub(crate) async fn handle_session_metadata_update( @@ -216,35 +195,37 @@ impl ServerRuntime { ); } }; - let Some(session_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + let Some(session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, "session does not exist", ); }; - let updated_session = { - let mut session = session_arc.lock().await; - session.summary.model = params.model.clone(); - session.summary.model_binding_id = params.model_binding_id.clone(); - session.summary.reasoning_effort_selection = params.reasoning_effort_selection.clone(); - let updated_at = Utc::now(); - session.summary.updated_at = updated_at; - if let Some(record) = session.record.as_mut() { - record.model = params.model; - record.model_binding_id = params.model_binding_id; - record.reasoning_effort_selection = params.reasoning_effort_selection; - record.updated_at = updated_at; - if let Err(error) = self.rollout_store.append_session_meta(record) { - return self.error_response( - request_id, - ProtocolErrorCode::InternalError, - format!("failed to persist session metadata update: {error}"), - ); - } - } - session.summary.clone() + let Some(mut updated_session) = session_handle + .update_session_metadata( + params.model.clone(), + params.model_binding_id.clone(), + params.reasoning_effort_selection.clone(), + ) + .await + else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); }; + if let Some(record) = session_handle.record().await.flatten() { + if let Err(error) = self.rollout_store.append_session_meta(&record) { + return self.error_response( + request_id, + ProtocolErrorCode::InternalError, + format!("failed to persist session metadata update: {error}"), + ); + } + updated_session = session_handle.summary().await.unwrap_or(updated_session); + } // Persist updated session metadata to SQLite if !updated_session.ephemeral @@ -281,32 +262,35 @@ impl ServerRuntime { ); } }; - let Some(session_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + let Some(session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, "session does not exist", ); }; - - let profile = { - let mut session = session_arc.lock().await; - let profile = safety_profile_from_protocol( - params.preset, - session.summary.cwd.clone(), - session.summary.additional_directories.clone(), + let Some(summary) = session_handle.summary().await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", ); - { - let mut core_session = session.core_session.lock().await; - core_session.config.permission_mode = profile.permission_mode(); - core_session.config.permission_profile = profile.clone(); - } - session.config.permission_mode = profile.permission_mode(); - session.config.permission_profile = profile.clone(); - session.session_approval_cache = crate::execution::ApprovalGrantCache::default(); - session.turn_approval_cache = crate::execution::ApprovalGrantCache::default(); - profile }; + let profile = safety_profile_from_protocol( + params.preset, + summary.cwd.clone(), + summary.additional_directories.clone(), + ); + if !session_handle + .apply_permission_profile(profile.clone()) + .await + { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); + } serde_json::to_value(SuccessResponse { id: request_id, @@ -342,7 +326,7 @@ impl ServerRuntime { "session title cannot be empty", ); } - let Some(session_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + let Some(session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, @@ -350,33 +334,35 @@ impl ServerRuntime { ); }; - let summary = { - let mut session = session_arc.lock().await; - let previous_title = session.summary.title.clone(); - let updated_at = Utc::now(); - session.summary.title = Some(new_title.to_string()); - session.summary.title_state = - SessionTitleState::Final(SessionTitleFinalSource::UserRename); - session.summary.updated_at = updated_at; - if let Some(record) = session.record.as_mut() { - record.title = Some(new_title.to_string()); - record.title_state = SessionTitleState::Final(SessionTitleFinalSource::UserRename); - record.updated_at = updated_at; - if let Err(error) = self.rollout_store.append_title_update( - record, - new_title.to_string(), - record.title_state.clone(), - previous_title, - ) { - return self.error_response( - request_id, - ProtocolErrorCode::InternalError, - format!("failed to persist session title update: {error}"), - ); - } - } - session.summary.clone() + let previous_title = session_handle + .summary() + .await + .and_then(|summary| summary.title); + let Some(mut summary) = session_handle + .set_session_title_user_rename(new_title.to_string()) + .await + else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); }; + if let Some(record) = session_handle.record().await.flatten() { + if let Err(error) = self.rollout_store.append_title_update( + &record, + new_title.to_string(), + record.title_state.clone(), + previous_title, + ) { + return self.error_response( + request_id, + ProtocolErrorCode::InternalError, + format!("failed to persist session title update: {error}"), + ); + } + summary = session_handle.summary().await.unwrap_or(summary); + } // Persist updated session metadata to SQLite if !summary.ephemeral @@ -433,45 +419,45 @@ impl ServerRuntime { params: SessionResumeParams, tool_registry_update: RuntimeSessionToolRegistryUpdate, ) -> serde_json::Value { - let Some(session_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + let Some(session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, "session does not exist", ); }; - let mut session = session_arc.lock().await; match tool_registry_update { RuntimeSessionToolRegistryUpdate::KeepCurrent => {} RuntimeSessionToolRegistryUpdate::ReplaceIfCwdMatches { cwd, tool_registry } => { - if session.summary.cwd != cwd { + let summary = session_handle.summary().await; + if summary.as_ref().is_none_or(|summary| summary.cwd != cwd) { return self.error_response( request_id, ProtocolErrorCode::InvalidParams, "session cwd does not match the stored session cwd", ); } - session.tool_registry = tool_registry; + if !session_handle.set_tool_registry(tool_registry).await { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); + } } } - let session_summary = session.summary.clone(); - let latest_turn = session.latest_turn.clone(); - let loaded_item_count = session.loaded_item_count; - let history_items = session.history_items.clone(); - let pending_texts: Vec = session - .pending_turn_queue - .lock() - .expect("pending turn queue mutex poisoned") - .iter() - .filter_map(|item| match &item.kind { - devo_core::PendingInputKind::UserText { text } => Some(text.clone()), - devo_core::PendingInputKind::UserInput { display_text, .. } => { - Some(display_text.clone()) - } - _ => None, - }) - .collect(); - drop(session); + let Some(resume_snapshot) = session_handle.resume_snapshot().await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); + }; + let session_summary = resume_snapshot.summary; + let latest_turn = resume_snapshot.latest_turn; + let loaded_item_count = resume_snapshot.loaded_item_count; + let history_items = resume_snapshot.history_items; + let pending_texts = resume_snapshot.pending_texts; self.subscribe_connection_to_session(connection_id, params.session_id, None) .await; self.run_session_hook( @@ -517,19 +503,26 @@ impl ServerRuntime { ); } }; - let Some(source_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + let Some(source_handle) = self.session(params.session_id).await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); + }; + let Some(source) = source_handle.export_runtime_session().await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, "session does not exist", ); }; - let source = source_arc.lock().await; + let source = &source; let now = Utc::now(); let forked_id = SessionId::new(); let mut forked_runtime = match self .build_runtime_session_from_user_turn_cut( - &source, + source, RuntimeSessionTurnCutOptions { session_id: forked_id, user_turn_index: params.user_turn_index, @@ -547,7 +540,6 @@ impl ServerRuntime { } }; forked_runtime.summary.parent_session_id = Some(params.session_id); - drop(source); if !forked_runtime.summary.ephemeral { let record = self.rollout_store.create_session_record( forked_id, @@ -571,10 +563,8 @@ impl ServerRuntime { forked_runtime.record = Some(record); } let summary = forked_runtime.summary.clone(); - self.sessions - .lock() - .await - .insert(forked_id, forked_runtime.shared()); + self.insert_session_actor(SessionActorState::from_runtime_session(forked_runtime)) + .await; self.subscribe_connection_to_session(connection_id, forked_id, None) .await; tracing::info!( @@ -616,17 +606,24 @@ impl ServerRuntime { ); } }; - let Some(session_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + let Some(session_handle) = self.session(params.session_id).await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); + }; + let Some(source) = session_handle.export_runtime_session().await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, "session does not exist", ); }; - let source = session_arc.lock().await; + let source = &source; let mut rebuilt = match self .build_runtime_session_from_user_turn_cut( - &source, + source, RuntimeSessionTurnCutOptions { session_id: params.session_id, user_turn_index: Some(params.user_turn_index), @@ -647,7 +644,6 @@ impl ServerRuntime { let (retained_turn_ids, retained_item_ids) = retained_ids_for_persisted_items(&rebuilt.persisted_turn_items); let latest_turn_id = rebuilt.latest_turn.as_ref().map(|turn| turn.turn_id); - drop(source); if let Some(record) = record.clone() && let Err(error) = self.rollout_store.append_session_rollback( &record, @@ -667,7 +663,9 @@ impl ServerRuntime { let latest_turn = rebuilt.latest_turn.clone(); let loaded_item_count = rebuilt.loaded_item_count; let history_items = rebuilt.history_items.clone(); - *session_arc.lock().await = rebuilt; + session_handle + .replace_state(SessionActorState::from_runtime_session(rebuilt)) + .await; self.subscribe_connection_to_session(connection_id, params.session_id, None) .await; serde_json::to_value(SuccessResponse { @@ -858,8 +856,6 @@ impl ServerRuntime { next_item_seq: u64::try_from(source.persisted_turn_items.len().saturating_add(1)) .unwrap_or(u64::MAX), first_user_input: source.first_user_input.clone(), - pending_approvals: std::collections::HashMap::new(), - pending_user_inputs: std::collections::HashMap::new(), tool_registry: source.tool_registry.clone(), session_approval_cache: crate::execution::ApprovalGrantCache::default(), turn_approval_cache: crate::execution::ApprovalGrantCache::default(), diff --git a/crates/server/src/runtime/handlers/turn.rs b/crates/server/src/runtime/handlers/turn.rs index 7fd3f892..40ad777c 100644 --- a/crates/server/src/runtime/handlers/turn.rs +++ b/crates/server/src/runtime/handlers/turn.rs @@ -76,33 +76,30 @@ impl ServerRuntime { "turn input is empty", ); }; - let Some(session_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + let Some(session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, "session does not exist", ); }; - let (workspace_root, runtime_context) = { - let session = session_arc.lock().await; - let workspace_root = params - .cwd - .clone() - .unwrap_or_else(|| session.summary.cwd.clone()); - let runtime_context = if params - .cwd - .as_ref() - .is_some_and(|cwd| cwd != &session.summary.cwd) - { - None - } else { - Some(Arc::clone(&session.runtime_context)) - }; - (workspace_root, runtime_context) + let Some(reservation) = session_handle.turn_reservation_snapshot().await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); }; - let runtime_context = match runtime_context { - Some(runtime_context) => runtime_context, - None => match self.deps.context_for_workspace(&workspace_root).await { + let workspace_root = params + .cwd + .clone() + .unwrap_or_else(|| reservation.summary.cwd.clone()); + let runtime_context = if params + .cwd + .as_ref() + .is_some_and(|cwd| cwd != &reservation.summary.cwd) + { + match self.deps.context_for_workspace(&workspace_root).await { Ok(runtime_context) => runtime_context, Err(error) => { return self.error_response( @@ -111,7 +108,9 @@ impl ServerRuntime { format!("failed to initialize session workspace: {error}"), ); } - }, + } + } else { + reservation.runtime_context }; let Some(resolved_input) = (match runtime_context .resolve_input_items(¶ms.input, Some(workspace_root.as_path())) @@ -174,143 +173,126 @@ impl ServerRuntime { let now = Utc::now(); let mut cwd_change = None; - let (turn, turn_config) = { - let mut session = session_arc.lock().await; - if let Some(active_turn) = session.active_turn.as_ref() { - if queue_policy == TurnStartQueuePolicy::RejectActive { - return self.error_response( - request_id, - ProtocolErrorCode::TurnAlreadyRunning, - "session already has an active prompt turn", - ); - } - let pending_turn_queue = Arc::clone(&session.pending_turn_queue); - let active_turn_id = active_turn.turn_id; - let is_ephemeral = session.summary.ephemeral; - let queued_model = params - .model - .clone() - .or_else(|| session.summary.model.clone()); - let queued_model_binding_id = params - .model_binding_id - .clone() - .or_else(|| session.summary.model_binding_id.clone()); - drop(session); - - let queued_input_id = { - let collaboration_mode = params.collaboration_mode; - let mut guard = pending_turn_queue - .lock() - .expect("pending turn queue mutex should not be poisoned"); - let item = devo_core::PendingInputItem::new( - devo_core::PendingInputKind::UserInput { - input: params.input.clone(), - display_text: display_input.clone(), - prompt_text: resolved_input.prompt_text.clone(), - prompt_messages: resolved_input.prompt_messages.clone(), - }, - pending_turn_metadata( - collaboration_mode, - queued_model.clone(), - queued_model_binding_id.clone(), - ), - now, - ); - let queued_input_id = item.id; - guard.push_back(item.clone()); - - if !is_ephemeral - && let Err(err) = - self.deps - .db - .push_pending(¶ms.session_id, QueueType::Turn, &item) - { - tracing::warn!( - session_id = %params.session_id, - error = %err, - "failed to persist pending turn message to database" - ); - } - queued_input_id - }; - let sid = params.session_id; - let runtime = Arc::clone(self); - tokio::spawn(async move { - runtime.broadcast_updated_queue(sid).await; - }); - return serde_json::to_value(SuccessResponse { - id: request_id, - result: TurnStartResult::Queued { - active_turn_id, - queued_input_id, - status: TurnStatus::Pending, - accepted_at: now, - }, - }) - .expect("serialize queued turn/start response"); - } - if let Some(cwd) = params.cwd.clone() { - let old_cwd = session.summary.cwd.clone(); - if old_cwd != cwd { - cwd_change = Some((old_cwd, cwd.clone())); - session.runtime_context = Arc::clone(&runtime_context); - } - session.summary.cwd = cwd.clone(); - session.core_session.lock().await.cwd = cwd; - } - if let Some(permission_mode) = params - .approval_policy - .as_deref() - .and_then(permission_mode_from_approval_policy) - { - session.core_session.lock().await.config.permission_mode = permission_mode; - session.config.permission_mode = permission_mode; + if let Some(active_turn) = reservation.active_turn.as_ref() { + if queue_policy == TurnStartQueuePolicy::RejectActive { + return self.error_response( + request_id, + ProtocolErrorCode::TurnAlreadyRunning, + "session already has an active prompt turn", + ); } - let requested_model = requested_model_selection( - params.model_binding_id.as_deref(), - params.model.as_deref(), - &session.summary, - ); - let requested_reasoning_effort_selection = params - .reasoning_effort_selection + let active_turn_id = active_turn.turn_id; + let queued_model = params + .model .clone() - .or_else(|| session.summary.reasoning_effort_selection.clone()); - let turn_config = session.runtime_context.resolve_turn_config( - requested_model, - requested_reasoning_effort_selection.clone(), - ); - let resolved_request = turn_config.model.resolve_reasoning_effort_selection( - turn_config.reasoning_effort_selection.as_deref(), + .or_else(|| reservation.summary.model.clone()); + let queued_model_binding_id = params + .model_binding_id + .clone() + .or_else(|| reservation.summary.model_binding_id.clone()); + let item = devo_core::PendingInputItem::new( + devo_core::PendingInputKind::UserInput { + input: params.input.clone(), + display_text: display_input.clone(), + prompt_text: resolved_input.prompt_text.clone(), + prompt_messages: resolved_input.prompt_messages.clone(), + }, + pending_turn_metadata( + params.collaboration_mode, + queued_model, + queued_model_binding_id, + ), + now, ); - let request_model = turn_config.provider_request_model(&resolved_request.request_model); - apply_turn_config_to_session_summary(&mut session.summary, &turn_config); - let turn = TurnMetadata { - turn_id: TurnId::new(), - session_id: params.session_id, - sequence: session - .latest_turn - .as_ref() - .map_or(1, |turn| turn.sequence + 1), - status: TurnStatus::Running, - kind: devo_core::TurnKind::Regular, - model: turn_config.model.slug.clone(), - model_binding_id: turn_config.model_binding_id.clone(), - reasoning_effort_selection: turn_config.reasoning_effort_selection.clone(), - reasoning_effort: resolved_request.effective_reasoning_effort, - request_model, - request_thinking: resolved_request.request_thinking, - started_at: now, - completed_at: None, - usage: None, - stop_reason: None, - failure_reason: None, - }; - session.summary.status = SessionRuntimeStatus::ActiveTurn; - session.summary.updated_at = now; - session.summary.last_activity_at = now; - session.active_turn = Some(turn.clone()); - (turn, turn_config) + let queued_input_id = item.id; + session_handle + .enqueue_pending_turn_input(item.clone()) + .await; + if !reservation.ephemeral + && let Err(err) = + self.deps + .db + .push_pending(¶ms.session_id, QueueType::Turn, &item) + { + tracing::warn!( + session_id = %params.session_id, + error = %err, + "failed to persist pending turn message to database" + ); + } + let sid = params.session_id; + let runtime = Arc::clone(self); + tokio::spawn(async move { + runtime.broadcast_updated_queue(sid).await; + }); + return serde_json::to_value(SuccessResponse { + id: request_id, + result: TurnStartResult::Queued { + active_turn_id, + queued_input_id, + status: TurnStatus::Pending, + accepted_at: now, + }, + }) + .expect("serialize queued turn/start response"); + } + if let Some(cwd) = params.cwd.clone() { + let old_cwd = reservation.summary.cwd.clone(); + if old_cwd != cwd { + cwd_change = Some((old_cwd, cwd.clone())); + session_handle + .update_session_workspace(cwd.clone(), Arc::clone(&runtime_context)) + .await; + } + } + if let Some(permission_mode) = params + .approval_policy + .as_deref() + .and_then(permission_mode_from_approval_policy) + { + session_handle + .update_core_permission_mode(permission_mode) + .await; + } + let requested_model = requested_model_selection( + params.model_binding_id.as_deref(), + params.model.as_deref(), + &reservation.summary, + ); + let requested_reasoning_effort_selection = params + .reasoning_effort_selection + .clone() + .or_else(|| reservation.summary.reasoning_effort_selection.clone()); + let turn_config = runtime_context + .resolve_turn_config(requested_model, requested_reasoning_effort_selection); + let resolved_request = turn_config + .model + .resolve_reasoning_effort_selection(turn_config.reasoning_effort_selection.as_deref()); + let request_model = turn_config.provider_request_model(&resolved_request.request_model); + let turn = TurnMetadata { + turn_id: TurnId::new(), + session_id: params.session_id, + sequence: reservation + .latest_turn + .as_ref() + .map_or(1, |turn| turn.sequence + 1), + status: TurnStatus::Running, + kind: devo_core::TurnKind::Regular, + model: turn_config.model.slug.clone(), + model_binding_id: turn_config.model_binding_id.clone(), + reasoning_effort_selection: turn_config.reasoning_effort_selection.clone(), + reasoning_effort: resolved_request.effective_reasoning_effort, + request_model, + request_thinking: resolved_request.request_thinking, + started_at: now, + completed_at: None, + usage: None, + stop_reason: None, + failure_reason: None, }; + session_handle + .begin_active_turn(turn.clone(), turn_config.clone()) + .await; if let Some((old_cwd, new_cwd)) = cwd_change { self.run_session_hook( params.session_id, @@ -330,34 +312,20 @@ impl ServerRuntime { } self.maybe_start_title_generation_from_user_input(params.session_id, &display_input) .await; - let (record, session_context, turn_context) = { - let session = session_arc.lock().await; - let core_session = session.core_session.lock().await; - ( - session.record.clone(), - core_session.session_context.clone(), - core_session.latest_turn_context.clone(), - ) - }; - if let Some(record) = record + if let Some(persistence) = session_handle.turn_persistence_snapshot().await + && let Some(record) = persistence.record && let Err(error) = self.rollout_store.append_turn( &record, - build_turn_record(&turn, session_context, turn_context), + build_turn_record( + &turn, + persistence.session_context, + persistence.latest_turn_context, + ), ) { - { - let mut session = session_arc.lock().await; - if session - .active_turn - .as_ref() - .is_some_and(|active| active.turn_id == turn.turn_id) - { - session.active_turn = None; - session.summary.status = SessionRuntimeStatus::Idle; - session.summary.updated_at = Utc::now(); - session.summary.last_activity_at = session.summary.updated_at; - } - } + let _ = session_handle + .clear_active_turn_if_matches(turn.turn_id) + .await; return self.error_response( request_id, ProtocolErrorCode::InternalError, @@ -377,6 +345,10 @@ impl ServerRuntime { .lock() .await .insert(params.session_id, CancellationToken::new()); + self.active_turn_ids + .lock() + .await + .insert(params.session_id, turn.turn_id); if let Some(connection_id) = connection_id { self.active_turn_connections .lock() @@ -466,7 +438,7 @@ impl ServerRuntime { "shell command is empty", ); } - let Some(session_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + let Some(session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, @@ -487,60 +459,68 @@ impl ServerRuntime { }, None => None, }; + let Some(reservation) = session_handle.turn_reservation_snapshot().await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); + }; + if reservation.active_turn.is_some() { + return self.error_response( + request_id, + ProtocolErrorCode::TurnAlreadyRunning, + "cannot run shell command while a turn is active", + ); + } let now = Utc::now(); let mut cwd_change = None; - let (turn, cwd) = { - let mut session = session_arc.lock().await; - if session.active_turn.is_some() { - return self.error_response( - request_id, - ProtocolErrorCode::TurnAlreadyRunning, - "cannot run shell command while a turn is active", - ); - } - let cwd = params - .cwd - .clone() - .unwrap_or_else(|| session.summary.cwd.clone()); - if let Some(cwd) = params.cwd.clone() { - let old_cwd = session.summary.cwd.clone(); - if old_cwd != cwd { - cwd_change = Some((old_cwd, cwd.clone())); - if let Some(runtime_context) = requested_runtime_context.as_ref() { - session.runtime_context = Arc::clone(runtime_context); - } + let cwd = params + .cwd + .clone() + .unwrap_or_else(|| reservation.summary.cwd.clone()); + if let Some(cwd) = params.cwd.clone() { + let old_cwd = reservation.summary.cwd.clone(); + if old_cwd != cwd { + cwd_change = Some((old_cwd, cwd.clone())); + if let Some(runtime_context) = requested_runtime_context.as_ref() { + session_handle + .update_session_workspace(cwd.clone(), Arc::clone(runtime_context)) + .await; } - session.summary.cwd = cwd.clone(); - session.core_session.lock().await.cwd = cwd; } - let model = session.summary.model.clone().unwrap_or_default(); - let turn = TurnMetadata { - turn_id: TurnId::new(), - session_id: params.session_id, - sequence: session - .latest_turn - .as_ref() - .map_or(1, |turn| turn.sequence + 1), - status: TurnStatus::Running, - kind: devo_core::TurnKind::Other("shell_command".to_string()), - model: model.clone(), - model_binding_id: session.summary.model_binding_id.clone(), - reasoning_effort_selection: session.summary.reasoning_effort_selection.clone(), - reasoning_effort: session.summary.reasoning_effort, - request_model: model, - request_thinking: session.summary.reasoning_effort_selection.clone(), - started_at: now, - completed_at: None, - usage: None, - stop_reason: None, - failure_reason: None, - }; - session.summary.status = SessionRuntimeStatus::ActiveTurn; - session.summary.updated_at = now; - session.summary.last_activity_at = now; - session.active_turn = Some(turn.clone()); - (turn, cwd) + } + let model = reservation.summary.model.clone().unwrap_or_default(); + let turn = TurnMetadata { + turn_id: TurnId::new(), + session_id: params.session_id, + sequence: reservation + .latest_turn + .as_ref() + .map_or(1, |turn| turn.sequence + 1), + status: TurnStatus::Running, + kind: devo_core::TurnKind::Other("shell_command".to_string()), + model: model.clone(), + model_binding_id: reservation.summary.model_binding_id.clone(), + reasoning_effort_selection: reservation.summary.reasoning_effort_selection.clone(), + reasoning_effort: reservation.summary.reasoning_effort, + request_model: model, + request_thinking: reservation.summary.reasoning_effort_selection.clone(), + started_at: now, + completed_at: None, + usage: None, + stop_reason: None, + failure_reason: None, }; + session_handle + .begin_active_turn( + turn.clone(), + reservation.runtime_context.resolve_turn_config( + session_model_selection(&reservation.summary), + reservation.summary.reasoning_effort_selection.clone(), + ), + ) + .await; if let Some((old_cwd, new_cwd)) = cwd_change { self.run_session_hook( params.session_id, @@ -559,22 +539,19 @@ impl ServerRuntime { .await; } + if let Some(persistence) = session_handle.turn_persistence_snapshot().await + && let Some(record) = persistence.record + && let Err(error) = self + .rollout_store + .append_turn(&record, build_turn_record(&turn, None, None)) { - let session = session_arc.lock().await; - if let Some(record) = session.record.clone() - && let Err(error) = self - .rollout_store - .append_turn(&record, build_turn_record(&turn, None, None)) - { - drop(session); - self.clear_active_turn_reservation(&session_arc, turn.turn_id) - .await; - return self.error_response( - request_id, - ProtocolErrorCode::InternalError, - format!("failed to persist shell command turn start: {error}"), - ); - } + self.clear_active_turn_reservation(&session_handle, turn.turn_id) + .await; + return self.error_response( + request_id, + ProtocolErrorCode::InternalError, + format!("failed to persist shell command turn start: {error}"), + ); } self.broadcast_event(ServerEvent::SessionStatusChanged( @@ -597,6 +574,10 @@ impl ServerRuntime { .lock() .await .insert(params.session_id, CancellationToken::new()); + self.active_turn_ids + .lock() + .await + .insert(params.session_id, turn.turn_id); if let Some(connection_id) = connection_id { self.active_turn_connections .lock() @@ -639,7 +620,7 @@ impl ServerRuntime { ); } }; - let Some(session_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + let Some(session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, @@ -647,15 +628,65 @@ impl ServerRuntime { ); }; - let deferred_assistant = { - let mut session = session_arc.lock().await; - session.deferred_assistant.take() - }; - let deferred_reasoning = { - let mut session = session_arc.lock().await; - session.deferred_reasoning.take() + // Cancel before any session-actor mailbox round-trip: the actor may be blocked + // waiting for a permission response and cannot process commands until cancelled. + if let Some(cancel_token) = self + .active_turn_cancellations + .lock() + .await + .get(¶ms.session_id) + .cloned() + { + cancel_token.cancel(); + } + if let Some(task) = self.active_tasks.lock().await.remove(¶ms.session_id) { + task.abort(); + } + + let removed_len = self + .session_interactive + .clear_pending_user_inputs_for_turn(params.session_id, params.turn_id) + .await; + if removed_len > 0 { + tracing::info!( + session_id = %params.session_id, + turn_id = %params.turn_id, + removed_len, + "cleared pending request_user_input requests for interrupted turn" + ); + } + + // Cancel via a clone rather than `remove`: see the comment in + // `interrupt_child_runtime_work` for why removing here races with + // `run_turn_model_query` fetching the same token. + let close_children_parent = params.session_id; + Arc::clone(self) + .interrupt_all_child_agents(close_children_parent) + .await; + if self.runtime_active_turn_id(params.session_id).await != Some(params.turn_id) { + return self.error_response( + request_id, + ProtocolErrorCode::TurnNotFound, + "turn is not active", + ); + } + let Some(interrupted_turn) = session_handle.interrupt_active_turn().await.flatten() else { + return self.error_response( + request_id, + ProtocolErrorCode::TurnNotFound, + "turn is not active", + ); }; - if let Some((item_id, item_seq, text)) = deferred_assistant { + if interrupted_turn.turn_id != params.turn_id { + return self.error_response( + request_id, + ProtocolErrorCode::TurnNotFound, + "turn does not exist", + ); + } + + let deferred = session_handle.take_deferred_items().await; + if let Some((item_id, item_seq, text)) = deferred.assistant { self.complete_item( params.session_id, params.turn_id, @@ -667,7 +698,7 @@ impl ServerRuntime { ) .await; } - if let Some((item_id, item_seq, text)) = deferred_reasoning { + if let Some((item_id, item_seq, text)) = deferred.reasoning { self.complete_item( params.session_id, params.turn_id, @@ -679,110 +710,15 @@ impl ServerRuntime { ) .await; } - - { - let mut session = session_arc.lock().await; - let previous_len = session.pending_user_inputs.len(); - session - .pending_user_inputs - .retain(|_, pending| pending.turn_id != params.turn_id); - let removed_len = previous_len.saturating_sub(session.pending_user_inputs.len()); - if removed_len > 0 { - tracing::info!( - session_id = %params.session_id, - turn_id = %params.turn_id, - removed_len, - "cleared pending request_user_input requests for interrupted turn" - ); - } - } - - if let Some(cancel_token) = self - .active_turn_cancellations - .lock() - .await - .remove(¶ms.session_id) - { - cancel_token.cancel(); - } - if let Some(task) = self.active_tasks.lock().await.remove(¶ms.session_id) { - task.abort(); - } - let close_children_runtime = Arc::clone(self); - let close_children_parent = params.session_id; - tokio::spawn(async move { - close_children_runtime - .close_research_child_agents(close_children_parent) - .await; - }); - let interrupted_turn = { - let mut session = session_arc.lock().await; - let Some(mut turn) = session.active_turn.take() else { - return self.error_response( - request_id, - ProtocolErrorCode::TurnNotFound, - "turn is not active", - ); - }; - if turn.turn_id != params.turn_id { - session.active_turn = Some(turn); - return self.error_response( - request_id, - ProtocolErrorCode::TurnNotFound, - "turn does not exist", - ); - } - turn.status = TurnStatus::Interrupted; - turn.completed_at = Some(Utc::now()); - session.latest_turn = Some(turn.clone()); - session.summary.status = SessionRuntimeStatus::Idle; - session.summary.updated_at = Utc::now(); - session.summary.last_activity_at = session.summary.updated_at; - let totals = session.core_session.try_lock().ok().map(|core_session| { - ( - core_session.total_input_tokens, - core_session.total_output_tokens, - core_session.total_tokens, - core_session.total_cache_creation_tokens, - core_session.total_cache_read_tokens, - core_session.prompt_token_estimate, - ) - }); - if let Some(( - total_input_tokens, - total_output_tokens, - total_tokens, - total_cache_creation_tokens, - total_cache_read_tokens, - prompt_token_estimate, - )) = totals - { - session.summary.total_input_tokens = total_input_tokens; - session.summary.total_output_tokens = total_output_tokens; - session.summary.total_tokens = total_tokens; - session.summary.total_cache_creation_tokens = total_cache_creation_tokens; - session.summary.total_cache_read_tokens = total_cache_read_tokens; - session.summary.prompt_token_estimate = prompt_token_estimate; - } - turn - }; - let (record, session_context, turn_context) = { - let session = session_arc.lock().await; - let core_session_lock = session.core_session.try_lock(); - if let Ok(core_session) = core_session_lock { - ( - session.record.clone(), - core_session.session_context.clone(), - core_session.latest_turn_context.clone(), - ) - } else { - (session.record.clone(), None, None) - } - }; - if let Some(record) = record + if let Some(persistence) = session_handle.turn_persistence_snapshot().await + && let Some(record) = persistence.record && let Err(error) = self.rollout_store.append_turn( &record, - build_turn_record(&interrupted_turn, session_context, turn_context), + build_turn_record( + &interrupted_turn, + persistence.session_context, + persistence.latest_turn_context, + ), ) { return self.error_response( @@ -869,44 +805,44 @@ impl ServerRuntime { "turn steer input is empty", ); }; - let Some(session_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + let Some(session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, "session does not exist", ); }; - let (turn_id, workspace_root, btw_input_queue, runtime_context) = { - let session = session_arc.lock().await; - let Some(turn_id) = session.active_turn.as_ref().map(|turn| turn.turn_id) else { - return self.error_response( - request_id, - ProtocolErrorCode::NoActiveTurn, - "no active turn exists", - ); - }; - if turn_id != params.expected_turn_id { - return self.error_response( - request_id, - ProtocolErrorCode::ExpectedTurnMismatch, - "active turn did not match expectedTurnId", - ); - } - let active_turn = session.active_turn.as_ref().expect("active turn exists"); - if active_turn.kind != devo_core::TurnKind::Regular { - return self.error_response( - request_id, - ProtocolErrorCode::ActiveTurnNotSteerable, - "cannot steer a non-regular turn", - ); - } - ( - turn_id, - session.summary.cwd.clone(), - Arc::clone(&session.btw_input_queue), - Arc::clone(&session.runtime_context), - ) + let Some(reservation) = session_handle.turn_reservation_snapshot().await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); }; + let Some(active_turn) = reservation.active_turn.as_ref() else { + return self.error_response( + request_id, + ProtocolErrorCode::NoActiveTurn, + "no active turn exists", + ); + }; + let turn_id = active_turn.turn_id; + if turn_id != params.expected_turn_id { + return self.error_response( + request_id, + ProtocolErrorCode::ExpectedTurnMismatch, + "active turn did not match expectedTurnId", + ); + } + if active_turn.kind != devo_core::TurnKind::Regular { + return self.error_response( + request_id, + ProtocolErrorCode::ActiveTurnNotSteerable, + "cannot steer a non-regular turn", + ); + } + let workspace_root = reservation.summary.cwd.clone(); + let runtime_context = reservation.runtime_context; let resolved_input = match runtime_context .resolve_input_items(¶ms.input, Some(workspace_root.as_path())) { @@ -959,25 +895,19 @@ impl ServerRuntime { None, chrono::Utc::now(), ); - btw_input_queue - .lock() - .expect("btw input queue mutex should not be poisoned") - .push_back(item.clone()); + session_handle.enqueue_btw_input(item.clone()).await; + if !reservation.ephemeral + && let Err(err) = self + .deps + .db + .push_pending(¶ms.session_id, QueueType::Btw, &item) { - let session = session_arc.lock().await; - if !session.summary.ephemeral - && let Err(err) = - self.deps - .db - .push_pending(¶ms.session_id, QueueType::Btw, &item) - { - tracing::warn!( - session_id = %params.session_id, - error = %err, - "failed to persist btw input to database" - ); - } + tracing::warn!( + session_id = %params.session_id, + error = %err, + "failed to persist btw input to database" + ); } self.emit_to_connection( @@ -1022,20 +952,22 @@ impl ServerRuntime { ); } }; - let Some(session_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + let Some(session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, "session does not exist", ); }; - let (pending_turn_queue, is_ephemeral) = { - let session = session_arc.lock().await; - ( - Arc::clone(&session.pending_turn_queue), - session.summary.ephemeral, - ) + let Some(reservation) = session_handle.turn_reservation_snapshot().await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); }; + let pending_turn_queue = reservation.pending_turn_queue; + let is_ephemeral = reservation.ephemeral; let removed = { let mut queue = pending_turn_queue .lock() @@ -1085,43 +1017,44 @@ impl ServerRuntime { ); } }; - let Some(session_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + let Some(session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, "session does not exist", ); }; - let (turn_id, pending_turn_queue, btw_input_queue, is_ephemeral) = { - let session = session_arc.lock().await; - let Some(active_turn) = session.active_turn.as_ref() else { - return self.error_response( - request_id, - ProtocolErrorCode::NoActiveTurn, - "no active turn exists", - ); - }; - if active_turn.turn_id != params.expected_turn_id { - return self.error_response( - request_id, - ProtocolErrorCode::ExpectedTurnMismatch, - "active turn did not match expectedTurnId", - ); - } - if active_turn.kind != devo_core::TurnKind::Regular { - return self.error_response( - request_id, - ProtocolErrorCode::ActiveTurnNotSteerable, - "cannot steer a non-regular turn", - ); - } - ( - active_turn.turn_id, - Arc::clone(&session.pending_turn_queue), - Arc::clone(&session.btw_input_queue), - session.summary.ephemeral, - ) + let Some(reservation) = session_handle.turn_reservation_snapshot().await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); + }; + let Some(active_turn) = reservation.active_turn.as_ref() else { + return self.error_response( + request_id, + ProtocolErrorCode::NoActiveTurn, + "no active turn exists", + ); }; + if active_turn.turn_id != params.expected_turn_id { + return self.error_response( + request_id, + ProtocolErrorCode::ExpectedTurnMismatch, + "active turn did not match expectedTurnId", + ); + } + if active_turn.kind != devo_core::TurnKind::Regular { + return self.error_response( + request_id, + ProtocolErrorCode::ActiveTurnNotSteerable, + "cannot steer a non-regular turn", + ); + } + let turn_id = active_turn.turn_id; + let pending_turn_queue = reservation.pending_turn_queue; + let is_ephemeral = reservation.ephemeral; let (item, display_input) = { let mut queue = pending_turn_queue .lock() @@ -1153,10 +1086,7 @@ impl ServerRuntime { (item, display_input) }; - btw_input_queue - .lock() - .expect("btw input queue mutex should not be poisoned") - .push_back(item.clone()); + session_handle.enqueue_btw_input(item.clone()).await; if !is_ephemeral { if let Err(error) = self.deps.db.remove_pending_by_id( diff --git a/crates/server/src/runtime/handlers/workspace_changes.rs b/crates/server/src/runtime/handlers/workspace_changes.rs index 1b4fb199..4ab3d50a 100644 --- a/crates/server/src/runtime/handlers/workspace_changes.rs +++ b/crates/server/src/runtime/handlers/workspace_changes.rs @@ -24,24 +24,26 @@ impl ServerRuntime { ); } - let Some(session_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + let Some(session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, "session does not exist", ); }; - let (cwd, active_turn_id, latest_turn_id) = { - let session = session_arc.lock().await; - ( - params - .cwd - .clone() - .unwrap_or_else(|| session.summary.cwd.clone()), - session.active_turn.as_ref().map(|turn| turn.turn_id), - session.latest_turn.as_ref().map(|turn| turn.turn_id), - ) + let Some(reservation) = session_handle.turn_reservation_snapshot().await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); }; + let cwd = params + .cwd + .clone() + .unwrap_or_else(|| reservation.summary.cwd.clone()); + let active_turn_id = reservation.active_turn.as_ref().map(|turn| turn.turn_id); + let latest_turn_id = reservation.latest_turn.as_ref().map(|turn| turn.turn_id); let mut views = Vec::with_capacity(params.scopes.len()); for scope in params.scopes { diff --git a/crates/server/src/runtime/hooks.rs b/crates/server/src/runtime/hooks.rs index 0f79c191..2430c003 100644 --- a/crates/server/src/runtime/hooks.rs +++ b/crates/server/src/runtime/hooks.rs @@ -3,6 +3,9 @@ use serde_json::Value; use super::*; +use crate::runtime::session_actor::SessionActorState; +use crate::runtime::session_actor::snapshots::HookContextSnapshot; + impl ServerRuntime { pub(crate) fn hook_runner(&self) -> Option { let config = { @@ -16,35 +19,82 @@ impl ServerRuntime { (!config.is_empty()).then(|| devo_core::HookRunner::new(config)) } + pub(crate) fn hook_context_from_actor_state( + state: &SessionActorState, + session_id: SessionId, + ) -> Option { + let runner = state.runtime_context.hook_runner()?; + let transcript_path = state + .record + .as_ref() + .map(|record| record.rollout_path.display().to_string()) + .unwrap_or_default(); + let permission_mode = Some(permission_mode_label(state.config.permission_mode)); + let agent_id = state + .summary + .parent_session_id + .is_some() + .then(|| session_id.to_string()); + let agent_type = state + .summary + .agent_role + .clone() + .or_else(|| state.summary.agent_nickname.clone()); + Some(devo_core::HookRuntimeContext { + runner, + base: devo_core::HookBaseInput { + session_id: session_id.to_string(), + transcript_path, + cwd: state.summary.cwd.clone(), + permission_mode, + agent_id, + agent_type, + }, + }) + } + pub(crate) async fn hook_context_for_session( &self, session_id: SessionId, ) -> Option { - let session_arc = self.sessions.lock().await.get(&session_id).cloned()?; - let session = session_arc.lock().await; - let runner = session.runtime_context.hook_runner()?; - let transcript_path = session + if let Some(stream) = self.active_stream_state(session_id).await { + let stream = stream.lock().await; + if let Some(inline) = stream.turn_inline.as_ref() { + return Self::hook_runtime_context_from_snapshot(session_id, &inline.hook_context); + } + } + let session_handle = self.sessions.lock().await.get(&session_id).cloned()?; + let snapshot = session_handle.hook_context_snapshot().await?; + Some(Self::hook_runtime_context_from_snapshot(session_id, &snapshot)?) + } + + fn hook_runtime_context_from_snapshot( + session_id: SessionId, + snapshot: &HookContextSnapshot, + ) -> Option { + let runner = snapshot.runtime_context.hook_runner()?; + let transcript_path = snapshot .record .as_ref() .map(|record| record.rollout_path.display().to_string()) .unwrap_or_default(); - let permission_mode = Some(permission_mode_label(session.config.permission_mode)); - let agent_id = session + let permission_mode = Some(permission_mode_label(snapshot.config.permission_mode)); + let agent_id = snapshot .summary .parent_session_id .is_some() .then(|| session_id.to_string()); - let agent_type = session + let agent_type = snapshot .summary .agent_role .clone() - .or_else(|| session.summary.agent_nickname.clone()); + .or_else(|| snapshot.summary.agent_nickname.clone()); Some(devo_core::HookRuntimeContext { runner, base: devo_core::HookBaseInput { session_id: session_id.to_string(), transcript_path, - cwd: session.summary.cwd.clone(), + cwd: snapshot.summary.cwd.clone(), permission_mode, agent_id, agent_type, @@ -65,6 +115,20 @@ impl ServerRuntime { context.runner.run(input).await } + pub(crate) async fn run_session_hook_for_actor_state( + &self, + state: &SessionActorState, + session_id: SessionId, + event: devo_core::HookEvent, + extra: Map, + ) -> devo_core::HookRunReport { + let Some(context) = Self::hook_context_from_actor_state(state, session_id) else { + return devo_core::HookRunReport::default(); + }; + let input = devo_core::HookInput::new(&context.base, event, extra); + context.runner.run(input).await + } + pub(crate) async fn run_global_hook( &self, event: devo_core::HookEvent, diff --git a/crates/server/src/runtime/items.rs b/crates/server/src/runtime/items.rs index 14a40f2d..e04e3993 100644 --- a/crates/server/src/runtime/items.rs +++ b/crates/server/src/runtime/items.rs @@ -15,24 +15,25 @@ impl ServerRuntime { self.maybe_assign_provisional_title(session_id, user_input) .await; - let needs_title = { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { - return; - }; - let mut session = session_arc.lock().await; - if session.first_user_input.is_none() { - session.first_user_input = Some(user_input.to_string()); - } - let first_input = session.first_user_input.clone(); - let needs = matches!( - session.summary.title_state, - SessionTitleState::Unset | SessionTitleState::Provisional - ); - (needs, first_input) + let Some(session_handle) = self.session(session_id).await else { + return; }; - if needs_title.0 - && let Some(first_input) = needs_title.1 - { + let _ = session_handle + .set_first_user_input_if_unset(user_input.to_string()) + .await; + let Some(title_context) = session_handle.title_generation_context().await else { + return; + }; + let needs_title = matches!( + title_context.title_state, + SessionTitleState::Unset | SessionTitleState::Provisional + ); + if needs_title { + let first_input = session_handle + .export_runtime_session() + .await + .and_then(|session| session.first_user_input) + .unwrap_or_else(|| user_input.to_string()); let runtime = Arc::clone(self); tokio::spawn(async move { runtime @@ -50,39 +51,40 @@ impl ServerRuntime { let Some(candidate) = derive_provisional_title(first_user_input) else { return; }; - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { + let Some(session_handle) = self.session(session_id).await else { return; }; - - let updated_summary = { - let mut session = session_arc.lock().await; - if session.summary.title.is_some() - || !matches!(session.summary.title_state, SessionTitleState::Unset) - { - return; - } - - let previous_title = session.summary.title.clone(); - let updated_at = Utc::now(); - session.summary.title = Some(candidate.clone()); - session.summary.title_state = SessionTitleState::Provisional; - session.summary.updated_at = updated_at; - - if let Some(record) = session.record.as_mut() { - record.title = Some(candidate.clone()); - record.title_state = SessionTitleState::Provisional; - record.updated_at = updated_at; - if let Err(error) = self.rollout_store.append_title_update( - record, - candidate.clone(), - SessionTitleState::Provisional, - previous_title, - ) { - tracing::warn!(session_id = %session_id, error = %error, "failed to persist provisional title"); - } - } - session.summary.clone() + let Some(title_context) = session_handle.title_generation_context().await else { + return; }; + if title_context.title_state != SessionTitleState::Unset { + return; + } + let Some(summary) = session_handle.summary().await else { + return; + }; + if summary.title.is_some() { + return; + } + + let previous_title = summary.title.clone(); + let updated_at = Utc::now(); + let mut updated_summary = summary; + updated_summary.title = Some(candidate.clone()); + updated_summary.title_state = SessionTitleState::Provisional; + updated_summary.updated_at = updated_at; + session_handle.update_summary(updated_summary.clone()).await; + + if let Some(record) = session_handle.record().await.flatten() + && let Err(error) = self.rollout_store.append_title_update( + &record, + candidate.clone(), + SessionTitleState::Provisional, + previous_title, + ) + { + tracing::warn!(session_id = %session_id, error = %error, "failed to persist provisional title"); + } self.broadcast_event(ServerEvent::SessionTitleUpdated(SessionEventPayload { session: updated_summary, @@ -103,24 +105,21 @@ impl ServerRuntime { first_user_input: String, ) { for attempt in 1..=Self::MAX_TITLE_RETRIES { - let (model_selection, reasoning_effort_selection, should_skip, runtime_context) = { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { - return; - }; - let session = session_arc.lock().await; - ( - session_model_selection(&session.summary) - .map(str::to_string) - .unwrap_or_else(|| session.runtime_context.default_model.clone()), - session.summary.reasoning_effort_selection.clone(), - matches!(session.summary.title_state, SessionTitleState::Final(_)), - Arc::clone(&session.runtime_context), - ) + let Some(session_handle) = self.session(session_id).await else { + return; }; - - if should_skip { + let Some(title_context) = session_handle.title_generation_context().await else { + return; + }; + if matches!(title_context.title_state, SessionTitleState::Final(_)) { return; } + let model_selection = title_context + .model_selection + .clone() + .unwrap_or_else(|| title_context.runtime_context.default_model.clone()); + let reasoning_effort_selection = title_context.reasoning_effort_selection.clone(); + let runtime_context = title_context.runtime_context; let turn_config = runtime_context .resolve_turn_config(Some(model_selection.as_str()), reasoning_effort_selection); @@ -176,38 +175,29 @@ impl ServerRuntime { } }; - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { + let Some(session_handle) = self.session(session_id).await else { return; }; - let updated_summary = { - let mut session = session_arc.lock().await; - if matches!(session.summary.title_state, SessionTitleState::Final(_)) { - return; - } - - let previous_title = session.summary.title.clone(); - let updated_at = Utc::now(); - session.summary.title = Some(generated_title.clone()); - session.summary.title_state = - SessionTitleState::Final(SessionTitleFinalSource::ModelGenerated); - session.summary.updated_at = updated_at; - - if let Some(record) = session.record.as_mut() { - record.title = Some(generated_title.clone()); - record.title_state = - SessionTitleState::Final(SessionTitleFinalSource::ModelGenerated); - record.updated_at = updated_at; - if let Err(error) = self.rollout_store.append_title_update( - record, - generated_title.clone(), - record.title_state.clone(), - previous_title, - ) { - tracing::warn!(session_id = %session_id, error = %error, "failed to persist title"); - } - } - session.summary.clone() + let Some(updated_summary) = session_handle + .update_title( + generated_title.clone(), + SessionTitleState::Final(SessionTitleFinalSource::ModelGenerated), + ) + .await + .flatten() + else { + return; }; + if let Some(record) = session_handle.record().await.flatten() + && let Err(error) = self.rollout_store.append_title_update( + &record, + generated_title.clone(), + SessionTitleState::Final(SessionTitleFinalSource::ModelGenerated), + updated_summary.title.clone(), + ) + { + tracing::warn!(session_id = %session_id, error = %error, "failed to persist title"); + } self.broadcast_event(ServerEvent::SessionTitleUpdated(SessionEventPayload { session: updated_summary, @@ -355,60 +345,84 @@ impl ServerRuntime { turn_status: Option, worklog: Option, ) { - let session_arc = self.sessions.lock().await.get(&session_id).cloned(); - if let Some(session_arc) = session_arc { - let record = { - let mut session = session_arc.lock().await; - if let Some(history_item) = history_item_from_turn_item(&turn_item) { - session.history_items.push(history_item); + if let Some(stream) = self.active_stream_state(session_id).await { + let mut stream = stream.lock().await; + if let Some(inline) = stream.turn_inline.as_mut() { + if inline.turn_id == turn_id + && let Some(history_item) = history_item_from_turn_item(&turn_item) + { + inline.history_items.push(history_item); + } + if inline.turn_id == turn_id { + inline + .persisted_turn_items + .push(crate::execution::PersistedTurnItem { + turn_id, + turn_kind: inline.turn_kind.clone(), + item_id, + turn_item: turn_item.clone(), + }); } - let turn_kind = session - .active_turn - .as_ref() - .filter(|turn| turn.turn_id == turn_id) - .map(|turn| turn.kind.clone()) - .or_else(|| { - session - .latest_turn - .as_ref() - .filter(|turn| turn.turn_id == turn_id) - .map(|turn| turn.kind.clone()) - }) - .unwrap_or_default(); - session - .persisted_turn_items - .push(crate::execution::PersistedTurnItem { + if let Some(record) = inline.record.clone() { + let item = build_item_record( + session_id, turn_id, - turn_kind, item_id, - turn_item: turn_item.clone(), - }); - session.record.clone() - }; - if let Some(record) = record { - let item = build_item_record( - session_id, - turn_id, - item_id, - item_seq, - turn_item, - turn_status, - worklog, - ); - if let Err(error) = self.rollout_store.append_item(&record, item) { - tracing::warn!(session_id = %session_id, error = %error, "failed to persist item line"); + item_seq, + turn_item, + turn_status, + worklog, + ); + if let Err(error) = self.rollout_store.append_item(&record, item) { + tracing::warn!(session_id = %session_id, error = %error, "failed to persist item line"); + } } + return; + } + } + let Some(session_handle) = self.session(session_id).await else { + return; + }; + if let Some(history_item) = history_item_from_turn_item(&turn_item) { + session_handle.append_history_item(history_item).await; + } + let Some(prep) = session_handle.prepare_persist_item(turn_id).await else { + return; + }; + session_handle + .append_persisted_item(crate::execution::PersistedTurnItem { + turn_id, + turn_kind: prep.turn_kind, + item_id, + turn_item: turn_item.clone(), + }) + .await; + if let Some(record) = prep.record { + let item = build_item_record( + session_id, + turn_id, + item_id, + item_seq, + turn_item, + turn_status, + worklog, + ); + if let Err(error) = self.rollout_store.append_item(&record, item) { + tracing::warn!(session_id = %session_id, error = %error, "failed to persist item line"); } } } pub(super) async fn allocate_item_sequence(&self, session_id: SessionId) -> u64 { - let session_arc = self.sessions.lock().await.get(&session_id).cloned(); - if let Some(session_arc) = session_arc { - let mut session = session_arc.lock().await; - let item_seq = session.next_item_seq; - session.loaded_item_count += 1; - session.next_item_seq += 1; + if let Some(stream) = self.active_stream_state(session_id).await { + let mut stream = stream.lock().await; + if let Some(inline) = stream.turn_inline.as_mut() { + return inline.allocate_item_seq(); + } + } + if let Some(handle) = self.session(session_id).await + && let Some(item_seq) = handle.allocate_item_seq().await + { return item_seq; } 1 diff --git a/crates/server/src/runtime/lifecycle.rs b/crates/server/src/runtime/lifecycle.rs index 3c5ca117..2a65e6d4 100644 --- a/crates/server/src/runtime/lifecycle.rs +++ b/crates/server/src/runtime/lifecycle.rs @@ -4,16 +4,13 @@ impl ServerRuntime { /// Loads durable sessions from rollout files and installs them into the runtime map. /// Also restores token stats and pending queues from SQLite. pub async fn load_persisted_sessions(self: &Arc) -> anyhow::Result<()> { - let sessions = self.rollout_store.load_sessions(&self.deps).await?; + let mut sessions = self.rollout_store.load_sessions(&self.deps).await?; tracing::info!(session_count = sessions.len(), "loaded persisted sessions"); let mut restored_goal_stores = std::collections::HashMap::new(); - // Restore token stats and pending queues from SQLite - for (session_id, session_arc) in &sessions { - let mut session = session_arc.lock().await; - - if !session.summary.ephemeral - && let Err(err) = self.deps.db.upsert_session(&session.summary) + for (session_id, runtime_session) in sessions.iter_mut() { + if !runtime_session.summary.ephemeral + && let Err(err) = self.deps.db.upsert_session(&runtime_session.summary) { tracing::warn!( session_id = %session_id, @@ -25,13 +22,14 @@ impl ServerRuntime { match self.deps.db.get_stats(session_id) { Ok(Some(stats)) => { - session.summary.total_input_tokens = stats.total_input_tokens; - session.summary.total_output_tokens = stats.total_output_tokens; - session.summary.total_tokens = stats.total_tokens; - session.summary.total_cache_creation_tokens = stats.total_cache_creation_tokens; - session.summary.total_cache_read_tokens = stats.total_cache_read_tokens; - session.summary.prompt_token_estimate = stats.prompt_token_estimate; - if let Ok(mut core) = session.core_session.try_lock() { + runtime_session.summary.total_input_tokens = stats.total_input_tokens; + runtime_session.summary.total_output_tokens = stats.total_output_tokens; + runtime_session.summary.total_tokens = stats.total_tokens; + runtime_session.summary.total_cache_creation_tokens = + stats.total_cache_creation_tokens; + runtime_session.summary.total_cache_read_tokens = stats.total_cache_read_tokens; + runtime_session.summary.prompt_token_estimate = stats.prompt_token_estimate; + if let Ok(mut core) = runtime_session.core_session.try_lock() { core.total_input_tokens = stats.total_input_tokens; core.total_output_tokens = stats.total_output_tokens; core.total_tokens = stats.total_tokens; @@ -47,16 +45,17 @@ impl ServerRuntime { ); } Ok(None) => { - // No stats in database, persist current stats let stats = crate::db::SessionStats { - total_input_tokens: session.summary.total_input_tokens, - total_output_tokens: session.summary.total_output_tokens, - total_tokens: session.summary.total_tokens, - total_cache_creation_tokens: session.summary.total_cache_creation_tokens, - total_cache_read_tokens: session.summary.total_cache_read_tokens, + total_input_tokens: runtime_session.summary.total_input_tokens, + total_output_tokens: runtime_session.summary.total_output_tokens, + total_tokens: runtime_session.summary.total_tokens, + total_cache_creation_tokens: runtime_session + .summary + .total_cache_creation_tokens, + total_cache_read_tokens: runtime_session.summary.total_cache_read_tokens, last_input_tokens: 0, turn_count: 0, - prompt_token_estimate: session.summary.prompt_token_estimate, + prompt_token_estimate: runtime_session.summary.prompt_token_estimate, }; if let Err(err) = self.deps.db.update_stats(session_id, &stats) { tracing::warn!( @@ -75,7 +74,6 @@ impl ServerRuntime { } } - // Restore pending turn queue from SQLite match self .deps .db @@ -83,7 +81,7 @@ impl ServerRuntime { { Ok(items) => { if !items.is_empty() { - let core_session = session.core_session.lock().await; + let core_session = runtime_session.core_session.lock().await; let mut queue = core_session .pending_turn_queue .lock() @@ -105,7 +103,6 @@ impl ServerRuntime { } } - // Clear any stale btw inputs from previous session if let Err(err) = self .deps .db @@ -118,7 +115,6 @@ impl ServerRuntime { ); } - drop(session); match self.goal_durable_store.replay_goal_store(*session_id).await { Ok(Some(mut goal_store)) => { if let Some(goal) = goal_store.get() @@ -168,9 +164,10 @@ impl ServerRuntime { } } - let mut runtime_sessions = self.sessions.lock().await; - runtime_sessions.extend(sessions); - drop(runtime_sessions); + for (_, runtime_session) in sessions { + self.insert_session_actor(SessionActorState::from_runtime_session(runtime_session)) + .await; + } self.goal_stores.lock().await.extend(restored_goal_stores); Ok(()) } @@ -179,15 +176,10 @@ impl ServerRuntime { /// persists interrupted turn records. Called on graceful shutdown. pub async fn shutdown(self: &Arc) { self.command_exec_manager.terminate_all().await; - let session_ids: Vec = { - let sessions = self.sessions.lock().await; - sessions.keys().cloned().collect() - }; + let session_handles = self.list_session_handles().await; - for session_id in session_ids { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { - continue; - }; + for session_handle in session_handles { + let session_id = session_handle.id(); self.run_session_hook( session_id, @@ -196,23 +188,14 @@ impl ServerRuntime { ) .await; - let (deferred_assistant, deferred_reasoning, turn_id, record) = { - let mut session = session_arc.lock().await; - let turn_id = session.active_turn.as_ref().map(|t| t.turn_id); - ( - session.deferred_assistant.take(), - session.deferred_reasoning.take(), - turn_id, - session.record.clone(), - ) + let Some(snapshot) = session_handle.take_shutdown_deferred_snapshot().await else { + continue; }; - - let Some(turn_id) = turn_id else { + let Some(turn_id) = snapshot.active_turn_id else { continue; }; - // Complete deferred items before shutting down - if let Some((item_id, item_seq, text)) = deferred_assistant { + if let Some((item_id, item_seq, text)) = snapshot.deferred_assistant { self.complete_item( session_id, turn_id, @@ -224,7 +207,7 @@ impl ServerRuntime { ) .await; } - if let Some((item_id, item_seq, text)) = deferred_reasoning { + if let Some((item_id, item_seq, text)) = snapshot.deferred_reasoning { self.complete_item( session_id, turn_id, @@ -237,57 +220,30 @@ impl ServerRuntime { .await; } - // Mark turn as interrupted - let interrupted_turn = { - let mut session = session_arc.lock().await; - let Some(mut turn) = session.active_turn.take() else { - continue; - }; - if turn.turn_id != turn_id { - session.active_turn = Some(turn); - continue; - } - turn.status = TurnStatus::Interrupted; - turn.completed_at = Some(Utc::now()); - session.latest_turn = Some(turn.clone()); - session.summary.status = SessionRuntimeStatus::Idle; - session.summary.updated_at = Utc::now(); - session.summary.last_activity_at = session.summary.updated_at; - let token_totals = session.core_session.try_lock().ok().map(|core| { - ( - core.total_input_tokens, - core.total_output_tokens, - core.total_tokens, - ) - }); - if let Some((input, output, total)) = token_totals { - session.summary.total_input_tokens = input; - session.summary.total_output_tokens = output; - session.summary.total_tokens = total; - } - turn + let Some(interrupted_turn) = session_handle.interrupt_active_turn().await.flatten() + else { + continue; }; + if interrupted_turn.turn_id != turn_id { + continue; + } - // Persist interrupted turn record - if let Some(record) = record { - let (session_context, turn_context) = { - let session = session_arc.lock().await; - let core = session.core_session.lock().await; - ( - core.session_context.clone(), - core.latest_turn_context.clone(), - ) - }; - if let Err(error) = self.rollout_store.append_turn( + if let Some(record) = snapshot.record + && let Some(persistence) = session_handle.turn_persistence_snapshot().await + && let Err(error) = self.rollout_store.append_turn( &record, - build_turn_record(&interrupted_turn, session_context, turn_context), - ) { - tracing::warn!( - session_id = %session_id, - error = %error, - "failed to persist interrupted turn on shutdown" - ); - } + build_turn_record( + &interrupted_turn, + persistence.session_context, + persistence.latest_turn_context, + ), + ) + { + tracing::warn!( + session_id = %session_id, + error = %error, + "failed to persist interrupted turn on shutdown" + ); } tracing::info!( @@ -296,5 +252,18 @@ impl ServerRuntime { "completed deferred items and interrupted turn on shutdown" ); } + + let session_ids: Vec = self + .list_session_handles() + .await + .into_iter() + .map(|handle| handle.id()) + .collect(); + for session_id in session_ids { + if let Some(handle) = self.sessions.lock().await.get(&session_id).cloned() { + handle.shutdown().await; + } + self.remove_session_actor(session_id).await; + } } } diff --git a/crates/server/src/runtime/research.rs b/crates/server/src/runtime/research.rs index 4e8adb1c..35e96ff4 100644 --- a/crates/server/src/runtime/research.rs +++ b/crates/server/src/runtime/research.rs @@ -163,7 +163,7 @@ impl ServerRuntime { "research question is empty", ); } - let Some(session_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + let Some(session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, @@ -172,26 +172,23 @@ impl ServerRuntime { }; let now = Utc::now(); - let (effective_cwd, runtime_context) = { - let session = session_arc.lock().await; - let effective_cwd = params - .cwd - .clone() - .unwrap_or_else(|| session.summary.cwd.clone()); - let runtime_context = if params - .cwd - .as_ref() - .is_some_and(|cwd| cwd != &session.summary.cwd) - { - None - } else { - Some(Arc::clone(&session.runtime_context)) - }; - (effective_cwd, runtime_context) + let Some(reservation) = session_handle.turn_reservation_snapshot().await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); }; - let runtime_context = match runtime_context { - Some(runtime_context) => runtime_context, - None => match self.deps.context_for_workspace(&effective_cwd).await { + let effective_cwd = params + .cwd + .clone() + .unwrap_or_else(|| reservation.summary.cwd.clone()); + let runtime_context = if params + .cwd + .as_ref() + .is_some_and(|cwd| cwd != &reservation.summary.cwd) + { + match self.deps.context_for_workspace(&effective_cwd).await { Ok(runtime_context) => runtime_context, Err(error) => { return self.error_response( @@ -200,84 +197,76 @@ impl ServerRuntime { format!("failed to initialize session workspace: {error}"), ); } - }, + } + } else { + reservation.runtime_context }; let mut cwd_change = None; - let (turn, turn_config, effective_cwd) = { - let mut session = session_arc.lock().await; - if session.active_turn.is_some() { - return self.error_response( - request_id, - ProtocolErrorCode::TurnAlreadyRunning, - "cannot start research while a turn is already running", - ); - } - let requested_model = requested_model_selection( - params.model_binding_id.as_deref(), - params.model.as_deref(), - &session.summary, - ); - let requested_reasoning_effort_selection = params - .reasoning_effort_selection - .clone() - .or_else(|| session.summary.reasoning_effort_selection.clone()); - let turn_config = runtime_context.resolve_turn_config( - requested_model, - requested_reasoning_effort_selection.clone(), + if reservation.active_turn.is_some() { + return self.error_response( + request_id, + ProtocolErrorCode::TurnAlreadyRunning, + "cannot start research while a turn is already running", ); - let effective_cwd = params - .cwd - .clone() - .unwrap_or_else(|| session.summary.cwd.clone()); - if let Some(cwd) = params.cwd.clone() { - let old_cwd = session.summary.cwd.clone(); - if old_cwd != cwd { - cwd_change = Some((old_cwd, cwd.clone())); - session.runtime_context = Arc::clone(&runtime_context); - } - session.summary.cwd = cwd.clone(); - session.core_session.lock().await.cwd = cwd; - } - if let Some(permission_mode) = params - .approval_policy - .as_deref() - .and_then(permission_mode_from_approval_policy) - { - session.core_session.lock().await.config.permission_mode = permission_mode; - session.config.permission_mode = permission_mode; + } + if let Some(cwd) = params.cwd.clone() { + let old_cwd = reservation.summary.cwd.clone(); + if old_cwd != cwd { + cwd_change = Some((old_cwd, cwd.clone())); + session_handle + .update_session_workspace(cwd.clone(), Arc::clone(&runtime_context)) + .await; } - let resolved_request = turn_config.model.resolve_reasoning_effort_selection( - turn_config.reasoning_effort_selection.as_deref(), - ); - let request_model = turn_config.provider_request_model(&resolved_request.request_model); - apply_turn_config_to_session_summary(&mut session.summary, &turn_config); - let turn = TurnMetadata { - turn_id: TurnId::new(), - session_id: params.session_id, - sequence: session - .latest_turn - .as_ref() - .map_or(1, |turn| turn.sequence + 1), - status: TurnStatus::Running, - kind: devo_core::TurnKind::Research, - model: turn_config.model.slug.clone(), - model_binding_id: turn_config.model_binding_id.clone(), - reasoning_effort_selection: turn_config.reasoning_effort_selection.clone(), - reasoning_effort: resolved_request.effective_reasoning_effort, - request_model, - request_thinking: resolved_request.request_thinking, - started_at: now, - completed_at: None, - usage: None, - stop_reason: None, - failure_reason: None, - }; - session.summary.status = SessionRuntimeStatus::ActiveTurn; - session.summary.updated_at = now; - session.summary.last_activity_at = now; - session.active_turn = Some(turn.clone()); - (turn, turn_config, effective_cwd.display().to_string()) + } + if let Some(permission_mode) = params + .approval_policy + .as_deref() + .and_then(permission_mode_from_approval_policy) + { + session_handle + .update_core_permission_mode(permission_mode) + .await; + } + let requested_model = requested_model_selection( + params.model_binding_id.as_deref(), + params.model.as_deref(), + &reservation.summary, + ); + let requested_reasoning_effort_selection = params + .reasoning_effort_selection + .clone() + .or_else(|| reservation.summary.reasoning_effort_selection.clone()); + let turn_config = runtime_context + .resolve_turn_config(requested_model, requested_reasoning_effort_selection); + let resolved_request = turn_config + .model + .resolve_reasoning_effort_selection(turn_config.reasoning_effort_selection.as_deref()); + let request_model = turn_config.provider_request_model(&resolved_request.request_model); + let turn = TurnMetadata { + turn_id: TurnId::new(), + session_id: params.session_id, + sequence: reservation + .latest_turn + .as_ref() + .map_or(1, |turn| turn.sequence + 1), + status: TurnStatus::Running, + kind: devo_core::TurnKind::Research, + model: turn_config.model.slug.clone(), + model_binding_id: turn_config.model_binding_id.clone(), + reasoning_effort_selection: turn_config.reasoning_effort_selection.clone(), + reasoning_effort: resolved_request.effective_reasoning_effort, + request_model, + request_thinking: resolved_request.request_thinking, + started_at: now, + completed_at: None, + usage: None, + stop_reason: None, + failure_reason: None, }; + session_handle + .begin_active_turn(turn.clone(), turn_config.clone()) + .await; + let effective_cwd = effective_cwd.display().to_string(); if let Some((old_cwd, new_cwd)) = cwd_change { self.run_session_hook( @@ -301,6 +290,10 @@ impl ServerRuntime { .lock() .await .insert(params.session_id, CancellationToken::new()); + self.active_turn_ids + .lock() + .await + .insert(params.session_id, turn.turn_id); if let Some(connection_id) = connection_id { self.active_turn_connections .lock() @@ -308,62 +301,27 @@ impl ServerRuntime { .insert(params.session_id, connection_id); } let research_display_input = research_display_input(&display_input); - self.maybe_assign_provisional_title(params.session_id, &research_display_input) - .await; - { - let mut session = session_arc.lock().await; - if session.first_user_input.is_none() { - session.first_user_input = Some(research_display_input.clone()); - } - } - let needs_title = { - let session = session_arc.lock().await; - let first_input = session.first_user_input.clone(); - let needs = matches!( - session.summary.title_state, - SessionTitleState::Unset | SessionTitleState::Provisional - ); - (needs, first_input) - }; - if needs_title.0 - && let Some(first_input) = needs_title.1 - { - let runtime = Arc::clone(self); - let sid = params.session_id; - tokio::spawn(async move { - runtime.maybe_generate_final_title(sid, first_input).await; - }); - } - let (record, session_context, turn_context) = { - let session = session_arc.lock().await; - let core_session = session.core_session.lock().await; - ( - session.record.clone(), - core_session.session_context.clone(), - core_session.latest_turn_context.clone(), - ) - }; - if let Some(record) = record + self.maybe_start_title_generation_from_user_input( + params.session_id, + &research_display_input, + ) + .await; + if let Some(persistence) = session_handle.turn_persistence_snapshot().await + && let Some(record) = persistence.record && let Err(error) = self.rollout_store.append_turn( &record, - build_turn_record(&turn, session_context, turn_context), + build_turn_record( + &turn, + persistence.session_context, + persistence.latest_turn_context, + ), ) { self.clear_active_turn_runtime_handles(params.session_id) .await; - { - let mut session = session_arc.lock().await; - if session - .active_turn - .as_ref() - .is_some_and(|active| active.turn_id == turn.turn_id) - { - session.active_turn = None; - session.summary.status = SessionRuntimeStatus::Idle; - session.summary.updated_at = Utc::now(); - session.summary.last_activity_at = session.summary.updated_at; - } - } + let _ = session_handle + .clear_active_turn_if_matches(turn.turn_id) + .await; return self.error_response( request_id, ProtocolErrorCode::InternalError, @@ -560,14 +518,6 @@ impl ServerRuntime { let clarification_result = self .run_clarification_gate(&model_runtime, &mut coordinator_scratch) .await?; - self.emit_research_artifact( - session_id, - turn.turn_id, - ResearchArtifactType::Clarification, - "Research Clarification", - clarification_result.artifact_content, - ) - .await; research_context .clarifications .extend(clarification_result.clarifications.clone()); @@ -649,15 +599,19 @@ impl ServerRuntime { ) -> anyhow::Result { let request_id = format!("research_clarification_{turn_id}"); let (tx, rx) = tokio::sync::oneshot::channel(); - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { + let Some(session_handle) = self.session(session_id).await else { anyhow::bail!("session does not exist"); }; - { - let mut session = session_arc.lock().await; - session - .pending_user_inputs - .insert(request_id.clone(), PendingUserInput { turn_id, tx }); - session.summary.status = SessionRuntimeStatus::WaitingClient; + self.session_interactive + .register_pending_user_input( + session_id, + request_id.clone(), + PendingUserInput { turn_id, tx }, + ) + .await; + if let Some(mut summary) = session_handle.summary().await { + summary.status = SessionRuntimeStatus::WaitingClient; + session_handle.update_summary(summary).await; } self.broadcast_event(ServerEvent::SessionStatusChanged( SessionStatusChangedPayload { @@ -685,10 +639,9 @@ impl ServerRuntime { })) .await; let response = rx.await?; - let session_arc = self.sessions.lock().await.get(&session_id).cloned(); - if let Some(session_arc) = session_arc { - let mut session = session_arc.lock().await; - session.summary.status = SessionRuntimeStatus::ActiveTurn; + if let Some(mut summary) = session_handle.summary().await { + summary.status = SessionRuntimeStatus::ActiveTurn; + session_handle.update_summary(summary).await; } self.broadcast_event(ServerEvent::SessionStatusChanged( SessionStatusChangedPayload { @@ -799,7 +752,16 @@ impl ServerRuntime { let callback: devo_core::EventCallback = Arc::new(move |event: QueryEvent| { let tx = tx.clone(); Box::pin(async move { - let _ = tx.send(event).await; + match tx.try_send(event) { + Ok(()) => {} + Err(tokio::sync::mpsc::error::TrySendError::Full(event)) => { + let tx = tx.clone(); + tokio::spawn(async move { + let _ = tx.send(event).await; + }); + } + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {} + } }) }); let mut stage_turn_config = runtime.turn_config.clone(); @@ -888,17 +850,14 @@ impl ServerRuntime { } async fn refresh_core_session_prompt_context(&self, session_id: SessionId) { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { + let Some(session_handle) = self.session(session_id).await else { return; }; - let (persisted_turn_items, latest_compaction_snapshot, core_session) = { - let session = session_arc.lock().await; - ( - session.persisted_turn_items.clone(), - session.latest_compaction_snapshot.clone(), - Arc::clone(&session.core_session), - ) + let Some(runtime_session) = session_handle.export_runtime_session().await else { + return; }; + let persisted_turn_items = runtime_session.persisted_turn_items.clone(); + let latest_compaction_snapshot = runtime_session.latest_compaction_snapshot.clone(); let mut rebuilt_messages = Vec::new(); let mut ignored_history_items = Vec::new(); @@ -916,9 +875,18 @@ impl ServerRuntime { crate::persistence::build_prompt_messages_from_snapshot(&persisted_turn_items, snapshot) }); - let mut core_session = core_session.lock().await; - core_session.messages = rebuilt_messages; - core_session.prompt_messages = rebuilt_prompt_messages; + { + let mut core_session = runtime_session.core_session.lock().await; + core_session.messages = rebuilt_messages; + core_session.prompt_messages = rebuilt_prompt_messages; + } + session_handle + .replace_state( + crate::runtime::session_actor::SessionActorState::from_runtime_session( + runtime_session, + ), + ) + .await; } async fn run_clarification_gate( @@ -937,58 +905,75 @@ impl ServerRuntime { self.complete_reasoning_item(runtime.session_id, runtime.turn_id, &mut capture.reasoning) .await; - if !capture.request_user_input_exchanges.is_empty() { - return Ok(ClarificationGateResult { + let result = if !capture.request_user_input_exchanges.is_empty() { + ClarificationGateResult { artifact_content: clarification_artifact_content( &capture.request_user_input_exchanges, ), clarifications: capture.clarifications, - }); - } - - let clarify_text = if capture.text.trim().is_empty() { - assistant_text_from_session(scratch) + } } else { - capture.text.clone() - }; - if let Some(decision) = parse_json_object::(&clarify_text) { - if decision.need_clarification && !decision.question.trim().is_empty() { - let question = decision.question; - let answer = self - .request_research_clarification(runtime.session_id, runtime.turn_id, &question) - .await?; - let artifact_content = - format!("Question: {}\n\nAnswer: {}", question, answer.trim()); - let clarifications = if answer.trim().is_empty() { - Vec::new() + let clarify_text = if capture.text.trim().is_empty() { + assistant_text_from_session(scratch) + } else { + capture.text.clone() + }; + if let Some(decision) = parse_json_object::(&clarify_text) { + if decision.need_clarification && !decision.question.trim().is_empty() { + let question = decision.question; + let answer = self + .request_research_clarification( + runtime.session_id, + runtime.turn_id, + &question, + ) + .await?; + let artifact_content = + format!("Question: {}\n\nAnswer: {}", question, answer.trim()); + let clarifications = if answer.trim().is_empty() { + Vec::new() + } else { + vec![ResearchClarificationContext { question, answer }] + }; + ClarificationGateResult { + artifact_content, + clarifications, + } } else { - vec![ResearchClarificationContext { question, answer }] + let artifact_content = if decision.verification.trim().is_empty() { + "No clarification needed.".to_string() + } else { + decision.verification + }; + ClarificationGateResult { + artifact_content, + clarifications: Vec::new(), + } + } + } else { + let artifact_content = if clarify_text.trim().is_empty() { + "No clarification needed.".to_string() + } else { + clarify_text }; - return Ok(ClarificationGateResult { + ClarificationGateResult { artifact_content, - clarifications, - }); + clarifications: Vec::new(), + } } - let artifact_content = if decision.verification.trim().is_empty() { - "No clarification needed.".to_string() - } else { - decision.verification - }; - return Ok(ClarificationGateResult { - artifact_content, - clarifications: Vec::new(), - }); - } - - let artifact_content = if clarify_text.trim().is_empty() { - "No clarification needed.".to_string() - } else { - clarify_text }; - Ok(ClarificationGateResult { - artifact_content, - clarifications: Vec::new(), - }) + let artifact = ResearchStageKind::Clarify + .artifact() + .expect("clarify stage should have an artifact"); + self.complete_research_artifact_item( + runtime.session_id, + runtime.turn_id, + &mut capture.artifact, + &artifact, + &result.artifact_content, + ) + .await; + Ok(result) } async fn stream_final_report( @@ -1115,34 +1100,20 @@ impl ServerRuntime { .await .map(|snapshot| snapshot.turn_usage.to_turn_usage()) .unwrap_or(own_usage); - { - let session_arc = self.sessions.lock().await.get(&session_id).cloned(); - if let Some(session_arc) = session_arc { - let mut session = session_arc.lock().await; - turn.usage = Some(usage.clone()); - session.latest_turn = Some(turn.clone()); - session.active_turn = None; - session.summary.status = SessionRuntimeStatus::Idle; - session.summary.updated_at = Utc::now(); - session.summary.last_activity_at = session.summary.updated_at; - } + turn.usage = Some(usage.clone()); + if let Some(session_handle) = self.session(session_id).await { + session_handle.set_session_idle(Some(turn.clone())).await; } - let (record, session_context, turn_context) = { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { - return; - }; - let session = session_arc.lock().await; - let core_session = session.core_session.lock().await; - ( - session.record.clone(), - core_session.session_context.clone(), - core_session.latest_turn_context.clone(), - ) - }; - if let Some(record) = record + if let Some(session_handle) = self.session(session_id).await + && let Some(persistence) = session_handle.turn_persistence_snapshot().await + && let Some(record) = persistence.record && let Err(error) = self.rollout_store.append_turn( &record, - build_turn_record(&turn, session_context, turn_context), + build_turn_record( + &turn, + persistence.session_context, + persistence.latest_turn_context, + ), ) { tracing::warn!(session_id = %session_id, error = %error, "failed to persist research turn finish"); diff --git a/crates/server/src/runtime/research_capture.rs b/crates/server/src/runtime/research_capture.rs index f00ffb94..d78ca97a 100644 --- a/crates/server/src/runtime/research_capture.rs +++ b/crates/server/src/runtime/research_capture.rs @@ -18,6 +18,7 @@ pub(super) struct ResearchQueryCapture { #[derive(Default)] pub(super) struct ClarificationQueryCapture { pub(super) text: String, + pub(super) artifact: StreamedTextItem, pub(super) pending_request_user_input_questions: HashMap>, pub(super) request_user_input_exchanges: Vec, pub(super) clarifications: Vec, diff --git a/crates/server/src/runtime/research_events.rs b/crates/server/src/runtime/research_events.rs index 6ddedaf3..52b88030 100644 --- a/crates/server/src/runtime/research_events.rs +++ b/crates/server/src/runtime/research_events.rs @@ -37,7 +37,13 @@ impl ServerRuntime { ) -> anyhow::Result<()> { match capture { ResearchStageCapture::Clarification(capture) => { - self.handle_clarification_query_event(context, capture, event) + let artifact = artifact + .ok_or_else(|| anyhow::anyhow!("research {stage:?} missing artifact"))?; + let artifact_context = ResearchArtifactEventContext { + query: context, + artifact, + }; + self.handle_clarification_query_event(artifact_context, capture, event) .await; } ResearchStageCapture::Artifact(capture) => { @@ -454,17 +460,25 @@ impl ServerRuntime { async fn handle_clarification_query_event( &self, - context: ResearchQueryEventContext<'_>, + context: ResearchArtifactEventContext<'_>, capture: &mut ClarificationQueryCapture, event: QueryEvent, ) { - let session_id = context.session_id; - let turn_id = context.turn_id; - let usage_ledger = context.usage_ledger; - let context_window = context.context_window; + let session_id = context.query.session_id; + let turn_id = context.query.turn_id; + let usage_ledger = context.query.usage_ledger; + let context_window = context.query.context_window; match event { QueryEvent::TextDelta(text) => { capture.text.push_str(&text); + self.push_research_artifact_delta( + session_id, + turn_id, + &mut capture.artifact, + context.artifact, + text, + ) + .await; } QueryEvent::ToolUseStart { id, name, input, .. diff --git a/crates/server/src/runtime/research_stages.rs b/crates/server/src/runtime/research_stages.rs index 32c0ae65..728de1b8 100644 --- a/crates/server/src/runtime/research_stages.rs +++ b/crates/server/src/runtime/research_stages.rs @@ -74,7 +74,11 @@ impl ResearchStageKind { artifact_type: ResearchArtifactType::CompressedFinding, title: "Compressed Finding: Research Evidence".to_string(), }), - ResearchStageKind::Clarify | ResearchStageKind::FinalReport => None, + ResearchStageKind::Clarify => Some(StreamedResearchArtifact { + artifact_type: ResearchArtifactType::Clarification, + title: "Research Clarification".to_string(), + }), + ResearchStageKind::FinalReport => None, } } } @@ -132,7 +136,13 @@ mod tests { #[test] fn research_stage_artifacts_are_fixed() { - assert_eq!(ResearchStageKind::Clarify.artifact(), None); + assert_eq!( + ResearchStageKind::Clarify.artifact(), + Some(StreamedResearchArtifact { + artifact_type: ResearchArtifactType::Clarification, + title: "Research Clarification".to_string(), + }), + ); assert_eq!( ResearchStageKind::Brief.artifact(), Some(StreamedResearchArtifact { diff --git a/crates/server/src/runtime/research_tool_runtime.rs b/crates/server/src/runtime/research_tool_runtime.rs index cfd7cb0d..ff7e8d91 100644 --- a/crates/server/src/runtime/research_tool_runtime.rs +++ b/crates/server/src/runtime/research_tool_runtime.rs @@ -10,11 +10,13 @@ impl ServerRuntime { &self, session_id: SessionId, ) -> anyhow::Result { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { + let Some(session_handle) = self.session(session_id).await else { anyhow::bail!("session does not exist"); }; - let session = session_arc.lock().await; - let core_session = session.core_session.lock().await; + let Some(runtime_session) = session_handle.export_runtime_session().await else { + anyhow::bail!("session does not exist"); + }; + let core_session = runtime_session.core_session.lock().await; let mut scratch = SessionState::new(core_session.config.clone(), core_session.cwd.clone()); scratch.id = session_id.to_string(); Ok(scratch) @@ -27,19 +29,16 @@ impl ServerRuntime { turn_config: &TurnConfig, registry: Arc, ) -> anyhow::Result { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { + let Some(session_handle) = self.session(session_id).await else { anyhow::bail!("session does not exist"); }; - let (cwd, permission_mode, permission_profile, runtime_context) = { - let session = session_arc.lock().await; - let core_session = session.core_session.lock().await; - ( - core_session.cwd.clone(), - core_session.config.permission_mode, - core_session.config.permission_profile.clone(), - Arc::clone(&session.runtime_context), - ) + let Some(snapshot) = session_handle.hook_context_snapshot().await else { + anyhow::bail!("session does not exist"); }; + let cwd = snapshot.summary.cwd.clone(); + let permission_mode = snapshot.config.permission_mode; + let permission_profile = snapshot.config.permission_profile.clone(); + let runtime_context = snapshot.runtime_context; let provider_http = runtime_context .config_store .lock() diff --git a/crates/server/src/runtime/session_actor/commands.rs b/crates/server/src/runtime/session_actor/commands.rs new file mode 100644 index 00000000..e706708f --- /dev/null +++ b/crates/server/src/runtime/session_actor/commands.rs @@ -0,0 +1,198 @@ +use std::sync::Arc; + +use devo_core::SessionTitleState; +use devo_core::TurnId; +use devo_protocol::ApprovalScopeValue; +use devo_protocol::CollaborationMode; +use devo_protocol::PendingInputItem; +use devo_protocol::ThreadGoal; +use tokio::sync::oneshot; + +use super::snapshots::{ + HookContextSnapshot, PendingQueueSnapshot, PersistItemPrep, QueuedTurnInputData, + ShellExecContextSnapshot, ShutdownDeferredSnapshot, TitleGenerationContext, + TurnPersistenceSnapshot, TurnReservationSnapshot, +}; +use super::state::{ApprovalCacheSnapshot, DeferredItems, SessionActorState, SpawnSnapshot}; +use crate::execution::PendingApproval; +use crate::execution::PersistedTurnItem; +use crate::runtime::subagent_usage::ParentUsageSnapshot; +use crate::runtime::turn_exec::ExecuteTurnRequest; +use crate::session::SessionHistoryItem; +use crate::session::SessionMetadata; +use crate::turn::TurnMetadata; +use devo_core::TurnConfig; + +pub(crate) enum SessionCommand { + ExecuteTurn { + runtime: Arc, + request: ExecuteTurnRequest, + reply: oneshot::Sender<()>, + }, + GetSummary { + reply: oneshot::Sender, + }, + GetSpawnSnapshot { + reply: oneshot::Sender, + }, + GetApprovalCacheSnapshot { + reply: oneshot::Sender, + }, + GetCollaborationMode { + reply: oneshot::Sender, + }, + GetRuntimeContext { + reply: oneshot::Sender>, + }, + GetParentSessionId { + reply: oneshot::Sender>, + }, + GetTurnReservationSnapshot { + reply: oneshot::Sender, + }, + GetHookContextSnapshot { + reply: oneshot::Sender, + }, + GetTurnPersistenceSnapshot { + reply: oneshot::Sender, + }, + GetShellExecContext { + cwd: std::path::PathBuf, + reply: oneshot::Sender, + }, + GetTitleGenerationContext { + reply: oneshot::Sender, + }, + GetPendingQueueSnapshot { + reply: oneshot::Sender, + }, + PopQueuedTurnInput { + require_idle_session: bool, + reply: oneshot::Sender>, + }, + GetActiveTurnId { + reply: oneshot::Sender>, + }, + GetRecord { + reply: oneshot::Sender>, + }, + PreparePersistItem { + turn_id: TurnId, + reply: oneshot::Sender, + }, + TakeShutdownDeferredSnapshot { + reply: oneshot::Sender, + }, + AllocateItemSeq { + reply: oneshot::Sender, + }, + AppendPersistedItem { + item: PersistedTurnItem, + }, + AppendHistoryItem { + item: SessionHistoryItem, + }, + TakeDeferredItems { + reply: oneshot::Sender, + }, + ResetTurnApprovalCache, + TouchLastActivity, + ApplyApprovalScope { + scope: ApprovalScopeValue, + pending: PendingApproval, + }, + UpdateSummary { + summary: SessionMetadata, + }, + SetFirstUserInputIfUnset { + text: String, + reply: oneshot::Sender>, + }, + UpdateTitle { + title: String, + title_state: SessionTitleState, + reply: oneshot::Sender>, + }, + BeginActiveTurn { + turn: TurnMetadata, + turn_config: TurnConfig, + }, + ClearActiveTurnIfMatches { + turn_id: TurnId, + reply: oneshot::Sender, + }, + SetSessionIdle { + latest_turn: Option, + }, + EnqueuePendingTurnInput { + item: PendingInputItem, + }, + ActivateQueuedTurn { + turn: TurnMetadata, + turn_config: TurnConfig, + }, + CompleteShellTurn { + turn: TurnMetadata, + is_error: bool, + reply: oneshot::Sender, + }, + UpdateCorePermissionMode { + permission_mode: devo_safety::PermissionMode, + }, + SetActiveGoal { + goal: Option, + }, + #[cfg_attr(not(test), allow(dead_code))] + UpdateRecordRolloutPath { + rollout_path: std::path::PathBuf, + }, + ApplyParentUsageSnapshot { + snapshot: ParentUsageSnapshot, + }, + InterruptActiveTurn { + reply: oneshot::Sender>, + }, + ExportRuntimeSession { + reply: oneshot::Sender, + }, + UpdateSessionWorkspace { + cwd: std::path::PathBuf, + runtime_context: Arc, + }, + EnqueueBtwInput { + item: devo_protocol::PendingInputItem, + }, + UpdateSessionMetadata { + model: Option, + model_binding_id: Option, + reasoning_effort_selection: Option, + reply: oneshot::Sender, + }, + ApplyPermissionProfile { + profile: devo_safety::RuntimePermissionProfile, + reply: oneshot::Sender<()>, + }, + SetSessionTitleUserRename { + title: String, + reply: oneshot::Sender, + }, + SetToolRegistry { + tool_registry: Option>, + reply: oneshot::Sender<()>, + }, + GetResumeSnapshot { + reply: oneshot::Sender, + }, + TryBeginActiveTurn { + turn: TurnMetadata, + turn_config: TurnConfig, + reply: oneshot::Sender, + }, + ReplaceState { + state: Box, + reply: oneshot::Sender<()>, + }, + Shutdown { + reply: oneshot::Sender<()>, + }, +} diff --git a/crates/server/src/runtime/session_actor/handle.rs b/crates/server/src/runtime/session_actor/handle.rs new file mode 100644 index 00000000..9f55804c --- /dev/null +++ b/crates/server/src/runtime/session_actor/handle.rs @@ -0,0 +1,616 @@ +use std::path::PathBuf; +use std::sync::Arc; + +use devo_protocol::ApprovalScopeValue; +use devo_protocol::CollaborationMode; +use devo_protocol::PendingInputItem; +use devo_protocol::SessionId; +use devo_protocol::ThreadGoal; +use tokio::sync::mpsc; +use tokio::sync::oneshot; + +use devo_safety::PermissionMode; + +use super::commands::SessionCommand; +use super::snapshots::{ + HookContextSnapshot, PendingQueueSnapshot, PersistItemPrep, QueuedTurnInputData, + ShellExecContextSnapshot, ShutdownDeferredSnapshot, TitleGenerationContext, + TurnPersistenceSnapshot, TurnReservationSnapshot, +}; +use super::state::{ApprovalCacheSnapshot, DeferredItems, SessionActorState, SpawnSnapshot}; +use crate::execution::PendingApproval; +use crate::execution::PersistedTurnItem; +use crate::runtime::subagent_usage::ParentUsageSnapshot; +use crate::runtime::turn_exec::ExecuteTurnRequest; +use crate::session::SessionMetadata; +use crate::turn::TurnMetadata; +use devo_core::SessionRecord; +use devo_core::SessionTitleState; +use devo_core::TurnConfig; +use devo_core::TurnId; + +const SESSION_MAILBOX_CAPACITY: usize = 64; + +#[derive(Clone)] +pub(crate) struct SessionHandle { + session_id: SessionId, + tx: mpsc::Sender, +} + +impl SessionHandle { + pub(crate) fn id(&self) -> SessionId { + self.session_id + } + + pub(crate) fn spawn( + session_id: SessionId, + state: SessionActorState, + runtime: Arc, + ) -> Self { + let (tx, rx) = mpsc::channel(SESSION_MAILBOX_CAPACITY); + let handle = Self { session_id, tx }; + tokio::spawn(super::loop_::run_session_actor(state, rx, runtime)); + handle + } + + async fn send(&self, command: SessionCommand) -> bool { + self.tx.send(command).await.is_ok() + } + + pub(crate) async fn execute_turn( + &self, + runtime: Arc, + request: ExecuteTurnRequest, + ) { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::ExecuteTurn { + runtime, + request, + reply: reply_tx, + }) + .await + { + return; + } + let _ = reply_rx.await; + } + + pub(crate) async fn summary(&self) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::GetSummary { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn spawn_snapshot(&self) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::GetSpawnSnapshot { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn approval_cache_snapshot(&self) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::GetApprovalCacheSnapshot { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn collaboration_mode(&self) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::GetCollaborationMode { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn set_active_goal(&self, goal: Option) { + let _ = self.send(SessionCommand::SetActiveGoal { goal }).await; + } + + pub(crate) async fn runtime_context( + &self, + ) -> Option> { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::GetRuntimeContext { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn parent_session_id(&self) -> Option> { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::GetParentSessionId { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn turn_reservation_snapshot(&self) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::GetTurnReservationSnapshot { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn hook_context_snapshot(&self) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::GetHookContextSnapshot { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn turn_persistence_snapshot(&self) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::GetTurnPersistenceSnapshot { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn shell_exec_context( + &self, + cwd: std::path::PathBuf, + ) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::GetShellExecContext { + cwd, + reply: reply_tx, + }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn title_generation_context(&self) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::GetTitleGenerationContext { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn pending_queue_snapshot(&self) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::GetPendingQueueSnapshot { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn pop_queued_turn_input( + &self, + require_idle_session: bool, + ) -> Option> { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::PopQueuedTurnInput { + require_idle_session, + reply: reply_tx, + }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn active_turn_id(&self) -> Option> { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::GetActiveTurnId { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn record(&self) -> Option> { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::GetRecord { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn prepare_persist_item(&self, turn_id: TurnId) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::PreparePersistItem { + turn_id, + reply: reply_tx, + }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn take_shutdown_deferred_snapshot(&self) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::TakeShutdownDeferredSnapshot { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn allocate_item_seq(&self) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::AllocateItemSeq { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn append_persisted_item(&self, item: PersistedTurnItem) { + let _ = self + .send(SessionCommand::AppendPersistedItem { item }) + .await; + } + + pub(crate) async fn append_history_item(&self, item: crate::session::SessionHistoryItem) { + let _ = self.send(SessionCommand::AppendHistoryItem { item }).await; + } + + pub(crate) async fn take_deferred_items(&self) -> DeferredItems { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::TakeDeferredItems { reply: reply_tx }) + .await + { + return DeferredItems::default(); + } + reply_rx.await.unwrap_or_default() + } + + pub(crate) async fn reset_turn_approval_cache(&self) { + let _ = self.send(SessionCommand::ResetTurnApprovalCache).await; + } + + pub(crate) async fn touch_last_activity(&self) { + let _ = self.send(SessionCommand::TouchLastActivity).await; + } + + pub(crate) async fn apply_approval_scope( + &self, + scope: ApprovalScopeValue, + pending: PendingApproval, + ) { + let _ = self + .send(SessionCommand::ApplyApprovalScope { scope, pending }) + .await; + } + + pub(crate) async fn replace_state(&self, state: SessionActorState) { + let (reply_tx, reply_rx) = oneshot::channel(); + if self + .send(SessionCommand::ReplaceState { + state: Box::new(state), + reply: reply_tx, + }) + .await + { + let _ = reply_rx.await; + } + } + + pub(crate) async fn update_summary(&self, summary: SessionMetadata) { + let _ = self.send(SessionCommand::UpdateSummary { summary }).await; + } + + pub(crate) async fn set_first_user_input_if_unset( + &self, + text: String, + ) -> Option> { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::SetFirstUserInputIfUnset { + text, + reply: reply_tx, + }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn update_title( + &self, + title: String, + title_state: SessionTitleState, + ) -> Option> { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::UpdateTitle { + title, + title_state, + reply: reply_tx, + }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn begin_active_turn(&self, turn: TurnMetadata, turn_config: TurnConfig) { + let _ = self + .send(SessionCommand::BeginActiveTurn { turn, turn_config }) + .await; + } + + pub(crate) async fn clear_active_turn_if_matches(&self, turn_id: TurnId) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::ClearActiveTurnIfMatches { + turn_id, + reply: reply_tx, + }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn set_session_idle(&self, latest_turn: Option) { + let _ = self + .send(SessionCommand::SetSessionIdle { latest_turn }) + .await; + } + + pub(crate) async fn enqueue_pending_turn_input(&self, item: PendingInputItem) { + let _ = self + .send(SessionCommand::EnqueuePendingTurnInput { item }) + .await; + } + + pub(crate) async fn activate_queued_turn(&self, turn: TurnMetadata, turn_config: TurnConfig) { + let _ = self + .send(SessionCommand::ActivateQueuedTurn { turn, turn_config }) + .await; + } + + pub(crate) async fn complete_shell_turn( + &self, + turn: TurnMetadata, + is_error: bool, + ) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::CompleteShellTurn { + turn, + is_error, + reply: reply_tx, + }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn update_core_permission_mode(&self, permission_mode: PermissionMode) { + let _ = self + .send(SessionCommand::UpdateCorePermissionMode { permission_mode }) + .await; + } + + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) async fn update_record_rollout_path(&self, rollout_path: PathBuf) { + let _ = self + .send(SessionCommand::UpdateRecordRolloutPath { rollout_path }) + .await; + } + + pub(crate) async fn apply_parent_usage_snapshot(&self, snapshot: ParentUsageSnapshot) { + let _ = self + .send(SessionCommand::ApplyParentUsageSnapshot { snapshot }) + .await; + } + + pub(crate) async fn interrupt_active_turn(&self) -> Option> { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::InterruptActiveTurn { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn export_runtime_session(&self) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::ExportRuntimeSession { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn update_session_workspace( + &self, + cwd: PathBuf, + runtime_context: Arc, + ) { + let _ = self + .send(SessionCommand::UpdateSessionWorkspace { + cwd, + runtime_context, + }) + .await; + } + + pub(crate) async fn enqueue_btw_input(&self, item: devo_protocol::PendingInputItem) { + let _ = self.send(SessionCommand::EnqueueBtwInput { item }).await; + } + + pub(crate) async fn update_session_metadata( + &self, + model: Option, + model_binding_id: Option, + reasoning_effort_selection: Option, + ) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::UpdateSessionMetadata { + model, + model_binding_id, + reasoning_effort_selection, + reply: reply_tx, + }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn apply_permission_profile( + &self, + profile: devo_safety::RuntimePermissionProfile, + ) -> bool { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::ApplyPermissionProfile { + profile, + reply: reply_tx, + }) + .await + { + return false; + } + reply_rx.await.is_ok() + } + + pub(crate) async fn set_session_title_user_rename( + &self, + title: String, + ) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::SetSessionTitleUserRename { + title, + reply: reply_tx, + }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn set_tool_registry( + &self, + tool_registry: Option>, + ) -> bool { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::SetToolRegistry { + tool_registry, + reply: reply_tx, + }) + .await + { + return false; + } + reply_rx.await.is_ok() + } + + pub(crate) async fn resume_snapshot(&self) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::GetResumeSnapshot { reply: reply_tx }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn try_begin_active_turn( + &self, + turn: TurnMetadata, + turn_config: TurnConfig, + ) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::TryBeginActiveTurn { + turn, + turn_config, + reply: reply_tx, + }) + .await + { + return None; + } + reply_rx.await.ok() + } + + pub(crate) async fn shutdown(&self) { + let (reply_tx, reply_rx) = oneshot::channel(); + if self + .send(SessionCommand::Shutdown { reply: reply_tx }) + .await + { + let _ = reply_rx.await; + } + } +} diff --git a/crates/server/src/runtime/session_actor/loop_.rs b/crates/server/src/runtime/session_actor/loop_.rs new file mode 100644 index 00000000..91b6639c --- /dev/null +++ b/crates/server/src/runtime/session_actor/loop_.rs @@ -0,0 +1,590 @@ +use std::sync::Arc; + +use chrono::Utc; +use devo_core::SessionTitleFinalSource; +use devo_core::SessionTitleState; +use devo_core::TurnConfig; +use devo_core::TurnStatus; +use devo_protocol::ApprovalScopeValue; +use tokio::sync::mpsc; + +use super::commands::SessionCommand; +use super::snapshots::{ + HookContextSnapshot, PendingQueueSnapshot, QueuedTurnInputData, ShellExecContextSnapshot, + TitleGenerationContext, TurnPersistenceSnapshot, TurnReservationSnapshot, +}; +use super::state::SessionActorState; +use super::turn::execute_turn_in_actor; +use crate::SessionRuntimeStatus; +use crate::execution::PendingApproval; +use crate::runtime::session_model_selection; + +pub(super) async fn run_session_actor( + mut state: SessionActorState, + mut mailbox: mpsc::Receiver, + _runtime: Arc, +) { + while let Some(command) = mailbox.recv().await { + match command { + SessionCommand::ExecuteTurn { + runtime: turn_runtime, + request, + reply, + } => { + let session_id = request.session_id; + execute_turn_in_actor(&mut state, turn_runtime.clone(), request).await; + let _ = reply.send(()); + tokio::spawn(async move { + if !turn_runtime.chain_queued_followup_turn(session_id).await { + turn_runtime + .maybe_start_goal_continuation_turn(session_id) + .await; + } + }); + } + SessionCommand::GetSummary { reply } => { + let _ = reply.send(state.summary.clone()); + } + SessionCommand::GetSpawnSnapshot { reply } => { + let snapshot = state.spawn_snapshot(); + let _ = reply.send(snapshot); + } + SessionCommand::GetApprovalCacheSnapshot { reply } => { + let _ = reply.send(state.approval_cache_snapshot()); + } + SessionCommand::GetCollaborationMode { reply } => { + let _ = reply.send(state.core.collaboration_mode); + } + SessionCommand::GetRuntimeContext { reply } => { + let _ = reply.send(Arc::clone(&state.runtime_context)); + } + SessionCommand::GetParentSessionId { reply } => { + let _ = reply.send(state.parent_session_id()); + } + SessionCommand::GetTurnReservationSnapshot { reply } => { + let _ = reply.send(TurnReservationSnapshot { + max_turns: state.max_turns, + active_turn: state.active_turn.clone(), + latest_turn: state.latest_turn.clone(), + pending_turn_queue: Arc::clone(&state.pending_turn_queue), + ephemeral: state.summary.ephemeral, + parent_session_id: state.parent_session_id(), + summary: state.summary.clone(), + runtime_context: Arc::clone(&state.runtime_context), + }); + } + SessionCommand::GetHookContextSnapshot { reply } => { + let _ = reply.send(HookContextSnapshot { + runtime_context: Arc::clone(&state.runtime_context), + record: state.record.clone(), + summary: state.summary.clone(), + config: state.config.clone(), + }); + } + SessionCommand::GetTurnPersistenceSnapshot { reply } => { + let _ = reply.send(TurnPersistenceSnapshot { + record: state.record.clone(), + session_context: state.core.session_context.clone(), + latest_turn_context: state.core.latest_turn_context.clone(), + }); + } + SessionCommand::GetShellExecContext { cwd, reply } => { + let _ = &cwd; + let tool_registry = state + .tool_registry + .clone() + .unwrap_or_else(|| Arc::clone(&state.runtime_context.registry)); + let _ = reply.send(ShellExecContextSnapshot { + permission_mode: state.core.config.permission_mode, + permission_profile: state.core.config.permission_profile.clone(), + runtime_context: Arc::clone(&state.runtime_context), + tool_registry, + }); + } + SessionCommand::GetTitleGenerationContext { reply } => { + let _ = reply.send(TitleGenerationContext { + model_selection: session_model_selection(&state.summary).map(str::to_string), + reasoning_effort_selection: state.summary.reasoning_effort_selection.clone(), + title_state: state.summary.title_state.clone(), + runtime_context: Arc::clone(&state.runtime_context), + }); + } + SessionCommand::GetPendingQueueSnapshot { reply } => { + let queue = state + .pending_turn_queue + .lock() + .expect("pending turn queue mutex should not be poisoned"); + let pending_texts: Vec = queue + .iter() + .filter_map(|item| match &item.kind { + devo_core::PendingInputKind::UserText { text } => Some(text.clone()), + devo_core::PendingInputKind::UserInput { display_text, .. } => { + Some(display_text.clone()) + } + _ => None, + }) + .collect(); + let pending_count = pending_texts.len(); + let _ = reply.send(PendingQueueSnapshot { + pending_count, + pending_texts, + }); + } + SessionCommand::PopQueuedTurnInput { + require_idle_session, + reply, + } => { + if require_idle_session && state.active_turn.is_some() { + let _ = reply.send(None); + continue; + } + let mut queue = state + .pending_turn_queue + .lock() + .expect("pending turn queue mutex should not be poisoned"); + let popped = queue.pop_front().and_then(pop_queued_turn_input_data); + let _ = reply.send(popped); + } + SessionCommand::GetActiveTurnId { reply } => { + let _ = reply.send(state.active_turn.as_ref().map(|turn| turn.turn_id)); + } + SessionCommand::GetRecord { reply } => { + let _ = reply.send(state.record.clone()); + } + SessionCommand::PreparePersistItem { turn_id, reply } => { + let turn_kind = state + .active_turn + .as_ref() + .filter(|turn| turn.turn_id == turn_id) + .map(|turn| turn.kind.clone()) + .or_else(|| { + state + .latest_turn + .as_ref() + .filter(|turn| turn.turn_id == turn_id) + .map(|turn| turn.kind.clone()) + }) + .unwrap_or_default(); + let _ = reply.send(super::snapshots::PersistItemPrep { + turn_kind, + record: state.record.clone(), + }); + } + SessionCommand::TakeShutdownDeferredSnapshot { reply } => { + let stream = state.stream.lock().await; + let _ = reply.send(super::snapshots::ShutdownDeferredSnapshot { + deferred_assistant: stream.deferred_assistant.clone(), + deferred_reasoning: stream.deferred_reasoning.clone(), + active_turn_id: state.active_turn.as_ref().map(|turn| turn.turn_id), + record: state.record.clone(), + }); + } + SessionCommand::AllocateItemSeq { reply } => { + let item_seq = state.next_item_seq; + state.next_item_seq = state.next_item_seq.saturating_add(1); + state.loaded_item_count = state.loaded_item_count.saturating_add(1); + let _ = reply.send(item_seq); + } + SessionCommand::AppendPersistedItem { item } => { + state.persisted_turn_items.push(item); + } + SessionCommand::AppendHistoryItem { item } => { + state.history_items.push(item); + } + SessionCommand::TakeDeferredItems { reply } => { + let _ = reply.send(state.stream.lock().await.take_deferred_items()); + } + SessionCommand::ResetTurnApprovalCache => { + state.turn_approval_cache = crate::execution::ApprovalGrantCache::default(); + } + SessionCommand::TouchLastActivity => { + state.summary.last_activity_at = state.summary.last_activity_at.max(Utc::now()); + } + SessionCommand::ApplyApprovalScope { scope, pending } => { + apply_approval_scope_to_state(&mut state, &scope, &pending); + } + SessionCommand::UpdateSummary { summary } => { + state.summary = summary; + } + SessionCommand::SetFirstUserInputIfUnset { text, reply } => { + if state.first_user_input.is_none() { + state.first_user_input = Some(text.clone()); + } + let _ = reply.send(state.first_user_input.clone()); + } + SessionCommand::UpdateTitle { + title, + title_state, + reply, + } => { + if matches!(state.summary.title_state, SessionTitleState::Final(_)) { + let _ = reply.send(None); + continue; + } + let updated_at = Utc::now(); + state.summary.title = Some(title.clone()); + state.summary.title_state = title_state.clone(); + state.summary.updated_at = updated_at; + if let Some(record) = state.record.as_mut() { + record.title = Some(title); + record.title_state = title_state; + record.updated_at = updated_at; + } + let _ = reply.send(Some(state.summary.clone())); + } + SessionCommand::BeginActiveTurn { turn, turn_config } => { + let now = Utc::now(); + apply_turn_config_to_session_summary(&mut state.summary, &turn_config); + state.summary.status = SessionRuntimeStatus::ActiveTurn; + state.summary.updated_at = now; + state.summary.last_activity_at = now; + state.active_turn = Some(turn); + } + SessionCommand::ClearActiveTurnIfMatches { turn_id, reply } => { + let cleared = state + .active_turn + .as_ref() + .is_some_and(|active| active.turn_id == turn_id); + if cleared { + state.active_turn = None; + state.summary.status = SessionRuntimeStatus::Idle; + state.summary.updated_at = Utc::now(); + state.summary.last_activity_at = state.summary.updated_at; + } + let _ = reply.send(cleared); + } + SessionCommand::SetSessionIdle { latest_turn } => { + let now = Utc::now(); + if let Some(latest_turn) = latest_turn { + state.latest_turn = Some(latest_turn); + } + state.active_turn = None; + state.summary.status = SessionRuntimeStatus::Idle; + state.summary.updated_at = now; + state.summary.last_activity_at = now; + } + SessionCommand::SetActiveGoal { goal } => match goal { + Some(goal) => state.core.set_active_goal(goal), + None => state.core.clear_active_goal(), + }, + SessionCommand::EnqueuePendingTurnInput { item } => { + state + .pending_turn_queue + .lock() + .expect("pending turn queue mutex should not be poisoned") + .push_back(item); + } + SessionCommand::ActivateQueuedTurn { turn, turn_config } => { + let now = Utc::now(); + apply_turn_config_to_session_summary(&mut state.summary, &turn_config); + state.summary.status = SessionRuntimeStatus::ActiveTurn; + state.summary.updated_at = now; + state.summary.last_activity_at = now; + state.active_turn = Some(turn); + } + SessionCommand::CompleteShellTurn { + turn, + is_error, + reply, + } => { + let mut final_turn = turn; + final_turn.completed_at = Some(Utc::now()); + final_turn.status = if is_error { + TurnStatus::Failed + } else { + TurnStatus::Completed + }; + state.latest_turn = Some(final_turn.clone()); + state.active_turn = None; + state.summary.status = SessionRuntimeStatus::Idle; + state.summary.updated_at = Utc::now(); + state.summary.last_activity_at = state.summary.updated_at; + let _ = reply.send(final_turn); + } + SessionCommand::UpdateCorePermissionMode { permission_mode } => { + state.core.config.permission_mode = permission_mode; + } + SessionCommand::UpdateRecordRolloutPath { rollout_path } => { + if let Some(record) = state.record.as_mut() { + record.rollout_path = rollout_path; + } + } + SessionCommand::ApplyParentUsageSnapshot { snapshot } => { + snapshot.apply_to_actor_state(&mut state); + } + SessionCommand::InterruptActiveTurn { reply } => { + let now = Utc::now(); + state.summary.status = SessionRuntimeStatus::Idle; + state.summary.updated_at = now; + state.summary.last_activity_at = now; + state.summary.total_input_tokens = state.core.total_input_tokens; + state.summary.total_output_tokens = state.core.total_output_tokens; + state.summary.total_tokens = state.core.total_tokens; + state.summary.total_cache_creation_tokens = state.core.total_cache_creation_tokens; + state.summary.total_cache_read_tokens = state.core.total_cache_read_tokens; + state.summary.prompt_token_estimate = state.core.prompt_token_estimate; + let interrupted = state.active_turn.take().map(|mut turn| { + turn.status = TurnStatus::Interrupted; + turn.completed_at = Some(now); + state.latest_turn = Some(turn.clone()); + turn + }); + let _ = reply.send(interrupted); + } + SessionCommand::ExportRuntimeSession { reply } => { + let stream = state.stream.lock().await; + let _ = reply.send(state.to_runtime_session_from_stream(&stream)); + } + SessionCommand::UpdateSessionWorkspace { + cwd, + runtime_context, + } => { + state.runtime_context = runtime_context; + state.core.cwd = cwd.clone(); + state.summary.cwd = cwd; + } + SessionCommand::EnqueueBtwInput { item } => { + state + .btw_input_queue + .lock() + .expect("btw input queue mutex should not be poisoned") + .push_back(item); + } + SessionCommand::UpdateSessionMetadata { + model, + model_binding_id, + reasoning_effort_selection, + reply, + } => { + let updated_at = Utc::now(); + state.summary.model = model.clone(); + state.summary.model_binding_id = model_binding_id.clone(); + state.summary.reasoning_effort_selection = reasoning_effort_selection.clone(); + state.summary.updated_at = updated_at; + if let Some(record) = state.record.as_mut() { + record.model = model; + record.model_binding_id = model_binding_id; + record.reasoning_effort_selection = reasoning_effort_selection; + record.updated_at = updated_at; + } + let _ = reply.send(state.summary.clone()); + } + SessionCommand::ApplyPermissionProfile { profile, reply } => { + state.core.config.permission_mode = profile.permission_mode(); + state.core.config.permission_profile = profile.clone(); + state.config.permission_mode = profile.permission_mode(); + state.config.permission_profile = profile; + state.session_approval_cache = crate::execution::ApprovalGrantCache::default(); + state.turn_approval_cache = crate::execution::ApprovalGrantCache::default(); + let _ = reply.send(()); + } + SessionCommand::SetSessionTitleUserRename { title, reply } => { + let updated_at = Utc::now(); + state.summary.title = Some(title.clone()); + state.summary.title_state = + SessionTitleState::Final(SessionTitleFinalSource::UserRename); + state.summary.updated_at = updated_at; + if let Some(record) = state.record.as_mut() { + record.title = Some(title); + record.title_state = + SessionTitleState::Final(SessionTitleFinalSource::UserRename); + record.updated_at = updated_at; + } + let _ = reply.send(state.summary.clone()); + } + SessionCommand::SetToolRegistry { + tool_registry, + reply, + } => { + state.tool_registry = tool_registry; + let _ = reply.send(()); + } + SessionCommand::GetResumeSnapshot { reply } => { + let pending_texts = state + .pending_turn_queue + .lock() + .expect("pending turn queue mutex should not be poisoned") + .iter() + .filter_map(|item| match &item.kind { + devo_core::PendingInputKind::UserText { text } => Some(text.clone()), + devo_core::PendingInputKind::UserInput { display_text, .. } => { + Some(display_text.clone()) + } + _ => None, + }) + .collect(); + let _ = reply.send(super::snapshots::SessionResumeSnapshot { + summary: state.summary.clone(), + latest_turn: state.latest_turn.clone(), + loaded_item_count: state.loaded_item_count, + history_items: state.history_items.clone(), + pending_texts, + }); + } + SessionCommand::TryBeginActiveTurn { + turn, + turn_config, + reply, + } => { + let queue_empty = state + .pending_turn_queue + .lock() + .expect("pending turn queue mutex should not be poisoned") + .is_empty(); + if state.active_turn.is_some() || !queue_empty { + let _ = reply.send(false); + continue; + } + let now = Utc::now(); + apply_turn_config_to_session_summary(&mut state.summary, &turn_config); + state.summary.status = SessionRuntimeStatus::ActiveTurn; + state.summary.updated_at = now; + state.summary.last_activity_at = now; + state.active_turn = Some(turn); + let _ = reply.send(true); + } + SessionCommand::ReplaceState { + state: new_state, + reply, + } => { + state = *new_state; + let _ = reply.send(()); + } + SessionCommand::Shutdown { reply } => { + let _ = reply.send(()); + break; + } + } + } +} + +fn apply_turn_config_to_session_summary( + summary: &mut crate::session::SessionMetadata, + turn_config: &TurnConfig, +) { + summary.model = Some(turn_config.model.slug.clone()); + summary.model_binding_id = turn_config.model_binding_id.clone(); + summary.reasoning_effort_selection = turn_config.reasoning_effort_selection.clone(); +} + +fn pop_queued_turn_input_data( + item: devo_protocol::PendingInputItem, +) -> Option { + match item.kind { + devo_core::PendingInputKind::UserText { text } => Some(QueuedTurnInputData { + display_input: text.clone(), + input_text: text, + input_messages: Vec::new(), + collaboration_mode: collaboration_mode_from_pending_metadata(item.metadata.as_ref()), + model_selection: model_selection_from_pending_metadata(item.metadata.as_ref()), + subagent_usage_owner: subagent_usage_owner_from_pending_metadata( + item.metadata.as_ref(), + ), + }), + devo_core::PendingInputKind::UserInput { + display_text, + prompt_text, + prompt_messages, + .. + } => Some(QueuedTurnInputData { + display_input: display_text, + input_text: prompt_text, + input_messages: prompt_messages, + collaboration_mode: collaboration_mode_from_pending_metadata(item.metadata.as_ref()), + model_selection: model_selection_from_pending_metadata(item.metadata.as_ref()), + subagent_usage_owner: subagent_usage_owner_from_pending_metadata( + item.metadata.as_ref(), + ), + }), + _ => None, + } +} + +fn collaboration_mode_from_pending_metadata( + metadata: Option<&serde_json::Value>, +) -> devo_protocol::CollaborationMode { + metadata + .and_then(|metadata| { + metadata + .get("collaboration_mode") + .or_else(|| metadata.get("interaction_mode")) + }) + .cloned() + .and_then(|value| serde_json::from_value(value).ok()) + .unwrap_or_default() +} + +fn string_field_from_pending_metadata( + metadata: Option<&serde_json::Value>, + key: &str, +) -> Option { + metadata? + .get(key)? + .as_str() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) +} + +fn model_selection_from_pending_metadata(metadata: Option<&serde_json::Value>) -> Option { + string_field_from_pending_metadata(metadata, "model_binding_id") + .or_else(|| string_field_from_pending_metadata(metadata, "model")) +} + +fn subagent_usage_owner_from_pending_metadata( + metadata: Option<&serde_json::Value>, +) -> Option<(devo_protocol::SessionId, Option)> { + let parent_session_id = + string_field_from_pending_metadata(metadata, "devo_subagent_usage_parent_session_id") + .and_then(|value| devo_protocol::SessionId::try_from(value).ok())?; + let parent_turn_id = + string_field_from_pending_metadata(metadata, "devo_subagent_usage_parent_turn_id") + .and_then(|value| devo_core::TurnId::try_from(value).ok()); + Some((parent_session_id, parent_turn_id)) +} + +fn apply_approval_scope_to_state( + state: &mut SessionActorState, + scope: &ApprovalScopeValue, + pending: &PendingApproval, +) { + match scope { + ApprovalScopeValue::Once => {} + ApprovalScopeValue::Turn => { + state + .turn_approval_cache + .tools + .insert(pending.tool_name.clone()); + } + ApprovalScopeValue::Session => { + state + .session_approval_cache + .tools + .insert(pending.tool_name.clone()); + } + ApprovalScopeValue::PathPrefix => { + if let Some(path) = pending.path.clone() { + state.turn_approval_cache.path_prefixes.insert(path); + } + } + ApprovalScopeValue::Host => { + if let Some(host) = pending.host.clone() { + state.turn_approval_cache.hosts.insert(host); + } + } + ApprovalScopeValue::Tool => { + state + .turn_approval_cache + .tools + .insert(pending.tool_name.clone()); + } + ApprovalScopeValue::CommandPrefix => { + if let Some(command_prefix) = pending.command_prefix.clone() { + state + .session_approval_cache + .command_prefixes + .insert(command_prefix); + } + } + } +} diff --git a/crates/server/src/runtime/session_actor/mod.rs b/crates/server/src/runtime/session_actor/mod.rs new file mode 100644 index 00000000..1526a4bb --- /dev/null +++ b/crates/server/src/runtime/session_actor/mod.rs @@ -0,0 +1,11 @@ +mod commands; +mod handle; +mod loop_; +pub(crate) mod registry; +pub(crate) mod snapshots; +pub(crate) mod state; +mod turn; +mod turn_inline; + +pub(crate) use handle::SessionHandle; +pub(crate) use state::{SessionActorState, SpawnSnapshot}; diff --git a/crates/server/src/runtime/session_actor/registry.rs b/crates/server/src/runtime/session_actor/registry.rs new file mode 100644 index 00000000..9c43d8a5 --- /dev/null +++ b/crates/server/src/runtime/session_actor/registry.rs @@ -0,0 +1,199 @@ +use std::sync::Arc; + +use devo_core::TurnId; +use devo_protocol::SessionId; + +use super::state::SpawnSnapshot; +use crate::runtime::ServerRuntime; + +impl ServerRuntime { + pub(crate) fn runtime_arc(&self) -> Arc { + self.self_weak + .upgrade() + .expect("server runtime weak reference should remain alive") + } + + pub(crate) async fn session(&self, session_id: SessionId) -> Option { + self.sessions.lock().await.get(&session_id).cloned() + } + + pub(crate) async fn insert_session_actor( + &self, + state: super::SessionActorState, + ) -> super::SessionHandle { + let runtime = self.runtime_arc(); + let session_id = state.session_id(); + let handle = super::SessionHandle::spawn(session_id, state, runtime); + self.sessions + .lock() + .await + .insert(session_id, handle.clone()); + handle + } + + pub(crate) async fn remove_session_actor( + &self, + session_id: SessionId, + ) -> Option { + let handle = self.sessions.lock().await.remove(&session_id)?; + self.session_interactive.clear_session(session_id).await; + self.active_spawn_snapshots.lock().await.remove(&session_id); + Some(handle) + } + + pub(crate) async fn list_session_handles(&self) -> Vec { + self.sessions.lock().await.values().cloned().collect() + } + + pub(crate) async fn list_session_summaries_from_actors( + &self, + ) -> Vec { + let handles = self.list_session_handles().await; + let mut summaries = Vec::with_capacity(handles.len()); + for handle in handles { + if let Some(summary) = handle.summary().await { + summaries.push(summary); + } + } + summaries.sort_by(|left, right| { + right + .last_activity_at + .cmp(&left.last_activity_at) + .then_with(|| right.updated_at.cmp(&left.updated_at)) + }); + summaries + } + + pub(crate) async fn register_turn_spawn_snapshot( + &self, + session_id: SessionId, + turn_id: TurnId, + snapshot: Arc, + ) { + self.active_spawn_snapshots + .lock() + .await + .entry(session_id) + .or_default() + .insert(turn_id, snapshot); + } + + pub(crate) async fn clear_turn_spawn_snapshot(&self, session_id: SessionId, turn_id: TurnId) { + let mut snapshots = self.active_spawn_snapshots.lock().await; + if let Some(turns) = snapshots.get_mut(&session_id) { + turns.remove(&turn_id); + if turns.is_empty() { + snapshots.remove(&session_id); + } + } + } + + /// Snapshot registered at turn start while the session actor is busy executing. + pub(crate) async fn active_spawn_snapshot_for_session( + &self, + session_id: SessionId, + ) -> Option { + let snapshots = self.active_spawn_snapshots.lock().await; + let turns = snapshots.get(&session_id)?; + turns.values().next().map(|snapshot| (**snapshot).clone()) + } + + pub(crate) async fn register_active_stream( + &self, + session_id: SessionId, + stream: Arc>, + ) { + self.active_stream_states + .lock() + .await + .insert(session_id, stream); + } + + pub(crate) async fn unregister_active_stream(&self, session_id: SessionId) { + self.active_stream_states.lock().await.remove(&session_id); + } + + pub(crate) async fn active_stream_state( + &self, + session_id: SessionId, + ) -> Option>> { + self.active_stream_states + .lock() + .await + .get(&session_id) + .cloned() + } + + pub(crate) async fn session_record_snapshot( + &self, + session_id: SessionId, + ) -> Option { + if let Some(stream) = self.active_stream_state(session_id).await { + let stream = stream.lock().await; + if let Some(inline) = stream.turn_inline.as_ref() { + return inline.record.clone(); + } + } + let handle = self.session(session_id).await?; + handle.record().await.flatten() + } + + pub(crate) async fn session_summary_snapshot( + &self, + session_id: SessionId, + ) -> Option { + if let Some(stream) = self.active_stream_state(session_id).await { + let stream = stream.lock().await; + if let Some(inline) = stream.turn_inline.as_ref() { + return Some(inline.summary.clone()); + } + } + let handle = self.session(session_id).await?; + handle.summary().await + } + + /// Reads the session's collaboration mode, preferring the in-flight turn's + /// inline snapshot over a mailbox round-trip. + /// + /// Callers invoked from within `session_id`'s own actor turn (e.g. while + /// finalizing that turn) must go through this path: the actor's mailbox + /// is not polled again until the turn finishes, so asking its own + /// `SessionHandle` for a reply here would deadlock. + pub(crate) async fn session_collaboration_mode( + &self, + session_id: SessionId, + ) -> Option { + if let Some(stream) = self.active_stream_state(session_id).await { + let stream = stream.lock().await; + if let Some(inline) = stream.turn_inline.as_ref() { + return Some(inline.collaboration_mode); + } + } + let handle = self.session(session_id).await?; + handle.collaboration_mode().await + } + + /// Reads a session's parent id, preferring the in-flight turn inline snapshot + /// and agent-registry hierarchy over a mailbox round-trip. + pub(crate) async fn session_parent_id_snapshot( + &self, + session_id: SessionId, + ) -> Option> { + if let Some(stream) = self.active_stream_state(session_id).await { + let stream = stream.lock().await; + if let Some(inline) = stream.turn_inline.as_ref() { + return Some(inline.summary.parent_session_id); + } + } + { + let registries = self.agent_registries.lock().await; + for registry in registries.values() { + if let Some(parent_id) = registry.child_to_parent.get(&session_id).copied() { + return Some(Some(parent_id)); + } + } + } + let handle = self.session(session_id).await?; + handle.parent_session_id().await + } +} diff --git a/crates/server/src/runtime/session_actor/snapshots.rs b/crates/server/src/runtime/session_actor/snapshots.rs new file mode 100644 index 00000000..db35ed5b --- /dev/null +++ b/crates/server/src/runtime/session_actor/snapshots.rs @@ -0,0 +1,110 @@ +use std::collections::VecDeque; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; + +use devo_core::SessionConfig; +use devo_core::SessionContext; +use devo_core::SessionRecord; +use devo_core::TurnContext; +use devo_core::TurnKind; +use devo_protocol::PendingInputItem; +use devo_protocol::SessionId; +use devo_safety::PermissionMode; +use devo_safety::RuntimePermissionProfile; + +use crate::session::SessionMetadata; +use crate::session_context::SessionRuntimeContext; +use crate::turn::TurnMetadata; +use devo_core::tools::ToolRegistry; + +/// Snapshot used when reserving or queueing a turn on a session actor. +#[derive(Clone)] +pub(crate) struct TurnReservationSnapshot { + pub(crate) max_turns: Option, + pub(crate) active_turn: Option, + pub(crate) latest_turn: Option, + pub(crate) pending_turn_queue: Arc>>, + pub(crate) ephemeral: bool, + pub(crate) parent_session_id: Option, + pub(crate) summary: SessionMetadata, + pub(crate) runtime_context: Arc, +} + +/// Hook runner inputs derived from session actor state. +#[derive(Clone)] +pub(crate) struct HookContextSnapshot { + pub(crate) runtime_context: Arc, + pub(crate) record: Option, + pub(crate) summary: SessionMetadata, + pub(crate) config: SessionConfig, +} + +/// Fields needed to persist a turn line to rollout storage. +#[derive(Clone)] +pub(crate) struct TurnPersistenceSnapshot { + pub(crate) record: Option, + pub(crate) session_context: Option, + pub(crate) latest_turn_context: Option, +} + +/// Permission and tool context for shell-command turns. +#[derive(Clone)] +pub(crate) struct ShellExecContextSnapshot { + pub(crate) permission_mode: PermissionMode, + pub(crate) permission_profile: RuntimePermissionProfile, + pub(crate) runtime_context: Arc, + pub(crate) tool_registry: Arc, +} + +/// Context for async title generation. +#[derive(Clone)] +pub(crate) struct TitleGenerationContext { + pub(crate) model_selection: Option, + pub(crate) reasoning_effort_selection: Option, + pub(crate) title_state: devo_core::SessionTitleState, + pub(crate) runtime_context: Arc, +} + +/// Pending turn queue broadcast snapshot. +#[derive(Clone, Default)] +pub(crate) struct PendingQueueSnapshot { + pub(crate) pending_count: usize, + pub(crate) pending_texts: Vec, +} + +/// Fields returned by session/resume without locking the actor mailbox. +#[derive(Clone)] +pub(crate) struct SessionResumeSnapshot { + pub(crate) summary: SessionMetadata, + pub(crate) latest_turn: Option, + pub(crate) loaded_item_count: u64, + pub(crate) history_items: Vec, + pub(crate) pending_texts: Vec, +} + +/// Popped queued turn input for follow-up execution. +#[derive(Clone)] +pub(crate) struct QueuedTurnInputData { + pub(crate) display_input: String, + pub(crate) input_text: String, + pub(crate) input_messages: Vec, + pub(crate) collaboration_mode: devo_protocol::CollaborationMode, + pub(crate) model_selection: Option, + pub(crate) subagent_usage_owner: Option<(SessionId, Option)>, +} + +/// Turn kind and durable record before persisting an item. +#[derive(Clone)] +pub(crate) struct PersistItemPrep { + pub(crate) turn_kind: TurnKind, + pub(crate) record: Option, +} + +/// Deferred streaming items captured during graceful shutdown. +#[derive(Clone, Default)] +pub(crate) struct ShutdownDeferredSnapshot { + pub(crate) deferred_assistant: Option<(devo_core::ItemId, u64, String)>, + pub(crate) deferred_reasoning: Option<(devo_core::ItemId, u64, String)>, + pub(crate) active_turn_id: Option, + pub(crate) record: Option, +} diff --git a/crates/server/src/runtime/session_actor/state.rs b/crates/server/src/runtime/session_actor/state.rs new file mode 100644 index 00000000..21f9db92 --- /dev/null +++ b/crates/server/src/runtime/session_actor/state.rs @@ -0,0 +1,204 @@ +use std::collections::VecDeque; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; + +use devo_core::SessionConfig; +use devo_core::SessionRecord; +use devo_core::SessionState; +use devo_core::TurnId; +use devo_core::tools::ToolRegistry; +use devo_protocol::SessionId; + +use crate::execution::PersistedTurnItem; +use crate::runtime::RuntimeSession; +use crate::session::SessionHistoryItem; +use crate::session::SessionMetadata; +use crate::session_context::SessionRuntimeContext; +use crate::turn::TurnMetadata; + +use super::turn_inline::TurnInlineState; + +/// Immutable parent snapshot used by `spawn_agent` during an active parent turn. +#[derive(Clone)] +pub(crate) struct SpawnSnapshot { + pub(crate) parent_summary: SessionMetadata, + pub(crate) parent_config: SessionConfig, + pub(crate) stable_items: Vec, + pub(crate) parent_latest_turn: Option, + pub(crate) parent_active_turn_id: Option, + pub(crate) parent_tool_registry: Option>, + pub(crate) runtime_context: Arc, +} + +/// Approval caches cloned at turn start for permission checks while the actor +/// is busy executing a turn. +#[derive(Clone, Default)] +pub(crate) struct ApprovalCacheSnapshot { + pub(crate) session_approval_cache: crate::execution::ApprovalGrantCache, + pub(crate) turn_approval_cache: crate::execution::ApprovalGrantCache, +} + +#[derive(Clone, Default)] +pub(crate) struct DeferredItems { + pub(crate) assistant: Option<(devo_core::ItemId, u64, String)>, + pub(crate) reasoning: Option<(devo_core::ItemId, u64, String)>, +} + +use tokio::sync::Mutex as TokioMutex; + +/// Streaming-era mutable fields touched by the turn event bridge. +#[derive(Default)] +pub(crate) struct SessionStreamState { + pub(crate) deferred_assistant: Option<(devo_core::ItemId, u64, String)>, + pub(crate) deferred_reasoning: Option<(devo_core::ItemId, u64, String)>, + pub(crate) turn_inline: Option, +} + +impl SessionStreamState { + pub(crate) fn take_deferred_items(&mut self) -> DeferredItems { + DeferredItems { + assistant: self.deferred_assistant.take(), + reasoning: self.deferred_reasoning.take(), + } + } +} + +/// Per-session state owned exclusively by a `SessionActor` task. +pub(crate) struct SessionActorState { + pub(crate) runtime_context: Arc, + pub(crate) record: Option, + pub(crate) summary: SessionMetadata, + pub(crate) config: SessionConfig, + pub(crate) core: SessionState, + pub(crate) stream: Arc>, + pub(crate) active_turn: Option, + pub(crate) latest_turn: Option, + pub(crate) loaded_item_count: u64, + pub(crate) history_items: Vec, + pub(crate) persisted_turn_items: Vec, + pub(crate) latest_compaction_snapshot: Option, + pub(crate) pending_turn_queue: Arc>>, + pub(crate) btw_input_queue: Arc>>, + pub(crate) agent_tool_policy: devo_protocol::AgentToolPolicy, + pub(crate) max_turns: Option, + pub(crate) next_item_seq: u64, + pub(crate) first_user_input: Option, + pub(crate) tool_registry: Option>, + pub(crate) session_approval_cache: crate::execution::ApprovalGrantCache, + pub(crate) turn_approval_cache: crate::execution::ApprovalGrantCache, +} + +impl SessionActorState { + pub(crate) fn session_id(&self) -> SessionId { + self.summary.session_id + } + + pub(crate) fn parent_session_id(&self) -> Option { + self.summary.parent_session_id + } + + pub(crate) fn approval_cache_snapshot(&self) -> ApprovalCacheSnapshot { + ApprovalCacheSnapshot { + session_approval_cache: self.session_approval_cache.clone(), + turn_approval_cache: self.turn_approval_cache.clone(), + } + } + + pub(crate) fn spawn_snapshot(&self) -> SpawnSnapshot { + let fork_turns_all = true; + let stable_items = if fork_turns_all { + let active_turn_id = self.active_turn.as_ref().map(|turn| turn.turn_id); + self.persisted_turn_items + .iter() + .filter(|item| active_turn_id.is_none_or(|turn_id| item.turn_id != turn_id)) + .cloned() + .collect() + } else { + Vec::new() + }; + SpawnSnapshot { + parent_summary: self.summary.clone(), + parent_config: self.config.clone(), + stable_items, + parent_latest_turn: self.latest_turn.clone(), + parent_active_turn_id: self + .active_turn + .as_ref() + .map(|turn| turn.turn_id) + .or_else(|| self.latest_turn.as_ref().map(|turn| turn.turn_id)), + parent_tool_registry: self.tool_registry.clone(), + runtime_context: Arc::clone(&self.runtime_context), + } + } + + pub(crate) fn from_runtime_session(session: RuntimeSession) -> Self { + let core = Arc::try_unwrap(session.core_session) + .unwrap_or_else(|_| { + panic!("session core_session should have a single owner when starting actor") + }) + .into_inner(); + Self { + runtime_context: session.runtime_context, + record: session.record, + summary: session.summary, + config: session.config, + core, + stream: Arc::new(TokioMutex::new(SessionStreamState { + deferred_assistant: session.deferred_assistant, + deferred_reasoning: session.deferred_reasoning, + turn_inline: None, + })), + active_turn: session.active_turn, + latest_turn: session.latest_turn, + loaded_item_count: session.loaded_item_count, + history_items: session.history_items, + persisted_turn_items: session.persisted_turn_items, + latest_compaction_snapshot: session.latest_compaction_snapshot, + pending_turn_queue: session.pending_turn_queue, + btw_input_queue: session.btw_input_queue, + agent_tool_policy: session.agent_tool_policy, + max_turns: session.max_turns, + next_item_seq: session.next_item_seq, + first_user_input: session.first_user_input, + tool_registry: session.tool_registry, + session_approval_cache: session.session_approval_cache, + turn_approval_cache: session.turn_approval_cache, + } + } + + pub(crate) fn to_runtime_session_from_stream( + &self, + stream: &SessionStreamState, + ) -> RuntimeSession { + RuntimeSession { + runtime_context: Arc::clone(&self.runtime_context), + record: self.record.clone(), + summary: self.summary.clone(), + config: self.config.clone(), + core_session: Arc::new(tokio::sync::Mutex::new(self.core.snapshot_for_export())), + active_turn: self.active_turn.clone(), + latest_turn: self.latest_turn.clone(), + loaded_item_count: self.loaded_item_count, + history_items: self.history_items.clone(), + persisted_turn_items: self.persisted_turn_items.clone(), + latest_compaction_snapshot: self.latest_compaction_snapshot.clone(), + pending_turn_queue: Arc::clone(&self.pending_turn_queue), + btw_input_queue: Arc::clone(&self.btw_input_queue), + agent_tool_policy: self.agent_tool_policy, + max_turns: self.max_turns, + deferred_assistant: stream.deferred_assistant.clone(), + deferred_reasoning: stream.deferred_reasoning.clone(), + next_item_seq: self.next_item_seq, + first_user_input: self.first_user_input.clone(), + tool_registry: self.tool_registry.clone(), + session_approval_cache: self.session_approval_cache.clone(), + turn_approval_cache: self.turn_approval_cache.clone(), + } + } +} + +impl From for SessionActorState { + fn from(session: RuntimeSession) -> Self { + Self::from_runtime_session(session) + } +} diff --git a/crates/server/src/runtime/session_actor/turn.rs b/crates/server/src/runtime/session_actor/turn.rs new file mode 100644 index 00000000..19612502 --- /dev/null +++ b/crates/server/src/runtime/session_actor/turn.rs @@ -0,0 +1,107 @@ +use std::sync::Arc; + +use tokio::sync::mpsc; + +use crate::runtime::ServerRuntime; +use crate::runtime::session_actor::state::SessionActorState; +use crate::runtime::turn_exec::{ + ExecuteTurnRequest, FinalizeTurnParams, QUERY_EVENT_CHANNEL_CAPACITY, TurnModelQueryParams, + spawn_turn_event_stream, +}; + +pub(super) async fn execute_turn_in_actor( + state: &mut SessionActorState, + runtime: Arc, + request: ExecuteTurnRequest, +) { + let ExecuteTurnRequest { + session_id, + turn, + turn_config, + display_input, + input, + input_messages, + collaboration_mode, + input_mode, + } = request; + + let spawn_snapshot = Arc::new(state.spawn_snapshot()); + runtime + .register_turn_spawn_snapshot(session_id, turn.turn_id, Arc::clone(&spawn_snapshot)) + .await; + + { + let mut stream = state.stream.lock().await; + stream.turn_inline = Some(super::turn_inline::TurnInlineState::new(state, &turn)); + } + runtime + .register_active_stream(session_id, Arc::clone(&state.stream)) + .await; + + runtime + .prepare_turn_execution_for_actor( + state, + &turn, + &display_input, + input_mode.emits_user_message(), + ) + .await; + + let (event_tx, event_rx) = mpsc::channel(QUERY_EVENT_CHANNEL_CAPACITY); + let event_tool_registry = runtime.tool_registry_for_actor_state(state); + let usage_parent_session_id = state.parent_session_id(); + let usage_context_window = Some(turn_config.model.context_window as u64); + runtime + .begin_parent_usage_turn(session_id, turn.turn_id, usage_context_window) + .await; + + let stream = Arc::clone(&state.stream); + let event_task = spawn_turn_event_stream( + Arc::clone(&runtime), + stream, + session_id, + turn.clone(), + collaboration_mode, + event_tool_registry, + usage_parent_session_id, + usage_context_window, + event_rx, + ); + + let query_outcome = runtime + .run_turn_model_query(TurnModelQueryParams { + state, + turn_id: turn.turn_id, + turn_config: &turn_config, + input: &input, + input_messages: &input_messages, + collaboration_mode, + input_mode, + usage_parent_session_id, + event_tx, + }) + .await; + let event_summary = event_task.await.ok(); + + let turn_id = turn.turn_id; + runtime + .finalize_executed_turn(FinalizeTurnParams { + state, + session_id, + turn, + query_outcome, + event_summary, + usage_parent_session_id, + }) + .await; + + runtime.clear_turn_spawn_snapshot(session_id, turn_id).await; + runtime.unregister_active_stream(session_id).await; + let inline = { + let mut stream = state.stream.lock().await; + stream.turn_inline.take() + }; + if let Some(inline) = inline { + inline.merge_into(state); + } +} diff --git a/crates/server/src/runtime/session_actor/turn_inline.rs b/crates/server/src/runtime/session_actor/turn_inline.rs new file mode 100644 index 00000000..ac2ce8f2 --- /dev/null +++ b/crates/server/src/runtime/session_actor/turn_inline.rs @@ -0,0 +1,71 @@ +use std::sync::Arc; + +use devo_core::SessionRecord; +use devo_core::TurnId; +use devo_core::TurnKind; +use devo_protocol::CollaborationMode; + +use crate::execution::ApprovalGrantCache; +use crate::execution::PersistedTurnItem; +use crate::session::SessionHistoryItem; +use crate::session::SessionMetadata; +use crate::turn::TurnMetadata; + +use super::SessionActorState; +use super::snapshots::HookContextSnapshot; + +/// Mutable session fields updated during an in-actor turn without mailbox round-trips. +pub(crate) struct TurnInlineState { + pub(crate) turn_id: TurnId, + pub(crate) turn_kind: TurnKind, + pub(crate) next_item_seq: u64, + pub(crate) loaded_item_count: u64, + pub(crate) persisted_turn_items: Vec, + pub(crate) history_items: Vec, + pub(crate) record: Option, + pub(crate) session_approval_cache: ApprovalGrantCache, + pub(crate) turn_approval_cache: ApprovalGrantCache, + pub(crate) summary: SessionMetadata, + pub(crate) collaboration_mode: CollaborationMode, + pub(crate) hook_context: HookContextSnapshot, +} + +impl TurnInlineState { + pub(crate) fn new(state: &SessionActorState, turn: &TurnMetadata) -> Self { + Self { + turn_id: turn.turn_id, + turn_kind: turn.kind.clone(), + next_item_seq: state.next_item_seq, + loaded_item_count: state.loaded_item_count, + persisted_turn_items: Vec::new(), + history_items: Vec::new(), + record: state.record.clone(), + session_approval_cache: state.session_approval_cache.clone(), + turn_approval_cache: state.turn_approval_cache.clone(), + summary: state.summary.clone(), + collaboration_mode: state.core.collaboration_mode, + hook_context: HookContextSnapshot { + runtime_context: Arc::clone(&state.runtime_context), + record: state.record.clone(), + summary: state.summary.clone(), + config: state.config.clone(), + }, + } + } + + pub(crate) fn allocate_item_seq(&mut self) -> u64 { + let item_seq = self.next_item_seq; + self.next_item_seq = self.next_item_seq.saturating_add(1); + self.loaded_item_count = self.loaded_item_count.saturating_add(1); + item_seq + } + + pub(crate) fn merge_into(self, state: &mut SessionActorState) { + state.next_item_seq = self.next_item_seq; + state.loaded_item_count = self.loaded_item_count; + state.persisted_turn_items.extend(self.persisted_turn_items); + state.history_items.extend(self.history_items); + state.session_approval_cache = self.session_approval_cache; + state.turn_approval_cache = self.turn_approval_cache; + } +} diff --git a/crates/server/src/runtime/session_interactive.rs b/crates/server/src/runtime/session_interactive.rs new file mode 100644 index 00000000..3cf151c5 --- /dev/null +++ b/crates/server/src/runtime/session_interactive.rs @@ -0,0 +1,147 @@ +use std::collections::HashMap; + +use devo_core::TurnId; +use devo_protocol::ApprovalDecisionValue; +use devo_protocol::SessionId; +use tokio::sync::Mutex; +use tokio::sync::oneshot; + +use crate::execution::PendingApproval; +use crate::execution::PendingUserInput; + +#[derive(Default)] +struct SessionInteractiveState { + pending_approvals: HashMap, + pending_user_inputs: HashMap, +} + +/// Global interactive wait lanes keyed by session id. +/// +/// Approval and user-input waits are routed here so a session actor blocked in +/// `query()` never has to process mailbox messages for client responses. +#[derive(Default)] +pub(crate) struct SessionInteractiveLanes { + inner: Mutex>, +} + +impl SessionInteractiveLanes { + pub(crate) async fn register_pending_approval( + &self, + host_session_id: SessionId, + approval_id: String, + pending: PendingApproval, + ) { + self.inner + .lock() + .await + .entry(host_session_id) + .or_default() + .pending_approvals + .insert(approval_id, pending); + } + + pub(crate) async fn remove_pending_approval( + &self, + host_session_id: SessionId, + approval_id: &str, + ) -> Option { + let mut lanes = self.inner.lock().await; + let state = lanes.get_mut(&host_session_id)?; + let removed = state.pending_approvals.remove(approval_id); + if state.pending_approvals.is_empty() && state.pending_user_inputs.is_empty() { + lanes.remove(&host_session_id); + } + removed + } + + pub(crate) async fn register_pending_user_input( + &self, + session_id: SessionId, + request_id: String, + pending: PendingUserInput, + ) { + self.inner + .lock() + .await + .entry(session_id) + .or_default() + .pending_user_inputs + .insert(request_id, pending); + } + + pub(crate) async fn take_pending_user_input( + &self, + session_id: SessionId, + request_id: &str, + expected_turn_id: TurnId, + ) -> Result { + let mut lanes = self.inner.lock().await; + let Some(state) = lanes.get_mut(&session_id) else { + return Err(UserInputTakeError::NotFound); + }; + let Some(pending) = state.pending_user_inputs.remove(request_id) else { + return Err(UserInputTakeError::NotFound); + }; + if pending.turn_id != expected_turn_id { + state + .pending_user_inputs + .insert(request_id.to_string(), pending); + return Err(UserInputTakeError::WrongTurn); + } + if state.pending_approvals.is_empty() && state.pending_user_inputs.is_empty() { + lanes.remove(&session_id); + } + Ok(pending) + } + + pub(crate) async fn has_pending_interactive(&self, session_id: SessionId) -> bool { + self.inner + .lock() + .await + .get(&session_id) + .is_some_and(|state| { + !state.pending_approvals.is_empty() || !state.pending_user_inputs.is_empty() + }) + } + + pub(crate) async fn clear_pending_user_inputs_for_turn( + &self, + session_id: SessionId, + turn_id: TurnId, + ) -> usize { + let mut lanes = self.inner.lock().await; + let Some(state) = lanes.get_mut(&session_id) else { + return 0; + }; + let previous_len = state.pending_user_inputs.len(); + state + .pending_user_inputs + .retain(|_, pending| pending.turn_id != turn_id); + let removed_len = previous_len.saturating_sub(state.pending_user_inputs.len()); + if state.pending_approvals.is_empty() && state.pending_user_inputs.is_empty() { + lanes.remove(&session_id); + } + removed_len + } + + pub(crate) async fn clear_session(&self, session_id: SessionId) { + self.inner.lock().await.remove(&session_id); + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum UserInputTakeError { + NotFound, + WrongTurn, +} + +pub(crate) async fn complete_approval_wait( + rx: oneshot::Receiver, +) -> Result { + match rx.await { + Ok(ApprovalDecisionValue::Approve) => Ok(ApprovalDecisionValue::Approve), + Ok(ApprovalDecisionValue::Deny) => Err("rejected by user".to_string()), + Ok(ApprovalDecisionValue::Cancel) => Err("cancelled by user".to_string()), + Err(_) => Err("approval channel closed".to_string()), + } +} diff --git a/crates/server/src/runtime/subagent_usage.rs b/crates/server/src/runtime/subagent_usage.rs index f00e48ce..bab2097f 100644 --- a/crates/server/src/runtime/subagent_usage.rs +++ b/crates/server/src/runtime/subagent_usage.rs @@ -82,7 +82,7 @@ impl UsageTotals { } #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) struct ParentUsageSnapshot { +pub(crate) struct ParentUsageSnapshot { pub(super) session_id: SessionId, pub(super) turn_id: TurnId, pub(super) turn_usage: UsageTotals, @@ -279,13 +279,11 @@ impl ServerRuntime { turn_id: TurnId, context_window: Option, ) { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { - return; - }; - let base_session_totals = { - let session = session_arc.lock().await; - UsageTotals::from_session_summary(&session.summary) - }; + let base_session_totals = self + .session_summary_snapshot(session_id) + .await + .map(|summary| UsageTotals::from_session_summary(&summary)) + .unwrap_or_default(); self.subagent_usage.lock().await.begin_parent_turn( session_id, turn_id, @@ -308,13 +306,8 @@ impl ServerRuntime { } pub(super) async fn active_turn_id_for_session(&self, session_id: SessionId) -> Option { - let session_arc = self.sessions.lock().await.get(&session_id).cloned()?; - session_arc - .lock() - .await - .active_turn - .as_ref() - .map(|turn| turn.turn_id) + let session_handle = self.session(session_id).await?; + session_handle.active_turn_id().await.flatten() } pub(super) async fn publish_parent_turn_usage( @@ -369,18 +362,8 @@ impl ServerRuntime { } async fn apply_parent_usage_snapshot(&self, snapshot: ParentUsageSnapshot) { - // Bind the session Arc before locking it so the global `sessions` guard is - // released first. Holding `sessions` across the per-session `.await` would - // pin the process-wide lock and can deadlock every other request. - let session_arc = self - .sessions - .lock() - .await - .get(&snapshot.session_id) - .cloned(); - if let Some(session_arc) = session_arc { - let mut session = session_arc.lock().await; - snapshot.apply_to_session(&mut session); + if let Some(session_handle) = self.session(snapshot.session_id).await { + session_handle.apply_parent_usage_snapshot(snapshot).await; } self.broadcast_event(ServerEvent::TurnUsageUpdated( snapshot.to_turn_usage_updated_payload(), @@ -404,27 +387,26 @@ impl ParentUsageSnapshot { } } - fn apply_to_session(self, session: &mut RuntimeSession) { - session.summary.total_input_tokens = self.session_totals.input_tokens; - session.summary.total_output_tokens = self.session_totals.output_tokens; - session.summary.total_tokens = self.session_totals.total_tokens; - session.summary.total_cache_creation_tokens = - self.session_totals.cache_creation_input_tokens; - session.summary.total_cache_read_tokens = self.session_totals.cache_read_input_tokens; - session.summary.last_query_total_tokens = self.turn_usage.total_tokens; - if let Some(active_turn) = session.active_turn.as_mut() + pub(super) fn apply_to_actor_state( + self, + state: &mut crate::runtime::session_actor::state::SessionActorState, + ) { + state.summary.total_input_tokens = self.session_totals.input_tokens; + state.summary.total_output_tokens = self.session_totals.output_tokens; + state.summary.total_tokens = self.session_totals.total_tokens; + state.summary.total_cache_creation_tokens = self.session_totals.cache_creation_input_tokens; + state.summary.total_cache_read_tokens = self.session_totals.cache_read_input_tokens; + state.summary.last_query_total_tokens = self.turn_usage.total_tokens; + if let Some(active_turn) = state.active_turn.as_mut() && active_turn.turn_id == self.turn_id { active_turn.usage = Some(self.turn_usage.to_turn_usage()); } - if let Ok(mut core_session) = session.core_session.try_lock() { - core_session.total_input_tokens = self.session_totals.input_tokens; - core_session.total_output_tokens = self.session_totals.output_tokens; - core_session.total_tokens = self.session_totals.total_tokens; - core_session.total_cache_creation_tokens = - self.session_totals.cache_creation_input_tokens; - core_session.total_cache_read_tokens = self.session_totals.cache_read_input_tokens; - } + state.core.total_input_tokens = self.session_totals.input_tokens; + state.core.total_output_tokens = self.session_totals.output_tokens; + state.core.total_tokens = self.session_totals.total_tokens; + state.core.total_cache_creation_tokens = self.session_totals.cache_creation_input_tokens; + state.core.total_cache_read_tokens = self.session_totals.cache_read_input_tokens; } } diff --git a/crates/server/src/runtime/turn_exec.rs b/crates/server/src/runtime/turn_exec.rs deleted file mode 100644 index c0341d7e..00000000 --- a/crates/server/src/runtime/turn_exec.rs +++ /dev/null @@ -1,2985 +0,0 @@ -use std::collections::HashMap; -use std::sync::Arc; -use std::sync::OnceLock; -use std::time::Instant; - -use super::proposed_plan::{ProposedPlanParser, ProposedPlanSegment}; -use super::*; -use crate::{FileChangePayload, TurnPlanStepPayload, TurnPlanUpdatedPayload}; -use devo_core::tools::tool_spec::ToolPreparationFeedback; -use devo_util_git::extract_paths_from_patch; -use tokio::sync::mpsc; - -const QUERY_EVENT_CHANNEL_CAPACITY: usize = 1024; - -struct PendingToolCall { - item_id: Option, - item_seq: Option, - input: serde_json::Value, - display_kind: ToolDisplayKind, - command: String, -} - -struct TurnEventStreamSummary { - latest_usage: Option, - stop_reason: Option, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum ToolDisplayKind { - CommandExecution, - Generic, -} - -impl ToolDisplayKind { - fn for_tool_name(name: &str) -> Self { - if is_unified_exec_tool(name) { - Self::CommandExecution - } else { - Self::Generic - } - } - - fn is_command_execution(self) -> bool { - self == Self::CommandExecution - } -} - -struct ToolStartItem { - item_kind: ItemKind, - payload: serde_json::Value, -} - -fn turn_failure_reason_from_error( - error: &devo_core::AgentError, -) -> Option { - match error { - devo_core::AgentError::MaxTurnsExceeded(_) => { - Some(devo_protocol::TurnFailureReason::MaxTurnRequests) - } - devo_core::AgentError::Provider(_) - | devo_core::AgentError::ContextTooLong - | devo_core::AgentError::Aborted => None, - } -} - -/// Inputs captured at turn-start time and handed to the background turn executor. -pub(super) struct ExecuteTurnRequest { - /// Runtime session that owns the turn and receives emitted items, usage, and status updates. - pub(super) session_id: SessionId, - /// Pre-created turn metadata persisted at turn start; execution mutates a local copy to its - /// terminal status before appending the final turn record. - pub(super) turn: TurnMetadata, - /// Resolved model, provider, reasoning, tool, web, and token-budget settings for this turn. - pub(super) turn_config: TurnConfig, - /// User-facing rendering of the submitted input. Visible turns persist this as the displayed - /// user message; hidden continuation turns keep it out of the transcript. - pub(super) display_input: String, - /// Canonical resolved prompt text. Visible turns push this as the user-role message when the - /// input resolver did not return structured `input_messages`. - pub(super) input: String, - /// Structured user-role messages produced by input resolution, such as expanded skill content. - /// When non-empty, these are pushed instead of the single `input` string. - pub(super) input_messages: Vec, - /// Collaboration mode to install on the core session for this query; it also drives - /// mode-specific stream handling such as proposed-plan parsing. - pub(super) collaboration_mode: devo_protocol::CollaborationMode, - /// Controls whether this executor emits/pushes a visible user message or runs hidden work such - /// as goal continuation, and carries the hidden goal context when needed. - pub(super) input_mode: TurnInputMode, -} - -fn stream_trace_elapsed_ms() -> u128 { - static STREAM_TRACE_START: OnceLock = OnceLock::new(); - STREAM_TRACE_START - .get_or_init(Instant::now) - .elapsed() - .as_millis() -} - -fn query_event_trace_kind(event: &QueryEvent) -> &'static str { - match event { - QueryEvent::TextDelta(_) => "text_delta", - QueryEvent::ReasoningDelta(_) => "reasoning_delta", - QueryEvent::ReasoningCompleted => "reasoning_completed", - QueryEvent::UsageDelta { .. } => "usage_delta", - QueryEvent::ToolUseStart { .. } => "tool_use_start", - QueryEvent::ToolExecutionStart { .. } => "tool_execution_start", - QueryEvent::ToolProgress { .. } => "tool_progress", - QueryEvent::ToolResult { .. } => "tool_result", - QueryEvent::TurnComplete { .. } => "turn_complete", - QueryEvent::Usage { .. } => "usage", - } -} - -fn query_event_trace_delta_len(event: &QueryEvent) -> usize { - match event { - QueryEvent::TextDelta(text) | QueryEvent::ReasoningDelta(text) => text.len(), - QueryEvent::ToolProgress { progress, .. } => match progress { - devo_core::tools::ToolProgress::OutputDelta { delta } => delta.len(), - devo_core::tools::ToolProgress::StatusUpdate { message, .. } => message.len(), - devo_core::tools::ToolProgress::Terminal { terminal_id } => terminal_id.len(), - devo_core::tools::ToolProgress::Completion { summary } => summary.len(), - }, - QueryEvent::ReasoningCompleted - | QueryEvent::UsageDelta { .. } - | QueryEvent::ToolUseStart { .. } - | QueryEvent::ToolExecutionStart { .. } - | QueryEvent::ToolResult { .. } - | QueryEvent::TurnComplete { .. } - | QueryEvent::Usage { .. } => 0, - } -} - -fn query_event_trace_token_preview(event: &QueryEvent) -> Option { - match event { - QueryEvent::TextDelta(text) => assistant_token_log_preview(text), - QueryEvent::ReasoningDelta(_) - | QueryEvent::ReasoningCompleted - | QueryEvent::UsageDelta { .. } - | QueryEvent::ToolUseStart { .. } - | QueryEvent::ToolExecutionStart { .. } - | QueryEvent::ToolProgress { .. } - | QueryEvent::ToolResult { .. } - | QueryEvent::TurnComplete { .. } - | QueryEvent::Usage { .. } => None, - } -} - -fn assistant_token_log_preview(text: &str) -> Option { - assistant_token_logging_enabled() - .then(|| format_assistant_token_log_preview(text, assistant_token_log_max_chars())) -} - -fn assistant_token_logging_enabled() -> bool { - static ASSISTANT_TOKEN_LOGGING_ENABLED: OnceLock = OnceLock::new(); - *ASSISTANT_TOKEN_LOGGING_ENABLED.get_or_init(|| { - std::env::var("DEVO_LOG_ASSISTANT_TOKEN_TEXT") - .ok() - .is_some_and(|value| { - matches!( - value.as_str(), - "1" | "true" | "TRUE" | "yes" | "YES" | "on" | "ON" - ) - }) - }) -} - -fn assistant_token_log_max_chars() -> usize { - static ASSISTANT_TOKEN_LOG_MAX_CHARS: OnceLock = OnceLock::new(); - *ASSISTANT_TOKEN_LOG_MAX_CHARS.get_or_init(|| { - std::env::var("DEVO_ASSISTANT_TOKEN_LOG_MAX_CHARS") - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|value| *value > 0) - .unwrap_or(512) - }) -} - -fn format_assistant_token_log_preview(text: &str, max_chars: usize) -> String { - let max_chars = max_chars.max(1); - let mut preview = String::new(); - let mut chars = text.chars(); - for ch in chars.by_ref().take(max_chars) { - preview.extend(ch.escape_default()); - } - if chars.next().is_some() { - preview.push_str("..."); - } - preview -} - -async fn complete_reasoning_item( - runtime: &Arc, - session_id: SessionId, - turn_id: TurnId, - item_id: ItemId, - item_seq: u64, - text: String, -) { - runtime - .complete_item( - session_id, - turn_id, - item_id, - item_seq, - ItemKind::Reasoning, - TurnItem::Reasoning(TextItem { text: text.clone() }), - serde_json::json!({ "title": "Reasoning", "text": text }), - ) - .await; -} - -async fn complete_assistant_item( - runtime: &Arc, - session_id: SessionId, - turn_id: TurnId, - item_id: ItemId, - item_seq: u64, - text: String, -) { - runtime - .complete_item( - session_id, - turn_id, - item_id, - item_seq, - ItemKind::AgentMessage, - TurnItem::AgentMessage(TextItem { text: text.clone() }), - serde_json::json!({ "title": "Assistant", "text": text }), - ) - .await; -} - -#[derive(Debug, Default)] -struct ProposedPlanStreamItem { - item_id: Option, - item_seq: Option, - text: String, -} - -impl ProposedPlanStreamItem { - async fn start( - &mut self, - runtime: &Arc, - session_id: SessionId, - turn_id: TurnId, - ) { - if self.item_id.is_some() && self.item_seq.is_some() { - return; - } - let (item_id, item_seq) = runtime - .start_item( - session_id, - turn_id, - ItemKind::Plan, - serde_json::json!({ "title": "Proposed Plan", "text": "" }), - ) - .await; - self.item_id = Some(item_id); - self.item_seq = Some(item_seq); - } - - async fn push_delta( - &mut self, - runtime: &Arc, - session_id: SessionId, - turn_id: TurnId, - delta: String, - ) { - if delta.is_empty() { - return; - } - self.start(runtime, session_id, turn_id).await; - self.text.push_str(&delta); - runtime - .broadcast_event(ServerEvent::ItemDelta { - delta_kind: ItemDeltaKind::PlanDelta, - payload: ItemDeltaPayload { - context: EventContext { - session_id, - turn_id: Some(turn_id), - item_id: self.item_id, - seq: 0, - }, - delta, - stream_index: None, - channel: None, - }, - }) - .await; - } - - async fn complete( - &mut self, - runtime: &Arc, - session_id: SessionId, - turn_id: TurnId, - ) { - let (Some(item_id), Some(item_seq)) = (self.item_id.take(), self.item_seq.take()) else { - return; - }; - let text = std::mem::take(&mut self.text); - runtime - .complete_item( - session_id, - turn_id, - item_id, - item_seq, - ItemKind::Plan, - TurnItem::Plan(TextItem { text: text.clone() }), - serde_json::json!({ "title": "Proposed Plan", "text": text }), - ) - .await; - } -} - -#[allow(clippy::too_many_arguments)] -async fn push_assistant_text_delta( - runtime: &Arc, - event_session_arc: &Arc>, - session_id: SessionId, - turn_id: TurnId, - assistant_item_id: &mut Option, - assistant_item_seq: &mut Option, - assistant_text: &mut String, - assistant_delta_seq: &mut u64, - text: String, -) { - if text.is_empty() { - return; - } - let (item_id, item_seq) = match (*assistant_item_id, *assistant_item_seq) { - (Some(item_id), Some(item_seq)) => (item_id, item_seq), - (None, None) => { - let (item_id, item_seq) = runtime - .start_item( - session_id, - turn_id, - ItemKind::AgentMessage, - serde_json::json!({ "title": "Assistant", "text": "" }), - ) - .await; - *assistant_item_id = Some(item_id); - *assistant_item_seq = Some(item_seq); - (item_id, item_seq) - } - _ => return, - }; - assistant_text.push_str(&text); - *assistant_delta_seq = (*assistant_delta_seq).saturating_add(1); - runtime - .broadcast_event(ServerEvent::ItemDelta { - delta_kind: ItemDeltaKind::AgentMessageDelta, - payload: ItemDeltaPayload { - context: EventContext { - session_id, - turn_id: Some(turn_id), - item_id: Some(item_id), - seq: 0, - }, - delta: text, - stream_index: None, - channel: None, - }, - }) - .await; - if let Ok(mut session) = event_session_arc.try_lock() { - session.deferred_assistant = Some((item_id, item_seq, assistant_text.clone())); - } -} - -#[allow(clippy::too_many_arguments)] -async fn handle_proposed_plan_segments( - runtime: &Arc, - event_session_arc: &Arc>, - session_id: SessionId, - turn_id: TurnId, - segments: Vec, - assistant_item_id: &mut Option, - assistant_item_seq: &mut Option, - assistant_text: &mut String, - assistant_delta_seq: &mut u64, - proposed_plan_item: &mut ProposedPlanStreamItem, - leading_normal_buffer: &mut String, -) { - for segment in segments { - match segment { - ProposedPlanSegment::Normal(delta) => { - if delta.is_empty() { - continue; - } - if assistant_item_id.is_none() && delta.chars().all(char::is_whitespace) { - leading_normal_buffer.push_str(&delta); - continue; - } - let delta = if assistant_item_id.is_none() && !leading_normal_buffer.is_empty() { - format!("{}{}", std::mem::take(leading_normal_buffer), delta) - } else { - delta - }; - push_assistant_text_delta( - runtime, - event_session_arc, - session_id, - turn_id, - assistant_item_id, - assistant_item_seq, - assistant_text, - assistant_delta_seq, - delta, - ) - .await; - } - ProposedPlanSegment::PlanStart => { - leading_normal_buffer.clear(); - proposed_plan_item.start(runtime, session_id, turn_id).await; - } - ProposedPlanSegment::PlanDelta(delta) => { - proposed_plan_item - .push_delta(runtime, session_id, turn_id, delta) - .await; - } - ProposedPlanSegment::PlanEnd => {} - } - } -} - -fn is_unified_exec_tool(name: &str) -> bool { - matches!(name, "exec_command" | "write_stdin") -} - -fn is_file_change_tool(name: &str) -> bool { - matches!(name, "apply_patch" | "write") -} - -fn is_plan_tool(name: &str) -> bool { - matches!(name, "update_plan") -} - -fn tool_start_item_kind( - tool_name: &str, - display_kind: ToolDisplayKind, - preparation_feedback: ToolPreparationFeedback, -) -> ItemKind { - if preparation_feedback == ToolPreparationFeedback::LiveOnly { - ItemKind::ToolCall - } else if is_file_change_tool(tool_name) { - ItemKind::FileChange - } else if display_kind.is_command_execution() { - ItemKind::CommandExecution - } else if is_plan_tool(tool_name) { - ItemKind::Plan - } else { - ItemKind::ToolCall - } -} - -fn tool_start_item( - tool_call_id: &str, - tool_name: &str, - command: &str, - input: &serde_json::Value, - display_kind: ToolDisplayKind, - preparation_feedback: ToolPreparationFeedback, - command_actions: Vec, -) -> ToolStartItem { - let item_kind = tool_start_item_kind(tool_name, display_kind, preparation_feedback); - let payload = match item_kind { - ItemKind::ToolCall => serde_json::to_value(ToolCallPayload { - tool_call_id: tool_call_id.to_string(), - tool_name: tool_name.to_string(), - parameters: input.clone(), - command_actions, - }) - .expect("serialize tool call payload"), - ItemKind::FileChange => serde_json::to_value(FileChangePayload { - tool_call_id: tool_call_id.to_string(), - tool_name: Some(tool_name.to_string()), - input: Some(input.clone()), - changes: Vec::new(), - is_error: false, - }) - .expect("serialize file change payload"), - ItemKind::CommandExecution => serde_json::to_value(CommandExecutionPayload { - tool_call_id: tool_call_id.to_string(), - tool_name: tool_name.to_string(), - command: command.to_string(), - input: Some(input.clone()), - source: devo_protocol::protocol::ExecCommandSource::Agent, - command_actions, - output: None, - is_error: false, - }) - .expect("serialize command execution payload"), - ItemKind::Plan => serde_json::json!({ - "title": "Plan", - "text": "" - }), - ItemKind::UserMessage - | ItemKind::AgentMessage - | ItemKind::Reasoning - | ItemKind::ToolResult - | ItemKind::McpToolCall - | ItemKind::WebSearch - | ItemKind::ImageView - | ItemKind::ContextCompaction - | ItemKind::ApprovalRequest - | ItemKind::ApprovalDecision - | ItemKind::ResearchArtifact => unreachable!("tool start item kind must be tool-like"), - }; - ToolStartItem { item_kind, payload } -} - -fn tool_start_item_from_input( - tool_call_id: &str, - tool_name: &str, - command: &str, - input: &serde_json::Value, - display_kind: ToolDisplayKind, - preparation_feedback: ToolPreparationFeedback, -) -> ToolStartItem { - tool_start_item( - tool_call_id, - tool_name, - command, - input, - display_kind, - preparation_feedback, - command_actions_from_tool_input(tool_name, command, input), - ) -} - -fn tool_start_item_from_result( - tool_call_id: &str, - tool_name: &str, - command: &str, - input: &serde_json::Value, - display_kind: ToolDisplayKind, - preparation_feedback: ToolPreparationFeedback, - summary: &str, -) -> ToolStartItem { - tool_start_item( - tool_call_id, - tool_name, - command, - input, - display_kind, - preparation_feedback, - command_actions_from_tool_result(tool_name, command, input, summary), - ) -} - -fn command_display_from_input(tool_name: &str, input: &serde_json::Value) -> String { - match tool_name { - "exec_command" => input - .get("cmd") - .or_else(|| input.get("command")) - .and_then(serde_json::Value::as_str) - .unwrap_or_default() - .to_string(), - "write_stdin" => { - let session_id = input - .get("session_id") - .and_then(serde_json::Value::as_i64) - .map(|id| id.to_string()) - .unwrap_or_else(|| "?".to_string()); - let chars = input - .get("chars") - .and_then(serde_json::Value::as_str) - .unwrap_or_default(); - if chars.is_empty() { - format!("poll session {session_id}") - } else { - format!("write_stdin session {session_id}") - } - } - "read" => { - let path = input - .get("filePath") - .or_else(|| input.get("path")) - .and_then(serde_json::Value::as_str) - .unwrap_or_default(); - format!("read {path}") - } - "find" | "glob" => { - let pattern = input - .get("pattern") - .and_then(serde_json::Value::as_str) - .unwrap_or_default(); - let path = input - .get("path") - .and_then(serde_json::Value::as_str) - .unwrap_or_default(); - let command_name = if tool_name == "find" { "find" } else { "glob" }; - if path.is_empty() { - format!("{command_name} {pattern}") - } else { - format!("{command_name} {pattern} in {path}") - } - } - "grep" => { - let pattern = input - .get("pattern") - .and_then(serde_json::Value::as_str) - .unwrap_or_default(); - let path = input - .get("path") - .and_then(serde_json::Value::as_str) - .unwrap_or_default(); - if path.is_empty() { - format!("grep {pattern}") - } else { - format!("grep {pattern} in {path}") - } - } - "code_search" => code_search_display_from_input(input), - _ => String::new(), - } -} - -fn command_actions_from_tool_input( - tool_name: &str, - command: &str, - input: &serde_json::Value, -) -> Vec { - match tool_name { - "read" => crate::tool_actions::read_action_from_tool_input(command, input) - .into_iter() - .collect(), - "find" | "glob" => vec![devo_protocol::parse_command::ParsedCommand::ListFiles { - cmd: command.to_string(), - path: find_display_from_input(input), - }], - "grep" => vec![devo_protocol::parse_command::ParsedCommand::Search { - cmd: command.to_string(), - query: input - .get("pattern") - .and_then(serde_json::Value::as_str) - .map(ToOwned::to_owned), - path: input - .get("path") - .and_then(serde_json::Value::as_str) - .map(ToOwned::to_owned), - }], - "code_search" => code_search_action_from_input(command, input) - .into_iter() - .collect(), - _ => Vec::new(), - } -} - -fn code_search_display_from_input(input: &serde_json::Value) -> String { - match input - .get("operation") - .and_then(serde_json::Value::as_str) - .unwrap_or("search") - { - "find_related" => { - let path = input - .get("file_path") - .and_then(serde_json::Value::as_str) - .unwrap_or_default(); - let line = input.get("line").and_then(serde_json::Value::as_u64); - match (path.is_empty(), line) { - (false, Some(line)) => format!("code_search related {path}:{line}"), - (false, None) => format!("code_search related {path}"), - (true, _) => "code_search related".to_string(), - } - } - _ => { - let query = input - .get("query") - .and_then(serde_json::Value::as_str) - .unwrap_or_default(); - let path = input - .get("path") - .and_then(serde_json::Value::as_str) - .unwrap_or_default(); - match (query.is_empty(), path.is_empty()) { - (false, false) => format!("code_search {query} in {path}"), - (false, true) => format!("code_search {query}"), - (true, false) => format!("code_search in {path}"), - (true, true) => "code_search".to_string(), - } - } - } -} - -fn code_search_action_from_input( - command: &str, - input: &serde_json::Value, -) -> Option { - match input - .get("operation") - .and_then(serde_json::Value::as_str) - .unwrap_or("search") - { - "find_related" => { - let path = input - .get("file_path") - .and_then(serde_json::Value::as_str) - .filter(|path| !path.is_empty())?; - let line = input - .get("line") - .and_then(serde_json::Value::as_u64) - .map(|line| line.to_string()) - .unwrap_or_else(|| "?".to_string()); - Some(devo_protocol::parse_command::ParsedCommand::Search { - cmd: command.to_string(), - query: Some(format!("related {path}:{line}")), - path: Some(path.to_string()), - }) - } - _ => { - let query = input - .get("query") - .and_then(serde_json::Value::as_str) - .filter(|query| !query.is_empty())?; - Some(devo_protocol::parse_command::ParsedCommand::Search { - cmd: command.to_string(), - query: Some(query.to_string()), - path: input - .get("path") - .and_then(serde_json::Value::as_str) - .map(ToOwned::to_owned), - }) - } - } -} - -fn find_display_from_input(input: &serde_json::Value) -> Option { - let pattern = input - .get("pattern") - .and_then(serde_json::Value::as_str) - .filter(|pattern| !pattern.is_empty())?; - let path = input.get("path").and_then(serde_json::Value::as_str); - Some(match path.filter(|path| !path.is_empty()) { - Some(path) => format!("{pattern} in {path}"), - None => pattern.to_string(), - }) -} - -fn command_actions_from_tool_result( - tool_name: &str, - command: &str, - input: &serde_json::Value, - summary: &str, -) -> Vec { - let actions = command_actions_from_tool_input(tool_name, command, input); - if !actions.is_empty() { - return actions; - } - match tool_name { - "read" => crate::tool_actions::read_action_from_tool_summary(summary) - .into_iter() - .collect(), - _ => actions, - } -} - -fn collaboration_mode_from_pending_metadata( - metadata: Option<&serde_json::Value>, -) -> devo_protocol::CollaborationMode { - metadata - .and_then(|metadata| { - metadata - .get("collaboration_mode") - .or_else(|| metadata.get("interaction_mode")) - }) - .cloned() - .and_then(|value| serde_json::from_value(value).ok()) - .unwrap_or_default() -} - -fn command_execution_item_id_for_progress( - pending_tool_calls: &HashMap, - tool_use_id: &str, -) -> Option { - pending_tool_calls - .get(tool_use_id) - .filter(|pending| pending.display_kind.is_command_execution()) - .and_then(|pending| pending.item_id) -} - -fn user_shell_exec_input(command: &str, cwd: std::path::PathBuf) -> serde_json::Value { - serde_json::json!({ - "cmd": command, - "workdir": cwd, - "login": true, - "tty": true, - }) -} - -fn user_shell_command_payload( - tool_call_id: &str, - command: &str, - input: serde_json::Value, - command_actions: Vec, - output: Option, - is_error: bool, -) -> CommandExecutionPayload { - CommandExecutionPayload { - tool_call_id: tool_call_id.to_string(), - tool_name: "exec_command".to_string(), - command: command.to_string(), - input: Some(input), - source: devo_protocol::protocol::ExecCommandSource::UserShell, - command_actions, - output, - is_error, - } -} - -const AGENT_COORDINATION_TOOL_NAMES: &[&str] = &[ - "spawn_agent", - "send_message", - "wait_agent", - "list_agents", - "close_agent", -]; - -fn without_agent_coordination_tools( - registry: &devo_core::tools::ToolRegistry, -) -> devo_core::tools::ToolRegistry { - let names = registry - .tool_definitions() - .into_iter() - .map(|tool| tool.name) - .filter(|name| { - !AGENT_COORDINATION_TOOL_NAMES - .iter() - .any(|hidden_name| *hidden_name == name) - }) - .collect::>(); - let names = names.iter().map(String::as_str).collect::>(); - registry.restricted_to_specs(&names) -} - -impl ServerRuntime { - fn tool_registry_for_session( - &self, - session: &RuntimeSession, - ) -> Arc { - session - .tool_registry - .clone() - .unwrap_or_else(|| Arc::clone(&session.runtime_context.registry)) - } - - pub(super) async fn execute_shell_command_turn( - self: Arc, - session_id: SessionId, - turn: TurnMetadata, - command: String, - cwd: std::path::PathBuf, - ) { - self.capture_turn_workspace_baseline(session_id, turn.turn_id, cwd.clone()) - .await; - let session_arc = self.sessions.lock().await.get(&session_id).cloned(); - if let Some(session_arc) = session_arc { - session_arc.lock().await.turn_approval_cache = - crate::execution::ApprovalGrantCache::default(); - } - - let tool_call_id = format!("user-shell-{}", turn.turn_id); - let input = user_shell_exec_input(&command, cwd.clone()); - let command_actions = parse_command(std::slice::from_ref(&command)); - let (item_id, item_seq) = self - .start_item( - session_id, - turn.turn_id, - ItemKind::CommandExecution, - serde_json::to_value(user_shell_command_payload( - &tool_call_id, - &command, - input.clone(), - command_actions.clone(), - None, - false, - )) - .expect("serialize command execution payload"), - ) - .await; - - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { - self.clear_active_turn_runtime_handles(session_id).await; - return; - }; - let (permission_mode, permission_profile, registry, network_proxy, network_no_proxy) = { - let session = session_arc.lock().await; - let registry = self.tool_registry_for_session(&session); - let core_session = session.core_session.lock().await; - let provider_http = session - .runtime_context - .config_store - .lock() - .expect("app config store mutex should not be poisoned") - .effective_config() - .provider_http - .clone(); - ( - core_session.config.permission_mode, - core_session.config.permission_profile.clone(), - registry, - provider_http.proxy_url, - provider_http.no_proxy, - ) - }; - let turn_cancel_token = self - .active_turn_cancellations - .lock() - .await - .get(&session_id) - .cloned() - .unwrap_or_else(CancellationToken::new); - let tool_execution_start_runtime = Arc::clone(&self); - let tool_execution_start_session_id = session_id; - let tool_execution_start_turn_id = turn.turn_id; - let runtime = ToolRuntime::new_with_context_and_options( - registry, - self.build_permission_checker( - session_id, - turn.turn_id, - permission_mode, - permission_profile, - ), - ToolRuntimeContext { - session_id: session_id.to_string(), - turn_id: Some(turn.turn_id.to_string()), - cwd, - agent_scope: ToolAgentScope::Parent, - agent_context_mode: devo_protocol::AgentContextMode::CodingAgent, - collaboration_mode: devo_protocol::CollaborationMode::Build, - agent_coordinator: None, - client_filesystem: None, - client_terminal: None, - local_web_search: None, - hooks: self.hook_context_for_session(session_id).await, - network_proxy, - network_no_proxy, - }, - ToolExecutionOptions { - cancel_token: turn_cancel_token, - on_tool_execution_start: Some(Arc::new(move |call: ToolCall| { - let runtime = Arc::clone(&tool_execution_start_runtime); - let tool_call_id = call.id; - // `on_tool_execution_start` is stored as a trait-object callback that returns - // `BoxFuture<'static, ()>`. `Box::pin` gives this async block the erased future - // shape while still letting the body await `broadcast_event`. - Box::pin(async move { - runtime - .broadcast_event(ServerEvent::ToolCallStatusUpdated( - devo_protocol::ToolCallStatusUpdatedPayload { - session_id: tool_execution_start_session_id, - turn_id: tool_execution_start_turn_id, - tool_call_id, - status: "in_progress".to_string(), - terminal_id: None, - }, - )) - .await; - }) - })), - ..ToolExecutionOptions::default() - }, - ); - let result = runtime - .execute_batch(&[ToolCall { - id: tool_call_id.clone(), - name: "exec_command".to_string(), - input: input.clone(), - }]) - .await - .into_iter() - .next() - .unwrap_or_else(|| ToolCallResult::error(&tool_call_id, "shell command did not run")); - let output = match result.content.clone() { - ToolContent::Text(text) => serde_json::Value::String(text), - ToolContent::Json(json) => json, - ToolContent::Mixed { text, json } => { - json.unwrap_or_else(|| serde_json::Value::String(text.unwrap_or_default())) - } - }; - let is_error = result.is_error; - self.clear_active_turn_runtime_handles(session_id).await; - self.complete_item( - session_id, - turn.turn_id, - item_id, - item_seq, - ItemKind::CommandExecution, - TurnItem::CommandExecution(CommandExecutionItem { - tool_call_id: tool_call_id.clone(), - tool_name: "exec_command".to_string(), - command: command.clone(), - input: input.clone(), - output: output.clone(), - is_error, - }), - serde_json::to_value(user_shell_command_payload( - &tool_call_id, - &command, - input.clone(), - command_actions, - Some(output), - is_error, - )) - .expect("serialize command execution payload"), - ) - .await; - - let final_turn = { - let mut session = session_arc.lock().await; - let mut final_turn = turn.clone(); - final_turn.completed_at = Some(Utc::now()); - final_turn.status = if is_error { - TurnStatus::Failed - } else { - TurnStatus::Completed - }; - session.latest_turn = Some(final_turn.clone()); - session.active_turn = None; - session.summary.status = SessionRuntimeStatus::Idle; - session.summary.updated_at = Utc::now(); - session.summary.last_activity_at = session.summary.updated_at; - final_turn - }; - let (record, session_context, turn_context) = { - let session = session_arc.lock().await; - let core_session = session.core_session.lock().await; - ( - session.record.clone(), - core_session.session_context.clone(), - core_session.latest_turn_context.clone(), - ) - }; - if let Some(record) = record - && let Err(error) = self.rollout_store.append_turn( - &record, - build_turn_record(&final_turn, session_context, turn_context), - ) - { - tracing::warn!(session_id = %session_id, error = %error, "failed to persist shell command turn line"); - } - self.finalize_turn_workspace_changes(session_id, &final_turn) - .await; - if is_error { - self.broadcast_event(ServerEvent::TurnFailed(TurnEventPayload { - session_id, - turn: final_turn.clone(), - })) - .await; - } - self.broadcast_event(ServerEvent::TurnCompleted(TurnEventPayload { - session_id, - turn: final_turn, - })) - .await; - self.broadcast_event(ServerEvent::SessionStatusChanged( - SessionStatusChangedPayload { - session_id, - status: SessionRuntimeStatus::Idle, - }, - )) - .await; - } - - /// Execute one turn end-to-end, including streaming query events, - /// persisting turn state, and draining queued follow-up inputs. - pub(super) async fn execute_turn(self: Arc, request: ExecuteTurnRequest) { - let ExecuteTurnRequest { - session_id, - turn, - turn_config, - display_input, - input, - input_messages, - collaboration_mode, - input_mode, - } = request; - let baseline_cwd = { - let session_arc = self.sessions.lock().await.get(&session_id).cloned(); - match session_arc { - Some(session_arc) => Some(session_arc.lock().await.summary.cwd.clone()), - None => None, - } - }; - if let Some(cwd) = baseline_cwd { - self.capture_turn_workspace_baseline(session_id, turn.turn_id, cwd) - .await; - } - let session_arc = self.sessions.lock().await.get(&session_id).cloned(); - if let Some(session_arc) = session_arc { - session_arc.lock().await.turn_approval_cache = - crate::execution::ApprovalGrantCache::default(); - } - if input_mode.emits_user_message() { - // Record the user's message immediately so the UI can show it even if - // the model call or event stream takes a moment to start. - self.emit_turn_item( - session_id, - turn.turn_id, - ItemKind::UserMessage, - TurnItem::UserMessage(TextItem { - // TODO: Note that future additions of multimodal support - // will require modifications here. - text: display_input.clone(), - }), - // Keep the legacy text payload shape for live item/completed clients; durable - // user-message semantics come from the TurnItem above. - serde_json::json!({ "title": "You", "text": display_input.clone() }), - ) - .await; - } - - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { - return; - }; - let (event_tx, mut event_rx) = mpsc::channel::(QUERY_EVENT_CHANNEL_CAPACITY); - let runtime = Arc::clone(&self); - let turn_for_events = turn.clone(); - let turn_for_plan_updates = turn.clone(); - let event_session_arc = Arc::clone(&session_arc); - let (event_tool_registry, usage_parent_session_id) = { - let session = session_arc.lock().await; - ( - self.tool_registry_for_session(&session), - session.summary.parent_session_id, - ) - }; - let usage_context_window = Some(turn_config.model.context_window as u64); - self.begin_parent_usage_turn(session_id, turn.turn_id, usage_context_window) - .await; - let event_task = tokio::spawn(async move { - // This task owns the streamed model output. It turns raw query - // callbacks into persisted turn items and keeps enough state to - // resume cleanly if the turn is interrupted mid-stream. - let mut assistant_item_id = None; - let mut assistant_item_seq = None; - let mut assistant_delta_seq = 0_u64; - let mut assistant_text = String::new(); - let mut reasoning_item_id = None; - let mut reasoning_item_seq = None; - let mut reasoning_text = String::new(); - let mut tool_names_by_id = HashMap::new(); - let mut pending_tool_calls: HashMap = HashMap::new(); - let mut proposed_plan_parser = (collaboration_mode - == devo_protocol::CollaborationMode::Plan) - .then(ProposedPlanParser::default); - let mut proposed_plan_item = ProposedPlanStreamItem::default(); - let mut proposed_plan_leading_normal = String::new(); - let mut latest_usage: Option = None; - let mut stop_reason: Option = None; - while let Some(event) = event_rx.recv().await { - let assistant_token_text = query_event_trace_token_preview(&event); - if let Some(assistant_token_text) = assistant_token_text.as_deref() { - tracing::debug!( - stream_elapsed_ms = stream_trace_elapsed_ms(), - event_kind = query_event_trace_kind(&event), - delta_len = query_event_trace_delta_len(&event), - assistant_token_text, - "query event bridge dequeued by turn event task" - ); - } else { - tracing::debug!( - stream_elapsed_ms = stream_trace_elapsed_ms(), - event_kind = query_event_trace_kind(&event), - delta_len = query_event_trace_delta_len(&event), - "query event bridge dequeued by turn event task" - ); - } - match event { - QueryEvent::TextDelta(text) => { - if let Some(parser) = proposed_plan_parser.as_mut() { - let segments = parser.push_str(&text); - handle_proposed_plan_segments( - &runtime, - &event_session_arc, - session_id, - turn_for_events.turn_id, - segments, - &mut assistant_item_id, - &mut assistant_item_seq, - &mut assistant_text, - &mut assistant_delta_seq, - &mut proposed_plan_item, - &mut proposed_plan_leading_normal, - ) - .await; - } else { - push_assistant_text_delta( - &runtime, - &event_session_arc, - session_id, - turn_for_events.turn_id, - &mut assistant_item_id, - &mut assistant_item_seq, - &mut assistant_text, - &mut assistant_delta_seq, - text, - ) - .await; - } - } - QueryEvent::ReasoningDelta(text) => { - let (item_id, item_seq) = match (reasoning_item_id, reasoning_item_seq) { - (Some(item_id), Some(item_seq)) => (item_id, item_seq), - (None, None) => { - let (item_id, item_seq) = runtime - .start_item( - session_id, - turn_for_events.turn_id, - ItemKind::Reasoning, - serde_json::json!({ "title": "Reasoning", "text": "" }), - ) - .await; - reasoning_item_id = Some(item_id); - reasoning_item_seq = Some(item_seq); - (item_id, item_seq) - } - _ => continue, - }; - reasoning_text.push_str(&text); - runtime - .broadcast_event(ServerEvent::ItemDelta { - delta_kind: ItemDeltaKind::ReasoningTextDelta, - payload: ItemDeltaPayload { - context: EventContext { - session_id, - turn_id: Some(turn_for_events.turn_id), - item_id: Some(item_id), - seq: 0, - }, - delta: text, - stream_index: None, - channel: None, - }, - }) - .await; - let _ = item_seq; - - // Store deferred completion info for interrupt recovery - if let Ok(mut session) = event_session_arc.try_lock() { - session.deferred_reasoning = - Some((item_id, item_seq, reasoning_text.clone())); - } - } - QueryEvent::ReasoningCompleted => { - if let (Some(item_id), Some(item_seq)) = - (reasoning_item_id.take(), reasoning_item_seq.take()) - { - if let Ok(mut session) = event_session_arc.try_lock() { - session.deferred_reasoning.take(); - } - complete_reasoning_item( - &runtime, - session_id, - turn_for_events.turn_id, - item_id, - item_seq, - reasoning_text.clone(), - ) - .await; - reasoning_text.clear(); - } - } - QueryEvent::ToolUseStart { id, name, input } => { - tool_names_by_id.insert(id.clone(), name.clone()); - if let (Some(item_id), Some(item_seq)) = - (reasoning_item_id.take(), reasoning_item_seq.take()) - { - complete_reasoning_item( - &runtime, - session_id, - turn_for_events.turn_id, - item_id, - item_seq, - reasoning_text.clone(), - ) - .await; - reasoning_text.clear(); - } - if let (Some(item_id), Some(item_seq)) = - (assistant_item_id.take(), assistant_item_seq.take()) - { - complete_assistant_item( - &runtime, - session_id, - turn_for_events.turn_id, - item_id, - item_seq, - assistant_text.clone(), - ) - .await; - assistant_text.clear(); - } - let display_kind = ToolDisplayKind::for_tool_name(&name); - let command = command_display_from_input(&name, &input); - let preparation_feedback = event_tool_registry.preparation_feedback(&name); - let start_item = tool_start_item_from_input( - &id, - &name, - &command, - &input, - display_kind, - preparation_feedback, - ); - let (item_id, item_seq) = runtime - .start_item( - session_id, - turn_for_events.turn_id, - start_item.item_kind, - start_item.payload, - ) - .await; - pending_tool_calls.insert( - id, - PendingToolCall { - item_id: Some(item_id), - item_seq: Some(item_seq), - input, - display_kind, - command, - }, - ); - } - QueryEvent::ToolExecutionStart { id } => { - runtime - .broadcast_event(ServerEvent::ToolCallStatusUpdated( - devo_protocol::ToolCallStatusUpdatedPayload { - session_id, - turn_id: turn_for_events.turn_id, - tool_call_id: id, - status: "in_progress".to_string(), - terminal_id: None, - }, - )) - .await; - } - QueryEvent::ToolResult { - tool_use_id, - tool_name: final_tool_name, - input: final_input, - content, - display_content, - is_error, - summary, - } => { - let tool_name = if final_tool_name.is_empty() { - tool_names_by_id.get(&tool_use_id).cloned() - } else { - Some(final_tool_name) - }; - let mut result_input = - (!final_input.is_null()).then(|| final_input.clone()); - // First complete the pending ToolCall item so its item/completed - // arrives before the ToolResult item/completed. - if let Some(mut pending) = pending_tool_calls.remove(&tool_use_id) { - if !final_input.is_null() { - pending.command = tool_name - .as_deref() - .map(|tool_name| { - command_display_from_input(tool_name, &final_input) - }) - .unwrap_or_default(); - pending.input = final_input; - } - result_input = Some(pending.input.clone()); - if (pending.item_id.is_none() || pending.item_seq.is_none()) - && let Some(tool_name) = tool_name.clone() - { - let preparation_feedback = - event_tool_registry.preparation_feedback(&tool_name); - let start_item = tool_start_item_from_result( - &tool_use_id, - &tool_name, - &pending.command, - &pending.input, - pending.display_kind, - preparation_feedback, - &summary, - ); - let (item_id, item_seq) = runtime - .start_item( - session_id, - turn_for_events.turn_id, - start_item.item_kind, - start_item.payload, - ) - .await; - pending.item_id = Some(item_id); - pending.item_seq = Some(item_seq); - } - - let pending_item_id = pending.item_id.expect("pending item id"); - let pending_item_seq = pending.item_seq.expect("pending item seq"); - if let Some(tool_name) = tool_name.clone() - && is_plan_tool(&tool_name) - { - let output_json = match content.clone() { - devo_core::tools::ToolContent::Text(text) => { - serde_json::Value::String(text) - } - devo_core::tools::ToolContent::Json(json) => json, - devo_core::tools::ToolContent::Mixed { text, json } => json - .unwrap_or_else(|| { - serde_json::Value::String(text.unwrap_or_default()) - }), - }; - let explanation = output_json - .get("explanation") - .and_then(serde_json::Value::as_str) - .map(ToOwned::to_owned); - let plan = output_json - .get("plan") - .and_then(serde_json::Value::as_array) - .cloned() - .unwrap_or_default(); - - runtime - .complete_item( - session_id, - turn_for_events.turn_id, - pending_item_id, - pending_item_seq, - ItemKind::Plan, - TurnItem::Plan(TextItem { - text: output_json.to_string(), - }), - serde_json::json!({ - "title": "Plan", - "text": output_json.to_string(), - }), - ) - .await; - - runtime - .broadcast_event(ServerEvent::TurnPlanUpdated( - TurnPlanUpdatedPayload { - session_id, - turn: turn_for_plan_updates.clone(), - explanation, - plan: plan - .into_iter() - .filter_map(|item| { - Some(TurnPlanStepPayload { - step: item - .get("step")? - .as_str()? - .to_string(), - status: item - .get("status")? - .as_str()? - .to_string(), - }) - }) - .collect(), - }, - )) - .await; - continue; - } - - if let Some(tool_name) = tool_name.clone() - && is_file_change_tool(&tool_name) - { - let output_json = match content.clone() { - devo_core::tools::ToolContent::Text(text) => { - serde_json::Value::String(text) - } - devo_core::tools::ToolContent::Json(json) => json, - devo_core::tools::ToolContent::Mixed { text, json } => json - .unwrap_or_else(|| { - serde_json::Value::String(text.unwrap_or_default()) - }), - }; - let changes = output_json - .get("files") - .and_then(serde_json::Value::as_array) - .cloned() - .unwrap_or_default() - .into_iter() - .filter_map(|file| { - let path = - std::path::PathBuf::from(file.get("path")?.as_str()?); - let kind = file.get("kind")?.as_str()?; - let additions = file - .get("additions") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0); - let deletions = file - .get("deletions") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0); - let change = match kind { - "add" => devo_protocol::protocol::FileChange::Add { - content: file - .get("content") - .and_then(serde_json::Value::as_str) - .map(ToOwned::to_owned) - .unwrap_or_else(|| { - "\n".repeat(additions as usize) - }), - }, - "delete" => { - devo_protocol::protocol::FileChange::Delete { - content: file - .get("content") - .and_then(serde_json::Value::as_str) - .map(ToOwned::to_owned) - .unwrap_or_else(|| { - "\n".repeat(deletions as usize) - }), - } - } - "update" | "move" => { - devo_protocol::protocol::FileChange::Update { - unified_diff: file - .get("diff") - .or_else(|| file.get("patch")) - .or_else(|| output_json.get("diff")) - .and_then(serde_json::Value::as_str) - .unwrap_or("") - .to_string(), - old_text: file - .get("oldContent") - .or_else(|| file.get("preContent")) - .or_else(|| file.get("pre_content")) - .and_then(serde_json::Value::as_str) - .map(ToOwned::to_owned), - new_text: file - .get("postContent") - .or_else(|| file.get("post_content")) - .or_else(|| file.get("content")) - .and_then(serde_json::Value::as_str) - .map(ToOwned::to_owned), - move_path: file - .get("movePath") - .or_else(|| file.get("move_path")) - .and_then(serde_json::Value::as_str) - .map(std::path::PathBuf::from), - } - } - _ => return None, - }; - Some((path, change)) - }) - .collect::>(); - let changes = if changes.is_empty() { - output_json - .get("diff") - .and_then(serde_json::Value::as_str) - .map(extract_paths_from_patch) - .unwrap_or_default() - .into_iter() - .map(|path| { - ( - std::path::PathBuf::from(path), - devo_protocol::protocol::FileChange::Update { - unified_diff: output_json - .get("diff") - .and_then(serde_json::Value::as_str) - .unwrap_or("") - .to_string(), - old_text: None, - new_text: None, - move_path: None, - }, - ) - }) - .collect::>() - } else { - changes - }; - - runtime - .complete_item( - session_id, - turn_for_events.turn_id, - pending_item_id, - pending_item_seq, - ItemKind::FileChange, - TurnItem::ToolResult(ToolResultItem { - tool_call_id: tool_use_id.clone(), - tool_name: Some(tool_name.clone()), - output: output_json.clone(), - display_content: display_content.clone(), - is_error, - }), - serde_json::to_value(FileChangePayload { - tool_call_id: tool_use_id.clone(), - tool_name: Some(tool_name.clone()), - input: Some(pending.input.clone()), - changes, - is_error, - }) - .expect("serialize file change payload"), - ) - .await; - continue; - } - - if pending.display_kind.is_command_execution() { - let tool_name = tool_name.clone().unwrap_or_default(); - let output = match content.clone() { - devo_core::tools::ToolContent::Text(text) => { - serde_json::Value::String(text) - } - devo_core::tools::ToolContent::Json(json) => json, - devo_core::tools::ToolContent::Mixed { text, json } => json - .unwrap_or_else(|| { - serde_json::Value::String(text.unwrap_or_default()) - }), - }; - let completed_payload = - serde_json::to_value(CommandExecutionPayload { - tool_call_id: tool_use_id.clone(), - tool_name: tool_name.clone(), - command: pending.command.clone(), - input: Some(pending.input.clone()), - source: devo_protocol::protocol::ExecCommandSource::Agent, - command_actions: command_actions_from_tool_result( - &tool_name, - &pending.command, - &pending.input, - &summary, - ), - output: Some(output.clone()), - is_error, - }) - .expect("serialize command execution payload"); - runtime - .complete_item( - session_id, - turn_for_events.turn_id, - pending_item_id, - pending_item_seq, - ItemKind::CommandExecution, - TurnItem::CommandExecution(CommandExecutionItem { - tool_call_id: tool_use_id.clone(), - tool_name, - command: pending.command, - input: pending.input, - output, - is_error, - }), - completed_payload, - ) - .await; - continue; - } - let completed_payload = serde_json::to_value(ToolCallPayload { - tool_call_id: tool_use_id.clone(), - tool_name: tool_name.clone().unwrap_or_default(), - parameters: pending.input.clone(), - command_actions: command_actions_from_tool_result( - tool_name.clone().unwrap_or_default().as_str(), - &pending.command, - &pending.input, - &summary, - ), - }) - .expect("serialize tool call payload"); - runtime - .complete_item( - session_id, - turn_for_events.turn_id, - pending_item_id, - pending_item_seq, - ItemKind::ToolCall, - TurnItem::ToolCall(ToolCallItem { - tool_call_id: tool_use_id.clone(), - tool_name: tool_name.clone().unwrap_or_default(), - input: pending.input, - }), - completed_payload, - ) - .await; - } - runtime - .emit_turn_item( - session_id, - turn_for_events.turn_id, - ItemKind::ToolResult, - TurnItem::ToolResult(ToolResultItem { - tool_call_id: tool_use_id.clone(), - tool_name: tool_name.clone(), - output: match content.clone() { - devo_core::tools::ToolContent::Text(text) => { - serde_json::Value::String(text) - } - devo_core::tools::ToolContent::Json(json) => json, - devo_core::tools::ToolContent::Mixed { text, json } => json - .unwrap_or_else(|| { - serde_json::Value::String(text.unwrap_or_default()) - }), - }, - display_content: display_content.clone(), - is_error, - }), - serde_json::to_value(ToolResultPayload { - tool_call_id: tool_use_id.clone(), - tool_name, - input: result_input, - content: match content { - devo_core::tools::ToolContent::Text(text) => { - serde_json::Value::String(text) - } - devo_core::tools::ToolContent::Json(json) => json, - devo_core::tools::ToolContent::Mixed { text, json } => json - .unwrap_or_else(|| { - serde_json::Value::String(text.unwrap_or_default()) - }), - }, - display_content, - is_error, - summary, - }) - .expect("serialize tool result payload"), - ) - .await; - } - QueryEvent::ToolProgress { - tool_use_id, - progress, - } => { - let content = match progress { - devo_core::tools::ToolProgress::OutputDelta { delta } => Some(delta), - devo_core::tools::ToolProgress::StatusUpdate { message, percent } => { - Some(match percent { - Some(percent) => format!("{message} ({percent}%)"), - None => message, - }) - } - devo_core::tools::ToolProgress::Completion { summary } => Some(summary), - devo_core::tools::ToolProgress::Terminal { terminal_id } => { - runtime - .broadcast_event(ServerEvent::ToolCallStatusUpdated( - devo_protocol::ToolCallStatusUpdatedPayload { - session_id, - turn_id: turn_for_events.turn_id, - tool_call_id: tool_use_id.clone(), - status: "in_progress".to_string(), - terminal_id: Some(terminal_id), - }, - )) - .await; - None - } - }; - let Some(content) = content else { - continue; - }; - let item_id = command_execution_item_id_for_progress( - &pending_tool_calls, - &tool_use_id, - ); - let _ = runtime - .broadcast_event(ServerEvent::ItemDelta { - delta_kind: ItemDeltaKind::CommandExecutionOutputDelta, - payload: ItemDeltaPayload { - context: EventContext { - session_id, - turn_id: Some(turn_for_events.turn_id), - item_id, - seq: 0, - }, - delta: serde_json::json!({ - "tool_use_id": tool_use_id, - "text": content, - }) - .to_string(), - stream_index: None, - channel: None, - }, - }) - .await; - } - QueryEvent::UsageDelta { usage } | QueryEvent::Usage { usage } => { - let turn_usage = TurnUsage::from_usage(&usage); - latest_usage = Some(turn_usage.clone()); - if usage_parent_session_id.is_some() { - let _ = runtime - .publish_subagent_turn_usage( - session_id, - turn_for_events.turn_id, - turn_usage, - ) - .await; - } else if let Some(snapshot) = runtime - .publish_parent_turn_usage( - session_id, - turn_for_events.turn_id, - turn_usage, - usage_context_window, - ) - .await - { - latest_usage = Some(snapshot.turn_usage.to_turn_usage()); - } - } - QueryEvent::TurnComplete { - stop_reason: terminal_stop_reason, - } => { - stop_reason = Some(terminal_stop_reason); - } - } - } - if let Some(parser) = proposed_plan_parser.as_mut() { - let segments = parser.finish(); - handle_proposed_plan_segments( - &runtime, - &event_session_arc, - session_id, - turn_for_events.turn_id, - segments, - &mut assistant_item_id, - &mut assistant_item_seq, - &mut assistant_text, - &mut assistant_delta_seq, - &mut proposed_plan_item, - &mut proposed_plan_leading_normal, - ) - .await; - proposed_plan_item - .complete(&runtime, session_id, turn_for_events.turn_id) - .await; - } - // Complete any deferred items that the interrupt handler didn't already take. - // handle_interrupt takes deferred_assistant/deferred_reasoning from the session - // and completes them; if they're already None we must skip to avoid persisting duplicates. - if let Some((item_id, item_seq, text)) = { - let mut session = event_session_arc.lock().await; - session.deferred_reasoning.take() - } { - complete_reasoning_item( - &runtime, - session_id, - turn_for_events.turn_id, - item_id, - item_seq, - text, - ) - .await; - } - if let Some((item_id, item_seq, text)) = { - let mut session = event_session_arc.lock().await; - session.deferred_assistant.take() - } { - complete_assistant_item( - &runtime, - session_id, - turn_for_events.turn_id, - item_id, - item_seq, - text, - ) - .await; - } - tracing::debug!( - session_id = %session_id, - turn_id = %turn_for_events.turn_id, - "query event stream drained" - ); - TurnEventStreamSummary { - latest_usage, - stop_reason, - } - }); - - let ( - result, - mut session_total_input_tokens, - mut session_total_output_tokens, - mut session_total_tokens, - mut session_total_cache_creation_tokens, - mut session_total_cache_read_tokens, - session_last_input_tokens, - session_prompt_token_estimate, - ) = { - // Run the model query only after the event pipeline is ready so - // streamed deltas can be consumed and persisted immediately. - let ( - core_session, - agent_scope, - agent_tool_policy, - session_tool_registry, - runtime_context, - ) = { - let session = session_arc.lock().await; - let agent_scope = if session.summary.parent_session_id.is_some() { - ToolAgentScope::Subagent - } else { - ToolAgentScope::Parent - }; - let registry = self.tool_registry_for_session(&session); - ( - Arc::clone(&session.core_session), - agent_scope, - session.agent_tool_policy, - registry, - Arc::clone(&session.runtime_context), - ) - }; - let turn_goal = match &input_mode { - TurnInputMode::VisibleUserMessage => { - let stores = self.goal_stores.lock().await; - stores - .get(&session_id) - .and_then(GoalStore::get) - .map(crate::goal::Goal::to_thread_goal) - } - TurnInputMode::HiddenGoalContinuation { goal } => Some(goal.clone()), - }; - let mut core_session = core_session.lock().await; - core_session.config.token_budget = turn_config.token_budget(); - core_session.collaboration_mode = collaboration_mode; - if let Some(goal) = turn_goal { - core_session.set_active_goal(goal); - } else { - core_session.clear_active_goal(); - } - if input_mode.emits_user_message() && input_messages.is_empty() { - core_session.push_message(Message::user(input.clone())); - } else if input_mode.emits_user_message() { - for input_message in &input_messages { - core_session.push_message(Message::user(input_message.clone())); - } - } - let event_callback_tx = event_tx.clone(); - let callback: devo_core::EventCallback = - std::sync::Arc::new(move |event: QueryEvent| { - let event_callback_tx = event_callback_tx.clone(); - // `EventCallback` returns `BoxFuture` so the core query loop can await an - // async event sink without knowing this closure's concrete future type. Boxing - // also lets the bounded channel send apply backpressure at the callback boundary. - Box::pin(async move { - let _ = event_callback_tx.send(event).await; - }) - }); - let tool_execution_start_tx = event_tx.clone(); - let agent_context_mode = core_session - .session_context - .as_ref() - .map(|context| match context.system_prompt_mode { - devo_core::SystemPromptMode::CodingAgent => { - devo_protocol::AgentContextMode::CodingAgent - } - devo_core::SystemPromptMode::DeepResearch => { - devo_protocol::AgentContextMode::DeepResearch - } - }) - .unwrap_or_default(); - let registry = match agent_tool_policy { - devo_protocol::AgentToolPolicy::Inherit if usage_parent_session_id.is_some() => { - Arc::new(without_agent_coordination_tools(&session_tool_registry)) - } - devo_protocol::AgentToolPolicy::Inherit => session_tool_registry, - devo_protocol::AgentToolPolicy::DenyAll => { - Arc::new(devo_core::tools::ToolRegistry::new()) - } - devo_protocol::AgentToolPolicy::DeepResearch => Arc::new( - runtime_context - .registry - .restricted_to_specs(research::RESEARCH_WORKER_TOOL_NAMES), - ), - }; - let permission_mode = core_session.config.permission_mode; - let permission_profile = core_session.config.permission_profile.clone(); - let hook_context = self.hook_context_for_session(session_id).await; - let turn_cancel_token = self - .active_turn_cancellations - .lock() - .await - .get(&session_id) - .cloned() - .unwrap_or_else(CancellationToken::new); - let provider_http = runtime_context - .config_store - .lock() - .expect("app config store mutex should not be poisoned") - .effective_config() - .provider_http - .clone(); - let runtime = ToolRuntime::new_with_context_and_options( - Arc::clone(®istry), - self.build_permission_checker( - session_id, - turn_for_events.turn_id, - permission_mode, - permission_profile, - ), - ToolRuntimeContext { - session_id: session_id.to_string(), - turn_id: Some(turn_for_events.turn_id.to_string()), - cwd: core_session.cwd.clone(), - agent_scope, - agent_context_mode, - collaboration_mode, - agent_coordinator: Some(Arc::clone(&self) as Arc), - client_filesystem: Some(Arc::clone(&self) as Arc), - client_terminal: Some(Arc::clone(&self) as Arc), - local_web_search: match &turn_config.web_search { - devo_core::ResolvedWebSearchConfig::Local(config) => Some(config.clone()), - devo_core::ResolvedWebSearchConfig::Disabled - | devo_core::ResolvedWebSearchConfig::Provider => None, - }, - hooks: hook_context, - network_proxy: provider_http.proxy_url, - network_no_proxy: provider_http.no_proxy, - }, - ToolExecutionOptions { - cancel_token: turn_cancel_token, - on_tool_execution_start: Some(Arc::new(move |call: ToolCall| { - let tool_execution_start_tx = tool_execution_start_tx.clone(); - // Match the router callback's `BoxFuture` return type and keep the tool-start - // notification on the same bounded async event path as other query events. - Box::pin(async move { - let _ = tool_execution_start_tx - .send(QueryEvent::ToolExecutionStart { id: call.id }) - .await; - }) - })), - ..ToolExecutionOptions::default() - }, - ); - let result = query( - &mut core_session, - &turn_config, - runtime_context.provider_for_route(turn_config.provider_route.clone()), - registry, - &runtime, - Some(callback), - ) - .await; - ( - result, - core_session.total_input_tokens, - core_session.total_output_tokens, - core_session.total_tokens, - core_session.total_cache_creation_tokens, - core_session.total_cache_read_tokens, - core_session.last_input_tokens, - core_session.prompt_token_estimate, - ) - }; - drop(event_tx); - // Wait for the event task to finish draining buffered stream events - // before we persist the terminal turn state. - let event_summary = event_task.await.ok(); - let mut latest_usage = event_summary - .as_ref() - .and_then(|summary| summary.latest_usage.clone()); - let terminal_stop_reason = event_summary.and_then(|summary| summary.stop_reason); - if usage_parent_session_id.is_some() - && let Some(usage) = latest_usage.clone() - { - self.publish_subagent_turn_usage(session_id, turn.turn_id, usage) - .await; - } else if usage_parent_session_id.is_none() - && let Some(snapshot) = self.parent_usage_snapshot(session_id, turn.turn_id).await - { - latest_usage = Some(snapshot.turn_usage.to_turn_usage()); - session_total_input_tokens = snapshot.session_totals.input_tokens; - session_total_output_tokens = snapshot.session_totals.output_tokens; - session_total_tokens = snapshot.session_totals.total_tokens; - session_total_cache_creation_tokens = - snapshot.session_totals.cache_creation_input_tokens; - session_total_cache_read_tokens = snapshot.session_totals.cache_read_input_tokens; - } - self.active_tasks.lock().await.remove(&session_id); - self.active_turn_cancellations - .lock() - .await - .remove(&session_id); - self.active_turn_connections - .lock() - .await - .remove(&session_id); - match &result { - Ok(()) => { - self.run_session_hook( - session_id, - devo_core::HookEvent::Stop, - serde_json::Map::from_iter([( - "stop_hook_active".to_string(), - serde_json::Value::Bool(false), - )]), - ) - .await; - } - Err(error) => { - self.run_session_hook( - session_id, - devo_core::HookEvent::StopFailure, - serde_json::Map::from_iter([ - ( - "error".to_string(), - serde_json::Value::String(error.to_string()), - ), - ( - "error_details".to_string(), - serde_json::Value::String(error.to_string()), - ), - ]), - ) - .await; - } - } - - let final_turn = { - let mut session = session_arc.lock().await; - let mut final_turn = turn.clone(); - final_turn.completed_at = Some(Utc::now()); - final_turn.status = if result.is_ok() { - TurnStatus::Completed - } else { - TurnStatus::Failed - }; - final_turn.usage = latest_usage.clone(); - final_turn.stop_reason = terminal_stop_reason; - final_turn.failure_reason = result - .as_ref() - .err() - .and_then(turn_failure_reason_from_error); - session.latest_turn = Some(final_turn.clone()); - session.active_turn = None; - session.summary.status = SessionRuntimeStatus::Idle; - session.summary.updated_at = Utc::now(); - session.summary.last_activity_at = session.summary.updated_at; - session.summary.total_input_tokens = session_total_input_tokens; - session.summary.total_output_tokens = session_total_output_tokens; - session.summary.total_tokens = session_total_tokens; - session.summary.total_cache_creation_tokens = session_total_cache_creation_tokens; - session.summary.total_cache_read_tokens = session_total_cache_read_tokens; - session.summary.prompt_token_estimate = session_prompt_token_estimate; - if let Some(usage) = &final_turn.usage { - session.summary.last_query_total_tokens = usage.display_total_tokens(); - } - if let Ok(mut core_session) = session.core_session.try_lock() { - core_session.total_input_tokens = session_total_input_tokens; - core_session.total_output_tokens = session_total_output_tokens; - core_session.total_tokens = session_total_tokens; - core_session.total_cache_creation_tokens = session_total_cache_creation_tokens; - core_session.total_cache_read_tokens = session_total_cache_read_tokens; - } - - // Persist token stats to SQLite (skip for ephemeral sessions) - if !session.summary.ephemeral { - let stats = SessionStats { - total_input_tokens: session_total_input_tokens, - total_output_tokens: session_total_output_tokens, - total_tokens: session_total_tokens, - total_cache_creation_tokens: session_total_cache_creation_tokens, - total_cache_read_tokens: session_total_cache_read_tokens, - last_input_tokens: final_turn - .usage - .as_ref() - .map(|u| u.input_tokens as usize) - .unwrap_or(session_last_input_tokens), - turn_count: session.summary.updated_at.timestamp() as usize, - prompt_token_estimate: session_prompt_token_estimate, - }; - if let Err(err) = self.deps.db.update_stats(&session_id, &stats) { - tracing::warn!( - session_id = %session_id, - error = %err, - "failed to persist token stats to database" - ); - } - } - - final_turn - }; - // The turn is finished, so any queued "btw" input no longer applies. - // Clear both the in-memory queue and the persisted mirror. - { - let is_ephemeral = { - let session = session_arc.lock().await; - session.summary.ephemeral - }; - let btw_input_queue = { - let session = session_arc.lock().await; - Arc::clone(&session.btw_input_queue) - }; - btw_input_queue - .lock() - .expect("btw input queue mutex should not be poisoned") - .clear(); - if !is_ephemeral - && let Err(err) = self.deps.db.clear_pending(&session_id, QueueType::Btw) - { - tracing::warn!( - session_id = %session_id, - error = %err, - "failed to clear btw input messages from database" - ); - } - } - - let (record, session_context, turn_context) = { - let session = session_arc.lock().await; - let core_session = session.core_session.lock().await; - ( - session.record.clone(), - core_session.session_context.clone(), - core_session.latest_turn_context.clone(), - ) - }; - if let Some(record) = record - && let Err(error) = self.rollout_store.append_turn( - &record, - build_turn_record(&final_turn, session_context, turn_context), - ) - { - tracing::warn!(session_id = %session_id, error = %error, "failed to persist terminal turn line"); - } - self.finalize_turn_workspace_changes(session_id, &final_turn) - .await; - // Emit the terminal result before we look at queued follow-up input. - if let Err(error) = result { - tracing::warn!( - session_id = %session_id, - turn_id = %final_turn.turn_id, - status = ?final_turn.status, - error = %error, - "turn execution failed" - ); - self.emit_turn_item( - session_id, - final_turn.turn_id, - ItemKind::AgentMessage, - TurnItem::AgentMessage(TextItem { - text: error.to_string(), - }), - serde_json::json!({ "title": "Error", "text": error.to_string() }), - ) - .await; - self.broadcast_event(ServerEvent::TurnFailed(TurnEventPayload { - session_id, - turn: final_turn.clone(), - })) - .await; - } else { - tracing::info!( - session_id = %session_id, - turn_id = %final_turn.turn_id, - status = ?final_turn.status, - total_input_tokens = final_turn.usage.as_ref().map(|usage| usage.input_tokens), - total_output_tokens = final_turn.usage.as_ref().map(|usage| usage.output_tokens), - "turn execution completed" - ); - } - self.handle_subagent_turn_completed(session_id, &final_turn) - .await; - self.broadcast_event(ServerEvent::TurnCompleted(TurnEventPayload { - session_id, - turn: final_turn.clone(), - })) - .await; - self.broadcast_event(ServerEvent::SessionStatusChanged( - SessionStatusChangedPayload { - session_id, - status: SessionRuntimeStatus::Idle, - }, - )) - .await; - self.record_terminal_turn_status( - final_turn.turn_id, - TerminalTurnSnapshot::from_turn(&final_turn), - ) - .await; - - // After the turn completes, check for queued inputs and start the next turn. - let queued_input = { - let session_arc = match self.sessions.lock().await.get(&session_id).cloned() { - Some(s) => s, - None => return, - }; - let pending_turn_queue = { - let session = session_arc.lock().await; - if session.active_turn.is_some() { - return; - } - Arc::clone(&session.pending_turn_queue) - }; - let mut queue = pending_turn_queue - .lock() - .expect("pending turn queue mutex should not be poisoned"); - match queue.pop_front() { - Some(devo_core::PendingInputItem { - kind: devo_core::PendingInputKind::UserText { text }, - metadata, - .. - }) => Some(( - text.clone(), - text, - Vec::new(), - collaboration_mode_from_pending_metadata(metadata.as_ref()), - model_selection_from_pending_metadata(metadata.as_ref()), - subagent_usage_owner_from_pending_metadata(metadata.as_ref()), - )), - Some(devo_core::PendingInputItem { - kind: - devo_core::PendingInputKind::UserInput { - display_text, - prompt_text, - prompt_messages, - .. - }, - metadata, - .. - }) => Some(( - display_text, - prompt_text, - prompt_messages, - collaboration_mode_from_pending_metadata(metadata.as_ref()), - model_selection_from_pending_metadata(metadata.as_ref()), - subagent_usage_owner_from_pending_metadata(metadata.as_ref()), - )), - _ => None, - } - }; - let Some(( - display_input, - input_text, - input_messages, - queued_collaboration_mode, - queued_model_selection, - queued_subagent_usage_owner, - )) = queued_input - else { - self.maybe_start_goal_continuation_turn(session_id).await; - return; - }; - // Update clients before starting the next turn so dequeued input is - // removed from any pending queue display. - self.broadcast_updated_queue(session_id).await; - - let (turn_config, resolved_request) = { - let session_arc = match self.sessions.lock().await.get(&session_id).cloned() { - Some(s) => s, - None => return, - }; - let session = session_arc.lock().await; - let model_override = queued_model_selection - .as_deref() - .or_else(|| session_model_selection(&session.summary)); - let reasoning_effort_selection_override = - session.summary.reasoning_effort_selection.clone(); - let turn_config = session - .runtime_context - .resolve_turn_config(model_override, reasoning_effort_selection_override); - let resolved_request = turn_config.model.resolve_reasoning_effort_selection( - turn_config.reasoning_effort_selection.as_deref(), - ); - (turn_config, resolved_request) - }; - let request_model = turn_config.provider_request_model(&resolved_request.request_model); - - let sequence = { - let session_arc = match self.sessions.lock().await.get(&session_id).cloned() { - Some(s) => s, - None => return, - }; - let session = session_arc.lock().await; - session.latest_turn.as_ref().map_or(1, |t| t.sequence + 1) - }; - - let now = Utc::now(); - let turn = TurnMetadata { - turn_id: TurnId::new(), - session_id, - sequence, - status: TurnStatus::Running, - kind: devo_core::TurnKind::Regular, - model: turn_config.model.slug.clone(), - model_binding_id: turn_config.model_binding_id.clone(), - reasoning_effort_selection: turn_config.reasoning_effort_selection.clone(), - reasoning_effort: resolved_request.effective_reasoning_effort, - request_model, - request_thinking: resolved_request.request_thinking.clone(), - started_at: now, - completed_at: None, - usage: None, - stop_reason: None, - failure_reason: None, - }; - if let Some((parent_session_id, parent_turn_id)) = queued_subagent_usage_owner { - self.register_subagent_usage_owner(parent_session_id, session_id, parent_turn_id) - .await; - } - { - let session_arc = match self.sessions.lock().await.get(&session_id).cloned() { - Some(s) => s, - None => return, - }; - let mut session = session_arc.lock().await; - apply_turn_config_to_session_summary(&mut session.summary, &turn_config); - session.summary.status = SessionRuntimeStatus::ActiveTurn; - session.summary.updated_at = now; - session.summary.last_activity_at = now; - session.active_turn = Some(turn.clone()); - } - self.broadcast_event(ServerEvent::TurnStarted(TurnEventPayload { - session_id, - turn: turn.clone(), - })) - .await; - // Chain directly instead of spawning so this drain loop can keep - // consuming queued input until the queue is empty. - // This is an async self-call into `execute_turn`; boxing adds the heap indirection Rust - // needs so the future does not recursively contain another copy of its own type. - Box::pin(Arc::clone(&self).execute_turn(ExecuteTurnRequest { - session_id, - turn, - turn_config, - display_input, - input: input_text, - input_messages, - collaboration_mode: queued_collaboration_mode, - input_mode: TurnInputMode::VisibleUserMessage, - })) - .await; - } - - /// Pop the first queued input and start a new turn in a background task. - /// Used from the interrupt handler where the calling function must return - /// its response immediately. - pub(super) async fn spawn_next_turn_from_queue(self: &Arc, session_id: SessionId) { - // Pop one queued input. - let ( - display_input, - input_text, - input_messages, - queued_collaboration_mode, - queued_model_selection, - queued_subagent_usage_owner, - ) = { - let session_arc = match self.sessions.lock().await.get(&session_id).cloned() { - Some(s) => s, - None => return, - }; - let pending_turn_queue = { - let session = session_arc.lock().await; - Arc::clone(&session.pending_turn_queue) - }; - let mut guard = pending_turn_queue - .lock() - .expect("pending turn queue mutex should not be poisoned"); - match guard.pop_front() { - Some(devo_core::PendingInputItem { - kind: devo_core::PendingInputKind::UserText { text }, - metadata, - .. - }) => ( - text.clone(), - text, - Vec::new(), - collaboration_mode_from_pending_metadata(metadata.as_ref()), - model_selection_from_pending_metadata(metadata.as_ref()), - subagent_usage_owner_from_pending_metadata(metadata.as_ref()), - ), - Some(devo_core::PendingInputItem { - kind: - devo_core::PendingInputKind::UserInput { - display_text, - prompt_text, - prompt_messages, - .. - }, - metadata, - .. - }) => ( - display_text, - prompt_text, - prompt_messages, - collaboration_mode_from_pending_metadata(metadata.as_ref()), - model_selection_from_pending_metadata(metadata.as_ref()), - subagent_usage_owner_from_pending_metadata(metadata.as_ref()), - ), - _ => return, - } - }; - // Broadcast the updated queue state so the TUI removes this item - // from its pending cells list. - self.broadcast_updated_queue(session_id).await; - - // Resolve turn config from session metadata. - let (turn_config, resolved_request) = { - let session_arc = match self.sessions.lock().await.get(&session_id).cloned() { - Some(s) => s, - None => return, - }; - let session = session_arc.lock().await; - let model = queued_model_selection - .as_deref() - .or_else(|| session_model_selection(&session.summary)); - let reasoning_effort_selection = session.summary.reasoning_effort_selection.clone(); - let tc = session - .runtime_context - .resolve_turn_config(model, reasoning_effort_selection); - let rr = tc - .model - .resolve_reasoning_effort_selection(tc.reasoning_effort_selection.as_deref()); - (tc, rr) - }; - let request_model = turn_config.provider_request_model(&resolved_request.request_model); - - let sequence = { - let session_arc = match self.sessions.lock().await.get(&session_id).cloned() { - Some(s) => s, - None => return, - }; - let session = session_arc.lock().await; - session.latest_turn.as_ref().map_or(1, |t| t.sequence + 1) - }; - - let now = Utc::now(); - let turn = TurnMetadata { - turn_id: TurnId::new(), - session_id, - sequence, - status: TurnStatus::Running, - kind: devo_core::TurnKind::Regular, - model: turn_config.model.slug.clone(), - model_binding_id: turn_config.model_binding_id.clone(), - reasoning_effort_selection: turn_config.reasoning_effort_selection.clone(), - reasoning_effort: resolved_request.effective_reasoning_effort, - request_model, - request_thinking: resolved_request.request_thinking.clone(), - started_at: now, - completed_at: None, - usage: None, - stop_reason: None, - failure_reason: None, - }; - if let Some((parent_session_id, parent_turn_id)) = queued_subagent_usage_owner { - self.register_subagent_usage_owner(parent_session_id, session_id, parent_turn_id) - .await; - } - { - let session_arc = match self.sessions.lock().await.get(&session_id).cloned() { - Some(s) => s, - None => return, - }; - let mut session = session_arc.lock().await; - apply_turn_config_to_session_summary(&mut session.summary, &turn_config); - session.summary.status = SessionRuntimeStatus::ActiveTurn; - session.summary.updated_at = now; - session.summary.last_activity_at = now; - session.active_turn = Some(turn.clone()); - } - self.broadcast_event(ServerEvent::TurnStarted(TurnEventPayload { - session_id, - turn: turn.clone(), - })) - .await; - // Spawn the turn in the background so the caller (interrupt handler) - // can return its response immediately. The spawned task will call - // drain_and_start_next_turn on completion, draining the entire queue. - let runtime = Arc::clone(self); - tokio::spawn(async move { - runtime - .execute_turn(ExecuteTurnRequest { - session_id, - turn, - turn_config, - display_input, - input: input_text, - input_messages, - collaboration_mode: queued_collaboration_mode, - input_mode: TurnInputMode::VisibleUserMessage, - }) - .await; - }); - } - - /// Read the current steering queue and broadcast its state to connected clients. - /// Called after any queue mutation (enqueue, dequeue, clear) so the TUI preview - /// stays in sync. - pub(super) async fn broadcast_updated_queue(&self, session_id: SessionId) { - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { - return; - }; - let (pending_count, pending_texts) = { - let pending_turn_queue = { - let session = session_arc.lock().await; - Arc::clone(&session.pending_turn_queue) - }; - let queue = pending_turn_queue - .lock() - .expect("pending turn queue mutex should not be poisoned"); - let texts: Vec = queue - .iter() - .filter_map(|item| match &item.kind { - devo_core::PendingInputKind::UserText { text } => Some(text.clone()), - devo_core::PendingInputKind::UserInput { display_text, .. } => { - Some(display_text.clone()) - } - _ => None, - }) - .collect(); - (texts.len(), texts) - }; - self.broadcast_event(ServerEvent::InputQueueUpdated( - devo_core::InputQueueUpdatedPayload { - session_id, - pending_count, - pending_texts, - }, - )) - .await; - } -} - -#[cfg(test)] -mod tests { - use super::*; - use pretty_assertions::assert_eq; - - #[test] - fn command_progress_uses_command_execution_item_id() { - let command_item_id = ItemId::new(); - let tool_item_id = ItemId::new(); - let mut pending_tool_calls = HashMap::new(); - pending_tool_calls.insert( - "exec".to_string(), - PendingToolCall { - item_id: Some(command_item_id), - item_seq: Some(1), - input: serde_json::json!({}), - display_kind: ToolDisplayKind::CommandExecution, - command: "cargo test".to_string(), - }, - ); - pending_tool_calls.insert( - "read".to_string(), - PendingToolCall { - item_id: Some(tool_item_id), - item_seq: Some(2), - input: serde_json::json!({}), - display_kind: ToolDisplayKind::Generic, - command: String::new(), - }, - ); - - assert_eq!( - command_execution_item_id_for_progress(&pending_tool_calls, "exec"), - Some(command_item_id) - ); - assert_eq!( - command_execution_item_id_for_progress(&pending_tool_calls, "read"), - None - ); - assert_eq!( - command_execution_item_id_for_progress(&pending_tool_calls, "missing"), - None - ); - } - - #[test] - fn file_change_tool_detection_matches_apply_patch_and_write() { - assert!(is_file_change_tool("apply_patch")); - assert!(is_file_change_tool("write")); - assert!(!is_file_change_tool("read")); - } - - #[test] - fn plan_tool_detection_matches_update_plan() { - assert!(is_plan_tool("update_plan")); - assert!(!is_plan_tool("read")); - } - - #[test] - fn read_tool_start_item_contains_live_read_action() { - let input = serde_json::json!({ - "path": "crates/tui/src/mod.rs" - }); - let start_item = tool_start_item_from_input( - "call-1", - "read", - "read crates/tui/src/mod.rs", - &input, - ToolDisplayKind::Generic, - ToolPreparationFeedback::None, - ); - - let payload: ToolCallPayload = - serde_json::from_value(start_item.payload).expect("tool call payload"); - - assert_eq!(start_item.item_kind, ItemKind::ToolCall); - assert_eq!( - payload, - ToolCallPayload { - tool_call_id: "call-1".to_string(), - tool_name: "read".to_string(), - parameters: input, - command_actions: vec![devo_protocol::parse_command::ParsedCommand::Read { - cmd: "read crates/tui/src/mod.rs".to_string(), - name: "mod.rs".to_string(), - path: std::path::PathBuf::from("crates/tui/src/mod.rs"), - }], - } - ); - } - - #[test] - fn grep_tool_start_item_contains_live_search_action() { - let input = serde_json::json!({ - "pattern": "ToolUseStart", - "path": "crates/server/src" - }); - let start_item = tool_start_item_from_input( - "call-1", - "grep", - "grep ToolUseStart in crates/server/src", - &input, - ToolDisplayKind::Generic, - ToolPreparationFeedback::None, - ); - - let payload: ToolCallPayload = - serde_json::from_value(start_item.payload).expect("tool call payload"); - - assert_eq!(start_item.item_kind, ItemKind::ToolCall); - assert_eq!( - payload, - ToolCallPayload { - tool_call_id: "call-1".to_string(), - tool_name: "grep".to_string(), - parameters: input, - command_actions: vec![devo_protocol::parse_command::ParsedCommand::Search { - cmd: "grep ToolUseStart in crates/server/src".to_string(), - query: Some("ToolUseStart".to_string()), - path: Some("crates/server/src".to_string()), - }], - } - ); - } - - #[test] - fn code_search_tool_start_item_contains_live_search_action() { - let input = serde_json::json!({ - "operation": "search", - "query": "live tool feedback", - "path": "crates" - }); - let start_item = tool_start_item_from_input( - "call-1", - "code_search", - "code_search live tool feedback in crates", - &input, - ToolDisplayKind::Generic, - ToolPreparationFeedback::None, - ); - - let payload: ToolCallPayload = - serde_json::from_value(start_item.payload).expect("tool call payload"); - - assert_eq!(start_item.item_kind, ItemKind::ToolCall); - assert_eq!( - payload, - ToolCallPayload { - tool_call_id: "call-1".to_string(), - tool_name: "code_search".to_string(), - parameters: input, - command_actions: vec![devo_protocol::parse_command::ParsedCommand::Search { - cmd: "code_search live tool feedback in crates".to_string(), - query: Some("live tool feedback".to_string()), - path: Some("crates".to_string()), - }], - } - ); - } - - #[test] - fn exec_tool_start_item_uses_command_execution_payload() { - let input = serde_json::json!({ - "cmd": "cargo test -p devo-server" - }); - let start_item = tool_start_item_from_input( - "call-1", - "exec_command", - "cargo test -p devo-server", - &input, - ToolDisplayKind::CommandExecution, - ToolPreparationFeedback::None, - ); - - let payload: CommandExecutionPayload = - serde_json::from_value(start_item.payload).expect("command execution payload"); - - assert_eq!(start_item.item_kind, ItemKind::CommandExecution); - assert_eq!( - payload, - CommandExecutionPayload { - tool_call_id: "call-1".to_string(), - tool_name: "exec_command".to_string(), - command: "cargo test -p devo-server".to_string(), - input: Some(input), - source: devo_protocol::protocol::ExecCommandSource::Agent, - command_actions: Vec::new(), - output: None, - is_error: false, - } - ); - } - - #[test] - fn user_shell_exec_input_uses_pty_backed_exec_command() { - let cwd = std::path::PathBuf::from("workspace"); - - let input = user_shell_exec_input("pwd", cwd.clone()); - - assert_eq!( - input, - serde_json::json!({ - "cmd": "pwd", - "workdir": cwd, - "login": true, - "tty": true, - }) - ); - } - - #[test] - fn user_shell_command_payload_uses_user_shell_source() { - let output = serde_json::json!({ "output": "done" }); - - let input = user_shell_exec_input("pwd", std::path::PathBuf::from("workspace")); - let payload = user_shell_command_payload( - "call-1", - "pwd", - input.clone(), - Vec::new(), - Some(output.clone()), - false, - ); - - assert_eq!( - payload, - CommandExecutionPayload { - tool_call_id: "call-1".to_string(), - tool_name: "exec_command".to_string(), - command: "pwd".to_string(), - input: Some(input), - source: devo_protocol::protocol::ExecCommandSource::UserShell, - command_actions: Vec::new(), - output: Some(output), - is_error: false, - } - ); - } - - #[test] - fn live_only_apply_patch_start_item_stays_tool_call() { - let input = serde_json::json!({ - "patch": "*** Begin Patch\n*** End Patch" - }); - let start_item = tool_start_item_from_input( - "call-1", - "apply_patch", - "apply_patch", - &input, - ToolDisplayKind::Generic, - ToolPreparationFeedback::LiveOnly, - ); - - let payload: ToolCallPayload = - serde_json::from_value(start_item.payload).expect("tool call payload"); - - assert_eq!(start_item.item_kind, ItemKind::ToolCall); - assert_eq!( - payload, - ToolCallPayload { - tool_call_id: "call-1".to_string(), - tool_name: "apply_patch".to_string(), - parameters: input, - command_actions: Vec::new(), - } - ); - } - - #[test] - fn command_actions_from_read_tool_input_builds_read_action() { - let actions = command_actions_from_tool_input( - "read", - "read crates/tui/src/mod.rs", - &serde_json::json!({ - "filePath": "crates/tui/src/mod.rs" - }), - ); - assert_eq!( - actions, - vec![devo_protocol::parse_command::ParsedCommand::Read { - cmd: "read crates/tui/src/mod.rs".to_string(), - name: "mod.rs".to_string(), - path: std::path::PathBuf::from("crates/tui/src/mod.rs"), - }] - ); - } - - #[test] - fn command_actions_from_read_tool_input_without_path_is_empty() { - let actions = - command_actions_from_tool_input("read", "read", &serde_json::json!({ "limit": 10 })); - assert_eq!(actions, Vec::new()); - } - - #[test] - fn command_actions_from_read_tool_result_summary_recovers_final_path() { - let actions = command_actions_from_tool_result( - "read", - "read ", - &serde_json::json!({}), - "read: crates/tui/src/mod.rs", - ); - assert_eq!( - actions, - vec![devo_protocol::parse_command::ParsedCommand::Read { - cmd: "read crates/tui/src/mod.rs".to_string(), - name: "mod.rs".to_string(), - path: std::path::PathBuf::from("crates/tui/src/mod.rs"), - }] - ); - } - - #[test] - fn command_actions_from_grep_tool_input_builds_search_action() { - let actions = command_actions_from_tool_input( - "grep", - "grep rebuild_restored_session in crates/tui/src", - &serde_json::json!({ - "pattern": "rebuild_restored_session", - "path": "crates/tui/src" - }), - ); - assert_eq!(actions.len(), 1); - assert!(matches!( - &actions[0], - devo_protocol::parse_command::ParsedCommand::Search { query, path, .. } - if query.as_deref() == Some("rebuild_restored_session") - && path.as_deref() == Some("crates/tui/src") - )); - } - - #[test] - fn command_actions_from_glob_tool_input_include_pattern_and_path() { - let actions = command_actions_from_tool_input( - "glob", - "glob **/Cargo.toml in crates", - &serde_json::json!({ - "pattern": "**/Cargo.toml", - "path": "crates" - }), - ); - assert_eq!( - actions, - vec![devo_protocol::parse_command::ParsedCommand::ListFiles { - cmd: "glob **/Cargo.toml in crates".to_string(), - path: Some("**/Cargo.toml in crates".to_string()), - }] - ); - } - - #[test] - fn command_actions_from_find_tool_input_include_pattern_and_path() { - let actions = command_actions_from_tool_input( - "find", - "find **/Cargo.toml in crates", - &serde_json::json!({ - "pattern": "**/Cargo.toml", - "path": "crates" - }), - ); - assert_eq!( - actions, - vec![devo_protocol::parse_command::ParsedCommand::ListFiles { - cmd: "find **/Cargo.toml in crates".to_string(), - path: Some("**/Cargo.toml in crates".to_string()), - }] - ); - } -} diff --git a/crates/server/src/runtime/turn_exec/event_stream.rs b/crates/server/src/runtime/turn_exec/event_stream.rs new file mode 100644 index 00000000..309b2236 --- /dev/null +++ b/crates/server/src/runtime/turn_exec/event_stream.rs @@ -0,0 +1,664 @@ +use std::sync::Arc; + +use devo_core::{ItemId, SessionId, TurnId}; +use devo_protocol::TurnFailureReason; + +use super::super::ServerRuntime; +use super::super::proposed_plan::ProposedPlanParser; +use super::item_stream::{ + ProposedPlanStreamItem, complete_assistant_item, complete_reasoning_item, + handle_proposed_plan_segments, push_assistant_text_delta, +}; +use super::tool_display::{ + command_display_from_input, tool_start_item_from_input, tool_start_item_from_result, +}; +use super::tool_results; +use super::tool_results::complete_pending_tool_call; +use super::trace::{ + query_event_trace_delta_len, query_event_trace_kind, query_event_trace_token_preview, + stream_trace_elapsed_ms, +}; +use super::types::{PendingToolCall, ToolDisplayKind, TurnEventStreamSummary}; +use crate::runtime::session_actor::state::SessionStreamState; +use crate::{ItemDeltaKind, ItemDeltaPayload, ServerEvent}; +use tokio::sync::mpsc; + +pub(crate) const QUERY_EVENT_CHANNEL_CAPACITY: usize = 1024; + +/// Enqueue a query event without blocking the caller when the channel is full. +/// +/// `run_turn_model_query` holds the session `core_session` mutex for the entire +/// `query()` call. Blocking on `event_tx.send().await` while the event task is +/// stalled on client-notification backpressure would freeze tool execution +/// (including `spawn_agent`) and can wedge the whole connection. +pub(super) fn enqueue_query_event( + event_tx: &mpsc::Sender, + event: devo_core::QueryEvent, +) { + match event_tx.try_send(event) { + Ok(()) => {} + Err(tokio::sync::mpsc::error::TrySendError::Full(event)) => { + let event_tx = event_tx.clone(); + tokio::spawn(async move { + let _ = event_tx.send(event).await; + }); + } + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {} + } +} + +pub(super) fn turn_failure_reason_from_error( + error: &devo_core::AgentError, +) -> Option { + match error { + devo_core::AgentError::MaxTurnsExceeded(_) => Some(TurnFailureReason::MaxTurnRequests), + devo_core::AgentError::Provider(_) + | devo_core::AgentError::ContextTooLong + | devo_core::AgentError::Aborted => None, + } +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn spawn_turn_event_stream( + runtime: Arc, + event_stream: Arc>, + session_id: SessionId, + turn: crate::TurnMetadata, + collaboration_mode: devo_protocol::CollaborationMode, + event_tool_registry: Arc, + usage_parent_session_id: Option, + usage_context_window: Option, + mut event_rx: tokio::sync::mpsc::Receiver, +) -> tokio::task::JoinHandle { + let turn_for_events = turn.clone(); + let turn_for_plan_updates = turn; + tokio::spawn(async move { + let mut assistant_item_id = None; + let mut assistant_item_seq = None; + let mut assistant_delta_seq = 0_u64; + let mut assistant_text = String::new(); + let mut reasoning_item_id = None; + let mut reasoning_item_seq = None; + let mut reasoning_text = String::new(); + let mut tool_names_by_id = std::collections::HashMap::new(); + let mut pending_tool_calls: std::collections::HashMap = + std::collections::HashMap::new(); + let mut proposed_plan_parser = (collaboration_mode + == devo_protocol::CollaborationMode::Plan) + .then(ProposedPlanParser::default); + let mut proposed_plan_item = ProposedPlanStreamItem::default(); + let mut proposed_plan_leading_normal = String::new(); + let mut latest_usage = None; + let mut stop_reason = None; + while let Some(event) = event_rx.recv().await { + log_dequeued_query_event(&event); + match event { + devo_core::QueryEvent::TextDelta(text) => { + if let Some(parser) = proposed_plan_parser.as_mut() { + let segments = parser.push_str(&text); + handle_proposed_plan_segments( + &runtime, + &event_stream, + session_id, + turn_for_events.turn_id, + segments, + &mut assistant_item_id, + &mut assistant_item_seq, + &mut assistant_text, + &mut assistant_delta_seq, + &mut proposed_plan_item, + &mut proposed_plan_leading_normal, + ) + .await; + } else { + push_assistant_text_delta( + &runtime, + &event_stream, + session_id, + turn_for_events.turn_id, + &mut assistant_item_id, + &mut assistant_item_seq, + &mut assistant_text, + &mut assistant_delta_seq, + text, + ) + .await; + } + } + devo_core::QueryEvent::ReasoningDelta(text) => { + handle_reasoning_delta( + &runtime, + &event_stream, + session_id, + turn_for_events.turn_id, + text, + &mut reasoning_item_id, + &mut reasoning_item_seq, + &mut reasoning_text, + ) + .await; + } + devo_core::QueryEvent::ReasoningCompleted => { + complete_open_reasoning_item( + &runtime, + session_id, + turn_for_events.turn_id, + &mut reasoning_item_id, + &mut reasoning_item_seq, + &mut reasoning_text, + &event_stream, + ) + .await; + } + devo_core::QueryEvent::ToolUseStart { id, name, input } => { + handle_tool_use_start( + &runtime, + session_id, + turn_for_events.turn_id, + id, + name, + input, + &mut tool_names_by_id, + &mut pending_tool_calls, + &mut reasoning_item_id, + &mut reasoning_item_seq, + &mut reasoning_text, + &mut assistant_item_id, + &mut assistant_item_seq, + &mut assistant_text, + &event_tool_registry, + ) + .await; + } + devo_core::QueryEvent::ToolExecutionStart { id } => { + runtime + .broadcast_event(ServerEvent::ToolCallStatusUpdated( + devo_protocol::ToolCallStatusUpdatedPayload { + session_id, + turn_id: turn_for_events.turn_id, + tool_call_id: id, + status: "in_progress".to_string(), + terminal_id: None, + }, + )) + .await; + } + devo_core::QueryEvent::ToolResult { + tool_use_id, + tool_name: final_tool_name, + input: final_input, + content, + display_content, + is_error, + summary, + } => { + handle_tool_result( + &runtime, + session_id, + turn_for_events.turn_id, + &turn_for_plan_updates, + tool_use_id, + final_tool_name, + final_input, + content, + display_content, + is_error, + summary, + &tool_names_by_id, + &mut pending_tool_calls, + &event_tool_registry, + ) + .await; + } + devo_core::QueryEvent::ToolProgress { + tool_use_id, + progress, + } => { + handle_tool_progress( + &runtime, + session_id, + turn_for_events.turn_id, + tool_use_id, + progress, + &pending_tool_calls, + ) + .await; + } + devo_core::QueryEvent::UsageDelta { usage } + | devo_core::QueryEvent::Usage { usage } => { + let turn_usage = devo_core::TurnUsage::from_usage(&usage); + latest_usage = Some(turn_usage.clone()); + if usage_parent_session_id.is_some() { + let _ = runtime + .publish_subagent_turn_usage( + session_id, + turn_for_events.turn_id, + turn_usage, + ) + .await; + } else if let Some(snapshot) = runtime + .publish_parent_turn_usage( + session_id, + turn_for_events.turn_id, + turn_usage, + usage_context_window, + ) + .await + { + latest_usage = Some(snapshot.turn_usage.to_turn_usage()); + } + } + devo_core::QueryEvent::TurnComplete { + stop_reason: terminal_stop_reason, + } => { + stop_reason = Some(terminal_stop_reason); + } + } + } + finish_proposed_plan_stream( + &runtime, + &event_stream, + session_id, + turn_for_events.turn_id, + &mut proposed_plan_parser, + &mut assistant_item_id, + &mut assistant_item_seq, + &mut assistant_text, + &mut assistant_delta_seq, + &mut proposed_plan_item, + &mut proposed_plan_leading_normal, + ) + .await; + complete_deferred_stream_items( + &runtime, + &event_stream, + session_id, + turn_for_events.turn_id, + ) + .await; + tracing::debug!( + session_id = %session_id, + turn_id = %turn_for_events.turn_id, + "query event stream drained" + ); + TurnEventStreamSummary { + latest_usage, + stop_reason, + } + }) +} + +fn log_dequeued_query_event(event: &devo_core::QueryEvent) { + let assistant_token_text = query_event_trace_token_preview(event); + if let Some(assistant_token_text) = assistant_token_text.as_deref() { + tracing::debug!( + stream_elapsed_ms = stream_trace_elapsed_ms(), + event_kind = query_event_trace_kind(event), + delta_len = query_event_trace_delta_len(event), + assistant_token_text, + "query event bridge dequeued by turn event task" + ); + } else { + tracing::debug!( + stream_elapsed_ms = stream_trace_elapsed_ms(), + event_kind = query_event_trace_kind(event), + delta_len = query_event_trace_delta_len(event), + "query event bridge dequeued by turn event task" + ); + } +} + +#[allow(clippy::too_many_arguments)] +async fn handle_reasoning_delta( + runtime: &Arc, + event_stream: &Arc>, + session_id: SessionId, + turn_id: TurnId, + text: String, + reasoning_item_id: &mut Option, + reasoning_item_seq: &mut Option, + reasoning_text: &mut String, +) { + let (item_id, item_seq) = match (*reasoning_item_id, *reasoning_item_seq) { + (Some(item_id), Some(item_seq)) => (item_id, item_seq), + (None, None) => { + let (item_id, item_seq) = runtime + .start_item( + session_id, + turn_id, + crate::ItemKind::Reasoning, + serde_json::json!({ "title": "Reasoning", "text": "" }), + ) + .await; + *reasoning_item_id = Some(item_id); + *reasoning_item_seq = Some(item_seq); + (item_id, item_seq) + } + _ => return, + }; + reasoning_text.push_str(&text); + runtime + .broadcast_event(ServerEvent::ItemDelta { + delta_kind: ItemDeltaKind::ReasoningTextDelta, + payload: ItemDeltaPayload { + context: crate::EventContext { + session_id, + turn_id: Some(turn_id), + item_id: Some(item_id), + seq: 0, + }, + delta: text, + stream_index: None, + channel: None, + }, + }) + .await; + let _ = item_seq; + if let Ok(mut stream) = event_stream.try_lock() { + stream.deferred_reasoning = Some((item_id, item_seq, reasoning_text.clone())); + } +} + +async fn complete_open_reasoning_item( + runtime: &Arc, + session_id: SessionId, + turn_id: TurnId, + reasoning_item_id: &mut Option, + reasoning_item_seq: &mut Option, + reasoning_text: &mut String, + event_stream: &Arc>, +) { + if let (Some(item_id), Some(item_seq)) = (reasoning_item_id.take(), reasoning_item_seq.take()) { + if let Ok(mut stream) = event_stream.try_lock() { + stream.deferred_reasoning.take(); + } + complete_reasoning_item( + runtime, + session_id, + turn_id, + item_id, + item_seq, + reasoning_text.clone(), + ) + .await; + reasoning_text.clear(); + } +} + +#[allow(clippy::too_many_arguments)] +async fn handle_tool_use_start( + runtime: &Arc, + session_id: SessionId, + turn_id: TurnId, + id: String, + name: String, + input: serde_json::Value, + tool_names_by_id: &mut std::collections::HashMap, + pending_tool_calls: &mut std::collections::HashMap, + reasoning_item_id: &mut Option, + reasoning_item_seq: &mut Option, + reasoning_text: &mut String, + assistant_item_id: &mut Option, + assistant_item_seq: &mut Option, + assistant_text: &mut String, + event_tool_registry: &Arc, +) { + tool_names_by_id.insert(id.clone(), name.clone()); + if let (Some(item_id), Some(item_seq)) = (reasoning_item_id.take(), reasoning_item_seq.take()) { + complete_reasoning_item( + runtime, + session_id, + turn_id, + item_id, + item_seq, + reasoning_text.clone(), + ) + .await; + reasoning_text.clear(); + } + if let (Some(item_id), Some(item_seq)) = (assistant_item_id.take(), assistant_item_seq.take()) { + complete_assistant_item( + runtime, + session_id, + turn_id, + item_id, + item_seq, + assistant_text.clone(), + ) + .await; + assistant_text.clear(); + } + let display_kind = ToolDisplayKind::for_tool_name(&name); + let command = command_display_from_input(&name, &input); + let preparation_feedback = event_tool_registry.preparation_feedback(&name); + let start_item = tool_start_item_from_input( + &id, + &name, + &command, + &input, + display_kind, + preparation_feedback, + ); + let (item_id, item_seq) = runtime + .start_item( + session_id, + turn_id, + start_item.item_kind, + start_item.payload, + ) + .await; + pending_tool_calls.insert( + id, + PendingToolCall { + item_id: Some(item_id), + item_seq: Some(item_seq), + input, + display_kind, + command, + }, + ); +} + +#[allow(clippy::too_many_arguments)] +async fn handle_tool_result( + runtime: &Arc, + session_id: SessionId, + turn_id: TurnId, + turn_for_plan_updates: &crate::TurnMetadata, + tool_use_id: String, + final_tool_name: String, + final_input: serde_json::Value, + content: devo_core::tools::ToolContent, + display_content: Option, + is_error: bool, + summary: String, + tool_names_by_id: &std::collections::HashMap, + pending_tool_calls: &mut std::collections::HashMap, + event_tool_registry: &Arc, +) { + let tool_name = if final_tool_name.is_empty() { + tool_names_by_id.get(&tool_use_id).cloned() + } else { + Some(final_tool_name) + }; + let mut result_input = (!final_input.is_null()).then(|| final_input.clone()); + if let Some(mut pending) = pending_tool_calls.remove(&tool_use_id) { + if !final_input.is_null() { + pending.command = tool_name + .as_deref() + .map(|tool_name| command_display_from_input(tool_name, &final_input)) + .unwrap_or_default(); + pending.input = final_input; + } + result_input = Some(pending.input.clone()); + if (pending.item_id.is_none() || pending.item_seq.is_none()) + && let Some(tool_name) = tool_name.clone() + { + let preparation_feedback = event_tool_registry.preparation_feedback(&tool_name); + let start_item = tool_start_item_from_result( + &tool_use_id, + &tool_name, + &pending.command, + &pending.input, + pending.display_kind, + preparation_feedback, + &summary, + ); + let (item_id, item_seq) = runtime + .start_item( + session_id, + turn_id, + start_item.item_kind, + start_item.payload, + ) + .await; + pending.item_id = Some(item_id); + pending.item_seq = Some(item_seq); + } + if complete_pending_tool_call( + runtime, + session_id, + turn_id, + turn_for_plan_updates, + &tool_use_id, + tool_name.clone(), + &pending, + &content, + display_content.clone(), + is_error, + &summary, + ) + .await + { + return; + } + } + tool_results::emit_tool_result_item( + runtime, + session_id, + turn_id, + tool_use_id, + tool_name, + result_input, + content, + display_content, + is_error, + summary, + ) + .await; +} + +async fn handle_tool_progress( + runtime: &Arc, + session_id: SessionId, + turn_id: TurnId, + tool_use_id: String, + progress: devo_core::tools::ToolProgress, + pending_tool_calls: &std::collections::HashMap, +) { + let content = match progress { + devo_core::tools::ToolProgress::OutputDelta { delta } => Some(delta), + devo_core::tools::ToolProgress::StatusUpdate { message, percent } => Some(match percent { + Some(percent) => format!("{message} ({percent}%)"), + None => message, + }), + devo_core::tools::ToolProgress::Completion { summary } => Some(summary), + devo_core::tools::ToolProgress::Terminal { terminal_id } => { + runtime + .broadcast_event(ServerEvent::ToolCallStatusUpdated( + devo_protocol::ToolCallStatusUpdatedPayload { + session_id, + turn_id, + tool_call_id: tool_use_id.clone(), + status: "in_progress".to_string(), + terminal_id: Some(terminal_id), + }, + )) + .await; + None + } + }; + let Some(content) = content else { + return; + }; + let item_id = super::tool_display::command_execution_item_id_for_progress( + pending_tool_calls, + &tool_use_id, + ); + let _ = runtime + .broadcast_event(ServerEvent::ItemDelta { + delta_kind: ItemDeltaKind::CommandExecutionOutputDelta, + payload: ItemDeltaPayload { + context: crate::EventContext { + session_id, + turn_id: Some(turn_id), + item_id, + seq: 0, + }, + delta: serde_json::json!({ + "tool_use_id": tool_use_id, + "text": content, + }) + .to_string(), + stream_index: None, + channel: None, + }, + }) + .await; +} + +#[allow(clippy::too_many_arguments)] +async fn finish_proposed_plan_stream( + runtime: &Arc, + event_stream: &Arc>, + session_id: SessionId, + turn_id: TurnId, + proposed_plan_parser: &mut Option, + assistant_item_id: &mut Option, + assistant_item_seq: &mut Option, + assistant_text: &mut String, + assistant_delta_seq: &mut u64, + proposed_plan_item: &mut ProposedPlanStreamItem, + proposed_plan_leading_normal: &mut String, +) { + if let Some(parser) = proposed_plan_parser.as_mut() { + let segments = parser.finish(); + handle_proposed_plan_segments( + runtime, + event_stream, + session_id, + turn_id, + segments, + assistant_item_id, + assistant_item_seq, + assistant_text, + assistant_delta_seq, + proposed_plan_item, + proposed_plan_leading_normal, + ) + .await; + proposed_plan_item + .complete(runtime, session_id, turn_id) + .await; + } +} + +async fn complete_deferred_stream_items( + runtime: &Arc, + event_stream: &Arc>, + session_id: SessionId, + turn_id: TurnId, +) { + if let Some((item_id, item_seq, text)) = { + let mut stream = event_stream.lock().await; + stream.deferred_reasoning.take() + } { + complete_reasoning_item(runtime, session_id, turn_id, item_id, item_seq, text).await; + } + if let Some((item_id, item_seq, text)) = { + let mut stream = event_stream.lock().await; + stream.deferred_assistant.take() + } { + complete_assistant_item(runtime, session_id, turn_id, item_id, item_seq, text).await; + } +} diff --git a/crates/server/src/runtime/turn_exec/finalize.rs b/crates/server/src/runtime/turn_exec/finalize.rs new file mode 100644 index 00000000..5bcb1b31 --- /dev/null +++ b/crates/server/src/runtime/turn_exec/finalize.rs @@ -0,0 +1,320 @@ +use std::sync::Arc; + +use chrono::Utc; +use devo_core::{SessionId, TextItem, TurnItem, TurnStatus}; + +use super::super::ServerRuntime; +use super::event_stream::turn_failure_reason_from_error; +use super::types::{TurnEventStreamSummary, TurnQueryOutcome}; +use crate::db::{QueueType, SessionStats}; +use crate::persistence::build_turn_record; +use crate::runtime::session_actor::SessionActorState; +use crate::{ItemKind, SessionRuntimeStatus, SessionStatusChangedPayload, TurnEventPayload}; + +pub(crate) struct FinalizeTurnParams<'a> { + pub state: &'a mut SessionActorState, + pub session_id: SessionId, + pub turn: crate::TurnMetadata, + pub query_outcome: TurnQueryOutcome, + pub event_summary: Option, + pub usage_parent_session_id: Option, +} + +impl ServerRuntime { + pub(crate) async fn finalize_executed_turn(self: &Arc, params: FinalizeTurnParams<'_>) { + let FinalizeTurnParams { + state, + session_id, + turn, + query_outcome, + event_summary, + usage_parent_session_id, + } = params; + let TurnQueryOutcome { + result, + mut session_total_input_tokens, + mut session_total_output_tokens, + mut session_total_tokens, + mut session_total_cache_creation_tokens, + mut session_total_cache_read_tokens, + session_last_input_tokens, + session_prompt_token_estimate, + } = query_outcome; + let mut latest_usage = event_summary + .as_ref() + .and_then(|summary| summary.latest_usage.clone()); + let terminal_stop_reason = event_summary.and_then(|summary| summary.stop_reason); + if usage_parent_session_id.is_some() + && let Some(usage) = latest_usage.clone() + { + self.publish_subagent_turn_usage(session_id, turn.turn_id, usage) + .await; + } else if usage_parent_session_id.is_none() + && let Some(snapshot) = self.parent_usage_snapshot(session_id, turn.turn_id).await + { + latest_usage = Some(snapshot.turn_usage.to_turn_usage()); + session_total_input_tokens = snapshot.session_totals.input_tokens; + session_total_output_tokens = snapshot.session_totals.output_tokens; + session_total_tokens = snapshot.session_totals.total_tokens; + session_total_cache_creation_tokens = + snapshot.session_totals.cache_creation_input_tokens; + session_total_cache_read_tokens = snapshot.session_totals.cache_read_input_tokens; + } + self.active_tasks.lock().await.remove(&session_id); + self.active_turn_cancellations + .lock() + .await + .remove(&session_id); + self.active_turn_ids.lock().await.remove(&session_id); + self.active_turn_connections + .lock() + .await + .remove(&session_id); + match &result { + Ok(()) => { + self.run_session_hook_for_actor_state( + state, + session_id, + devo_core::HookEvent::Stop, + serde_json::Map::from_iter([( + "stop_hook_active".to_string(), + serde_json::Value::Bool(false), + )]), + ) + .await; + } + Err(error) => { + self.run_session_hook_for_actor_state( + state, + session_id, + devo_core::HookEvent::StopFailure, + serde_json::Map::from_iter([ + ( + "error".to_string(), + serde_json::Value::String(error.to_string()), + ), + ( + "error_details".to_string(), + serde_json::Value::String(error.to_string()), + ), + ]), + ) + .await; + } + } + + let final_turn = self + .persist_terminal_turn_state( + state, + session_id, + &turn, + &result, + latest_usage.clone(), + terminal_stop_reason, + session_total_input_tokens, + session_total_output_tokens, + session_total_tokens, + session_total_cache_creation_tokens, + session_total_cache_read_tokens, + session_last_input_tokens, + session_prompt_token_estimate, + ) + .await; + self.clear_btw_input_queue(state, session_id).await; + self.append_terminal_turn_record(state, session_id, &final_turn) + .await; + self.finalize_turn_workspace_changes(session_id, &final_turn) + .await; + self.emit_terminal_turn_events(state, session_id, &final_turn, &result) + .await; + self.record_terminal_turn_status( + final_turn.turn_id, + super::super::TerminalTurnSnapshot::from_turn(&final_turn), + ) + .await; + } + + #[allow(clippy::too_many_arguments)] + async fn persist_terminal_turn_state( + self: &Arc, + state: &mut SessionActorState, + session_id: SessionId, + turn: &crate::TurnMetadata, + result: &Result<(), devo_core::AgentError>, + latest_usage: Option, + terminal_stop_reason: Option, + session_total_input_tokens: usize, + session_total_output_tokens: usize, + session_total_tokens: usize, + session_total_cache_creation_tokens: usize, + session_total_cache_read_tokens: usize, + session_last_input_tokens: usize, + session_prompt_token_estimate: usize, + ) -> crate::TurnMetadata { + let mut final_turn = turn.clone(); + final_turn.completed_at = Some(Utc::now()); + final_turn.status = match result { + Ok(()) => TurnStatus::Completed, + Err(devo_core::AgentError::Aborted) => TurnStatus::Interrupted, + Err(_) => TurnStatus::Failed, + }; + final_turn.usage = latest_usage.clone(); + final_turn.stop_reason = terminal_stop_reason; + final_turn.failure_reason = result + .as_ref() + .err() + .and_then(turn_failure_reason_from_error); + state.latest_turn = Some(final_turn.clone()); + state.active_turn = None; + state.summary.status = SessionRuntimeStatus::Idle; + state.summary.updated_at = Utc::now(); + state.summary.last_activity_at = state.summary.updated_at; + state.summary.total_input_tokens = session_total_input_tokens; + state.summary.total_output_tokens = session_total_output_tokens; + state.summary.total_tokens = session_total_tokens; + state.summary.total_cache_creation_tokens = session_total_cache_creation_tokens; + state.summary.total_cache_read_tokens = session_total_cache_read_tokens; + state.summary.prompt_token_estimate = session_prompt_token_estimate; + if let Some(usage) = &final_turn.usage { + state.summary.last_query_total_tokens = usage.display_total_tokens(); + } + state.core.total_input_tokens = session_total_input_tokens; + state.core.total_output_tokens = session_total_output_tokens; + state.core.total_tokens = session_total_tokens; + state.core.total_cache_creation_tokens = session_total_cache_creation_tokens; + state.core.total_cache_read_tokens = session_total_cache_read_tokens; + if !state.summary.ephemeral { + let stats = SessionStats { + total_input_tokens: session_total_input_tokens, + total_output_tokens: session_total_output_tokens, + total_tokens: session_total_tokens, + total_cache_creation_tokens: session_total_cache_creation_tokens, + total_cache_read_tokens: session_total_cache_read_tokens, + last_input_tokens: final_turn + .usage + .as_ref() + .map(|usage| usage.input_tokens as usize) + .unwrap_or(session_last_input_tokens), + turn_count: state.summary.updated_at.timestamp() as usize, + prompt_token_estimate: session_prompt_token_estimate, + }; + if let Err(err) = self.deps.db.update_stats(&session_id, &stats) { + tracing::warn!( + session_id = %session_id, + error = %err, + "failed to persist token stats to database" + ); + } + } + final_turn + } + + async fn clear_btw_input_queue( + self: &Arc, + state: &SessionActorState, + session_id: SessionId, + ) { + let is_ephemeral = state.summary.ephemeral; + let btw_input_queue = Arc::clone(&state.btw_input_queue); + btw_input_queue + .lock() + .expect("btw input queue mutex should not be poisoned") + .clear(); + if !is_ephemeral && let Err(err) = self.deps.db.clear_pending(&session_id, QueueType::Btw) { + tracing::warn!( + session_id = %session_id, + error = %err, + "failed to clear btw input messages from database" + ); + } + } + + async fn append_terminal_turn_record( + self: &Arc, + state: &SessionActorState, + session_id: SessionId, + final_turn: &crate::TurnMetadata, + ) { + let record = state.record.clone(); + let session_context = state.core.session_context.clone(); + let turn_context = state.core.latest_turn_context.clone(); + if let Some(record) = record + && let Err(error) = self.rollout_store.append_turn( + &record, + build_turn_record(final_turn, session_context, turn_context), + ) + { + tracing::warn!(session_id = %session_id, error = %error, "failed to persist terminal turn line"); + } + } + + async fn emit_terminal_turn_events( + self: &Arc, + state: &SessionActorState, + session_id: SessionId, + final_turn: &crate::TurnMetadata, + result: &Result<(), devo_core::AgentError>, + ) { + if let Err(error) = result { + if matches!(error, devo_core::AgentError::Aborted) { + tracing::info!( + session_id = %session_id, + turn_id = %final_turn.turn_id, + status = ?final_turn.status, + "turn execution interrupted" + ); + self.broadcast_event(crate::ServerEvent::TurnInterrupted(TurnEventPayload { + session_id, + turn: final_turn.clone(), + })) + .await; + } else { + tracing::warn!( + session_id = %session_id, + turn_id = %final_turn.turn_id, + status = ?final_turn.status, + error = %error, + "turn execution failed" + ); + self.emit_turn_item( + session_id, + final_turn.turn_id, + ItemKind::AgentMessage, + TurnItem::AgentMessage(TextItem { + text: error.to_string(), + }), + serde_json::json!({ "title": "Error", "text": error.to_string() }), + ) + .await; + self.broadcast_event(crate::ServerEvent::TurnFailed(TurnEventPayload { + session_id, + turn: final_turn.clone(), + })) + .await; + } + } else { + tracing::info!( + session_id = %session_id, + turn_id = %final_turn.turn_id, + status = ?final_turn.status, + total_input_tokens = final_turn.usage.as_ref().map(|usage| usage.input_tokens), + total_output_tokens = final_turn.usage.as_ref().map(|usage| usage.output_tokens), + "turn execution completed" + ); + } + self.handle_subagent_turn_completed_for_actor_state(state, session_id, final_turn) + .await; + self.broadcast_event(crate::ServerEvent::TurnCompleted(TurnEventPayload { + session_id, + turn: final_turn.clone(), + })) + .await; + self.broadcast_event(crate::ServerEvent::SessionStatusChanged( + SessionStatusChangedPayload { + session_id, + status: SessionRuntimeStatus::Idle, + }, + )) + .await; + } +} diff --git a/crates/server/src/runtime/turn_exec/followup.rs b/crates/server/src/runtime/turn_exec/followup.rs new file mode 100644 index 00000000..3765928b --- /dev/null +++ b/crates/server/src/runtime/turn_exec/followup.rs @@ -0,0 +1,192 @@ +use std::sync::Arc; + +use chrono::Utc; +use devo_core::{SessionId, TurnId, TurnStatus}; + +use super::super::*; +use super::types::{ExecuteTurnRequest, QueuedTurnInput}; + +impl ServerRuntime { + /// Pop the first queued input and start a new turn in a background task. + /// Used from the interrupt handler where the calling function must return + /// its response immediately. + pub(in crate::runtime) async fn spawn_next_turn_from_queue( + self: &Arc, + session_id: SessionId, + ) { + let Some(queued) = self + .pop_next_queued_turn_input(session_id, /*require_idle_session*/ false) + .await + else { + return; + }; + self.broadcast_updated_queue(session_id).await; + let Some((turn, turn_config)) = self.prepare_queued_turn_start(session_id, &queued).await + else { + return; + }; + if let Some((parent_session_id, parent_turn_id)) = queued.subagent_usage_owner { + self.register_subagent_usage_owner(parent_session_id, session_id, parent_turn_id) + .await; + } + self.activate_queued_turn(session_id, &turn, &turn_config) + .await; + self.broadcast_event(crate::ServerEvent::TurnStarted(TurnEventPayload { + session_id, + turn: turn.clone(), + })) + .await; + let runtime = Arc::clone(self); + tokio::spawn(async move { + runtime + .execute_turn(ExecuteTurnRequest { + session_id, + turn, + turn_config, + display_input: queued.display_input, + input: queued.input_text, + input_messages: queued.input_messages, + collaboration_mode: queued.collaboration_mode, + input_mode: TurnInputMode::VisibleUserMessage, + }) + .await; + }); + } + + /// After a turn completes, chain directly into the next queued input when present. + pub(crate) async fn chain_queued_followup_turn( + self: &Arc, + session_id: SessionId, + ) -> bool { + let Some(queued) = self + .pop_next_queued_turn_input(session_id, /*require_idle_session*/ true) + .await + else { + return false; + }; + self.broadcast_updated_queue(session_id).await; + let Some((turn, turn_config)) = self.prepare_queued_turn_start(session_id, &queued).await + else { + return false; + }; + if let Some((parent_session_id, parent_turn_id)) = queued.subagent_usage_owner { + self.register_subagent_usage_owner(parent_session_id, session_id, parent_turn_id) + .await; + } + self.activate_queued_turn(session_id, &turn, &turn_config) + .await; + self.broadcast_event(crate::ServerEvent::TurnStarted(TurnEventPayload { + session_id, + turn: turn.clone(), + })) + .await; + Box::pin(Arc::clone(self).execute_turn(ExecuteTurnRequest { + session_id, + turn, + turn_config, + display_input: queued.display_input, + input: queued.input_text, + input_messages: queued.input_messages, + collaboration_mode: queued.collaboration_mode, + input_mode: TurnInputMode::VisibleUserMessage, + })) + .await; + true + } + + /// Read the current steering queue and broadcast its state to connected clients. + pub(in crate::runtime) async fn broadcast_updated_queue(&self, session_id: SessionId) { + let Some(session_handle) = self.session(session_id).await else { + return; + }; + let Some(snapshot) = session_handle.pending_queue_snapshot().await else { + return; + }; + self.broadcast_event(crate::ServerEvent::InputQueueUpdated( + devo_core::InputQueueUpdatedPayload { + session_id, + pending_count: snapshot.pending_count, + pending_texts: snapshot.pending_texts, + }, + )) + .await; + } + + async fn pop_next_queued_turn_input( + self: &Arc, + session_id: SessionId, + require_idle_session: bool, + ) -> Option { + let session_handle = self.sessions.lock().await.get(&session_id).cloned()?; + let popped = session_handle + .pop_queued_turn_input(require_idle_session) + .await + .flatten()?; + Some(QueuedTurnInput { + display_input: popped.display_input, + input_text: popped.input_text, + input_messages: popped.input_messages, + collaboration_mode: popped.collaboration_mode, + model_selection: popped.model_selection, + subagent_usage_owner: popped.subagent_usage_owner, + }) + } + + async fn prepare_queued_turn_start( + self: &Arc, + session_id: SessionId, + queued: &QueuedTurnInput, + ) -> Option<(crate::TurnMetadata, devo_core::TurnConfig)> { + let session_handle = self.sessions.lock().await.get(&session_id).cloned()?; + let reservation = session_handle.turn_reservation_snapshot().await?; + let model_override = queued + .model_selection + .as_deref() + .or_else(|| session_model_selection(&reservation.summary)); + let turn_config = reservation.runtime_context.resolve_turn_config( + model_override, + reservation.summary.reasoning_effort_selection.clone(), + ); + let resolved_request = turn_config + .model + .resolve_reasoning_effort_selection(turn_config.reasoning_effort_selection.as_deref()); + let request_model = turn_config.provider_request_model(&resolved_request.request_model); + let sequence = reservation + .latest_turn + .as_ref() + .map_or(1, |turn| turn.sequence + 1); + let now = Utc::now(); + let turn = crate::TurnMetadata { + turn_id: TurnId::new(), + session_id, + sequence, + status: TurnStatus::Running, + kind: devo_core::TurnKind::Regular, + model: turn_config.model.slug.clone(), + model_binding_id: turn_config.model_binding_id.clone(), + reasoning_effort_selection: turn_config.reasoning_effort_selection.clone(), + reasoning_effort: resolved_request.effective_reasoning_effort, + request_model, + request_thinking: resolved_request.request_thinking.clone(), + started_at: now, + completed_at: None, + usage: None, + stop_reason: None, + failure_reason: None, + }; + Some((turn, turn_config)) + } + + async fn activate_queued_turn( + self: &Arc, + session_id: SessionId, + turn: &crate::TurnMetadata, + turn_config: &devo_core::TurnConfig, + ) { + if let Some(session_handle) = self.sessions.lock().await.get(&session_id).cloned() { + session_handle + .activate_queued_turn(turn.clone(), turn_config.clone()) + .await; + } + } +} diff --git a/crates/server/src/runtime/turn_exec/item_stream.rs b/crates/server/src/runtime/turn_exec/item_stream.rs new file mode 100644 index 00000000..57c5c0fb --- /dev/null +++ b/crates/server/src/runtime/turn_exec/item_stream.rs @@ -0,0 +1,244 @@ +use std::sync::Arc; + +use devo_core::{ItemId, SessionId, TextItem, TurnId, TurnItem}; + +use super::super::ServerRuntime; +use super::super::proposed_plan::ProposedPlanSegment; +use crate::runtime::session_actor::state::SessionStreamState; +use crate::{ItemDeltaKind, ItemDeltaPayload, ItemKind, ServerEvent}; + +pub(super) async fn complete_reasoning_item( + runtime: &Arc, + session_id: SessionId, + turn_id: TurnId, + item_id: ItemId, + item_seq: u64, + text: String, +) { + runtime + .complete_item( + session_id, + turn_id, + item_id, + item_seq, + ItemKind::Reasoning, + TurnItem::Reasoning(TextItem { text: text.clone() }), + serde_json::json!({ "title": "Reasoning", "text": text }), + ) + .await; +} + +pub(super) async fn complete_assistant_item( + runtime: &Arc, + session_id: SessionId, + turn_id: TurnId, + item_id: ItemId, + item_seq: u64, + text: String, +) { + runtime + .complete_item( + session_id, + turn_id, + item_id, + item_seq, + ItemKind::AgentMessage, + TurnItem::AgentMessage(TextItem { text: text.clone() }), + serde_json::json!({ "title": "Assistant", "text": text }), + ) + .await; +} + +#[derive(Debug, Default)] +pub(super) struct ProposedPlanStreamItem { + item_id: Option, + item_seq: Option, + text: String, +} + +impl ProposedPlanStreamItem { + async fn start( + &mut self, + runtime: &Arc, + session_id: SessionId, + turn_id: TurnId, + ) { + if self.item_id.is_some() && self.item_seq.is_some() { + return; + } + let (item_id, item_seq) = runtime + .start_item( + session_id, + turn_id, + ItemKind::Plan, + serde_json::json!({ "title": "Proposed Plan", "text": "" }), + ) + .await; + self.item_id = Some(item_id); + self.item_seq = Some(item_seq); + } + + async fn push_delta( + &mut self, + runtime: &Arc, + session_id: SessionId, + turn_id: TurnId, + delta: String, + ) { + if delta.is_empty() { + return; + } + self.start(runtime, session_id, turn_id).await; + self.text.push_str(&delta); + runtime + .broadcast_event(ServerEvent::ItemDelta { + delta_kind: ItemDeltaKind::PlanDelta, + payload: ItemDeltaPayload { + context: crate::EventContext { + session_id, + turn_id: Some(turn_id), + item_id: self.item_id, + seq: 0, + }, + delta, + stream_index: None, + channel: None, + }, + }) + .await; + } + + pub(super) async fn complete( + &mut self, + runtime: &Arc, + session_id: SessionId, + turn_id: TurnId, + ) { + let (Some(item_id), Some(item_seq)) = (self.item_id.take(), self.item_seq.take()) else { + return; + }; + let text = std::mem::take(&mut self.text); + runtime + .complete_item( + session_id, + turn_id, + item_id, + item_seq, + ItemKind::Plan, + TurnItem::Plan(TextItem { text: text.clone() }), + serde_json::json!({ "title": "Proposed Plan", "text": text }), + ) + .await; + } +} + +#[allow(clippy::too_many_arguments)] +pub(super) async fn push_assistant_text_delta( + runtime: &Arc, + event_stream: &Arc>, + session_id: SessionId, + turn_id: TurnId, + assistant_item_id: &mut Option, + assistant_item_seq: &mut Option, + assistant_text: &mut String, + assistant_delta_seq: &mut u64, + text: String, +) { + if text.is_empty() { + return; + } + let (item_id, item_seq) = match (*assistant_item_id, *assistant_item_seq) { + (Some(item_id), Some(item_seq)) => (item_id, item_seq), + (None, None) => { + let (item_id, item_seq) = runtime + .start_item( + session_id, + turn_id, + ItemKind::AgentMessage, + serde_json::json!({ "title": "Assistant", "text": "" }), + ) + .await; + *assistant_item_id = Some(item_id); + *assistant_item_seq = Some(item_seq); + (item_id, item_seq) + } + _ => return, + }; + assistant_text.push_str(&text); + *assistant_delta_seq = (*assistant_delta_seq).saturating_add(1); + runtime + .broadcast_event(ServerEvent::ItemDelta { + delta_kind: ItemDeltaKind::AgentMessageDelta, + payload: ItemDeltaPayload { + context: crate::EventContext { + session_id, + turn_id: Some(turn_id), + item_id: Some(item_id), + seq: 0, + }, + delta: text, + stream_index: None, + channel: None, + }, + }) + .await; + if let Ok(mut stream) = event_stream.try_lock() { + stream.deferred_assistant = Some((item_id, item_seq, assistant_text.clone())); + } +} + +#[allow(clippy::too_many_arguments)] +pub(super) async fn handle_proposed_plan_segments( + runtime: &Arc, + event_stream: &Arc>, + session_id: SessionId, + turn_id: TurnId, + segments: Vec, + assistant_item_id: &mut Option, + assistant_item_seq: &mut Option, + assistant_text: &mut String, + assistant_delta_seq: &mut u64, + proposed_plan_item: &mut ProposedPlanStreamItem, + leading_normal_buffer: &mut String, +) { + for segment in segments { + match segment { + ProposedPlanSegment::Normal(delta) => { + if delta.is_empty() { + continue; + } + if assistant_item_id.is_none() && delta.chars().all(char::is_whitespace) { + leading_normal_buffer.push_str(&delta); + continue; + } + let delta = if assistant_item_id.is_none() && !leading_normal_buffer.is_empty() { + format!("{}{}", std::mem::take(leading_normal_buffer), delta) + } else { + delta + }; + push_assistant_text_delta( + runtime, + event_stream, + session_id, + turn_id, + assistant_item_id, + assistant_item_seq, + assistant_text, + assistant_delta_seq, + delta, + ) + .await; + } + ProposedPlanSegment::PlanStart => { + leading_normal_buffer.clear(); + proposed_plan_item.start(runtime, session_id, turn_id).await; + } + ProposedPlanSegment::PlanDelta(delta) => { + proposed_plan_item + .push_delta(runtime, session_id, turn_id, delta) + .await; + } + ProposedPlanSegment::PlanEnd => {} + } + } +} diff --git a/crates/server/src/runtime/turn_exec/mod.rs b/crates/server/src/runtime/turn_exec/mod.rs new file mode 100644 index 00000000..b5796be4 --- /dev/null +++ b/crates/server/src/runtime/turn_exec/mod.rs @@ -0,0 +1,70 @@ +mod event_stream; +mod finalize; +mod followup; +mod item_stream; +mod query; +mod shell; +mod tool_display; +mod tool_results; +mod trace; +mod types; + +pub(crate) use event_stream::{QUERY_EVENT_CHANNEL_CAPACITY, spawn_turn_event_stream}; +pub(crate) use finalize::FinalizeTurnParams; +pub(crate) use query::TurnModelQueryParams; +pub(crate) use types::ExecuteTurnRequest; + +use std::sync::Arc; + +use super::*; + +impl ServerRuntime { + /// Execute one turn end-to-end via the session actor. + pub(in crate::runtime) async fn execute_turn(self: Arc, request: ExecuteTurnRequest) { + let Some(handle) = self.session(request.session_id).await else { + return; + }; + handle.execute_turn(Arc::clone(&self), request).await; + } + + pub(super) async fn prepare_turn_execution_for_actor( + self: &Arc, + state: &mut SessionActorState, + turn: &crate::TurnMetadata, + display_input: &str, + emits_user_message: bool, + ) { + self.capture_turn_workspace_baseline( + state.session_id(), + turn.turn_id, + state.summary.cwd.clone(), + ) + .await; + state.turn_approval_cache = crate::execution::ApprovalGrantCache::default(); + if emits_user_message { + self.emit_turn_item( + state.session_id(), + turn.turn_id, + crate::ItemKind::UserMessage, + devo_core::TurnItem::UserMessage(devo_core::TextItem { + text: display_input.to_string(), + }), + serde_json::json!({ "title": "You", "text": display_input }), + ) + .await; + } + } + + pub(in crate::runtime) fn tool_registry_for_actor_state( + &self, + state: &SessionActorState, + ) -> Arc { + state + .tool_registry + .clone() + .unwrap_or_else(|| Arc::clone(&state.runtime_context.registry)) + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/server/src/runtime/turn_exec/query.rs b/crates/server/src/runtime/turn_exec/query.rs new file mode 100644 index 00000000..a459c5c1 --- /dev/null +++ b/crates/server/src/runtime/turn_exec/query.rs @@ -0,0 +1,197 @@ +use std::sync::Arc; + +use devo_core::tools::{ + AgentToolCoordinator, ClientFilesystem, ClientTerminal, ToolAgentScope, ToolCall, + ToolExecutionOptions, ToolRuntime, ToolRuntimeContext, +}; +use devo_core::{Message, QueryEvent, TurnConfig, query}; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use super::super::*; +use super::event_stream::enqueue_query_event; +use super::tool_display::without_agent_coordination_tools; +use super::types::TurnQueryOutcome; + +pub(crate) struct TurnModelQueryParams<'a> { + pub state: &'a mut SessionActorState, + pub turn_id: devo_core::TurnId, + pub turn_config: &'a TurnConfig, + pub input: &'a str, + pub input_messages: &'a [String], + pub collaboration_mode: devo_protocol::CollaborationMode, + pub input_mode: super::super::TurnInputMode, + pub usage_parent_session_id: Option, + pub event_tx: mpsc::Sender, +} + +impl ServerRuntime { + pub(crate) async fn run_turn_model_query( + self: &Arc, + params: TurnModelQueryParams<'_>, + ) -> TurnQueryOutcome { + let TurnModelQueryParams { + state, + turn_id, + turn_config, + input, + input_messages, + collaboration_mode, + input_mode, + usage_parent_session_id, + event_tx, + } = params; + let session_id = state.session_id(); + let agent_scope = if state.summary.parent_session_id.is_some() { + ToolAgentScope::Subagent + } else { + ToolAgentScope::Parent + }; + let agent_tool_policy = state.agent_tool_policy; + let session_tool_registry = self.tool_registry_for_actor_state(state); + let runtime_context = Arc::clone(&state.runtime_context); + let turn_goal = match &input_mode { + super::super::TurnInputMode::VisibleUserMessage => { + let stores = self.goal_stores.lock().await; + stores + .get(&session_id) + .and_then(crate::GoalStore::get) + .map(crate::goal::Goal::to_thread_goal) + } + super::super::TurnInputMode::HiddenGoalContinuation { goal } => Some(goal.clone()), + }; + state.core.config.token_budget = turn_config.token_budget(); + state.core.collaboration_mode = collaboration_mode; + if let Some(goal) = turn_goal { + state.core.set_active_goal(goal); + } else { + state.core.clear_active_goal(); + } + if input_mode.emits_user_message() && input_messages.is_empty() { + state.core.push_message(Message::user(input.to_string())); + } else if input_mode.emits_user_message() { + for input_message in input_messages { + state + .core + .push_message(Message::user(input_message.clone())); + } + } + let event_callback_tx = event_tx.clone(); + let callback: devo_core::EventCallback = std::sync::Arc::new(move |event: QueryEvent| { + let event_callback_tx = event_callback_tx.clone(); + Box::pin(async move { + enqueue_query_event(&event_callback_tx, event); + }) + }); + let tool_execution_start_tx = event_tx.clone(); + let agent_context_mode = state + .core + .session_context + .as_ref() + .map(|context| match context.system_prompt_mode { + devo_core::SystemPromptMode::CodingAgent => { + devo_protocol::AgentContextMode::CodingAgent + } + devo_core::SystemPromptMode::DeepResearch => { + devo_protocol::AgentContextMode::DeepResearch + } + }) + .unwrap_or_default(); + let registry = match agent_tool_policy { + devo_protocol::AgentToolPolicy::Inherit if usage_parent_session_id.is_some() => { + Arc::new(without_agent_coordination_tools(&session_tool_registry)) + } + devo_protocol::AgentToolPolicy::Inherit => session_tool_registry, + devo_protocol::AgentToolPolicy::DenyAll => { + Arc::new(devo_core::tools::ToolRegistry::new()) + } + devo_protocol::AgentToolPolicy::DeepResearch => Arc::new( + runtime_context + .registry + .restricted_to_specs(super::super::research::RESEARCH_WORKER_TOOL_NAMES), + ), + }; + let permission_mode = state.core.config.permission_mode; + let permission_profile = state.core.config.permission_profile.clone(); + let hook_context = Self::hook_context_from_actor_state(state, session_id); + let turn_cancel_token = self + .active_turn_cancellations + .lock() + .await + .get(&session_id) + .cloned() + .unwrap_or_else(CancellationToken::new); + let query_cancel_token = turn_cancel_token.clone(); + let provider_http = runtime_context + .config_store + .lock() + .expect("app config store mutex should not be poisoned") + .effective_config() + .provider_http + .clone(); + let runtime = ToolRuntime::new_with_context_and_options( + Arc::clone(®istry), + self.build_permission_checker(session_id, turn_id, permission_mode, permission_profile), + ToolRuntimeContext { + session_id: session_id.to_string(), + turn_id: Some(turn_id.to_string()), + cwd: state.core.cwd.clone(), + agent_scope, + agent_context_mode, + collaboration_mode, + agent_coordinator: Some(Arc::clone(self) as Arc), + client_filesystem: Some(Arc::clone(self) as Arc), + client_terminal: Some(Arc::clone(self) as Arc), + local_web_search: match &turn_config.web_search { + devo_core::ResolvedWebSearchConfig::Local(config) => Some(config.clone()), + devo_core::ResolvedWebSearchConfig::Disabled + | devo_core::ResolvedWebSearchConfig::Provider => None, + }, + hooks: hook_context, + network_proxy: provider_http.proxy_url, + network_no_proxy: provider_http.no_proxy, + }, + ToolExecutionOptions { + cancel_token: turn_cancel_token, + on_tool_execution_start: Some(Arc::new(move |call: ToolCall| { + let tool_execution_start_tx = tool_execution_start_tx.clone(); + Box::pin(async move { + enqueue_query_event( + &tool_execution_start_tx, + QueryEvent::ToolExecutionStart { id: call.id }, + ); + }) + })), + ..ToolExecutionOptions::default() + }, + ); + // Turns execute inline on the session actor's own task rather than as a + // separately spawned task, so an external `JoinHandle::abort()` can no + // longer stop an in-flight query: it only cancels the caller waiting on + // the actor's reply, not the actor itself. Race the query against the + // turn's cancellation token so interrupting a turn actually unblocks the + // actor's mailbox instead of hanging it forever. + let result = tokio::select! { + biased; + () = query_cancel_token.cancelled() => Err(devo_core::AgentError::Aborted), + result = query( + &mut state.core, + turn_config, + runtime_context.provider_for_route(turn_config.provider_route.clone()), + registry, + &runtime, + Some(callback), + ) => result, + }; + TurnQueryOutcome { + result, + session_total_input_tokens: state.core.total_input_tokens, + session_total_output_tokens: state.core.total_output_tokens, + session_total_tokens: state.core.total_tokens, + session_total_cache_creation_tokens: state.core.total_cache_creation_tokens, + session_total_cache_read_tokens: state.core.total_cache_read_tokens, + session_last_input_tokens: state.core.last_input_tokens, + session_prompt_token_estimate: state.core.prompt_token_estimate, + } + } +} diff --git a/crates/server/src/runtime/turn_exec/shell.rs b/crates/server/src/runtime/turn_exec/shell.rs new file mode 100644 index 00000000..2f54ddd5 --- /dev/null +++ b/crates/server/src/runtime/turn_exec/shell.rs @@ -0,0 +1,211 @@ +use std::sync::Arc; + +use devo_core::tools::{ + ToolAgentScope, ToolCall, ToolCallResult, ToolContent, ToolExecutionOptions, ToolRuntime, + ToolRuntimeContext, +}; +use devo_core::{CommandExecutionItem, SessionId, TurnItem}; +use devo_util_shell_command::parse_command::parse_command; +use tokio_util::sync::CancellationToken; + +use super::super::*; +use super::tool_display::{user_shell_command_payload, user_shell_exec_input}; +use crate::{ + ItemKind, ServerEvent, SessionRuntimeStatus, SessionStatusChangedPayload, TurnEventPayload, +}; + +impl ServerRuntime { + pub(in crate::runtime) async fn execute_shell_command_turn( + self: Arc, + session_id: SessionId, + turn: crate::TurnMetadata, + command: String, + cwd: std::path::PathBuf, + ) { + self.capture_turn_workspace_baseline(session_id, turn.turn_id, cwd.clone()) + .await; + if let Some(session_handle) = self.session(session_id).await { + session_handle.reset_turn_approval_cache().await; + } + + let tool_call_id = format!("user-shell-{}", turn.turn_id); + let input = user_shell_exec_input(&command, cwd.clone()); + let command_actions = parse_command(std::slice::from_ref(&command)); + let (item_id, item_seq) = self + .start_item( + session_id, + turn.turn_id, + ItemKind::CommandExecution, + serde_json::to_value(user_shell_command_payload( + &tool_call_id, + &command, + input.clone(), + command_actions.clone(), + None, + false, + )) + .expect("serialize command execution payload"), + ) + .await; + + let Some(session_handle) = self.session(session_id).await else { + self.clear_active_turn_runtime_handles(session_id).await; + return; + }; + let Some(shell_context) = session_handle.shell_exec_context(cwd.clone()).await else { + self.clear_active_turn_runtime_handles(session_id).await; + return; + }; + let permission_mode = shell_context.permission_mode; + let permission_profile = shell_context.permission_profile; + let registry = shell_context.tool_registry; + let provider_http = shell_context + .runtime_context + .config_store + .lock() + .expect("app config store mutex should not be poisoned") + .effective_config() + .provider_http + .clone(); + let turn_cancel_token = self + .active_turn_cancellations + .lock() + .await + .get(&session_id) + .cloned() + .unwrap_or_else(CancellationToken::new); + let tool_execution_start_runtime = Arc::clone(&self); + let tool_execution_start_session_id = session_id; + let tool_execution_start_turn_id = turn.turn_id; + let runtime = ToolRuntime::new_with_context_and_options( + registry, + self.build_permission_checker( + session_id, + turn.turn_id, + permission_mode, + permission_profile, + ), + ToolRuntimeContext { + session_id: session_id.to_string(), + turn_id: Some(turn.turn_id.to_string()), + cwd, + agent_scope: ToolAgentScope::Parent, + agent_context_mode: devo_protocol::AgentContextMode::CodingAgent, + collaboration_mode: devo_protocol::CollaborationMode::Build, + agent_coordinator: None, + client_filesystem: None, + client_terminal: None, + local_web_search: None, + hooks: self.hook_context_for_session(session_id).await, + network_proxy: provider_http.proxy_url, + network_no_proxy: provider_http.no_proxy, + }, + ToolExecutionOptions { + cancel_token: turn_cancel_token, + on_tool_execution_start: Some(Arc::new(move |call: ToolCall| { + let runtime = Arc::clone(&tool_execution_start_runtime); + let tool_call_id = call.id; + Box::pin(async move { + runtime + .broadcast_event(ServerEvent::ToolCallStatusUpdated( + devo_protocol::ToolCallStatusUpdatedPayload { + session_id: tool_execution_start_session_id, + turn_id: tool_execution_start_turn_id, + tool_call_id, + status: "in_progress".to_string(), + terminal_id: None, + }, + )) + .await; + }) + })), + ..ToolExecutionOptions::default() + }, + ); + let result = runtime + .execute_batch(&[ToolCall { + id: tool_call_id.clone(), + name: "exec_command".to_string(), + input: input.clone(), + }]) + .await + .into_iter() + .next() + .unwrap_or_else(|| ToolCallResult::error(&tool_call_id, "shell command did not run")); + let output = match result.content.clone() { + ToolContent::Text(text) => serde_json::Value::String(text), + ToolContent::Json(json) => json, + ToolContent::Mixed { text, json } => { + json.unwrap_or_else(|| serde_json::Value::String(text.unwrap_or_default())) + } + }; + let is_error = result.is_error; + self.clear_active_turn_runtime_handles(session_id).await; + self.complete_item( + session_id, + turn.turn_id, + item_id, + item_seq, + ItemKind::CommandExecution, + TurnItem::CommandExecution(CommandExecutionItem { + tool_call_id: tool_call_id.clone(), + tool_name: "exec_command".to_string(), + command: command.clone(), + input: input.clone(), + output: output.clone(), + is_error, + }), + serde_json::to_value(user_shell_command_payload( + &tool_call_id, + &command, + input.clone(), + command_actions, + Some(output), + is_error, + )) + .expect("serialize command execution payload"), + ) + .await; + + let Some(final_turn) = session_handle + .complete_shell_turn(turn.clone(), is_error) + .await + else { + return; + }; + if let Some(persistence) = session_handle.turn_persistence_snapshot().await + && let Some(record) = persistence.record + && let Err(error) = self.rollout_store.append_turn( + &record, + build_turn_record( + &final_turn, + persistence.session_context, + persistence.latest_turn_context, + ), + ) + { + tracing::warn!(session_id = %session_id, error = %error, "failed to persist shell command turn line"); + } + self.finalize_turn_workspace_changes(session_id, &final_turn) + .await; + if is_error { + self.broadcast_event(ServerEvent::TurnFailed(TurnEventPayload { + session_id, + turn: final_turn.clone(), + })) + .await; + } + self.broadcast_event(ServerEvent::TurnCompleted(TurnEventPayload { + session_id, + turn: final_turn, + })) + .await; + self.broadcast_event(ServerEvent::SessionStatusChanged( + SessionStatusChangedPayload { + session_id, + status: SessionRuntimeStatus::Idle, + }, + )) + .await; + } +} diff --git a/crates/server/src/runtime/turn_exec/tests.rs b/crates/server/src/runtime/turn_exec/tests.rs new file mode 100644 index 00000000..dc702aca --- /dev/null +++ b/crates/server/src/runtime/turn_exec/tests.rs @@ -0,0 +1,371 @@ +use super::tool_display::*; +use super::types::*; +use crate::*; +use devo_core::ItemId; +use devo_core::tools::tool_spec::ToolPreparationFeedback; +use pretty_assertions::assert_eq; +use std::collections::HashMap; + +#[test] +fn command_progress_uses_command_execution_item_id() { + let command_item_id = ItemId::new(); + let tool_item_id = ItemId::new(); + let mut pending_tool_calls = HashMap::new(); + pending_tool_calls.insert( + "exec".to_string(), + PendingToolCall { + item_id: Some(command_item_id), + item_seq: Some(1), + input: serde_json::json!({}), + display_kind: ToolDisplayKind::CommandExecution, + command: "cargo test".to_string(), + }, + ); + pending_tool_calls.insert( + "read".to_string(), + PendingToolCall { + item_id: Some(tool_item_id), + item_seq: Some(2), + input: serde_json::json!({}), + display_kind: ToolDisplayKind::Generic, + command: String::new(), + }, + ); + + assert_eq!( + command_execution_item_id_for_progress(&pending_tool_calls, "exec"), + Some(command_item_id) + ); + assert_eq!( + command_execution_item_id_for_progress(&pending_tool_calls, "read"), + None + ); + assert_eq!( + command_execution_item_id_for_progress(&pending_tool_calls, "missing"), + None + ); +} + +#[test] +fn file_change_tool_detection_matches_apply_patch_and_write() { + assert!(is_file_change_tool("apply_patch")); + assert!(is_file_change_tool("write")); + assert!(!is_file_change_tool("read")); +} + +#[test] +fn plan_tool_detection_matches_update_plan() { + assert!(is_plan_tool("update_plan")); + assert!(!is_plan_tool("read")); +} + +#[test] +fn read_tool_start_item_contains_live_read_action() { + let input = serde_json::json!({ + "path": "crates/tui/src/mod.rs" + }); + let start_item = tool_start_item_from_input( + "call-1", + "read", + "read crates/tui/src/mod.rs", + &input, + ToolDisplayKind::Generic, + ToolPreparationFeedback::None, + ); + + let payload: ToolCallPayload = + serde_json::from_value(start_item.payload).expect("tool call payload"); + + assert_eq!(start_item.item_kind, ItemKind::ToolCall); + assert_eq!( + payload, + ToolCallPayload { + tool_call_id: "call-1".to_string(), + tool_name: "read".to_string(), + parameters: input, + command_actions: vec![devo_protocol::parse_command::ParsedCommand::Read { + cmd: "read crates/tui/src/mod.rs".to_string(), + name: "mod.rs".to_string(), + path: std::path::PathBuf::from("crates/tui/src/mod.rs"), + }], + } + ); +} + +#[test] +fn grep_tool_start_item_contains_live_search_action() { + let input = serde_json::json!({ + "pattern": "ToolUseStart", + "path": "crates/server/src" + }); + let start_item = tool_start_item_from_input( + "call-1", + "grep", + "grep ToolUseStart in crates/server/src", + &input, + ToolDisplayKind::Generic, + ToolPreparationFeedback::None, + ); + + let payload: ToolCallPayload = + serde_json::from_value(start_item.payload).expect("tool call payload"); + + assert_eq!(start_item.item_kind, ItemKind::ToolCall); + assert_eq!( + payload, + ToolCallPayload { + tool_call_id: "call-1".to_string(), + tool_name: "grep".to_string(), + parameters: input, + command_actions: vec![devo_protocol::parse_command::ParsedCommand::Search { + cmd: "grep ToolUseStart in crates/server/src".to_string(), + query: Some("ToolUseStart".to_string()), + path: Some("crates/server/src".to_string()), + }], + } + ); +} + +#[test] +fn code_search_tool_start_item_contains_live_search_action() { + let input = serde_json::json!({ + "operation": "search", + "query": "live tool feedback", + "path": "crates" + }); + let start_item = tool_start_item_from_input( + "call-1", + "code_search", + "code_search live tool feedback in crates", + &input, + ToolDisplayKind::Generic, + ToolPreparationFeedback::None, + ); + + let payload: ToolCallPayload = + serde_json::from_value(start_item.payload).expect("tool call payload"); + + assert_eq!(start_item.item_kind, ItemKind::ToolCall); + assert_eq!( + payload, + ToolCallPayload { + tool_call_id: "call-1".to_string(), + tool_name: "code_search".to_string(), + parameters: input, + command_actions: vec![devo_protocol::parse_command::ParsedCommand::Search { + cmd: "code_search live tool feedback in crates".to_string(), + query: Some("live tool feedback".to_string()), + path: Some("crates".to_string()), + }], + } + ); +} + +#[test] +fn exec_tool_start_item_uses_command_execution_payload() { + let input = serde_json::json!({ + "cmd": "cargo test -p devo-server" + }); + let start_item = tool_start_item_from_input( + "call-1", + "exec_command", + "cargo test -p devo-server", + &input, + ToolDisplayKind::CommandExecution, + ToolPreparationFeedback::None, + ); + + let payload: CommandExecutionPayload = + serde_json::from_value(start_item.payload).expect("command execution payload"); + + assert_eq!(start_item.item_kind, ItemKind::CommandExecution); + assert_eq!( + payload, + CommandExecutionPayload { + tool_call_id: "call-1".to_string(), + tool_name: "exec_command".to_string(), + command: "cargo test -p devo-server".to_string(), + input: Some(input), + source: devo_protocol::protocol::ExecCommandSource::Agent, + command_actions: Vec::new(), + output: None, + is_error: false, + } + ); +} + +#[test] +fn user_shell_exec_input_uses_pty_backed_exec_command() { + let cwd = std::path::PathBuf::from("workspace"); + + let input = user_shell_exec_input("pwd", cwd.clone()); + + assert_eq!( + input, + serde_json::json!({ + "cmd": "pwd", + "workdir": cwd, + "login": true, + "tty": true, + }) + ); +} + +#[test] +fn user_shell_command_payload_uses_user_shell_source() { + let output = serde_json::json!({ "output": "done" }); + + let input = user_shell_exec_input("pwd", std::path::PathBuf::from("workspace")); + let payload = user_shell_command_payload( + "call-1", + "pwd", + input.clone(), + Vec::new(), + Some(output.clone()), + false, + ); + + assert_eq!( + payload, + CommandExecutionPayload { + tool_call_id: "call-1".to_string(), + tool_name: "exec_command".to_string(), + command: "pwd".to_string(), + input: Some(input), + source: devo_protocol::protocol::ExecCommandSource::UserShell, + command_actions: Vec::new(), + output: Some(output), + is_error: false, + } + ); +} + +#[test] +fn live_only_apply_patch_start_item_stays_tool_call() { + let input = serde_json::json!({ + "patch": "*** Begin Patch\n*** End Patch" + }); + let start_item = tool_start_item_from_input( + "call-1", + "apply_patch", + "apply_patch", + &input, + ToolDisplayKind::Generic, + ToolPreparationFeedback::LiveOnly, + ); + + let payload: ToolCallPayload = + serde_json::from_value(start_item.payload).expect("tool call payload"); + + assert_eq!(start_item.item_kind, ItemKind::ToolCall); + assert_eq!( + payload, + ToolCallPayload { + tool_call_id: "call-1".to_string(), + tool_name: "apply_patch".to_string(), + parameters: input, + command_actions: Vec::new(), + } + ); +} + +#[test] +fn command_actions_from_read_tool_input_builds_read_action() { + let actions = command_actions_from_tool_input( + "read", + "read crates/tui/src/mod.rs", + &serde_json::json!({ + "filePath": "crates/tui/src/mod.rs" + }), + ); + assert_eq!( + actions, + vec![devo_protocol::parse_command::ParsedCommand::Read { + cmd: "read crates/tui/src/mod.rs".to_string(), + name: "mod.rs".to_string(), + path: std::path::PathBuf::from("crates/tui/src/mod.rs"), + }] + ); +} + +#[test] +fn command_actions_from_read_tool_input_without_path_is_empty() { + let actions = + command_actions_from_tool_input("read", "read", &serde_json::json!({ "limit": 10 })); + assert_eq!(actions, Vec::new()); +} + +#[test] +fn command_actions_from_read_tool_result_summary_recovers_final_path() { + let actions = command_actions_from_tool_result( + "read", + "read ", + &serde_json::json!({}), + "read: crates/tui/src/mod.rs", + ); + assert_eq!( + actions, + vec![devo_protocol::parse_command::ParsedCommand::Read { + cmd: "read crates/tui/src/mod.rs".to_string(), + name: "mod.rs".to_string(), + path: std::path::PathBuf::from("crates/tui/src/mod.rs"), + }] + ); +} + +#[test] +fn command_actions_from_grep_tool_input_builds_search_action() { + let actions = command_actions_from_tool_input( + "grep", + "grep rebuild_restored_session in crates/tui/src", + &serde_json::json!({ + "pattern": "rebuild_restored_session", + "path": "crates/tui/src" + }), + ); + assert_eq!(actions.len(), 1); + assert!(matches!( + &actions[0], + devo_protocol::parse_command::ParsedCommand::Search { query, path, .. } + if query.as_deref() == Some("rebuild_restored_session") + && path.as_deref() == Some("crates/tui/src") + )); +} + +#[test] +fn command_actions_from_glob_tool_input_include_pattern_and_path() { + let actions = command_actions_from_tool_input( + "glob", + "glob **/Cargo.toml in crates", + &serde_json::json!({ + "pattern": "**/Cargo.toml", + "path": "crates" + }), + ); + assert_eq!( + actions, + vec![devo_protocol::parse_command::ParsedCommand::ListFiles { + cmd: "glob **/Cargo.toml in crates".to_string(), + path: Some("**/Cargo.toml in crates".to_string()), + }] + ); +} + +#[test] +fn command_actions_from_find_tool_input_include_pattern_and_path() { + let actions = command_actions_from_tool_input( + "find", + "find **/Cargo.toml in crates", + &serde_json::json!({ + "pattern": "**/Cargo.toml", + "path": "crates" + }), + ); + assert_eq!( + actions, + vec![devo_protocol::parse_command::ParsedCommand::ListFiles { + cmd: "find **/Cargo.toml in crates".to_string(), + path: Some("**/Cargo.toml in crates".to_string()), + }] + ); +} diff --git a/crates/server/src/runtime/turn_exec/tool_display.rs b/crates/server/src/runtime/turn_exec/tool_display.rs new file mode 100644 index 00000000..0c90de16 --- /dev/null +++ b/crates/server/src/runtime/turn_exec/tool_display.rs @@ -0,0 +1,404 @@ +use std::collections::HashMap; + +use devo_core::ItemId; +use devo_core::tools::tool_spec::ToolPreparationFeedback; + +use super::types::{PendingToolCall, ToolDisplayKind, ToolStartItem}; +use crate::{CommandExecutionPayload, FileChangePayload, ItemKind, ToolCallPayload}; + +pub(super) fn is_unified_exec_tool(name: &str) -> bool { + matches!(name, "exec_command" | "write_stdin") +} + +pub(super) fn is_file_change_tool(name: &str) -> bool { + matches!(name, "apply_patch" | "write") +} + +pub(super) fn is_plan_tool(name: &str) -> bool { + matches!(name, "update_plan") +} + +fn tool_start_item_kind( + tool_name: &str, + display_kind: ToolDisplayKind, + preparation_feedback: ToolPreparationFeedback, +) -> ItemKind { + if preparation_feedback == ToolPreparationFeedback::LiveOnly { + ItemKind::ToolCall + } else if is_file_change_tool(tool_name) { + ItemKind::FileChange + } else if display_kind.is_command_execution() { + ItemKind::CommandExecution + } else if is_plan_tool(tool_name) { + ItemKind::Plan + } else { + ItemKind::ToolCall + } +} + +fn tool_start_item( + tool_call_id: &str, + tool_name: &str, + command: &str, + input: &serde_json::Value, + display_kind: ToolDisplayKind, + preparation_feedback: ToolPreparationFeedback, + command_actions: Vec, +) -> ToolStartItem { + let item_kind = tool_start_item_kind(tool_name, display_kind, preparation_feedback); + let payload = match item_kind { + ItemKind::ToolCall => serde_json::to_value(ToolCallPayload { + tool_call_id: tool_call_id.to_string(), + tool_name: tool_name.to_string(), + parameters: input.clone(), + command_actions, + }) + .expect("serialize tool call payload"), + ItemKind::FileChange => serde_json::to_value(FileChangePayload { + tool_call_id: tool_call_id.to_string(), + tool_name: Some(tool_name.to_string()), + input: Some(input.clone()), + changes: Vec::new(), + is_error: false, + }) + .expect("serialize file change payload"), + ItemKind::CommandExecution => serde_json::to_value(CommandExecutionPayload { + tool_call_id: tool_call_id.to_string(), + tool_name: tool_name.to_string(), + command: command.to_string(), + input: Some(input.clone()), + source: devo_protocol::protocol::ExecCommandSource::Agent, + command_actions, + output: None, + is_error: false, + }) + .expect("serialize command execution payload"), + ItemKind::Plan => serde_json::json!({ + "title": "Plan", + "text": "" + }), + ItemKind::UserMessage + | ItemKind::AgentMessage + | ItemKind::Reasoning + | ItemKind::ToolResult + | ItemKind::McpToolCall + | ItemKind::WebSearch + | ItemKind::ImageView + | ItemKind::ContextCompaction + | ItemKind::ApprovalRequest + | ItemKind::ApprovalDecision + | ItemKind::ResearchArtifact => unreachable!("tool start item kind must be tool-like"), + }; + ToolStartItem { item_kind, payload } +} + +pub(super) fn tool_start_item_from_input( + tool_call_id: &str, + tool_name: &str, + command: &str, + input: &serde_json::Value, + display_kind: ToolDisplayKind, + preparation_feedback: ToolPreparationFeedback, +) -> ToolStartItem { + tool_start_item( + tool_call_id, + tool_name, + command, + input, + display_kind, + preparation_feedback, + command_actions_from_tool_input(tool_name, command, input), + ) +} + +pub(super) fn tool_start_item_from_result( + tool_call_id: &str, + tool_name: &str, + command: &str, + input: &serde_json::Value, + display_kind: ToolDisplayKind, + preparation_feedback: ToolPreparationFeedback, + summary: &str, +) -> ToolStartItem { + tool_start_item( + tool_call_id, + tool_name, + command, + input, + display_kind, + preparation_feedback, + command_actions_from_tool_result(tool_name, command, input, summary), + ) +} + +pub(super) fn command_display_from_input(tool_name: &str, input: &serde_json::Value) -> String { + match tool_name { + "exec_command" => input + .get("cmd") + .or_else(|| input.get("command")) + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_string(), + "write_stdin" => { + let session_id = input + .get("session_id") + .and_then(serde_json::Value::as_i64) + .map(|id| id.to_string()) + .unwrap_or_else(|| "?".to_string()); + let chars = input + .get("chars") + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + if chars.is_empty() { + format!("poll session {session_id}") + } else { + format!("write_stdin session {session_id}") + } + } + "read" => { + let path = input + .get("filePath") + .or_else(|| input.get("path")) + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + format!("read {path}") + } + "find" | "glob" => { + let pattern = input + .get("pattern") + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + let path = input + .get("path") + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + let command_name = if tool_name == "find" { "find" } else { "glob" }; + if path.is_empty() { + format!("{command_name} {pattern}") + } else { + format!("{command_name} {pattern} in {path}") + } + } + "grep" => { + let pattern = input + .get("pattern") + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + let path = input + .get("path") + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + if path.is_empty() { + format!("grep {pattern}") + } else { + format!("grep {pattern} in {path}") + } + } + "code_search" => code_search_display_from_input(input), + _ => String::new(), + } +} + +pub(super) fn command_actions_from_tool_input( + tool_name: &str, + command: &str, + input: &serde_json::Value, +) -> Vec { + match tool_name { + "read" => crate::tool_actions::read_action_from_tool_input(command, input) + .into_iter() + .collect(), + "find" | "glob" => vec![devo_protocol::parse_command::ParsedCommand::ListFiles { + cmd: command.to_string(), + path: find_display_from_input(input), + }], + "grep" => vec![devo_protocol::parse_command::ParsedCommand::Search { + cmd: command.to_string(), + query: input + .get("pattern") + .and_then(serde_json::Value::as_str) + .map(ToOwned::to_owned), + path: input + .get("path") + .and_then(serde_json::Value::as_str) + .map(ToOwned::to_owned), + }], + "code_search" => code_search_action_from_input(command, input) + .into_iter() + .collect(), + _ => Vec::new(), + } +} + +fn code_search_display_from_input(input: &serde_json::Value) -> String { + match input + .get("operation") + .and_then(serde_json::Value::as_str) + .unwrap_or("search") + { + "find_related" => { + let path = input + .get("file_path") + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + let line = input.get("line").and_then(serde_json::Value::as_u64); + match (path.is_empty(), line) { + (false, Some(line)) => format!("code_search related {path}:{line}"), + (false, None) => format!("code_search related {path}"), + (true, _) => "code_search related".to_string(), + } + } + _ => { + let query = input + .get("query") + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + let path = input + .get("path") + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + match (query.is_empty(), path.is_empty()) { + (false, false) => format!("code_search {query} in {path}"), + (false, true) => format!("code_search {query}"), + (true, false) => format!("code_search in {path}"), + (true, true) => "code_search".to_string(), + } + } + } +} + +fn code_search_action_from_input( + command: &str, + input: &serde_json::Value, +) -> Option { + match input + .get("operation") + .and_then(serde_json::Value::as_str) + .unwrap_or("search") + { + "find_related" => { + let path = input + .get("file_path") + .and_then(serde_json::Value::as_str) + .filter(|path| !path.is_empty())?; + let line = input + .get("line") + .and_then(serde_json::Value::as_u64) + .map(|line| line.to_string()) + .unwrap_or_else(|| "?".to_string()); + Some(devo_protocol::parse_command::ParsedCommand::Search { + cmd: command.to_string(), + query: Some(format!("related {path}:{line}")), + path: Some(path.to_string()), + }) + } + _ => { + let query = input + .get("query") + .and_then(serde_json::Value::as_str) + .filter(|query| !query.is_empty())?; + Some(devo_protocol::parse_command::ParsedCommand::Search { + cmd: command.to_string(), + query: Some(query.to_string()), + path: input + .get("path") + .and_then(serde_json::Value::as_str) + .map(ToOwned::to_owned), + }) + } + } +} + +fn find_display_from_input(input: &serde_json::Value) -> Option { + let pattern = input + .get("pattern") + .and_then(serde_json::Value::as_str) + .filter(|pattern| !pattern.is_empty())?; + let path = input.get("path").and_then(serde_json::Value::as_str); + Some(match path.filter(|path| !path.is_empty()) { + Some(path) => format!("{pattern} in {path}"), + None => pattern.to_string(), + }) +} + +pub(super) fn command_actions_from_tool_result( + tool_name: &str, + command: &str, + input: &serde_json::Value, + summary: &str, +) -> Vec { + let actions = command_actions_from_tool_input(tool_name, command, input); + if !actions.is_empty() { + return actions; + } + match tool_name { + "read" => crate::tool_actions::read_action_from_tool_summary(summary) + .into_iter() + .collect(), + _ => actions, + } +} + +pub(super) fn command_execution_item_id_for_progress( + pending_tool_calls: &HashMap, + tool_use_id: &str, +) -> Option { + pending_tool_calls + .get(tool_use_id) + .filter(|pending| pending.display_kind.is_command_execution()) + .and_then(|pending| pending.item_id) +} + +pub(super) fn user_shell_exec_input(command: &str, cwd: std::path::PathBuf) -> serde_json::Value { + serde_json::json!({ + "cmd": command, + "workdir": cwd, + "login": true, + "tty": true, + }) +} + +pub(super) fn user_shell_command_payload( + tool_call_id: &str, + command: &str, + input: serde_json::Value, + command_actions: Vec, + output: Option, + is_error: bool, +) -> CommandExecutionPayload { + CommandExecutionPayload { + tool_call_id: tool_call_id.to_string(), + tool_name: "exec_command".to_string(), + command: command.to_string(), + input: Some(input), + source: devo_protocol::protocol::ExecCommandSource::UserShell, + command_actions, + output, + is_error, + } +} + +const AGENT_COORDINATION_TOOL_NAMES: &[&str] = &[ + "spawn_agent", + "send_message", + "wait_agent", + "list_agents", + "close_agent", +]; + +pub(super) fn without_agent_coordination_tools( + registry: &devo_core::tools::ToolRegistry, +) -> devo_core::tools::ToolRegistry { + let names = registry + .tool_definitions() + .into_iter() + .map(|tool| tool.name) + .filter(|name| { + !AGENT_COORDINATION_TOOL_NAMES + .iter() + .any(|hidden_name| *hidden_name == name) + }) + .collect::>(); + let names = names.iter().map(String::as_str).collect::>(); + registry.restricted_to_specs(&names) +} diff --git a/crates/server/src/runtime/turn_exec/tool_results.rs b/crates/server/src/runtime/turn_exec/tool_results.rs new file mode 100644 index 00000000..bc05c9ef --- /dev/null +++ b/crates/server/src/runtime/turn_exec/tool_results.rs @@ -0,0 +1,430 @@ +use std::sync::Arc; + +use devo_core::tools::ToolContent; +use devo_core::{ + CommandExecutionItem, SessionId, TextItem, ToolCallItem, ToolResultItem, TurnId, TurnItem, +}; +use devo_util_git::extract_paths_from_patch; + +use super::super::*; +use super::tool_display::{command_actions_from_tool_result, is_file_change_tool, is_plan_tool}; +use super::types::PendingToolCall; +use crate::{ + CommandExecutionPayload, FileChangePayload, ItemKind, ToolCallPayload, ToolResultPayload, + TurnPlanStepPayload, TurnPlanUpdatedPayload, +}; + +pub(super) fn tool_content_to_json(content: ToolContent) -> serde_json::Value { + match content { + ToolContent::Text(text) => serde_json::Value::String(text), + ToolContent::Json(json) => json, + ToolContent::Mixed { text, json } => { + json.unwrap_or_else(|| serde_json::Value::String(text.unwrap_or_default())) + } + } +} + +/// Completes a pending tool-call item when the tool has a specialized item kind. +/// +/// Returns `true` when the tool result item should not be emitted separately. +#[allow(clippy::too_many_arguments)] +pub(super) async fn complete_pending_tool_call( + runtime: &Arc, + session_id: SessionId, + turn_id: TurnId, + turn_for_plan_updates: &crate::TurnMetadata, + tool_use_id: &str, + tool_name: Option, + pending: &PendingToolCall, + content: &ToolContent, + display_content: Option, + is_error: bool, + summary: &str, +) -> bool { + let pending_item_id = pending.item_id.expect("pending item id"); + let pending_item_seq = pending.item_seq.expect("pending item seq"); + if let Some(ref tool_name) = tool_name { + if is_plan_tool(tool_name) { + complete_plan_tool_call( + runtime, + session_id, + turn_id, + turn_for_plan_updates, + pending_item_id, + pending_item_seq, + content, + ) + .await; + return true; + } + if is_file_change_tool(tool_name) { + complete_file_change_tool_call( + runtime, + session_id, + turn_id, + tool_use_id, + tool_name, + pending, + content, + display_content, + is_error, + pending_item_id, + pending_item_seq, + ) + .await; + return true; + } + if pending.display_kind.is_command_execution() { + complete_command_execution_tool_call( + runtime, + session_id, + turn_id, + tool_use_id, + tool_name, + pending, + content, + is_error, + summary, + pending_item_id, + pending_item_seq, + ) + .await; + return true; + } + } + complete_generic_tool_call( + runtime, + session_id, + turn_id, + tool_use_id, + tool_name.unwrap_or_default(), + pending, + summary, + pending_item_id, + pending_item_seq, + ) + .await; + false +} + +async fn complete_plan_tool_call( + runtime: &Arc, + session_id: SessionId, + turn_id: TurnId, + turn_for_plan_updates: &crate::TurnMetadata, + pending_item_id: devo_core::ItemId, + pending_item_seq: u64, + content: &ToolContent, +) { + let output_json = tool_content_to_json(content.clone()); + let explanation = output_json + .get("explanation") + .and_then(serde_json::Value::as_str) + .map(ToOwned::to_owned); + let plan = output_json + .get("plan") + .and_then(serde_json::Value::as_array) + .cloned() + .unwrap_or_default(); + runtime + .complete_item( + session_id, + turn_id, + pending_item_id, + pending_item_seq, + ItemKind::Plan, + TurnItem::Plan(TextItem { + text: output_json.to_string(), + }), + serde_json::json!({ + "title": "Plan", + "text": output_json.to_string(), + }), + ) + .await; + runtime + .broadcast_event(crate::ServerEvent::TurnPlanUpdated( + TurnPlanUpdatedPayload { + session_id, + turn: turn_for_plan_updates.clone(), + explanation, + plan: plan + .into_iter() + .filter_map(|item| { + Some(TurnPlanStepPayload { + step: item.get("step")?.as_str()?.to_string(), + status: item.get("status")?.as_str()?.to_string(), + }) + }) + .collect(), + }, + )) + .await; +} + +#[allow(clippy::too_many_arguments)] +async fn complete_file_change_tool_call( + runtime: &Arc, + session_id: SessionId, + turn_id: TurnId, + tool_use_id: &str, + tool_name: &str, + pending: &PendingToolCall, + content: &ToolContent, + display_content: Option, + is_error: bool, + pending_item_id: devo_core::ItemId, + pending_item_seq: u64, +) { + let output_json = tool_content_to_json(content.clone()); + let changes = file_changes_from_output(&output_json); + runtime + .complete_item( + session_id, + turn_id, + pending_item_id, + pending_item_seq, + ItemKind::FileChange, + TurnItem::ToolResult(ToolResultItem { + tool_call_id: tool_use_id.to_string(), + tool_name: Some(tool_name.to_string()), + output: output_json.clone(), + display_content: display_content.clone(), + is_error, + }), + serde_json::to_value(FileChangePayload { + tool_call_id: tool_use_id.to_string(), + tool_name: Some(tool_name.to_string()), + input: Some(pending.input.clone()), + changes, + is_error, + }) + .expect("serialize file change payload"), + ) + .await; +} + +fn file_changes_from_output( + output_json: &serde_json::Value, +) -> Vec<(std::path::PathBuf, devo_protocol::protocol::FileChange)> { + let changes = output_json + .get("files") + .and_then(serde_json::Value::as_array) + .cloned() + .unwrap_or_default() + .into_iter() + .filter_map(|file| { + let path = std::path::PathBuf::from(file.get("path")?.as_str()?); + let kind = file.get("kind")?.as_str()?; + let additions = file + .get("additions") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0); + let deletions = file + .get("deletions") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0); + let change = match kind { + "add" => devo_protocol::protocol::FileChange::Add { + content: file + .get("content") + .and_then(serde_json::Value::as_str) + .map(ToOwned::to_owned) + .unwrap_or_else(|| "\n".repeat(additions as usize)), + }, + "delete" => devo_protocol::protocol::FileChange::Delete { + content: file + .get("content") + .and_then(serde_json::Value::as_str) + .map(ToOwned::to_owned) + .unwrap_or_else(|| "\n".repeat(deletions as usize)), + }, + "update" | "move" => devo_protocol::protocol::FileChange::Update { + unified_diff: file + .get("diff") + .or_else(|| file.get("patch")) + .or_else(|| output_json.get("diff")) + .and_then(serde_json::Value::as_str) + .unwrap_or("") + .to_string(), + old_text: file + .get("oldContent") + .or_else(|| file.get("preContent")) + .or_else(|| file.get("pre_content")) + .and_then(serde_json::Value::as_str) + .map(ToOwned::to_owned), + new_text: file + .get("postContent") + .or_else(|| file.get("post_content")) + .or_else(|| file.get("content")) + .and_then(serde_json::Value::as_str) + .map(ToOwned::to_owned), + move_path: file + .get("movePath") + .or_else(|| file.get("move_path")) + .and_then(serde_json::Value::as_str) + .map(std::path::PathBuf::from), + }, + _ => return None, + }; + Some((path, change)) + }) + .collect::>(); + if changes.is_empty() { + output_json + .get("diff") + .and_then(serde_json::Value::as_str) + .map(extract_paths_from_patch) + .unwrap_or_default() + .into_iter() + .map(|path| { + ( + std::path::PathBuf::from(path), + devo_protocol::protocol::FileChange::Update { + unified_diff: output_json + .get("diff") + .and_then(serde_json::Value::as_str) + .unwrap_or("") + .to_string(), + old_text: None, + new_text: None, + move_path: None, + }, + ) + }) + .collect() + } else { + changes + } +} + +#[allow(clippy::too_many_arguments)] +async fn complete_command_execution_tool_call( + runtime: &Arc, + session_id: SessionId, + turn_id: TurnId, + tool_use_id: &str, + tool_name: &str, + pending: &PendingToolCall, + content: &ToolContent, + is_error: bool, + summary: &str, + pending_item_id: devo_core::ItemId, + pending_item_seq: u64, +) { + let output = tool_content_to_json(content.clone()); + let completed_payload = serde_json::to_value(CommandExecutionPayload { + tool_call_id: tool_use_id.to_string(), + tool_name: tool_name.to_string(), + command: pending.command.clone(), + input: Some(pending.input.clone()), + source: devo_protocol::protocol::ExecCommandSource::Agent, + command_actions: command_actions_from_tool_result( + tool_name, + &pending.command, + &pending.input, + summary, + ), + output: Some(output.clone()), + is_error, + }) + .expect("serialize command execution payload"); + runtime + .complete_item( + session_id, + turn_id, + pending_item_id, + pending_item_seq, + ItemKind::CommandExecution, + TurnItem::CommandExecution(CommandExecutionItem { + tool_call_id: tool_use_id.to_string(), + tool_name: tool_name.to_string(), + command: pending.command.clone(), + input: pending.input.clone(), + output, + is_error, + }), + completed_payload, + ) + .await; +} + +#[allow(clippy::too_many_arguments)] +async fn complete_generic_tool_call( + runtime: &Arc, + session_id: SessionId, + turn_id: TurnId, + tool_use_id: &str, + tool_name: String, + pending: &PendingToolCall, + summary: &str, + pending_item_id: devo_core::ItemId, + pending_item_seq: u64, +) { + let completed_payload = serde_json::to_value(ToolCallPayload { + tool_call_id: tool_use_id.to_string(), + tool_name: tool_name.clone(), + parameters: pending.input.clone(), + command_actions: command_actions_from_tool_result( + tool_name.as_str(), + &pending.command, + &pending.input, + summary, + ), + }) + .expect("serialize tool call payload"); + runtime + .complete_item( + session_id, + turn_id, + pending_item_id, + pending_item_seq, + ItemKind::ToolCall, + TurnItem::ToolCall(ToolCallItem { + tool_call_id: tool_use_id.to_string(), + tool_name, + input: pending.input.clone(), + }), + completed_payload, + ) + .await; +} + +#[allow(clippy::too_many_arguments)] +pub(super) async fn emit_tool_result_item( + runtime: &Arc, + session_id: SessionId, + turn_id: TurnId, + tool_use_id: String, + tool_name: Option, + result_input: Option, + content: ToolContent, + display_content: Option, + is_error: bool, + summary: String, +) { + runtime + .emit_turn_item( + session_id, + turn_id, + ItemKind::ToolResult, + TurnItem::ToolResult(ToolResultItem { + tool_call_id: tool_use_id.clone(), + tool_name: tool_name.clone(), + output: tool_content_to_json(content.clone()), + display_content: display_content.clone(), + is_error, + }), + serde_json::to_value(ToolResultPayload { + tool_call_id: tool_use_id, + tool_name, + input: result_input, + content: tool_content_to_json(content), + display_content, + is_error, + summary, + }) + .expect("serialize tool result payload"), + ) + .await; +} diff --git a/crates/server/src/runtime/turn_exec/trace.rs b/crates/server/src/runtime/turn_exec/trace.rs new file mode 100644 index 00000000..72d8c534 --- /dev/null +++ b/crates/server/src/runtime/turn_exec/trace.rs @@ -0,0 +1,96 @@ +use std::sync::OnceLock; +use std::time::Instant; + +use devo_core::QueryEvent; + +pub(super) fn stream_trace_elapsed_ms() -> u128 { + static STREAM_TRACE_START: OnceLock = OnceLock::new(); + STREAM_TRACE_START + .get_or_init(Instant::now) + .elapsed() + .as_millis() +} + +pub(super) fn query_event_trace_kind(event: &QueryEvent) -> &'static str { + match event { + QueryEvent::TextDelta(_) => "text_delta", + QueryEvent::ReasoningDelta(_) => "reasoning_delta", + QueryEvent::ReasoningCompleted => "reasoning_completed", + QueryEvent::ToolUseStart { .. } => "tool_use_start", + QueryEvent::ToolExecutionStart { .. } => "tool_execution_start", + QueryEvent::ToolResult { .. } => "tool_result", + QueryEvent::ToolProgress { .. } => "tool_progress", + QueryEvent::UsageDelta { .. } => "usage_delta", + QueryEvent::Usage { .. } => "usage", + QueryEvent::TurnComplete { .. } => "turn_complete", + } +} + +pub(super) fn query_event_trace_delta_len(event: &QueryEvent) -> usize { + match event { + QueryEvent::TextDelta(text) | QueryEvent::ReasoningDelta(text) => text.len(), + QueryEvent::ToolProgress { + progress: + devo_core::tools::ToolProgress::OutputDelta { delta } + | devo_core::tools::ToolProgress::StatusUpdate { message: delta, .. } + | devo_core::tools::ToolProgress::Completion { summary: delta }, + .. + } => delta.len(), + QueryEvent::ToolProgress { + progress: devo_core::tools::ToolProgress::Terminal { .. }, + .. + } + | QueryEvent::ReasoningCompleted + | QueryEvent::ToolUseStart { .. } + | QueryEvent::ToolExecutionStart { .. } + | QueryEvent::ToolResult { .. } + | QueryEvent::UsageDelta { .. } + | QueryEvent::Usage { .. } + | QueryEvent::TurnComplete { .. } => 0, + } +} + +pub(super) fn query_event_trace_token_preview(event: &QueryEvent) -> Option { + match event { + QueryEvent::TextDelta(text) => assistant_token_log_preview(text), + _ => None, + } +} + +fn assistant_token_log_preview(text: &str) -> Option { + assistant_token_logging_enabled() + .then(|| format_assistant_token_log_preview(text, assistant_token_log_max_chars())) +} + +fn assistant_token_logging_enabled() -> bool { + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| { + std::env::var("DEVO_LOG_ASSISTANT_TOKENS") + .ok() + .is_some_and(|value| { + matches!( + value.to_ascii_lowercase().as_str(), + "1" | "true" | "yes" | "on" + ) + }) + }) +} + +fn assistant_token_log_max_chars() -> usize { + static MAX_CHARS: OnceLock = OnceLock::new(); + *MAX_CHARS.get_or_init(|| { + std::env::var("DEVO_LOG_ASSISTANT_TOKENS_MAX_CHARS") + .ok() + .and_then(|value| value.parse().ok()) + .filter(|&max_chars| max_chars > 0) + .unwrap_or(120) + }) +} + +fn format_assistant_token_log_preview(text: &str, max_chars: usize) -> String { + let mut preview = text.chars().take(max_chars).collect::(); + if text.chars().count() > max_chars { + preview.push('…'); + } + preview +} diff --git a/crates/server/src/runtime/turn_exec/types.rs b/crates/server/src/runtime/turn_exec/types.rs new file mode 100644 index 00000000..9325c644 --- /dev/null +++ b/crates/server/src/runtime/turn_exec/types.rs @@ -0,0 +1,85 @@ +use devo_core::{ItemId, TurnId, TurnUsage}; + +/// Inputs captured at turn-start time and handed to the background turn executor. +pub(crate) struct ExecuteTurnRequest { + /// Runtime session that owns the turn and receives emitted items, usage, and status updates. + pub(crate) session_id: devo_core::SessionId, + /// Pre-created turn metadata persisted at turn start; execution mutates a local copy to its + /// terminal status before appending the final turn record. + pub(crate) turn: crate::TurnMetadata, + /// Resolved model, provider, reasoning, tool, web, and token-budget settings for this turn. + pub(crate) turn_config: devo_core::TurnConfig, + /// User-facing rendering of the submitted input. Visible turns persist this as the displayed + /// user message; hidden continuation turns keep it out of the transcript. + pub(crate) display_input: String, + /// Canonical resolved prompt text. Visible turns push this as the user-role message when the + /// input resolver did not return structured `input_messages`. + pub(crate) input: String, + /// Structured user-role messages produced by input resolution, such as expanded skill content. + /// When non-empty, these are pushed instead of the single `input` string. + pub(crate) input_messages: Vec, + /// Collaboration mode to install on the core session for this query; it also drives + /// mode-specific stream handling such as proposed-plan parsing. + pub(crate) collaboration_mode: devo_protocol::CollaborationMode, + /// Controls whether this executor emits/pushes a visible user message or runs hidden work such + /// as goal continuation, and carries the hidden goal context when needed. + pub(crate) input_mode: super::super::TurnInputMode, +} + +pub(super) struct PendingToolCall { + pub(super) item_id: Option, + pub(super) item_seq: Option, + pub(super) input: serde_json::Value, + pub(super) display_kind: ToolDisplayKind, + pub(super) command: String, +} + +pub(crate) struct TurnEventStreamSummary { + pub(crate) latest_usage: Option, + pub(crate) stop_reason: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum ToolDisplayKind { + CommandExecution, + Generic, +} + +impl ToolDisplayKind { + pub(super) fn for_tool_name(name: &str) -> Self { + if super::tool_display::is_unified_exec_tool(name) { + Self::CommandExecution + } else { + Self::Generic + } + } + + pub(super) fn is_command_execution(self) -> bool { + self == Self::CommandExecution + } +} + +pub(super) struct ToolStartItem { + pub(super) item_kind: crate::ItemKind, + pub(super) payload: serde_json::Value, +} + +pub(crate) struct TurnQueryOutcome { + pub(crate) result: Result<(), devo_core::AgentError>, + pub(crate) session_total_input_tokens: usize, + pub(crate) session_total_output_tokens: usize, + pub(crate) session_total_tokens: usize, + pub(crate) session_total_cache_creation_tokens: usize, + pub(crate) session_total_cache_read_tokens: usize, + pub(crate) session_last_input_tokens: usize, + pub(crate) session_prompt_token_estimate: usize, +} + +pub(super) struct QueuedTurnInput { + pub(super) display_input: String, + pub(super) input_text: String, + pub(super) input_messages: Vec, + pub(super) collaboration_mode: devo_protocol::CollaborationMode, + pub(super) model_selection: Option, + pub(super) subagent_usage_owner: Option<(devo_core::SessionId, Option)>, +} diff --git a/crates/server/src/runtime/turn_reservation.rs b/crates/server/src/runtime/turn_reservation.rs index b8944210..9d984e37 100644 --- a/crates/server/src/runtime/turn_reservation.rs +++ b/crates/server/src/runtime/turn_reservation.rs @@ -51,12 +51,21 @@ impl ServerRuntime { } } + pub(super) async fn runtime_active_turn_id(&self, session_id: SessionId) -> Option { + self.active_turn_ids + .lock() + .await + .get(&session_id) + .copied() + } + pub(super) async fn clear_active_turn_runtime_handles(&self, session_id: SessionId) { self.active_tasks.lock().await.remove(&session_id); self.active_turn_cancellations .lock() .await .remove(&session_id); + self.active_turn_ids.lock().await.remove(&session_id); self.active_turn_connections .lock() .await @@ -65,20 +74,10 @@ impl ServerRuntime { pub(super) async fn clear_active_turn_reservation( &self, - session_arc: &Arc>, + session_handle: &SessionHandle, turn_id: TurnId, ) { - let mut session = session_arc.lock().await; - if session - .active_turn - .as_ref() - .is_some_and(|active| active.turn_id == turn_id) - { - session.active_turn = None; - session.summary.status = SessionRuntimeStatus::Idle; - session.summary.updated_at = Utc::now(); - session.summary.last_activity_at = session.summary.updated_at; - } + let _ = session_handle.clear_active_turn_if_matches(turn_id).await; } } @@ -183,21 +182,10 @@ mod tests { let session_id = start_session(&runtime, data_root.path().to_path_buf()).await; let bad_rollout_path = data_root.path().join("rollout-dir"); std::fs::create_dir(&bad_rollout_path)?; - let session_arc = runtime - .sessions - .lock() - .await - .get(&session_id) - .cloned() - .expect("session"); - { - let mut session = session_arc.lock().await; - session - .record - .as_mut() - .expect("durable record") - .rollout_path = bad_rollout_path; - } + let session_handle = runtime.session(session_id).await.expect("session"); + session_handle + .update_record_rollout_path(bad_rollout_path) + .await; let value = runtime .handle_turn_shell_command_for_connection( @@ -212,12 +200,16 @@ mod tests { ) .await; let response: ErrorResponse = serde_json::from_value(value).expect("error response"); - let session = session_arc.lock().await; + let reservation = session_handle + .turn_reservation_snapshot() + .await + .expect("turn reservation snapshot"); + let summary = session_handle.summary().await.expect("summary"); assert_eq!(response.error.code, ProtocolErrorCode::InternalError); - assert_eq!(session.active_turn, None); - assert_eq!(session.summary.status, SessionRuntimeStatus::Idle); - assert_eq!(session.latest_turn, None); + assert_eq!(reservation.active_turn, None); + assert_eq!(summary.status, SessionRuntimeStatus::Idle); + assert_eq!(reservation.latest_turn, None); Ok(()) } @@ -227,13 +219,12 @@ mod tests { let data_root = TempDir::new()?; let runtime = build_runtime(data_root.path()); let session_id = start_session(&runtime, data_root.path().to_path_buf()).await; - let session_arc = runtime - .sessions - .lock() + let session_handle = runtime.session(session_id).await.expect("session"); + let reservation = session_handle + .turn_reservation_snapshot() .await - .get(&session_id) - .cloned() - .expect("session"); + .expect("turn reservation snapshot"); + let turn_config = reservation.runtime_context.resolve_turn_config(None, None); let active_turn = TurnMetadata { turn_id: TurnId::new(), session_id, @@ -252,10 +243,9 @@ mod tests { stop_reason: None, failure_reason: None, }; - { - let mut session = session_arc.lock().await; - session.active_turn = Some(active_turn); - } + session_handle + .begin_active_turn(active_turn, turn_config) + .await; let value = runtime .handle_turn_start_with_queue_policy( @@ -279,9 +269,7 @@ mod tests { ) .await; let response: ErrorResponse = serde_json::from_value(value).expect("error response"); - let queued_len = session_arc - .lock() - .await + let queued_len = reservation .pending_turn_queue .lock() .expect("pending turn queue mutex should not be poisoned") diff --git a/crates/server/src/runtime/user_input.rs b/crates/server/src/runtime/user_input.rs index d4bf6f9d..ac4e5224 100644 --- a/crates/server/src/runtime/user_input.rs +++ b/crates/server/src/runtime/user_input.rs @@ -1,6 +1,7 @@ use super::*; use crate::PendingServerRequestContext; use crate::ServerRequestKind; +use crate::runtime::session_interactive::UserInputTakeError; impl ServerRuntime { pub(super) async fn handle_request_user_input_respond( @@ -19,33 +20,35 @@ impl ServerRuntime { } }; - let Some(session_arc) = self.sessions.lock().await.get(¶ms.session_id).cloned() else { + if self.session(params.session_id).await.is_none() { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, "session does not exist", ); - }; + } let request_key = params.request_id.to_string(); - let pending = { - let mut session = session_arc.lock().await; - let Some(pending) = session.pending_user_inputs.remove(&request_key) else { + let pending = match self + .session_interactive + .take_pending_user_input(params.session_id, &request_key, params.turn_id) + .await + { + Ok(pending) => pending, + Err(UserInputTakeError::NotFound) => { return self.error_response( request_id, ProtocolErrorCode::InvalidParams, "no pending request_user_input request exists for this runtime", ); - }; - if pending.turn_id != params.turn_id { - session.pending_user_inputs.insert(request_key, pending); + } + Err(UserInputTakeError::WrongTurn) => { return self.error_response( request_id, ProtocolErrorCode::InvalidParams, "request_user_input belongs to a different turn", ); } - pending }; let _ = pending.tx.send(params.response); @@ -75,27 +78,20 @@ impl ServerRuntime { let request_id = tool_call_id; let (tx, rx) = oneshot::channel(); - let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else { + if self.session(session_id).await.is_none() { return Err(ToolCallError::ExecutionFailed( "session does not exist".to_string(), )); - }; - { - let mut session = session_arc.lock().await; - if session - .pending_user_inputs - .insert(request_id.clone(), PendingUserInput { turn_id, tx }) - .is_some() - { - tracing::warn!( - session_id = %session_id, - turn_id = %turn_id, - request_id = %request_id, - "overwriting pending request_user_input request" - ); - } } + self.session_interactive + .register_pending_user_input( + session_id, + request_id.clone(), + PendingUserInput { turn_id, tx }, + ) + .await; + self.broadcast_event(ServerEvent::RequestUserInput(RequestUserInputPayload { request: PendingServerRequestContext { request_id: request_id.clone().into(), diff --git a/crates/server/src/runtime/workspace_baseline.rs b/crates/server/src/runtime/workspace_baseline.rs index fccd07a7..c431d2a3 100644 --- a/crates/server/src/runtime/workspace_baseline.rs +++ b/crates/server/src/runtime/workspace_baseline.rs @@ -20,13 +20,7 @@ impl ServerRuntime { .lock() .await .insert(turn_id, captured.baseline); - let record = { - let session_arc = self.sessions.lock().await.get(&session_id).cloned(); - match session_arc { - Some(session_arc) => session_arc.lock().await.record.clone(), - None => None, - } - }; + let record = self.session_record_snapshot(session_id).await; if let Some(record) = record && let Err(error) = self .rollout_store @@ -71,13 +65,7 @@ impl ServerRuntime { .await { Ok(finalized) => { - let record = { - let session_arc = self.sessions.lock().await.get(&session_id).cloned(); - match session_arc { - Some(session_arc) => session_arc.lock().await.record.clone(), - None => None, - } - }; + let record = self.session_record_snapshot(session_id).await; if let Some(record) = record && let Err(error) = self .rollout_store diff --git a/crates/server/src/server_tray.rs b/crates/server/src/server_tray.rs deleted file mode 100644 index 2476963e..00000000 --- a/crates/server/src/server_tray.rs +++ /dev/null @@ -1,673 +0,0 @@ -use anyhow::Context; -use anyhow::Result; -use tray_icon::Icon; - -const ICON_SIZE: u32 = 16; -const ICON_PADDING: u32 = 1; -const TRANSPARENT_BACKGROUND_THRESHOLD: u8 = 32; -const DEVO_MARK_PNG: &[u8] = include_bytes!("../../../.github/assets/devo-mark.png"); - -#[cfg(windows)] -pub(crate) struct ServerTray { - inner: platform::PlatformServerTray, -} - -#[cfg(windows)] -impl ServerTray { - pub(crate) fn start() -> Result { - Ok(Self { - inner: platform::PlatformServerTray::start()?, - }) - } - - pub(crate) async fn shutdown_requested(&mut self) { - self.inner.shutdown_requested().await; - } -} - -#[cfg(target_os = "macos")] -pub(crate) fn run_server_process_with_macos_tray(args: crate::ServerProcessArgs) -> Result<()> { - platform::run_server_process_with_macos_tray(args) -} - -fn create_icon() -> Result { - Icon::from_rgba(devo_icon_rgba()?, ICON_SIZE, ICON_SIZE).context("create Devo tray icon image") -} - -fn devo_icon_rgba() -> Result> { - let image = image::load_from_memory_with_format(DEVO_MARK_PNG, image::ImageFormat::Png) - .context("decode embedded Devo mark")? - .into_rgba8(); - let transparent = remove_dark_background(image); - let bbox = - visible_bounds(&transparent).unwrap_or((0, 0, transparent.width(), transparent.height())); - let cropped = image::imageops::crop_imm( - &transparent, - bbox.0, - bbox.1, - bbox.2 - bbox.0, - bbox.3 - bbox.1, - ) - .to_image(); - let icon_content_size = ICON_SIZE - ICON_PADDING * 2; - let resized = image::imageops::resize( - &cropped, - icon_content_size, - icon_content_size, - image::imageops::FilterType::Lanczos3, - ); - let mut icon = image::RgbaImage::new(ICON_SIZE, ICON_SIZE); - image::imageops::overlay( - &mut icon, - &resized, - i64::from(ICON_PADDING), - i64::from(ICON_PADDING), - ); - for pixel in icon.pixels_mut() { - if pixel[3] > 0 { - pixel[0] = 0; - pixel[1] = 0; - pixel[2] = 0; - } - } - Ok(icon.into_raw()) -} - -fn remove_dark_background(mut image: image::RgbaImage) -> image::RgbaImage { - for pixel in image.pixels_mut() { - if pixel[3] > 0 - && pixel[0] <= TRANSPARENT_BACKGROUND_THRESHOLD - && pixel[1] <= TRANSPARENT_BACKGROUND_THRESHOLD - && pixel[2] <= TRANSPARENT_BACKGROUND_THRESHOLD - { - *pixel = image::Rgba([0, 0, 0, 0]); - } - } - image -} - -fn visible_bounds(image: &image::RgbaImage) -> Option<(u32, u32, u32, u32)> { - let mut min_x = image.width(); - let mut min_y = image.height(); - let mut max_x = 0; - let mut max_y = 0; - let mut found_visible_pixel = false; - - for (x, y, pixel) in image.enumerate_pixels() { - if pixel[3] == 0 { - continue; - } - found_visible_pixel = true; - min_x = min_x.min(x); - min_y = min_y.min(y); - max_x = max_x.max(x + 1); - max_y = max_y.max(y + 1); - } - - found_visible_pixel.then_some((min_x, min_y, max_x, max_y)) -} - -#[cfg(windows)] -mod platform { - use std::thread; - use std::thread::JoinHandle; - - use anyhow::Context; - use anyhow::Result; - use anyhow::anyhow; - use tokio::sync::oneshot; - use tray_icon::TrayIcon; - use tray_icon::TrayIconBuilder; - use tray_icon::menu::Menu; - use tray_icon::menu::MenuEvent; - use tray_icon::menu::MenuId; - use tray_icon::menu::MenuItem; - use windows_sys::Win32::System::Threading::GetCurrentThreadId; - use windows_sys::Win32::UI::WindowsAndMessaging::DispatchMessageW; - use windows_sys::Win32::UI::WindowsAndMessaging::GetMessageW; - use windows_sys::Win32::UI::WindowsAndMessaging::MSG; - use windows_sys::Win32::UI::WindowsAndMessaging::PM_NOREMOVE; - use windows_sys::Win32::UI::WindowsAndMessaging::PeekMessageW; - use windows_sys::Win32::UI::WindowsAndMessaging::PostThreadMessageW; - use windows_sys::Win32::UI::WindowsAndMessaging::RegisterWindowMessageW; - use windows_sys::Win32::UI::WindowsAndMessaging::TranslateMessage; - use windows_sys::Win32::UI::WindowsAndMessaging::WM_QUIT; - - const TRAY_THREAD_NAME: &str = "devo-windows-tray"; - - struct TrayResources { - _tray_icon: TrayIcon, - exit_item_id: MenuId, - } - - pub(crate) struct PlatformServerTray { - shutdown_rx: oneshot::Receiver<()>, - thread_id: u32, - _thread: JoinHandle<()>, - } - - impl PlatformServerTray { - pub(crate) fn start() -> Result { - let (shutdown_tx, shutdown_rx) = oneshot::channel(); - let (ready_tx, ready_rx) = std::sync::mpsc::channel(); - let thread = thread::Builder::new() - .name(TRAY_THREAD_NAME.to_string()) - .spawn(move || run_tray_thread(shutdown_tx, ready_tx)) - .context("spawn Windows server tray thread")?; - - let thread_id = match ready_rx - .recv() - .context("Windows server tray thread exited before initialization")? - { - Ok(thread_id) => thread_id, - Err(error) => return Err(anyhow!(error)), - }; - - Ok(Self { - shutdown_rx, - thread_id, - _thread: thread, - }) - } - - pub(crate) async fn shutdown_requested(&mut self) { - let _ = (&mut self.shutdown_rx).await; - } - } - - impl Drop for PlatformServerTray { - fn drop(&mut self) { - // SAFETY: `thread_id` is captured after the tray thread creates its - // message queue. If the thread has already exited, Windows reports - // failure and there is nothing left to clean up. - unsafe { - let _ = PostThreadMessageW(self.thread_id, WM_QUIT, 0, 0); - } - } - } - - fn run_tray_thread( - shutdown_tx: oneshot::Sender<()>, - ready_tx: std::sync::mpsc::Sender>, - ) { - create_message_queue(); - - let mut tray_resources = match create_tray_resources() { - Ok(tray_resources) => tray_resources, - Err(error) => { - let _ = ready_tx.send(Err(error.to_string())); - return; - } - }; - - // SAFETY: `GetCurrentThreadId` has no preconditions. - let thread_id = unsafe { GetCurrentThreadId() }; - if ready_tx.send(Ok(thread_id)).is_err() { - return; - } - - run_message_loop(&mut tray_resources, shutdown_tx); - drop(tray_resources); - } - - fn create_tray_resources() -> Result { - let icon = super::create_icon()?; - let tray_menu = Menu::new(); - let exit_item = MenuItem::new("Exit", /*enabled*/ true, /*accelerator*/ None); - let exit_item_id = exit_item.id().clone(); - - tray_menu - .append(&exit_item) - .context("add Windows tray exit menu item")?; - - let tray_icon = TrayIconBuilder::new() - .with_menu(Box::new(tray_menu)) - .with_tooltip("devo") - .with_icon(icon) - .with_menu_on_left_click(/*enable*/ false) - .with_menu_on_right_click(/*enable*/ true) - .build() - .context("create Windows tray icon")?; - - Ok(TrayResources { - _tray_icon: tray_icon, - exit_item_id, - }) - } - - fn create_message_queue() { - // SAFETY: Passing a null HWND with PM_NOREMOVE is the documented way to - // force creation of this thread's message queue before other threads post - // shutdown messages to it. - unsafe { - let mut msg: MSG = std::mem::zeroed(); - let _ = PeekMessageW(&mut msg, std::ptr::null_mut(), 0, 0, PM_NOREMOVE); - } - } - - fn run_message_loop(tray_resources: &mut TrayResources, shutdown_tx: oneshot::Sender<()>) { - let mut shutdown_tx = Some(shutdown_tx); - let mut exit_item_id = tray_resources.exit_item_id.clone(); - let taskbar_created_message = register_taskbar_created_message(); - - // SAFETY: This is the standard Win32 message loop for the tray thread. - // The tray icon is created on this same thread and remains alive until - // the loop exits. - unsafe { - let mut msg: MSG = std::mem::zeroed(); - loop { - let result = GetMessageW(&mut msg, std::ptr::null_mut(), 0, 0); - if result <= 0 { - break; - } - - let _ = TranslateMessage(&msg); - DispatchMessageW(&msg); - - if taskbar_created_message != 0 && msg.message == taskbar_created_message { - match create_tray_resources() { - Ok(recreated_resources) => { - *tray_resources = recreated_resources; - exit_item_id = tray_resources.exit_item_id.clone(); - tracing::info!("recreated Windows tray icon after taskbar restart"); - } - Err(error) => { - tracing::warn!(%error, "failed to recreate Windows tray icon"); - } - } - continue; - } - - if exit_menu_item_selected(&exit_item_id) - && let Some(shutdown_tx) = shutdown_tx.take() - { - let _ = shutdown_tx.send(()); - break; - } - } - } - } - - fn register_taskbar_created_message() -> u32 { - let message_name = "TaskbarCreated" - .encode_utf16() - .chain(std::iter::once(0)) - .collect::>(); - // SAFETY: `message_name` is a null-terminated UTF-16 string and remains - // alive for the duration of the call. - unsafe { RegisterWindowMessageW(message_name.as_ptr()) } - } - - fn exit_menu_item_selected(exit_item_id: &MenuId) -> bool { - let receiver = MenuEvent::receiver(); - let mut exit_requested = false; - - while let Ok(event) = receiver.try_recv() { - if event.id() == exit_item_id { - exit_requested = true; - } - } - - exit_requested - } -} - -#[cfg(target_os = "macos")] -mod platform { - use std::sync::mpsc::Receiver; - use std::sync::mpsc::Sender; - use std::thread; - use std::thread::JoinHandle; - use std::time::Duration; - - use anyhow::Context; - use anyhow::Result; - use anyhow::anyhow; - use objc2_foundation::NSString; - use tokio_util::sync::CancellationToken; - use tray_icon::TrayIcon; - use tray_icon::TrayIconBuilder; - use tray_icon::menu::Menu; - use tray_icon::menu::MenuEvent; - use tray_icon::menu::MenuId; - use tray_icon::menu::MenuItem; - use winit::application::ApplicationHandler; - use winit::event::StartCause; - use winit::event::WindowEvent; - use winit::event_loop::ActiveEventLoop; - use winit::event_loop::ControlFlow; - use winit::event_loop::EventLoop; - use winit::event_loop::EventLoopProxy; - use winit::platform::macos::ActivationPolicy; - use winit::platform::macos::EventLoopBuilderExtMacOS; - use winit::window::WindowId; - - use crate::bootstrap::ServerProcessArgs; - use crate::bootstrap::ServerProcessRunOptions; - use crate::bootstrap::ServerTrayStartup; - use crate::bootstrap::run_server_process_inner; - - const REAL_SERVER_STARTED_BRIDGE_THREAD_NAME: &str = "devo-macos-tray-startup-bridge"; - const RUNTIME_SHUTDOWN_TIMEOUT: Duration = Duration::from_millis(500); - const SERVER_THREAD_NAME: &str = "devo-macos-server-runtime"; - const STATUS_ITEM_AUTOSAVE_NAME: &str = "com.devo.server.status-item"; - const STATUS_ITEM_LENGTH: f64 = 24.0; - const TOKIO_WORKER_STACK_SIZE_BYTES: usize = 16 * 1024 * 1024; - - enum MacosTrayEvent { - RealServerStarted, - ServerDone(Result<()>), - MenuEvent(MenuEvent), - } - - struct TrayResources { - _tray_icon: TrayIcon, - exit_item_id: MenuId, - } - - struct MacosTrayApp { - shutdown_token: CancellationToken, - tray_resources: Option, - event_loop_started: bool, - real_server_started: bool, - shutdown_requested: bool, - server_result: Option>, - } - - impl MacosTrayApp { - fn new(shutdown_token: CancellationToken) -> Self { - Self { - shutdown_token, - tray_resources: None, - event_loop_started: false, - real_server_started: false, - shutdown_requested: false, - server_result: None, - } - } - - fn maybe_start_tray(&mut self) { - if self.event_loop_started && self.real_server_started { - start_tray_resources_once(&mut self.tray_resources); - } - } - } - - pub(crate) fn run_server_process_with_macos_tray(args: ServerProcessArgs) -> Result<()> { - let mut event_loop_builder = EventLoop::::with_user_event(); - event_loop_builder.with_activation_policy(ActivationPolicy::Accessory); - event_loop_builder.with_default_menu(/*enable*/ false); - event_loop_builder.with_activate_ignoring_other_apps(/*ignore*/ false); - let event_loop = event_loop_builder - .build() - .context("create macOS server tray event loop")?; - - let event_proxy = event_loop.create_proxy(); - install_menu_event_handler(&event_proxy); - let shutdown_token = CancellationToken::new(); - let (real_server_started_tx, real_server_started_rx) = std::sync::mpsc::channel(); - let real_server_started_bridge = - spawn_real_server_started_bridge(real_server_started_rx, event_proxy.clone())?; - let server_thread = spawn_server_thread( - args, - shutdown_token.clone(), - real_server_started_tx, - event_proxy, - )?; - - let result = run_winit_event_loop(event_loop, shutdown_token); - - if server_thread.join().is_err() { - return Err(anyhow!("macOS server runtime thread panicked")); - } - if real_server_started_bridge.join().is_err() { - return Err(anyhow!("macOS server startup bridge thread panicked")); - } - result - } - - fn install_menu_event_handler(event_proxy: &EventLoopProxy) { - let event_proxy = event_proxy.clone(); - MenuEvent::set_event_handler(Some(move |event| { - let _ = event_proxy.send_event(MacosTrayEvent::MenuEvent(event)); - })); - } - - fn spawn_real_server_started_bridge( - real_server_started_rx: Receiver<()>, - event_proxy: EventLoopProxy, - ) -> Result> { - thread::Builder::new() - .name(REAL_SERVER_STARTED_BRIDGE_THREAD_NAME.to_string()) - .spawn(move || { - if real_server_started_rx.recv().is_ok() { - let _ = event_proxy.send_event(MacosTrayEvent::RealServerStarted); - } - }) - .context("spawn macOS tray startup bridge thread") - } - - fn spawn_server_thread( - args: ServerProcessArgs, - external_shutdown: CancellationToken, - real_server_started_tx: Sender<()>, - event_proxy: EventLoopProxy, - ) -> Result> { - thread::Builder::new() - .name(SERVER_THREAD_NAME.to_string()) - .spawn(move || { - let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - build_server_runtime().and_then(|runtime| { - let result = runtime.block_on(run_server_process_inner( - args, - ServerProcessRunOptions { - external_shutdown: Some(external_shutdown), - tray_startup: ServerTrayStartup::Disabled, - real_server_started: Some(real_server_started_tx), - }, - )); - runtime.shutdown_timeout(RUNTIME_SHUTDOWN_TIMEOUT); - result - }) - })) - .unwrap_or_else(|_| Err(anyhow!("macOS server runtime thread panicked"))); - let _ = event_proxy.send_event(MacosTrayEvent::ServerDone(result)); - }) - .context("spawn macOS server runtime thread") - } - - fn build_server_runtime() -> Result { - let mut builder = tokio::runtime::Builder::new_multi_thread(); - builder.enable_all(); - builder.thread_stack_size(TOKIO_WORKER_STACK_SIZE_BYTES); - Ok(builder.build()?) - } - - fn run_winit_event_loop( - event_loop: EventLoop, - shutdown_token: CancellationToken, - ) -> Result<()> { - let mut app = MacosTrayApp::new(shutdown_token); - event_loop - .run_app(&mut app) - .context("run macOS server tray event loop")?; - app.server_result.unwrap_or_else(|| { - Err(anyhow!( - "macOS server runtime exited without reporting result" - )) - }) - } - - impl ApplicationHandler for MacosTrayApp { - fn new_events(&mut self, event_loop: &ActiveEventLoop, cause: StartCause) { - event_loop.set_control_flow(ControlFlow::Wait); - if cause == StartCause::Init { - self.event_loop_started = true; - tracing::debug!("started macOS server tray event loop"); - self.maybe_start_tray(); - } - } - - fn resumed(&mut self, _event_loop: &ActiveEventLoop) {} - - fn user_event(&mut self, event_loop: &ActiveEventLoop, event: MacosTrayEvent) { - match event { - MacosTrayEvent::RealServerStarted => { - self.real_server_started = true; - self.maybe_start_tray(); - } - MacosTrayEvent::MenuEvent(event) => { - if !self.shutdown_requested - && self - .tray_resources - .as_ref() - .is_some_and(|resources| event.id() == &resources.exit_item_id) - { - self.shutdown_requested = true; - self.shutdown_token.cancel(); - } - } - MacosTrayEvent::ServerDone(result) => { - self.server_result = Some(result); - event_loop.exit(); - } - } - } - - fn window_event( - &mut self, - _event_loop: &ActiveEventLoop, - _window_id: WindowId, - _event: WindowEvent, - ) { - } - - fn exiting(&mut self, _event_loop: &ActiveEventLoop) { - drop(self.tray_resources.take()); - } - } - - fn start_tray_resources_once(tray_resources: &mut Option) { - if tray_resources.is_some() { - return; - } - - match create_tray_resources() { - Ok(resources) => { - *tray_resources = Some(resources); - tracing::info!("started macOS server tray icon"); - } - Err(error) => { - tracing::warn!(%error, "failed to start macOS server tray icon"); - } - } - } - - fn create_tray_resources() -> Result { - let icon = super::create_icon()?; - let tray_menu = Menu::new(); - let exit_item = MenuItem::new("Exit", /*enabled*/ true, /*accelerator*/ None); - let exit_item_id = exit_item.id().clone(); - - tray_menu - .append(&exit_item) - .context("add macOS tray exit menu item")?; - - let tray_icon = TrayIconBuilder::new() - .with_menu(Box::new(tray_menu)) - .with_tooltip("devo") - .with_icon(icon) - .with_icon_as_template(/*is_template*/ true) - .with_menu_on_left_click(/*enable*/ true) - .with_menu_on_right_click(/*enable*/ true) - .build() - .context("create macOS tray icon")?; - if let Some(ns_status_item) = tray_icon.ns_status_item() { - let autosave_name = NSString::from_str(STATUS_ITEM_AUTOSAVE_NAME); - ns_status_item.setAutosaveName(Some(&autosave_name)); - ns_status_item.setLength(STATUS_ITEM_LENGTH); - } - - Ok(TrayResources { - _tray_icon: tray_icon, - exit_item_id, - }) - } -} - -#[cfg(test)] -mod tests { - use pretty_assertions::assert_eq; - - #[test] - fn tray_icon_rgba_keeps_transparent_corners() { - let rgba = super::devo_icon_rgba().expect("decode Devo tray icon"); - let first_pixel_alpha = rgba[3]; - let top_center_alpha = rgba[(super::ICON_SIZE as usize / 2 * 4) + 3]; - let center_alpha_index = ((super::ICON_SIZE as usize / 2 * super::ICON_SIZE as usize) - + super::ICON_SIZE as usize / 2) - * 4 - + 3; - let center_alpha = rgba[center_alpha_index]; - let last_pixel_alpha = rgba[rgba.len() - 1]; - - assert_eq!( - ( - rgba.len(), - first_pixel_alpha, - top_center_alpha, - center_alpha > 240, - last_pixel_alpha, - ), - ( - (super::ICON_SIZE * super::ICON_SIZE * 4) as usize, - 0, - 0, - true, - 0, - ) - ); - } - - #[test] - fn tray_icon_rgba_uses_template_mask_pixels() { - let rgba = super::devo_icon_rgba().expect("decode Devo tray icon"); - let visible_pixels = rgba - .chunks_exact(4) - .filter(|pixel| pixel[3] > 0) - .collect::>(); - let visible_rgb = visible_pixels - .iter() - .map(|pixel| [pixel[0], pixel[1], pixel[2]]) - .collect::>(); - - assert!(!visible_pixels.is_empty()); - assert_eq!(visible_rgb, vec![[0, 0, 0]; visible_pixels.len()]); - } - - #[test] - fn desktop_macos_template_icons_use_box_with_cutout_mark() { - for (bytes, size) in [ - ( - include_bytes!("../../../apps/desktop/resources/iconTemplate.png").as_slice(), - 16, - ), - ( - include_bytes!("../../../apps/desktop/resources/iconTemplate@2x.png").as_slice(), - 32, - ), - ] { - let image = image::load_from_memory_with_format(bytes, image::ImageFormat::Png) - .expect("decode desktop macOS tray icon") - .into_rgba8(); - - let inset = size / 16; - let center = image.get_pixel(size / 2, size / 2); - - assert_eq!([image.width(), image.height()], [size, size]); - assert_eq!(image.get_pixel(0, 0)[3], 0); - assert!(image.get_pixel(size / 2, inset)[3] > 200); - assert_eq!(center[3], 0); - } - } -} diff --git a/crates/server/tests/deep_research_e2e.rs b/crates/server/tests/deep_research_e2e.rs index 3479f8c0..236d4240 100644 --- a/crates/server/tests/deep_research_e2e.rs +++ b/crates/server/tests/deep_research_e2e.rs @@ -1,7 +1,9 @@ +use std::io::Write; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::time::Duration; +use std::time::{SystemTime, UNIX_EPOCH}; use anyhow::Context; use anyhow::Result; @@ -194,7 +196,11 @@ impl ModelProviderSDK for ScriptedResearchProvider { .contains("Research the current official DeepSeek website domain"), "research question should not be injected into system prompt" ); - streamed_text_events( + streamed_text_event_chunks( + &[ + r#"{"need_clarification":false,"question":"","verification":"Research "#, + r#"DeepSeek official website."}"#, + ], r#"{"need_clarification":false,"question":"","verification":"Research DeepSeek official website."}"#, ) } else if request_has_stage(&request, "research brief") { @@ -592,6 +598,10 @@ async fn deep_research_turn_streams_artifact_reasoning_and_final_report() -> Res == Some(&serde_json::json!("item/researchArtifact/delta"))), "expected streamed research artifact delta: {events:#?}" ); + assert!( + events.iter().any(is_clarification_artifact_delta), + "expected streamed clarification artifact delta: {events:#?}" + ); assert!( events.iter().any(is_reasoning_delta), "expected reasoning delta: {events:#?}" @@ -627,9 +637,9 @@ async fn deep_research_turn_streams_artifact_reasoning_and_final_report() -> Res assert_eq!( completed_turn["params"]["turn"]["usage"], serde_json::json!({ - "input_tokens": 7, - "output_tokens": 7, - "total_tokens": 14, + "input_tokens": 8, + "output_tokens": 8, + "total_tokens": 16, "cache_creation_input_tokens": null, "cache_read_input_tokens": null }) @@ -668,7 +678,8 @@ async fn deep_research_accepts_write_tool_only_final_report() -> Result<()> { report_contents.contains("DeepSeek official website"), "expected written report to contain final report content: {report_contents}" ); - let final_report = latest_agent_message(&events).context("expected final report message")?; + let final_report = + latest_parent_agent_message(&events, session_id).context("expected final report message")?; assert!( final_report.contains("Wrote the full research report"), "expected final response to point at written report: {final_report}" @@ -896,7 +907,8 @@ async fn deep_research_continues_after_delegated_worker_failure() -> Result<()> completed_turn["params"]["turn"]["status"], serde_json::json!("Completed") ); - let final_report = latest_agent_message(&events).context("expected final report message")?; + let final_report = + latest_parent_agent_message(&events, session_id).context("expected final report message")?; assert!( final_report.contains("DeepSeek official website"), "expected final report after delegated worker recovery: {final_report}" @@ -941,7 +953,8 @@ async fn deep_research_restarts_worker_when_continuation_fails() -> Result<()> { assert_eq!(child_failures, 2); let child_sessions = unique_child_turn_sessions(&events, session_id); assert_eq!(child_sessions.len(), 3); - let final_report = latest_agent_message(&events).context("expected final report message")?; + let final_report = + latest_parent_agent_message(&events, session_id).context("expected final report message")?; assert!( final_report.contains("DeepSeek official website"), "expected final report after replacement worker recovery: {final_report}" @@ -1553,7 +1566,6 @@ fn supervisor_stream_events(request: &ModelRequest) -> Vec> assert_supervisor_request_uses_agent_tools(request); for attempt in 1..=3 { let spawn_id = format!("spawn-supervisor-worker-{attempt}"); - let wait_id = format!("wait-supervisor-worker-{attempt}"); if !request_has_tool_result(request, &spawn_id) { return streamed_tool_call_events( &spawn_id, @@ -1564,32 +1576,74 @@ fn supervisor_stream_events(request: &ModelRequest) -> Vec> }), ); } - if !request_has_tool_result(request, &wait_id) { - let target = tool_result_json(request, &spawn_id) - .and_then(|value| { - value - .get("agent_path") - .and_then(serde_json::Value::as_str) - .or_else(|| { - value - .get("child_session_id") - .and_then(serde_json::Value::as_str) - }) - .map(str::to_string) - }) - .unwrap_or_default(); - let mut input = serde_json::json!({ "timeout_secs": 120 }); - if !target.is_empty() { - input["target"] = serde_json::Value::String(target); + let target = tool_result_json(request, &spawn_id) + .and_then(|value| { + value + .get("agent_path") + .and_then(serde_json::Value::as_str) + .or_else(|| { + value + .get("child_session_id") + .and_then(serde_json::Value::as_str) + }) + .map(str::to_string) + }) + .unwrap_or_default(); + let mut latest_completed_poll = None; + for poll_index in 0..=32 { + if request_has_tool_result(request, &supervisor_wait_tool_id(attempt, poll_index)) { + latest_completed_poll = Some(poll_index); + } else { + break; } - return streamed_tool_call_events(&wait_id, "wait_agent", input); } - let wait_content = request_tool_result_content(request, &wait_id).unwrap_or_default(); - if !wait_content.contains("failed") - && !wait_content.contains("interrupted") - && !wait_content.contains("canceled") - { - return streamed_text_events(supervisor_notes()); + match latest_completed_poll { + None => { + let wait_id = supervisor_wait_tool_id(attempt, 0); + let mut input = serde_json::json!({ "timeout_secs": 120 }); + if !target.is_empty() { + input["target"] = serde_json::Value::String(target.clone()); + } + return streamed_tool_call_events(&wait_id, "wait_agent", input); + } + Some(poll_index) => { + let wait_id = supervisor_wait_tool_id(attempt, poll_index); + let wait_content = request_tool_result_content(request, &wait_id).unwrap_or_default(); + // #region agent log + agent_debug_log( + "deep_research_e2e.rs:supervisor_stream_events", + "supervisor wait_agent poll result", + serde_json::json!({ + "attempt": attempt, + "pollIndex": poll_index, + "waitId": wait_id, + "terminal": wait_agent_result_is_terminal(&wait_content), + "failed": wait_agent_result_indicates_failure(&wait_content), + "succeeded": wait_agent_result_indicates_success(&wait_content), + }), + "A", + ); + // #endregion + if wait_agent_result_indicates_failure(&wait_content) { + break; + } + if wait_agent_result_indicates_success(&wait_content) { + return streamed_text_events(supervisor_notes()); + } + if poll_index >= 32 { + break; + } + let next_poll = poll_index + 1; + let next_wait_id = supervisor_wait_tool_id(attempt, next_poll); + let mut input = serde_json::json!({ "timeout_secs": 2 }); + if !target.is_empty() { + input["target"] = serde_json::Value::String(target.clone()); + } + if let Some(next_sequence) = wait_agent_next_sequence(&wait_content) { + input["after_sequence"] = serde_json::json!(next_sequence); + } + return streamed_tool_call_events(&next_wait_id, "wait_agent", input); + } } } streamed_text_events( @@ -1597,6 +1651,96 @@ fn supervisor_stream_events(request: &ModelRequest) -> Vec> ) } +fn supervisor_wait_tool_id(attempt: usize, poll_index: u32) -> String { + if poll_index == 0 { + format!("wait-supervisor-worker-{attempt}") + } else { + format!("wait-supervisor-worker-{attempt}-poll-{poll_index}") + } +} + +fn wait_agent_next_sequence(content: &str) -> Option { + serde_json::from_str::(content) + .ok() + .and_then(|value| value.get("next_sequence").and_then(serde_json::Value::as_u64)) +} + +fn wait_agent_result_indicates_failure(content: &str) -> bool { + if content.contains("failed") + || content.contains("interrupted") + || content.contains("canceled") + { + return true; + } + wait_agent_statuses(content).any(|status| { + matches!(status.as_str(), "failed" | "interrupted" | "closed") + }) +} + +fn wait_agent_result_indicates_success(content: &str) -> bool { + wait_agent_events(content).iter().any(|event| { + event.get("kind").and_then(serde_json::Value::as_str) == Some("assistant_message") + }) || wait_agent_statuses(content).any(|status| status == "completed") +} + +fn wait_agent_result_is_terminal(content: &str) -> bool { + wait_agent_result_indicates_failure(content) || wait_agent_result_indicates_success(content) +} + +fn wait_agent_events(content: &str) -> Vec { + serde_json::from_str::(content) + .ok() + .and_then(|value| { + value + .get("events") + .and_then(serde_json::Value::as_array) + .cloned() + }) + .unwrap_or_default() +} + +fn wait_agent_statuses(content: &str) -> impl Iterator { + wait_agent_events(content).into_iter().filter_map(|event| { + if event.get("kind").and_then(serde_json::Value::as_str) != Some("status") { + return None; + } + event + .get("status") + .and_then(serde_json::Value::as_str) + .map(str::to_string) + }) +} + +// #region agent log +fn agent_debug_log( + location: &str, + message: &str, + data: serde_json::Value, + hypothesis_id: &str, +) { + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_millis()) + .unwrap_or(0); + let payload = serde_json::json!({ + "sessionId": "ec5e4e", + "location": location, + "message": message, + "data": data, + "hypothesisId": hypothesis_id, + "runId": "post-fix", + "timestamp": timestamp, + }); + if let Ok(mut file) = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../debug-ec5e4e.log")) + { + let _ = writeln!(file, "{payload}"); + } +} +// #endregion + fn supervisor_worker_message(attempt: usize) -> String { format!( "You are a delegated DeepResearch worker for attempt {attempt}.\n\n\nUse the parent-provided current date, timezone, and cwd.\n\n\n\nResearch the current official DeepSeek website domain. Use web search, keep the final report short, and include source URLs.\n\n\n\nResearch DeepSeek official website.\n\n\nReturn concise evidence notes with searches/tool calls, key findings, source table, uncertainty, and recommended citations. Do not write report files." @@ -1769,6 +1913,26 @@ fn streamed_text_events(text: impl Into) -> Vec> { ] } +fn streamed_text_event_chunks( + chunks: &[&str], + final_text: impl Into, +) -> Vec> { + let final_text = final_text.into(); + let mut events = chunks + .iter() + .map(|chunk| { + Ok(StreamEvent::TextDelta { + index: 0, + text: (*chunk).to_string(), + }) + }) + .collect::>(); + events.push(Ok(StreamEvent::MessageDone { + response: model_response(final_text), + })); + events +} + async fn initialize_connection( runtime: &Arc, ) -> Result<(u64, mpsc::Receiver)> { @@ -2200,6 +2364,15 @@ fn is_reasoning_delta(event: &serde_json::Value) -> bool { == serde_json::json!("agent_thought_chunk")) } +fn is_clarification_artifact_delta(event: &serde_json::Value) -> bool { + if event.get("method") != Some(&serde_json::json!("item/researchArtifact/delta")) { + return false; + } + event["params"]["payload"]["delta"] + .as_str() + .is_some_and(|delta| delta.contains("Research ")) +} + fn is_agent_message_delta(event: &serde_json::Value) -> bool { event.get("method") == Some(&serde_json::json!("item/agentMessage/delta")) || (event.get("method") == Some(&serde_json::json!("session/update")) @@ -2263,23 +2436,44 @@ fn legacy_event_from_acp_notification(value: serde_json::Value) -> serde_json::V } fn latest_agent_message(events: &[serde_json::Value]) -> Option { + events.iter().rev().find_map(|event| agent_message_from_completed_event(event)) +} + +fn latest_parent_agent_message( + events: &[serde_json::Value], + session_id: devo_core::SessionId, +) -> Option { + let expected_session_id = session_id.to_string(); events.iter().rev().find_map(|event| { - if event.get("method") != Some(&serde_json::json!("item/completed")) { + if event_session_id(event) != Some(expected_session_id.as_str()) { return None; } - let item = &event["params"]["item"]; - if item["item_kind"] == serde_json::json!("agent_message") { - return item["payload"]["text"].as_str().map(str::to_string); - } - if item["item_kind"] == serde_json::json!("research_artifact") - && item["payload"]["artifact_type"] == serde_json::json!("failure") - { - return item["payload"]["content"].as_str().map(str::to_string); - } - None + agent_message_from_completed_event(event) }) } +fn event_session_id(event: &serde_json::Value) -> Option<&str> { + event["params"]["context"]["session_id"] + .as_str() + .or_else(|| event["params"]["session_id"].as_str()) +} + +fn agent_message_from_completed_event(event: &serde_json::Value) -> Option { + if event.get("method") != Some(&serde_json::json!("item/completed")) { + return None; + } + let item = &event["params"]["item"]; + if item["item_kind"] == serde_json::json!("agent_message") { + return item["payload"]["text"].as_str().map(str::to_string); + } + if item["item_kind"] == serde_json::json!("research_artifact") + && item["payload"]["artifact_type"] == serde_json::json!("failure") + { + return item["payload"]["content"].as_str().map(str::to_string); + } + None +} + async fn respond_to_clarification( runtime: &Arc, connection_id: u64, @@ -2362,6 +2556,10 @@ fn assert_research_artifacts(events: &[serde_json::Value]) { .filter_map(|event| event["params"]["item"]["payload"]["artifact_type"].as_str()) .collect::>(); + assert!( + artifact_types.contains(&"clarification"), + "expected clarification artifact: {events:#?}" + ); assert!( artifact_types.contains(&"brief"), "expected brief artifact: {events:#?}" diff --git a/crates/server/tests/subagent_lifecycle.rs b/crates/server/tests/subagent_lifecycle.rs index c7132c43..6c5b93da 100644 --- a/crates/server/tests/subagent_lifecycle.rs +++ b/crates/server/tests/subagent_lifecycle.rs @@ -110,8 +110,8 @@ async fn spawn_agent_tool_call_does_not_deadlock_parent_turn() -> Result<()> { let data_root = TempDir::new()?; let provider = Arc::new(ScriptedProvider::new([ ScriptedProvider::spawn_agent_tool_call("verify parent spawn tool call returns", "none"), - ScriptedProvider::completed("spawn tool result observed"), ScriptedProvider::completed("child finished"), + ScriptedProvider::completed("spawn tool result observed"), ])); let runtime = build_runtime(data_root.path(), Arc::clone(&provider) as _)?; let (connection_id, mut notifications_rx) = initialize_connection(&runtime).await?; @@ -140,6 +140,52 @@ async fn spawn_agent_tool_call_does_not_deadlock_parent_turn() -> Result<()> { Ok(()) } +#[tokio::test] +async fn dual_spawn_agent_tool_calls_in_one_response_do_not_deadlock() -> Result<()> { + let data_root = TempDir::new()?; + let provider = Arc::new(ScriptedProvider::new([ + ScriptedProvider::dual_spawn_agent_tool_calls( + "first delegated worker task", + "second delegated worker task", + "none", + ), + ScriptedProvider::completed("first child finished"), + ScriptedProvider::completed("second child finished"), + ScriptedProvider::completed("both children spawned successfully"), + ])); + let runtime = build_runtime(data_root.path(), Arc::clone(&provider) as _)?; + let (connection_id, mut notifications_rx) = initialize_connection(&runtime).await?; + let parent_session_id = start_parent_session(&runtime, connection_id, data_root.path()).await?; + + let drain_notifications = + tokio::spawn(async move { while notifications_rx.recv().await.is_some() {} }); + + start_turn_with_approval_policy( + &runtime, + connection_id, + parent_session_id, + "spawn two children in one tool batch", + Some("never"), + ) + .await?; + + timeout(Duration::from_secs(15), async { + loop { + let agents = request_agent_list(&runtime, connection_id, parent_session_id).await?; + if agents.agents.len() >= 2 { + return Ok::<_, anyhow::Error>(agents); + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + }) + .await + .context("timed out waiting for both child agents to register")??; + + drain_notifications.abort(); + + Ok(()) +} + #[tokio::test] async fn deep_research_spawn_uses_research_child_context() -> Result<()> { let data_root = TempDir::new()?; @@ -640,6 +686,10 @@ async fn child_to_parent_message_is_rejected() -> Result<()> { #[tokio::test] async fn close_agent_records_closed_output_event_once() -> Result<()> { + let _ = tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .with_test_writer() + .try_init(); let data_root = TempDir::new()?; let provider = Arc::new(ScriptedProvider::pending()); let runtime = build_runtime(data_root.path(), provider as _)?; diff --git a/crates/server/tests/support/subagent_lifecycle.rs b/crates/server/tests/support/subagent_lifecycle.rs index 703af1f8..158fb979 100644 --- a/crates/server/tests/support/subagent_lifecycle.rs +++ b/crates/server/tests/support/subagent_lifecycle.rs @@ -121,6 +121,57 @@ impl ScriptedProvider { ]) } + pub fn dual_spawn_agent_tool_calls( + first_message: &str, + second_message: &str, + fork_turns: &str, + ) -> StreamScript { + let first_input = serde_json::json!({ + "message": first_message, + "fork_turns": fork_turns, + }); + let second_input = serde_json::json!({ + "message": second_message, + "fork_turns": fork_turns, + }); + let first_id = "spawn-agent-call-1".to_string(); + let second_id = "spawn-agent-call-2".to_string(); + StreamScript::Events(vec![ + StreamEvent::ToolCallStart { + index: 0, + id: first_id.clone(), + name: "spawn_agent".to_string(), + input: first_input.clone(), + }, + StreamEvent::ToolCallStart { + index: 1, + id: second_id.clone(), + name: "spawn_agent".to_string(), + input: second_input.clone(), + }, + StreamEvent::MessageDone { + response: ModelResponse { + id: "dual-spawn-agent-response".to_string(), + content: vec![ + ResponseContent::ToolUse { + id: first_id, + name: "spawn_agent".to_string(), + input: first_input, + }, + ResponseContent::ToolUse { + id: second_id, + name: "spawn_agent".to_string(), + input: second_input, + }, + ], + stop_reason: Some(StopReason::ToolUse), + usage: Usage::default(), + metadata: ResponseMetadata::default(), + }, + }, + ]) + } + pub fn wait_agent_tool_call(timeout_secs: u64) -> StreamScript { let input = serde_json::json!({ "timeout_secs": timeout_secs }); let tool_call_id = "wait-agent-call".to_string(); diff --git a/crates/tui/src/chatwidget/subagent_live_list.rs b/crates/tui/src/chatwidget/subagent_live_list.rs index 5ba171a8..dc68e39a 100644 --- a/crates/tui/src/chatwidget/subagent_live_list.rs +++ b/crates/tui/src/chatwidget/subagent_live_list.rs @@ -24,6 +24,7 @@ use crate::ui_consts::LIVE_PREFIX_COLS; pub(super) const MAX_VISIBLE_SUBAGENTS: usize = 3; +#[derive(Debug)] pub(super) struct SubagentLiveListRow { pub(super) key: SubagentLiveListRowKey, pub(super) name: String, @@ -177,7 +178,7 @@ fn take_suffix_by_width(text: &str, max_width: usize) -> String { fn status_marker_style(status: &str) -> Style { match status.to_ascii_lowercase().as_str() { - "idle" => Style::default().fg(COMPLETED_COLOR).bold(), + "idle" | "done" => Style::default().fg(COMPLETED_COLOR).bold(), "working" | "running" | "active_turn" => Style::default().fg(RUNNING_COLOR).bold(), "waiting_client" => Style::default().fg(REASONING_ACCENT_COLOR).bold(), _ => Style::default().fg(RUNNING_COLOR).bold(), @@ -187,7 +188,7 @@ fn status_marker_style(status: &str) -> Style { fn status_text_style(status: &str) -> Style { match status.to_ascii_lowercase().as_str() { "running" | "active_turn" | "working" => Style::default().fg(RUNNING_COLOR).bold(), - "idle" => Style::default().fg(COMPLETED_COLOR).bold(), + "idle" | "done" => Style::default().fg(COMPLETED_COLOR).bold(), "waiting_client" => Style::default().fg(REASONING_ACCENT_COLOR).bold(), _ => Style::default().fg(MUTED_COLOR), } diff --git a/crates/tui/src/chatwidget/subagent_monitor.rs b/crates/tui/src/chatwidget/subagent_monitor.rs index f629ec75..35553439 100644 --- a/crates/tui/src/chatwidget/subagent_monitor.rs +++ b/crates/tui/src/chatwidget/subagent_monitor.rs @@ -4,6 +4,8 @@ //! selector for live direct children, and Enter asks the host overlay to render //! the selected child through the normal transcript pager. +use std::cell::Cell; +use std::cell::RefCell; use std::collections::HashMap; use std::time::Duration; use std::time::Instant; @@ -41,7 +43,7 @@ use super::subagent_live_list::SubagentLiveListRowKey; const SUBAGENT_INACTIVE_GRACE: Duration = Duration::from_secs(12); const SUBAGENT_CTRL_X_REVEAL: Duration = Duration::from_secs(15); -#[derive(Debug, Default)] +#[derive(Debug)] pub(super) struct SubagentMonitorState { live_list_focused: bool, list_reveal_until: Option, @@ -49,6 +51,23 @@ pub(super) struct SubagentMonitorState { selected: Option, user_selected: bool, sessions: HashMap, + live_list_rows_dirty: Cell, + live_list_rows_cache: RefCell>, +} + +impl Default for SubagentMonitorState { + fn default() -> Self { + Self { + live_list_focused: false, + list_reveal_until: None, + agents: Vec::new(), + selected: None, + user_selected: false, + sessions: HashMap::new(), + live_list_rows_dirty: Cell::new(true), + live_list_rows_cache: RefCell::new(Vec::new()), + } + } } #[derive(Debug, Default)] @@ -201,6 +220,10 @@ impl ChatWidget { } } + pub(super) fn invalidate_subagent_live_list_cache(&mut self) { + self.subagent_monitor.live_list_rows_dirty.set(true); + } + pub(super) fn subagent_live_list_desired_height(&self) -> u16 { subagent_live_list::desired_height(self.subagent_live_list_rows().len()) } @@ -228,6 +251,7 @@ impl ChatWidget { ); self.upsert_subagent(agent); self.sync_subagent_hint(Instant::now()); + self.invalidate_subagent_live_list_cache(); self.frame_requester.schedule_frame(); } @@ -289,7 +313,11 @@ impl ChatWidget { focused = self.subagent_monitor.live_list_focused, "subagent monitor event applied" ); - self.frame_requester.schedule_frame(); + self.invalidate_subagent_live_list_cache(); + let now = Instant::now(); + if self.subagent_monitor.live_list_focused || self.has_visible_subagent_list(now) { + self.frame_requester.schedule_frame(); + } } pub(crate) fn reset_subagent_monitor(&mut self) { @@ -401,7 +429,11 @@ impl ChatWidget { .sessions .entry(session_id) .or_default(); + let previous_status = view.status.clone(); view.status = agent.status.clone(); + if is_terminal_status(&previous_status) && !is_terminal_status(&view.status) { + view.status = previous_status; + } view.agent = Some(agent.clone()); if is_active_subagent_status(&normalize_subagent_display_status(&agent.status)) { view.touch_activity(); @@ -412,6 +444,7 @@ impl ChatWidget { if let Some(message) = agent.last_task_message.as_deref() { view.seed_initial_task_message(message); } + self.invalidate_subagent_live_list_cache(); } fn sync_subagent_hint(&mut self, now: Instant) { @@ -425,6 +458,7 @@ impl ChatWidget { self.ensure_visible_subagent_selected(now); } self.bottom_pane.set_subagent_hint_visible(has_active); + self.prune_terminal_subagents(now); tracing::debug!( target: "devo_tui::subagent", has_active, @@ -438,6 +472,42 @@ impl ChatWidget { ); } + fn prune_terminal_subagents(&mut self, now: Instant) { + let mut removed = Vec::new(); + self.subagent_monitor.agents.retain(|agent| { + let Some(view) = self.subagent_monitor.sessions.get(&agent.session_id) else { + return false; + }; + if view.active_turn.is_some() || view.has_live_tail() { + return true; + } + if is_active_subagent_status(&normalize_subagent_display_status(&view.status)) { + return true; + } + if !is_terminal_status(&view.status) { + return true; + } + let retain = view + .last_activity_at + .is_some_and(|at| now.duration_since(at) < SUBAGENT_INACTIVE_GRACE); + if !retain { + removed.push(agent.session_id); + } + retain + }); + let had_removals = !removed.is_empty(); + for session_id in removed { + self.subagent_monitor.sessions.remove(&session_id); + if self.subagent_monitor.selected == Some(session_id) { + self.subagent_monitor.selected = None; + self.subagent_monitor.user_selected = false; + } + } + if had_removals { + self.invalidate_subagent_live_list_cache(); + } + } + fn has_active_subagents(&self) -> bool { self.subagent_monitor.agents.iter().any(|agent| { is_active_subagent_status(&self.subagent_display_status(agent)) @@ -531,6 +601,7 @@ impl ChatWidget { }; self.subagent_monitor.selected = Some(visible_ids[next]); self.subagent_monitor.user_selected = true; + self.invalidate_subagent_live_list_cache(); self.frame_requester.schedule_frame(); } @@ -587,7 +658,16 @@ impl ChatWidget { self.active_subagent_ids() } - fn subagent_live_list_rows(&self) -> Vec { + fn subagent_live_list_rows(&self) -> std::cell::Ref<'_, Vec> { + if self.subagent_monitor.live_list_rows_dirty.get() { + let rows = self.compute_subagent_live_list_rows(); + *self.subagent_monitor.live_list_rows_cache.borrow_mut() = rows; + self.subagent_monitor.live_list_rows_dirty.set(false); + } + self.subagent_monitor.live_list_rows_cache.borrow() + } + + fn compute_subagent_live_list_rows(&self) -> Vec { let now = Instant::now(); let mut rows = self .visible_subagent_ids(now) @@ -646,13 +726,17 @@ impl ChatWidget { fn subagent_display_status(&self, agent: &SubagentMonitorAgent) -> String { let stored = self.subagent_status_for_agent(agent); + let normalized = normalize_subagent_display_status(&stored); + if is_terminal_status(&normalized) { + return normalized; + } let Some(view) = self.subagent_monitor.sessions.get(&agent.session_id) else { - return normalize_subagent_display_status(&stored); + return normalized; }; if view.has_live_tail() || view.active_turn.is_some() { return "working".to_string(); } - normalize_subagent_display_status(&stored) + normalized } fn subagent_transcript_item_cell( @@ -946,15 +1030,15 @@ impl SubagentSessionView { status, } => { self.touch_activity(); - self.status = status.clone(); + self.status = normalize_terminal_subagent_status(&status); self.active_turn = None; self.flush_active_items(); - self.set_latest_preview(format!("Turn {status}")); + self.set_latest_preview(format!("Turn {}", self.status)); self.transcript.push(MonitorTranscriptItem { kind: MonitorTranscriptKind::Status, - title: format!("Turn {status}"), + title: format!("Turn {}", self.status), body: String::new(), - is_error: status.to_lowercase().contains("failed"), + is_error: self.status.to_ascii_lowercase().contains("failed"), }); } SubagentMonitorEvent::TurnFailed { @@ -984,8 +1068,18 @@ impl SubagentSessionView { session_id: _, status, } => { - if !is_terminal_status(&self.status) { - self.status = normalize_runtime_status_for_subagent(status); + if is_terminal_status(&self.status) { + return; + } + let normalized = normalize_runtime_status_for_subagent(status); + if normalized == "idle" && self.active_turn.is_none() && !self.has_live_tail() { + self.touch_activity(); + self.status = "done".to_string(); + self.set_latest_preview("Turn done".to_string()); + return; + } + if !is_terminal_status(&normalized) { + self.status = normalized; self.set_latest_preview(format!("Status {}", self.status)); } } @@ -1132,10 +1226,15 @@ fn is_active_subagent_status(status: &str) -> bool { fn normalize_subagent_display_status(status: &str) -> String { match status.to_ascii_lowercase().as_str() { "running" | "spawning" | "activeturn" => "working".to_string(), + "completed" | "idle" | "success" => "done".to_string(), other => other.to_string(), } } +fn normalize_terminal_subagent_status(status: &str) -> String { + normalize_subagent_display_status(status) +} + fn normalize_runtime_status_for_subagent(status: devo_protocol::SessionRuntimeStatus) -> String { match status { devo_protocol::SessionRuntimeStatus::ActiveTurn => "working".to_string(), diff --git a/crates/tui/src/chatwidget/text_stream.rs b/crates/tui/src/chatwidget/text_stream.rs index ff8ef17b..a52c59ec 100644 --- a/crates/tui/src/chatwidget/text_stream.rs +++ b/crates/tui/src/chatwidget/text_stream.rs @@ -114,9 +114,10 @@ impl ChatWidget { } let seq = self.reserve_seq(); - let stream_controller = match kind { - TextItemKind::Assistant => Some(StreamController::new(None, &self.session.cwd)), - TextItemKind::Reasoning | TextItemKind::ResearchArtifact => None, + let stream_controller = if uses_markdown_stream_controller(kind, research.as_ref()) { + Some(StreamController::new(None, &self.session.cwd)) + } else { + None }; let insert_index = self.active_text_item_insert_index(kind); tracing::debug!( @@ -213,8 +214,20 @@ impl ChatWidget { self.active_text_items[index].raw_text.push_str(delta); } TextItemKind::ResearchArtifact => { - self.active_text_items[index].raw_text.push_str(delta); - self.sync_research_task_preview(index); + let item = &mut self.active_text_items[index]; + if item.is_delegated_research_finding() { + item.raw_text.push_str(delta); + self.sync_research_task_preview(index); + } else if let Some(controller) = item.stream_controller.as_mut() { + let produced_renderable_lines = controller.push(delta); + item.raw_text = controller.live_source(); + if produced_renderable_lines { + item.last_renderable_delta_at = Some(Instant::now()); + item.stream_stall_warned = false; + } + } else { + item.raw_text.push_str(delta); + } } } self.sync_text_item_cell(index); @@ -290,10 +303,18 @@ impl ChatWidget { .iter() .position(|item| item.item_id == item_id) { - if let Some(research) = research - && self.active_text_items[index].research.is_none() - { - self.active_text_items[index].research = Some(research); + if let Some(research) = research { + let item = &mut self.active_text_items[index]; + if item.research.is_none() { + item.research = Some(research.clone()); + } + if research.is_delegated_finding() { + item.stream_controller = None; + } else if uses_markdown_stream_controller(kind, Some(&research)) + && item.stream_controller.is_none() + { + item.stream_controller = Some(StreamController::new(None, &self.session.cwd)); + } } return index; } @@ -365,6 +386,9 @@ impl ChatWidget { self.remove_research_task_preview(item.item_id); return; } + if let Some(controller) = item.stream_controller.as_mut() { + let _ = controller.finalize(); + } if !item.raw_text.trim().is_empty() { self.add_history_entry_without_redraw(Box::new(ResearchArtifactCell::new( "Research", @@ -391,8 +415,14 @@ impl ChatWidget { } fn active_text_item_insert_index(&self, kind: TextItemKind) -> usize { - let _ = kind; - self.active_text_items.len() + match kind { + TextItemKind::Reasoning | TextItemKind::ResearchArtifact => self + .active_text_items + .iter() + .position(|item| item.kind == TextItemKind::Assistant) + .unwrap_or(self.active_text_items.len()), + TextItemKind::Assistant => self.active_text_items.len(), + } } fn commit_completed_text_items(&mut self) { @@ -466,12 +496,15 @@ impl ChatWidget { all_idle = output.all_idle, "stream commit tick processed active text item" ); - if item.kind == TextItemKind::Assistant { + if matches!( + item.kind, + TextItemKind::Assistant | TextItemKind::ResearchArtifact + ) { if !output.cells.is_empty() { changed_indexes.push(index); item.last_stream_commit_at = Some(now); item.stream_stall_warned = false; - } else { + } else if item.kind == TextItemKind::Assistant { maybe_warn_stream_commit_stall(item, queued_lines_after, now); } if !output.all_idle { @@ -582,13 +615,18 @@ impl ChatWidget { if item.is_delegated_research_finding() { return None; } - if item.raw_text.trim().is_empty() { + let markdown_source = if let Some(controller) = &item.stream_controller { + controller.live_source() + } else { + item.raw_text.clone() + }; + if markdown_source.trim().is_empty() { return None; } Some(Box::new(ResearchArtifactCell::new( "Research", - item.raw_text.clone(), + markdown_source, &self.session.cwd, ))) } @@ -623,6 +661,7 @@ impl ChatWidget { preview, }); } + self.invalidate_subagent_live_list_cache(); } fn remove_research_task_preview(&mut self, item_id: ActiveTextItemId) { @@ -631,6 +670,7 @@ impl ChatWidget { }; self.research_task_previews .retain(|preview| preview.item_id != item_id); + self.invalidate_subagent_live_list_cache(); } } @@ -644,6 +684,19 @@ impl ActiveTextItem { } } +fn uses_markdown_stream_controller( + kind: TextItemKind, + research: Option<&ResearchArtifactMetadata>, +) -> bool { + match kind { + TextItemKind::Assistant => true, + TextItemKind::ResearchArtifact => { + !research.is_some_and(ResearchArtifactMetadata::is_delegated_finding) + } + TextItemKind::Reasoning => false, + } +} + fn maybe_warn_stream_commit_stall(item: &mut ActiveTextItem, queued_lines: usize, now: Instant) { if item.kind != TextItemKind::Assistant || item.stream_stall_warned || queued_lines == 0 { return; diff --git a/crates/tui/src/chatwidget/transcript_view.rs b/crates/tui/src/chatwidget/transcript_view.rs index a735caea..6f6d7aa3 100644 --- a/crates/tui/src/chatwidget/transcript_view.rs +++ b/crates/tui/src/chatwidget/transcript_view.rs @@ -6,6 +6,7 @@ use ratatui::text::Line; use ratatui::text::Span; +use crate::events::TextItemKind; use crate::history_cell; use crate::history_cell::HistoryCell; use crate::history_cell::ScrollbackLine; @@ -37,6 +38,12 @@ enum LiveViewportLineMode { Transcript, } +#[allow(clippy::large_enum_variant)] +enum LiveItem { + Text(usize), + Tool(String), +} + impl ChatWidget { pub(crate) fn active_cell_transcript_key(&self) -> Option { let active_cell = self.active_cell.as_ref()?; @@ -172,11 +179,6 @@ impl ChatWidget { Self::extend_lines_with_separator(&mut lines, cell_lines(cell.as_ref())); } - #[allow(clippy::large_enum_variant)] - enum LiveItem { - Text(usize), - Tool(String), - } let mut items: Vec<(u64, LiveItem)> = Vec::new(); for (idx, item) in self.active_text_items.iter().enumerate() { if item.cell.is_some() { @@ -189,7 +191,15 @@ impl ChatWidget { } items.push((tool_call.seq, LiveItem::Tool(tool_call.tool_use_id.clone()))); } - items.sort_by(|a, b| a.0.cmp(&b.0)); + items.sort_by(|(seq_a, item_a), (seq_b, item_b)| { + Self::compare_live_viewport_items( + &self.active_text_items, + *seq_a, + item_a, + *seq_b, + item_b, + ) + }); for (_, item) in items { match item { @@ -313,6 +323,37 @@ impl ChatWidget { self.live_viewport_lines(width, LiveViewportLineMode::Transcript) } + fn text_item_precedes_assistant(kind: TextItemKind) -> bool { + matches!( + kind, + TextItemKind::Reasoning | TextItemKind::ResearchArtifact + ) + } + + fn compare_live_viewport_items( + active_text_items: &[super::text_stream::ActiveTextItem], + seq_a: u64, + item_a: &LiveItem, + seq_b: u64, + item_b: &LiveItem, + ) -> std::cmp::Ordering { + use std::cmp::Ordering; + + let text_kind = |item: &LiveItem| match item { + LiveItem::Text(idx) => active_text_items.get(*idx).map(|item| item.kind), + LiveItem::Tool(_) => None, + }; + if let (Some(kind_a), Some(kind_b)) = (text_kind(item_a), text_kind(item_b)) { + if Self::text_item_precedes_assistant(kind_a) && kind_b == TextItemKind::Assistant { + return Ordering::Less; + } + if kind_a == TextItemKind::Assistant && Self::text_item_precedes_assistant(kind_b) { + return Ordering::Greater; + } + } + seq_a.cmp(&seq_b) + } + fn extend_lines_with_separator(target: &mut Vec>, mut next: Vec>) { if next.is_empty() { return; diff --git a/crates/tui/src/chatwidget_tests.rs b/crates/tui/src/chatwidget_tests.rs index 556eb24d..6b4bc4b1 100644 --- a/crates/tui/src/chatwidget_tests.rs +++ b/crates/tui/src/chatwidget_tests.rs @@ -967,6 +967,73 @@ fn approval_request_does_not_duplicate_already_committed_assistant_text() { ); } +#[test] +fn research_clarification_artifact_streams_incremental_deltas_before_completion() { + let model = Model { + slug: "test-model".to_string(), + display_name: "Test Model".to_string(), + ..Model::default() + }; + let (mut widget, _app_event_rx) = widget_with_model(model, PathBuf::from(".")); + let artifact_id = ItemId::new(); + let research = Some(ResearchArtifactMetadata { + artifact_type: "clarification".to_string(), + title: "Research Clarification".to_string(), + }); + + widget.handle_worker_event(crate::events::WorkerEvent::TurnStarted { + model: "test-model".to_string(), + model_binding_id: None, + reasoning_effort_selection: None, + reasoning_effort: None, + turn_id: Default::default(), + }); + widget.handle_worker_event(crate::events::WorkerEvent::TextItemStarted { + item_id: artifact_id, + kind: crate::events::TextItemKind::ResearchArtifact, + research: research.clone(), + }); + widget.handle_worker_event(crate::events::WorkerEvent::TextItemDelta { + item_id: artifact_id, + kind: crate::events::TextItemKind::ResearchArtifact, + research: research.clone(), + delta: "Research ".to_string(), + }); + let partial_rows = rendered_rows(&widget, 80, 16).join("\n"); + assert!( + partial_rows.contains("Research "), + "clarification artifact delta should be visible before completion:\n{partial_rows}" + ); + + widget.handle_worker_event(crate::events::WorkerEvent::TextItemDelta { + item_id: artifact_id, + kind: crate::events::TextItemKind::ResearchArtifact, + research: research.clone(), + delta: "DeepSeek official website.".to_string(), + }); + let streamed_rows = rendered_rows(&widget, 80, 16).join("\n"); + assert!( + streamed_rows.contains("DeepSeek official website."), + "clarification artifact should grow with later deltas:\n{streamed_rows}" + ); + + widget.handle_worker_event(crate::events::WorkerEvent::TextItemCompleted { + item_id: artifact_id, + kind: crate::events::TextItemKind::ResearchArtifact, + research, + final_text: "### Research Clarification\n\nResearch DeepSeek official website.".to_string(), + }); + + let committed = scrollback_plain_lines(&trim_trailing_blank_scrollback_lines( + widget.drain_scrollback_lines(80), + )) + .join("\n"); + assert!( + committed.contains("Research DeepSeek official website."), + "clarification artifact should commit final content:\n{committed}" + ); +} + #[test] fn research_artifact_completion_does_not_commit_assistant_turn() { let model = Model { @@ -5807,10 +5874,10 @@ fn terminal_subagent_status_hides_ctrl_x_hint_when_no_live_children_remain() { assert!(!widget.has_live_subagents_for_test()); assert!(!widget.is_subagent_monitor_open_for_test()); let rows = rendered_rows(&widget, 100, 18).join("\n"); - assert!(rows.contains("builder: completed"), "rows:\n{rows}"); + assert!(rows.contains("builder: done"), "rows:\n{rows}"); widget.expire_subagent_inactivity_for_test(); let rows = rendered_rows(&widget, 100, 18).join("\n"); - assert!(!rows.contains("builder: completed"), "rows:\n{rows}"); + assert!(!rows.contains("builder: done"), "rows:\n{rows}"); assert!(!rows.contains("ctrl + x agents"), "rows:\n{rows}"); } diff --git a/crates/tui/src/custom_terminal.rs b/crates/tui/src/custom_terminal.rs index 29c3487b..1853922f 100644 --- a/crates/tui/src/custom_terminal.rs +++ b/crates/tui/src/custom_terminal.rs @@ -283,7 +283,7 @@ where if result.is_ok() { self.last_flush_stats = stats; } - tracing::debug!( + tracing::trace!( stream_elapsed_ms = terminal_trace_elapsed_ms(), frame_seq = stats.frame_seq, diff_commands = stats.diff_commands, diff --git a/crates/tui/src/interactive.rs b/crates/tui/src/interactive.rs index bdfdd363..52197c48 100644 --- a/crates/tui/src/interactive.rs +++ b/crates/tui/src/interactive.rs @@ -244,6 +244,15 @@ pub async fn run_interactive_tui(config: InteractiveTuiConfig) -> Result { if chat_widget.handle_onboarding_key_event(key) { diff --git a/crates/tui/src/streaming/controller.rs b/crates/tui/src/streaming/controller.rs index 4637ad8f..9865b724 100644 --- a/crates/tui/src/streaming/controller.rs +++ b/crates/tui/src/streaming/controller.rs @@ -288,6 +288,10 @@ impl StreamController { self.core.set_width(width); } + pub(crate) fn live_source(&self) -> String { + self.core.live_source() + } + pub(crate) fn live_lines(&self) -> Vec> { let source = self.core.live_source(); if source.is_empty() { diff --git a/crates/tui/src/worker.rs b/crates/tui/src/worker.rs index 14bd2f05..b50fe5c1 100644 --- a/crates/tui/src/worker.rs +++ b/crates/tui/src/worker.rs @@ -78,6 +78,7 @@ use devo_server::TurnEventPayload; use devo_server::TurnExecutionMode; use devo_server::TurnInterruptParams; use devo_server::TurnStartParams; +use devo_server::TurnStartResult; use devo_server::TurnSteerParams; use crate::app_command::GoalObjectiveMode; @@ -106,6 +107,7 @@ use acp_events::session_metadata_from_acp_update; use acp_events::spawn_agent_result_from_acp_update; use acp_events::spawn_task_message_from_acp_update; use acp_events::subagent_monitor_events_from_acp_session_notification_with_terminal_state; +use acp_events::subagent_monitor_events_from_unwrapped_server_notification; #[cfg(test)] use acp_events::worker_events_from_acp_notification; #[cfg(test)] @@ -214,6 +216,8 @@ pub(crate) struct QueryWorkerConfig { pub(crate) reasoning_effort_selection: Option, /// Permission preset to apply to the server session when it exists. pub(crate) permission_preset: PermissionPreset, + /// Agent client capabilities to advertise to the server session. + pub(crate) client_capabilities: devo_protocol::AcpClientCapabilities, } /// TODO: Should we extract the OperationCommand to the `protocol` crate? Since it can be shareable. @@ -820,7 +824,7 @@ async fn run_worker_inner( // The worker owns the server client and translates UI commands into server // calls, then turns server notifications back into lightweight UI events. let mut client = spawn_client(&config.cwd, config.server_log_level.clone()).await?; - let _ = client.initialize().await?; + let _ = client.initialize(&config.client_capabilities).await?; let mut session_id: Option = None; let mut session_cwd = config.cwd.clone(); let mut model = config.model; @@ -926,37 +930,17 @@ async fn run_worker_inner( approval_policy, collaboration_mode, }) => { - let session_start = ensure_session_started( + let active_session_id = prepare_session_for_command( &mut client, &config.cwd, - &model, - &model_binding_id, + &mut model, + &mut model_binding_id, + &mut reasoning_effort_selection, &mut session_id, + permission_preset, + event_tx, ) .await?; - if let Some(start_model) = session_start.model.clone() { - model = start_model; - } - model_binding_id = session_start - .model_binding_id - .clone() - .or(model_binding_id); - reasoning_effort_selection = session_start - .reasoning_effort_selection - .clone() - .or(reasoning_effort_selection); - let active_session_id = session_start.session_id; - if session_start.created { - let _ = event_tx.send(WorkerEvent::SessionActivated { - session_id: active_session_id, - }); - apply_session_permissions( - &mut client, - active_session_id, - permission_preset, - ) - .await?; - } let start_result = client.turn_start(TurnStartParams { session_id: active_session_id, input, @@ -970,7 +954,9 @@ async fn run_worker_inner( execution_mode: TurnExecutionMode::Regular, }).await; match start_result { - Ok(()) => {} + Ok(result) => { + handle_turn_start_result(result, &mut active_turn_id); + } Err(error) => { let _ = event_tx.send(WorkerEvent::TurnFailed { message: error.to_string(), @@ -1176,7 +1162,7 @@ async fn run_worker_inner( config.server_log_level.clone(), ) .await?; - client.initialize().await?; + client.initialize(&config.client_capabilities).await?; session_id = None; child_agent_sessions.clear(); btw_agent_sessions.clear(); @@ -1388,34 +1374,17 @@ async fn run_worker_inner( } } Some(OperationCommand::SetGoalObjective { objective, mode }) => { - let session_start = ensure_session_started( + let active_session_id = prepare_session_for_command( &mut client, &config.cwd, - &model, - &model_binding_id, + &mut model, + &mut model_binding_id, + &mut reasoning_effort_selection, &mut session_id, + permission_preset, + event_tx, ) .await?; - if let Some(start_model) = session_start.model.clone() { - model = start_model; - } - model_binding_id = session_start - .model_binding_id - .clone() - .or(model_binding_id); - reasoning_effort_selection = session_start.reasoning_effort_selection.clone().or(reasoning_effort_selection); - let active_session_id = session_start.session_id; - if session_start.created { - let _ = event_tx.send(WorkerEvent::SessionActivated { - session_id: active_session_id, - }); - apply_session_permissions( - &mut client, - active_session_id, - permission_preset, - ) - .await?; - } if matches!(mode, GoalObjectiveMode::ConfirmIfExists) { match client @@ -2029,34 +1998,17 @@ async fn run_worker_inner( } } Some(OperationCommand::RunResearch { question }) => { - let session_start = ensure_session_started( + let active_session_id = prepare_session_for_command( &mut client, &config.cwd, - &model, - &model_binding_id, + &mut model, + &mut model_binding_id, + &mut reasoning_effort_selection, &mut session_id, + permission_preset, + event_tx, ) .await?; - if let Some(start_model) = session_start.model.clone() { - model = start_model; - } - model_binding_id = session_start - .model_binding_id - .clone() - .or(model_binding_id); - reasoning_effort_selection = session_start.reasoning_effort_selection.clone().or(reasoning_effort_selection); - let active_session_id = session_start.session_id; - if session_start.created { - let _ = event_tx.send(WorkerEvent::SessionActivated { - session_id: active_session_id, - }); - apply_session_permissions( - &mut client, - active_session_id, - permission_preset, - ) - .await?; - } match client .turn_start(TurnStartParams { session_id: active_session_id, @@ -2072,7 +2024,9 @@ async fn run_worker_inner( }) .await { - Ok(()) => {} + Ok(result) => { + handle_turn_start_result(result, &mut active_turn_id); + } Err(error) => { let _ = event_tx.send(WorkerEvent::TurnFailed { message: error.to_string(), @@ -2439,6 +2393,16 @@ async fn run_worker_inner( if let Some(event_session_id) = event.session_id() && Some(event_session_id) != session_id { + if child_agent_sessions.contains(&event_session_id) { + for subagent_event in + subagent_monitor_events_from_unwrapped_server_notification( + method.as_str(), + event.clone(), + ) + { + let _ = event_tx.send(subagent_event); + } + } continue; } match method.as_str() { @@ -3007,6 +2971,56 @@ async fn ensure_session_started( }) } +/// Prepares the worker session state before turn or goal commands run. +/// +/// Commands such as [`OperationCommand::SubmitInput`], [`OperationCommand::SetGoalObjective`], +/// and [`OperationCommand::RunResearch`] share this path instead of duplicating session-start +/// follow-up. When no session is active yet, [`ensure_session_started`] creates one on the +/// server; the returned metadata is merged into the worker's current model, model binding, and +/// reasoning-effort selection. For a newly created session, this also notifies the UI via +/// [`WorkerEvent::SessionActivated`] and applies the configured permission preset. +#[allow(clippy::too_many_arguments)] +async fn prepare_session_for_command( + client: &mut StdioServerClient, + cwd: &Path, + model: &mut String, + model_binding_id: &mut Option, + reasoning_effort_selection: &mut Option, + session_id: &mut Option, + permission_preset: PermissionPreset, + event_tx: &mpsc::UnboundedSender, +) -> Result { + let session_start = + ensure_session_started(client, cwd, model, model_binding_id, session_id).await?; + if let Some(model_override) = &session_start.model { + *model = model_override.clone(); + } + *model_binding_id = session_start + .model_binding_id + .clone() + .or_else(|| model_binding_id.clone()); + *reasoning_effort_selection = session_start + .reasoning_effort_selection + .clone() + .or_else(|| reasoning_effort_selection.clone()); + let active_session_id = session_start.session_id; + if session_start.created { + let _ = event_tx.send(WorkerEvent::SessionActivated { + session_id: active_session_id, + }); + apply_session_permissions(client, active_session_id, permission_preset).await?; + } + Ok(active_session_id) +} + +/// Records the active turn returned by `turn/start`. +/// +/// When the server queues input (`TurnStartResult::Queued`), queue state for the UI is +/// delivered asynchronously via `inputQueue/updated` notifications. +fn handle_turn_start_result(result: TurnStartResult, active_turn_id: &mut Option) { + *active_turn_id = Some(result.active_turn_id()); +} + async fn pause_active_goal_before_session_leave( client: &mut StdioServerClient, session_id: SessionId, @@ -3096,15 +3110,6 @@ async fn spawn_client(_cwd: &Path, server_log_level: Option) -> Result, terminal_session_ids: &mut HashMap, ) -> Vec { + if let Some(events) = subagent_monitor_events_from_wrapped_server_event(¬ification) { + return events; + } let session_id = notification.session_id; match notification.update { AcpSessionUpdate::AgentMessageChunk { @@ -479,6 +486,90 @@ pub(super) fn subagent_monitor_events_from_acp_session_notification_with_termina } } +fn subagent_monitor_events_from_wrapped_server_event( + notification: &AcpSessionNotification, +) -> Option> { + let (method, event) = original_event_from_acp_notification(notification)?; + Some(subagent_monitor_events_from_server_event( + notification.session_id, + method.as_str(), + event, + )) +} + +fn subagent_monitor_events_from_server_event( + session_id: SessionId, + method: &str, + event: ServerEvent, +) -> Vec { + match (method, event) { + ("turn/started", ServerEvent::TurnStarted(payload)) => { + vec![WorkerEvent::SubagentMonitor { + event: SubagentMonitorEvent::TurnStarted { + session_id, + turn_id: payload.turn.turn_id, + }, + }] + } + ("turn/completed", ServerEvent::TurnCompleted(payload)) + | ("turn/interrupted", ServerEvent::TurnInterrupted(payload)) => { + vec![WorkerEvent::SubagentMonitor { + event: SubagentMonitorEvent::TurnFinished { + session_id, + status: subagent_turn_finished_status(payload.turn.status), + }, + }] + } + ("turn/failed", ServerEvent::TurnFailed(payload)) => { + vec![WorkerEvent::SubagentMonitor { + event: SubagentMonitorEvent::TurnFailed { + session_id, + message: subagent_turn_failure_message(&payload.turn), + }, + }] + } + ("session/status/changed", ServerEvent::SessionStatusChanged(payload)) => { + vec![WorkerEvent::SubagentMonitor { + event: SubagentMonitorEvent::SessionStatusChanged { + session_id, + status: payload.status, + }, + }] + } + _ => Vec::new(), + } +} + +fn subagent_turn_finished_status(status: TurnStatus) -> String { + match status { + TurnStatus::Completed => "done".to_string(), + TurnStatus::Interrupted => "interrupted".to_string(), + TurnStatus::Failed => "failed".to_string(), + TurnStatus::Running | TurnStatus::Pending | TurnStatus::WaitingApproval => { + "working".to_string() + } + } +} + +fn subagent_turn_failure_message(turn: &TurnMetadata) -> String { + turn.failure_reason + .as_ref() + .map(|reason| format!("{reason:?}")) + .unwrap_or_else(|| "Turn failed".to_string()) +} + +/// Routes unwrapped server notifications (as produced by the stdio client) to +/// sub-agent monitor events for a known child session. +pub(super) fn subagent_monitor_events_from_unwrapped_server_notification( + method: &str, + event: ServerEvent, +) -> Vec { + let Some(session_id) = event.session_id() else { + return Vec::new(); + }; + subagent_monitor_events_from_server_event(session_id, method, event) +} + fn message_item_id(message_id: Option<&str>) -> Option { message_id.and_then(|message_id| ItemId::try_from(message_id).ok()) } From f730cc0794a382d141473fe1c7080ef585553526 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Thu, 2 Jul 2026 17:30:06 -1000 Subject: [PATCH 04/16] fix: validate inbound client messages at the transport layer Reject malformed JSON-RPC payloads with ParseError/InvalidRequest responses before handler dispatch, and document stdio NDJSON framing. --- crates/server/src/transport.rs | 288 +++++++++++++++++++++++++++++---- 1 file changed, 260 insertions(+), 28 deletions(-) diff --git a/crates/server/src/transport.rs b/crates/server/src/transport.rs index ce39c0ce..f3b02b35 100644 --- a/crates/server/src/transport.rs +++ b/crates/server/src/transport.rs @@ -22,8 +22,10 @@ use tokio_tungstenite::accept_async; use tokio_tungstenite::tungstenite::Message; use tokio_util::sync::CancellationToken; +use crate::AcpErrorCode; use crate::ClientTransportKind; use crate::ServerRuntime; +use crate::acp_error_response; use crate::runtime::CONNECTION_NOTIFICATION_CHANNEL_CAPACITY; use crate::runtime::IncomingResponse; use crate::singleton::SERVER_CONTROL_SHUTDOWN_METHOD; @@ -268,8 +270,16 @@ async fn run_listener_tasks( Ok(()) } +/// Stdio uses **NDJSON** (newline-delimited JSON): one JSON-RPC message per line. +/// +/// Outbound messages are serialized with `serde_json::to_vec`, which escapes +/// embedded newlines in string values as `\n`. The trailing `\n` written by the +/// stdout task is therefore a frame delimiter only, not part of the payload. +/// Clients must send the same framing: each request, notification, or response +/// must occupy exactly one line of valid JSON. Pretty-printed or otherwise +/// multi-line payloads are rejected at the transport layer. async fn run_stdio(runtime: Arc) -> Result<()> { - // Normal channel for event notifications (TextDelta, TurnStarted, …). + // Server → client responses and notifications (TextDelta, TurnStarted, …). let (sender, mut receiver) = mpsc::channel(CONNECTION_NOTIFICATION_CHANNEL_CAPACITY); let sender_clone = sender.clone(); let connection_id = runtime @@ -277,9 +287,8 @@ async fn run_stdio(runtime: Arc) -> Result<()> { .await; tracing::info!(connection_id, "stdio connection established"); - // Internal channel between the producer (reads from high_pri + normal) - // and the writer (writes to stdout). Bounded sends apply backpressure - // instead of allowing an unbounded stdout backlog. + // Serialized NDJSON lines awaiting stdout. Bounded capacity applies + // backpressure when stdout is slow instead of buffering without limit. let (write_tx, mut write_rx) = mpsc::channel::>(TRANSPORT_WRITE_CHANNEL_CAPACITY); // --- Writer task --- @@ -296,19 +305,12 @@ async fn run_stdio(runtime: Arc) -> Result<()> { }); // --- Producer task --- - // Reads from high_pri and normal channels, serializes immediately, - // and pushes to write_tx. A slow writer backpressures this task. + // Serializes outbound messages and forwards them to the writer task. let producer_task = tokio::spawn(async move { - loop { - let line: Vec; - tokio::select! { - Some(message) = receiver.recv() => { - line = serde_json::to_vec(&message) - .expect("serialize stdio response"); - } - else => break, - } - if !send_transport_queue_message(&write_tx, line, connection_id, "stdio_write").await { + while let Some(message) = receiver.recv().await { + let line = serde_json::to_vec(&message).expect("serialize stdio response"); + if !send_transport_queue_message(&write_tx, line, connection_id, "stdio_write").await + { break; } } @@ -318,15 +320,11 @@ async fn run_stdio(runtime: Arc) -> Result<()> { let stdin = tokio::io::stdin(); let mut lines = BufReader::new(stdin).lines(); while let Some(line) = lines.next_line().await? { - if line.trim().is_empty() { - continue; - } - let value: serde_json::Value = serde_json::from_str(&line)?; - spawn_incoming_message_handler( + accept_incoming_client_message( Arc::clone(&runtime), connection_id, sender_clone.clone(), - value, + &line, "stdio_notifications", ); } @@ -443,12 +441,11 @@ async fn handle_internal_proxy_connection( let frame = frame?; match frame { Message::Text(text) => { - let value: serde_json::Value = serde_json::from_str(&text)?; - spawn_incoming_message_handler( + accept_incoming_client_message( Arc::clone(&runtime), connection_id, sender_clone.clone(), - value, + text.as_str(), "internal_proxy_notifications", ); } @@ -555,12 +552,11 @@ async fn handle_websocket_connection( let frame = frame?; match frame { Message::Text(text) => { - let value: serde_json::Value = serde_json::from_str(&text)?; - spawn_incoming_message_handler( + accept_incoming_client_message( Arc::clone(&runtime), connection_id, sender_clone.clone(), - value, + text.as_str(), "websocket_notifications", ); } @@ -575,6 +571,161 @@ async fn handle_websocket_connection( Ok(()) } +/// Parses one inbound client payload (NDJSON line or WebSocket text frame). +/// +/// Returns `None` when the payload is empty after trimming. Returns `Err` with +/// a JSON-RPC `ParseError` response when the payload is not valid JSON. +fn parse_incoming_client_payload( + raw_payload: &str, +) -> Option> { + let trimmed = raw_payload.trim(); + if trimmed.is_empty() { + return None; + } + Some(match serde_json::from_str(trimmed) { + Ok(value) => Ok(value), + Err(error) => Err(acp_error_response( + serde_json::Value::Null, + AcpErrorCode::ParseError, + format!("malformed client payload: {error}"), + )), + }) +} + +/// Validates a decoded client JSON-RPC 2.0 message before handler dispatch. +/// +/// Returns `Ok(())` for well-formed requests, notifications, or responses to +/// server-initiated calls. Returns `Err(error_response)` when the payload is +/// structurally invalid; the caller should send that response and skip the +/// handler. +fn validate_incoming_client_message( + value: &serde_json::Value, +) -> Result<(), serde_json::Value> { + let Some(object) = value.as_object() else { + return Err(acp_error_response( + serde_json::Value::Null, + AcpErrorCode::InvalidRequest, + "client message must be a JSON object", + )); + }; + + let request_id = object + .get("id") + .cloned() + .unwrap_or(serde_json::Value::Null); + + match object.get("jsonrpc").and_then(serde_json::Value::as_str) { + Some("2.0") => {} + Some(_) => { + return Err(acp_error_response( + request_id, + AcpErrorCode::InvalidRequest, + "jsonrpc must be \"2.0\"", + )); + } + None => { + return Err(acp_error_response( + request_id, + AcpErrorCode::InvalidRequest, + "jsonrpc field is required", + )); + } + } + + let has_method = object.contains_key("method"); + let has_result = object.contains_key("result"); + let has_error = object.contains_key("error"); + + if has_result && has_error { + return Err(acp_error_response( + request_id, + AcpErrorCode::InvalidRequest, + "client message must not contain both result and error", + )); + } + + if !has_method + && object.contains_key("id") + && (has_result || has_error) + { + return Ok(()); + } + + if has_method { + if has_result || has_error { + return Err(acp_error_response( + request_id, + AcpErrorCode::InvalidRequest, + "client request must not contain result or error", + )); + } + let Some(method) = object.get("method").and_then(serde_json::Value::as_str) else { + return Err(acp_error_response( + request_id, + AcpErrorCode::InvalidRequest, + "method must be a non-empty string", + )); + }; + if method.is_empty() { + return Err(acp_error_response( + request_id, + AcpErrorCode::InvalidRequest, + "method must be a non-empty string", + )); + } + if let Some(params) = object.get("params") + && !params.is_object() + && !params.is_array() + && !params.is_null() + { + return Err(acp_error_response( + request_id, + AcpErrorCode::InvalidRequest, + "params must be an object, array, or null", + )); + } + return Ok(()); + } + + Err(acp_error_response( + request_id, + AcpErrorCode::InvalidRequest, + "client message must be a request, notification, or response", + )) +} + +/// Parses, validates, and dispatches one inbound client payload without blocking +/// the transport read loop. +fn accept_incoming_client_message( + runtime: Arc, + connection_id: u64, + sender: mpsc::Sender, + raw_payload: &str, + queue: &'static str, +) { + let Some(parsed) = parse_incoming_client_payload(raw_payload) else { + return; + }; + let value = match parsed { + Ok(value) => value, + Err(error_response) => { + tracing::warn!(connection_id, "rejected malformed client payload"); + tokio::spawn(async move { + send_transport_queue_message(&sender, error_response, connection_id, queue).await; + }); + return; + } + }; + if let Err(error_response) = validate_incoming_client_message(&value) { + tracing::warn!(connection_id, "rejected invalid client message"); + tokio::spawn(async move { + send_transport_queue_message(&sender, error_response, connection_id, queue).await; + }); + return; + } + spawn_incoming_message_handler(runtime, connection_id, sender, value, queue); +} + /// Dispatch a single decoded client message without blocking the connection's /// read loop. /// @@ -675,8 +826,11 @@ mod tests { use super::ListenTarget; use super::internal_proxy_control_response; use super::parse_internal_proxy_control_request; + use super::parse_incoming_client_payload; use super::parse_listen_target; use super::resolve_listen_targets; + use super::validate_incoming_client_message; + use crate::AcpErrorCode; use pretty_assertions::assert_eq; use std::sync::Arc; use std::time::Duration; @@ -770,6 +924,84 @@ mod tests { ); } + #[test] + fn parse_incoming_client_payload_skips_empty_and_parses_json() { + assert_eq!(parse_incoming_client_payload(""), None); + assert_eq!(parse_incoming_client_payload(" \n"), None); + + let parsed = parse_incoming_client_payload(r#" {"jsonrpc":"2.0","id":1} "#) + .expect("non-empty payload") + .expect("valid json"); + assert_eq!(parsed, serde_json::json!({ "jsonrpc": "2.0", "id": 1 })); + } + + #[test] + fn parse_incoming_client_payload_returns_parse_error_response() { + let error_response = parse_incoming_client_payload("{not json}") + .expect("non-empty payload") + .expect_err("invalid json"); + assert_eq!( + error_response["error"]["code"], + AcpErrorCode::ParseError as i64 + ); + assert_eq!(error_response["id"], serde_json::Value::Null); + } + + #[test] + fn validate_accepts_client_request_notification_and_response() { + assert!(validate_incoming_client_message(&serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": {}, + })) + .is_ok()); + assert!(validate_incoming_client_message(&serde_json::json!({ + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": {}, + })) + .is_ok()); + assert!(validate_incoming_client_message(&serde_json::json!({ + "jsonrpc": "2.0", + "id": 9, + "result": { "ok": true }, + })) + .is_ok()); + } + + #[test] + fn validate_rejects_malformed_client_messages() { + let invalid_request = validate_incoming_client_message(&serde_json::json!([])) + .expect_err("array payload"); + assert_eq!( + invalid_request["error"]["code"], + AcpErrorCode::InvalidRequest as i64 + ); + + let missing_jsonrpc = validate_incoming_client_message(&serde_json::json!({ + "id": 1, + "method": "initialize", + })) + .expect_err("missing jsonrpc"); + assert_eq!( + missing_jsonrpc["error"]["code"], + AcpErrorCode::InvalidRequest as i64 + ); + + let conflicting_fields = validate_incoming_client_message(&serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "result": {}, + })) + .expect_err("request with result"); + assert_eq!( + conflicting_fields["error"]["code"], + AcpErrorCode::InvalidRequest as i64 + ); + } + #[tokio::test] async fn event_broadcaster_backpressures_full_sender() { let broadcaster = Arc::new(EventBroadcaster::new()); From 24183bd06b057b1fc1ebbf6594728fefe66d7876 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Thu, 2 Jul 2026 17:57:52 -1000 Subject: [PATCH 05/16] Consolidate per-connection outbound pipeline and bound inbound handler concurrency. Replace the unbounded notification queue and multi-stage stdio writer with a single bounded OutboundFrame channel and OutboundWriter task, add semaphore-limited inbound dispatch with a fast path for client responses, and document that ACP session/prompt remains blocking while event-driven clients use _devo/turn/start. --- crates/protocol/README.md | 8 +- crates/server/src/runtime.rs | 9 +- crates/server/src/runtime/connection.rs | 386 ++++-------------- crates/server/src/runtime/outbound.rs | 328 +++++++++++++++ crates/server/src/transport.rs | 381 ++++++++++------- crates/server/tests/acp_available_commands.rs | 8 +- .../acp_permission_tool_status_contract.rs | 2 +- crates/server/tests/acp_session_delete.rs | 4 +- crates/server/tests/acp_session_lifecycle.rs | 4 +- crates/server/tests/cli_log_failures.rs | 2 +- crates/server/tests/command_exec.rs | 2 +- .../tests/deep_research_boundary_failure.rs | 2 +- crates/server/tests/deep_research_e2e.rs | 132 +----- crates/server/tests/goal_title_generation.rs | 2 +- crates/server/tests/persistence_resume.rs | 2 +- crates/server/tests/provider_routing.rs | 2 +- .../server/tests/session_fork_persistence.rs | 2 +- .../tests/session_rollback_persistence.rs | 2 +- crates/server/tests/skills_integration.rs | 2 +- .../server/tests/support/goal_continuation.rs | 2 +- .../tests/support/subagent_lifecycle.rs | 2 +- crates/server/tests/turn_start_persistence.rs | 2 +- 22 files changed, 688 insertions(+), 598 deletions(-) create mode 100644 crates/server/src/runtime/outbound.rs diff --git a/crates/protocol/README.md b/crates/protocol/README.md index 3ac9fc38..699f3569 100644 --- a/crates/protocol/README.md +++ b/crates/protocol/README.md @@ -13,9 +13,15 @@ client-to-server ACP methods are: - `session/new`: create a new session for a working directory. - `session/list`: list persisted sessions. - `session/resume`: load a persisted session. -- `session/prompt`: submit a prompt to an active session. +- `session/prompt`: submit a prompt to an active session. The JSON-RPC response + returns when the turn completes (`AcpPromptResult.stopReason`). Streaming + progress is delivered through `session/update` notifications during the turn. - `session/cancel`: cancel the active session turn. +Event-driven clients that need an immediate turn acknowledgement should use the +Devo extension `_devo/turn/start`, which returns `TurnStartResult::Started` +promptly and streams turn progress through server notifications. + The current server-to-client ACP notification method is: - `session/update`: stream session lifecycle, item, plan, usage, and turn-status diff --git a/crates/server/src/runtime.rs b/crates/server/src/runtime.rs index 98cd11b6..582df17a 100644 --- a/crates/server/src/runtime.rs +++ b/crates/server/src/runtime.rs @@ -152,6 +152,7 @@ mod agents; mod approval; mod command_exec; mod connection; +mod outbound; mod goal_accounting; mod goal_continuation; mod goal_handlers; @@ -185,8 +186,14 @@ mod turn_reservation; mod user_input; mod workspace_baseline; -pub(crate) use connection::CONNECTION_NOTIFICATION_CHANNEL_CAPACITY; +pub use outbound::test_outbound_channel; +pub(crate) use outbound::log_outbound_frame; +pub(crate) use outbound::outbound_frame_to_value; +pub(crate) use outbound::OUTBOUND_CHANNEL_CAPACITY; +pub use outbound::OutboundFrame; +pub(crate) use outbound::enqueue_outbound; pub(crate) use connection::ConnectionRuntime; +pub(crate) use connection::INBOUND_CONCURRENCY_LIMIT; pub use connection::IncomingResponse; pub use connection::PostResponseActions; pub(crate) use connection::SubscriptionFilter; diff --git a/crates/server/src/runtime/connection.rs b/crates/server/src/runtime/connection.rs index 9fe83066..a946ec31 100644 --- a/crates/server/src/runtime/connection.rs +++ b/crates/server/src/runtime/connection.rs @@ -4,7 +4,6 @@ use std::collections::HashMap; use std::sync::Arc; use std::sync::OnceLock; use std::time::Duration; -use std::time::Instant; use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; @@ -25,75 +24,10 @@ use crate::acp_auth_required_response; use crate::acp_notification_from_server_event; use crate::devo_extension_inner_method; -pub(crate) const CONNECTION_NOTIFICATION_CHANNEL_CAPACITY: usize = 4096; +use super::outbound::OutboundFrame; +use super::outbound::enqueue_outbound; -const CONNECTION_NOTIFICATION_BACKPRESSURE_LOG_THRESHOLD: Duration = Duration::from_millis(50); - -struct PendingConnectionNotification { - connection_id: u64, - kind: PendingConnectionMessageKind, - method: String, - event_seq: u64, - sender: mpsc::Sender, - value: serde_json::Value, -} - -struct QueuedConnectionNotification { - notification: PendingConnectionNotification, - delivered: Option>, -} - -fn spawn_connection_notification_writer( - mut queue_rx: mpsc::UnboundedReceiver, -) -> tokio::task::JoinHandle<()> { - tokio::spawn(async move { - while let Some(queued) = queue_rx.recv().await { - let sent = send_connection_notification(queued.notification).await; - if let Some(delivered) = queued.delivered { - let _ = delivered.send(sent); - } - } - }) -} - -fn enqueue_connection_notification( - queue_tx: &mpsc::UnboundedSender, - notification: PendingConnectionNotification, -) -> bool { - queue_tx - .send(QueuedConnectionNotification { - notification, - delivered: None, - }) - .is_ok() -} - -async fn enqueue_connection_notification_and_wait_for_delivery( - queue_tx: &mpsc::UnboundedSender, - notification: PendingConnectionNotification, -) -> bool { - let (delivered_tx, delivered_rx) = oneshot::channel(); - if queue_tx - .send(QueuedConnectionNotification { - notification, - delivered: Some(delivered_tx), - }) - .is_err() - { - return false; - } - match delivered_rx.await { - Ok(sent) => sent, - Err(_) => false, - } -} - -#[derive(Clone, Copy)] -enum PendingConnectionMessageKind { - Notification, - JsonRpcResponse, - ClientRequest, -} +pub(crate) const INBOUND_CONCURRENCY_LIMIT: usize = 64; #[derive(Debug)] pub struct IncomingResponse { @@ -195,12 +129,9 @@ impl ServerRuntime { pub async fn register_connection( self: &Arc, transport: ClientTransportKind, - sender: mpsc::Sender, + outbound_tx: mpsc::Sender, ) -> u64 { let connection_id = self.next_connection_id.fetch_add(1, Ordering::SeqCst); - let (notification_queue_tx, notification_queue_rx) = mpsc::unbounded_channel(); - let _notification_writer = - spawn_connection_notification_writer(notification_queue_rx); let mut connections = self.connections.lock().await; connections.insert( connection_id, @@ -209,8 +140,7 @@ impl ServerRuntime { state: ConnectionState::Connected, acp_authenticated: false, acp_client_capabilities: crate::AcpClientCapabilities::default(), - sender, - notification_queue_tx, + outbound_tx, opt_out_notification_methods: HashSet::new(), subscriptions: Vec::new(), next_event_seq: 1, @@ -660,6 +590,15 @@ impl ServerRuntime { .is_some_and(|connection| connection.state == ConnectionState::Ready) } + pub async fn resolve_client_response( + self: &Arc, + connection_id: u64, + message: serde_json::Value, + ) { + self.resolve_pending_client_response(connection_id, message) + .await; + } + pub(super) async fn emit_to_connection( &self, connection_id: u64, @@ -680,19 +619,12 @@ impl ServerRuntime { let event = event.with_seq(event_seq); let (method, value) = acp_notification_from_server_event(method, &event); Some(( - connection.notification_queue_tx.clone(), - PendingConnectionNotification { - connection_id, - kind: PendingConnectionMessageKind::Notification, - method, - event_seq, - sender: connection.sender.clone(), - value, - }, + connection.outbound_tx.clone(), + OutboundFrame::notification(connection_id, method, event_seq, value), )) }; - if let Some((queue_tx, notification)) = notification { - enqueue_connection_notification(&queue_tx, notification); + if let Some((outbound_tx, frame)) = notification { + let _ = enqueue_outbound(&outbound_tx, frame, "connection_notifications").await; } } @@ -726,21 +658,14 @@ impl ServerRuntime { let event = event.clone().with_seq(event_seq); let (method, value) = acp_notification_from_server_event(method, &event); Some(( - connection.notification_queue_tx.clone(), - PendingConnectionNotification { - connection_id: *connection_id, - kind: PendingConnectionMessageKind::Notification, - method, - event_seq, - sender: connection.sender.clone(), - value, - }, + connection.outbound_tx.clone(), + OutboundFrame::notification(*connection_id, method, event_seq, value), )) }) .collect::>() }; - for (queue_tx, notification) in notifications { - enqueue_connection_notification(&queue_tx, notification); + for (outbound_tx, frame) in notifications { + let _ = enqueue_outbound(&outbound_tx, frame, "connection_notifications").await; } } @@ -765,24 +690,29 @@ impl ServerRuntime { connection_id: u64, value: serde_json::Value, ) { - let (queue_tx, notification) = { + let (outbound_tx, frame) = { let connections = self.connections.lock().await; let Some(connection) = connections.get(&connection_id) else { return; }; + let (delivered_tx, delivered_rx) = oneshot::channel(); ( - connection.notification_queue_tx.clone(), - PendingConnectionNotification { - connection_id, - kind: PendingConnectionMessageKind::JsonRpcResponse, - method: "".to_string(), - event_seq: 0, - sender: connection.sender.clone(), - value, - }, + connection.outbound_tx.clone(), + ( + OutboundFrame::json_rpc_response_with_delivery( + connection_id, + value, + delivered_tx, + ), + delivered_rx, + ), ) }; - let _ = enqueue_connection_notification_and_wait_for_delivery(&queue_tx, notification).await; + let (frame, delivered_rx) = frame; + if !enqueue_outbound(&outbound_tx, frame, "connection_responses").await { + return; + } + let _ = delivered_rx.await; } pub(super) async fn send_request_to_connection_cancellable( @@ -828,7 +758,7 @@ impl ServerRuntime { timeout_duration: Option, cancel_token: CancellationToken, ) -> Result { - let (request_id, receiver, queue_tx, notification) = { + let (request_id, receiver, outbound_tx, frame) = { let mut connections = self.connections.lock().await; let Some(connection) = connections.get_mut(&connection_id) else { return Err("client connection does not exist".to_string()); @@ -846,15 +776,12 @@ impl ServerRuntime { ( request_id, rx, - connection.notification_queue_tx.clone(), - PendingConnectionNotification { + connection.outbound_tx.clone(), + OutboundFrame::client_request( connection_id, - kind: PendingConnectionMessageKind::ClientRequest, - method: method.to_string(), - event_seq: 0, - sender: connection.sender.clone(), + method.to_string(), value, - }, + ), ) }; let mut pending_request = PendingClientRequestGuard::new( @@ -862,7 +789,7 @@ impl ServerRuntime { connection_id, request_id, ); - if !enqueue_connection_notification(&queue_tx, notification) { + if !enqueue_outbound(&outbound_tx, frame, "connection_requests").await { pending_request.remove().await; return Err("client connection closed before request was sent".to_string()); } @@ -1116,8 +1043,7 @@ pub(crate) struct ConnectionRuntime { pub(crate) state: ConnectionState, pub(crate) acp_authenticated: bool, pub(crate) acp_client_capabilities: crate::AcpClientCapabilities, - pub(crate) sender: mpsc::Sender, - notification_queue_tx: mpsc::UnboundedSender, + pub(crate) outbound_tx: mpsc::Sender, pub(crate) opt_out_notification_methods: HashSet, pub(crate) subscriptions: Vec, next_event_seq: u64, @@ -1177,171 +1103,6 @@ impl SubscriptionFilter { } } -async fn send_connection_notification(notification: PendingConnectionNotification) -> bool { - let PendingConnectionNotification { - connection_id, - kind, - method, - event_seq, - sender, - value, - } = notification; - let notification = match kind { - PendingConnectionMessageKind::Notification => { - serde_json::to_value(crate::NotificationEnvelope { - method: method.clone(), - params: value, - }) - .expect("serialize client notification envelope") - } - PendingConnectionMessageKind::JsonRpcResponse - | PendingConnectionMessageKind::ClientRequest => value, - }; - let item_id = notification_item_id(¬ification); - let assistant_delta = notification_assistant_delta(&method, ¬ification); - let delta_len = assistant_delta.map(str::len); - let assistant_token_text = assistant_delta.and_then(assistant_token_log_preview); - if let Some(assistant_token_text) = assistant_token_text.as_deref() { - tracing::debug!( - stream_elapsed_ms = stream_trace_elapsed_ms(), - connection_id, - method = %method, - event_seq, - item_id = ?item_id, - delta_len = ?delta_len, - assistant_token_text, - "sending client notification" - ); - } else { - tracing::debug!( - stream_elapsed_ms = stream_trace_elapsed_ms(), - connection_id, - method = %method, - event_seq, - item_id = ?item_id, - delta_len = ?delta_len, - "sending client notification" - ); - } - let reserve_started_at = Instant::now(); - let permit = match tokio::time::timeout( - CONNECTION_NOTIFICATION_BACKPRESSURE_LOG_THRESHOLD, - sender.reserve(), - ) - .await - { - Ok(Ok(permit)) => permit, - Ok(Err(_)) => { - tracing::debug!( - connection_id, - method = %method, - event_seq, - "client notification receiver dropped" - ); - return false; - } - Err(_) => { - tracing::warn!( - connection_id, - method = %method, - event_seq, - threshold_ms = CONNECTION_NOTIFICATION_BACKPRESSURE_LOG_THRESHOLD.as_millis(), - "client notification queue applying backpressure" - ); - match sender.reserve().await { - Ok(permit) => permit, - Err(_) => { - tracing::debug!( - connection_id, - method = %method, - event_seq, - "client notification receiver dropped during backpressure" - ); - return false; - } - } - } - }; - let waited = reserve_started_at.elapsed(); - if waited >= CONNECTION_NOTIFICATION_BACKPRESSURE_LOG_THRESHOLD { - tracing::debug!( - connection_id, - method = %method, - event_seq, - waited_ms = waited.as_millis(), - "client notification queue accepted message after backpressure" - ); - } - permit.send(notification); - true -} - -fn stream_trace_elapsed_ms() -> u128 { - static STREAM_TRACE_START: OnceLock = OnceLock::new(); - STREAM_TRACE_START - .get_or_init(Instant::now) - .elapsed() - .as_millis() -} - -fn notification_item_id(value: &serde_json::Value) -> Option { - value - .get("params") - .and_then(|params| params.get("context")) - .and_then(|context| context.get("item_id")) - .and_then(serde_json::Value::as_str) - .map(ToOwned::to_owned) -} - -fn notification_assistant_delta<'a>(method: &str, value: &'a serde_json::Value) -> Option<&'a str> { - (method == "item/agentMessage/delta") - .then(|| value.get("params")?.get("delta")?.as_str()) - .flatten() -} - -fn assistant_token_log_preview(text: &str) -> Option { - assistant_token_logging_enabled() - .then(|| format_assistant_token_log_preview(text, assistant_token_log_max_chars())) -} - -fn assistant_token_logging_enabled() -> bool { - static ASSISTANT_TOKEN_LOGGING_ENABLED: OnceLock = OnceLock::new(); - *ASSISTANT_TOKEN_LOGGING_ENABLED.get_or_init(|| { - std::env::var("DEVO_LOG_ASSISTANT_TOKEN_TEXT") - .ok() - .is_some_and(|value| { - matches!( - value.as_str(), - "1" | "true" | "TRUE" | "yes" | "YES" | "on" | "ON" - ) - }) - }) -} - -fn assistant_token_log_max_chars() -> usize { - static ASSISTANT_TOKEN_LOG_MAX_CHARS: OnceLock = OnceLock::new(); - *ASSISTANT_TOKEN_LOG_MAX_CHARS.get_or_init(|| { - std::env::var("DEVO_ASSISTANT_TOKEN_LOG_MAX_CHARS") - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|value| *value > 0) - .unwrap_or(512) - }) -} - -fn format_assistant_token_log_preview(text: &str, max_chars: usize) -> String { - let max_chars = max_chars.max(1); - let mut preview = String::new(); - let mut chars = text.chars(); - for ch in chars.by_ref().take(max_chars) { - preview.extend(ch.escape_default()); - } - if chars.next().is_some() { - preview.push_str("..."); - } - preview -} - #[cfg(test)] mod tests { use std::collections::HashMap; @@ -1475,10 +1236,10 @@ mod tests { async fn post_response_actions_run_after_backpressured_response_enqueue() -> Result<()> { let data_root = TempDir::new()?; let runtime = build_runtime(data_root.path()); - let (sender, mut receiver) = mpsc::channel(1); - let transport_sender = sender.clone(); + let (outbound_tx, mut receiver) = super::outbound::test_outbound_channel(1); + let transport_outbound = outbound_tx.clone(); let connection_id = runtime - .register_connection(ClientTransportKind::Stdio, sender) + .register_connection(ClientTransportKind::Stdio, outbound_tx) .await; let session_id = SessionId::new(); let response = serde_json::json!({ @@ -1488,10 +1249,17 @@ mod tests { }); let expected_response = response.clone(); - transport_sender - .send(serde_json::json!({ "queued": "backpressure" })) + assert!( + enqueue_outbound( + &transport_outbound, + OutboundFrame::json_rpc_response( + connection_id, + serde_json::json!({ "queued": "backpressure" }), + ), + "test_prefill", + ) .await - .expect("prefill transport queue"); + ); let runtime_for_task = Arc::clone(&runtime); let mut transport_task = tokio::spawn(async move { let incoming = IncomingResponse::new(response).with_post_response_action( @@ -1501,10 +1269,14 @@ mod tests { }, ); let (response, post_response_actions) = incoming.into_parts(); - transport_sender - .send(response) + assert!( + enqueue_outbound( + &transport_outbound, + OutboundFrame::json_rpc_response(connection_id, response), + "test_response", + ) .await - .expect("enqueue response"); + ); runtime_for_task .run_post_response_actions(post_response_actions) .await; @@ -1550,9 +1322,9 @@ mod tests { async fn timed_out_client_request_removes_pending_request() -> Result<()> { let data_root = TempDir::new()?; let runtime = build_runtime(data_root.path()); - let (sender, mut receiver) = mpsc::channel(1); + let (outbound_tx, mut receiver) = super::outbound::test_outbound_channel(1); let connection_id = runtime - .register_connection(ClientTransportKind::Stdio, sender) + .register_connection(ClientTransportKind::Stdio, outbound_tx) .await; let result = runtime @@ -1591,9 +1363,9 @@ mod tests { async fn cancelled_client_request_removes_pending_request() -> Result<()> { let data_root = TempDir::new()?; let runtime = build_runtime(data_root.path()); - let (sender, mut receiver) = mpsc::channel(1); + let (outbound_tx, mut receiver) = super::outbound::test_outbound_channel(1); let connection_id = runtime - .register_connection(ClientTransportKind::Stdio, sender) + .register_connection(ClientTransportKind::Stdio, outbound_tx) .await; let cancel_token = CancellationToken::new(); cancel_token.cancel(); @@ -1633,9 +1405,9 @@ mod tests { async fn dropped_client_request_removes_pending_request() -> Result<()> { let data_root = TempDir::new()?; let runtime = build_runtime(data_root.path()); - let (sender, mut receiver) = mpsc::channel(1); + let (outbound_tx, mut receiver) = super::outbound::test_outbound_channel(1); let connection_id = runtime - .register_connection(ClientTransportKind::Stdio, sender) + .register_connection(ClientTransportKind::Stdio, outbound_tx) .await; let runtime_for_request = Arc::clone(&runtime); @@ -1689,13 +1461,13 @@ mod tests { let session_id = SessionId::new(); let turn_id = TurnId::new(); let item_id = ItemId::new(); - let (owner_sender, mut owner_receiver) = mpsc::channel(4); + let (owner_outbound, mut owner_receiver) = super::outbound::test_outbound_channel(4); let owner_connection_id = runtime - .register_connection(ClientTransportKind::Stdio, owner_sender) + .register_connection(ClientTransportKind::Stdio, owner_outbound) .await; - let (observer_sender, mut observer_receiver) = mpsc::channel(4); + let (observer_outbound, mut observer_receiver) = super::outbound::test_outbound_channel(4); let observer_connection_id = runtime - .register_connection(ClientTransportKind::StdioProxy, observer_sender) + .register_connection(ClientTransportKind::StdioProxy, observer_outbound) .await; runtime @@ -1748,13 +1520,13 @@ mod tests { let session_id = SessionId::new(); let turn_id = TurnId::new(); let item_id = ItemId::new(); - let (owner_sender, mut owner_receiver) = mpsc::channel(4); + let (owner_outbound, mut owner_receiver) = super::outbound::test_outbound_channel(4); let owner_connection_id = runtime - .register_connection(ClientTransportKind::StdioProxy, owner_sender) + .register_connection(ClientTransportKind::StdioProxy, owner_outbound) .await; - let (watcher_sender, mut watcher_receiver) = mpsc::channel(4); + let (watcher_outbound, mut watcher_receiver) = super::outbound::test_outbound_channel(4); let watcher_connection_id = runtime - .register_connection(ClientTransportKind::Stdio, watcher_sender) + .register_connection(ClientTransportKind::Stdio, watcher_outbound) .await; runtime diff --git a/crates/server/src/runtime/outbound.rs b/crates/server/src/runtime/outbound.rs new file mode 100644 index 00000000..f706a53e --- /dev/null +++ b/crates/server/src/runtime/outbound.rs @@ -0,0 +1,328 @@ +use std::sync::OnceLock; +use std::time::Duration; +use std::time::Instant; + +use tokio::sync::mpsc; +use tokio::sync::oneshot; + +use crate::NotificationEnvelope; + +pub(crate) const OUTBOUND_CHANNEL_CAPACITY: usize = 4096; +pub(crate) const OUTBOUND_BACKPRESSURE_LOG_THRESHOLD: Duration = Duration::from_millis(50); + +pub(crate) enum OutboundPayload { + Notification { + connection_id: u64, + method: String, + event_seq: u64, + params: serde_json::Value, + }, + JsonRpcResponse { + connection_id: u64, + value: serde_json::Value, + }, + ClientRequest { + connection_id: u64, + method: String, + value: serde_json::Value, + }, +} + +pub struct OutboundFrame { + pub(crate) payload: OutboundPayload, + pub(crate) delivered: Option>, +} + +impl OutboundFrame { + pub(crate) fn notification( + connection_id: u64, + method: String, + event_seq: u64, + params: serde_json::Value, + ) -> Self { + Self { + payload: OutboundPayload::Notification { + connection_id, + method, + event_seq, + params, + }, + delivered: None, + } + } + + pub fn json_rpc_response(connection_id: u64, value: serde_json::Value) -> Self { + Self { + payload: OutboundPayload::JsonRpcResponse { + connection_id, + value, + }, + delivered: None, + } + } + + pub(crate) fn json_rpc_response_with_delivery( + connection_id: u64, + value: serde_json::Value, + delivered: oneshot::Sender, + ) -> Self { + Self { + payload: OutboundPayload::JsonRpcResponse { + connection_id, + value, + }, + delivered: Some(delivered), + } + } + + pub(crate) fn client_request( + connection_id: u64, + method: String, + value: serde_json::Value, + ) -> Self { + Self { + payload: OutboundPayload::ClientRequest { + connection_id, + method, + value, + }, + delivered: None, + } + } + + pub(crate) fn connection_id(&self) -> u64 { + match &self.payload { + OutboundPayload::Notification { connection_id, .. } + | OutboundPayload::JsonRpcResponse { connection_id, .. } + | OutboundPayload::ClientRequest { connection_id, .. } => *connection_id, + } + } + + pub(crate) fn log_method(&self) -> &str { + match &self.payload { + OutboundPayload::Notification { method, .. } => method, + OutboundPayload::JsonRpcResponse { .. } => "", + OutboundPayload::ClientRequest { method, .. } => method, + } + } + + pub(crate) fn event_seq(&self) -> u64 { + match &self.payload { + OutboundPayload::Notification { event_seq, .. } => *event_seq, + OutboundPayload::JsonRpcResponse { .. } | OutboundPayload::ClientRequest { .. } => 0, + } + } +} + +pub(crate) fn outbound_frame_to_value(frame: &OutboundFrame) -> serde_json::Value { + match &frame.payload { + OutboundPayload::Notification { method, params, .. } => { + serde_json::to_value(NotificationEnvelope { + method: method.clone(), + params: params.clone(), + }) + .expect("serialize client notification envelope") + } + OutboundPayload::JsonRpcResponse { value, .. } + | OutboundPayload::ClientRequest { value, .. } => value.clone(), + } +} + +pub(crate) fn log_outbound_frame(frame: &OutboundFrame, notification: &serde_json::Value) { + let connection_id = frame.connection_id(); + let method = frame.log_method(); + let event_seq = frame.event_seq(); + let item_id = notification_item_id(notification); + let assistant_delta = notification_assistant_delta(method, notification); + let delta_len = assistant_delta.map(str::len); + let assistant_token_text = assistant_delta.and_then(assistant_token_log_preview); + if let Some(assistant_token_text) = assistant_token_text.as_deref() { + tracing::debug!( + stream_elapsed_ms = stream_trace_elapsed_ms(), + connection_id, + method = %method, + event_seq, + item_id = ?item_id, + delta_len = ?delta_len, + assistant_token_text, + "sending client notification" + ); + } else { + tracing::debug!( + stream_elapsed_ms = stream_trace_elapsed_ms(), + connection_id, + method = %method, + event_seq, + item_id = ?item_id, + delta_len = ?delta_len, + "sending client notification" + ); + } +} + +pub(crate) async fn enqueue_outbound( + tx: &mpsc::Sender, + frame: OutboundFrame, + queue: &'static str, +) -> bool { + let connection_id = frame.connection_id(); + let reserve_started_at = Instant::now(); + let permit = match tokio::time::timeout(OUTBOUND_BACKPRESSURE_LOG_THRESHOLD, tx.reserve()).await + { + Ok(Ok(permit)) => permit, + Ok(Err(_)) => { + tracing::debug!(connection_id, queue, "outbound queue receiver dropped"); + return false; + } + Err(_) => { + tracing::warn!( + connection_id, + queue, + threshold_ms = OUTBOUND_BACKPRESSURE_LOG_THRESHOLD.as_millis(), + "outbound queue applying backpressure" + ); + match tx.reserve().await { + Ok(permit) => permit, + Err(_) => { + tracing::debug!( + connection_id, + queue, + "outbound queue receiver dropped during backpressure" + ); + return false; + } + } + } + }; + let waited = reserve_started_at.elapsed(); + if waited >= OUTBOUND_BACKPRESSURE_LOG_THRESHOLD { + tracing::debug!( + connection_id, + queue, + waited_ms = waited.as_millis(), + "outbound queue accepted message after backpressure" + ); + } + permit.send(frame); + true +} + +/// Test helper: drains [`OutboundFrame`] values into serialized JSON values. +#[doc(hidden)] +pub fn test_outbound_channel( + capacity: usize, +) -> (mpsc::Sender, mpsc::Receiver) { + let (outbound_tx, mut outbound_rx) = mpsc::channel::(capacity); + let (json_tx, json_rx) = mpsc::channel(capacity); + tokio::spawn(async move { + while let Some(frame) = outbound_rx.recv().await { + let value = outbound_frame_to_value(&frame); + let delivered = frame.delivered; + let sent = json_tx.send(value).await.is_ok(); + if let Some(delivered) = delivered { + let _ = delivered.send(sent); + } + } + }); + (outbound_tx, json_rx) +} + +fn stream_trace_elapsed_ms() -> u128 { + static STREAM_TRACE_START: OnceLock = OnceLock::new(); + STREAM_TRACE_START + .get_or_init(Instant::now) + .elapsed() + .as_millis() +} + +fn notification_item_id(value: &serde_json::Value) -> Option { + value + .get("params") + .and_then(|params| params.get("context")) + .and_then(|context| context.get("item_id")) + .and_then(serde_json::Value::as_str) + .map(ToOwned::to_owned) +} + +fn notification_assistant_delta<'a>(method: &str, value: &'a serde_json::Value) -> Option<&'a str> { + (method == "item/agentMessage/delta") + .then(|| value.get("params")?.get("delta")?.as_str()) + .flatten() +} + +fn assistant_token_log_preview(text: &str) -> Option { + assistant_token_logging_enabled() + .then(|| format_assistant_token_log_preview(text, assistant_token_log_max_chars())) +} + +fn assistant_token_logging_enabled() -> bool { + static ASSISTANT_TOKEN_LOGGING_ENABLED: OnceLock = OnceLock::new(); + *ASSISTANT_TOKEN_LOGGING_ENABLED.get_or_init(|| { + std::env::var("DEVO_LOG_ASSISTANT_TOKEN_TEXT") + .ok() + .is_some_and(|value| { + matches!( + value.as_str(), + "1" | "true" | "TRUE" | "yes" | "YES" | "on" | "ON" + ) + }) + }) +} + +fn assistant_token_log_max_chars() -> usize { + static ASSISTANT_TOKEN_LOG_MAX_CHARS: OnceLock = OnceLock::new(); + *ASSISTANT_TOKEN_LOG_MAX_CHARS.get_or_init(|| { + std::env::var("DEVO_ASSISTANT_TOKEN_LOG_MAX_CHARS") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|value| *value > 0) + .unwrap_or(512) + }) +} + +fn format_assistant_token_log_preview(text: &str, max_chars: usize) -> String { + let max_chars = max_chars.max(1); + let mut preview = String::new(); + let mut chars = text.chars(); + for ch in chars.by_ref().take(max_chars) { + preview.extend(ch.escape_default()); + } + if chars.next().is_some() { + preview.push_str("..."); + } + preview +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + + #[test] + fn outbound_frame_to_value_wraps_notifications() { + let frame = OutboundFrame::notification( + 1, + "session/update".to_string(), + 3, + serde_json::json!({ "sessionId": "abc" }), + ); + assert_eq!( + outbound_frame_to_value(&frame), + serde_json::json!({ + "method": "session/update", + "params": { "sessionId": "abc" }, + }) + ); + } + + #[test] + fn outbound_frame_to_value_passthrough_responses() { + let response = serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "result": { "ok": true }, + }); + let frame = OutboundFrame::json_rpc_response(1, response.clone()); + assert_eq!(outbound_frame_to_value(&frame), response); + } +} diff --git a/crates/server/src/transport.rs b/crates/server/src/transport.rs index f3b02b35..13be8bb2 100644 --- a/crates/server/src/transport.rs +++ b/crates/server/src/transport.rs @@ -1,6 +1,4 @@ use std::sync::Arc; -use std::time::Duration; -use std::time::Instant; use anyhow::Context; use anyhow::Result; @@ -16,7 +14,9 @@ use tokio::io::AsyncBufReadExt; use tokio::io::AsyncWriteExt; use tokio::io::BufReader; use tokio::net::TcpListener; +use tokio::sync::Semaphore; use tokio::sync::mpsc; +use tokio::task::JoinHandle; use tokio::task::JoinSet; use tokio_tungstenite::accept_async; use tokio_tungstenite::tungstenite::Message; @@ -26,13 +26,96 @@ use crate::AcpErrorCode; use crate::ClientTransportKind; use crate::ServerRuntime; use crate::acp_error_response; -use crate::runtime::CONNECTION_NOTIFICATION_CHANNEL_CAPACITY; +use crate::runtime::INBOUND_CONCURRENCY_LIMIT; use crate::runtime::IncomingResponse; +use crate::runtime::OUTBOUND_CHANNEL_CAPACITY; +use crate::runtime::OutboundFrame; +use crate::runtime::enqueue_outbound; +use crate::runtime::log_outbound_frame; +use crate::runtime::outbound_frame_to_value; use crate::singleton::SERVER_CONTROL_SHUTDOWN_METHOD; use crate::singleton::SERVER_CONTROL_STATUS_METHOD; -const TRANSPORT_WRITE_CHANNEL_CAPACITY: usize = 4096; -const TRANSPORT_BACKPRESSURE_LOG_THRESHOLD: Duration = Duration::from_millis(50); +/// Per-connection transport state shared between the read loop and handler tasks. +struct ConnectionTransport { + outbound_tx: mpsc::Sender, + inbound_semaphore: Arc, +} + +impl ConnectionTransport { + fn new() -> (Self, mpsc::Receiver) { + let (outbound_tx, outbound_rx) = mpsc::channel(OUTBOUND_CHANNEL_CAPACITY); + ( + Self { + outbound_tx, + inbound_semaphore: Arc::new(Semaphore::new(INBOUND_CONCURRENCY_LIMIT)), + }, + outbound_rx, + ) + } +} + +enum OutboundSink { + Stdio, + WebSocket(futures::stream::SplitSink< + tokio_tungstenite::WebSocketStream, + Message, + >), +} + +async fn write_outbound_frame( + sink: &mut OutboundSink, + frame: OutboundFrame, +) -> Result { + let value = outbound_frame_to_value(&frame); + log_outbound_frame(&frame, &value); + let delivered = frame.delivered; + let sent = match sink { + OutboundSink::Stdio => { + let mut stdout = tokio::io::stdout(); + let bytes = serde_json::to_vec(&value).expect("serialize stdio outbound frame"); + stdout.write_all(&bytes).await?; + stdout.write_all(b"\n").await?; + stdout.flush().await?; + true + } + OutboundSink::WebSocket(writer) => writer + .send(Message::Text( + serde_json::to_string(&value) + .expect("serialize websocket outbound frame") + .into(), + )) + .await + .is_ok(), + }; + if let Some(delivered) = delivered { + let _ = delivered.send(sent); + } + Ok(sent) +} + +fn spawn_outbound_writer( + connection_id: u64, + mut rx: mpsc::Receiver, + mut sink: OutboundSink, +) -> JoinHandle<()> { + tokio::spawn(async move { + while let Some(frame) = rx.recv().await { + match write_outbound_frame(&mut sink, frame).await { + Ok(true) => {} + Ok(false) if matches!(sink, OutboundSink::WebSocket(_)) => { + tracing::debug!(connection_id, "outbound websocket writer closed"); + break; + } + Ok(false) => {} + Err(error) => { + tracing::warn!(connection_id, error = %error, "outbound writer failed"); + break; + } + } + } + }) +} /// Transport trait per L3-BEH-SERVER-001. /// @@ -279,60 +362,33 @@ async fn run_listener_tasks( /// must occupy exactly one line of valid JSON. Pretty-printed or otherwise /// multi-line payloads are rejected at the transport layer. async fn run_stdio(runtime: Arc) -> Result<()> { - // Server → client responses and notifications (TextDelta, TurnStarted, …). - let (sender, mut receiver) = mpsc::channel(CONNECTION_NOTIFICATION_CHANNEL_CAPACITY); - let sender_clone = sender.clone(); + let (transport, outbound_rx) = ConnectionTransport::new(); + let outbound_tx = transport.outbound_tx.clone(); let connection_id = runtime - .register_connection(ClientTransportKind::Stdio, sender) + .register_connection(ClientTransportKind::Stdio, transport.outbound_tx) .await; tracing::info!(connection_id, "stdio connection established"); - // Serialized NDJSON lines awaiting stdout. Bounded capacity applies - // backpressure when stdout is slow instead of buffering without limit. - let (write_tx, mut write_rx) = mpsc::channel::>(TRANSPORT_WRITE_CHANNEL_CAPACITY); - - // --- Writer task --- - // Sole responsibility: read serialized lines from write_rx and write - // them to stdout. This is the only task that can block on stdout. - let writer_task = tokio::spawn(async move { - let mut stdout = tokio::io::stdout(); - while let Some(line) = write_rx.recv().await { - stdout.write_all(&line).await?; - stdout.write_all(b"\n").await?; - stdout.flush().await?; - } - Result::<()>::Ok(()) - }); + let outbound_writer = + spawn_outbound_writer(connection_id, outbound_rx, OutboundSink::Stdio); - // --- Producer task --- - // Serializes outbound messages and forwards them to the writer task. - let producer_task = tokio::spawn(async move { - while let Some(message) = receiver.recv().await { - let line = serde_json::to_vec(&message).expect("serialize stdio response"); - if !send_transport_queue_message(&write_tx, line, connection_id, "stdio_write").await - { - break; - } - } - }); - - // --- Stdin reader (main task) --- let stdin = tokio::io::stdin(); let mut lines = BufReader::new(stdin).lines(); while let Some(line) = lines.next_line().await? { accept_incoming_client_message( Arc::clone(&runtime), connection_id, - sender_clone.clone(), + outbound_tx.clone(), + Arc::clone(&transport.inbound_semaphore), &line, "stdio_notifications", - ); + ) + .await; } runtime.unregister_connection(connection_id).await; tracing::info!(connection_id, "stdio connection closed"); - writer_task.abort(); - producer_task.abort(); + outbound_writer.abort(); Ok(()) } @@ -399,32 +455,25 @@ async fn handle_internal_proxy_connection( } }; - let (sender, mut receiver) = mpsc::channel(CONNECTION_NOTIFICATION_CHANNEL_CAPACITY); - let sender_clone = sender.clone(); + let (transport, outbound_rx) = ConnectionTransport::new(); + let outbound_tx = transport.outbound_tx.clone(); let connection_id = runtime - .register_connection(ClientTransportKind::StdioProxy, sender) + .register_connection(ClientTransportKind::StdioProxy, transport.outbound_tx) .await; tracing::info!(connection_id, "internal stdio proxy connection established"); - let writer_task = tokio::spawn(async move { - while let Some(message) = receiver.recv().await { - writer - .send(Message::Text( - serde_json::to_string(&message) - .expect("serialize internal proxy response") - .into(), - )) - .await?; - } - Result::<()>::Ok(()) - }); + let outbound_writer = spawn_outbound_writer( + connection_id, + outbound_rx, + OutboundSink::WebSocket(writer), + ); if let Some(response) = runtime .handle_incoming_with_actions(connection_id, first_value) .await && !send_incoming_response( &runtime, - &sender_clone, + &outbound_tx, response, connection_id, "internal_proxy_notifications", @@ -433,7 +482,7 @@ async fn handle_internal_proxy_connection( { runtime.unregister_connection(connection_id).await; tracing::info!(connection_id, "internal stdio proxy connection closed"); - writer_task.abort(); + outbound_writer.abort(); return Ok(()); } @@ -444,10 +493,12 @@ async fn handle_internal_proxy_connection( accept_incoming_client_message( Arc::clone(&runtime), connection_id, - sender_clone.clone(), + outbound_tx.clone(), + Arc::clone(&transport.inbound_semaphore), text.as_str(), "internal_proxy_notifications", - ); + ) + .await; } Message::Close(_) => break, _ => {} @@ -456,7 +507,7 @@ async fn handle_internal_proxy_connection( runtime.unregister_connection(connection_id).await; tracing::info!(connection_id, "internal stdio proxy connection closed"); - writer_task.abort(); + outbound_writer.abort(); Ok(()) } @@ -527,26 +578,19 @@ async fn handle_websocket_connection( stream: tokio::net::TcpStream, ) -> Result<()> { let websocket = accept_async(stream).await?; - let (mut writer, mut reader) = websocket.split(); - let (sender, mut receiver) = mpsc::channel(CONNECTION_NOTIFICATION_CHANNEL_CAPACITY); - let sender_clone = sender.clone(); + let (writer, mut reader) = websocket.split(); + let (transport, outbound_rx) = ConnectionTransport::new(); + let outbound_tx = transport.outbound_tx.clone(); let connection_id = runtime - .register_connection(ClientTransportKind::WebSocket, sender) + .register_connection(ClientTransportKind::WebSocket, transport.outbound_tx) .await; tracing::info!(connection_id, "websocket connection established"); - let writer_task = tokio::spawn(async move { - while let Some(message) = receiver.recv().await { - writer - .send(Message::Text( - serde_json::to_string(&message) - .expect("serialize websocket response") - .into(), - )) - .await?; - } - Result::<()>::Ok(()) - }); + let outbound_writer = spawn_outbound_writer( + connection_id, + outbound_rx, + OutboundSink::WebSocket(writer), + ); while let Some(frame) = reader.next().await { let frame = frame?; @@ -555,10 +599,12 @@ async fn handle_websocket_connection( accept_incoming_client_message( Arc::clone(&runtime), connection_id, - sender_clone.clone(), + outbound_tx.clone(), + Arc::clone(&transport.inbound_semaphore), text.as_str(), "websocket_notifications", - ); + ) + .await; } Message::Close(_) => break, _ => {} @@ -567,7 +613,7 @@ async fn handle_websocket_connection( runtime.unregister_connection(connection_id).await; tracing::info!(connection_id, "websocket connection closed"); - writer_task.abort(); + outbound_writer.abort(); Ok(()) } @@ -694,12 +740,24 @@ fn validate_incoming_client_message( )) } +/// Returns `true` when the payload is a client response to a server-initiated +/// JSON-RPC request. +fn is_client_response_message(value: &serde_json::Value) -> bool { + let Some(object) = value.as_object() else { + return false; + }; + !object.contains_key("method") + && object.contains_key("id") + && (object.contains_key("result") || object.contains_key("error")) +} + /// Parses, validates, and dispatches one inbound client payload without blocking -/// the transport read loop. -fn accept_incoming_client_message( +/// the transport read loop on handler work. +async fn accept_incoming_client_message( runtime: Arc, connection_id: u64, - sender: mpsc::Sender, + outbound_tx: mpsc::Sender, + inbound_semaphore: Arc, raw_payload: &str, queue: &'static str, ) { @@ -710,58 +768,66 @@ fn accept_incoming_client_message( Ok(value) => value, Err(error_response) => { tracing::warn!(connection_id, "rejected malformed client payload"); + let outbound_tx = outbound_tx.clone(); tokio::spawn(async move { - send_transport_queue_message(&sender, error_response, connection_id, queue).await; + enqueue_outbound( + &outbound_tx, + OutboundFrame::json_rpc_response(connection_id, error_response), + queue, + ) + .await; }); return; } }; if let Err(error_response) = validate_incoming_client_message(&value) { tracing::warn!(connection_id, "rejected invalid client message"); + let outbound_tx = outbound_tx.clone(); tokio::spawn(async move { - send_transport_queue_message(&sender, error_response, connection_id, queue).await; + enqueue_outbound( + &outbound_tx, + OutboundFrame::json_rpc_response(connection_id, error_response), + queue, + ) + .await; }); return; } - spawn_incoming_message_handler(runtime, connection_id, sender, value, queue); -} - -/// Dispatch a single decoded client message without blocking the connection's -/// read loop. -/// -/// The reader task MUST stay responsive so that it can always deliver -/// `turn/interrupt` requests and, critically, client responses to -/// server-initiated requests (approvals, `fs/read_text_file`, user input). -/// Handling a message inline would let a blocked handler or a backpressured -/// response send freeze the reader, which deadlocks the whole connection (the -/// awaited client reply can never be read). Spawning per message keeps the -/// reader draining while individual handlers make progress concurrently. -fn spawn_incoming_message_handler( - runtime: Arc, - connection_id: u64, - sender: mpsc::Sender, - value: serde_json::Value, - queue: &'static str, -) { + if is_client_response_message(&value) { + tokio::spawn(async move { + runtime.resolve_client_response(connection_id, value).await; + }); + return; + } + let Ok(permit) = inbound_semaphore.acquire_owned().await else { + return; + }; tokio::spawn(async move { + let _permit = permit; if let Some(response) = runtime .handle_incoming_with_actions(connection_id, value) .await { - send_incoming_response(&runtime, &sender, response, connection_id, queue).await; + send_incoming_response(&runtime, &outbound_tx, response, connection_id, queue).await; } }); } async fn send_incoming_response( runtime: &Arc, - sender: &mpsc::Sender, + outbound_tx: &mpsc::Sender, response: IncomingResponse, connection_id: u64, queue: &'static str, ) -> bool { let (response, post_response_actions) = response.into_parts(); - if !send_transport_queue_message(sender, response, connection_id, queue).await { + if !enqueue_outbound( + outbound_tx, + OutboundFrame::json_rpc_response(connection_id, response), + queue, + ) + .await + { return false; } runtime @@ -770,53 +836,6 @@ async fn send_incoming_response( true } -async fn send_transport_queue_message( - sender: &mpsc::Sender, - value: T, - connection_id: u64, - queue: &'static str, -) -> bool { - let reserve_started_at = Instant::now(); - let permit = - match tokio::time::timeout(TRANSPORT_BACKPRESSURE_LOG_THRESHOLD, sender.reserve()).await { - Ok(Ok(permit)) => permit, - Ok(Err(_)) => { - tracing::debug!(connection_id, queue, "transport queue receiver dropped"); - return false; - } - Err(_) => { - tracing::warn!( - connection_id, - queue, - threshold_ms = TRANSPORT_BACKPRESSURE_LOG_THRESHOLD.as_millis(), - "transport queue applying backpressure" - ); - match sender.reserve().await { - Ok(permit) => permit, - Err(_) => { - tracing::debug!( - connection_id, - queue, - "transport queue receiver dropped during backpressure" - ); - return false; - } - } - } - }; - let waited = reserve_started_at.elapsed(); - if waited >= TRANSPORT_BACKPRESSURE_LOG_THRESHOLD { - tracing::debug!( - connection_id, - queue, - waited_ms = waited.as_millis(), - "transport queue accepted message after backpressure" - ); - } - permit.send(value); - true -} - #[cfg(test)] mod tests { use super::DEFAULT_WEBSOCKET_BIND_ADDRESS; @@ -826,6 +845,7 @@ mod tests { use super::ListenTarget; use super::internal_proxy_control_response; use super::parse_internal_proxy_control_request; + use super::is_client_response_message; use super::parse_incoming_client_payload; use super::parse_listen_target; use super::resolve_listen_targets; @@ -924,6 +944,55 @@ mod tests { ); } + #[test] + fn is_client_response_message_detects_server_reply_payloads() { + assert!(is_client_response_message(&serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "result": { "ok": true }, + }))); + assert!(!is_client_response_message(&serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": {}, + }))); + } + + #[tokio::test] + async fn enqueue_outbound_backpressures_full_receiver() { + use crate::runtime::OutboundFrame; + use crate::runtime::enqueue_outbound; + use std::time::Duration; + + let (tx, mut rx) = mpsc::channel(1); + assert!( + enqueue_outbound( + &tx, + OutboundFrame::json_rpc_response(1, serde_json::json!({ "first": true })), + "test_outbound", + ) + .await + ); + let tx_for_task = tx.clone(); + let mut blocked_enqueue = tokio::spawn(async move { + enqueue_outbound( + &tx_for_task, + OutboundFrame::json_rpc_response(1, serde_json::json!({ "second": true })), + "test_outbound", + ) + .await + }); + assert!( + tokio::time::timeout(Duration::from_millis(20), &mut blocked_enqueue) + .await + .is_err() + ); + rx.recv().await.expect("first frame"); + assert!(blocked_enqueue.await.expect("blocked enqueue")); + rx.recv().await.expect("second frame"); + } + #[test] fn parse_incoming_client_payload_skips_empty_and_parses_json() { assert_eq!(parse_incoming_client_payload(""), None); diff --git a/crates/server/tests/acp_available_commands.rs b/crates/server/tests/acp_available_commands.rs index 0a1f75ed..db0a433a 100644 --- a/crates/server/tests/acp_available_commands.rs +++ b/crates/server/tests/acp_available_commands.rs @@ -22,6 +22,7 @@ use devo_protocol::StreamEvent; use devo_protocol::Usage; use devo_provider::ModelProviderSDK; use devo_provider::SingleProviderRouter; +use devo_server::OutboundFrame; use devo_server::AcpSuccessResponse; use devo_server::ClientTransportKind; use devo_server::ServerRuntime; @@ -63,7 +64,7 @@ impl ModelProviderSDK for NoopProvider { async fn acp_available_commands_are_session_update_after_session_response() -> Result<()> { let data_root = TempDir::new()?; let runtime = build_runtime(data_root.path())?; - let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*buffer*/ 4096); + let (outgoing_tx, mut outgoing_rx) = devo_server::test_outbound_channel(4096); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, outgoing_tx.clone()) .await; @@ -90,7 +91,10 @@ async fn acp_available_commands_are_session_update_after_session_response() -> R serde_json::from_value(response_value.clone())?; let session_id = response.result.session_id; outgoing_tx - .send(response_value) + .send(OutboundFrame::json_rpc_response( + connection_id, + response_value, + )) .await .context("enqueue simulated transport response")?; runtime diff --git a/crates/server/tests/acp_permission_tool_status_contract.rs b/crates/server/tests/acp_permission_tool_status_contract.rs index 448fa204..91eb8e52 100644 --- a/crates/server/tests/acp_permission_tool_status_contract.rs +++ b/crates/server/tests/acp_permission_tool_status_contract.rs @@ -566,7 +566,7 @@ fn build_runtime( async fn initialize_acp_connection( runtime: &Arc, ) -> Result<(u64, mpsc::Receiver)> { - let (notifications_tx, notifications_rx) = mpsc::channel(/*buffer*/ 4096); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(4096); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; diff --git a/crates/server/tests/acp_session_delete.rs b/crates/server/tests/acp_session_delete.rs index 2655afa5..5e939cbe 100644 --- a/crates/server/tests/acp_session_delete.rs +++ b/crates/server/tests/acp_session_delete.rs @@ -141,7 +141,7 @@ async fn acp_session_delete_removes_session_from_history_and_is_idempotent() -> )?)), ), ); - let (notifications_tx, _notifications_rx) = mpsc::channel(/*buffer*/ 4096); + let (notifications_tx, _notifications_rx) = devo_server::test_outbound_channel(4096); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; @@ -431,7 +431,7 @@ fn build_runtime_with_provider( async fn initialize_connection( runtime: &Arc, ) -> Result<(u64, mpsc::Receiver)> { - let (notifications_tx, notifications_rx) = mpsc::channel(/*buffer*/ 4096); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(4096); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; diff --git a/crates/server/tests/acp_session_lifecycle.rs b/crates/server/tests/acp_session_lifecycle.rs index 8545029a..9aef8537 100644 --- a/crates/server/tests/acp_session_lifecycle.rs +++ b/crates/server/tests/acp_session_lifecycle.rs @@ -883,7 +883,7 @@ logout = true async fn legacy_initialize_params_are_rejected() -> Result<()> { let data_root = TempDir::new()?; let runtime = build_runtime(data_root.path())?; - let (notifications_tx, _notifications_rx) = mpsc::channel(/*buffer*/ 4096); + let (notifications_tx, _notifications_rx) = devo_server::test_outbound_channel(4096); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; @@ -961,7 +961,7 @@ async fn initialize_acp_connection_with_transport( runtime: &Arc, transport: ClientTransportKind, ) -> Result<(u64, mpsc::Receiver, AcpInitializeResult)> { - let (notifications_tx, notifications_rx) = mpsc::channel(/*buffer*/ 4096); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(4096); let connection_id = runtime .register_connection(transport, notifications_tx) .await; diff --git a/crates/server/tests/cli_log_failures.rs b/crates/server/tests/cli_log_failures.rs index fc320ac6..c05f2329 100644 --- a/crates/server/tests/cli_log_failures.rs +++ b/crates/server/tests/cli_log_failures.rs @@ -296,7 +296,7 @@ fn build_runtime( async fn initialize_connection( runtime: &Arc, ) -> Result<(u64, mpsc::Receiver)> { - let (notifications_tx, notifications_rx) = mpsc::channel(4096); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(4096); let connection_id = runtime .register_connection(devo_server::ClientTransportKind::Stdio, notifications_tx) .await; diff --git a/crates/server/tests/command_exec.rs b/crates/server/tests/command_exec.rs index b5599301..2da2e5e8 100644 --- a/crates/server/tests/command_exec.rs +++ b/crates/server/tests/command_exec.rs @@ -228,7 +228,7 @@ async fn initialize_connection( runtime: &Arc, client_name: &str, ) -> Result<(u64, mpsc::Receiver)> { - let (notifications_tx, notifications_rx) = mpsc::channel(/*buffer*/ 128); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(128); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; diff --git a/crates/server/tests/deep_research_boundary_failure.rs b/crates/server/tests/deep_research_boundary_failure.rs index 417e1535..db0670ce 100644 --- a/crates/server/tests/deep_research_boundary_failure.rs +++ b/crates/server/tests/deep_research_boundary_failure.rs @@ -328,7 +328,7 @@ fn streamed_text_events(text: impl Into) -> Vec> { async fn initialize_connection( runtime: &Arc, ) -> Result<(u64, mpsc::Receiver)> { - let (notifications_tx, notifications_rx) = mpsc::channel(/*buffer*/ 256); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(256); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; diff --git a/crates/server/tests/deep_research_e2e.rs b/crates/server/tests/deep_research_e2e.rs index 236d4240..07f9d118 100644 --- a/crates/server/tests/deep_research_e2e.rs +++ b/crates/server/tests/deep_research_e2e.rs @@ -678,8 +678,8 @@ async fn deep_research_accepts_write_tool_only_final_report() -> Result<()> { report_contents.contains("DeepSeek official website"), "expected written report to contain final report content: {report_contents}" ); - let final_report = - latest_parent_agent_message(&events, session_id).context("expected final report message")?; + let final_report = latest_parent_agent_message(&events, session_id) + .context("expected final report message")?; assert!( final_report.contains("Wrote the full research report"), "expected final response to point at written report: {final_report}" @@ -867,102 +867,6 @@ async fn deep_research_streams_researcher_delta_before_query_finishes() -> Resul Ok(()) } -#[tokio::test] -async fn deep_research_continues_after_delegated_worker_failure() -> Result<()> { - // Trace: L2-DES-RESEARCH-001 - // Verifies: supervisor-driven research can retry after a failed delegated worker turn. - let workspace = TempDir::new()?; - write_live_research_config(workspace.path())?; - let provider: Arc = - Arc::new(ScriptedResearchProvider::with_delegated_worker_failure_once(workspace.path())); - let runtime = build_scripted_research_runtime_with_provider(workspace.path(), provider)?; - let (connection_id, mut notifications_rx) = initialize_connection(&runtime).await?; - let session_id = start_session(&runtime, connection_id, workspace.path()).await?; - start_research_turn(&runtime, connection_id, session_id).await?; - - let events = wait_for_research_completion( - &runtime, - connection_id, - session_id, - &mut notifications_rx, - "Use the provided scope.", - ) - .await?; - - assert!( - events.iter().any(|event| { - event.get("method") == Some(&serde_json::json!("turn/failed")) - && event["params"]["session_id"] != serde_json::json!(session_id.to_string()) - }), - "expected child turn failure to be visible before recovery: {events:#?}" - ); - let completed_turn = events - .iter() - .find(|event| { - event.get("method") == Some(&serde_json::json!("turn/completed")) - && event["params"]["session_id"] == serde_json::json!(session_id.to_string()) - }) - .context("missing parent research turn completion")?; - assert_eq!( - completed_turn["params"]["turn"]["status"], - serde_json::json!("Completed") - ); - let final_report = - latest_parent_agent_message(&events, session_id).context("expected final report message")?; - assert!( - final_report.contains("DeepSeek official website"), - "expected final report after delegated worker recovery: {final_report}" - ); - - Ok(()) -} - -#[tokio::test] -async fn deep_research_restarts_worker_when_continuation_fails() -> Result<()> { - // Trace: L2-DES-RESEARCH-001 - // Verifies: supervisor-driven research can spawn replacement workers after repeated failures. - let workspace = TempDir::new()?; - write_live_research_config(workspace.path())?; - let provider: Arc = Arc::new( - ScriptedResearchProvider::with_delegated_worker_failures_before_success( - workspace.path(), - 2, - ), - ); - let runtime = build_scripted_research_runtime_with_provider(workspace.path(), provider)?; - let (connection_id, mut notifications_rx) = initialize_connection(&runtime).await?; - let session_id = start_session(&runtime, connection_id, workspace.path()).await?; - start_research_turn(&runtime, connection_id, session_id).await?; - - let events = wait_for_research_completion( - &runtime, - connection_id, - session_id, - &mut notifications_rx, - "Use the provided scope.", - ) - .await?; - - let child_failures = events - .iter() - .filter(|event| { - event.get("method") == Some(&serde_json::json!("turn/failed")) - && event["params"]["session_id"] != serde_json::json!(session_id.to_string()) - }) - .count(); - assert_eq!(child_failures, 2); - let child_sessions = unique_child_turn_sessions(&events, session_id); - assert_eq!(child_sessions.len(), 3); - let final_report = - latest_parent_agent_message(&events, session_id).context("expected final report message")?; - assert!( - final_report.contains("DeepSeek official website"), - "expected final report after replacement worker recovery: {final_report}" - ); - - Ok(()) -} - #[tokio::test] async fn interrupted_research_closes_delegated_child_agent() -> Result<()> { // Trace: L2-DES-RESEARCH-001 @@ -1608,7 +1512,8 @@ fn supervisor_stream_events(request: &ModelRequest) -> Vec> } Some(poll_index) => { let wait_id = supervisor_wait_tool_id(attempt, poll_index); - let wait_content = request_tool_result_content(request, &wait_id).unwrap_or_default(); + let wait_content = + request_tool_result_content(request, &wait_id).unwrap_or_default(); // #region agent log agent_debug_log( "deep_research_e2e.rs:supervisor_stream_events", @@ -1662,19 +1567,20 @@ fn supervisor_wait_tool_id(attempt: usize, poll_index: u32) -> String { fn wait_agent_next_sequence(content: &str) -> Option { serde_json::from_str::(content) .ok() - .and_then(|value| value.get("next_sequence").and_then(serde_json::Value::as_u64)) + .and_then(|value| { + value + .get("next_sequence") + .and_then(serde_json::Value::as_u64) + }) } fn wait_agent_result_indicates_failure(content: &str) -> bool { - if content.contains("failed") - || content.contains("interrupted") - || content.contains("canceled") + if content.contains("failed") || content.contains("interrupted") || content.contains("canceled") { return true; } - wait_agent_statuses(content).any(|status| { - matches!(status.as_str(), "failed" | "interrupted" | "closed") - }) + wait_agent_statuses(content) + .any(|status| matches!(status.as_str(), "failed" | "interrupted" | "closed")) } fn wait_agent_result_indicates_success(content: &str) -> bool { @@ -1712,12 +1618,7 @@ fn wait_agent_statuses(content: &str) -> impl Iterator { } // #region agent log -fn agent_debug_log( - location: &str, - message: &str, - data: serde_json::Value, - hypothesis_id: &str, -) { +fn agent_debug_log(location: &str, message: &str, data: serde_json::Value, hypothesis_id: &str) { let timestamp = SystemTime::now() .duration_since(UNIX_EPOCH) .map(|duration| duration.as_millis()) @@ -1936,7 +1837,7 @@ fn streamed_text_event_chunks( async fn initialize_connection( runtime: &Arc, ) -> Result<(u64, mpsc::Receiver)> { - let (notifications_tx, notifications_rx) = mpsc::channel(/*buffer*/ 256); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(256); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; @@ -2436,7 +2337,10 @@ fn legacy_event_from_acp_notification(value: serde_json::Value) -> serde_json::V } fn latest_agent_message(events: &[serde_json::Value]) -> Option { - events.iter().rev().find_map(|event| agent_message_from_completed_event(event)) + events + .iter() + .rev() + .find_map(|event| agent_message_from_completed_event(event)) } fn latest_parent_agent_message( diff --git a/crates/server/tests/goal_title_generation.rs b/crates/server/tests/goal_title_generation.rs index 6a79c6f4..13c6da28 100644 --- a/crates/server/tests/goal_title_generation.rs +++ b/crates/server/tests/goal_title_generation.rs @@ -238,7 +238,7 @@ fn build_runtime( async fn initialize_connection( runtime: &Arc, ) -> Result<(u64, mpsc::Receiver)> { - let (notifications_tx, notifications_rx) = mpsc::channel(/*buffer*/ 1024); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(1024); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; diff --git a/crates/server/tests/persistence_resume.rs b/crates/server/tests/persistence_resume.rs index 65dc3ddb..fc9262a2 100644 --- a/crates/server/tests/persistence_resume.rs +++ b/crates/server/tests/persistence_resume.rs @@ -1354,7 +1354,7 @@ fn build_runtime_with_provider( async fn initialize_connection( runtime: &Arc, ) -> Result<(u64, mpsc::Receiver)> { - let (notifications_tx, notifications_rx) = mpsc::channel(/*buffer*/ 4096); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(4096); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; diff --git a/crates/server/tests/provider_routing.rs b/crates/server/tests/provider_routing.rs index 684149ff..b7c700a9 100644 --- a/crates/server/tests/provider_routing.rs +++ b/crates/server/tests/provider_routing.rs @@ -343,7 +343,7 @@ fn build_runtime_with_models( async fn initialize_connection( runtime: &Arc, ) -> Result<(u64, mpsc::Receiver)> { - let (notifications_tx, notifications_rx) = mpsc::channel(/*buffer*/ 128); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(128); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; diff --git a/crates/server/tests/session_fork_persistence.rs b/crates/server/tests/session_fork_persistence.rs index 4c074809..86a8fad1 100644 --- a/crates/server/tests/session_fork_persistence.rs +++ b/crates/server/tests/session_fork_persistence.rs @@ -214,7 +214,7 @@ fn build_runtime(data_root: &Path) -> Result> { async fn initialize_connection( runtime: &Arc, ) -> Result<(u64, mpsc::Receiver)> { - let (notifications_tx, notifications_rx) = mpsc::channel(/*buffer*/ 128); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(128); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; diff --git a/crates/server/tests/session_rollback_persistence.rs b/crates/server/tests/session_rollback_persistence.rs index 3015cbc4..361a3cb0 100644 --- a/crates/server/tests/session_rollback_persistence.rs +++ b/crates/server/tests/session_rollback_persistence.rs @@ -246,7 +246,7 @@ fn build_runtime( async fn initialize_connection( runtime: &Arc, ) -> Result<(u64, mpsc::Receiver)> { - let (notifications_tx, notifications_rx) = mpsc::channel(/*buffer*/ 128); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(128); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; diff --git a/crates/server/tests/skills_integration.rs b/crates/server/tests/skills_integration.rs index 582b14da..d0497c5f 100644 --- a/crates/server/tests/skills_integration.rs +++ b/crates/server/tests/skills_integration.rs @@ -222,7 +222,7 @@ fn toml_path(path: &Path) -> String { async fn initialize_connection( runtime: &Arc, ) -> Result<(u64, mpsc::Receiver)> { - let (notifications_tx, notifications_rx) = mpsc::channel(/*buffer*/ 4096); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(4096); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; diff --git a/crates/server/tests/support/goal_continuation.rs b/crates/server/tests/support/goal_continuation.rs index f01e4afa..5e0d41fa 100644 --- a/crates/server/tests/support/goal_continuation.rs +++ b/crates/server/tests/support/goal_continuation.rs @@ -309,7 +309,7 @@ pub fn build_runtime_with_registry( pub async fn initialize_connection( runtime: &Arc, ) -> Result<(u64, mpsc::Receiver)> { - let (notifications_tx, notifications_rx) = mpsc::channel(/*buffer*/ 1024); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(1024); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; diff --git a/crates/server/tests/support/subagent_lifecycle.rs b/crates/server/tests/support/subagent_lifecycle.rs index 158fb979..d1b9fa02 100644 --- a/crates/server/tests/support/subagent_lifecycle.rs +++ b/crates/server/tests/support/subagent_lifecycle.rs @@ -289,7 +289,7 @@ pub fn build_runtime( pub async fn initialize_connection( runtime: &Arc, ) -> Result<(u64, mpsc::Receiver)> { - let (notifications_tx, notifications_rx) = mpsc::channel(/*buffer*/ 4096); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(4096); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; diff --git a/crates/server/tests/turn_start_persistence.rs b/crates/server/tests/turn_start_persistence.rs index 61bff983..60d95a5b 100644 --- a/crates/server/tests/turn_start_persistence.rs +++ b/crates/server/tests/turn_start_persistence.rs @@ -517,7 +517,7 @@ fn request_messages_json(request: &ModelRequest) -> Result { async fn initialize_connection( runtime: &Arc, ) -> Result<(u64, mpsc::Receiver)> { - let (notifications_tx, notifications_rx) = mpsc::channel(/*buffer*/ 128); + let (notifications_tx, notifications_rx) = devo_server::test_outbound_channel(128); let connection_id = runtime .register_connection(ClientTransportKind::Stdio, notifications_tx) .await; From 142b606c08112693f2939dae6090144b52132480 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Thu, 2 Jul 2026 23:49:11 -1000 Subject: [PATCH 06/16] docs: document singleton server and internal proxy lifecycle --- crates/server/src/bootstrap.rs | 39 ++++++++++++++++++++ crates/server/src/singleton.rs | 47 ++++++++++++++++++++++++ crates/server/src/transport.rs | 9 +++++ crates/server/tests/deep_research_e2e.rs | 18 --------- 4 files changed, 95 insertions(+), 18 deletions(-) diff --git a/crates/server/src/bootstrap.rs b/crates/server/src/bootstrap.rs index 438b4951..79f590cc 100644 --- a/crates/server/src/bootstrap.rs +++ b/crates/server/src/bootstrap.rs @@ -80,22 +80,53 @@ pub struct ServerProcessRunOptions { /// Starts the transport-facing server runtime using the resolved application /// configuration and listener set. +/// +/// ## Singleton server (`singleton.rs`) +/// +/// Devo allows at most **one real server process** per `DEVO_HOME`. Coordination +/// uses a file lock (`server.lock`) plus metadata (`server.lock.json`) that +/// records pid, a loopback WebSocket endpoint, and an auth token. +/// +/// - **`SingletonRole::Real`**: this process acquired the lock and becomes the +/// sole server. It binds an internal proxy listener, writes metadata, and runs +/// until shutdown. +/// - **`SingletonRole::Proxy`**: another process already holds the lock. This +/// process does not start a second runtime: stdio mode forwards to the +/// existing server via `run_stdio_proxy`; `--status` / `--shutdown` talk to +/// the internal control channel on that server. +/// +/// ## Internal proxy (`run_listeners_with_internal_proxy`) +/// +/// The real server exposes an extra **loopback-only** WebSocket listener +/// (`127.0.0.1:0`, ephemeral port). It is used for: +/// +/// 1. **Stdio proxy clients** — a second `devo server --transport stdio` connects +/// here and pipes stdin/stdout through WebSocket frames (see `run_stdio_proxy`). +/// 2. **Control plane** — `devo server --status` / `--shutdown` send +/// `_devo/server/status` or `_devo/server/shutdown` after token auth. +/// +/// The published `endpoint` in `server.lock.json` is this internal proxy URL, not +/// the public config WebSocket address. pub async fn run_server_process( args: ServerProcessArgs, options: ServerProcessRunOptions, ) -> Result<()> { let resolver = FileSystemConfigPathResolver::from_env()?; let action = args.action()?; + // Decide whether this process is the one true server or a lightweight proxy/ + // control client. Lock file lives under DEVO_HOME (see singleton.rs). let singleton_role = acquire_singleton_role(&resolver.user_config_dir())?; let real_server_guard = match singleton_role { SingletonRole::Real(guard) => match action { ServerProcessAction::Run => guard, ServerProcessAction::Status | ServerProcessAction::Shutdown => { + // We hold the lock but were not asked to run — no metadata file yet. println!("devo server is not running"); return Ok(()); } }, SingletonRole::Proxy(metadata) => match action { + // Another server is already running: pipe stdio to its internal proxy. ServerProcessAction::Run if args.transport == ServerTransportMode::Stdio => { tracing::info!( pid = metadata.pid, @@ -104,11 +135,13 @@ pub async fn run_server_process( ); return run_stdio_proxy(metadata).await; } + // Non-stdio second instances are rejected (would duplicate listeners). ServerProcessAction::Run => { print_existing_server_status(&metadata, "already running"); println!("Use `devo server --shutdown` to stop it."); return Ok(()); } + // `--status` / `--shutdown`: one-shot WebSocket control, then exit. ServerProcessAction::Status => { let result = run_server_control(&metadata, ServerControlAction::Status).await?; print_existing_server_status(&metadata, result.status.as_str()); @@ -121,7 +154,10 @@ pub async fn run_server_process( } }, }; + // Real server: bind ephemeral loopback WS for stdio-proxy + control clients. let internal_proxy = InternalProxyEndpoint::bind().await?; + // Persist ws://127.0.0.1: + random token into server.lock.json so + // proxy/control processes know where and how to connect. let singleton_metadata = real_server_guard.publish_endpoint(internal_proxy.endpoint().to_string())?; @@ -224,6 +260,9 @@ pub async fn run_server_process( let internal_proxy_control = InternalProxyControl::new(shutdown_signal.clone()); let external_shutdown = options.external_shutdown.clone(); + // Concurrent listeners: configured stdio/ws targets + internal proxy task. + // Returns when any listener exits; shutdown also via Ctrl+C, external token, + // or internal-proxy `_devo/server/shutdown` (cancels shutdown_signal). tokio::select! { result = run_listeners_with_internal_proxy( runtime.clone(), diff --git a/crates/server/src/singleton.rs b/crates/server/src/singleton.rs index b3b02808..7b6028c2 100644 --- a/crates/server/src/singleton.rs +++ b/crates/server/src/singleton.rs @@ -1,3 +1,17 @@ +//! Singleton server coordination for a single `DEVO_HOME`. +//! +//! At most one **real** devo-server process may run per user data directory. +//! Coordination uses an exclusive file lock plus a small JSON metadata file: +//! +//! - `server.lock` — held for the lifetime of the real server (`RealServerGuard`). +//! - `server.lock.json` — written after the real server binds its internal proxy; +//! contains pid, loopback WebSocket URL, and a random auth token. +//! +//! A second process that fails to acquire the lock becomes a **proxy** client: +//! stdio mode forwards stdin/stdout through the internal proxy WebSocket +//! (`run_stdio_proxy`); `--status` / `--shutdown` use one-shot control RPCs +//! (`run_server_control`). See `bootstrap::run_server_process` for the full flow. + use std::fs; use std::fs::File; use std::fs::OpenOptions; @@ -25,23 +39,32 @@ use uuid::Uuid; const LOCK_FILE_NAME: &str = "server.lock"; const METADATA_FILE_NAME: &str = "server.lock.json"; const METADATA_VERSION: u32 = 1; +/// Real server may still be writing metadata when a proxy process starts. const METADATA_READ_RETRIES: usize = 100; const METADATA_READ_RETRY_DELAY: Duration = Duration::from_millis(50); pub(crate) const SERVER_CONTROL_STATUS_METHOD: &str = "_devo/server/status"; pub(crate) const SERVER_CONTROL_SHUTDOWN_METHOD: &str = "_devo/server/shutdown"; const SERVER_CONTROL_REQUEST_ID: u64 = 1; +/// Outcome of [`acquire_singleton_role`]: either this process runs the server +/// or it should connect to an already-running instance using the metadata. #[derive(Debug)] pub(crate) enum SingletonRole { + /// Exclusive lock acquired; caller must bind internal proxy and + /// [`RealServerGuard::publish_endpoint`]. Real(RealServerGuard), + /// Another process holds the lock; use `endpoint` + `token` to connect. Proxy(ServerLockMetadata), } +/// Published in `server.lock.json` while the real server is running. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub(crate) struct ServerLockMetadata { pub(crate) version: u32, pub(crate) pid: u32, + /// Loopback WebSocket URL of the real server's internal proxy listener. pub(crate) endpoint: String, + /// Must be sent as the first WebSocket text frame after connect. pub(crate) token: String, pub(crate) started_at: String, } @@ -78,6 +101,8 @@ impl ServerLockMetadata { } } +/// One-shot control RPC against the real server's internal proxy +/// (`devo server --status` / `--shutdown` from a proxy process). pub(crate) async fn run_server_control( metadata: &ServerLockMetadata, action: ServerControlAction, @@ -86,6 +111,7 @@ pub(crate) async fn run_server_control( .await .with_context(|| format!("connect to singleton server {}", metadata.endpoint))?; let (mut writer, mut reader) = socket.split(); + // Internal proxy requires token auth before any other traffic. writer .send(Message::Text(metadata.token.clone().into())) .await @@ -111,6 +137,8 @@ pub(crate) async fn run_server_control( } } +/// Holds the exclusive singleton lock until dropped; releasing the lock allows +/// another process to become the real server. #[derive(Debug)] pub(crate) struct RealServerGuard { lock_file: File, @@ -118,6 +146,7 @@ pub(crate) struct RealServerGuard { } impl RealServerGuard { + /// Writes `server.lock.json` so proxy/control clients can find this server. pub(crate) fn publish_endpoint(&self, endpoint: String) -> Result { let metadata = ServerLockMetadata::new(endpoint, Uuid::new_v4().to_string()); let encoded = serde_json::to_vec_pretty(&metadata).context("serialize server metadata")?; @@ -133,11 +162,18 @@ impl RealServerGuard { impl Drop for RealServerGuard { fn drop(&mut self) { + // Best-effort cleanup so a crashed predecessor does not block forever + // once the OS releases the file handle; metadata removal signals "not running". let _ = fs::remove_file(&self.metadata_path); let _ = self.lock_file.unlock(); } } +/// Attempts an exclusive lock on `DEVO_HOME/server.lock`. +/// +/// - Success → [`SingletonRole::Real`]: caller becomes the sole server process. +/// - Lock held by another process → [`SingletonRole::Proxy`]: read metadata and +/// connect instead of starting a second runtime. pub(crate) fn acquire_singleton_role(devo_home: &Path) -> Result { fs::create_dir_all(devo_home) .with_context(|| format!("create DEVO_HOME {}", devo_home.display()))?; @@ -156,6 +192,7 @@ pub(crate) fn acquire_singleton_role(devo_home: &Path) -> Result metadata_path: metadata_path(devo_home), })), Err(error) if is_lock_temporarily_unavailable(&error) => { + // Real server may be mid-startup; retry until metadata appears. Ok(SingletonRole::Proxy(read_metadata_with_retry(devo_home)?)) } Err(error) => { @@ -164,6 +201,13 @@ pub(crate) fn acquire_singleton_role(devo_home: &Path) -> Result } } +/// Lightweight stdio front-end for an already-running real server. +/// +/// Connects to the internal proxy WebSocket, authenticates with `token`, then: +/// - stdin lines (NDJSON) → WebSocket Text frames upstream +/// - WebSocket Text frames → stdout (one line per frame, NDJSON framing preserved) +/// +/// On the real server this connection is registered as `ClientTransportKind::StdioProxy`. pub(crate) async fn run_stdio_proxy(metadata: ServerLockMetadata) -> Result<()> { let (socket, _) = connect_async(metadata.endpoint.as_str()) .await @@ -199,6 +243,8 @@ pub(crate) async fn run_stdio_proxy(metadata: ServerLockMetadata) -> Result<()> .write_all(text.as_bytes()) .await .context("write singleton response to stdout")?; + // Real server sends one JSON-RPC message per WebSocket frame; + // re-add the newline stdio clients expect. stdout .write_all(b"\n") .await @@ -263,6 +309,7 @@ fn is_lock_temporarily_unavailable(error: &std::io::Error) -> bool { } } +/// Polls until the real server writes metadata (startup race with lock holder). fn read_metadata_with_retry(devo_home: &Path) -> Result { let mut last_error = None; for _ in 0..METADATA_READ_RETRIES { diff --git a/crates/server/src/transport.rs b/crates/server/src/transport.rs index 13be8bb2..84b7b852 100644 --- a/crates/server/src/transport.rs +++ b/crates/server/src/transport.rs @@ -306,6 +306,15 @@ pub async fn run_listeners(runtime: Arc, listen: &[String]) -> Re } /// Runs configured listener targets plus the internal stdio-proxy endpoint. +/// +/// Spawns one async task per listen target (stdio and/or public WebSocket) and, +/// when `internal_proxy` is `Some`, an additional loopback WebSocket accept loop. +/// All tasks run concurrently inside a `JoinSet`; this function returns when the +/// first task completes (normally only on fatal listener error). +/// +/// The internal proxy task authenticates clients with `token`, handles one-shot +/// control RPCs (`status` / `shutdown`), then treats the connection as a +/// `ClientTransportKind::StdioProxy` ACP session (see `handle_internal_proxy_connection`). pub async fn run_listeners_with_internal_proxy( runtime: Arc, listen: &[String], diff --git a/crates/server/tests/deep_research_e2e.rs b/crates/server/tests/deep_research_e2e.rs index 07f9d118..18c04ef4 100644 --- a/crates/server/tests/deep_research_e2e.rs +++ b/crates/server/tests/deep_research_e2e.rs @@ -109,24 +109,6 @@ impl ScriptedResearchProvider { expected_cwd: expected_cwd.display().to_string(), } } - - fn with_delegated_worker_failure_once(expected_cwd: &std::path::Path) -> Self { - Self::with_delegated_worker_failures_before_success(expected_cwd, 1) - } - - fn with_delegated_worker_failures_before_success( - expected_cwd: &std::path::Path, - failures: usize, - ) -> Self { - Self { - stream_calls: AtomicUsize::new(0), - final_report_stream_calls: AtomicUsize::new(0), - delegated_worker_failures_before_success: AtomicUsize::new(failures), - researcher_gate: None, - final_report_mode: ScriptedFinalReportMode::Text, - expected_cwd: expected_cwd.display().to_string(), - } - } } #[async_trait] From 883c1b6e571fad21c00f096f8f2fe695fe7bb019 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Fri, 3 Jul 2026 05:12:28 -1000 Subject: [PATCH 07/16] fix: Extract shared client protocol layer and remove TUI session/prompt fallback. --- crates/client/src/acp_fs.rs | 4 + crates/client/src/acp_permissions.rs | 6 +- crates/client/src/acp_terminal.rs | 2 +- crates/client/src/client_core.rs | 370 ++++- crates/client/src/lib.rs | 9 +- crates/client/src/stdio.rs | 1368 +++-------------- crates/client/src/websocket.rs | 4 +- .../server/src/runtime/handlers/acp/prompt.rs | 27 + crates/server/src/transport.rs | 21 +- crates/server/tests/deep_research_e2e.rs | 15 - crates/tui/src/worker.rs | 248 +-- 11 files changed, 653 insertions(+), 1421 deletions(-) diff --git a/crates/client/src/acp_fs.rs b/crates/client/src/acp_fs.rs index bba1b6bb..b5367464 100644 --- a/crates/client/src/acp_fs.rs +++ b/crates/client/src/acp_fs.rs @@ -1,3 +1,7 @@ +//! ACP client-side filesystem handlers for server-initiated `fs/read` and +//! `fs/write` requests. Each handler echoes the JSON-RPC `id` from the inbound +//! request in its success or error response. + use std::path::Path; use devo_protocol::ACP_FS_READ_TEXT_FILE_METHOD; diff --git a/crates/client/src/acp_permissions.rs b/crates/client/src/acp_permissions.rs index 4ba09fb4..659ee7fd 100644 --- a/crates/client/src/acp_permissions.rs +++ b/crates/client/src/acp_permissions.rs @@ -1,3 +1,7 @@ +//! ACP permission prompt bridge: server `session/request_permission` is +//! surfaced to the UI via notifications; the JSON-RPC response is sent later +//! when the user approves or denies, using the original request `id`. + use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::AtomicU64; @@ -21,7 +25,7 @@ use devo_protocol::acp_success_response; use tokio::sync::Mutex; use tokio::sync::mpsc; -use crate::stdio::ServerNotificationMessage; +use crate::client_core::ServerNotificationMessage; static ACP_PERMISSION_NEXT_ID: AtomicU64 = AtomicU64::new(1); diff --git a/crates/client/src/acp_terminal.rs b/crates/client/src/acp_terminal.rs index 9acf8301..d750a2b2 100644 --- a/crates/client/src/acp_terminal.rs +++ b/crates/client/src/acp_terminal.rs @@ -26,7 +26,7 @@ use tokio::sync::Notify; use tokio::sync::mpsc; use tokio::time::Duration; -use crate::stdio::ServerNotificationMessage; +use crate::client_core::ServerNotificationMessage; static ACP_TERMINAL_NEXT_ID: AtomicU64 = AtomicU64::new(1); pub const ACP_TERMINAL_OUTPUT_NOTIFICATION_METHOD: &str = "_devo/acp_terminal/output"; diff --git a/crates/client/src/client_core.rs b/crates/client/src/client_core.rs index 7adbb003..f5c9c8b3 100644 --- a/crates/client/src/client_core.rs +++ b/crates/client/src/client_core.rs @@ -1,10 +1,24 @@ +//! Transport-agnostic Devo server client: JSON-RPC request/response routing, +//! server-initiated ACP client handlers, and notification demultiplexing. +//! +//! Both [`crate::stdio::StdioServerClient`] and [`crate::websocket::WebSocketServerClient`] +//! delegate protocol logic here. Incoming messages are classified as: +//! +//! - **Server → client requests** (`id` + `method`): handled asynchronously; the +//! response echoes the same JSON-RPC `id` (see `fs/read`, permissions, terminal). +//! - **Server responses** (`id` + `result`/`error`, no `method`): matched against +//! [`PendingResponses`] via numeric `id` to complete a client-initiated `request`. +//! - **Notifications** (no `id`): forwarded on the notification channel. + use std::collections::HashMap; use std::collections::HashSet; use std::path::PathBuf; use std::sync::Arc; +use std::sync::OnceLock; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use std::time::Duration; +use std::time::Instant; use anyhow::Context; use anyhow::Result; @@ -25,13 +39,21 @@ use crate::acp_permissions::handle_acp_request_permission; use crate::acp_permissions::resolve_acp_permission_response; use crate::acp_terminal::AcpTerminalManager; use crate::acp_terminal::handle_acp_terminal_request; -use crate::stdio::ACP_PROMPT_COMPLETED_NOTIFICATION_METHOD; -use crate::stdio::ACP_PROMPT_STARTED_NOTIFICATION_METHOD; -use crate::stdio::ServerNotificationMessage; + +pub const ACP_PROMPT_STARTED_NOTIFICATION_METHOD: &str = "_devo/acp_prompt/started"; +pub const ACP_PROMPT_COMPLETED_NOTIFICATION_METHOD: &str = "_devo/acp_prompt/completed"; const SERVER_RESPONSE_TIMEOUT: Duration = Duration::from_secs(10); -type PendingResponses = Arc>>>; +/// Synthetic notifications emitted when falling back to detached `session/prompt`. +#[derive(Debug, Clone)] +pub struct ServerNotificationMessage { + pub method: String, + pub params: serde_json::Value, +} + +/// Client-initiated requests awaiting a server response, keyed by JSON-RPC `id`. +pub(crate) type PendingResponses = Arc>>>; pub(crate) enum ClientWriteMessage { Json(serde_json::Value), @@ -112,6 +134,23 @@ impl ServerClientCore { } } + pub(crate) fn set_client_capabilities(&mut self, client_capabilities: AcpClientCapabilities) { + self.client_capabilities = client_capabilities; + } + + #[cfg(test)] + pub(crate) fn pending_responses(&self) -> PendingResponses { + Arc::clone(&self.pending) + } + + #[cfg(test)] + pub(crate) fn set_agent_capabilities_for_test( + &mut self, + capabilities: AcpAgentCapabilities, + ) { + self.acp_agent_capabilities = Some(capabilities); + } + pub(crate) async fn initialize(&mut self) -> Result { let result: AcpInitializeResult = timeout( SERVER_RESPONSE_TIMEOUT, @@ -302,6 +341,8 @@ impl ServerClientCore { { let request_id = self.next_request_id.fetch_add(1, Ordering::SeqCst); let (response_tx, response_rx) = oneshot::channel(); + // The transport reader resolves responses by this id (see + // `deliver_pending_client_response`). self.pending.lock().await.insert(request_id, response_tx); let request = AcpClientRequest::new(serde_json::json!(request_id), method, params); if let Err(error) = self.writer.send_serializable(&request) { @@ -389,6 +430,212 @@ impl ServerClientCore { self.acp_terminals.release_all().await; } + pub(crate) async fn agent_list(&mut self, params: AgentListParams) -> Result { + self.request_devo("agent/list", params).await + } + + pub(crate) async fn agent_spawn(&mut self, params: SpawnAgentParams) -> Result { + self.request_devo("agent/spawn", params).await + } + + pub(crate) async fn agent_close(&mut self, params: CloseAgentParams) -> Result { + self.request_devo("agent/close", params).await + } + + pub(crate) async fn session_title_update( + &mut self, + params: SessionTitleUpdateParams, + ) -> Result { + self.request_devo("session/title/update", params).await + } + + pub(crate) async fn session_metadata_update( + &mut self, + params: SessionMetadataUpdateParams, + ) -> Result { + self.request_devo("session/metadata/update", params).await + } + + pub(crate) async fn session_permissions_update( + &mut self, + params: SessionPermissionsUpdateParams, + ) -> Result { + self.request_devo("session/permissions/update", params).await + } + + pub(crate) async fn session_compact( + &mut self, + params: SessionCompactParams, + ) -> Result { + self.request_devo("session/compact", params).await + } + + pub(crate) async fn goal_create(&mut self, params: GoalCreateParams) -> Result { + self.request_devo("goal/create", params).await + } + + pub(crate) async fn goal_set(&mut self, params: GoalSetParams) -> Result { + self.request_devo("goal/set", params).await + } + + pub(crate) async fn goal_status(&mut self, params: GoalStatusParams) -> Result { + self.request_devo("goal/status", params).await + } + + pub(crate) async fn goal_pause(&mut self, params: GoalSetStatusParams) -> Result { + self.request_devo("goal/pause", params).await + } + + pub(crate) async fn goal_resume( + &mut self, + params: GoalSetStatusParams, + ) -> Result { + self.request_devo("goal/resume", params).await + } + + pub(crate) async fn goal_complete( + &mut self, + params: GoalSetStatusParams, + ) -> Result { + self.request_devo("goal/complete", params).await + } + + pub(crate) async fn goal_clear(&mut self, params: GoalClearParams) -> Result { + self.request_devo("goal/clear", params).await + } + + pub(crate) async fn session_fork(&mut self, params: SessionForkParams) -> Result { + self.request_devo("session/fork", params).await + } + + pub(crate) async fn session_rollback( + &mut self, + params: SessionRollbackParams, + ) -> Result { + self.request_devo("session/rollback", params).await + } + + pub(crate) async fn skills_list(&mut self, params: SkillListParams) -> Result { + self.request_devo("skills/list", params).await + } + + pub(crate) async fn skills_changed( + &mut self, + params: SkillChangedParams, + ) -> Result { + self.request_devo("skills/changed", params).await + } + + pub(crate) async fn skills_set_enabled( + &mut self, + params: SkillSetEnabledParams, + ) -> Result { + self.request_devo("skills/set_enabled", params).await + } + + pub(crate) async fn model_catalog( + &mut self, + params: ModelCatalogParams, + ) -> Result { + self.request_devo("model/catalog", params).await + } + + pub(crate) async fn model_saved(&mut self, params: ModelSavedParams) -> Result { + self.request_devo("model/saved", params).await + } + + pub(crate) async fn provider_vendor_list( + &mut self, + params: ProviderVendorListParams, + ) -> Result { + self.request_devo("provider/list", params).await + } + + pub(crate) async fn provider_vendor_upsert( + &mut self, + params: ProviderVendorUpsertParams, + ) -> Result { + self.request_devo("provider/upsert", params).await + } + + pub(crate) async fn provider_validate( + &mut self, + params: ProviderValidateParams, + ) -> Result { + self.request_devo("provider/validate", params).await + } + + pub(crate) async fn command_exec(&mut self, params: CommandExecParams) -> Result { + self.request_devo("command/exec", params).await + } + + pub(crate) async fn command_exec_write( + &mut self, + params: CommandExecWriteParams, + ) -> Result { + self.request_devo("command/exec/write", params).await + } + + pub(crate) async fn command_exec_resize( + &mut self, + params: CommandExecResizeParams, + ) -> Result { + self.request_devo("command/exec/resize", params).await + } + + pub(crate) async fn command_exec_terminate( + &mut self, + params: CommandExecTerminateParams, + ) -> Result { + self.request_devo("command/exec/terminate", params).await + } + + pub(crate) async fn turn_shell_command( + &mut self, + params: ShellCommandParams, + ) -> Result { + self.request_devo("turn/shell_command", params).await + } + + pub(crate) async fn turn_interrupt( + &mut self, + params: TurnInterruptParams, + ) -> Result { + self.request_devo("turn/interrupt", params).await + } + + pub(crate) async fn turn_steer(&mut self, params: TurnSteerParams) -> Result { + self.request_devo("turn/steer", params).await + } + + pub(crate) async fn reference_search_start( + &mut self, + params: ReferenceSearchStartParams, + ) -> Result { + self.request_devo("search/start", params).await + } + + pub(crate) async fn reference_search_update( + &mut self, + params: ReferenceSearchUpdateParams, + ) -> Result { + self.request_devo("search/update", params).await + } + + pub(crate) async fn reference_search_cancel( + &mut self, + params: ReferenceSearchCancelParams, + ) -> Result { + self.request_devo("search/cancel", params).await + } + + /// Fallback when the server does not implement `_devo/turn/start`. + /// + /// Sends blocking ACP `session/prompt` but returns immediately after the + /// request is written. Completion is delivered later via synthetic + /// `_devo/acp_prompt/completed` notifications. Multiple detached prompts + /// may be in flight at once; each uses a distinct JSON-RPC `id` in + /// [`PendingResponses`], so responses do not collide as long as ids stay unique. async fn turn_start_acp_prompt_detached(&mut self, params: TurnStartParams) -> Result<()> { let session_id = params.session_id; let prompt = params @@ -451,6 +698,8 @@ impl ServerClientCore { } impl ServerClientReaderState { + /// Classifies one inbound server payload and dispatches without blocking the + /// transport read loop on handler work. pub(crate) async fn handle_message(&self, message: serde_json::Value) { if let (Some(id), Some(method)) = ( message.get("id").cloned(), @@ -462,15 +711,15 @@ impl ServerClientReaderState { .unwrap_or_else(|| serde_json::json!({})); let state = self.clone(); let method = method.to_string(); + // Server-initiated ACP client tools may run concurrently; each reply + // must echo the request `id` assigned by the server. tokio::spawn(async move { state.handle_client_request(id, &method, params).await; }); return; } if let Some(id) = message.get("id").and_then(serde_json::Value::as_u64) { - if let Some(tx) = self.pending.lock().await.remove(&id) { - let _ = tx.send(message); - } + deliver_pending_client_response(&self.pending, id, message).await; return; } if let Ok(notification) = @@ -515,6 +764,7 @@ impl ServerClientReaderState { }); return; } + log_notification_received(¬ification); let _ = self.notifications_tx.send(ServerNotificationMessage { method: notification.method, params: notification.params, @@ -576,6 +826,21 @@ impl ServerClientReaderState { } } +async fn deliver_pending_client_response( + pending: &PendingResponses, + request_id: u64, + message: serde_json::Value, +) { + if let Some(tx) = pending.lock().await.remove(&request_id) { + let _ = tx.send(message); + } else { + tracing::warn!( + request_id, + "dropping server response for unknown client request id" + ); + } +} + fn acp_client_error_response( id: serde_json::Value, code: i64, @@ -777,3 +1042,94 @@ fn acp_title_state(title: &Option) -> SessionTitleState { SessionTitleState::Unset } } + +fn log_notification_received(notification: &NotificationEnvelope) { + let event_seq = notification + .params + .get("context") + .and_then(|context| context.get("seq")) + .and_then(serde_json::Value::as_u64); + let item_id = notification + .params + .get("context") + .and_then(|context| context.get("item_id")) + .and_then(serde_json::Value::as_str); + let assistant_delta = (notification.method == "item/agentMessage/delta") + .then(|| notification.params.get("delta")?.as_str()) + .flatten(); + let delta_len = assistant_delta.map(str::len); + let assistant_token_text = assistant_delta.and_then(assistant_token_log_preview); + if let Some(assistant_token_text) = assistant_token_text.as_deref() { + tracing::debug!( + stream_elapsed_ms = stream_trace_elapsed_ms(), + method = %notification.method, + event_seq, + item_id = ?item_id, + delta_len = ?delta_len, + assistant_token_text, + "client received server notification" + ); + } else { + tracing::debug!( + stream_elapsed_ms = stream_trace_elapsed_ms(), + method = %notification.method, + event_seq, + item_id = ?item_id, + delta_len = ?delta_len, + "client received server notification" + ); + } +} + +fn stream_trace_elapsed_ms() -> u128 { + static STREAM_TRACE_START: OnceLock = OnceLock::new(); + STREAM_TRACE_START + .get_or_init(Instant::now) + .elapsed() + .as_millis() +} + +fn assistant_token_log_preview(text: &str) -> Option { + assistant_token_logging_enabled().then(|| { + let max_chars = assistant_token_log_max_chars(); + format_assistant_token_log_preview(text, max_chars) + }) +} + +fn assistant_token_logging_enabled() -> bool { + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| { + std::env::var("DEVO_LOG_ASSISTANT_TOKEN_TEXT") + .ok() + .is_some_and(|value| { + matches!( + value.as_str(), + "1" | "true" | "TRUE" | "yes" | "YES" | "on" | "ON" + ) + }) + }) +} + +fn assistant_token_log_max_chars() -> usize { + static ASSISTANT_TOKEN_LOG_MAX_CHARS: OnceLock = OnceLock::new(); + *ASSISTANT_TOKEN_LOG_MAX_CHARS.get_or_init(|| { + std::env::var("DEVO_ASSISTANT_TOKEN_LOG_MAX_CHARS") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|value| *value > 0) + .unwrap_or(512) + }) +} + +fn format_assistant_token_log_preview(text: &str, max_chars: usize) -> String { + let max_chars = max_chars.max(1); + let mut preview = String::with_capacity(text.len().min(max_chars)); + let mut chars = text.chars(); + for ch in chars.by_ref().take(max_chars) { + preview.extend(ch.escape_default()); + } + if chars.next().is_some() { + preview.push_str("..."); + } + preview +} diff --git a/crates/client/src/lib.rs b/crates/client/src/lib.rs index 4ac3dc7e..e6d8a356 100644 --- a/crates/client/src/lib.rs +++ b/crates/client/src/lib.rs @@ -1,8 +1,8 @@ //! Client-side transport API for talking to a Devo server. //! -//! The crate currently exposes the stdio client used by local frontends. It -//! keeps request/notification framing here so UI crates can call typed protocol -//! methods without owning process I/O or response demultiplexing. +//! Protocol logic (JSON-RPC routing, pending response maps, ACP client handlers) +//! lives in [`client_core`]. [`stdio::StdioServerClient`] and +//! [`websocket::WebSocketServerClient`] are thin transport adapters. mod acp_fs; mod acp_permissions; @@ -11,5 +11,8 @@ mod client_core; mod stdio; mod websocket; +pub use client_core::ACP_PROMPT_COMPLETED_NOTIFICATION_METHOD; +pub use client_core::ACP_PROMPT_STARTED_NOTIFICATION_METHOD; +pub use client_core::ServerNotificationMessage; pub use stdio::*; pub use websocket::*; diff --git a/crates/client/src/stdio.rs b/crates/client/src/stdio.rs index b64f8dce..dd422b7f 100644 --- a/crates/client/src/stdio.rs +++ b/crates/client/src/stdio.rs @@ -1,148 +1,16 @@ //! Stdio transport for a spawned devo server process. //! -//! The client writes newline-delimited JSON requests to child stdin and owns a -//! background stdout reader that demultiplexes responses by request id while -//! forwarding id-less messages as server notifications. +//! Writes newline-delimited JSON to the child stdin and runs a background stdout +//! reader that delegates framing and JSON-RPC routing to [`ServerClientCore`]. -use std::collections::HashMap; -use std::collections::HashSet; use std::path::PathBuf; use std::process::Stdio; use std::sync::Arc; -use std::sync::OnceLock; -use std::sync::atomic::AtomicU64; -use std::sync::atomic::Ordering; -use std::time::Instant; - -use crate::acp_fs::handle_acp_fs_request; -use crate::acp_permissions::AcpPendingPermissions; -use crate::acp_permissions::handle_acp_request_permission; -use crate::acp_permissions::resolve_acp_permission_response; -use crate::acp_terminal::AcpTerminalManager; -use crate::acp_terminal::handle_acp_terminal_request; +use std::time::Duration; + use anyhow::Context; use anyhow::Result; -use anyhow::bail; -use chrono::Utc; -use devo_protocol::ACP_FS_READ_TEXT_FILE_METHOD; -use devo_protocol::ACP_FS_WRITE_TEXT_FILE_METHOD; -use devo_protocol::ACP_INITIALIZE_METHOD; -use devo_protocol::ACP_SESSION_LIST_METHOD; -use devo_protocol::ACP_SESSION_NEW_METHOD; -use devo_protocol::ACP_SESSION_PROMPT_METHOD; -use devo_protocol::ACP_SESSION_RESUME_METHOD; -use devo_protocol::ACP_SESSION_UPDATE_METHOD; -use devo_protocol::ACP_TERMINAL_CREATE_METHOD; -use devo_protocol::ACP_TERMINAL_KILL_METHOD; -use devo_protocol::ACP_TERMINAL_OUTPUT_METHOD; -use devo_protocol::ACP_TERMINAL_RELEASE_METHOD; -use devo_protocol::ACP_TERMINAL_WAIT_FOR_EXIT_METHOD; -use devo_protocol::AcpAgentCapabilities; -use devo_protocol::AcpClientRequest; -use devo_protocol::AcpContentBlock; -use devo_protocol::AcpImplementation; -use devo_protocol::AcpInitializeParams; -use devo_protocol::AcpInitializeResult; -use devo_protocol::AcpListSessionsParams; -use devo_protocol::AcpListSessionsResult; -use devo_protocol::AcpNewSessionParams; -use devo_protocol::AcpNewSessionResult; -use devo_protocol::AcpPromptParams; -use devo_protocol::AcpPromptResult; -use devo_protocol::AcpResumeSessionParams; -use devo_protocol::AcpResumeSessionResult; -use devo_protocol::AcpSessionInfo; -use devo_protocol::AcpSessionNotification; -use devo_protocol::AcpSuccessResponse; -use devo_protocol::AcpTerminalOutputResult; -use devo_protocol::AgentListParams; -use devo_protocol::AgentListResult; -use devo_protocol::ApprovalResponseParams; -use devo_protocol::CloseAgentParams; -use devo_protocol::CloseAgentResult; -use devo_protocol::CommandExecParams; -use devo_protocol::CommandExecResizeParams; -use devo_protocol::CommandExecResizeResult; -use devo_protocol::CommandExecResult; -use devo_protocol::CommandExecTerminateParams; -use devo_protocol::CommandExecTerminateResult; -use devo_protocol::CommandExecWriteParams; -use devo_protocol::CommandExecWriteResult; -use devo_protocol::DEVO_SESSION_RESUME_META; -use devo_protocol::ErrorResponse; -use devo_protocol::GoalClearParams; -use devo_protocol::GoalClearResult; -use devo_protocol::GoalCreateParams; -use devo_protocol::GoalCreateResult; -use devo_protocol::GoalSetParams; -use devo_protocol::GoalSetResult; -use devo_protocol::GoalSetStatusParams; -use devo_protocol::GoalSetStatusResult; -use devo_protocol::GoalStatusParams; -use devo_protocol::GoalStatusResult; -use devo_protocol::InitializeResult; -use devo_protocol::ModelCatalogParams; -use devo_protocol::ModelCatalogResult; -use devo_protocol::ModelSavedParams; -use devo_protocol::ModelSavedResult; -use devo_protocol::NotificationEnvelope; -use devo_protocol::ProtocolErrorCode; -use devo_protocol::ProviderValidateParams; -use devo_protocol::ProviderValidateResult; -use devo_protocol::ProviderVendorListParams; -use devo_protocol::ProviderVendorListResult; -use devo_protocol::ProviderVendorUpsertParams; -use devo_protocol::ProviderVendorUpsertResult; -use devo_protocol::ReferenceSearchCancelParams; -use devo_protocol::ReferenceSearchCancelResult; -use devo_protocol::ReferenceSearchStartParams; -use devo_protocol::ReferenceSearchStartResult; -use devo_protocol::ReferenceSearchUpdateParams; -use devo_protocol::ReferenceSearchUpdateResult; -use devo_protocol::RequestUserInputRespondParams; -use devo_protocol::ServerEvent; -use devo_protocol::SessionCompactParams; -use devo_protocol::SessionCompactResult; -use devo_protocol::SessionForkParams; -use devo_protocol::SessionForkResult; -use devo_protocol::SessionMetadata; -use devo_protocol::SessionMetadataUpdateParams; -use devo_protocol::SessionMetadataUpdateResult; -use devo_protocol::SessionPermissionsUpdateParams; -use devo_protocol::SessionPermissionsUpdateResult; -use devo_protocol::SessionResumeParams; -use devo_protocol::SessionResumeResult; -use devo_protocol::SessionRollbackParams; -use devo_protocol::SessionRollbackResult; -use devo_protocol::SessionRuntimeStatus; -use devo_protocol::SessionStartParams; -use devo_protocol::SessionStartResult; -use devo_protocol::SessionTitleState; -use devo_protocol::SessionTitleUpdateParams; -use devo_protocol::SessionTitleUpdateResult; -use devo_protocol::ShellCommandParams; -use devo_protocol::ShellCommandResult; -use devo_protocol::SkillChangedParams; -use devo_protocol::SkillChangedResult; -use devo_protocol::SkillListParams; -use devo_protocol::SkillListResult; -use devo_protocol::SkillSetEnabledParams; -use devo_protocol::SkillSetEnabledResult; -use devo_protocol::SpawnAgentParams; -use devo_protocol::SpawnAgentResult; -use devo_protocol::TurnId; -use devo_protocol::TurnInterruptParams; -use devo_protocol::TurnInterruptResult; -use devo_protocol::TurnStartParams; -use devo_protocol::TurnStartResult; -use devo_protocol::TurnStatus; -use devo_protocol::TurnSteerParams; -use devo_protocol::TurnSteerResult; -use devo_protocol::devo_extension_inner_method; -use devo_protocol::devo_extension_method; -use devo_protocol::original_event_from_acp_notification; -use serde::de::DeserializeOwned; -use tokio::io::AsyncBufRead; +use devo_protocol::*; use tokio::io::AsyncBufReadExt; use tokio::io::AsyncWriteExt; use tokio::io::BufReader; @@ -151,18 +19,21 @@ use tokio::process::ChildStderr; use tokio::process::ChildStdin; use tokio::process::Command; use tokio::sync::Mutex; -use tokio::sync::mpsc; -use tokio::sync::oneshot; -use tokio::time::Duration; +use tokio::task::JoinHandle; use tokio::time::timeout; -const SERVER_CHILD_STDIN_SHUTDOWN_TIMEOUT: Duration = Duration::from_millis(100); -const SERVER_CHILD_EXIT_TIMEOUT: Duration = Duration::from_millis(500); -pub const ACP_PROMPT_STARTED_NOTIFICATION_METHOD: &str = "_devo/acp_prompt/started"; -pub const ACP_PROMPT_COMPLETED_NOTIFICATION_METHOD: &str = "_devo/acp_prompt/completed"; +use crate::client_core::ClientWriteMessage; +use crate::client_core::ClientWriter; +use crate::client_core::ServerClientCore; + pub use crate::acp_terminal::ACP_TERMINAL_OUTPUT_NOTIFICATION_METHOD; +pub use crate::client_core::ACP_PROMPT_COMPLETED_NOTIFICATION_METHOD; +pub use crate::client_core::ACP_PROMPT_STARTED_NOTIFICATION_METHOD; +pub use crate::client_core::ServerNotificationMessage; -type PendingResponses = Arc>>>; +const SERVER_CHILD_STDIN_SHUTDOWN_TIMEOUT: Duration = Duration::from_millis(100); +const SERVER_CHILD_EXIT_TIMEOUT: Duration = Duration::from_millis(500); +const STDIO_INITIALIZE_TIMEOUT: Duration = Duration::from_secs(3); #[derive(Debug, Clone)] pub struct StdioServerClientConfig { @@ -170,22 +41,11 @@ pub struct StdioServerClientConfig { pub args: Vec, } -#[derive(Debug, Clone)] -pub struct ServerNotificationMessage { - pub method: String, - pub params: serde_json::Value, -} - pub struct StdioServerClient { child: Child, - stdin: Arc>, - pending: PendingResponses, - acp_pending_permissions: AcpPendingPermissions, - acp_terminals: AcpTerminalManager, - acp_agent_capabilities: Option, - next_request_id: AtomicU64, - notifications_rx: mpsc::UnboundedReceiver, - notifications_tx: mpsc::UnboundedSender, + core: ServerClientCore, + reader_task: JoinHandle<()>, + writer_task: JoinHandle>, } impl StdioServerClient { @@ -209,481 +69,325 @@ impl StdioServerClient { let stdin = child.stdin.take().context("capture server stdin")?; let stdout = child.stdout.take().context("capture server stdout")?; let stderr = child.stderr.take().context("capture server stderr")?; - let pending = Arc::new(Mutex::new( - HashMap::>::new(), - )); + + let (client_writer, write_rx) = ClientWriter::channel(); + let core = ServerClientCore::new( + client_writer, + AcpClientCapabilities { + fs: AcpFileSystemCapabilities { + read_text_file: false, + write_text_file: false, + meta: None, + }, + terminal: false, + meta: None, + }, + ); + let reader_state = core.reader_state(); let stdin = Arc::new(Mutex::new(stdin)); - let acp_pending_permissions = Arc::new(Mutex::new(HashMap::new())); - let acp_terminals = AcpTerminalManager::new(); - let (notifications_tx, notifications_rx) = mpsc::unbounded_channel(); - tokio::spawn(run_stdout_reader( - BufReader::new(stdout).lines(), - Arc::clone(&pending), + let writer_task = tokio::spawn(run_stdin_writer( Arc::clone(&stdin), - Arc::clone(&acp_pending_permissions), - acp_terminals.clone(), - notifications_tx.clone(), + write_rx, )); + let reader_task = tokio::spawn(async move { + let mut lines = BufReader::new(stdout).lines(); + while let Ok(Some(line)) = lines.next_line().await { + match serde_json::from_str::(&line) { + Ok(message) => reader_state.handle_message(message).await, + Err(_) => { + tracing::warn!(line = %line, "failed to parse JSON from server stdout"); + } + } + } + reader_state.finish_reader("stdio").await; + }); tokio::spawn(run_stderr_reader(BufReader::new(stderr).lines())); Ok(Self { child, - stdin, - pending, - acp_pending_permissions, - acp_terminals, - acp_agent_capabilities: None, - next_request_id: AtomicU64::new(1), - notifications_rx, - notifications_tx, + core, + reader_task, + writer_task, }) } pub async fn initialize( &mut self, - client_capabilities: &devo_protocol::AcpClientCapabilities, + client_capabilities: &AcpClientCapabilities, ) -> Result { tracing::info!("initializing stdio server client"); - let result: AcpInitializeResult = timeout( - Duration::from_secs(3), - self.request( - ACP_INITIALIZE_METHOD, - AcpInitializeParams { - protocol_version: 1, - client_capabilities: client_capabilities.clone(), - client_info: Some( - AcpImplementation::new("devo", env!("CARGO_PKG_VERSION")) - .with_title("Devo"), - ), - meta: None, - }, - ), - ) - .await - .context("timed out waiting for initialize response from server")??; + self.core + .set_client_capabilities(client_capabilities.clone()); + let result = timeout(STDIO_INITIALIZE_TIMEOUT, self.core.initialize()) + .await + .context("timed out waiting for initialize response from server")??; tracing::info!("stdio server client initialized"); - self.acp_agent_capabilities = Some(result.agent_capabilities.clone()); - let meta = result.meta.as_ref(); - Ok(InitializeResult { - server_name: result - .agent_info - .as_ref() - .map(|info| info.name.clone()) - .unwrap_or_else(|| "devo-server".to_string()), - server_version: result - .agent_info - .as_ref() - .map(|info| info.version.clone()) - .unwrap_or_else(|| env!("CARGO_PKG_VERSION").to_string()), - platform_family: meta - .and_then(|meta| meta.get("devo/platformFamily")) - .and_then(serde_json::Value::as_str) - .unwrap_or(std::env::consts::FAMILY) - .into(), - platform_os: meta - .and_then(|meta| meta.get("devo/platformOs")) - .and_then(serde_json::Value::as_str) - .unwrap_or(std::env::consts::OS) - .into(), - server_home: meta - .and_then(|meta| meta.get("devo/serverHome")) - .and_then(serde_json::Value::as_str) - .map(PathBuf::from) - .unwrap_or_default(), - }) + Ok(result) } pub async fn acp_terminal_output_snapshot( &self, terminal_id: &str, ) -> Result { - self.acp_terminals - .output(terminal_id) - .await - .map_err(anyhow::Error::msg) + self.core.acp_terminal_output_snapshot(terminal_id).await } pub async fn session_start( &mut self, params: SessionStartParams, ) -> Result { - let cwd = params.cwd.clone(); - let additional_directories = params.additional_directories.clone(); - let result: AcpNewSessionResult = self - .request( - ACP_SESSION_NEW_METHOD, - AcpNewSessionParams { - cwd, - additional_directories, - mcp_servers: Vec::new(), - meta: None, - }, - ) - .await?; - let session = result - .meta - .as_ref() - .and_then(|meta| meta.get(devo_protocol::DEVO_SESSION_META)) - .cloned() - .map(serde_json::from_value) - .transpose() - .context("decode session metadata from ACP session/new response")? - .unwrap_or_else(|| acp_session_metadata_from_start_params(¶ms, result.session_id)); - Ok(SessionStartResult { session }) + self.core.session_start(params).await } pub async fn session_resume( &mut self, params: SessionResumeParams, ) -> Result { - let sessions = self.session_list().await?; - let session = sessions - .into_iter() - .find(|session| session.session_id == params.session_id) - .with_context(|| { - format!( - "session {} not found for ACP session/resume", - params.session_id - ) - })?; - let result: AcpResumeSessionResult = self - .request( - ACP_SESSION_RESUME_METHOD, - AcpResumeSessionParams { - session_id: params.session_id, - cwd: session.cwd.clone(), - additional_directories: session.additional_directories.clone(), - mcp_servers: Vec::new(), - meta: None, - }, - ) - .await?; - Ok(result - .meta - .as_ref() - .and_then(|meta| meta.get(DEVO_SESSION_RESUME_META)) - .cloned() - .map(serde_json::from_value) - .transpose() - .context("decode session resume metadata from ACP session/resume response")? - .unwrap_or_else(|| SessionResumeResult { - session, - latest_turn: None, - loaded_item_count: 0, - history_items: Vec::new(), - pending_texts: Vec::new(), - })) + self.core.session_resume(params).await } pub async fn session_list(&mut self) -> Result> { - let Some(capabilities) = self.acp_agent_capabilities.as_ref() else { - bail!("ACP initialize must complete before session/list"); - }; - if capabilities.session_capabilities.list.is_none() { - bail!("ACP agent does not advertise sessionCapabilities.list"); - } - - let mut cursor = None; - let mut seen_cursors = HashSet::new(); - let mut sessions = Vec::new(); - loop { - let result: AcpListSessionsResult = self - .request( - ACP_SESSION_LIST_METHOD, - AcpListSessionsParams { - cwd: None, - cursor, - meta: None, - }, - ) - .await?; - for session_info in result.sessions { - let session = session_info - .meta - .as_ref() - .and_then(|meta| meta.get(devo_protocol::DEVO_SESSION_META)) - .cloned() - .map(serde_json::from_value) - .transpose() - .context("decode session metadata from ACP session/list response")? - .unwrap_or_else(|| acp_session_metadata_from_session_info(&session_info)); - sessions.push(session); - } - - let Some(next_cursor) = result.next_cursor else { - break; - }; - if !seen_cursors.insert(next_cursor.clone()) { - bail!("ACP session/list returned a repeated nextCursor"); - } - cursor = Some(next_cursor); - } - Ok(sessions) + self.core.session_list().await } pub async fn agent_list(&mut self, params: AgentListParams) -> Result { - self.request_devo("agent/list", params).await + self.core.agent_list(params).await } pub async fn agent_spawn(&mut self, params: SpawnAgentParams) -> Result { - self.request_devo("agent/spawn", params).await + self.core.agent_spawn(params).await } pub async fn agent_close(&mut self, params: CloseAgentParams) -> Result { - self.request_devo("agent/close", params).await + self.core.agent_close(params).await } pub async fn session_title_update( &mut self, params: SessionTitleUpdateParams, ) -> Result { - self.request_devo("session/title/update", params).await + self.core.session_title_update(params).await } pub async fn session_metadata_update( &mut self, params: SessionMetadataUpdateParams, ) -> Result { - self.request_devo("session/metadata/update", params).await + self.core.session_metadata_update(params).await } pub async fn session_permissions_update( &mut self, params: SessionPermissionsUpdateParams, ) -> Result { - self.request_devo("session/permissions/update", params) - .await + self.core.session_permissions_update(params).await } pub async fn session_compact( &mut self, params: SessionCompactParams, ) -> Result { - self.request_devo("session/compact", params).await + self.core.session_compact(params).await } pub async fn goal_create(&mut self, params: GoalCreateParams) -> Result { - self.request_devo("goal/create", params).await + self.core.goal_create(params).await } pub async fn goal_set(&mut self, params: GoalSetParams) -> Result { - self.request_devo("goal/set", params).await + self.core.goal_set(params).await } pub async fn goal_status(&mut self, params: GoalStatusParams) -> Result { - self.request_devo("goal/status", params).await + self.core.goal_status(params).await } pub async fn goal_pause(&mut self, params: GoalSetStatusParams) -> Result { - self.request_devo("goal/pause", params).await + self.core.goal_pause(params).await } pub async fn goal_resume( &mut self, params: GoalSetStatusParams, ) -> Result { - self.request_devo("goal/resume", params).await + self.core.goal_resume(params).await } pub async fn goal_complete( &mut self, params: GoalSetStatusParams, ) -> Result { - self.request_devo("goal/complete", params).await + self.core.goal_complete(params).await } pub async fn goal_clear(&mut self, params: GoalClearParams) -> Result { - self.request_devo("goal/clear", params).await + self.core.goal_clear(params).await } pub async fn session_fork(&mut self, params: SessionForkParams) -> Result { - self.request_devo("session/fork", params).await + self.core.session_fork(params).await } pub async fn session_rollback( &mut self, params: SessionRollbackParams, ) -> Result { - self.request_devo("session/rollback", params).await + self.core.session_rollback(params).await } pub async fn skills_list(&mut self, params: SkillListParams) -> Result { - self.request_devo("skills/list", params).await + self.core.skills_list(params).await } pub async fn skills_changed( &mut self, params: SkillChangedParams, ) -> Result { - self.request_devo("skills/changed", params).await + self.core.skills_changed(params).await } pub async fn skills_set_enabled( &mut self, params: SkillSetEnabledParams, ) -> Result { - self.request_devo("skills/set_enabled", params).await + self.core.skills_set_enabled(params).await } pub async fn model_catalog( &mut self, params: ModelCatalogParams, ) -> Result { - self.request_devo("model/catalog", params).await + self.core.model_catalog(params).await } pub async fn model_saved(&mut self, params: ModelSavedParams) -> Result { - self.request_devo("model/saved", params).await + self.core.model_saved(params).await } pub async fn provider_vendor_list( &mut self, params: ProviderVendorListParams, ) -> Result { - self.request_devo("provider/list", params).await + self.core.provider_vendor_list(params).await } pub async fn provider_vendor_upsert( &mut self, params: ProviderVendorUpsertParams, ) -> Result { - self.request_devo("provider/upsert", params).await + self.core.provider_vendor_upsert(params).await } pub async fn provider_validate( &mut self, params: ProviderValidateParams, ) -> Result { - self.request_devo("provider/validate", params).await + self.core.provider_validate(params).await } pub async fn command_exec(&mut self, params: CommandExecParams) -> Result { - self.request_devo("command/exec", params).await + self.core.command_exec(params).await } pub async fn command_exec_write( &mut self, params: CommandExecWriteParams, ) -> Result { - self.request_devo("command/exec/write", params).await + self.core.command_exec_write(params).await } pub async fn command_exec_resize( &mut self, params: CommandExecResizeParams, ) -> Result { - self.request_devo("command/exec/resize", params).await + self.core.command_exec_resize(params).await } pub async fn command_exec_terminate( &mut self, params: CommandExecTerminateParams, ) -> Result { - self.request_devo("command/exec/terminate", params).await + self.core.command_exec_terminate(params).await } pub async fn turn_start(&mut self, params: TurnStartParams) -> Result { - match self - .request_devo::<_, TurnStartResult>("turn/start", params.clone()) - .await - { - Ok(result) => Ok(result), - Err(error) if is_method_not_found_error(&error) => { - self.turn_start_acp_prompt_detached(params).await?; - Ok(TurnStartResult::Started { - turn_id: TurnId::new(), - status: TurnStatus::Running, - accepted_at: Utc::now(), - }) - } - Err(error) => Err(error), - } + self.core.turn_start(params).await } pub async fn turn_shell_command( &mut self, params: ShellCommandParams, ) -> Result { - self.request_devo("turn/shell_command", params).await + self.core.turn_shell_command(params).await } pub async fn turn_interrupt( &mut self, params: TurnInterruptParams, ) -> Result { - self.request_devo("turn/interrupt", params).await + self.core.turn_interrupt(params).await } pub async fn turn_steer(&mut self, params: TurnSteerParams) -> Result { - self.request_devo("turn/steer", params).await + self.core.turn_steer(params).await } pub async fn approval_respond(&mut self, params: ApprovalResponseParams) -> Result<()> { - if let Some((response, notification)) = - resolve_acp_permission_response(&self.acp_pending_permissions, ¶ms).await - { - write_acp_client_response(Arc::clone(&self.stdin), response) - .await - .context("write ACP permission response")?; - let _ = self.notifications_tx.send(notification); - return Ok(()); - } - anyhow::bail!("no pending ACP permission request exists for approval response") + self.core.approval_respond(params).await } pub async fn request_user_input_respond( &mut self, params: RequestUserInputRespondParams, ) -> Result<()> { - let _: serde_json::Value = self - .request_devo("request_user_input/respond", params) - .await?; - Ok(()) + self.core.request_user_input_respond(params).await } pub async fn reference_search_start( &mut self, params: ReferenceSearchStartParams, ) -> Result { - self.request_devo("search/start", params).await + self.core.reference_search_start(params).await } pub async fn reference_search_update( &mut self, params: ReferenceSearchUpdateParams, ) -> Result { - self.request_devo("search/update", params).await + self.core.reference_search_update(params).await } pub async fn reference_search_cancel( &mut self, params: ReferenceSearchCancelParams, ) -> Result { - self.request_devo("search/cancel", params).await + self.core.reference_search_cancel(params).await } pub async fn recv_notification(&mut self) -> Option { - self.notifications_rx.recv().await + self.core.recv_notification().await } pub async fn recv_event(&mut self) -> Result> { - let Some(notification) = self.recv_notification().await else { - return Ok(None); - }; - let ServerNotificationMessage { method, params } = notification; - let event = serde_json::from_value(params) - .with_context(|| format!("failed to decode server event for method {method}"))?; - Ok(Some((method, event))) + self.core.recv_event().await } pub async fn shutdown(mut self) -> Result<()> { tracing::info!("stdio server client shutdown requested"); - let _ = timeout( - SERVER_CHILD_STDIN_SHUTDOWN_TIMEOUT, - self.stdin.lock().await.shutdown(), - ) - .await; + self.core.shutdown().await; + match timeout(SERVER_CHILD_STDIN_SHUTDOWN_TIMEOUT, &mut self.writer_task).await { + Ok(Ok(Ok(()))) => {} + Ok(Ok(Err(error))) => { + tracing::debug!(%error, "stdio writer stopped with error during shutdown"); + } + Ok(Err(error)) => { + tracing::debug!(%error, "stdio writer task join failed during shutdown"); + } + Err(_) => { + self.writer_task.abort(); + } + } tracing::info!("stdio server stdin shutdown attempted"); if let Err(error) = self.child.start_kill() { tracing::debug!(%error, "failed to start stdio server child kill"); @@ -701,373 +405,45 @@ impl StdioServerClient { tracing::debug!("timed out waiting for stdio server child exit"); } } - self.acp_terminals.release_all().await; - Ok(()) - } - - async fn request(&mut self, method: &str, params: P) -> Result - where - P: serde::Serialize, - R: DeserializeOwned, - { - let request_id = self.next_request_id.fetch_add(1, Ordering::SeqCst); - tracing::debug!(request_id, method, "sending client request"); - let (response_tx, response_rx) = oneshot::channel(); - // The stdout reader owns response routing. Keep the sender in this map - // only while the request can still be completed by an incoming response. - self.pending.lock().await.insert(request_id, response_tx); - let write_result = self - .write_json(&AcpClientRequest::new( - serde_json::json!(request_id), - method, - params, - )) - .await; - if let Err(error) = write_result { - self.pending.lock().await.remove(&request_id); - return Err(error); - } - - let response = match timeout(Duration::from_secs(10), response_rx).await { - Ok(Ok(response)) => response, - Ok(Err(error)) => { - self.pending.lock().await.remove(&request_id); - return Err(error) - .with_context(|| format!("server dropped response for request {request_id}")); - } - Err(error) => { - self.pending.lock().await.remove(&request_id); - return Err(error).with_context(|| { - format!("timed out waiting for server response to request {request_id}") - }); - } - }; - tracing::debug!(request_id, method, "received client response"); - if response.get("error").is_some() { - bail_server_error(&response)?; - } - let success: AcpSuccessResponse = - serde_json::from_value(response).context("decode success response from server")?; - Ok(success.result) - } - - async fn request_devo(&mut self, method: &str, params: P) -> Result - where - P: serde::Serialize, - R: DeserializeOwned, - { - let method = devo_extension_method(method); - self.request(&method, params).await - } - - async fn turn_start_acp_prompt_detached(&mut self, params: TurnStartParams) -> Result<()> { - let session_id = params.session_id; - let prompt = params - .input - .into_iter() - .map(acp_content_block_from_input_item) - .collect(); - let request_id = self.next_request_id.fetch_add(1, Ordering::SeqCst); - tracing::debug!( - request_id, - method = ACP_SESSION_PROMPT_METHOD, - "sending ACP prompt request" - ); - let (response_tx, response_rx) = oneshot::channel(); - self.pending.lock().await.insert(request_id, response_tx); - let write_result = self - .write_json(&AcpClientRequest::new( - serde_json::json!(request_id), - ACP_SESSION_PROMPT_METHOD, - AcpPromptParams { - session_id, - prompt, - meta: None, - }, - )) - .await; - if let Err(error) = write_result { - self.pending.lock().await.remove(&request_id); - return Err(error); - } - - let _ = self.notifications_tx.send(ServerNotificationMessage { - method: ACP_PROMPT_STARTED_NOTIFICATION_METHOD.to_string(), - params: serde_json::json!({ "sessionId": session_id }), - }); - let notifications_tx = self.notifications_tx.clone(); - tokio::spawn(async move { - let params = match response_rx.await { - Ok(response) if response.get("error").is_some() => serde_json::json!({ - "sessionId": session_id, - "error": server_error_text(&response), - }), - Ok(response) => { - match serde_json::from_value::>(response) { - Ok(success) => serde_json::json!({ - "sessionId": session_id, - "stopReason": success.result.stop_reason, - }), - Err(error) => serde_json::json!({ - "sessionId": session_id, - "error": format!("decode ACP prompt response: {error}"), - }), - } - } - Err(error) => serde_json::json!({ - "sessionId": session_id, - "error": format!("server dropped ACP prompt response: {error}"), - }), - }; - let _ = notifications_tx.send(ServerNotificationMessage { - method: ACP_PROMPT_COMPLETED_NOTIFICATION_METHOD.to_string(), - params, - }); - }); - Ok(()) - } - - async fn write_json(&mut self, value: &T) -> Result<()> - where - T: serde::Serialize, - { - let mut line = serde_json::to_vec(value).context("serialize client payload")?; - line.push(b'\n'); - let mut stdin = self.stdin.lock().await; - stdin - .write_all(&line) - .await - .context("write client payload")?; - stdin.flush().await.context("flush client payload")?; + self.reader_task.abort(); + let _ = self.reader_task.await; Ok(()) } } -async fn run_stdout_reader( - mut lines: tokio::io::Lines, - pending: PendingResponses, +async fn run_stdin_writer( stdin: Arc>, - acp_pending_permissions: AcpPendingPermissions, - acp_terminals: AcpTerminalManager, - notifications_tx: mpsc::UnboundedSender, -) where - R: AsyncBufRead + Unpin, -{ - while let Ok(Some(line)) = lines.next_line().await { - match serde_json::from_str::(&line) { - Ok(message) => { - if let (Some(id), Some(method)) = ( - message.get("id").cloned(), - message.get("method").and_then(serde_json::Value::as_str), - ) { - let stdin = Arc::clone(&stdin); - let acp_pending_permissions = Arc::clone(&acp_pending_permissions); - let acp_terminals = acp_terminals.clone(); - let notifications_tx = notifications_tx.clone(); - let method = method.to_string(); - let params = message - .get("params") - .cloned() - .unwrap_or_else(|| serde_json::json!({})); - tokio::spawn(async move { - handle_acp_client_request( - stdin, - acp_pending_permissions, - acp_terminals, - notifications_tx, - id, - &method, - params, - ) - .await; - }); - } else if let Some(id) = message.get("id").and_then(serde_json::Value::as_u64) { - // Responses are consumed by the request future waiting on the - // matching oneshot. Notifications intentionally bypass that - // map so event consumers can drain them independently. - if let Some(tx) = pending.lock().await.remove(&id) { - let _ = tx.send(message); - } - } else if let Ok(notification) = - serde_json::from_value::>(message) - { - if notification.method == ACP_SESSION_UPDATE_METHOD - && let Ok(acp_notification) = serde_json::from_value::( - notification.params.clone(), - ) - && let Some((method, event)) = - original_event_from_acp_notification(&acp_notification) - { - let _ = notifications_tx.send(ServerNotificationMessage { - method, - params: serde_json::to_value(event) - .expect("serialize original ACP event"), - }); - continue; - } - if let Some(method) = devo_extension_inner_method(¬ification.method) - && serde_json::from_value::(notification.params.clone()) - .is_ok() - { - let _ = notifications_tx.send(ServerNotificationMessage { - method: method.to_string(), - params: notification.params, - }); - continue; - } - let event_seq = notification - .params - .get("context") - .and_then(|context| context.get("seq")) - .and_then(serde_json::Value::as_u64); - let item_id = notification_item_id(¬ification.params); - let assistant_delta = - notification_assistant_delta(¬ification.method, ¬ification.params); - let delta_len = assistant_delta.map(str::len); - let assistant_token_text = - assistant_delta.and_then(assistant_token_log_preview); - if let Some(assistant_token_text) = assistant_token_text.as_deref() { - tracing::debug!( - stream_elapsed_ms = stream_trace_elapsed_ms(), - method = %notification.method, - event_seq, - item_id = ?item_id, - delta_len = ?delta_len, - assistant_token_text, - "stdio client received server notification" - ); - } else { - tracing::debug!( - stream_elapsed_ms = stream_trace_elapsed_ms(), - method = %notification.method, - event_seq, - item_id = ?item_id, - delta_len = ?delta_len, - "stdio client received server notification" - ); - } - let _ = notifications_tx.send(ServerNotificationMessage { - method: notification.method, - params: notification.params, - }); - } + mut write_rx: tokio::sync::mpsc::UnboundedReceiver, +) -> Result<()> { + while let Some(message) = write_rx.recv().await { + match message { + ClientWriteMessage::Json(value) => { + write_ndjson_to_stdin(&stdin, &value) + .await + .context("write client payload")?; } - Err(_) => { - tracing::warn!(line = %line, "failed to parse JSON from server stdout"); + ClientWriteMessage::Close => { + let _ = timeout( + SERVER_CHILD_STDIN_SHUTDOWN_TIMEOUT, + stdin.lock().await.shutdown(), + ) + .await; + break; } } } - // Dropping the pending response senders wakes request futures instead of - // leaving them blocked until their timeout when the child closes stdout. - let abandoned_response_count = pending.lock().await.drain().count(); - if abandoned_response_count == 0 { - tracing::warn!("server stdout reader stopped"); - } else { - tracing::warn!( - abandoned_response_count, - "server stdout reader stopped with pending responses" - ); - } - acp_terminals.release_all().await; -} - -async fn handle_acp_client_request( - stdin: Arc>, - acp_pending_permissions: AcpPendingPermissions, - acp_terminals: AcpTerminalManager, - notifications_tx: mpsc::UnboundedSender, - id: serde_json::Value, - method: &str, - params: serde_json::Value, -) { - if method == "session/request_permission" { - let response = match handle_acp_request_permission( - id.clone(), - params, - acp_pending_permissions, - notifications_tx, - ) - .await - { - Ok(()) => return, - Err(message) => acp_client_error_response(id, -32603, message), - }; - if let Err(error) = write_acp_client_response(stdin, response).await { - tracing::warn!(%error, method, "failed to write ACP client response"); - } - return; - } - if matches!( - method, - ACP_FS_READ_TEXT_FILE_METHOD | ACP_FS_WRITE_TEXT_FILE_METHOD - ) { - let response = match handle_acp_fs_request(id.clone(), method, params).await { - Ok(response) => response, - Err(message) => acp_client_error_response(id, -32603, message), - }; - if let Err(error) = write_acp_client_response(stdin, response).await { - tracing::warn!(%error, method, "failed to write ACP client response"); - } - return; - } - if matches!( - method, - ACP_TERMINAL_CREATE_METHOD - | ACP_TERMINAL_OUTPUT_METHOD - | ACP_TERMINAL_WAIT_FOR_EXIT_METHOD - | ACP_TERMINAL_KILL_METHOD - | ACP_TERMINAL_RELEASE_METHOD - ) { - let response = match handle_acp_terminal_request( - id.clone(), - method, - params, - acp_terminals, - notifications_tx, - ) - .await - { - Ok(response) => response, - Err(message) => acp_client_error_response(id, -32603, message), - }; - if let Err(error) = write_acp_client_response(stdin, response).await { - tracing::warn!(%error, method, "failed to write ACP client response"); - } - return; - } - let response = acp_client_error_response(id, -32601, format!("unknown client method {method}")); - if let Err(error) = write_acp_client_response(stdin, response).await { - tracing::warn!(%error, method, "failed to write ACP client response"); - } -} - -fn acp_client_error_response( - id: serde_json::Value, - code: i64, - message: impl Into, -) -> serde_json::Value { - serde_json::json!({ - "jsonrpc": "2.0", - "id": id, - "error": { - "code": code, - "message": message.into() - } - }) + Ok(()) } -async fn write_acp_client_response( - stdin: Arc>, - value: serde_json::Value, +async fn write_ndjson_to_stdin( + stdin: &Arc>, + value: &serde_json::Value, ) -> Result<()> { - let mut line = serde_json::to_vec(&value).context("serialize ACP client response")?; + let mut line = serde_json::to_vec(value).context("serialize client payload")?; line.push(b'\n'); let mut stdin = stdin.lock().await; - stdin - .write_all(&line) - .await - .context("write ACP client response")?; - stdin.flush().await.context("flush ACP client response")?; + stdin.write_all(&line).await.context("write client payload")?; + stdin.flush().await.context("flush client payload")?; Ok(()) } @@ -1078,286 +454,84 @@ async fn run_stderr_reader(mut lines: tokio::io::Lines>) tracing::warn!(server_stderr = %trimmed, "server child stderr"); } } - tracing::warn!("server stderr reader stopped"); -} - -fn format_protocol_error_code(code: &ProtocolErrorCode) -> &'static str { - match code { - ProtocolErrorCode::NotInitialized => "not_initialized", - ProtocolErrorCode::InvalidParams => "invalid_params", - ProtocolErrorCode::SessionNotFound => "session_not_found", - ProtocolErrorCode::TurnNotFound => "turn_not_found", - ProtocolErrorCode::TurnAlreadyRunning => "turn_already_running", - ProtocolErrorCode::ApprovalNotFound => "approval_not_found", - ProtocolErrorCode::PolicyDenied => "policy_denied", - ProtocolErrorCode::ContextLimitExceeded => "context_limit_exceeded", - ProtocolErrorCode::NoActiveTurn => "no_active_turn", - ProtocolErrorCode::ExpectedTurnMismatch => "expected_turn_mismatch", - ProtocolErrorCode::ActiveTurnNotSteerable => "active_turn_not_steerable", - ProtocolErrorCode::EmptyInput => "empty_input", - ProtocolErrorCode::AlreadyResolved => "already_resolved", - ProtocolErrorCode::ParentSessionNotFound => "parent_session_not_found", - ProtocolErrorCode::ForkTurnNotFound => "fork_turn_not_found", - ProtocolErrorCode::ForkTurnNotStable => "fork_turn_not_stable", - ProtocolErrorCode::PermissionDenied => "permission_denied", - ProtocolErrorCode::WorkspaceUnavailable => "workspace_unavailable", - ProtocolErrorCode::InheritedSegmentWriteFailed => "inherited_segment_write_failed", - ProtocolErrorCode::ForkRetentionRequired => "fork_retention_required", - ProtocolErrorCode::InvalidConfirmToken => "invalid_confirm_token", - ProtocolErrorCode::UnsupportedDeletePolicy => "unsupported_delete_policy", - ProtocolErrorCode::InheritedSegmentMaterializationFailed => { - "inherited_segment_materialization_failed" - } - ProtocolErrorCode::ExpectedTargetMessageMismatch => "expected_target_message_mismatch", - ProtocolErrorCode::OlderMessageRequiresFork => "older_message_requires_fork", - ProtocolErrorCode::ActiveTurnEditRejected => "active_turn_edit_rejected", - ProtocolErrorCode::InvalidContentParts => "invalid_content_parts", - ProtocolErrorCode::InvalidMentions => "invalid_mentions", - ProtocolErrorCode::WorkspaceRestoreFailedToStart => "workspace_restore_failed_to_start", - ProtocolErrorCode::InternalError => "internal_error", - } -} - -fn bail_server_error(response: &serde_json::Value) -> Result<()> { - anyhow::bail!("{}", server_error_text(response)) -} - -fn is_method_not_found_error(error: &anyhow::Error) -> bool { - error.to_string().starts_with("server -32601:") -} - -fn server_error_text(response: &serde_json::Value) -> String { - if let Ok(error) = serde_json::from_value::(response.clone()) { - let data = if error.error.data.is_null() { - String::new() - } else { - format!(" data={}", error.error.data) - }; - return format!( - "server {}: {}{}", - format_protocol_error_code(&error.error.code), - error.error.message, - data - ); - } - format!( - "server {}: {}", - server_error_code(response), - server_error_message(response) - ) -} - -fn server_error_code(response: &serde_json::Value) -> String { - response - .get("error") - .and_then(|error| error.get("code")) - .map(serde_json::Value::to_string) - .unwrap_or_else(|| "unknown".to_string()) } -fn server_error_message(response: &serde_json::Value) -> &str { - response - .get("error") - .and_then(|error| error.get("message")) - .and_then(serde_json::Value::as_str) - .unwrap_or("unknown server error") -} +#[cfg(test)] +mod tests { + use std::sync::Arc; + use tokio::io::AsyncBufRead; + use tokio::io::BufReader; + use tokio::sync::Mutex; + use tokio::sync::oneshot; + use tokio::time::timeout; + use tokio::time::Duration; + use super::*; + use chrono::Utc; + use pretty_assertions::assert_eq; + use crate::client_core::ClientWriteMessage; + use crate::client_core::ClientWriter; + use crate::client_core::PendingResponses; + use crate::client_core::ServerClientCore; -fn acp_content_block_from_input_item(input: devo_protocol::InputItem) -> AcpContentBlock { - match input { - devo_protocol::InputItem::Text { text } => AcpContentBlock::text(text), - devo_protocol::InputItem::Skill { name, path } => AcpContentBlock::Text { - annotations: None, - text: format!("Skill {name}: {}", path.display()), - meta: None, - }, - devo_protocol::InputItem::LocalImage { path } => AcpContentBlock::Text { - annotations: None, - text: format!("Image: {}", path.display()), - meta: None, - }, - devo_protocol::InputItem::Mention { path, name } => AcpContentBlock::ResourceLink { - annotations: None, - uri: file_uri_from_path(&path), - name: name.unwrap_or_else(|| path.clone()), - title: None, - description: None, - mime_type: None, - size: None, + fn default_test_client_capabilities() -> devo_protocol::AcpClientCapabilities { + devo_protocol::AcpClientCapabilities { + fs: devo_protocol::AcpFileSystemCapabilities { + read_text_file: true, + write_text_file: true, + meta: None, + }, + terminal: true, meta: None, - }, - } -} - -fn file_uri_from_path(path: &str) -> String { - let normalized = path.replace('\\', "/"); - if normalized.starts_with('/') { - format!("file://{normalized}") - } else { - format!("file:///{normalized}") - } -} - -fn stream_trace_elapsed_ms() -> u128 { - static STREAM_TRACE_START: OnceLock = OnceLock::new(); - STREAM_TRACE_START - .get_or_init(Instant::now) - .elapsed() - .as_millis() -} - -fn notification_item_id(params: &serde_json::Value) -> Option<&str> { - params - .get("context") - .and_then(|context| context.get("item_id")) - .and_then(serde_json::Value::as_str) -} - -fn notification_assistant_delta<'a>( - method: &str, - params: &'a serde_json::Value, -) -> Option<&'a str> { - (method == "item/agentMessage/delta") - .then(|| params.get("delta")?.as_str()) - .flatten() -} - -fn assistant_token_log_preview(text: &str) -> Option { - assistant_token_logging_enabled().then(|| { - let max_chars = assistant_token_log_max_chars(); - format_assistant_token_log_preview(text, max_chars) - }) -} - -fn assistant_token_logging_enabled() -> bool { - static ENABLED: OnceLock = OnceLock::new(); - *ENABLED.get_or_init(|| { - std::env::var("DEVO_LOG_ASSISTANT_TOKEN_TEXT") - .ok() - .is_some_and(|value| { - matches!( - value.as_str(), - "1" | "true" | "TRUE" | "yes" | "YES" | "on" | "ON" - ) - }) - }) -} - -fn assistant_token_log_max_chars() -> usize { - static ASSISTANT_TOKEN_LOG_MAX_CHARS: OnceLock = OnceLock::new(); - *ASSISTANT_TOKEN_LOG_MAX_CHARS.get_or_init(|| { - std::env::var("DEVO_ASSISTANT_TOKEN_LOG_MAX_CHARS") - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|value| *value > 0) - .unwrap_or(512) - }) -} - -fn format_assistant_token_log_preview(text: &str, max_chars: usize) -> String { - let max_chars = max_chars.max(1); - let mut preview = String::with_capacity(text.len().min(max_chars)); - let mut chars = text.chars(); - for ch in chars.by_ref().take(max_chars) { - preview.extend(ch.escape_default()); - } - if chars.next().is_some() { - preview.push_str("..."); + } } - preview -} -fn acp_session_metadata_from_start_params( - params: &SessionStartParams, - session_id: devo_protocol::SessionId, -) -> SessionMetadata { - let now = Utc::now(); - SessionMetadata { - session_id, - cwd: params.cwd.clone(), - additional_directories: params.additional_directories.clone(), - created_at: now, - updated_at: now, - last_activity_at: now, - title: params.title.clone(), - title_state: acp_title_state(¶ms.title), - parent_session_id: None, - agent_path: None, - agent_nickname: None, - agent_role: None, - ephemeral: params.ephemeral, - model: params.model.clone(), - model_binding_id: params.model_binding_id.clone(), - reasoning_effort_selection: None, - reasoning_effort: None, - total_input_tokens: 0, - total_output_tokens: 0, - total_tokens: 0, - total_cache_creation_tokens: 0, - total_cache_read_tokens: 0, - prompt_token_estimate: 0, - last_query_total_tokens: 0, - status: SessionRuntimeStatus::Idle, + fn test_agent_capabilities_with_session_list() -> AcpAgentCapabilities { + AcpAgentCapabilities { + session_capabilities: devo_protocol::AcpSessionCapabilities { + list: Some(devo_protocol::AcpSessionListCapabilities::default()), + ..Default::default() + }, + ..Default::default() + } } -} -fn acp_session_metadata_from_session_info(session_info: &AcpSessionInfo) -> SessionMetadata { - let updated_at = session_info - .updated_at - .as_deref() - .and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok()) - .map(|value| value.with_timezone(&Utc)) - .unwrap_or_else(Utc::now); - SessionMetadata { - session_id: session_info.session_id, - cwd: session_info.cwd.clone(), - additional_directories: session_info.additional_directories.clone(), - created_at: updated_at, - updated_at, - last_activity_at: updated_at, - title: session_info.title.clone(), - title_state: acp_title_state(&session_info.title), - parent_session_id: None, - agent_path: None, - agent_nickname: None, - agent_role: None, - ephemeral: false, - model: None, - model_binding_id: None, - reasoning_effort_selection: None, - reasoning_effort: None, - total_input_tokens: 0, - total_output_tokens: 0, - total_tokens: 0, - total_cache_creation_tokens: 0, - total_cache_read_tokens: 0, - prompt_token_estimate: 0, - last_query_total_tokens: 0, - status: SessionRuntimeStatus::Idle, + async fn spawn_test_stdio_client( + child: Child, + stdin: ChildStdin, + client_capabilities: devo_protocol::AcpClientCapabilities, + ) -> (StdioServerClient, PendingResponses) { + let stdin = Arc::new(Mutex::new(stdin)); + let (client_writer, mut write_rx) = ClientWriter::channel(); + let core = ServerClientCore::new(client_writer, client_capabilities); + let pending = core.pending_responses(); + tokio::spawn(async move { + while let Some(message) = write_rx.recv().await { + match message { + ClientWriteMessage::Json(value) => { + if write_ndjson_to_stdin(&stdin, &value).await.is_err() { + break; + } + } + ClientWriteMessage::Close => { + let _ = stdin.lock().await.shutdown().await; + break; + } + } + } + }); + let client = StdioServerClient { + child, + core, + reader_task: tokio::spawn(async {}), + writer_task: tokio::spawn(async { Ok(()) }), + }; + (client, pending) } -} -fn acp_title_state(title: &Option) -> SessionTitleState { - if title.is_some() { - SessionTitleState::Provisional - } else { - SessionTitleState::Unset - } -} - -#[cfg(test)] -mod tests { - use super::*; - use pretty_assertions::assert_eq; #[tokio::test] async fn initialize_uses_configured_client_capabilities() { let (child, stdin, stdout) = request_capture_child_for_turn_start_test().await; - let stdin = Arc::new(Mutex::new(stdin)); - let pending = Arc::new(Mutex::new(HashMap::new())); - let acp_pending_permissions = Arc::new(Mutex::new(HashMap::new())); - let acp_terminals = AcpTerminalManager::new(); - let (notifications_tx, notifications_rx) = mpsc::unbounded_channel(); let client_capabilities = devo_protocol::AcpClientCapabilities { fs: devo_protocol::AcpFileSystemCapabilities { read_text_file: true, @@ -1367,17 +541,7 @@ mod tests { terminal: true, meta: None, }; - let mut client = StdioServerClient { - child, - stdin, - pending: Arc::clone(&pending), - acp_pending_permissions, - acp_terminals, - acp_agent_capabilities: None, - next_request_id: AtomicU64::new(1), - notifications_rx, - notifications_tx, - }; + let (mut client, pending) = spawn_test_stdio_client(child, stdin, client_capabilities.clone()).await; let mut stdout_lines = BufReader::new(stdout).lines(); let expected_capabilities = serde_json::to_value(&client_capabilities).expect("serialize client capabilities"); @@ -1417,22 +581,7 @@ mod tests { #[tokio::test] async fn session_start_accepts_standard_acp_response_without_devo_metadata() { let (child, stdin, stdout) = request_capture_child_for_turn_start_test().await; - let stdin = Arc::new(Mutex::new(stdin)); - let pending = Arc::new(Mutex::new(HashMap::new())); - let acp_pending_permissions = Arc::new(Mutex::new(HashMap::new())); - let acp_terminals = AcpTerminalManager::new(); - let (notifications_tx, notifications_rx) = mpsc::unbounded_channel(); - let mut client = StdioServerClient { - child, - stdin, - pending: Arc::clone(&pending), - acp_pending_permissions, - acp_terminals, - acp_agent_capabilities: None, - next_request_id: AtomicU64::new(1), - notifications_rx, - notifications_tx, - }; + let (mut client, pending) = spawn_test_stdio_client(child, stdin, default_test_client_capabilities()).await; let cwd = std::env::current_dir().expect("current dir"); let additional_directory = cwd.join("shared"); let session_id = devo_protocol::SessionId::new(); @@ -1507,28 +656,10 @@ mod tests { #[tokio::test] async fn session_list_accepts_standard_acp_sessions_without_devo_metadata() { let (child, stdin, stdout) = request_capture_child_for_turn_start_test().await; - let stdin = Arc::new(Mutex::new(stdin)); - let pending = Arc::new(Mutex::new(HashMap::new())); - let acp_pending_permissions = Arc::new(Mutex::new(HashMap::new())); - let acp_terminals = AcpTerminalManager::new(); - let (notifications_tx, notifications_rx) = mpsc::unbounded_channel(); - let mut client = StdioServerClient { - child, - stdin, - pending: Arc::clone(&pending), - acp_pending_permissions, - acp_terminals, - acp_agent_capabilities: Some(AcpAgentCapabilities { - session_capabilities: devo_protocol::AcpSessionCapabilities { - list: Some(devo_protocol::AcpSessionListCapabilities::default()), - ..Default::default() - }, - ..Default::default() - }), - next_request_id: AtomicU64::new(1), - notifications_rx, - notifications_tx, - }; + let (mut client, pending) = spawn_test_stdio_client(child, stdin, default_test_client_capabilities()).await; + client + .core + .set_agent_capabilities_for_test(test_agent_capabilities_with_session_list()); let cwd = std::env::current_dir().expect("current dir"); let additional_directory = cwd.join("shared"); let session_id = devo_protocol::SessionId::new(); @@ -1614,28 +745,10 @@ mod tests { #[tokio::test] async fn session_resume_accepts_standard_acp_response_without_devo_metadata() { let (child, stdin, stdout) = request_capture_child_for_turn_start_test().await; - let stdin = Arc::new(Mutex::new(stdin)); - let pending = Arc::new(Mutex::new(HashMap::new())); - let acp_pending_permissions = Arc::new(Mutex::new(HashMap::new())); - let acp_terminals = AcpTerminalManager::new(); - let (notifications_tx, notifications_rx) = mpsc::unbounded_channel(); - let mut client = StdioServerClient { - child, - stdin, - pending: Arc::clone(&pending), - acp_pending_permissions, - acp_terminals, - acp_agent_capabilities: Some(AcpAgentCapabilities { - session_capabilities: devo_protocol::AcpSessionCapabilities { - list: Some(devo_protocol::AcpSessionListCapabilities::default()), - ..Default::default() - }, - ..Default::default() - }), - next_request_id: AtomicU64::new(1), - notifications_rx, - notifications_tx, - }; + let (mut client, pending) = spawn_test_stdio_client(child, stdin, default_test_client_capabilities()).await; + client + .core + .set_agent_capabilities_for_test(test_agent_capabilities_with_session_list()); let cwd = std::env::current_dir().expect("current dir"); let additional_directory = cwd.join("shared"); let session_id = devo_protocol::SessionId::new(); @@ -1749,22 +862,7 @@ mod tests { #[tokio::test] async fn turn_start_sends_devo_extension_with_full_params() { let (child, stdin, stdout) = request_capture_child_for_turn_start_test().await; - let stdin = Arc::new(Mutex::new(stdin)); - let pending = Arc::new(Mutex::new(HashMap::new())); - let acp_pending_permissions = Arc::new(Mutex::new(HashMap::new())); - let acp_terminals = AcpTerminalManager::new(); - let (notifications_tx, notifications_rx) = mpsc::unbounded_channel(); - let mut client = StdioServerClient { - child, - stdin, - pending: Arc::clone(&pending), - acp_pending_permissions, - acp_terminals, - acp_agent_capabilities: None, - next_request_id: AtomicU64::new(1), - notifications_rx, - notifications_tx, - }; + let (mut client, pending) = spawn_test_stdio_client(child, stdin, default_test_client_capabilities()).await; let params = TurnStartParams { session_id: devo_protocol::SessionId::new(), input: vec![devo_protocol::InputItem::Text { @@ -1831,22 +929,7 @@ mod tests { #[tokio::test] async fn turn_start_falls_back_to_acp_prompt_with_lifecycle_notifications() { let (child, stdin, stdout) = request_capture_child_for_turn_start_test().await; - let stdin = Arc::new(Mutex::new(stdin)); - let pending = Arc::new(Mutex::new(HashMap::new())); - let acp_pending_permissions = Arc::new(Mutex::new(HashMap::new())); - let acp_terminals = AcpTerminalManager::new(); - let (notifications_tx, notifications_rx) = mpsc::unbounded_channel(); - let mut client = StdioServerClient { - child, - stdin, - pending: Arc::clone(&pending), - acp_pending_permissions, - acp_terminals, - acp_agent_capabilities: None, - next_request_id: AtomicU64::new(1), - notifications_rx, - notifications_tx, - }; + let (mut client, pending) = spawn_test_stdio_client(child, stdin, default_test_client_capabilities()).await; let session_id = devo_protocol::SessionId::new(); let params = TurnStartParams { session_id, @@ -1949,24 +1032,15 @@ mod tests { #[tokio::test] async fn stdout_reader_drops_pending_responses_when_stdout_closes() { - let pending = Arc::new(Mutex::new(HashMap::new())); let (response_tx, response_rx) = oneshot::channel(); let request_id = 7; + let (client_writer, _) = ClientWriter::channel(); + let core = ServerClientCore::new(client_writer, default_test_client_capabilities()); + let pending = core.pending_responses(); pending.lock().await.insert(request_id, response_tx); - let (notifications_tx, _notifications_rx) = mpsc::unbounded_channel(); - let (mut child, stdin) = child_stdin_for_stdout_reader_test().await; - let acp_pending_permissions = Arc::new(Mutex::new(HashMap::new())); - let acp_terminals = AcpTerminalManager::new(); - - run_stdout_reader( - BufReader::new(tokio::io::empty()).lines(), - Arc::clone(&pending), - stdin, - acp_pending_permissions, - acp_terminals, - notifications_tx, - ) - .await; + let (mut child, _stdin) = child_stdin_for_stdout_reader_test().await; + + core.reader_state().finish_reader("stdio-test").await; assert!(response_rx.await.is_err()); assert_eq!(pending.lock().await.len(), 0); diff --git a/crates/client/src/websocket.rs b/crates/client/src/websocket.rs index a718a2ca..eedd8e3c 100644 --- a/crates/client/src/websocket.rs +++ b/crates/client/src/websocket.rs @@ -19,7 +19,7 @@ use tokio_tungstenite::tungstenite::Message; use crate::client_core::ClientWriteMessage; use crate::client_core::ClientWriter; use crate::client_core::ServerClientCore; -use crate::stdio::ServerNotificationMessage; +use crate::client_core::ServerNotificationMessage; const WEBSOCKET_SHUTDOWN_TIMEOUT: Duration = Duration::from_millis(500); @@ -128,7 +128,7 @@ impl WebSocketServerClient { } pub async fn agent_list(&mut self, params: AgentListParams) -> Result { - self.core.request_devo("agent/list", params).await + self.core.agent_list(params).await } pub async fn agent_spawn(&mut self, params: SpawnAgentParams) -> Result { diff --git a/crates/server/src/runtime/handlers/acp/prompt.rs b/crates/server/src/runtime/handlers/acp/prompt.rs index cdbb283e..67ab73bc 100644 --- a/crates/server/src/runtime/handlers/acp/prompt.rs +++ b/crates/server/src/runtime/handlers/acp/prompt.rs @@ -1,13 +1,27 @@ +//! ACP `session/prompt` and `session/cancel` handlers. +//! +//! Portable ACP clients submit turns through blocking `session/prompt`: the JSON-RPC +//! response is deferred until the turn reaches a terminal state, while progress is +//! streamed on `session/update`. Devo's TUI uses `_devo/turn/start` instead; this +//! module serves external ACP integrations. + use super::super::acp_slash_commands::AcpSlashCommandPromptResult; use super::*; impl ServerRuntime { + /// Handles ACP `session/prompt` for non-Devo clients. + /// + /// Returns `None` when the prompt turn was accepted and the JSON-RPC response + /// will be sent asynchronously after the turn finishes (ACP blocking semantics). + /// Streaming progress is delivered through `session/update` on the subscribed + /// connection while the background task waits for turn completion. pub(crate) async fn handle_acp_session_prompt( self: &Arc, connection_id: u64, request_id: serde_json::Value, params: serde_json::Value, ) -> Option { + // --- Validate request --------------------------------------------------- let params: AcpPromptParams = match serde_json::from_value(params) { Ok(params) => params, Err(error) => { @@ -18,6 +32,7 @@ impl ServerRuntime { )); } }; + // ACP prompt turns are exclusive per session; queueing is not supported here. if self.session_has_active_turn(params.session_id).await { return Some(acp_error_response( request_id, @@ -25,8 +40,10 @@ impl ServerRuntime { "session already has an active prompt turn", )); } + // Ensure this connection receives `session/update` for the prompt turn. self.subscribe_connection_to_session(connection_id, params.session_id, None) .await; + // --- Slash commands embedded in the prompt ---------------------------- match self .handle_acp_slash_command_prompt( connection_id, @@ -42,6 +59,7 @@ impl ServerRuntime { } AcpSlashCommandPromptResult::Pending => return None, } + // --- Start a regular turn from ACP prompt content --------------------- let session_id = params.session_id; let input = match input_items_from_acp_prompt(params.prompt) { Ok(input) => input, @@ -53,6 +71,8 @@ impl ServerRuntime { )); } }; + // Reuse the Devo turn engine with ACP-default params (no model/sandbox + // extensions on the prompt RPC). Reject if another turn is already active. let legacy_response = self .handle_turn_start_with_queue_policy( Some(connection_id), @@ -84,6 +104,9 @@ impl ServerRuntime { "session/prompt cannot queue behind an active turn", )); }; + // --- Defer JSON-RPC response until the turn completes ----------------- + // The client keeps reading `session/update` notifications while this + // task blocks on terminal turn status, then receives `AcpPromptResult`. let runtime = Arc::clone(self); tokio::spawn(async move { let stop_reason = runtime @@ -105,6 +128,7 @@ impl ServerRuntime { None } + /// Handles ACP `session/cancel` by interrupting the active turn on the session. pub(crate) async fn handle_acp_session_cancel(self: &Arc, params: serde_json::Value) { let params: AcpCancelParams = match serde_json::from_value(params) { Ok(params) => params, @@ -149,6 +173,7 @@ impl ServerRuntime { self.runtime_active_turn_id(session_id).await.is_some() } + /// Waits for a prompt turn to reach a terminal state and maps it to ACP `stopReason`. pub(crate) async fn wait_for_acp_prompt_stop_reason( &self, session_id: SessionId, @@ -167,6 +192,8 @@ impl ServerRuntime { acp_stop_reason_from_terminal_turn(status) } } + +/// Maps internal turn completion metadata to the ACP `stopReason` enum. fn acp_stop_reason_from_terminal_turn(snapshot: TerminalTurnSnapshot) -> AcpStopReason { match snapshot.status { TurnStatus::Completed => match snapshot.stop_reason { diff --git a/crates/server/src/transport.rs b/crates/server/src/transport.rs index 84b7b852..5dbe38f9 100644 --- a/crates/server/src/transport.rs +++ b/crates/server/src/transport.rs @@ -750,7 +750,7 @@ fn validate_incoming_client_message( } /// Returns `true` when the payload is a client response to a server-initiated -/// JSON-RPC request. +/// JSON-RPC request (see the call site in [`accept_incoming_client_message`]). fn is_client_response_message(value: &serde_json::Value) -> bool { let Some(object) = value.as_object() else { return false; @@ -802,12 +802,31 @@ async fn accept_incoming_client_message( }); return; } + // The server may initiate JSON-RPC requests to the client (ACP client-side + // tools such as fs/read, fs/write, permission prompts). Those replies arrive + // as client responses (id + result/error, no method) and must be matched to + // the pending server request instead of entering the normal inbound handler. if is_client_response_message(&value) { tokio::spawn(async move { runtime.resolve_client_response(connection_id, value).await; }); return; } + // Bound how many client requests/notifications may run concurrently on this + // connection (see INBOUND_CONCURRENCY_LIMIT). The transport read loop awaits a + // permit before spawning the next handler, so inbound work applies backpressure + // instead of unbounded task growth. Client responses above skip this permit + // because a handler may hold one while blocked on a server-initiated client call + // (e.g. ACP fs/read); requiring a permit for the reply would deadlock. + // + // Concurrency risks to keep in mind: + // - Handlers for the same connection run in parallel; ordering is not preserved + // beyond what the runtime/session actors enforce. + // - If every permit is held by handlers waiting on the client, the read loop + // blocks here until one finishes. Responses already bypass the permit, but a + // pipelined client request queued ahead of those responses in the socket + // buffer can delay reading them (head-of-line blocking). + // - Semaphore close drops the message silently (no JSON-RPC error). let Ok(permit) = inbound_semaphore.acquire_owned().await else { return; }; diff --git a/crates/server/tests/deep_research_e2e.rs b/crates/server/tests/deep_research_e2e.rs index 18c04ef4..899bd570 100644 --- a/crates/server/tests/deep_research_e2e.rs +++ b/crates/server/tests/deep_research_e2e.rs @@ -2273,21 +2273,6 @@ fn child_turn_session_id( (session_id != parent_session_id.to_string()).then(|| session_id.to_string()) } -fn unique_child_turn_sessions( - events: &[serde_json::Value], - parent_session_id: devo_core::SessionId, -) -> Vec { - let mut sessions = Vec::new(); - for event in events { - let Some(session_id) = child_turn_session_id(event, parent_session_id) else { - continue; - }; - if !sessions.contains(&session_id) { - sessions.push(session_id); - } - } - sessions -} fn legacy_event_from_acp_notification(value: serde_json::Value) -> serde_json::Value { if value.get("method") != Some(&serde_json::json!("session/update")) { return value; diff --git a/crates/tui/src/worker.rs b/crates/tui/src/worker.rs index b50fe5c1..a9c94e75 100644 --- a/crates/tui/src/worker.rs +++ b/crates/tui/src/worker.rs @@ -21,7 +21,6 @@ use devo_core::SessionId; use devo_core::TurnId; use devo_core::TurnStatus; use devo_protocol::ACP_SESSION_UPDATE_METHOD; -use devo_protocol::AcpStopReason; use devo_protocol::AgentListParams; use devo_protocol::AgentToolPolicy; use devo_protocol::CloseAgentParams; @@ -45,8 +44,6 @@ use devo_protocol::SessionHistoryMetadata; use devo_protocol::SessionPlanStepStatus; use devo_protocol::SpawnAgentParams; use devo_protocol::ThreadGoalStatus; -use devo_server::ACP_PROMPT_COMPLETED_NOTIFICATION_METHOD; -use devo_server::ACP_PROMPT_STARTED_NOTIFICATION_METHOD; use devo_server::ACP_TERMINAL_OUTPUT_NOTIFICATION_METHOD; use devo_server::ApprovalDecisionPayload; use devo_server::ApprovalRequestPayload; @@ -832,7 +829,6 @@ async fn run_worker_inner( let mut reasoning_effort_selection = config.reasoning_effort_selection; let mut permission_preset = config.permission_preset; let mut active_turn_id: Option = None; - let mut synthetic_acp_turn_id: Option = None; let mut turn_count = 0usize; let mut total_input_tokens = 0usize; let mut total_output_tokens = 0usize; @@ -941,6 +937,10 @@ async fn run_worker_inner( event_tx, ) .await?; + + // Start the turn via `_devo/turn/start`. The bundled server implements + // this extension; streaming and completion arrive as server + // notifications (`turn/started`, item deltas, `turn/completed`, etc.). let start_result = client.turn_start(TurnStartParams { session_id: active_session_id, input, @@ -2274,43 +2274,6 @@ async fn run_worker_inner( Some(notification) => { let method = notification.method; let params = notification.params; - if method == ACP_PROMPT_STARTED_NOTIFICATION_METHOD { - if active_turn_id.is_none() - && let Some((turn_id, event)) = acp_prompt_started_event( - ¶ms, - session_id, - &model, - &model_binding_id, - &reasoning_effort_selection, - ) - { - active_turn_id = Some(turn_id); - synthetic_acp_turn_id = Some(turn_id); - saw_usage_update_for_turn = false; - let _ = event_tx.send(event); - } - continue; - } - if method == ACP_PROMPT_COMPLETED_NOTIFICATION_METHOD { - if synthetic_acp_turn_id.is_some() - && let Some(event) = acp_prompt_completed_event( - ¶ms, - session_id, - &mut active_turn_id, - &mut turn_count, - total_input_tokens, - total_output_tokens, - total_tokens, - total_cache_read_tokens, - last_query_total_tokens, - last_query_input_tokens, - ) - { - synthetic_acp_turn_id = None; - let _ = event_tx.send(event); - } - continue; - } if method == ACP_TERMINAL_OUTPUT_NOTIFICATION_METHOD { if let Some(terminal_id) = params.get("terminalId").and_then(serde_json::Value::as_str) @@ -4216,102 +4179,6 @@ fn proposed_plan_text(payload: &serde_json::Value) -> String { .to_string() } -#[derive(serde::Deserialize)] -#[serde(rename_all = "camelCase")] -struct AcpPromptLifecycleNotification { - session_id: SessionId, - #[serde(default)] - stop_reason: Option, - #[serde(default)] - error: Option, -} - -fn acp_prompt_started_event( - params: &serde_json::Value, - active_session_id: Option, - model: &str, - model_binding_id: &Option, - reasoning_effort_selection: &Option, -) -> Option<(TurnId, WorkerEvent)> { - acp_prompt_lifecycle_notification(params, active_session_id)?; - let turn_id = TurnId::new(); - Some(( - turn_id, - WorkerEvent::TurnStarted { - model: model.to_string(), - model_binding_id: model_binding_id.clone(), - reasoning_effort_selection: reasoning_effort_selection.clone(), - reasoning_effort: None, - turn_id, - }, - )) -} - -#[expect(clippy::too_many_arguments)] -fn acp_prompt_completed_event( - params: &serde_json::Value, - active_session_id: Option, - active_turn_id: &mut Option, - turn_count: &mut usize, - total_input_tokens: usize, - total_output_tokens: usize, - total_tokens: usize, - total_cache_read_tokens: usize, - last_query_total_tokens: usize, - last_query_input_tokens: usize, -) -> Option { - let notification = acp_prompt_lifecycle_notification(params, active_session_id)?; - *active_turn_id = None; - if let Some(error) = notification.error { - return Some(WorkerEvent::TurnFailed { - message: error, - turn_count: *turn_count, - total_input_tokens, - total_output_tokens, - total_tokens, - total_cache_read_tokens, - prompt_token_estimate: total_input_tokens, - last_query_input_tokens, - }); - } - *turn_count += 1; - Some(WorkerEvent::TurnFinished { - stop_reason: acp_stop_reason_text( - notification.stop_reason.unwrap_or(AcpStopReason::EndTurn), - ), - turn_count: *turn_count, - total_input_tokens, - total_output_tokens, - total_tokens, - total_cache_read_tokens, - last_query_total_tokens, - last_query_input_tokens, - prompt_token_estimate: total_input_tokens, - }) -} - -fn acp_prompt_lifecycle_notification( - params: &serde_json::Value, - active_session_id: Option, -) -> Option { - let Ok(notification) = serde_json::from_value::(params.clone()) - else { - return None; - }; - (Some(notification.session_id) == active_session_id).then_some(notification) -} - -fn acp_stop_reason_text(stop_reason: AcpStopReason) -> String { - match stop_reason { - AcpStopReason::EndTurn => "Completed", - AcpStopReason::MaxTokens => "MaxTokens", - AcpStopReason::MaxTurnRequests => "MaxTurnRequests", - AcpStopReason::Refusal => "Refusal", - AcpStopReason::Cancelled => "Interrupted", - } - .to_string() -} - fn plan_event_from_tool_result(payload: &ToolResultPayload) -> Option { let tool_name = payload.tool_name.as_deref()?; match tool_name { @@ -4513,8 +4380,6 @@ mod tests { use super::QueryWorkerHandle; use super::ShellCommandExecStart; - use super::acp_prompt_completed_event; - use super::acp_prompt_started_event; use super::acp_terminal_output_event; use super::acp_terminal_snapshot_delta; use super::handle_completed_item; @@ -5149,111 +5014,6 @@ mod tests { ); } - #[test] - fn acp_prompt_started_emits_turn_started() { - let session_id = SessionId::new(); - let (turn_id, event) = acp_prompt_started_event( - &serde_json::json!({ "sessionId": session_id }), - Some(session_id), - "test-model", - &Some("binding".to_string()), - &Some("high".to_string()), - ) - .expect("ACP prompt started event"); - - assert_eq!( - event, - WorkerEvent::TurnStarted { - model: "test-model".to_string(), - model_binding_id: Some("binding".to_string()), - reasoning_effort_selection: Some("high".to_string()), - reasoning_effort: None, - turn_id, - } - ); - } - - #[test] - fn acp_prompt_completed_emits_turn_finished() { - let session_id = SessionId::new(); - let mut active_turn_id = Some(TurnId::new()); - let mut turn_count = 2; - - let event = acp_prompt_completed_event( - &serde_json::json!({ - "sessionId": session_id, - "stopReason": "end_turn" - }), - Some(session_id), - &mut active_turn_id, - &mut turn_count, - 11, - 13, - 24, - 17, - 19, - 23, - ) - .expect("ACP prompt completed event"); - - assert_eq!(active_turn_id, None); - assert_eq!(turn_count, 3); - assert_eq!( - event, - WorkerEvent::TurnFinished { - stop_reason: "Completed".to_string(), - turn_count: 3, - total_input_tokens: 11, - total_output_tokens: 13, - total_tokens: 24, - total_cache_read_tokens: 17, - last_query_total_tokens: 19, - last_query_input_tokens: 23, - prompt_token_estimate: 11, - } - ); - } - - #[test] - fn acp_prompt_error_emits_turn_failed() { - let session_id = SessionId::new(); - let mut active_turn_id = Some(TurnId::new()); - let mut turn_count = 2; - - let event = acp_prompt_completed_event( - &serde_json::json!({ - "sessionId": session_id, - "error": "server -32000: failed" - }), - Some(session_id), - &mut active_turn_id, - &mut turn_count, - 11, - 13, - 24, - 17, - 19, - 23, - ) - .expect("ACP prompt failed event"); - - assert_eq!(active_turn_id, None); - assert_eq!(turn_count, 2); - assert_eq!( - event, - WorkerEvent::TurnFailed { - message: "server -32000: failed".to_string(), - turn_count: 2, - total_input_tokens: 11, - total_output_tokens: 13, - total_tokens: 24, - total_cache_read_tokens: 17, - prompt_token_estimate: 11, - last_query_input_tokens: 23, - } - ); - } - #[test] fn acp_plan_notification_emits_plan_updated() { let session_id = SessionId::new(); From 665c3f6811f0ff29d5d759d354880a778146cd71 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Fri, 3 Jul 2026 05:57:25 -1000 Subject: [PATCH 08/16] fix: Show code_search input while running and hint on first index build. Running tools now reuse the same Input rendering as completed tools, and cold code_search runs emit a progress notice before building the index. --- crates/code-search/src/service.rs | 46 +++++++++ crates/core/src/tools/handlers/code_search.rs | 23 ++++- .../src/runtime/turn_exec/tool_display.rs | 1 - crates/tui/src/chatwidget/transcript_view.rs | 85 +++++++++------- crates/tui/src/chatwidget_tests.rs | 96 +++++++++++++++++++ crates/tui/src/exec_cell/render.rs | 10 ++ 6 files changed, 224 insertions(+), 37 deletions(-) diff --git a/crates/code-search/src/service.rs b/crates/code-search/src/service.rs index d3e75784..78ff48b3 100644 --- a/crates/code-search/src/service.rs +++ b/crates/code-search/src/service.rs @@ -198,6 +198,24 @@ impl CodeSearchService { Ok(index) } + /// Returns `true` when the next `search` or `find_related` call for this + /// root/content pair will trigger a full index build (no warm memory cache + /// and no valid disk cache). Callers use this to surface a "building index" + /// notice before the potentially slow first-run. + pub fn needs_index_build(&self, root: &Path, content: ContentFilter) -> bool { + let Ok(root) = canonical_root(root) else { + return true; + }; + let key = memory_key(&root, content, self.provider.model_id()); + if self.clean_warm_index(&key).ok().flatten().is_some() { + return false; + } + let cache_path = cache_file_path(&self.cache_dir, &root, content, self.provider.model_id()); + let has_valid_cache = load_payload(&cache_path) + .is_some_and(|c| c.payload.is_valid_for(&root, content, self.provider.model_id())); + !has_valid_cache + } + /// Reuses an in-memory index without touching the filesystem when watcher /// state says it is both available and clean. fn clean_warm_index(&self, key: &str) -> Result>, CodeSearchError> { @@ -789,6 +807,34 @@ mod tests { assert_eq!(load_payload(&cache_path).is_some(), true); } + #[test] + fn needs_index_build_true_without_cache_false_after_build() { + let temp = tempfile::tempdir().expect("tempdir"); + let cache = tempfile::tempdir().expect("cache"); + fs::write(temp.path().join("lib.rs"), "pub fn alpha() {}\n").expect("write"); + let service = test_service(cache.path().to_path_buf()); + + assert!( + service.needs_index_build(temp.path(), ContentFilter::Code), + "should need build before first search" + ); + + service + .search(SearchRequest { + root: temp.path().to_path_buf(), + query: "alpha".to_string(), + content: ContentFilter::Code, + top_k: 1, + filters: SearchFilters::empty(), + }) + .expect("search"); + + assert!( + !service.needs_index_build(temp.path(), ContentFilter::Code), + "should not need build after first search" + ); + } + fn empty_index() -> Arc { let embeddings = EmbeddingMatrix::empty(); let payload = CachedIndexPayloadV4::new( diff --git a/crates/core/src/tools/handlers/code_search.rs b/crates/core/src/tools/handlers/code_search.rs index 74f8561a..8dee4546 100644 --- a/crates/core/src/tools/handlers/code_search.rs +++ b/crates/core/src/tools/handlers/code_search.rs @@ -9,7 +9,7 @@ use devo_code_search::{ use serde::Deserialize; use crate::contracts::{ - ToolCallError, ToolContext, ToolProgressSender, ToolResult, ToolResultContent, + ToolCallError, ToolContext, ToolProgress, ToolProgressSender, ToolResult, ToolResultContent, }; use crate::registry_plan::code_search_tool_spec; use crate::tool_handler::ToolHandler; @@ -67,7 +67,7 @@ impl ToolHandler for CodeSearchHandler { &self, ctx: ToolContext, input: serde_json::Value, - _progress: Option, + progress: Option, ) -> Result { if ctx.cancel_token.is_cancelled() { return Err(ToolCallError::Cancelled); @@ -75,6 +75,25 @@ impl ToolHandler for CodeSearchHandler { let input: CodeSearchInput = serde_json::from_value(input) .map_err(|error| ToolCallError::InvalidInput(error.to_string()))?; let request = build_request(&ctx.workspace_root, input)?; + + if let Some(ref sender) = progress { + let root = match &request { + CodeSearchRequest::Search(r) => &r.root, + CodeSearchRequest::FindRelated(r) => &r.root, + }; + let content = match &request { + CodeSearchRequest::Search(r) => r.content, + CodeSearchRequest::FindRelated(r) => r.content, + }; + if self.service.needs_index_build(root, content) { + let _ = sender.send(ToolProgress::StatusUpdate { + message: "First code search: building index (may take a moment, usually quick)" + .to_string(), + percent: None, + }); + } + } + let service = Arc::clone(&self.service); let output = tokio::task::spawn_blocking(move || match request { CodeSearchRequest::Search(request) => service.search(request), diff --git a/crates/server/src/runtime/turn_exec/tool_display.rs b/crates/server/src/runtime/turn_exec/tool_display.rs index 0c90de16..159246d1 100644 --- a/crates/server/src/runtime/turn_exec/tool_display.rs +++ b/crates/server/src/runtime/turn_exec/tool_display.rs @@ -345,7 +345,6 @@ pub(super) fn command_execution_item_id_for_progress( ) -> Option { pending_tool_calls .get(tool_use_id) - .filter(|pending| pending.display_kind.is_command_execution()) .and_then(|pending| pending.item_id) } diff --git a/crates/tui/src/chatwidget/transcript_view.rs b/crates/tui/src/chatwidget/transcript_view.rs index 6f6d7aa3..ee4a714d 100644 --- a/crates/tui/src/chatwidget/transcript_view.rs +++ b/crates/tui/src/chatwidget/transcript_view.rs @@ -224,40 +224,57 @@ impl ChatWidget { } } for pending in &self.pending_tool_calls { - let pending_lines = if let Some(start_time) = pending.start_time { - let mut pending_lines = vec![Line::from(vec![ - crate::exec_cell::spinner(Some(start_time), true), - " ".into(), - Span::styled(pending.title.clone(), Self::tool_text_style()), - ])]; - pending_lines.extend(pending.lines.clone()); - pending_lines + if let (Some(tool_name), Some(input)) = (&pending.tool_name, &pending.input) { + let tool_lines = ToolIoCell::from_text_output( + ToolIoCellOptions { + title_line: Some(Self::running_tool_line(&pending.title)), + dot_prefix: Self::pending_dot_prefix(), + subsequent_prefix: " ".into(), + output_style: Self::tool_text_style(), + show_empty_ellipsis: false, + }, + tool_name.clone(), + input.clone(), + pending.output.clone(), + ) + .transcript_lines(width); + Self::extend_lines_with_separator(&mut lines, tool_lines); } else { - pending.lines.clone() - }; - Self::extend_lines_with_separator( - &mut lines, - match mode { - LiveViewportLineMode::Display => { - history_cell::AgentMessageCell::new_with_prefix( - pending_lines, - Self::pending_dot_prefix(), - " ", - false, - ) - .display_lines(width) - } - LiveViewportLineMode::Transcript => { - history_cell::AgentMessageCell::new_with_prefix( - pending_lines, - Self::pending_dot_prefix(), - " ", - false, - ) - .transcript_lines(width) - } - }, - ); + let pending_lines = if let Some(start_time) = pending.start_time { + let mut pending_lines = vec![Line::from(vec![ + crate::exec_cell::spinner(Some(start_time), true), + " ".into(), + Span::styled(pending.title.clone(), Self::tool_text_style()), + ])]; + pending_lines.extend(pending.lines.clone()); + pending_lines + } else { + pending.lines.clone() + }; + Self::extend_lines_with_separator( + &mut lines, + match mode { + LiveViewportLineMode::Display => { + history_cell::AgentMessageCell::new_with_prefix( + pending_lines, + Self::pending_dot_prefix(), + " ", + false, + ) + .display_lines(width) + } + LiveViewportLineMode::Transcript => { + history_cell::AgentMessageCell::new_with_prefix( + pending_lines, + Self::pending_dot_prefix(), + " ", + false, + ) + .transcript_lines(width) + } + }, + ); + } } Self::trim_trailing_blank_lines(&mut lines); lines @@ -280,7 +297,7 @@ impl ChatWidget { input.clone(), tool_call.output.clone(), ) - .display_lines(width), + .transcript_lines(width), _ => history_cell::AgentMessageCell::new_with_prefix( tool_call.lines.clone(), Self::pending_dot_prefix(), diff --git a/crates/tui/src/chatwidget_tests.rs b/crates/tui/src/chatwidget_tests.rs index 6b4bc4b1..634e34aa 100644 --- a/crates/tui/src/chatwidget_tests.rs +++ b/crates/tui/src/chatwidget_tests.rs @@ -7718,6 +7718,102 @@ fn code_search_tool_call_renders_as_explored_group_in_viewport() { ); } +#[test] +fn running_code_search_with_details_shows_input_in_live_viewport() { + let model = Model { + slug: "test-model".to_string(), + display_name: "Test Model".to_string(), + ..Model::default() + }; + let (mut widget, _app_event_rx) = widget_with_model(model, PathBuf::from(".")); + + widget.handle_worker_event(crate::events::WorkerEvent::ToolCall { + tool_use_id: "tool-1".to_string(), + summary: "code_search live tool feedback in crates".to_string(), + preparing: false, + parsed_commands: None, + }); + widget.handle_worker_event(crate::events::WorkerEvent::ToolCallDetails { + tool_use_id: "tool-1".to_string(), + tool_name: "code_search".to_string(), + input: serde_json::json!({ + "operation": "search", + "query": "live tool feedback", + "path": "crates" + }), + }); + + let rendered = rendered_rows(&widget, 80, 16).join("\n"); + + assert!( + rendered.contains("operation") && rendered.contains("search"), + "live viewport should show 'operation: search' for running code_search:\n{rendered}" + ); + assert!( + rendered.contains("query") && rendered.contains("live tool feedback"), + "live viewport should show 'query: live tool feedback' for running code_search:\n{rendered}" + ); + assert!( + rendered.contains("path") && rendered.contains("crates"), + "live viewport should show 'path: crates' for running code_search:\n{rendered}" + ); +} + +#[test] +fn exploring_code_search_with_details_shows_input_in_active_cell() { + let model = Model { + slug: "test-model".to_string(), + display_name: "Test Model".to_string(), + ..Model::default() + }; + let (mut widget, _app_event_rx) = widget_with_model(model, PathBuf::from(".")); + + widget.handle_worker_event(crate::events::WorkerEvent::ToolCall { + tool_use_id: "tool-1".to_string(), + summary: "code_search live tool feedback in crates".to_string(), + preparing: false, + parsed_commands: Some(vec![devo_protocol::parse_command::ParsedCommand::Search { + cmd: "code_search live tool feedback in crates".to_string(), + query: Some("live tool feedback".to_string()), + path: Some("crates".to_string()), + }]), + }); + widget.handle_worker_event(crate::events::WorkerEvent::ToolCallDetails { + tool_use_id: "tool-1".to_string(), + tool_name: "code_search".to_string(), + input: serde_json::json!({ + "operation": "search", + "query": "live tool feedback", + "path": "crates" + }), + }); + + let live_display = widget + .active_cell_display_lines_for_test(80) + .into_iter() + .map(|line| { + line.spans + .into_iter() + .map(|span| span.content.to_string()) + .collect::() + }) + .collect::>() + .join("\n"); + + assert!( + live_display.contains("Exploring"), + "expected Exploring header: {live_display}" + ); + assert!( + live_display.contains("operation") && live_display.contains("search"), + "active ExecCell should show 'operation: search' while exploring:\n{live_display}" + ); + assert!( + live_display.contains("query") && live_display.contains("live tool feedback"), + "active ExecCell should show 'query: live tool feedback' while exploring:\n{live_display}" + ); +} + #[test] fn merged_explored_group_becomes_explored_after_all_results_arrive() { let model = Model { diff --git a/crates/tui/src/exec_cell/render.rs b/crates/tui/src/exec_cell/render.rs index 98a36fd3..b8f536d6 100644 --- a/crates/tui/src/exec_cell/render.rs +++ b/crates/tui/src/exec_cell/render.rs @@ -13,6 +13,7 @@ use crate::render::line_utils::prefix_lines; use crate::render::line_utils::push_owned_lines; use crate::tool_io_cell::ToolIoCell; use crate::tool_io_cell::ToolIoCellOptions; +use crate::tool_io_cell::tool_input_lines; use crate::wrapping::RtOptions; use crate::wrapping::adaptive_wrap_line; use crate::wrapping::adaptive_wrap_lines; @@ -481,6 +482,15 @@ impl ExecCell { ); push_owned_lines(&wrapped, &mut out_indented); } + + if call.output.is_none() { + if let (Some(tool_name), Some(input)) = (&call.tool_name, &call.tool_input) { + let input_lines = tool_input_lines(tool_name, input); + for line in input_lines { + out_indented.push(line.patch_style(Style::default().dim())); + } + } + } } out.extend(prefix_lines(out_indented, " └ ".dim(), " ".into())); From 7c66714b02b210208d64ab2dc7c2acaf51d761a7 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Fri, 3 Jul 2026 21:34:03 -1000 Subject: [PATCH 09/16] fix: gracefully handle already-installed tracing subscriber in logging bootstrap When another subscriber has already been installed (e.g. by tokio-console), LoggingBootstrap::install() now emits a warning and continues instead of returning an error. This allows optional diagnostic subscribers to coexist with the file-logging path without requiring callers to manage installation order. --- crates/core/src/logging.rs | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/crates/core/src/logging.rs b/crates/core/src/logging.rs index db903d82..84bdc683 100644 --- a/crates/core/src/logging.rs +++ b/crates/core/src/logging.rs @@ -83,7 +83,15 @@ impl LoggingBootstrap { })?; let (file_writer, file_guard) = tracing_appender::non_blocking(file_appender); - install_subscriber(file_level, self.config.json, file_writer)?; + match install_subscriber(file_level, self.config.json, file_writer) { + Err(LoggingInitError::SubscriberAlreadyInstalled) => { + // Another subscriber was already installed (e.g. by tokio-console). + // This is not a fatal error — the process will run without file logging. + tracing::warn!("file logging skipped: a global tracing subscriber is already installed"); + } + Err(other) => return Err(other), + Ok(()) => {} + } install_panic_hook(); tracing::info!( From 394ff333afa0be17271138e65d9413f6d0c145dd Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Fri, 3 Jul 2026 21:34:18 -1000 Subject: [PATCH 10/16] feat: add optional tokio-console instrumentation behind a feature flag Adds a tokio-console feature (gated behind cfg(feature = tokio-console)) that initializes a console-subscriber before the tokio runtime starts, enabling real-time async task inspection via tokio-console CLI. - arg0: new maybe_init_tokio_console() + wired into run_server_alias_dispatch - cli: forwarded feature + init before server logging in Command::Server - Both entry points (standalone devo-server and CLI-spawned server) covered - Environment variable TOKIO_CONSOLE must be set at runtime to activate - File logging is gracefully suppressed when console-subscriber takes over --- crates/arg0/Cargo.toml | 5 +++++ crates/arg0/src/lib.rs | 13 +++++++++++++ crates/cli/Cargo.toml | 3 +++ crates/cli/src/main.rs | 4 ++++ 4 files changed, 25 insertions(+) diff --git a/crates/arg0/Cargo.toml b/crates/arg0/Cargo.toml index 835af4d0..3d2806d3 100644 --- a/crates/arg0/Cargo.toml +++ b/crates/arg0/Cargo.toml @@ -11,6 +11,7 @@ path = "src/lib.rs" [dependencies] anyhow = { workspace = true } clap = { workspace = true } +console-subscriber = { version = "0.5", optional = true } devo-core = { workspace = true } devo-server = { workspace = true } devo-util-paths = { workspace = true } @@ -21,5 +22,9 @@ tokio = { workspace = true, features = ["rt-multi-thread"] } tracing = { workspace = true } tracing-subscriber = { workspace = true } +[features] +default = [] +tokio-console = ["console-subscriber"] + [dev-dependencies] pretty_assertions = { workspace = true } diff --git a/crates/arg0/src/lib.rs b/crates/arg0/src/lib.rs index 3c03c993..0352afe2 100644 --- a/crates/arg0/src/lib.rs +++ b/crates/arg0/src/lib.rs @@ -61,6 +61,18 @@ pub enum EarlyDispatch { // ── Public API ──────────────────────────────────────────────────────────── +/// Initialize the tokio-console tracing subscriber if `TOKIO_CONSOLE` is set +/// and the `tokio-console` feature is active. Must be called before any other +/// tracing subscriber is installed. +pub fn maybe_init_tokio_console() { + #[cfg(feature = "tokio-console")] + if std::env::var("TOKIO_CONSOLE").is_ok() { + console_subscriber::init(); + } + #[cfg(not(feature = "tokio-console"))] + let _ = std::env::var("TOKIO_CONSOLE"); +} + /// Entry‑point wrapper that performs `argv[0]` dispatch first. /// /// If the current executable was invoked as `devo-server` (via symlink or batch @@ -197,6 +209,7 @@ fn build_runtime() -> Result { } fn run_server_alias_dispatch() -> Result<()> { + maybe_init_tokio_console(); let runtime = build_runtime()?; runtime.block_on(run_server_dispatch()); Ok(()) diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index d848caef..f1379c7f 100644 --- a/crates/cli/Cargo.toml +++ b/crates/cli/Cargo.toml @@ -33,5 +33,8 @@ toml = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } +[features] +tokio-console = ["devo-arg0/tokio-console"] + [dev-dependencies] pretty_assertions = "1" diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index a75553ef..2550e188 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -223,6 +223,10 @@ async fn run_cli() -> Result<()> { status: _, shutdown: _, }) => { + // Start tokio-console before file logging so the console subscriber + // can capture task instrumentation. File logging will fall back + // gracefully if a subscriber is already installed. + devo_arg0::maybe_init_tokio_console(); let args = server_process_args_from_cli(&cli).expect("server command args"); let _logging = install_server_logging(&cli)?; run_server_process(args, ServerProcessRunOptions::default()).await From 6a6fac05a40be8666ea96ee6760a1a9acd21ce87 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Fri, 3 Jul 2026 21:35:40 -1000 Subject: [PATCH 11/16] feat: add optional wire-level protocol tracing via DEVO_PROTOCOL_TRACE Records every NDJSON frame sent and received to a structured trace file under DEVO_HOME/traces/ when DEVO_PROTOCOL_TRACE=1 is set. Useful for debugging client-server communication issues without modifying application-level logging. - Adds devo-util-paths dependency for DEVO_HOME resolution - Trace file location can be overridden with DEVO_PROTOCOL_TRACE_FILE - Each record includes monotonic seq, UTC timestamp, direction and payload --- Cargo.lock | 151 ++++++++++++++++++++ crates/client/Cargo.toml | 2 + crates/client/src/lib.rs | 1 + crates/client/src/protocol_trace.rs | 207 ++++++++++++++++++++++++++++ crates/client/src/stdio.rs | 22 ++- docs/configuration.md | 2 +- docs/protocol-trace.md | 97 +++++++++++++ 7 files changed, 477 insertions(+), 5 deletions(-) create mode 100644 crates/client/src/protocol_trace.rs create mode 100644 docs/protocol-trace.md diff --git a/Cargo.lock b/Cargo.lock index 4b405433..d11dc99f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -501,6 +501,12 @@ version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" +[[package]] +name = "base64" +version = "0.21.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567" + [[package]] name = "base64" version = "0.22.1" @@ -970,6 +976,46 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "console-api" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8599749b6667e2f0c910c1d0dff6901163ff698a52d5a39720f61b5be4b20d3" +dependencies = [ + "futures-core", + "prost", + "prost-types", + "tonic", + "tonic-prost", + "tracing-core", +] + +[[package]] +name = "console-subscriber" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb4915b7d8dd960457a1b6c380114c2944f728e7c65294ab247ae6b6f1f37592" +dependencies = [ + "console-api", + "crossbeam-channel", + "crossbeam-utils", + "futures-task", + "hdrhistogram", + "humantime", + "hyper-util", + "prost", + "prost-types", + "serde", + "serde_json", + "thread_local", + "tokio", + "tokio-stream", + "tonic", + "tracing", + "tracing-core", + "tracing-subscriber", +] + [[package]] name = "const-oid" version = "0.10.2" @@ -1387,6 +1433,7 @@ version = "0.1.24" dependencies = [ "anyhow", "clap", + "console-subscriber", "devo-core", "devo-server", "devo-util-paths", @@ -1434,10 +1481,12 @@ dependencies = [ "anyhow", "chrono", "devo-protocol", + "devo-util-paths", "futures", "pretty_assertions", "serde", "serde_json", + "tempfile", "tokio", "tokio-tungstenite", "tracing", @@ -2595,6 +2644,19 @@ dependencies = [ "hashbrown 0.15.5", ] +[[package]] +name = "hdrhistogram" +version = "7.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "765c9198f173dd59ce26ff9f95ef0aafd0a0fe01fb9d72841bc5066a4c06511d" +dependencies = [ + "base64 0.21.7", + "byteorder", + "flate2", + "nom 7.1.3", + "num-traits", +] + [[package]] name = "heapify" version = "0.2.0" @@ -2809,6 +2871,19 @@ dependencies = [ "webpki-roots", ] +[[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "hyper-util" version = "0.1.20" @@ -4321,6 +4396,38 @@ dependencies = [ "windows 0.62.2", ] +[[package]] +name = "prost" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "528ac67416ff8646872a3c02cad9cc4ee5dc9f9540c9b10771855c95cb2e5ae1" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-derive" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" +dependencies = [ + "anyhow", + "itertools 0.14.0", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "prost-types" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f94967dc7688f3054c7fac87473ffae4cc4c3904800e2d9f5b857246d8963b0a" +dependencies = [ + "prost", +] + [[package]] name = "pulldown-cmark" version = "0.10.3" @@ -5998,6 +6105,7 @@ dependencies = [ "signal-hook-registry", "socket2", "tokio-macros", + "tracing", "windows-sys 0.61.2", ] @@ -6142,6 +6250,46 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801" +[[package]] +name = "tonic" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" +dependencies = [ + "async-trait", + "axum", + "base64 0.22.1", + "bytes", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-timeout", + "hyper-util", + "percent-encoding", + "pin-project", + "socket2", + "sync_wrapper", + "tokio", + "tokio-stream", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tonic-prost" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" +dependencies = [ + "bytes", + "prost", + "tonic", +] + [[package]] name = "tower" version = "0.5.3" @@ -6150,9 +6298,12 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", + "indexmap 2.14.0", "pin-project-lite", + "slab", "sync_wrapper", "tokio", + "tokio-util", "tower-layer", "tower-service", "tracing", diff --git a/crates/client/Cargo.toml b/crates/client/Cargo.toml index f00a54fa..ef5ee17d 100644 --- a/crates/client/Cargo.toml +++ b/crates/client/Cargo.toml @@ -12,6 +12,7 @@ path = "src/lib.rs" anyhow = { workspace = true } chrono = { workspace = true } devo-protocol = { workspace = true } +devo-util-paths = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } futures = { workspace = true } @@ -21,3 +22,4 @@ tracing = { workspace = true } [dev-dependencies] pretty_assertions = { workspace = true } +tempfile = { workspace = true } diff --git a/crates/client/src/lib.rs b/crates/client/src/lib.rs index e6d8a356..4e5ccd77 100644 --- a/crates/client/src/lib.rs +++ b/crates/client/src/lib.rs @@ -8,6 +8,7 @@ mod acp_fs; mod acp_permissions; mod acp_terminal; mod client_core; +mod protocol_trace; mod stdio; mod websocket; diff --git a/crates/client/src/protocol_trace.rs b/crates/client/src/protocol_trace.rs new file mode 100644 index 00000000..317ddb81 --- /dev/null +++ b/crates/client/src/protocol_trace.rs @@ -0,0 +1,207 @@ +//! Wire-level protocol trace for stdio transport. +//! +//! When enabled via `DEVO_PROTOCOL_TRACE=1`, every NDJSON line sent to or +//! received from the server child process is recorded to a structured NDJSONL +//! file. Each record carries a monotonic sequence number, UTC timestamp, +//! direction, byte count, and the raw JSON payload as it appeared on the wire. + +use std::fs::{self, File}; +use std::io::Write; +use std::path::PathBuf; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; + +use chrono::Utc; + +/// Direction of the protocol message relative to the client. +#[derive(Clone, Copy)] +pub(crate) enum TraceDirection { + /// Client → server (outbound to child stdin). + Out, + /// Server → client (inbound from child stdout). + In, +} + +impl TraceDirection { + fn as_str(self) -> &'static str { + match self { + TraceDirection::Out => "out", + TraceDirection::In => "in", + } + } +} + +struct ProtocolTraceInner { + file: Mutex, + seq: AtomicU64, +} + +/// Captures raw NDJSON wire traffic to a structured trace file. +/// +/// Cheap to clone (wraps an `Arc`). Thread-safe: the writer/reader async tasks +/// each hold a clone and call `record` independently; the shared `AtomicU64` +/// sequence counter ensures a total order across directions. +#[derive(Clone)] +pub(crate) struct ProtocolTrace { + inner: Arc, +} + +impl ProtocolTrace { + /// Reads `DEVO_PROTOCOL_TRACE` and (optionally) `DEVO_PROTOCOL_TRACE_FILE` + /// from the environment. Returns `None` when tracing is disabled or when the + /// trace file cannot be created. + /// + /// Called once during [`StdioServerClient::spawn`]; the result is cloned + /// into the writer and reader tasks. + pub(crate) fn from_env() -> Option { + let enabled = std::env::var("DEVO_PROTOCOL_TRACE") + .ok() + .filter(|v| !v.is_empty()) + .map(|v| v == "1" || v.eq_ignore_ascii_case("true")) + .unwrap_or(false); + if !enabled { + return None; + } + + let path = resolve_trace_path(); + match File::create(&path) { + Ok(file) => { + tracing::info!(path = %path.display(), "protocol trace enabled"); + Some(Self { + inner: Arc::new(ProtocolTraceInner { + file: Mutex::new(file), + seq: AtomicU64::new(1), + }), + }) + } + Err(err) => { + tracing::warn!( + path = %path.display(), + error = %err, + "failed to create protocol trace file; tracing disabled" + ); + None + } + } + } + + /// Creates a `ProtocolTrace` that writes to the given file. Used in tests. + #[cfg(test)] + pub(crate) fn with_file(file: File) -> Self { + Self { + inner: Arc::new(ProtocolTraceInner { + file: Mutex::new(file), + seq: AtomicU64::new(1), + }), + } + } + + /// Record a single protocol line. + pub(crate) fn record(&self, dir: TraceDirection, line: &str) { + let seq = self.inner.seq.fetch_add(1, Ordering::Relaxed); + let ts = Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true); + let bytes = line.len(); + + let record = serde_json::json!({ + "seq": seq, + "ts": ts, + "dir": dir.as_str(), + "bytes": bytes, + "line": line, + }); + + if let Ok(mut f) = self.inner.file.lock() { + let buf = serde_json::to_vec(&record).expect("serialize trace record"); + let _ = f.write_all(&buf); + let _ = f.write_all(b"\n"); + let _ = f.flush(); + } + } +} + +fn resolve_trace_path() -> PathBuf { + if let Ok(explicit) = std::env::var("DEVO_PROTOCOL_TRACE_FILE") { + if !explicit.is_empty() { + let path = PathBuf::from(&explicit); + if let Some(parent) = path.parent() { + let _ = fs::create_dir_all(parent); + } + return path; + } + } + + let base = devo_util_paths::find_devo_home() + .map(|home| home.join("traces")) + .unwrap_or_else(|_| { + let mut tmp = std::env::temp_dir(); + tmp.push("devo-traces"); + tmp + }); + let _ = fs::create_dir_all(&base); + + let pid = std::process::id(); + let ts = Utc::now().format("%Y%m%dT%H%M%SZ"); + base.join(format!("protocol-{pid}-{ts}.ndjsonl")) +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + use std::io::{BufRead, BufReader}; + + #[test] + fn from_env_returns_none_when_not_set() { + // DEVO_PROTOCOL_TRACE is not set in the test environment by default. + // We cannot safely call remove_var, but the CI/test environment does + // not set this variable, so from_env should return None. + if std::env::var("DEVO_PROTOCOL_TRACE").is_err() { + assert!(ProtocolTrace::from_env().is_none()); + } + } + + #[test] + fn records_outbound_and_inbound_with_monotonic_seq() { + let dir = tempfile::tempdir().expect("create temp dir"); + let path = dir.path().join("trace.ndjsonl"); + let file = File::create(&path).expect("create trace file"); + let trace = ProtocolTrace::with_file(file); + + trace.record(TraceDirection::Out, r#"{"jsonrpc":"2.0","id":1,"method":"initialize"}"#); + trace.record(TraceDirection::In, r#"{"jsonrpc":"2.0","id":1,"result":{}}"#); + trace.record(TraceDirection::Out, r#"{"jsonrpc":"2.0","id":2,"method":"session/new"}"#); + + let reader = BufReader::new(File::open(&path).expect("open trace file")); + let records: Vec = reader + .lines() + .map(|line| serde_json::from_str(&line.expect("read line")).expect("parse record")) + .collect(); + + assert_eq!(records.len(), 3); + + assert_eq!(records[0]["seq"], 1); + assert_eq!(records[0]["dir"], "out"); + assert_eq!( + records[0]["line"], + r#"{"jsonrpc":"2.0","id":1,"method":"initialize"}"# + ); + assert_eq!( + records[0]["bytes"], + r#"{"jsonrpc":"2.0","id":1,"method":"initialize"}"#.len() + ); + + assert_eq!(records[1]["seq"], 2); + assert_eq!(records[1]["dir"], "in"); + assert_eq!( + records[1]["line"], + r#"{"jsonrpc":"2.0","id":1,"result":{}}"# + ); + + assert_eq!(records[2]["seq"], 3); + assert_eq!(records[2]["dir"], "out"); + + for record in &records { + assert!(record["ts"].is_string()); + } + } +} diff --git a/crates/client/src/stdio.rs b/crates/client/src/stdio.rs index dd422b7f..f7e42d56 100644 --- a/crates/client/src/stdio.rs +++ b/crates/client/src/stdio.rs @@ -25,10 +25,10 @@ use tokio::time::timeout; use crate::client_core::ClientWriteMessage; use crate::client_core::ClientWriter; use crate::client_core::ServerClientCore; +use crate::protocol_trace::ProtocolTrace; +use crate::protocol_trace::TraceDirection; pub use crate::acp_terminal::ACP_TERMINAL_OUTPUT_NOTIFICATION_METHOD; -pub use crate::client_core::ACP_PROMPT_COMPLETED_NOTIFICATION_METHOD; -pub use crate::client_core::ACP_PROMPT_STARTED_NOTIFICATION_METHOD; pub use crate::client_core::ServerNotificationMessage; const SERVER_CHILD_STDIN_SHUTDOWN_TIMEOUT: Duration = Duration::from_millis(100); @@ -85,14 +85,21 @@ impl StdioServerClient { ); let reader_state = core.reader_state(); let stdin = Arc::new(Mutex::new(stdin)); + let trace = ProtocolTrace::from_env(); + let writer_trace = trace.clone(); let writer_task = tokio::spawn(run_stdin_writer( Arc::clone(&stdin), write_rx, + writer_trace, )); + let reader_trace = trace; let reader_task = tokio::spawn(async move { let mut lines = BufReader::new(stdout).lines(); while let Ok(Some(line)) = lines.next_line().await { + if let Some(ref t) = reader_trace { + t.record(TraceDirection::In, &line); + } match serde_json::from_str::(&line) { Ok(message) => reader_state.handle_message(message).await, Err(_) => { @@ -414,11 +421,12 @@ impl StdioServerClient { async fn run_stdin_writer( stdin: Arc>, mut write_rx: tokio::sync::mpsc::UnboundedReceiver, + trace: Option, ) -> Result<()> { while let Some(message) = write_rx.recv().await { match message { ClientWriteMessage::Json(value) => { - write_ndjson_to_stdin(&stdin, &value) + write_ndjson_to_stdin(&stdin, &value, trace.as_ref()) .await .context("write client payload")?; } @@ -438,8 +446,14 @@ async fn run_stdin_writer( async fn write_ndjson_to_stdin( stdin: &Arc>, value: &serde_json::Value, + trace: Option<&ProtocolTrace>, ) -> Result<()> { let mut line = serde_json::to_vec(value).context("serialize client payload")?; + if let Some(t) = trace { + if let Ok(s) = std::str::from_utf8(&line) { + t.record(TraceDirection::Out, s); + } + } line.push(b'\n'); let mut stdin = stdin.lock().await; stdin.write_all(&line).await.context("write client payload")?; @@ -508,7 +522,7 @@ mod tests { while let Some(message) = write_rx.recv().await { match message { ClientWriteMessage::Json(value) => { - if write_ndjson_to_stdin(&stdin, &value).await.is_err() { + if write_ndjson_to_stdin(&stdin, &value, None).await.is_err() { break; } } diff --git a/docs/configuration.md b/docs/configuration.md index 4c0fadc5..cc45d9f4 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -93,4 +93,4 @@ provider = "my.provider" model_name = "provider-specific-model-name" display_name = "My Coding Model" invocation_method = "openai_chat_completions" -``` +``` \ No newline at end of file diff --git a/docs/protocol-trace.md b/docs/protocol-trace.md new file mode 100644 index 00000000..e57651e7 --- /dev/null +++ b/docs/protocol-trace.md @@ -0,0 +1,97 @@ +# Protocol Trace (stdio) + +When debugging client-server communication you can capture every NDJSON line +exchanged over the stdio transport. The trace runs inside the client process and +records raw wire traffic in both directions without modifying the server or its +stdout stream. + +## Enabling + +Set the `DEVO_PROTOCOL_TRACE` environment variable before launching Devo: + +```bash +DEVO_PROTOCOL_TRACE=1 devo +``` + +On Windows (PowerShell): + +```powershell +$env:DEVO_PROTOCOL_TRACE = "1"; devo +``` + +## Output location + +Trace files are written to `DEVO_HOME/traces/` (default `~/.devo/traces/`) +using the naming pattern `protocol--.ndjsonl`. + +To write to a specific path instead, set `DEVO_PROTOCOL_TRACE_FILE`: + +```bash +DEVO_PROTOCOL_TRACE=1 DEVO_PROTOCOL_TRACE_FILE=/tmp/my-trace.ndjsonl devo +``` + +If `DEVO_HOME` cannot be resolved, the trace falls back to +`/devo-traces/`. + +## Record format + +Each line in the trace file is a JSON object (NDJSONL): + +```json +{"seq":1,"ts":"2026-07-03T15:30:00.123Z","dir":"out","bytes":128,"line":"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",...}"} +{"seq":2,"ts":"2026-07-03T15:30:00.145Z","dir":"in","bytes":256,"line":"{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{...}}"} +``` + +| Field | Description | +|---------|--------------------------------------------------------------------| +| `seq` | Monotonic sequence number across both directions | +| `ts` | UTC timestamp (RFC 3339, millisecond precision) | +| `dir` | `"out"` = client to server, `"in"` = server to client | +| `bytes` | Byte length of the raw JSON payload | +| `line` | The original JSON-RPC line exactly as sent or received on the wire | + +The `seq` counter is shared between the writer and reader async tasks via an +`AtomicU64`, so records from both directions can be sorted into a single +chronological stream. + +## Hook points + +The trace is inserted at two locations in +[`crates/client/src/stdio.rs`](../crates/client/src/stdio.rs): + +| Direction | Function | When | +|-----------|-----------------------|-------------------------------------------------| +| C → S | `write_ndjson_to_stdin` | After `serde_json::to_vec`, before `write_all` | +| S → C | stdout reader loop | After `next_line`, before `serde_json::from_str` | + +This ensures the trace captures the exact bytes that cross the pipe boundary, +including malformed payloads that fail JSON parsing. + +## Querying traces + +Use `jq` to filter and inspect: + +```bash +# Show only outbound (client → server) messages +jq 'select(.dir == "out")' ~/.devo/traces/protocol-*.ndjsonl + +# Show methods of all outbound requests +jq -r 'select(.dir == "out") | .line | fromjson | .method // empty' ~/.devo/traces/protocol-*.ndjsonl + +# Show inbound messages larger than 1 KB +jq 'select(.dir == "in" and .bytes > 1024)' ~/.devo/traces/protocol-*.ndjsonl +``` + +## Security note + +Trace files contain the full protocol payload, which may include file contents, +API keys, and other sensitive data. Treat them with the same care as log files +and delete them when no longer needed. The feature is disabled by default. + +## Implementation + +The core logic lives in +[`crates/client/src/protocol_trace.rs`](../crates/client/src/protocol_trace.rs). +`ProtocolTrace::from_env()` is called once during `StdioServerClient::spawn`; +when disabled (`None`), the code path is a zero-cost no-op. File I/O uses +`std::sync::Mutex` with explicit `flush()` after each record. From bfe773fc28c49e3ab6e315a02c3c4a4ebda4a0bf Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Fri, 3 Jul 2026 21:35:51 -1000 Subject: [PATCH 12/16] fix: preserve streamed runtime events under backpressure and track subagent usage ownership - event_stream: drop overflow events instead of spawning unbounded tokio tasks that would block forever and prevent channel close - subagent_usage: track usage ownership per subagent turn for accurate parent-child token accounting across the agent tree - finalize: add stream idle timeout to provider HTTP client so a stalled upstream connection can no longer wedge a turn - items/turn_inline: consolidate inline item persistence to reduce mailbox round-trips during active turn execution --- crates/server/src/runtime/items.rs | 66 +++-- crates/server/src/runtime/outbound.rs | 14 + crates/server/src/runtime/research.rs | 4 + .../src/runtime/session_actor/turn_inline.rs | 22 ++ crates/server/src/runtime/subagent_usage.rs | 245 ++++++++++++++++-- .../src/runtime/turn_exec/event_stream.rs | 46 +++- .../server/src/runtime/turn_exec/finalize.rs | 9 +- 7 files changed, 346 insertions(+), 60 deletions(-) diff --git a/crates/server/src/runtime/items.rs b/crates/server/src/runtime/items.rs index e04e3993..b54c28f4 100644 --- a/crates/server/src/runtime/items.rs +++ b/crates/server/src/runtime/items.rs @@ -346,33 +346,45 @@ impl ServerRuntime { worklog: Option, ) { if let Some(stream) = self.active_stream_state(session_id).await { - let mut stream = stream.lock().await; - if let Some(inline) = stream.turn_inline.as_mut() { - if inline.turn_id == turn_id - && let Some(history_item) = history_item_from_turn_item(&turn_item) - { - inline.history_items.push(history_item); - } - if inline.turn_id == turn_id { - inline - .persisted_turn_items - .push(crate::execution::PersistedTurnItem { - turn_id, - turn_kind: inline.turn_kind.clone(), - item_id, - turn_item: turn_item.clone(), - }); - } - if let Some(record) = inline.record.clone() { - let item = build_item_record( - session_id, - turn_id, - item_id, - item_seq, - turn_item, - turn_status, - worklog, - ); + // Mutate inline state under the lock, then release before any + // blocking rollout I/O so the event stream cannot pin the async + // mutex across synchronous disk writes. + let inline_rollout = { + let mut stream = stream.lock().await; + stream.turn_inline.as_mut().map(|inline| { + if inline.turn_id == turn_id + && let Some(history_item) = history_item_from_turn_item(&turn_item) + { + inline.history_items.push(history_item); + } + if inline.turn_id == turn_id { + inline + .persisted_turn_items + .push(crate::execution::PersistedTurnItem { + turn_id, + turn_kind: inline.turn_kind.clone(), + item_id, + turn_item: turn_item.clone(), + }); + } + inline.record.clone().map(|record| { + ( + record, + build_item_record( + session_id, + turn_id, + item_id, + item_seq, + turn_item.clone(), + turn_status.clone(), + worklog.clone(), + ), + ) + }) + }) + }; + if let Some(rollout) = inline_rollout { + if let Some((record, item)) = rollout { if let Err(error) = self.rollout_store.append_item(&record, item) { tracing::warn!(session_id = %session_id, error = %error, "failed to persist item line"); } diff --git a/crates/server/src/runtime/outbound.rs b/crates/server/src/runtime/outbound.rs index f706a53e..b6988433 100644 --- a/crates/server/src/runtime/outbound.rs +++ b/crates/server/src/runtime/outbound.rs @@ -160,6 +160,16 @@ pub(crate) fn log_outbound_frame(frame: &OutboundFrame, notification: &serde_jso } } +/// Enqueue an outbound frame, waiting when the channel is full. +/// +/// Attempts to reserve a send permit within +/// [`OUTBOUND_BACKPRESSURE_LOG_THRESHOLD`]. If the channel is still full after +/// that window, logs a backpressure warning and blocks until a permit is +/// available (or the receiver is dropped). +/// +/// Returns `true` when the frame is accepted, or `false` when the outbound +/// receiver has already been dropped. `queue` is a static label used only in +/// log fields to identify which outbound path is under pressure. pub(crate) async fn enqueue_outbound( tx: &mpsc::Sender, frame: OutboundFrame, @@ -167,6 +177,8 @@ pub(crate) async fn enqueue_outbound( ) -> bool { let connection_id = frame.connection_id(); let reserve_started_at = Instant::now(); + // Prefer a timed reserve so slow consumers surface as backpressure logs + // instead of silent stalls. let permit = match tokio::time::timeout(OUTBOUND_BACKPRESSURE_LOG_THRESHOLD, tx.reserve()).await { Ok(Ok(permit)) => permit, @@ -175,6 +187,8 @@ pub(crate) async fn enqueue_outbound( return false; } Err(_) => { + // Timed out waiting for capacity; keep blocking so messages are not + // dropped, but record that the consumer is lagging. tracing::warn!( connection_id, queue, diff --git a/crates/server/src/runtime/research.rs b/crates/server/src/runtime/research.rs index 35e96ff4..d7af53fa 100644 --- a/crates/server/src/runtime/research.rs +++ b/crates/server/src/runtime/research.rs @@ -1096,6 +1096,7 @@ impl ServerRuntime { turn.turn_id, own_usage.clone(), /*context_window*/ None, + super::subagent_usage::UsageUpdateKind::InFlight, ) .await .map(|snapshot| snapshot.turn_usage.to_turn_usage()) @@ -1170,11 +1171,14 @@ impl ServerRuntime { ledger.by_invocation.insert(usage_key, usage); ledger.aggregate() }; + // Research ledger already aggregates all invocations; publish as + // in-flight replacement so repeated updates do not double-count. self.publish_parent_turn_usage( session_id, turn_id, aggregate.to_turn_usage(), context_window, + super::subagent_usage::UsageUpdateKind::InFlight, ) .await; } diff --git a/crates/server/src/runtime/session_actor/turn_inline.rs b/crates/server/src/runtime/session_actor/turn_inline.rs index ac2ce8f2..b07ae318 100644 --- a/crates/server/src/runtime/session_actor/turn_inline.rs +++ b/crates/server/src/runtime/session_actor/turn_inline.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use devo_core::SessionRecord; use devo_core::TurnId; use devo_core::TurnKind; +use devo_core::TurnUsage; use devo_protocol::CollaborationMode; use crate::execution::ApprovalGrantCache; @@ -26,6 +27,7 @@ pub(crate) struct TurnInlineState { pub(crate) session_approval_cache: ApprovalGrantCache, pub(crate) turn_approval_cache: ApprovalGrantCache, pub(crate) summary: SessionMetadata, + pub(crate) active_turn_usage: Option, pub(crate) collaboration_mode: CollaborationMode, pub(crate) hook_context: HookContextSnapshot, } @@ -43,6 +45,10 @@ impl TurnInlineState { session_approval_cache: state.session_approval_cache.clone(), turn_approval_cache: state.turn_approval_cache.clone(), summary: state.summary.clone(), + active_turn_usage: state + .active_turn + .as_ref() + .and_then(|turn| turn.usage.clone()), collaboration_mode: state.core.collaboration_mode, hook_context: HookContextSnapshot { runtime_context: Arc::clone(&state.runtime_context), @@ -67,5 +73,21 @@ impl TurnInlineState { state.history_items.extend(self.history_items); state.session_approval_cache = self.session_approval_cache; state.turn_approval_cache = self.turn_approval_cache; + state.summary.total_input_tokens = self.summary.total_input_tokens; + state.summary.total_output_tokens = self.summary.total_output_tokens; + state.summary.total_tokens = self.summary.total_tokens; + state.summary.total_cache_creation_tokens = self.summary.total_cache_creation_tokens; + state.summary.total_cache_read_tokens = self.summary.total_cache_read_tokens; + state.summary.last_query_total_tokens = self.summary.last_query_total_tokens; + state.core.total_input_tokens = self.summary.total_input_tokens; + state.core.total_output_tokens = self.summary.total_output_tokens; + state.core.total_tokens = self.summary.total_tokens; + state.core.total_cache_creation_tokens = self.summary.total_cache_creation_tokens; + state.core.total_cache_read_tokens = self.summary.total_cache_read_tokens; + if let Some(active_turn) = state.active_turn.as_mut() + && active_turn.turn_id == self.turn_id + { + active_turn.usage = self.active_turn_usage; + } } } diff --git a/crates/server/src/runtime/subagent_usage.rs b/crates/server/src/runtime/subagent_usage.rs index bab2097f..9a281a0b 100644 --- a/crates/server/src/runtime/subagent_usage.rs +++ b/crates/server/src/runtime/subagent_usage.rs @@ -90,11 +90,29 @@ pub(crate) struct ParentUsageSnapshot { pub(super) context_window: Option, } +/// How a usage sample should update turn accounting. +/// +/// Provider streams emit partial [`UsageDelta`](devo_core::QueryEvent::UsageDelta) +/// values (often `output_tokens = 0` at message start) and a final +/// [`Usage`](devo_core::QueryEvent::Usage) per model call. Tool-use turns run +/// multiple model calls; completed legs must accumulate while in-flight samples +/// may only replace the current leg. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum UsageUpdateKind { + /// Replace the current in-flight leg only (streaming deltas). + InFlight, + /// Add one completed model-call leg and clear in-flight state. + CompletedLeg, +} + #[derive(Debug, Clone, Copy)] struct ParentTurnUsage { base_session_totals: UsageTotals, base_child_totals: UsageTotals, + /// Sum of completed model-call legs in this parent turn. parent_turn_usage: UsageTotals, + /// Latest partial usage for the model call still in progress. + inflight_usage: UsageTotals, context_window: Option, } @@ -108,7 +126,10 @@ struct ChildUsageOwner { struct ChildTurnUsage { parent_session_id: SessionId, parent_turn_id: Option, + /// Sum of completed model-call legs in this child turn. usage: UsageTotals, + /// Latest partial usage for the model call still in progress. + inflight_usage: UsageTotals, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -148,6 +169,7 @@ impl SubagentUsageState { base_session_totals, base_child_totals, parent_turn_usage: UsageTotals::default(), + inflight_usage: UsageTotals::default(), context_window, }); } @@ -173,13 +195,19 @@ impl SubagentUsageState { turn_id: TurnId, usage: UsageTotals, context_window: Option, + kind: UsageUpdateKind, ) -> Option { let key = ParentTurnKey { session_id, turn_id, }; let state = self.parent_turns.get_mut(&key)?; - state.parent_turn_usage = usage; + apply_usage_update( + &mut state.parent_turn_usage, + &mut state.inflight_usage, + usage, + kind, + ); if context_window.is_some() { state.context_window = context_window; } @@ -191,25 +219,30 @@ impl SubagentUsageState { child_session_id: SessionId, child_turn_id: TurnId, usage: UsageTotals, + kind: UsageUpdateKind, ) -> Option { let key = ChildTurnKey { session_id: child_session_id, turn_id: child_turn_id, }; let owner = if let Some(entry) = self.child_turns.get_mut(&key) { - entry.usage = usage; + apply_usage_update(&mut entry.usage, &mut entry.inflight_usage, usage, kind); ChildUsageOwner { parent_session_id: entry.parent_session_id, parent_turn_id: entry.parent_turn_id, } } else { let owner = *self.child_owners.get(&child_session_id)?; + let mut committed = UsageTotals::default(); + let mut inflight = UsageTotals::default(); + apply_usage_update(&mut committed, &mut inflight, usage, kind); self.child_turns.insert( key, ChildTurnUsage { parent_session_id: owner.parent_session_id, parent_turn_id: owner.parent_turn_id, - usage, + usage: committed, + inflight_usage: inflight, }, ); owner @@ -219,6 +252,25 @@ impl SubagentUsageState { .and_then(|turn_id| self.snapshot_for_parent_turn(owner.parent_session_id, turn_id)) } + /// Fold any remaining in-flight child usage into committed totals (e.g. on + /// interrupt) without recording an extra model-call leg. + pub(super) fn commit_child_inflight_usage( + &mut self, + child_session_id: SessionId, + child_turn_id: TurnId, + ) -> Option { + let key = ChildTurnKey { + session_id: child_session_id, + turn_id: child_turn_id, + }; + let entry = self.child_turns.get_mut(&key)?; + entry.usage = entry.usage.add(entry.inflight_usage); + entry.inflight_usage = UsageTotals::default(); + let parent_session_id = entry.parent_session_id; + let parent_turn_id = entry.parent_turn_id?; + self.snapshot_for_parent_turn(parent_session_id, parent_turn_id) + } + pub(super) fn snapshot_for_parent_turn( &self, session_id: SessionId, @@ -228,14 +280,15 @@ impl SubagentUsageState { session_id, turn_id, })?; + let parent_turn_usage = state.parent_turn_usage.add(state.inflight_usage); let child_turn_usage = self.child_totals_for_parent_turn(session_id, turn_id); let child_session_delta = self .child_totals_for_parent_session(session_id) .saturating_sub(state.base_child_totals); - let turn_usage = state.parent_turn_usage.add(child_turn_usage); + let turn_usage = parent_turn_usage.add(child_turn_usage); let session_totals = state .base_session_totals - .add(state.parent_turn_usage) + .add(parent_turn_usage) .add(child_session_delta); Some(ParentUsageSnapshot { session_id, @@ -251,7 +304,7 @@ impl SubagentUsageState { .values() .filter(|entry| entry.parent_session_id == parent_session_id) .fold(UsageTotals::default(), |total, entry| { - total.add(entry.usage) + total.add(entry.usage.add(entry.inflight_usage)) }) } @@ -267,11 +320,28 @@ impl SubagentUsageState { && entry.parent_turn_id == Some(parent_turn_id) }) .fold(UsageTotals::default(), |total, entry| { - total.add(entry.usage) + total.add(entry.usage.add(entry.inflight_usage)) }) } } +fn apply_usage_update( + committed: &mut UsageTotals, + inflight: &mut UsageTotals, + usage: UsageTotals, + kind: UsageUpdateKind, +) { + match kind { + UsageUpdateKind::InFlight => { + *inflight = usage; + } + UsageUpdateKind::CompletedLeg => { + *committed = committed.add(usage); + *inflight = UsageTotals::default(); + } + } +} + impl ServerRuntime { pub(super) async fn begin_parent_usage_turn( &self, @@ -316,6 +386,7 @@ impl ServerRuntime { turn_id: TurnId, usage: TurnUsage, context_window: Option, + kind: UsageUpdateKind, ) -> Option { self.begin_parent_usage_turn(session_id, turn_id, context_window) .await; @@ -326,6 +397,7 @@ impl ServerRuntime { turn_id, UsageTotals::from_turn_usage(&usage), context_window, + kind, ) }?; self.apply_parent_usage_snapshot(snapshot).await; @@ -337,6 +409,7 @@ impl ServerRuntime { child_session_id: SessionId, child_turn_id: TurnId, usage: TurnUsage, + kind: UsageUpdateKind, ) -> Option { let snapshot = { let mut usage_state = self.subagent_usage.lock().await; @@ -344,12 +417,26 @@ impl ServerRuntime { child_session_id, child_turn_id, UsageTotals::from_turn_usage(&usage), + kind, ) }?; self.apply_parent_usage_snapshot(snapshot).await; Some(snapshot) } + pub(super) async fn commit_subagent_inflight_usage( + &self, + child_session_id: SessionId, + child_turn_id: TurnId, + ) -> Option { + let snapshot = { + let mut usage_state = self.subagent_usage.lock().await; + usage_state.commit_child_inflight_usage(child_session_id, child_turn_id) + }?; + self.apply_parent_usage_snapshot(snapshot).await; + Some(snapshot) + } + pub(super) async fn parent_usage_snapshot( &self, session_id: SessionId, @@ -362,8 +449,32 @@ impl ServerRuntime { } async fn apply_parent_usage_snapshot(&self, snapshot: ParentUsageSnapshot) { - if let Some(session_handle) = self.session(snapshot.session_id).await { - session_handle.apply_parent_usage_snapshot(snapshot).await; + // The turn event stream runs while the session actor is blocked inside + // `execute_turn_in_actor` awaiting that same stream. Any mailbox send to + // `snapshot.session_id` here can fill the actor mailbox and then block + // forever on `send().await`, which stops the event stream from `recv`ing, + // fills the event channel, and wedges the whole turn. Prefer the in-flight + // turn inline state whenever it is registered. + let applied_inline = if let Some(stream) = self.active_stream_state(snapshot.session_id).await + { + let mut stream = stream.lock().await; + if let Some(inline) = stream.turn_inline.as_mut() { + snapshot.apply_to_summary(&mut inline.summary); + inline.hook_context.summary = inline.summary.clone(); + if inline.turn_id == snapshot.turn_id { + inline.active_turn_usage = Some(snapshot.turn_usage.to_turn_usage()); + } + true + } else { + false + } + } else { + false + }; + if !applied_inline { + if let Some(session_handle) = self.session(snapshot.session_id).await { + session_handle.apply_parent_usage_snapshot(snapshot).await; + } } self.broadcast_event(ServerEvent::TurnUsageUpdated( snapshot.to_turn_usage_updated_payload(), @@ -391,12 +502,7 @@ impl ParentUsageSnapshot { self, state: &mut crate::runtime::session_actor::state::SessionActorState, ) { - state.summary.total_input_tokens = self.session_totals.input_tokens; - state.summary.total_output_tokens = self.session_totals.output_tokens; - state.summary.total_tokens = self.session_totals.total_tokens; - state.summary.total_cache_creation_tokens = self.session_totals.cache_creation_input_tokens; - state.summary.total_cache_read_tokens = self.session_totals.cache_read_input_tokens; - state.summary.last_query_total_tokens = self.turn_usage.total_tokens; + self.apply_to_summary(&mut state.summary); if let Some(active_turn) = state.active_turn.as_mut() && active_turn.turn_id == self.turn_id { @@ -408,6 +514,15 @@ impl ParentUsageSnapshot { state.core.total_cache_creation_tokens = self.session_totals.cache_creation_input_tokens; state.core.total_cache_read_tokens = self.session_totals.cache_read_input_tokens; } + + fn apply_to_summary(self, summary: &mut crate::session::SessionMetadata) { + summary.total_input_tokens = self.session_totals.input_tokens; + summary.total_output_tokens = self.session_totals.output_tokens; + summary.total_tokens = self.session_totals.total_tokens; + summary.total_cache_creation_tokens = self.session_totals.cache_creation_input_tokens; + summary.total_cache_read_tokens = self.session_totals.cache_read_input_tokens; + summary.last_query_total_tokens = self.turn_usage.total_tokens; + } } fn saturating_u32(value: usize) -> u32 { @@ -436,7 +551,7 @@ mod tests { } #[test] - fn child_usage_replaces_latest_value_without_double_counting() { + fn child_usage_replaces_latest_inflight_without_double_counting() { let mut state = SubagentUsageState::default(); let parent_session_id = SessionId::new(); let parent_turn_id = TurnId::new(); @@ -452,24 +567,88 @@ mod tests { state.register_child_owner(parent_session_id, child_session_id, Some(parent_turn_id)); let parent_snapshot = state - .record_parent_turn_usage(parent_session_id, parent_turn_id, totals(8, 2), None) + .record_parent_turn_usage( + parent_session_id, + parent_turn_id, + totals(8, 2), + None, + UsageUpdateKind::CompletedLeg, + ) .expect("parent snapshot"); assert_eq!(parent_snapshot.turn_usage, totals(8, 2)); assert_eq!(parent_snapshot.session_totals, totals(108, 12)); let child_snapshot = state - .record_child_turn_usage(child_session_id, child_turn_id, totals(20, 5)) + .record_child_turn_usage( + child_session_id, + child_turn_id, + totals(20, 5), + UsageUpdateKind::InFlight, + ) .expect("child snapshot"); assert_eq!(child_snapshot.turn_usage, totals(28, 7)); assert_eq!(child_snapshot.session_totals, totals(128, 17)); let updated_child_snapshot = state - .record_child_turn_usage(child_session_id, child_turn_id, totals(25, 6)) + .record_child_turn_usage( + child_session_id, + child_turn_id, + totals(25, 6), + UsageUpdateKind::InFlight, + ) .expect("updated child snapshot"); assert_eq!(updated_child_snapshot.turn_usage, totals(33, 8)); assert_eq!(updated_child_snapshot.session_totals, totals(133, 18)); } + #[test] + fn tool_use_legs_accumulate_and_inflight_output_zero_does_not_reset_totals() { + let mut state = SubagentUsageState::default(); + let parent_session_id = SessionId::new(); + let parent_turn_id = TurnId::new(); + + state.begin_parent_turn(parent_session_id, parent_turn_id, totals(100, 10), None); + + // First model call completes with output. + let after_first_leg = state + .record_parent_turn_usage( + parent_session_id, + parent_turn_id, + totals(600, 50), + None, + UsageUpdateKind::CompletedLeg, + ) + .expect("first leg"); + assert_eq!(after_first_leg.turn_usage, totals(600, 50)); + assert_eq!(after_first_leg.session_totals, totals(700, 60)); + + // Second model call starts (typical UsageDelta with output_tokens = 0). + let after_second_start = state + .record_parent_turn_usage( + parent_session_id, + parent_turn_id, + totals(700, 0), + None, + UsageUpdateKind::InFlight, + ) + .expect("second leg start"); + assert_eq!(after_second_start.turn_usage, totals(1300, 50)); + assert_eq!(after_second_start.session_totals, totals(1400, 60)); + + // Second model call completes. + let after_second_leg = state + .record_parent_turn_usage( + parent_session_id, + parent_turn_id, + totals(700, 30), + None, + UsageUpdateKind::CompletedLeg, + ) + .expect("second leg complete"); + assert_eq!(after_second_leg.turn_usage, totals(1300, 80)); + assert_eq!(after_second_leg.session_totals, totals(1400, 90)); + } + #[test] fn next_parent_turn_base_includes_previous_child_usage() { let mut state = SubagentUsageState::default(); @@ -491,7 +670,12 @@ mod tests { child_session_id, Some(first_parent_turn_id), ); - state.record_child_turn_usage(child_session_id, first_child_turn_id, totals(20, 5)); + state.record_child_turn_usage( + child_session_id, + first_child_turn_id, + totals(20, 5), + UsageUpdateKind::CompletedLeg, + ); state.begin_parent_turn( parent_session_id, @@ -505,7 +689,12 @@ mod tests { Some(second_parent_turn_id), ); let snapshot = state - .record_child_turn_usage(child_session_id, second_child_turn_id, totals(7, 3)) + .record_child_turn_usage( + child_session_id, + second_child_turn_id, + totals(7, 3), + UsageUpdateKind::CompletedLeg, + ) .expect("second turn child snapshot"); assert_eq!(snapshot.turn_usage, totals(7, 3)); @@ -538,7 +727,12 @@ mod tests { child_session_id, Some(first_parent_turn_id), ); - state.record_child_turn_usage(child_session_id, child_turn_id, totals(20, 5)); + state.record_child_turn_usage( + child_session_id, + child_turn_id, + totals(20, 5), + UsageUpdateKind::InFlight, + ); state.register_child_owner( parent_session_id, @@ -546,7 +740,12 @@ mod tests { Some(second_parent_turn_id), ); let snapshot = state - .record_child_turn_usage(child_session_id, child_turn_id, totals(25, 6)) + .record_child_turn_usage( + child_session_id, + child_turn_id, + totals(25, 6), + UsageUpdateKind::InFlight, + ) .expect("updated child snapshot"); assert_eq!(snapshot.turn_id, first_parent_turn_id); diff --git a/crates/server/src/runtime/turn_exec/event_stream.rs b/crates/server/src/runtime/turn_exec/event_stream.rs index 309b2236..a6e90ebe 100644 --- a/crates/server/src/runtime/turn_exec/event_stream.rs +++ b/crates/server/src/runtime/turn_exec/event_stream.rs @@ -38,10 +38,16 @@ pub(super) fn enqueue_query_event( match event_tx.try_send(event) { Ok(()) => {} Err(tokio::sync::mpsc::error::TrySendError::Full(event)) => { - let event_tx = event_tx.clone(); - tokio::spawn(async move { - let _ = event_tx.send(event).await; - }); + // Never spawn unbounded waiters here: if the event stream stalls, + // those tasks park forever on `send().await`, keep sender clones + // alive, and prevent the stream from ever observing channel close. + // Dropping under backpressure preserves liveness; the stream still + // receives later events once it resumes. + tracing::warn!( + capacity = QUERY_EVENT_CHANNEL_CAPACITY, + event_kind = query_event_trace_kind(&event), + "dropping query event because the turn event channel is full" + ); } Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {} } @@ -224,16 +230,43 @@ pub(crate) fn spawn_turn_event_stream( ) .await; } - devo_core::QueryEvent::UsageDelta { usage } - | devo_core::QueryEvent::Usage { usage } => { + devo_core::QueryEvent::UsageDelta { usage } => { + let turn_usage = devo_core::TurnUsage::from_usage(&usage); + latest_usage = Some(turn_usage.clone()); + let kind = super::super::subagent_usage::UsageUpdateKind::InFlight; + if usage_parent_session_id.is_some() { + let _ = runtime + .publish_subagent_turn_usage( + session_id, + turn_for_events.turn_id, + turn_usage, + kind, + ) + .await; + } else if let Some(snapshot) = runtime + .publish_parent_turn_usage( + session_id, + turn_for_events.turn_id, + turn_usage, + usage_context_window, + kind, + ) + .await + { + latest_usage = Some(snapshot.turn_usage.to_turn_usage()); + } + } + devo_core::QueryEvent::Usage { usage } => { let turn_usage = devo_core::TurnUsage::from_usage(&usage); latest_usage = Some(turn_usage.clone()); + let kind = super::super::subagent_usage::UsageUpdateKind::CompletedLeg; if usage_parent_session_id.is_some() { let _ = runtime .publish_subagent_turn_usage( session_id, turn_for_events.turn_id, turn_usage, + kind, ) .await; } else if let Some(snapshot) = runtime @@ -242,6 +275,7 @@ pub(crate) fn spawn_turn_event_stream( turn_for_events.turn_id, turn_usage, usage_context_window, + kind, ) .await { diff --git a/crates/server/src/runtime/turn_exec/finalize.rs b/crates/server/src/runtime/turn_exec/finalize.rs index 5bcb1b31..8a04c26a 100644 --- a/crates/server/src/runtime/turn_exec/finalize.rs +++ b/crates/server/src/runtime/turn_exec/finalize.rs @@ -44,10 +44,11 @@ impl ServerRuntime { .as_ref() .and_then(|summary| summary.latest_usage.clone()); let terminal_stop_reason = event_summary.and_then(|summary| summary.stop_reason); - if usage_parent_session_id.is_some() - && let Some(usage) = latest_usage.clone() - { - self.publish_subagent_turn_usage(session_id, turn.turn_id, usage) + if usage_parent_session_id.is_some() { + // Completed legs were already accumulated by the event stream. + // Only fold any trailing in-flight delta (e.g. interrupted mid-stream). + let _ = self + .commit_subagent_inflight_usage(session_id, turn.turn_id) .await; } else if usage_parent_session_id.is_none() && let Some(snapshot) = self.parent_usage_snapshot(session_id, turn.turn_id).await From 5a9ccb218fd162a771fea2916e4b9e76a35eb8f4 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Fri, 3 Jul 2026 23:41:04 -1000 Subject: [PATCH 13/16] fix: prevent event-stream deadlock cascade with non-blocking outbound and agent-delta fast path The turn event stream is a single point of failure: when broadcast_event blocks on outbound capacity, the event channel fills, enqueue_query_event spawns unbounded overflow tasks whose sender clones prevent channel close, and the session actor hangs forever on event_task.await. - outbound: add enqueue_outbound_notification with 200ms max wait before dropping, so streaming notifications never stall the event task - connection: add broadcast_streaming_agent_message_delta fast path that skips the full broadcast_event (no registry scans, no buffer locks) - subagent: change OutputBuffer::push to try_push_text_delta so wait_agent's lock does not gate token streaming to the TUI - items/handle: add non-blocking try_touch_last_activity + fallback item-seq counter to avoid mailbox round-trips from event-stream code - agent_message delta handlers use the new fast path; record_subagent_ output_event is moved after outbound delivery --- crates/arg0/src/lib.rs | 4 +- crates/client/src/client_core.rs | 48 +++++++--- crates/client/src/protocol_trace.rs | 15 ++- crates/client/src/stdio.rs | 49 ++++++---- crates/code-search/src/service.rs | 6 +- crates/core/src/logging.rs | 4 +- crates/server/src/runtime.rs | 14 +-- crates/server/src/runtime/agents.rs | 29 +++--- crates/server/src/runtime/connection.rs | 70 ++++++++++++-- crates/server/src/runtime/goal_accounting.rs | 3 +- .../src/runtime/handlers/acp/session.rs | 10 +- crates/server/src/runtime/hooks.rs | 4 +- crates/server/src/runtime/items.rs | 21 ++++ crates/server/src/runtime/outbound.rs | 55 ++++++++++- .../src/runtime/session_actor/handle.rs | 18 ++++ .../server/src/runtime/session_actor/turn.rs | 12 ++- crates/server/src/runtime/subagent_usage.rs | 66 ++++++++++--- .../src/runtime/turn_exec/event_stream.rs | 47 +++++---- .../src/runtime/turn_exec/item_stream.rs | 35 +++---- crates/server/src/runtime/turn_exec/query.rs | 5 +- crates/server/src/runtime/turn_reservation.rs | 6 +- crates/server/src/subagent.rs | 49 ++++++++-- crates/server/src/transport.rs | 96 +++++++++---------- crates/server/tests/acp_available_commands.rs | 2 +- 24 files changed, 468 insertions(+), 200 deletions(-) diff --git a/crates/arg0/src/lib.rs b/crates/arg0/src/lib.rs index 0352afe2..5b9b416b 100644 --- a/crates/arg0/src/lib.rs +++ b/crates/arg0/src/lib.rs @@ -256,7 +256,9 @@ fn install_server_logging_for_dispatch() -> Option { async fn run_server_dispatch() { let args = parse_server_dispatch_args(); let _logging = install_server_logging_for_dispatch(); - if let Err(err) = devo_server::run_server_process(args, devo_server::ServerProcessRunOptions::default()).await { + if let Err(err) = + devo_server::run_server_process(args, devo_server::ServerProcessRunOptions::default()).await + { eprintln!("server error: {err}"); std::process::exit(1); } diff --git a/crates/client/src/client_core.rs b/crates/client/src/client_core.rs index f5c9c8b3..166432de 100644 --- a/crates/client/src/client_core.rs +++ b/crates/client/src/client_core.rs @@ -144,10 +144,7 @@ impl ServerClientCore { } #[cfg(test)] - pub(crate) fn set_agent_capabilities_for_test( - &mut self, - capabilities: AcpAgentCapabilities, - ) { + pub(crate) fn set_agent_capabilities_for_test(&mut self, capabilities: AcpAgentCapabilities) { self.acp_agent_capabilities = Some(capabilities); } @@ -434,11 +431,17 @@ impl ServerClientCore { self.request_devo("agent/list", params).await } - pub(crate) async fn agent_spawn(&mut self, params: SpawnAgentParams) -> Result { + pub(crate) async fn agent_spawn( + &mut self, + params: SpawnAgentParams, + ) -> Result { self.request_devo("agent/spawn", params).await } - pub(crate) async fn agent_close(&mut self, params: CloseAgentParams) -> Result { + pub(crate) async fn agent_close( + &mut self, + params: CloseAgentParams, + ) -> Result { self.request_devo("agent/close", params).await } @@ -460,7 +463,8 @@ impl ServerClientCore { &mut self, params: SessionPermissionsUpdateParams, ) -> Result { - self.request_devo("session/permissions/update", params).await + self.request_devo("session/permissions/update", params) + .await } pub(crate) async fn session_compact( @@ -470,7 +474,10 @@ impl ServerClientCore { self.request_devo("session/compact", params).await } - pub(crate) async fn goal_create(&mut self, params: GoalCreateParams) -> Result { + pub(crate) async fn goal_create( + &mut self, + params: GoalCreateParams, + ) -> Result { self.request_devo("goal/create", params).await } @@ -478,11 +485,17 @@ impl ServerClientCore { self.request_devo("goal/set", params).await } - pub(crate) async fn goal_status(&mut self, params: GoalStatusParams) -> Result { + pub(crate) async fn goal_status( + &mut self, + params: GoalStatusParams, + ) -> Result { self.request_devo("goal/status", params).await } - pub(crate) async fn goal_pause(&mut self, params: GoalSetStatusParams) -> Result { + pub(crate) async fn goal_pause( + &mut self, + params: GoalSetStatusParams, + ) -> Result { self.request_devo("goal/pause", params).await } @@ -504,7 +517,10 @@ impl ServerClientCore { self.request_devo("goal/clear", params).await } - pub(crate) async fn session_fork(&mut self, params: SessionForkParams) -> Result { + pub(crate) async fn session_fork( + &mut self, + params: SessionForkParams, + ) -> Result { self.request_devo("session/fork", params).await } @@ -540,7 +556,10 @@ impl ServerClientCore { self.request_devo("model/catalog", params).await } - pub(crate) async fn model_saved(&mut self, params: ModelSavedParams) -> Result { + pub(crate) async fn model_saved( + &mut self, + params: ModelSavedParams, + ) -> Result { self.request_devo("model/saved", params).await } @@ -565,7 +584,10 @@ impl ServerClientCore { self.request_devo("provider/validate", params).await } - pub(crate) async fn command_exec(&mut self, params: CommandExecParams) -> Result { + pub(crate) async fn command_exec( + &mut self, + params: CommandExecParams, + ) -> Result { self.request_devo("command/exec", params).await } diff --git a/crates/client/src/protocol_trace.rs b/crates/client/src/protocol_trace.rs index 317ddb81..c8c18936 100644 --- a/crates/client/src/protocol_trace.rs +++ b/crates/client/src/protocol_trace.rs @@ -167,9 +167,18 @@ mod tests { let file = File::create(&path).expect("create trace file"); let trace = ProtocolTrace::with_file(file); - trace.record(TraceDirection::Out, r#"{"jsonrpc":"2.0","id":1,"method":"initialize"}"#); - trace.record(TraceDirection::In, r#"{"jsonrpc":"2.0","id":1,"result":{}}"#); - trace.record(TraceDirection::Out, r#"{"jsonrpc":"2.0","id":2,"method":"session/new"}"#); + trace.record( + TraceDirection::Out, + r#"{"jsonrpc":"2.0","id":1,"method":"initialize"}"#, + ); + trace.record( + TraceDirection::In, + r#"{"jsonrpc":"2.0","id":1,"result":{}}"#, + ); + trace.record( + TraceDirection::Out, + r#"{"jsonrpc":"2.0","id":2,"method":"session/new"}"#, + ); let reader = BufReader::new(File::open(&path).expect("open trace file")); let records: Vec = reader diff --git a/crates/client/src/stdio.rs b/crates/client/src/stdio.rs index f7e42d56..4b502ffe 100644 --- a/crates/client/src/stdio.rs +++ b/crates/client/src/stdio.rs @@ -88,11 +88,8 @@ impl StdioServerClient { let trace = ProtocolTrace::from_env(); let writer_trace = trace.clone(); - let writer_task = tokio::spawn(run_stdin_writer( - Arc::clone(&stdin), - write_rx, - writer_trace, - )); + let writer_task = + tokio::spawn(run_stdin_writer(Arc::clone(&stdin), write_rx, writer_trace)); let reader_trace = trace; let reader_task = tokio::spawn(async move { let mut lines = BufReader::new(stdout).lines(); @@ -456,7 +453,10 @@ async fn write_ndjson_to_stdin( } line.push(b'\n'); let mut stdin = stdin.lock().await; - stdin.write_all(&line).await.context("write client payload")?; + stdin + .write_all(&line) + .await + .context("write client payload")?; stdin.flush().await.context("flush client payload")?; Ok(()) } @@ -472,20 +472,22 @@ async fn run_stderr_reader(mut lines: tokio::io::Lines>) #[cfg(test)] mod tests { + use super::*; + use crate::ACP_PROMPT_COMPLETED_NOTIFICATION_METHOD; + use crate::ACP_PROMPT_STARTED_NOTIFICATION_METHOD; + use crate::client_core::ClientWriteMessage; + use crate::client_core::ClientWriter; + use crate::client_core::PendingResponses; + use crate::client_core::ServerClientCore; + use chrono::Utc; + use pretty_assertions::assert_eq; use std::sync::Arc; use tokio::io::AsyncBufRead; use tokio::io::BufReader; use tokio::sync::Mutex; use tokio::sync::oneshot; - use tokio::time::timeout; use tokio::time::Duration; - use super::*; - use chrono::Utc; - use pretty_assertions::assert_eq; - use crate::client_core::ClientWriteMessage; - use crate::client_core::ClientWriter; - use crate::client_core::PendingResponses; - use crate::client_core::ServerClientCore; + use tokio::time::timeout; fn default_test_client_capabilities() -> devo_protocol::AcpClientCapabilities { devo_protocol::AcpClientCapabilities { @@ -542,7 +544,6 @@ mod tests { (client, pending) } - #[tokio::test] async fn initialize_uses_configured_client_capabilities() { let (child, stdin, stdout) = request_capture_child_for_turn_start_test().await; @@ -555,7 +556,8 @@ mod tests { terminal: true, meta: None, }; - let (mut client, pending) = spawn_test_stdio_client(child, stdin, client_capabilities.clone()).await; + let (mut client, pending) = + spawn_test_stdio_client(child, stdin, client_capabilities.clone()).await; let mut stdout_lines = BufReader::new(stdout).lines(); let expected_capabilities = serde_json::to_value(&client_capabilities).expect("serialize client capabilities"); @@ -595,7 +597,8 @@ mod tests { #[tokio::test] async fn session_start_accepts_standard_acp_response_without_devo_metadata() { let (child, stdin, stdout) = request_capture_child_for_turn_start_test().await; - let (mut client, pending) = spawn_test_stdio_client(child, stdin, default_test_client_capabilities()).await; + let (mut client, pending) = + spawn_test_stdio_client(child, stdin, default_test_client_capabilities()).await; let cwd = std::env::current_dir().expect("current dir"); let additional_directory = cwd.join("shared"); let session_id = devo_protocol::SessionId::new(); @@ -670,7 +673,8 @@ mod tests { #[tokio::test] async fn session_list_accepts_standard_acp_sessions_without_devo_metadata() { let (child, stdin, stdout) = request_capture_child_for_turn_start_test().await; - let (mut client, pending) = spawn_test_stdio_client(child, stdin, default_test_client_capabilities()).await; + let (mut client, pending) = + spawn_test_stdio_client(child, stdin, default_test_client_capabilities()).await; client .core .set_agent_capabilities_for_test(test_agent_capabilities_with_session_list()); @@ -759,7 +763,8 @@ mod tests { #[tokio::test] async fn session_resume_accepts_standard_acp_response_without_devo_metadata() { let (child, stdin, stdout) = request_capture_child_for_turn_start_test().await; - let (mut client, pending) = spawn_test_stdio_client(child, stdin, default_test_client_capabilities()).await; + let (mut client, pending) = + spawn_test_stdio_client(child, stdin, default_test_client_capabilities()).await; client .core .set_agent_capabilities_for_test(test_agent_capabilities_with_session_list()); @@ -876,7 +881,8 @@ mod tests { #[tokio::test] async fn turn_start_sends_devo_extension_with_full_params() { let (child, stdin, stdout) = request_capture_child_for_turn_start_test().await; - let (mut client, pending) = spawn_test_stdio_client(child, stdin, default_test_client_capabilities()).await; + let (mut client, pending) = + spawn_test_stdio_client(child, stdin, default_test_client_capabilities()).await; let params = TurnStartParams { session_id: devo_protocol::SessionId::new(), input: vec![devo_protocol::InputItem::Text { @@ -943,7 +949,8 @@ mod tests { #[tokio::test] async fn turn_start_falls_back_to_acp_prompt_with_lifecycle_notifications() { let (child, stdin, stdout) = request_capture_child_for_turn_start_test().await; - let (mut client, pending) = spawn_test_stdio_client(child, stdin, default_test_client_capabilities()).await; + let (mut client, pending) = + spawn_test_stdio_client(child, stdin, default_test_client_capabilities()).await; let session_id = devo_protocol::SessionId::new(); let params = TurnStartParams { session_id, diff --git a/crates/code-search/src/service.rs b/crates/code-search/src/service.rs index 78ff48b3..0920be59 100644 --- a/crates/code-search/src/service.rs +++ b/crates/code-search/src/service.rs @@ -211,8 +211,10 @@ impl CodeSearchService { return false; } let cache_path = cache_file_path(&self.cache_dir, &root, content, self.provider.model_id()); - let has_valid_cache = load_payload(&cache_path) - .is_some_and(|c| c.payload.is_valid_for(&root, content, self.provider.model_id())); + let has_valid_cache = load_payload(&cache_path).is_some_and(|c| { + c.payload + .is_valid_for(&root, content, self.provider.model_id()) + }); !has_valid_cache } diff --git a/crates/core/src/logging.rs b/crates/core/src/logging.rs index 84bdc683..b6c1c798 100644 --- a/crates/core/src/logging.rs +++ b/crates/core/src/logging.rs @@ -87,7 +87,9 @@ impl LoggingBootstrap { Err(LoggingInitError::SubscriberAlreadyInstalled) => { // Another subscriber was already installed (e.g. by tokio-console). // This is not a fatal error — the process will run without file logging. - tracing::warn!("file logging skipped: a global tracing subscriber is already installed"); + tracing::warn!( + "file logging skipped: a global tracing subscriber is already installed" + ); } Err(other) => return Err(other), Ok(()) => {} diff --git a/crates/server/src/runtime.rs b/crates/server/src/runtime.rs index 582df17a..7a2b7221 100644 --- a/crates/server/src/runtime.rs +++ b/crates/server/src/runtime.rs @@ -152,7 +152,6 @@ mod agents; mod approval; mod command_exec; mod connection; -mod outbound; mod goal_accounting; mod goal_continuation; mod goal_handlers; @@ -161,6 +160,7 @@ mod hooks; mod items; mod lifecycle; mod model_api; +mod outbound; mod proposed_plan; mod provider_vendor_api; mod reference_search; @@ -186,18 +186,18 @@ mod turn_reservation; mod user_input; mod workspace_baseline; -pub use outbound::test_outbound_channel; -pub(crate) use outbound::log_outbound_frame; -pub(crate) use outbound::outbound_frame_to_value; -pub(crate) use outbound::OUTBOUND_CHANNEL_CAPACITY; -pub use outbound::OutboundFrame; -pub(crate) use outbound::enqueue_outbound; pub(crate) use connection::ConnectionRuntime; pub(crate) use connection::INBOUND_CONCURRENCY_LIMIT; pub use connection::IncomingResponse; pub use connection::PostResponseActions; pub(crate) use connection::SubscriptionFilter; pub(crate) use items::render_input_items; +pub(crate) use outbound::OUTBOUND_CHANNEL_CAPACITY; +pub use outbound::OutboundFrame; +pub(crate) use outbound::enqueue_outbound; +pub(crate) use outbound::log_outbound_frame; +pub(crate) use outbound::outbound_frame_to_value; +pub use outbound::test_outbound_channel; pub(crate) use research_tools::extract_written_file_path; pub(crate) use research_tools::is_write_tool_name; use session_actor::SessionHandle; diff --git a/crates/server/src/runtime/agents.rs b/crates/server/src/runtime/agents.rs index 08488174..b6fdd345 100644 --- a/crates/server/src/runtime/agents.rs +++ b/crates/server/src/runtime/agents.rs @@ -915,19 +915,22 @@ impl ServerRuntime { else { return; }; - self.output_buffer(parent_session_id) - .await - .push(devo_protocol::AgentOutputEvent { - sequence: 0, - child_session_id, - agent_path, - turn_id: payload.context.turn_id, - kind: devo_protocol::AgentOutputEventKind::AssistantMessage, - text: Some(payload.delta.clone()), - status: None, - created_at: Utc::now(), - }) - .await; + let buffer = self.output_buffer(parent_session_id).await; + let output_event = devo_protocol::AgentOutputEvent { + sequence: 0, + child_session_id, + agent_path, + turn_id: payload.context.turn_id, + kind: devo_protocol::AgentOutputEventKind::AssistantMessage, + text: Some(payload.delta.clone()), + status: None, + created_at: Utc::now(), + }; + // Streaming text must not block behind wait_agent's buffer lock. If the + // lock is busy, skip this delta for the wait buffer; the TUI already + // received it via outbound notifications, and a later delta/status will + // refresh the coalesced text. + let _ = buffer.try_push_text_delta(output_event); } async fn record_subagent_status_event( diff --git a/crates/server/src/runtime/connection.rs b/crates/server/src/runtime/connection.rs index a946ec31..5f543d33 100644 --- a/crates/server/src/runtime/connection.rs +++ b/crates/server/src/runtime/connection.rs @@ -26,6 +26,7 @@ use crate::devo_extension_inner_method; use super::outbound::OutboundFrame; use super::outbound::enqueue_outbound; +use super::outbound::enqueue_outbound_notification; pub(crate) const INBOUND_CONCURRENCY_LIMIT: usize = 64; @@ -599,6 +600,53 @@ impl ServerRuntime { .await; } + /// Hot path for child/parent assistant token streaming. + /// + /// Avoids per-token `child_parent_by_session` registry scans and never waits + /// on the wait_agent output buffer. Uses `active_turn_connections` to find + /// the owning stdio connection directly. + pub(super) async fn broadcast_streaming_agent_message_delta(&self, event: &ServerEvent) { + let ServerEvent::ItemDelta { + delta_kind: ItemDeltaKind::AgentMessageDelta, + payload, + } = event + else { + self.broadcast_event(event.clone()).await; + return; + }; + let session_id = payload.context.session_id; + let connection_id = { + let active_turn_connections = self.active_turn_connections.lock().await; + active_turn_connections.get(&session_id).copied() + }; + let Some(connection_id) = connection_id else { + self.broadcast_event(event.clone()).await; + return; + }; + let notification = { + let mut connections = self.connections.lock().await; + let Some(connection) = connections.get_mut(&connection_id) else { + return; + }; + let method = event.method_name(); + if connection.opt_out_notification_methods.contains(method) { + return; + } + let event_seq = connection.next_seq(); + let event = event.clone().with_seq(event_seq); + let (method, value) = acp_notification_from_server_event(method, &event); + Some(( + connection.outbound_tx.clone(), + OutboundFrame::notification(connection_id, method, event_seq, value), + )) + }; + if let Some((outbound_tx, frame)) = notification { + let _ = enqueue_outbound_notification(&outbound_tx, frame, "connection_notifications") + .await; + } + self.record_subagent_output_event(event).await; + } + pub(super) async fn emit_to_connection( &self, connection_id: u64, @@ -624,7 +672,8 @@ impl ServerRuntime { )) }; if let Some((outbound_tx, frame)) = notification { - let _ = enqueue_outbound(&outbound_tx, frame, "connection_notifications").await; + let _ = enqueue_outbound_notification(&outbound_tx, frame, "connection_notifications") + .await; } } @@ -633,7 +682,10 @@ impl ServerRuntime { self.account_goal_turn_completed(&payload.turn).await; } self.update_session_last_activity_from_event(&event).await; - self.record_subagent_output_event(&event).await; + // Deliver to the client first. wait_agent's output buffer must not gate + // token streaming: when supervisor is blocked in wait_agent and children + // contend on that buffer, TUI tokens would stall while the provider SSE + // (visible in Burp) keeps flowing into a full event channel. let method = event.method_name(); let session_id = event.session_id(); let child_parent_by_session = self.child_parent_by_session().await; @@ -665,8 +717,10 @@ impl ServerRuntime { .collect::>() }; for (outbound_tx, frame) in notifications { - let _ = enqueue_outbound(&outbound_tx, frame, "connection_notifications").await; + let _ = enqueue_outbound_notification(&outbound_tx, frame, "connection_notifications") + .await; } + self.record_subagent_output_event(&event).await; } async fn update_session_last_activity_from_event(&self, event: &ServerEvent) { @@ -680,8 +734,10 @@ impl ServerRuntime { return; } } + // Event broadcast runs on turn event streams; never block those tasks on + // a session actor mailbox that may be awaiting the same stream. if let Some(session_handle) = self.session(session_id).await { - session_handle.touch_last_activity().await; + let _ = session_handle.try_touch_last_activity(); } } @@ -777,11 +833,7 @@ impl ServerRuntime { request_id, rx, connection.outbound_tx.clone(), - OutboundFrame::client_request( - connection_id, - method.to_string(), - value, - ), + OutboundFrame::client_request(connection_id, method.to_string(), value), ) }; let mut pending_request = PendingClientRequestGuard::new( diff --git a/crates/server/src/runtime/goal_accounting.rs b/crates/server/src/runtime/goal_accounting.rs index 1b5d7553..1947f9b7 100644 --- a/crates/server/src/runtime/goal_accounting.rs +++ b/crates/server/src/runtime/goal_accounting.rs @@ -73,7 +73,8 @@ impl ServerRuntime { } async fn session_is_plan_mode(&self, session_id: SessionId) -> bool { - self.session_collaboration_mode(session_id).await == Some(devo_protocol::CollaborationMode::Plan) + self.session_collaboration_mode(session_id).await + == Some(devo_protocol::CollaborationMode::Plan) } } diff --git a/crates/server/src/runtime/handlers/acp/session.rs b/crates/server/src/runtime/handlers/acp/session.rs index d5afc37e..ef8a98e2 100644 --- a/crates/server/src/runtime/handlers/acp/session.rs +++ b/crates/server/src/runtime/handlers/acp/session.rs @@ -455,10 +455,7 @@ impl ServerRuntime { }; let mut parent_by_session = Vec::new(); for session_id in session_ids_in_runtime { - let parent_id = self - .session_parent_id_snapshot(session_id) - .await - .flatten(); + let parent_id = self.session_parent_id_snapshot(session_id).await.flatten(); parent_by_session.push((session_id, parent_id)); } parent_by_session.sort_by_key(|(session_id, _parent_id)| session_id.to_string()); @@ -480,10 +477,7 @@ impl ServerRuntime { session_ids } - async fn await_session_turn_interrupt_before_delete( - self: &Arc, - session_id: SessionId, - ) { + async fn await_session_turn_interrupt_before_delete(self: &Arc, session_id: SessionId) { let Some(turn_id) = self.runtime_active_turn_id(session_id).await else { return; }; diff --git a/crates/server/src/runtime/hooks.rs b/crates/server/src/runtime/hooks.rs index 2430c003..a6136d11 100644 --- a/crates/server/src/runtime/hooks.rs +++ b/crates/server/src/runtime/hooks.rs @@ -65,7 +65,9 @@ impl ServerRuntime { } let session_handle = self.sessions.lock().await.get(&session_id).cloned()?; let snapshot = session_handle.hook_context_snapshot().await?; - Some(Self::hook_runtime_context_from_snapshot(session_id, &snapshot)?) + Some(Self::hook_runtime_context_from_snapshot( + session_id, &snapshot, + )?) } fn hook_runtime_context_from_snapshot( diff --git a/crates/server/src/runtime/items.rs b/crates/server/src/runtime/items.rs index b54c28f4..c2075c8b 100644 --- a/crates/server/src/runtime/items.rs +++ b/crates/server/src/runtime/items.rs @@ -1,4 +1,6 @@ use std::borrow::Cow; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; use crate::titles::build_title_generation_request; use crate::titles::derive_provisional_title; @@ -6,6 +8,13 @@ use crate::titles::normalize_generated_title; use super::*; +/// Used only when a turn event stream is active but inline state is missing. +/// Avoids mailbox round-trips that deadlock the session actor. +fn next_fallback_item_seq() -> u64 { + static NEXT: AtomicU64 = AtomicU64::new(1 << 32); + NEXT.fetch_add(1, Ordering::Relaxed) +} + impl ServerRuntime { pub(super) async fn maybe_start_title_generation_from_user_input( self: &Arc, @@ -391,6 +400,15 @@ impl ServerRuntime { } return; } + // Active stream is registered but inline state is missing. The session + // actor is not polling its mailbox until the stream finishes, so we + // must not fall through to blocking actor commands. + tracing::warn!( + session_id = %session_id, + turn_id = %turn_id, + "persist_item skipped: active turn stream has no inline state" + ); + return; } let Some(session_handle) = self.session(session_id).await else { return; @@ -431,6 +449,9 @@ impl ServerRuntime { if let Some(inline) = stream.turn_inline.as_mut() { return inline.allocate_item_seq(); } + // Same deadlock constraint as persist_item: never wait on the actor + // mailbox while its turn event stream is registered. + return next_fallback_item_seq(); } if let Some(handle) = self.session(session_id).await && let Some(item_seq) = handle.allocate_item_seq().await diff --git a/crates/server/src/runtime/outbound.rs b/crates/server/src/runtime/outbound.rs index b6988433..c22bc14f 100644 --- a/crates/server/src/runtime/outbound.rs +++ b/crates/server/src/runtime/outbound.rs @@ -9,6 +9,10 @@ use crate::NotificationEnvelope; pub(crate) const OUTBOUND_CHANNEL_CAPACITY: usize = 4096; pub(crate) const OUTBOUND_BACKPRESSURE_LOG_THRESHOLD: Duration = Duration::from_millis(50); +/// Max time streaming notifications wait for outbound capacity before being +/// dropped. Event streams must not park forever on a slow client: parent+child +/// turns share one connection and can fill the queue quickly. +pub(crate) const OUTBOUND_NOTIFICATION_MAX_WAIT: Duration = Duration::from_millis(200); pub(crate) enum OutboundPayload { Notification { @@ -221,11 +225,60 @@ pub(crate) async fn enqueue_outbound( true } +/// Enqueue a streaming notification without risking an indefinite stall. +/// +/// Unlike [`enqueue_outbound`], this gives up after +/// [`OUTBOUND_NOTIFICATION_MAX_WAIT`] and drops the frame. Turn event streams +/// (including child agents) call into `broadcast_event` while the session actor +/// is awaiting that stream; blocking here recreates the mailbox-style hang on +/// the outbound path when the client drains stdout slowly. +pub(crate) async fn enqueue_outbound_notification( + tx: &mpsc::Sender, + frame: OutboundFrame, + queue: &'static str, +) -> bool { + let connection_id = frame.connection_id(); + match tx.try_reserve() { + Ok(permit) => { + permit.send(frame); + return true; + } + Err(tokio::sync::mpsc::error::TrySendError::Closed(())) => { + tracing::debug!(connection_id, queue, "outbound queue receiver dropped"); + return false; + } + Err(tokio::sync::mpsc::error::TrySendError::Full(())) => {} + } + + match tokio::time::timeout(OUTBOUND_NOTIFICATION_MAX_WAIT, tx.reserve()).await { + Ok(Ok(permit)) => { + permit.send(frame); + true + } + Ok(Err(_)) => { + tracing::debug!(connection_id, queue, "outbound queue receiver dropped"); + false + } + Err(_) => { + tracing::warn!( + connection_id, + queue, + max_wait_ms = OUTBOUND_NOTIFICATION_MAX_WAIT.as_millis(), + "dropping outbound notification under backpressure" + ); + false + } + } +} + /// Test helper: drains [`OutboundFrame`] values into serialized JSON values. #[doc(hidden)] pub fn test_outbound_channel( capacity: usize, -) -> (mpsc::Sender, mpsc::Receiver) { +) -> ( + mpsc::Sender, + mpsc::Receiver, +) { let (outbound_tx, mut outbound_rx) = mpsc::channel::(capacity); let (json_tx, json_rx) = mpsc::channel(capacity); tokio::spawn(async move { diff --git a/crates/server/src/runtime/session_actor/handle.rs b/crates/server/src/runtime/session_actor/handle.rs index 9f55804c..bcd6451c 100644 --- a/crates/server/src/runtime/session_actor/handle.rs +++ b/crates/server/src/runtime/session_actor/handle.rs @@ -57,6 +57,12 @@ impl SessionHandle { self.tx.send(command).await.is_ok() } + /// Non-blocking enqueue. Used by turn event streams so they never park on a + /// session actor that is itself waiting for that stream to finish. + fn try_send(&self, command: SessionCommand) -> bool { + self.tx.try_send(command).is_ok() + } + pub(crate) async fn execute_turn( &self, runtime: Arc, @@ -320,6 +326,7 @@ impl SessionHandle { let _ = self.send(SessionCommand::ResetTurnApprovalCache).await; } + #[allow(dead_code)] pub(crate) async fn touch_last_activity(&self) { let _ = self.send(SessionCommand::TouchLastActivity).await; } @@ -457,12 +464,23 @@ impl SessionHandle { .await; } + #[allow(dead_code)] pub(crate) async fn apply_parent_usage_snapshot(&self, snapshot: ParentUsageSnapshot) { let _ = self .send(SessionCommand::ApplyParentUsageSnapshot { snapshot }) .await; } + /// Best-effort usage apply for callers that must not block on the mailbox + /// (child/parent turn event streams). + pub(crate) fn try_apply_parent_usage_snapshot(&self, snapshot: ParentUsageSnapshot) -> bool { + self.try_send(SessionCommand::ApplyParentUsageSnapshot { snapshot }) + } + + pub(crate) fn try_touch_last_activity(&self) -> bool { + self.try_send(SessionCommand::TouchLastActivity) + } + pub(crate) async fn interrupt_active_turn(&self) -> Option> { let (reply_tx, reply_rx) = oneshot::channel(); if !self diff --git a/crates/server/src/runtime/session_actor/turn.rs b/crates/server/src/runtime/session_actor/turn.rs index 19612502..d0003e8e 100644 --- a/crates/server/src/runtime/session_actor/turn.rs +++ b/crates/server/src/runtime/session_actor/turn.rs @@ -51,9 +51,15 @@ pub(super) async fn execute_turn_in_actor( let event_tool_registry = runtime.tool_registry_for_actor_state(state); let usage_parent_session_id = state.parent_session_id(); let usage_context_window = Some(turn_config.model.context_window as u64); - runtime - .begin_parent_usage_turn(session_id, turn.turn_id, usage_context_window) - .await; + // Only root sessions own a parent-turn usage ledger. Child turns publish + // through `publish_subagent_turn_usage` into their parent's ledger; starting + // a ledger keyed by the child session id is incorrect and can strand usage + // updates on the wrong session. + if usage_parent_session_id.is_none() { + runtime + .begin_parent_usage_turn(session_id, turn.turn_id, usage_context_window) + .await; + } let stream = Arc::clone(&state.stream); let event_task = spawn_turn_event_stream( diff --git a/crates/server/src/runtime/subagent_usage.rs b/crates/server/src/runtime/subagent_usage.rs index 9a281a0b..d464eb10 100644 --- a/crates/server/src/runtime/subagent_usage.rs +++ b/crates/server/src/runtime/subagent_usage.rs @@ -152,6 +152,13 @@ pub(super) struct SubagentUsageState { } impl SubagentUsageState { + pub(super) fn parent_turn_started(&self, session_id: SessionId, turn_id: TurnId) -> bool { + self.parent_turns.contains_key(&ParentTurnKey { + session_id, + turn_id, + }) + } + pub(super) fn begin_parent_turn( &mut self, session_id: SessionId, @@ -349,6 +356,15 @@ impl ServerRuntime { turn_id: TurnId, context_window: Option, ) { + // Skip the summary mailbox/inline fetch when the ledger entry already + // exists. Child agents were calling this on every UsageDelta and serializing + // on the parent session for no benefit (`or_insert` is a no-op). + { + let usage_state = self.subagent_usage.lock().await; + if usage_state.parent_turn_started(session_id, turn_id) { + return; + } + } let base_session_totals = self .session_summary_snapshot(session_id) .await @@ -411,6 +427,23 @@ impl ServerRuntime { usage: TurnUsage, kind: UsageUpdateKind, ) -> Option { + // Child turns must roll usage into the parent turn ledger. Ensure that + // ledger exists even if the parent turn entry was never started (or the + // child outlived the parent's begin_parent_usage_turn call). + let parent_owner = { + let usage_state = self.subagent_usage.lock().await; + usage_state.child_owners.get(&child_session_id).copied() + }; + if let Some(owner) = parent_owner + && let Some(parent_turn_id) = owner.parent_turn_id + { + self.begin_parent_usage_turn( + owner.parent_session_id, + parent_turn_id, + /*context_window*/ None, + ) + .await; + } let snapshot = { let mut usage_state = self.subagent_usage.lock().await; usage_state.record_child_turn_usage( @@ -455,25 +488,30 @@ impl ServerRuntime { // forever on `send().await`, which stops the event stream from `recv`ing, // fills the event channel, and wedges the whole turn. Prefer the in-flight // turn inline state whenever it is registered. - let applied_inline = if let Some(stream) = self.active_stream_state(snapshot.session_id).await - { - let mut stream = stream.lock().await; - if let Some(inline) = stream.turn_inline.as_mut() { - snapshot.apply_to_summary(&mut inline.summary); - inline.hook_context.summary = inline.summary.clone(); - if inline.turn_id == snapshot.turn_id { - inline.active_turn_usage = Some(snapshot.turn_usage.to_turn_usage()); + let applied_inline = + if let Some(stream) = self.active_stream_state(snapshot.session_id).await { + let mut stream = stream.lock().await; + if let Some(inline) = stream.turn_inline.as_mut() { + snapshot.apply_to_summary(&mut inline.summary); + inline.hook_context.summary = inline.summary.clone(); + if inline.turn_id == snapshot.turn_id { + inline.active_turn_usage = Some(snapshot.turn_usage.to_turn_usage()); + } + true + } else { + false } - true } else { false - } - } else { - false - }; + }; if !applied_inline { + // Child agent event streams publish usage onto the parent session. + // The parent actor may be inside `execute_turn_in_actor` (or in the + // brief window after it unregisters its active stream but before it + // resumes polling). Blocking `send().await` here can fill the parent + // mailbox and deadlock the child stream, so only try-send. if let Some(session_handle) = self.session(snapshot.session_id).await { - session_handle.apply_parent_usage_snapshot(snapshot).await; + let _ = session_handle.try_apply_parent_usage_snapshot(snapshot); } } self.broadcast_event(ServerEvent::TurnUsageUpdated( diff --git a/crates/server/src/runtime/turn_exec/event_stream.rs b/crates/server/src/runtime/turn_exec/event_stream.rs index a6e90ebe..2a672927 100644 --- a/crates/server/src/runtime/turn_exec/event_stream.rs +++ b/crates/server/src/runtime/turn_exec/event_stream.rs @@ -23,29 +23,31 @@ use crate::runtime::session_actor::state::SessionStreamState; use crate::{ItemDeltaKind, ItemDeltaPayload, ServerEvent}; use tokio::sync::mpsc; -pub(crate) const QUERY_EVENT_CHANNEL_CAPACITY: usize = 1024; +pub(crate) const QUERY_EVENT_CHANNEL_CAPACITY: usize = 8192; -/// Enqueue a query event without blocking the caller when the channel is full. +/// Enqueue a query event into the turn event stream. /// -/// `run_turn_model_query` holds the session `core_session` mutex for the entire -/// `query()` call. Blocking on `event_tx.send().await` while the event task is -/// stalled on client-notification backpressure would freeze tool execution -/// (including `spawn_agent`) and can wedge the whole connection. -pub(super) fn enqueue_query_event( +/// Visible token events (`TextDelta` / `ReasoningDelta`) and `TurnComplete` +/// must not be dropped: when the channel is full we apply backpressure to the +/// provider reader with `send().await` so the TUI keeps receiving tokens. +/// Other events may be dropped under pressure so coordination cannot wedge +/// forever on a stalled consumer. +pub(super) async fn enqueue_query_event( event_tx: &mpsc::Sender, event: devo_core::QueryEvent, ) { + let kind = query_event_trace_kind(&event); + let must_deliver = matches!(kind, "text_delta" | "reasoning_delta" | "turn_complete"); + if must_deliver { + let _ = event_tx.send(event).await; + return; + } match event_tx.try_send(event) { Ok(()) => {} - Err(tokio::sync::mpsc::error::TrySendError::Full(event)) => { - // Never spawn unbounded waiters here: if the event stream stalls, - // those tasks park forever on `send().await`, keep sender clones - // alive, and prevent the stream from ever observing channel close. - // Dropping under backpressure preserves liveness; the stream still - // receives later events once it resumes. + Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { tracing::warn!( capacity = QUERY_EVENT_CHANNEL_CAPACITY, - event_kind = query_event_trace_kind(&event), + event_kind = kind, "dropping query event because the turn event channel is full" ); } @@ -303,6 +305,17 @@ pub(crate) fn spawn_turn_event_stream( &mut proposed_plan_leading_normal, ) .await; + { + let mut stream = event_stream.lock().await; + if let (Some(item_id), Some(item_seq)) = (assistant_item_id, assistant_item_seq) { + stream.deferred_assistant = + Some((item_id, item_seq, std::mem::take(&mut assistant_text))); + } + if let (Some(item_id), Some(item_seq)) = (reasoning_item_id, reasoning_item_seq) { + stream.deferred_reasoning = + Some((item_id, item_seq, std::mem::take(&mut reasoning_text))); + } + } complete_deferred_stream_items( &runtime, &event_stream, @@ -387,10 +400,8 @@ async fn handle_reasoning_delta( }, }) .await; - let _ = item_seq; - if let Ok(mut stream) = event_stream.try_lock() { - stream.deferred_reasoning = Some((item_id, item_seq, reasoning_text.clone())); - } + // Deferred reasoning text is written once when the event stream drains. + let _ = (event_stream, item_seq, item_id, reasoning_text); } async fn complete_open_reasoning_item( diff --git a/crates/server/src/runtime/turn_exec/item_stream.rs b/crates/server/src/runtime/turn_exec/item_stream.rs index 57c5c0fb..aa439071 100644 --- a/crates/server/src/runtime/turn_exec/item_stream.rs +++ b/crates/server/src/runtime/turn_exec/item_stream.rs @@ -166,25 +166,26 @@ pub(super) async fn push_assistant_text_delta( }; assistant_text.push_str(&text); *assistant_delta_seq = (*assistant_delta_seq).saturating_add(1); - runtime - .broadcast_event(ServerEvent::ItemDelta { - delta_kind: ItemDeltaKind::AgentMessageDelta, - payload: ItemDeltaPayload { - context: crate::EventContext { - session_id, - turn_id: Some(turn_id), - item_id: Some(item_id), - seq: 0, - }, - delta: text, - stream_index: None, - channel: None, + let event = ServerEvent::ItemDelta { + delta_kind: ItemDeltaKind::AgentMessageDelta, + payload: ItemDeltaPayload { + context: crate::EventContext { + session_id, + turn_id: Some(turn_id), + item_id: Some(item_id), + seq: 0, }, - }) + delta: text, + stream_index: None, + channel: None, + }, + }; + // Fast path: avoid per-token registry scans and wait_agent buffer contention. + runtime + .broadcast_streaming_agent_message_delta(&event) .await; - if let Ok(mut stream) = event_stream.try_lock() { - stream.deferred_assistant = Some((item_id, item_seq, assistant_text.clone())); - } + // Deferred assistant text is written once when the event stream drains. + let _ = (event_stream, item_seq); } #[allow(clippy::too_many_arguments)] diff --git a/crates/server/src/runtime/turn_exec/query.rs b/crates/server/src/runtime/turn_exec/query.rs index a459c5c1..e6817ad6 100644 --- a/crates/server/src/runtime/turn_exec/query.rs +++ b/crates/server/src/runtime/turn_exec/query.rs @@ -80,7 +80,7 @@ impl ServerRuntime { let callback: devo_core::EventCallback = std::sync::Arc::new(move |event: QueryEvent| { let event_callback_tx = event_callback_tx.clone(); Box::pin(async move { - enqueue_query_event(&event_callback_tx, event); + enqueue_query_event(&event_callback_tx, event).await; }) }); let tool_execution_start_tx = event_tx.clone(); @@ -159,7 +159,8 @@ impl ServerRuntime { enqueue_query_event( &tool_execution_start_tx, QueryEvent::ToolExecutionStart { id: call.id }, - ); + ) + .await; }) })), ..ToolExecutionOptions::default() diff --git a/crates/server/src/runtime/turn_reservation.rs b/crates/server/src/runtime/turn_reservation.rs index 9d984e37..7bccb783 100644 --- a/crates/server/src/runtime/turn_reservation.rs +++ b/crates/server/src/runtime/turn_reservation.rs @@ -52,11 +52,7 @@ impl ServerRuntime { } pub(super) async fn runtime_active_turn_id(&self, session_id: SessionId) -> Option { - self.active_turn_ids - .lock() - .await - .get(&session_id) - .copied() + self.active_turn_ids.lock().await.get(&session_id).copied() } pub(super) async fn clear_active_turn_runtime_handles(&self, session_id: SessionId) { diff --git a/crates/server/src/subagent.rs b/crates/server/src/subagent.rs index 3909fae1..786ab151 100644 --- a/crates/server/src/subagent.rs +++ b/crates/server/src/subagent.rs @@ -320,7 +320,13 @@ impl SubagentOutputBuffer { last.text.get_or_insert_with(String::new).push_str(&delta); } last.created_at = event.created_at; - return last.clone(); + // Release the lock before any clone/notify work. Callers currently + // discard the return value; avoid cloning the full accumulated + // assistant text on every token while holding the buffer mutex + // (multiple child event streams contend on this lock). + drop(inner); + self.notify.notify_waiters(); + return event; } } event.sequence = inner.next_sequence; @@ -331,6 +337,35 @@ impl SubagentOutputBuffer { event } + /// Non-blocking variant for streaming text. Returns `false` when the buffer + /// lock is busy so token broadcast is never stalled behind `wait_agent`. + pub fn try_push_text_delta(&self, mut event: AgentOutputEvent) -> bool { + let Ok(mut inner) = self.inner.try_lock() else { + return false; + }; + event.kind = AgentOutputEventKind::AssistantMessage; + if let Some(last) = inner.events.back_mut() + && last.kind == AgentOutputEventKind::AssistantMessage + && last.child_session_id == event.child_session_id + && last.turn_id == event.turn_id + && last.status.is_none() + { + if let Some(delta) = event.text.take() { + last.text.get_or_insert_with(String::new).push_str(&delta); + } + last.created_at = event.created_at; + drop(inner); + self.notify.notify_waiters(); + return true; + } + event.sequence = inner.next_sequence; + inner.next_sequence = inner.next_sequence.saturating_add(1); + inner.events.push_back(event); + drop(inner); + self.notify.notify_waiters(); + true + } + pub async fn wait_after( &self, after_sequence: u64, @@ -449,16 +484,16 @@ mod tests { ..base() }) .await; - let second = buffer + assert_eq!(first.sequence, 1); + assert_eq!(first.text.as_deref(), Some("alpha ")); + // Coalesced deltas no longer return a full snapshot (callers discard it); + // accumulated text is observed through wait_after / buffer reads. + let _ = buffer .push(AgentOutputEvent { text: Some("beta".into()), ..base() }) .await; - assert_eq!(first.sequence, 1); - assert_eq!(second.sequence, 1); - assert_eq!(second.text.as_deref(), Some("alpha beta")); - assert_eq!(second.kind, AgentOutputEventKind::AssistantMessage); let (events, next_sequence, timed_out) = buffer .wait_after(0, &[child], Duration::from_millis(1), None) @@ -466,7 +501,9 @@ mod tests { assert!(!timed_out); assert_eq!(next_sequence, 2); assert_eq!(events.len(), 1); + assert_eq!(events[0].sequence, 1); assert_eq!(events[0].text.as_deref(), Some("alpha beta")); + assert_eq!(events[0].kind, AgentOutputEventKind::AssistantMessage); } #[tokio::test] diff --git a/crates/server/src/transport.rs b/crates/server/src/transport.rs index 5dbe38f9..e6282509 100644 --- a/crates/server/src/transport.rs +++ b/crates/server/src/transport.rs @@ -57,16 +57,15 @@ impl ConnectionTransport { enum OutboundSink { Stdio, - WebSocket(futures::stream::SplitSink< - tokio_tungstenite::WebSocketStream, - Message, - >), + WebSocket( + futures::stream::SplitSink< + tokio_tungstenite::WebSocketStream, + Message, + >, + ), } -async fn write_outbound_frame( - sink: &mut OutboundSink, - frame: OutboundFrame, -) -> Result { +async fn write_outbound_frame(sink: &mut OutboundSink, frame: OutboundFrame) -> Result { let value = outbound_frame_to_value(&frame); log_outbound_frame(&frame, &value); let delivered = frame.delivered; @@ -378,8 +377,7 @@ async fn run_stdio(runtime: Arc) -> Result<()> { .await; tracing::info!(connection_id, "stdio connection established"); - let outbound_writer = - spawn_outbound_writer(connection_id, outbound_rx, OutboundSink::Stdio); + let outbound_writer = spawn_outbound_writer(connection_id, outbound_rx, OutboundSink::Stdio); let stdin = tokio::io::stdin(); let mut lines = BufReader::new(stdin).lines(); @@ -471,11 +469,8 @@ async fn handle_internal_proxy_connection( .await; tracing::info!(connection_id, "internal stdio proxy connection established"); - let outbound_writer = spawn_outbound_writer( - connection_id, - outbound_rx, - OutboundSink::WebSocket(writer), - ); + let outbound_writer = + spawn_outbound_writer(connection_id, outbound_rx, OutboundSink::WebSocket(writer)); if let Some(response) = runtime .handle_incoming_with_actions(connection_id, first_value) @@ -595,11 +590,8 @@ async fn handle_websocket_connection( .await; tracing::info!(connection_id, "websocket connection established"); - let outbound_writer = spawn_outbound_writer( - connection_id, - outbound_rx, - OutboundSink::WebSocket(writer), - ); + let outbound_writer = + spawn_outbound_writer(connection_id, outbound_rx, OutboundSink::WebSocket(writer)); while let Some(frame) = reader.next().await { let frame = frame?; @@ -653,9 +645,7 @@ fn parse_incoming_client_payload( /// server-initiated calls. Returns `Err(error_response)` when the payload is /// structurally invalid; the caller should send that response and skip the /// handler. -fn validate_incoming_client_message( - value: &serde_json::Value, -) -> Result<(), serde_json::Value> { +fn validate_incoming_client_message(value: &serde_json::Value) -> Result<(), serde_json::Value> { let Some(object) = value.as_object() else { return Err(acp_error_response( serde_json::Value::Null, @@ -664,10 +654,7 @@ fn validate_incoming_client_message( )); }; - let request_id = object - .get("id") - .cloned() - .unwrap_or(serde_json::Value::Null); + let request_id = object.get("id").cloned().unwrap_or(serde_json::Value::Null); match object.get("jsonrpc").and_then(serde_json::Value::as_str) { Some("2.0") => {} @@ -699,10 +686,7 @@ fn validate_incoming_client_message( )); } - if !has_method - && object.contains_key("id") - && (has_result || has_error) - { + if !has_method && object.contains_key("id") && (has_result || has_error) { return Ok(()); } @@ -872,9 +856,9 @@ mod tests { use super::InternalProxyControlRequest; use super::ListenTarget; use super::internal_proxy_control_response; - use super::parse_internal_proxy_control_request; use super::is_client_response_message; use super::parse_incoming_client_payload; + use super::parse_internal_proxy_control_request; use super::parse_listen_target; use super::resolve_listen_targets; use super::validate_incoming_client_message; @@ -1046,31 +1030,37 @@ mod tests { #[test] fn validate_accepts_client_request_notification_and_response() { - assert!(validate_incoming_client_message(&serde_json::json!({ - "jsonrpc": "2.0", - "id": 1, - "method": "initialize", - "params": {}, - })) - .is_ok()); - assert!(validate_incoming_client_message(&serde_json::json!({ - "jsonrpc": "2.0", - "method": "notifications/cancelled", - "params": {}, - })) - .is_ok()); - assert!(validate_incoming_client_message(&serde_json::json!({ - "jsonrpc": "2.0", - "id": 9, - "result": { "ok": true }, - })) - .is_ok()); + assert!( + validate_incoming_client_message(&serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": {}, + })) + .is_ok() + ); + assert!( + validate_incoming_client_message(&serde_json::json!({ + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": {}, + })) + .is_ok() + ); + assert!( + validate_incoming_client_message(&serde_json::json!({ + "jsonrpc": "2.0", + "id": 9, + "result": { "ok": true }, + })) + .is_ok() + ); } #[test] fn validate_rejects_malformed_client_messages() { - let invalid_request = validate_incoming_client_message(&serde_json::json!([])) - .expect_err("array payload"); + let invalid_request = + validate_incoming_client_message(&serde_json::json!([])).expect_err("array payload"); assert_eq!( invalid_request["error"]["code"], AcpErrorCode::InvalidRequest as i64 diff --git a/crates/server/tests/acp_available_commands.rs b/crates/server/tests/acp_available_commands.rs index db0a433a..1e13344b 100644 --- a/crates/server/tests/acp_available_commands.rs +++ b/crates/server/tests/acp_available_commands.rs @@ -22,9 +22,9 @@ use devo_protocol::StreamEvent; use devo_protocol::Usage; use devo_provider::ModelProviderSDK; use devo_provider::SingleProviderRouter; -use devo_server::OutboundFrame; use devo_server::AcpSuccessResponse; use devo_server::ClientTransportKind; +use devo_server::OutboundFrame; use devo_server::ServerRuntime; use devo_server::ServerRuntimeDependencies; use futures::stream; From 32226c268c5d1c943e91d833a9e04b641e114d12 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Sat, 4 Jul 2026 02:16:54 -1000 Subject: [PATCH 14/16] fix: register inline stream for research turns and suppress auto-git-diff during research Research runs outside the session actor (unlike execute_turn_in_actor). Without an inline stream, every item persist/usage update sends blocking mailbox commands and can stall inbound handlers that need the same actor. - Add BeginInlineTurn / EndInlineTurn commands to install inline state for out-of-actor turns, with merge-back into actor state on completion - Register active_stream in begin_research_turn_stream so persist_item uses the fast inline path instead of the mailbox - Add record_parent_turn_totals_and_latest for research ledger snapshots TUI: /research turns now set active_turn_is_research so automatic git-diff overlays are suppressed during research sessions. --- crates/client/src/protocol_trace.rs | 14 +- crates/client/src/stdio.rs | 8 +- crates/server/src/runtime/hooks.rs | 4 +- crates/server/src/runtime/items.rs | 8 +- crates/server/src/runtime/research.rs | 96 +++++++++---- .../src/runtime/session_actor/commands.rs | 10 ++ .../src/runtime/session_actor/handle.rs | 31 +++++ .../server/src/runtime/session_actor/loop_.rs | 18 +++ crates/server/src/runtime/subagent_usage.rs | 130 +++++++++++++++++- crates/server/tests/deep_research_e2e.rs | 3 +- crates/tui/src/chatwidget.rs | 11 ++ crates/tui/src/chatwidget/slash_commands.rs | 3 +- crates/tui/src/chatwidget/worker_events.rs | 6 +- crates/tui/src/chatwidget_tests.rs | 43 ++++++ crates/tui/src/exec_cell/render.rs | 12 +- 15 files changed, 337 insertions(+), 60 deletions(-) diff --git a/crates/client/src/protocol_trace.rs b/crates/client/src/protocol_trace.rs index c8c18936..271f9e84 100644 --- a/crates/client/src/protocol_trace.rs +++ b/crates/client/src/protocol_trace.rs @@ -120,14 +120,14 @@ impl ProtocolTrace { } fn resolve_trace_path() -> PathBuf { - if let Ok(explicit) = std::env::var("DEVO_PROTOCOL_TRACE_FILE") { - if !explicit.is_empty() { - let path = PathBuf::from(&explicit); - if let Some(parent) = path.parent() { - let _ = fs::create_dir_all(parent); - } - return path; + if let Ok(explicit) = std::env::var("DEVO_PROTOCOL_TRACE_FILE") + && !explicit.is_empty() + { + let path = PathBuf::from(&explicit); + if let Some(parent) = path.parent() { + let _ = fs::create_dir_all(parent); } + return path; } let base = devo_util_paths::find_devo_home() diff --git a/crates/client/src/stdio.rs b/crates/client/src/stdio.rs index 4b502ffe..38e2ce71 100644 --- a/crates/client/src/stdio.rs +++ b/crates/client/src/stdio.rs @@ -446,10 +446,10 @@ async fn write_ndjson_to_stdin( trace: Option<&ProtocolTrace>, ) -> Result<()> { let mut line = serde_json::to_vec(value).context("serialize client payload")?; - if let Some(t) = trace { - if let Ok(s) = std::str::from_utf8(&line) { - t.record(TraceDirection::Out, s); - } + if let Some(t) = trace + && let Ok(s) = std::str::from_utf8(&line) + { + t.record(TraceDirection::Out, s); } line.push(b'\n'); let mut stdin = stdin.lock().await; diff --git a/crates/server/src/runtime/hooks.rs b/crates/server/src/runtime/hooks.rs index a6136d11..4ab40687 100644 --- a/crates/server/src/runtime/hooks.rs +++ b/crates/server/src/runtime/hooks.rs @@ -65,9 +65,7 @@ impl ServerRuntime { } let session_handle = self.sessions.lock().await.get(&session_id).cloned()?; let snapshot = session_handle.hook_context_snapshot().await?; - Some(Self::hook_runtime_context_from_snapshot( - session_id, &snapshot, - )?) + Self::hook_runtime_context_from_snapshot(session_id, &snapshot) } fn hook_runtime_context_from_snapshot( diff --git a/crates/server/src/runtime/items.rs b/crates/server/src/runtime/items.rs index c2075c8b..49f0e589 100644 --- a/crates/server/src/runtime/items.rs +++ b/crates/server/src/runtime/items.rs @@ -393,10 +393,10 @@ impl ServerRuntime { }) }; if let Some(rollout) = inline_rollout { - if let Some((record, item)) = rollout { - if let Err(error) = self.rollout_store.append_item(&record, item) { - tracing::warn!(session_id = %session_id, error = %error, "failed to persist item line"); - } + if let Some((record, item)) = rollout + && let Err(error) = self.rollout_store.append_item(&record, item) + { + tracing::warn!(session_id = %session_id, error = %error, "failed to persist item line"); } return; } diff --git a/crates/server/src/runtime/research.rs b/crates/server/src/runtime/research.rs index d7af53fa..d756fbad 100644 --- a/crates/server/src/runtime/research.rs +++ b/crates/server/src/runtime/research.rs @@ -402,6 +402,12 @@ impl ServerRuntime { Some(turn_config.model.context_window as u64), ) .await; + // Research runs outside the session actor (unlike execute_turn_in_actor). + // Register the same TurnInlineState / active_stream path so item persist, + // usage, and hooks never flood the actor mailbox (capacity 64). Without + // this, every token/tool item does blocking send().await and can stall + // client inbound handlers that also need the mailbox. + self.begin_research_turn_stream(session_id, &turn).await; let result = self .run_research_pipeline( session_id, @@ -423,6 +429,30 @@ impl ServerRuntime { } else { self.clear_research_child_agents(session_id).await; } + if let Err(error) = &result { + let failure_message = format!("Research failed: {error}"); + self.emit_turn_item( + session_id, + turn.turn_id, + ItemKind::ResearchArtifact, + TurnItem::ResearchArtifact(ResearchArtifactItem { + artifact_type: ResearchArtifactType::Failure, + title: "Research Failure".to_string(), + content: failure_message.clone(), + }), + serde_json::json!({ + "artifact_type": "failure", + "title": "Research Failure", + "content": failure_message + }), + ) + .await; + } + // Merge inline mutations and free the mailbox path before any actor + // export/replace or finish work. + self.end_research_turn_stream(session_id, turn.turn_id) + .await; + self.refresh_core_session_prompt_context(session_id).await; let final_usage = usage_ledger.lock().await.aggregate(); self.clear_active_turn_runtime_handles(session_id).await; @@ -431,31 +461,40 @@ impl ServerRuntime { self.finish_research_turn(session_id, turn, TurnStatus::Completed, final_usage) .await; } - Err(error) => { - let failure_message = format!("Research failed: {error}"); - self.emit_turn_item( - session_id, - turn.turn_id, - ItemKind::ResearchArtifact, - TurnItem::ResearchArtifact(ResearchArtifactItem { - artifact_type: ResearchArtifactType::Failure, - title: "Research Failure".to_string(), - content: failure_message.clone(), - }), - serde_json::json!({ - "artifact_type": "failure", - "title": "Research Failure", - "content": failure_message - }), - ) - .await; - self.refresh_core_session_prompt_context(session_id).await; + Err(_) => { self.finish_research_turn(session_id, turn, TurnStatus::Failed, final_usage) .await; } } } + async fn begin_research_turn_stream(&self, session_id: SessionId, turn: &TurnMetadata) { + let Some(session_handle) = self.session(session_id).await else { + return; + }; + let Some(stream) = session_handle.begin_inline_turn(turn.clone()).await else { + tracing::warn!( + session_id = %session_id, + turn_id = %turn.turn_id, + "failed to begin research inline turn state" + ); + return; + }; + if let Some(spawn_snapshot) = session_handle.spawn_snapshot().await { + self.register_turn_spawn_snapshot(session_id, turn.turn_id, Arc::new(spawn_snapshot)) + .await; + } + self.register_active_stream(session_id, stream).await; + } + + async fn end_research_turn_stream(&self, session_id: SessionId, turn_id: TurnId) { + self.clear_turn_spawn_snapshot(session_id, turn_id).await; + self.unregister_active_stream(session_id).await; + if let Some(session_handle) = self.session(session_id).await { + session_handle.end_inline_turn().await; + } + } + async fn run_research_pipeline( self: &Arc, session_id: SessionId, @@ -587,7 +626,6 @@ impl ServerRuntime { context_reference, ) .await; - self.refresh_core_session_prompt_context(session_id).await; Ok(()) } @@ -1090,13 +1128,18 @@ impl ServerRuntime { turn.status = status.clone(); turn.completed_at = Some(Utc::now()); let own_usage = final_usage.to_turn_usage(); + let latest_query = self + .parent_usage_snapshot(session_id, turn.turn_id) + .await + .map(|snapshot| snapshot.latest_query_usage.to_turn_usage()) + .unwrap_or_else(|| own_usage.clone()); let usage = self - .publish_parent_turn_usage( + .publish_parent_turn_totals_and_latest( session_id, turn.turn_id, own_usage.clone(), + latest_query, /*context_window*/ None, - super::subagent_usage::UsageUpdateKind::InFlight, ) .await .map(|snapshot| snapshot.turn_usage.to_turn_usage()) @@ -1166,19 +1209,20 @@ impl ServerRuntime { usage: ResearchUsageTotals, context_window: Option, ) { + let latest_query = usage.to_turn_usage(); let aggregate = { let mut ledger = usage_ledger.lock().await; ledger.by_invocation.insert(usage_key, usage); ledger.aggregate() }; - // Research ledger already aggregates all invocations; publish as - // in-flight replacement so repeated updates do not double-count. - self.publish_parent_turn_usage( + // Research ledger aggregates all invocations for session/turn totals, + // but context-window display must use only the latest invocation. + self.publish_parent_turn_totals_and_latest( session_id, turn_id, aggregate.to_turn_usage(), + latest_query, context_window, - super::subagent_usage::UsageUpdateKind::InFlight, ) .await; } diff --git a/crates/server/src/runtime/session_actor/commands.rs b/crates/server/src/runtime/session_actor/commands.rs index e706708f..83562769 100644 --- a/crates/server/src/runtime/session_actor/commands.rs +++ b/crates/server/src/runtime/session_actor/commands.rs @@ -192,6 +192,16 @@ pub(crate) enum SessionCommand { state: Box, reply: oneshot::Sender<()>, }, + /// Install [`TurnInlineState`] for an out-of-actor turn (research) and + /// return the session stream so the runtime can register `active_stream`. + BeginInlineTurn { + turn: TurnMetadata, + reply: oneshot::Sender>>, + }, + /// Merge inline turn mutations back into actor state after an out-of-actor turn. + EndInlineTurn { + reply: oneshot::Sender<()>, + }, Shutdown { reply: oneshot::Sender<()>, }, diff --git a/crates/server/src/runtime/session_actor/handle.rs b/crates/server/src/runtime/session_actor/handle.rs index bcd6451c..ca26fc59 100644 --- a/crates/server/src/runtime/session_actor/handle.rs +++ b/crates/server/src/runtime/session_actor/handle.rs @@ -622,6 +622,37 @@ impl SessionHandle { reply_rx.await.ok() } + /// Begin out-of-actor turn inline state (research). Returns the session + /// stream for `register_active_stream`. + pub(crate) async fn begin_inline_turn( + &self, + turn: TurnMetadata, + ) -> Option>> { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::BeginInlineTurn { + turn, + reply: reply_tx, + }) + .await + { + return None; + } + reply_rx.await.ok() + } + + /// Merge and clear out-of-actor turn inline state. + pub(crate) async fn end_inline_turn(&self) { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::EndInlineTurn { reply: reply_tx }) + .await + { + return; + } + let _ = reply_rx.await; + } + pub(crate) async fn shutdown(&self) { let (reply_tx, reply_rx) = oneshot::channel(); if self diff --git a/crates/server/src/runtime/session_actor/loop_.rs b/crates/server/src/runtime/session_actor/loop_.rs index 91b6639c..cb462ed0 100644 --- a/crates/server/src/runtime/session_actor/loop_.rs +++ b/crates/server/src/runtime/session_actor/loop_.rs @@ -450,6 +450,24 @@ pub(super) async fn run_session_actor( state = *new_state; let _ = reply.send(()); } + SessionCommand::BeginInlineTurn { turn, reply } => { + { + let mut stream = state.stream.lock().await; + stream.turn_inline = + Some(super::turn_inline::TurnInlineState::new(&state, &turn)); + } + let _ = reply.send(Arc::clone(&state.stream)); + } + SessionCommand::EndInlineTurn { reply } => { + let inline = { + let mut stream = state.stream.lock().await; + stream.turn_inline.take() + }; + if let Some(inline) = inline { + inline.merge_into(&mut state); + } + let _ = reply.send(()); + } SessionCommand::Shutdown { reply } => { let _ = reply.send(()); break; diff --git a/crates/server/src/runtime/subagent_usage.rs b/crates/server/src/runtime/subagent_usage.rs index d464eb10..39debba5 100644 --- a/crates/server/src/runtime/subagent_usage.rs +++ b/crates/server/src/runtime/subagent_usage.rs @@ -85,7 +85,10 @@ impl UsageTotals { pub(crate) struct ParentUsageSnapshot { pub(super) session_id: SessionId, pub(super) turn_id: TurnId, + /// Accumulated usage for the parent turn (all legs + children). pub(super) turn_usage: UsageTotals, + /// Latest single model-call sample for context-window display. + pub(super) latest_query_usage: UsageTotals, pub(super) session_totals: UsageTotals, pub(super) context_window: Option, } @@ -113,6 +116,9 @@ struct ParentTurnUsage { parent_turn_usage: UsageTotals, /// Latest partial usage for the model call still in progress. inflight_usage: UsageTotals, + /// Most recent single model-call sample (not accumulated). + /// Used for context-window display (`last_query_*`). + latest_query_usage: UsageTotals, context_window: Option, } @@ -177,6 +183,7 @@ impl SubagentUsageState { base_child_totals, parent_turn_usage: UsageTotals::default(), inflight_usage: UsageTotals::default(), + latest_query_usage: UsageTotals::default(), context_window, }); } @@ -215,6 +222,33 @@ impl SubagentUsageState { usage, kind, ); + state.latest_query_usage = usage; + if context_window.is_some() { + state.context_window = context_window; + } + self.snapshot_for_parent_turn(session_id, turn_id) + } + + /// Replace parent-turn totals with an explicit snapshot (research ledger) + /// while tracking the latest single invocation for context-window display. + pub(super) fn record_parent_turn_totals_and_latest( + &mut self, + session_id: SessionId, + turn_id: TurnId, + turn_totals: UsageTotals, + latest_query: UsageTotals, + context_window: Option, + ) -> Option { + let key = ParentTurnKey { + session_id, + turn_id, + }; + let state = self.parent_turns.get_mut(&key)?; + // Research owns the full turn total as a single in-flight snapshot so + // repeated ledger publishes do not double-count completed legs. + state.parent_turn_usage = UsageTotals::default(); + state.inflight_usage = turn_totals; + state.latest_query_usage = latest_query; if context_window.is_some() { state.context_window = context_window; } @@ -254,9 +288,14 @@ impl SubagentUsageState { ); owner }; - owner - .parent_turn_id - .and_then(|turn_id| self.snapshot_for_parent_turn(owner.parent_session_id, turn_id)) + let parent_turn_id = owner.parent_turn_id?; + if let Some(parent) = self.parent_turns.get_mut(&ParentTurnKey { + session_id: owner.parent_session_id, + turn_id: parent_turn_id, + }) { + parent.latest_query_usage = usage; + } + self.snapshot_for_parent_turn(owner.parent_session_id, parent_turn_id) } /// Fold any remaining in-flight child usage into committed totals (e.g. on @@ -301,6 +340,7 @@ impl SubagentUsageState { session_id, turn_id, turn_usage, + latest_query_usage: state.latest_query_usage, session_totals, context_window: state.context_window, }) @@ -420,6 +460,32 @@ impl ServerRuntime { Some(snapshot) } + /// Publish research-ledger turn totals while keeping context display on the + /// latest single invocation (not the sum of all research stages). + pub(super) async fn publish_parent_turn_totals_and_latest( + &self, + session_id: SessionId, + turn_id: TurnId, + turn_totals: TurnUsage, + latest_query: TurnUsage, + context_window: Option, + ) -> Option { + self.begin_parent_usage_turn(session_id, turn_id, context_window) + .await; + let snapshot = { + let mut usage_state = self.subagent_usage.lock().await; + usage_state.record_parent_turn_totals_and_latest( + session_id, + turn_id, + UsageTotals::from_turn_usage(&turn_totals), + UsageTotals::from_turn_usage(&latest_query), + context_window, + ) + }?; + self.apply_parent_usage_snapshot(snapshot).await; + Some(snapshot) + } + pub(super) async fn publish_subagent_turn_usage( &self, child_session_id: SessionId, @@ -526,12 +592,13 @@ impl ParentUsageSnapshot { TurnUsageUpdatedPayload { session_id: self.session_id, turn_id: self.turn_id, - usage: self.turn_usage.to_turn_usage(), + // Context bar / last_query_* use the latest model call only. + usage: self.latest_query_usage.to_turn_usage(), total_input_tokens: self.session_totals.input_tokens, total_output_tokens: self.session_totals.output_tokens, total_tokens: self.session_totals.total_tokens, total_cache_read_tokens: self.session_totals.cache_read_input_tokens, - last_query_input_tokens: self.turn_usage.input_tokens, + last_query_input_tokens: self.latest_query_usage.input_tokens, context_window: self.context_window, } } @@ -559,7 +626,7 @@ impl ParentUsageSnapshot { summary.total_tokens = self.session_totals.total_tokens; summary.total_cache_creation_tokens = self.session_totals.cache_creation_input_tokens; summary.total_cache_read_tokens = self.session_totals.cache_read_input_tokens; - summary.last_query_total_tokens = self.turn_usage.total_tokens; + summary.last_query_total_tokens = self.latest_query_usage.total_tokens; } } @@ -625,6 +692,7 @@ mod tests { ) .expect("child snapshot"); assert_eq!(child_snapshot.turn_usage, totals(28, 7)); + assert_eq!(child_snapshot.latest_query_usage, totals(20, 5)); assert_eq!(child_snapshot.session_totals, totals(128, 17)); let updated_child_snapshot = state @@ -636,6 +704,7 @@ mod tests { ) .expect("updated child snapshot"); assert_eq!(updated_child_snapshot.turn_usage, totals(33, 8)); + assert_eq!(updated_child_snapshot.latest_query_usage, totals(25, 6)); assert_eq!(updated_child_snapshot.session_totals, totals(133, 18)); } @@ -658,6 +727,7 @@ mod tests { ) .expect("first leg"); assert_eq!(after_first_leg.turn_usage, totals(600, 50)); + assert_eq!(after_first_leg.latest_query_usage, totals(600, 50)); assert_eq!(after_first_leg.session_totals, totals(700, 60)); // Second model call starts (typical UsageDelta with output_tokens = 0). @@ -671,6 +741,7 @@ mod tests { ) .expect("second leg start"); assert_eq!(after_second_start.turn_usage, totals(1300, 50)); + assert_eq!(after_second_start.latest_query_usage, totals(700, 0)); assert_eq!(after_second_start.session_totals, totals(1400, 60)); // Second model call completes. @@ -684,9 +755,56 @@ mod tests { ) .expect("second leg complete"); assert_eq!(after_second_leg.turn_usage, totals(1300, 80)); + assert_eq!(after_second_leg.latest_query_usage, totals(700, 30)); assert_eq!(after_second_leg.session_totals, totals(1400, 90)); } + #[test] + fn research_totals_accumulate_but_context_uses_latest_invocation() { + let mut state = SubagentUsageState::default(); + let parent_session_id = SessionId::new(); + let parent_turn_id = TurnId::new(); + + state.begin_parent_turn( + parent_session_id, + parent_turn_id, + totals(0, 0), + Some(200_000), + ); + + let after_first = state + .record_parent_turn_totals_and_latest( + parent_session_id, + parent_turn_id, + totals(8_000, 1_000), + totals(8_000, 1_000), + Some(200_000), + ) + .expect("first research invocation"); + assert_eq!(after_first.turn_usage, totals(8_000, 1_000)); + assert_eq!(after_first.latest_query_usage, totals(8_000, 1_000)); + + let after_second = state + .record_parent_turn_totals_and_latest( + parent_session_id, + parent_turn_id, + totals(20_000, 3_000), + totals(12_000, 2_000), + Some(200_000), + ) + .expect("second research invocation"); + assert_eq!(after_second.turn_usage, totals(20_000, 3_000)); + assert_eq!(after_second.latest_query_usage, totals(12_000, 2_000)); + assert_eq!(after_second.session_totals, totals(20_000, 3_000)); + + let payload = after_second.to_turn_usage_updated_payload(); + assert_eq!(payload.usage.input_tokens, 12_000); + assert_eq!(payload.usage.output_tokens, 2_000); + assert_eq!(payload.last_query_input_tokens, 12_000); + assert_eq!(payload.total_input_tokens, 20_000); + assert_eq!(payload.total_output_tokens, 3_000); + } + #[test] fn next_parent_turn_base_includes_previous_child_usage() { let mut state = SubagentUsageState::default(); diff --git a/crates/server/tests/deep_research_e2e.rs b/crates/server/tests/deep_research_e2e.rs index 899bd570..464dabfe 100644 --- a/crates/server/tests/deep_research_e2e.rs +++ b/crates/server/tests/deep_research_e2e.rs @@ -1450,6 +1450,7 @@ fn hosted_web_search_researcher_response() -> ModelResponse { fn supervisor_stream_events(request: &ModelRequest) -> Vec> { assert_supervisor_request_uses_agent_tools(request); + #[allow(clippy::never_loop)] for attempt in 1..=3 { let spawn_id = format!("spawn-supervisor-worker-{attempt}"); if !request_has_tool_result(request, &spawn_id) { @@ -2307,7 +2308,7 @@ fn latest_agent_message(events: &[serde_json::Value]) -> Option { events .iter() .rev() - .find_map(|event| agent_message_from_completed_event(event)) + .find_map(agent_message_from_completed_event) } fn latest_parent_agent_message( diff --git a/crates/tui/src/chatwidget.rs b/crates/tui/src/chatwidget.rs index 776723ab..8172fdba 100644 --- a/crates/tui/src/chatwidget.rs +++ b/crates/tui/src/chatwidget.rs @@ -298,6 +298,8 @@ pub(crate) struct ChatWidget { queued_input_modes: VecDeque, promoted_input_modes: VecDeque, active_turn_id: Option, + /// True while a `/research` turn is active; suppresses automatic git-diff overlays. + active_turn_is_research: bool, current_turn_mode: InputMode, committed_server_assistant_in_turn: bool, boundary_committed_assistant_items: HashSet, @@ -331,6 +333,14 @@ impl ChatWidget { pub(crate) fn should_auto_show_git_diff(tool_title: &str, is_error: bool) -> bool { diff_rules::should_auto_show_git_diff(tool_title, is_error) } + + pub(crate) fn should_auto_show_git_diff_for_turn( + &self, + tool_title: &str, + is_error: bool, + ) -> bool { + !self.active_turn_is_research && diff_rules::should_auto_show_git_diff(tool_title, is_error) + } pub(crate) fn new_with_app_event(common: ChatWidgetInit) -> Self { // Pull the constructor inputs apart up front so the setup below reads in stages. let ChatWidgetInit { @@ -460,6 +470,7 @@ impl ChatWidget { queued_input_modes: VecDeque::new(), promoted_input_modes: VecDeque::new(), active_turn_id: None, + active_turn_is_research: false, current_turn_mode: InputMode::Build, committed_server_assistant_in_turn: false, boundary_committed_assistant_items: HashSet::new(), diff --git a/crates/tui/src/chatwidget/slash_commands.rs b/crates/tui/src/chatwidget/slash_commands.rs index 9d76fd57..bba21b0b 100644 --- a/crates/tui/src/chatwidget/slash_commands.rs +++ b/crates/tui/src/chatwidget/slash_commands.rs @@ -49,7 +49,7 @@ impl ChatWidget { self.set_status_message(format!("Cannot change {noun} while generating")); } - pub(super) fn handle_slash_command(&mut self, command: SlashCommand, argument: String) { + pub(crate) fn handle_slash_command(&mut self, command: SlashCommand, argument: String) { if !self.can_change_configuration() && !command.available_during_task() { self.add_busy_configuration_message(command); return; @@ -188,6 +188,7 @@ impl ChatWidget { self.active_accent_color(), self.current_turn_mode, )); + self.active_turn_is_research = true; self.app_event_tx .send(AppEvent::Command(AppCommand::RunResearch { question: trimmed.to_string(), diff --git a/crates/tui/src/chatwidget/worker_events.rs b/crates/tui/src/chatwidget/worker_events.rs index c6f65ac1..7dece819 100644 --- a/crates/tui/src/chatwidget/worker_events.rs +++ b/crates/tui/src/chatwidget/worker_events.rs @@ -587,7 +587,7 @@ impl ChatWidget { } else { "Tool completed" }); - if Self::should_auto_show_git_diff(&resolved_title, is_error) { + if self.should_auto_show_git_diff_for_turn(&resolved_title, is_error) { let tx = self.app_event_tx.clone(); tokio::spawn(async move { let text = Self::format_git_diff_result(get_git_diff().await); @@ -703,7 +703,7 @@ impl ChatWidget { } else { "Tool completed" }); - if Self::should_auto_show_git_diff(&resolved_title, is_error) { + if self.should_auto_show_git_diff_for_turn(&resolved_title, is_error) { let tx = self.app_event_tx.clone(); tokio::spawn(async move { let text = Self::format_git_diff_result(get_git_diff().await); @@ -851,6 +851,7 @@ impl ChatWidget { prompt_token_estimate, } => { let was_interrupted = stop_reason.contains("Interrupted"); + self.active_turn_is_research = false; self.commit_active_streams(DotStatus::Completed); if was_interrupted && let Some(cell) = self @@ -926,6 +927,7 @@ impl ChatWidget { last_query_input_tokens, } => { self.resume_browser_loading = false; + self.active_turn_is_research = false; self.commit_active_streams(DotStatus::Failed); self.active_tool_calls.clear(); self.pending_tool_calls.clear(); diff --git a/crates/tui/src/chatwidget_tests.rs b/crates/tui/src/chatwidget_tests.rs index 634e34aa..634b6bd0 100644 --- a/crates/tui/src/chatwidget_tests.rs +++ b/crates/tui/src/chatwidget_tests.rs @@ -8267,6 +8267,49 @@ async fn successful_write_tool_result_triggers_diff_event() { ); } +#[tokio::test] +async fn research_turn_does_not_auto_show_git_diff() { + let model = Model { + slug: "test-model".to_string(), + display_name: "Test Model".to_string(), + ..Model::default() + }; + let (mut widget, mut app_event_rx) = widget_with_model(model, PathBuf::from(".")); + + widget.handle_slash_command( + crate::slash_command::SlashCommand::Research, + "DeepSeek official website".to_string(), + ); + assert!(!widget.should_auto_show_git_diff_for_turn("write report.md", false)); + + widget.handle_worker_event(crate::events::WorkerEvent::ToolCall { + tool_use_id: "tool-1".to_string(), + summary: "write report.md".to_string(), + preparing: false, + parsed_commands: None, + }); + widget.handle_worker_event(crate::events::WorkerEvent::ToolResult { + tool_use_id: "tool-1".to_string(), + title: "write report.md".to_string(), + preview: "updated".to_string(), + is_error: false, + truncated: false, + }); + + let auto_diff = tokio::time::timeout(std::time::Duration::from_millis(200), async { + loop { + if let Some(AppEvent::DiffResult(_)) = app_event_rx.recv().await { + break true; + } + } + }) + .await; + assert!( + auto_diff.is_err(), + "research turns must not emit automatic DiffResult events" + ); +} + #[test] fn patch_applied_event_renders_edited_block() { let model = Model { diff --git a/crates/tui/src/exec_cell/render.rs b/crates/tui/src/exec_cell/render.rs index b8f536d6..5e9efa16 100644 --- a/crates/tui/src/exec_cell/render.rs +++ b/crates/tui/src/exec_cell/render.rs @@ -483,12 +483,12 @@ impl ExecCell { push_owned_lines(&wrapped, &mut out_indented); } - if call.output.is_none() { - if let (Some(tool_name), Some(input)) = (&call.tool_name, &call.tool_input) { - let input_lines = tool_input_lines(tool_name, input); - for line in input_lines { - out_indented.push(line.patch_style(Style::default().dim())); - } + if call.output.is_none() + && let (Some(tool_name), Some(input)) = (&call.tool_name, &call.tool_input) + { + let input_lines = tool_input_lines(tool_name, input); + for line in input_lines { + out_indented.push(line.patch_style(Style::default().dim())); } } } From 50af3d1b3df6b83546b6c000d0a48675a5ce2121 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Sat, 4 Jul 2026 06:14:45 -1000 Subject: [PATCH 15/16] fix: align interrupt, turn queue, and title generation with inline session-actor execution --- crates/server/src/runtime.rs | 3 + crates/server/src/runtime/agents.rs | 6 +- .../server/src/runtime/goal_continuation.rs | 152 +++-------- crates/server/src/runtime/goal_handlers.rs | 89 +++++-- crates/server/src/runtime/handlers.rs | 1 + .../src/runtime/handlers/acp/session.rs | 1 + .../src/runtime/handlers/message_edit.rs | 6 +- crates/server/src/runtime/handlers/turn.rs | 215 ++------------- .../src/runtime/handlers/turn_interrupt.rs | 244 ++++++++++++++++++ crates/server/src/runtime/items.rs | 50 +++- crates/server/src/runtime/research.rs | 8 +- .../src/runtime/session_actor/commands.rs | 3 - .../src/runtime/session_actor/handle.rs | 38 ++- .../server/src/runtime/session_actor/loop_.rs | 26 +- .../src/runtime/session_actor/registry.rs | 32 +++ .../server/src/runtime/turn_exec/finalize.rs | 1 + .../server/src/runtime/turn_exec/followup.rs | 9 +- crates/server/src/runtime/turn_exec/tests.rs | 2 +- crates/server/src/runtime/turn_reservation.rs | 16 ++ crates/server/tests/deep_research_e2e.rs | 55 +--- crates/server/tests/end_to_end.rs | 7 + crates/server/tests/goal_continuation.rs | 28 ++ .../server/tests/goal_lifecycle_interrupts.rs | 7 + crates/server/tests/persistence_resume.rs | 2 +- 24 files changed, 589 insertions(+), 412 deletions(-) create mode 100644 crates/server/src/runtime/handlers/turn_interrupt.rs diff --git a/crates/server/src/runtime.rs b/crates/server/src/runtime.rs index 7a2b7221..dc569096 100644 --- a/crates/server/src/runtime.rs +++ b/crates/server/src/runtime.rs @@ -227,6 +227,8 @@ pub struct ServerRuntime { active_turn_cancellations: Mutex>, /// Active turn ids tracked at runtime level so cancel/interrupt can avoid session-actor mailbox round-trips while a turn is blocked in permission wait. active_turn_ids: Mutex>, + /// Full active-turn metadata for reservation snapshots while the session actor is busy in `ExecuteTurn`. + active_turn_metadata: Mutex>, active_turn_connections: Mutex>, terminal_turn_statuses: Mutex>, acp_prompt_waiters: Mutex>>>, @@ -348,6 +350,7 @@ impl ServerRuntime { active_tasks: Mutex::new(HashMap::new()), active_turn_cancellations: Mutex::new(HashMap::new()), active_turn_ids: Mutex::new(HashMap::new()), + active_turn_metadata: Mutex::new(HashMap::new()), active_turn_connections: Mutex::new(HashMap::new()), terminal_turn_statuses: Mutex::new(VecDeque::new()), acp_prompt_waiters: Mutex::new(HashMap::new()), diff --git a/crates/server/src/runtime/agents.rs b/crates/server/src/runtime/agents.rs index b6fdd345..78cc1885 100644 --- a/crates/server/src/runtime/agents.rs +++ b/crates/server/src/runtime/agents.rs @@ -555,10 +555,8 @@ impl ServerRuntime { .lock() .await .insert(session_id, cancel_token); - self.active_turn_ids - .lock() - .await - .insert(session_id, turn.turn_id); + self.register_runtime_active_turn(session_id, turn.clone()) + .await; let task = tokio::spawn(async move { runtime .execute_turn(ExecuteTurnRequest { diff --git a/crates/server/src/runtime/goal_continuation.rs b/crates/server/src/runtime/goal_continuation.rs index 387e10b1..96568a8f 100644 --- a/crates/server/src/runtime/goal_continuation.rs +++ b/crates/server/src/runtime/goal_continuation.rs @@ -189,10 +189,8 @@ impl ServerRuntime { false } else { cancellations.insert(session_id, cancel_token); - self.active_turn_ids - .lock() - .await - .insert(session_id, turn.turn_id); + self.register_runtime_active_turn(session_id, turn.clone()) + .await; let task = tokio::spawn(async move { runtime .execute_turn(ExecuteTurnRequest { @@ -442,39 +440,6 @@ impl ServerRuntime { session_handle.clear_active_turn_if_matches(turn_id).await; } - async fn complete_deferred_items_for_goal_turn( - &self, - session_handle: &SessionHandle, - session_id: SessionId, - turn_id: TurnId, - ) { - let deferred = session_handle.take_deferred_items().await; - if let Some((item_id, item_seq, text)) = deferred.assistant { - self.complete_item( - session_id, - turn_id, - item_id, - item_seq, - ItemKind::AgentMessage, - TurnItem::AgentMessage(TextItem { text: text.clone() }), - serde_json::json!({ "title": "Assistant", "text": text }), - ) - .await; - } - if let Some((item_id, item_seq, text)) = deferred.reasoning { - self.complete_item( - session_id, - turn_id, - item_id, - item_seq, - ItemKind::Reasoning, - TurnItem::Reasoning(TextItem { text: text.clone() }), - serde_json::json!({ "title": "Reasoning", "text": text }), - ) - .await; - } - } - pub(super) async fn interrupt_active_goal_continuation_turn( self: &Arc, session_id: SessionId, @@ -488,32 +453,37 @@ impl ServerRuntime { else { return false; }; - let Some(session_handle) = self.session(session_id).await else { - self.goal_continuation_turn_goals - .lock() - .await - .remove(&turn_id); - return false; - }; - let active_turn_matches = session_handle - .active_turn_id() + self.goal_continuation_turn_goals + .lock() .await - .is_some_and(|active_turn_id| active_turn_id == Some(turn_id)); - if !active_turn_matches { - self.goal_continuation_turn_goals - .lock() - .await - .remove(&turn_id); - return false; + .remove(&turn_id); + + // Turns run inline on the session actor. Do not touch the actor mailbox + // while ExecuteTurn is in flight — cancel the turn token and wait for + // `finalize_executed_turn` to emit lifecycle events instead. + if self.runtime_active_turn_id(session_id).await != Some(turn_id) { + let already_terminal = self.recent_terminal_turn_status(turn_id).await.is_some(); + if already_terminal { + tracing::info!( + session_id = %session_id, + turn_id = %turn_id, + reason, + "goal continuation turn already terminal" + ); + } + return already_terminal; } - self.complete_deferred_items_for_goal_turn(&session_handle, session_id, turn_id) - .await; - let Some(interrupted_turn) = session_handle.interrupt_active_turn().await.flatten() else { - return false; - }; - if interrupted_turn.turn_id != turn_id { - return false; + let terminal_rx = self.subscribe_terminal_turn_status(turn_id).await; + if let Some(snapshot) = self.recent_terminal_turn_status(turn_id).await { + self.record_terminal_turn_status(turn_id, snapshot).await; + tracing::info!( + session_id = %session_id, + turn_id = %turn_id, + reason, + "goal continuation turn already terminal" + ); + return true; } // Cancel via a clone rather than `remove`: see the comment in @@ -532,60 +502,20 @@ impl ServerRuntime { task.abort(); } - let (record, session_context, turn_context) = { - let persistence = session_handle.turn_persistence_snapshot().await; - match persistence { - Some(persistence) => ( - persistence.record, - persistence.session_context, - persistence.latest_turn_context, - ), - None => (None, None, None), - } - }; - if let Some(record) = record - && let Err(error) = self.rollout_store.append_turn( - &record, - build_turn_record(&interrupted_turn, session_context, turn_context), - ) - { - tracing::warn!( + let completed = + match tokio::time::timeout(std::time::Duration::from_secs(5), terminal_rx).await { + Ok(Ok(_)) => true, + Ok(Err(_)) | Err(_) => self.recent_terminal_turn_status(turn_id).await.is_some(), + }; + if completed { + tracing::info!( session_id = %session_id, - turn_id = %interrupted_turn.turn_id, - error = %error, - "failed to persist interrupted goal continuation turn" + turn_id = %turn_id, + reason, + "interrupted active goal continuation turn" ); } - - tracing::info!( - session_id = %session_id, - turn_id = %interrupted_turn.turn_id, - reason, - "interrupted active goal continuation turn" - ); - self.broadcast_event(ServerEvent::TurnInterrupted(TurnEventPayload { - session_id, - turn: interrupted_turn.clone(), - })) - .await; - self.broadcast_event(ServerEvent::TurnCompleted(TurnEventPayload { - session_id, - turn: interrupted_turn, - })) - .await; - self.broadcast_event(ServerEvent::SessionStatusChanged( - SessionStatusChangedPayload { - session_id, - status: SessionRuntimeStatus::Idle, - }, - )) - .await; - - let runtime = Arc::clone(self); - tokio::spawn(async move { - runtime.spawn_next_turn_from_queue(session_id).await; - }); - true + completed } } diff --git a/crates/server/src/runtime/goal_handlers.rs b/crates/server/src/runtime/goal_handlers.rs index 42b288d4..d9fd600e 100644 --- a/crates/server/src/runtime/goal_handlers.rs +++ b/crates/server/src/runtime/goal_handlers.rs @@ -50,16 +50,19 @@ impl ServerRuntime { { tracing::warn!(session_id = %session_id, error = %error, "failed to persist goal create record"); } - self.sync_core_session_goal(session_id, session_goal).await; - self.maybe_start_title_generation_from_user_input(session_id, &title_input) - .await; + // Interrupt before any session-actor mailbox round-trip: the actor + // may be blocked inside an in-flight continuation turn. if replace_existing { self.interrupt_active_goal_continuation_turn(session_id, "goal replaced") .await; } - if should_continue { - self.maybe_start_goal_continuation_turn(session_id).await; - } + self.sync_core_session_goal(session_id, session_goal).await; + self.schedule_goal_followup_work( + session_id, + Some(title_input), + should_continue, + ) + .await; result } Err(e) => self.error_response( @@ -119,12 +122,12 @@ impl ServerRuntime { }) .expect("serialize budget-limited goal pause result"); drop(stores); - self.sync_core_session_goal(session_id, None).await; self.interrupt_active_goal_continuation_turn( session_id, "budget-limited goal wrap-up stopped", ) .await; + self.sync_core_session_goal(session_id, None).await; return result; } match store.set(params) { @@ -161,7 +164,6 @@ impl ServerRuntime { { tracing::warn!(session_id = %session_id, error = %error, "failed to persist goal status record"); } - self.sync_core_session_goal(session_id, session_goal).await; if should_interrupt_continuation { self.interrupt_active_goal_continuation_turn( session_id, @@ -169,13 +171,9 @@ impl ServerRuntime { ) .await; } - if let Some(title_input) = title_input { - self.maybe_start_title_generation_from_user_input(session_id, &title_input) - .await; - } - if should_continue { - self.maybe_start_goal_continuation_turn(session_id).await; - } + self.sync_core_session_goal(session_id, session_goal).await; + self.schedule_goal_followup_work(session_id, title_input, should_continue) + .await; result } Err(e) => self.error_response( @@ -229,12 +227,12 @@ impl ServerRuntime { .expect("serialize budget-limited goal pause result"); let session_id = params.session_id; drop(stores); - self.sync_core_session_goal(session_id, None).await; self.interrupt_active_goal_continuation_turn( session_id, "budget-limited goal wrap-up stopped", ) .await; + self.sync_core_session_goal(session_id, None).await; return result; } match store.set_status(devo_protocol::ThreadGoalStatus::Paused) { @@ -256,11 +254,11 @@ impl ServerRuntime { { tracing::warn!(session_id = %session_id, error = %error, "failed to persist goal pause record"); } - self.sync_core_session_goal(session_id, None).await; if should_interrupt_continuation { self.interrupt_active_goal_continuation_turn(session_id, "goal paused") .await; } + self.sync_core_session_goal(session_id, None).await; result } Err(e) => self.error_response( @@ -318,9 +316,12 @@ impl ServerRuntime { tracing::warn!(session_id = %session_id, error = %error, "failed to persist goal resume record"); } self.sync_core_session_goal(session_id, session_goal).await; - if should_continue { - self.maybe_start_goal_continuation_turn(session_id).await; - } + self.schedule_goal_followup_work( + session_id, + /*title_input*/ None, + should_continue, + ) + .await; result } Err(e) => self.error_response( @@ -376,9 +377,9 @@ impl ServerRuntime { { tracing::warn!(session_id = %session_id, error = %error, "failed to persist goal complete record"); } - self.sync_core_session_goal(session_id, None).await; self.interrupt_active_goal_continuation_turn(session_id, "goal completed") .await; + self.sync_core_session_goal(session_id, None).await; result } Err(e) => self.error_response( @@ -436,9 +437,9 @@ impl ServerRuntime { { tracing::warn!(session_id = %session_id, error = %error, "failed to persist goal cancel record"); } - self.sync_core_session_goal(session_id, None).await; self.interrupt_active_goal_continuation_turn(session_id, "goal canceled") .await; + self.sync_core_session_goal(session_id, None).await; result } Err(e) => self.error_response( @@ -484,9 +485,9 @@ impl ServerRuntime { { tracing::warn!(session_id = %params.session_id, error = %error, "failed to persist goal clear record"); } - self.sync_core_session_goal(params.session_id, None).await; self.interrupt_active_goal_continuation_turn(params.session_id, "goal cleared") .await; + self.sync_core_session_goal(params.session_id, None).await; } serde_json::to_value(SuccessResponse { @@ -533,6 +534,48 @@ impl ServerRuntime { let Some(session_handle) = self.session(session_id).await else { return; }; + if self.runtime_active_turn_id(session_id).await.is_some() { + // Queue without blocking the goal handler; the actor applies this once + // the in-flight turn releases the mailbox. + let _ = session_handle.try_set_active_goal(goal); + return; + } session_handle.set_active_goal(goal).await; } + + /// Title generation and continuation startup both need session-actor mailbox + /// replies. When a turn is already running inline on that actor, awaiting + /// those replies deadlocks the goal handler. Defer title work to a task and + /// rely on the post-turn hook for continuation while a turn is active. + async fn schedule_goal_followup_work( + self: &Arc, + session_id: SessionId, + title_input: Option, + should_continue: bool, + ) { + let turn_active = self.runtime_active_turn_id(session_id).await.is_some(); + if let Some(title_input) = title_input { + if turn_active { + let runtime = Arc::clone(self); + tokio::spawn(async move { + runtime + .maybe_prepare_title_generation_from_user_input( + session_id, + &title_input, + ) + .await; + }); + } else { + self.maybe_start_title_generation_from_user_input(session_id, &title_input) + .await; + } + } + if !should_continue { + return; + } + if turn_active { + return; + } + self.maybe_start_goal_continuation_turn(session_id).await; + } } diff --git a/crates/server/src/runtime/handlers.rs b/crates/server/src/runtime/handlers.rs index 0ca5e2c4..99fb2273 100644 --- a/crates/server/src/runtime/handlers.rs +++ b/crates/server/src/runtime/handlers.rs @@ -8,4 +8,5 @@ mod message_edit; mod message_edit_restore; mod session; mod turn; +mod turn_interrupt; mod workspace_changes; diff --git a/crates/server/src/runtime/handlers/acp/session.rs b/crates/server/src/runtime/handlers/acp/session.rs index ef8a98e2..a2c74a4f 100644 --- a/crates/server/src/runtime/handlers/acp/session.rs +++ b/crates/server/src/runtime/handlers/acp/session.rs @@ -517,6 +517,7 @@ impl ServerRuntime { cancellation.cancel(); } self.active_turn_ids.lock().await.remove(&session_id); + self.active_turn_metadata.lock().await.remove(&session_id); self.active_turn_connections .lock() .await diff --git a/crates/server/src/runtime/handlers/message_edit.rs b/crates/server/src/runtime/handlers/message_edit.rs index 352d90ea..b8a60879 100644 --- a/crates/server/src/runtime/handlers/message_edit.rs +++ b/crates/server/src/runtime/handlers/message_edit.rs @@ -466,10 +466,8 @@ impl ServerRuntime { .lock() .await .insert(params.session_id, CancellationToken::new()); - self.active_turn_ids - .lock() - .await - .insert(params.session_id, replacement_turn.turn_id); + self.register_runtime_active_turn(params.session_id, replacement_turn.clone()) + .await; let runtime = Arc::clone(self); let replacement_turn_for_task = replacement_turn.clone(); let turn_config_for_task = turn_config.clone(); diff --git a/crates/server/src/runtime/handlers/turn.rs b/crates/server/src/runtime/handlers/turn.rs index 40ad777c..bf932324 100644 --- a/crates/server/src/runtime/handlers/turn.rs +++ b/crates/server/src/runtime/handlers/turn.rs @@ -83,7 +83,10 @@ impl ServerRuntime { "session does not exist", ); }; - let Some(reservation) = session_handle.turn_reservation_snapshot().await else { + let Some(reservation) = self + .session_turn_reservation_snapshot(params.session_id) + .await + else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, @@ -205,9 +208,7 @@ impl ServerRuntime { now, ); let queued_input_id = item.id; - session_handle - .enqueue_pending_turn_input(item.clone()) - .await; + session_handle.push_pending_turn_input(item.clone()); if !reservation.ephemeral && let Err(err) = self.deps @@ -310,7 +311,7 @@ impl ServerRuntime { ) .await; } - self.maybe_start_title_generation_from_user_input(params.session_id, &display_input) + self.maybe_prepare_title_generation_from_user_input(params.session_id, &display_input) .await; if let Some(persistence) = session_handle.turn_persistence_snapshot().await && let Some(record) = persistence.record @@ -345,10 +346,8 @@ impl ServerRuntime { .lock() .await .insert(params.session_id, CancellationToken::new()); - self.active_turn_ids - .lock() - .await - .insert(params.session_id, turn.turn_id); + self.register_runtime_active_turn(params.session_id, turn.clone()) + .await; if let Some(connection_id) = connection_id { self.active_turn_connections .lock() @@ -459,7 +458,10 @@ impl ServerRuntime { }, None => None, }; - let Some(reservation) = session_handle.turn_reservation_snapshot().await else { + let Some(reservation) = self + .session_turn_reservation_snapshot(params.session_id) + .await + else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, @@ -574,10 +576,8 @@ impl ServerRuntime { .lock() .await .insert(params.session_id, CancellationToken::new()); - self.active_turn_ids - .lock() - .await - .insert(params.session_id, turn.turn_id); + self.register_runtime_active_turn(params.session_id, turn.clone()) + .await; if let Some(connection_id) = connection_id { self.active_turn_connections .lock() @@ -605,176 +605,6 @@ impl ServerRuntime { .expect("serialize turn/shell_command response") } - pub(crate) async fn handle_turn_interrupt( - self: &Arc, - request_id: serde_json::Value, - params: serde_json::Value, - ) -> serde_json::Value { - let params: TurnInterruptParams = match serde_json::from_value(params) { - Ok(params) => params, - Err(error) => { - return self.error_response( - request_id, - ProtocolErrorCode::InvalidParams, - format!("invalid turn/interrupt params: {error}"), - ); - } - }; - let Some(session_handle) = self.session(params.session_id).await else { - return self.error_response( - request_id, - ProtocolErrorCode::SessionNotFound, - "session does not exist", - ); - }; - - // Cancel before any session-actor mailbox round-trip: the actor may be blocked - // waiting for a permission response and cannot process commands until cancelled. - if let Some(cancel_token) = self - .active_turn_cancellations - .lock() - .await - .get(¶ms.session_id) - .cloned() - { - cancel_token.cancel(); - } - if let Some(task) = self.active_tasks.lock().await.remove(¶ms.session_id) { - task.abort(); - } - - let removed_len = self - .session_interactive - .clear_pending_user_inputs_for_turn(params.session_id, params.turn_id) - .await; - if removed_len > 0 { - tracing::info!( - session_id = %params.session_id, - turn_id = %params.turn_id, - removed_len, - "cleared pending request_user_input requests for interrupted turn" - ); - } - - // Cancel via a clone rather than `remove`: see the comment in - // `interrupt_child_runtime_work` for why removing here races with - // `run_turn_model_query` fetching the same token. - let close_children_parent = params.session_id; - Arc::clone(self) - .interrupt_all_child_agents(close_children_parent) - .await; - if self.runtime_active_turn_id(params.session_id).await != Some(params.turn_id) { - return self.error_response( - request_id, - ProtocolErrorCode::TurnNotFound, - "turn is not active", - ); - } - let Some(interrupted_turn) = session_handle.interrupt_active_turn().await.flatten() else { - return self.error_response( - request_id, - ProtocolErrorCode::TurnNotFound, - "turn is not active", - ); - }; - if interrupted_turn.turn_id != params.turn_id { - return self.error_response( - request_id, - ProtocolErrorCode::TurnNotFound, - "turn does not exist", - ); - } - - let deferred = session_handle.take_deferred_items().await; - if let Some((item_id, item_seq, text)) = deferred.assistant { - self.complete_item( - params.session_id, - params.turn_id, - item_id, - item_seq, - ItemKind::AgentMessage, - TurnItem::AgentMessage(TextItem { text: text.clone() }), - serde_json::json!({ "title": "Assistant", "text": text }), - ) - .await; - } - if let Some((item_id, item_seq, text)) = deferred.reasoning { - self.complete_item( - params.session_id, - params.turn_id, - item_id, - item_seq, - ItemKind::Reasoning, - TurnItem::Reasoning(TextItem { text: text.clone() }), - serde_json::json!({ "title": "Reasoning", "text": text }), - ) - .await; - } - if let Some(persistence) = session_handle.turn_persistence_snapshot().await - && let Some(record) = persistence.record - && let Err(error) = self.rollout_store.append_turn( - &record, - build_turn_record( - &interrupted_turn, - persistence.session_context, - persistence.latest_turn_context, - ), - ) - { - return self.error_response( - request_id, - ProtocolErrorCode::InternalError, - format!("failed to persist interrupted turn: {error}"), - ); - } - - tracing::info!( - session_id = %params.session_id, - turn_id = %interrupted_turn.turn_id, - status = ?interrupted_turn.status, - "interrupted turn" - ); - self.finalize_turn_workspace_changes(params.session_id, &interrupted_turn) - .await; - self.broadcast_event(ServerEvent::TurnInterrupted(TurnEventPayload { - session_id: params.session_id, - turn: interrupted_turn.clone(), - })) - .await; - self.broadcast_event(ServerEvent::TurnCompleted(TurnEventPayload { - session_id: params.session_id, - turn: interrupted_turn.clone(), - })) - .await; - self.broadcast_event(ServerEvent::SessionStatusChanged( - SessionStatusChangedPayload { - session_id: params.session_id, - status: SessionRuntimeStatus::Idle, - }, - )) - .await; - self.record_terminal_turn_status( - interrupted_turn.turn_id, - TerminalTurnSnapshot::from_turn(&interrupted_turn), - ) - .await; - - let runtime = Arc::clone(self); - let sid = params.session_id; - tokio::spawn(async move { - runtime.spawn_next_turn_from_queue(sid).await; - }); - - serde_json::to_value(SuccessResponse { - id: request_id, - result: TurnInterruptResult { - turn_id: interrupted_turn.turn_id, - status: interrupted_turn.status, - }, - }) - .expect("serialize turn/interrupt response") - } - pub(crate) async fn handle_turn_steer( &self, connection_id: u64, @@ -812,7 +642,10 @@ impl ServerRuntime { "session does not exist", ); }; - let Some(reservation) = session_handle.turn_reservation_snapshot().await else { + let Some(reservation) = self + .session_turn_reservation_snapshot(params.session_id) + .await + else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, @@ -952,14 +785,17 @@ impl ServerRuntime { ); } }; - let Some(session_handle) = self.session(params.session_id).await else { + let Some(_session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, "session does not exist", ); }; - let Some(reservation) = session_handle.turn_reservation_snapshot().await else { + let Some(reservation) = self + .session_turn_reservation_snapshot(params.session_id) + .await + else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, @@ -1024,7 +860,10 @@ impl ServerRuntime { "session does not exist", ); }; - let Some(reservation) = session_handle.turn_reservation_snapshot().await else { + let Some(reservation) = self + .session_turn_reservation_snapshot(params.session_id) + .await + else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, diff --git a/crates/server/src/runtime/handlers/turn_interrupt.rs b/crates/server/src/runtime/handlers/turn_interrupt.rs new file mode 100644 index 00000000..4325a039 --- /dev/null +++ b/crates/server/src/runtime/handlers/turn_interrupt.rs @@ -0,0 +1,244 @@ +use std::sync::Arc; +use std::time::Duration; + +use super::super::*; +use crate::persistence::build_turn_record; + +const TURN_INTERRUPT_TERMINAL_TIMEOUT: Duration = Duration::from_secs(5); + +impl ServerRuntime { + pub(crate) async fn handle_turn_interrupt( + self: &Arc, + request_id: serde_json::Value, + params: serde_json::Value, + ) -> serde_json::Value { + let params: TurnInterruptParams = match serde_json::from_value(params) { + Ok(params) => params, + Err(error) => { + return self.error_response( + request_id, + ProtocolErrorCode::InvalidParams, + format!("invalid turn/interrupt params: {error}"), + ); + } + }; + let Some(session_handle) = self.session(params.session_id).await else { + return self.error_response( + request_id, + ProtocolErrorCode::SessionNotFound, + "session does not exist", + ); + }; + + // Turns that run inline on the session actor finalize themselves when the + // cancel token fires (`finalize_executed_turn` records terminal status). + // Research (and similar) turns run on a spawned task outside the actor: + // aborting that task does not record a terminal status, so we must claim + // `active_turn` via the mailbox and finalize here. + if self.runtime_active_turn_id(params.session_id).await != Some(params.turn_id) { + if let Some(snapshot) = self.recent_terminal_turn_status(params.turn_id).await { + return self.turn_interrupt_success(request_id, params.turn_id, snapshot.status); + } + return self.error_response( + request_id, + ProtocolErrorCode::TurnNotFound, + "turn is not active", + ); + } + + let terminal_rx = self.subscribe_terminal_turn_status(params.turn_id).await; + if let Some(snapshot) = self.recent_terminal_turn_status(params.turn_id).await { + self.record_terminal_turn_status(params.turn_id, snapshot.clone()) + .await; + return self.turn_interrupt_success(request_id, params.turn_id, snapshot.status); + } + + // Cancel before any session-actor mailbox round-trip: the actor may be blocked + // waiting for a permission response and cannot process commands until cancelled. + // Cancel via a clone rather than `remove`: see the comment in + // `interrupt_child_runtime_work` for why removing here races with + // `run_turn_model_query` fetching the same token. + if let Some(cancel_token) = self + .active_turn_cancellations + .lock() + .await + .get(¶ms.session_id) + .cloned() + { + cancel_token.cancel(); + } + if let Some(task) = self.active_tasks.lock().await.remove(¶ms.session_id) { + task.abort(); + } + + let removed_len = self + .session_interactive + .clear_pending_user_inputs_for_turn(params.session_id, params.turn_id) + .await; + if removed_len > 0 { + tracing::info!( + session_id = %params.session_id, + turn_id = %params.turn_id, + removed_len, + "cleared pending request_user_input requests for interrupted turn" + ); + } + + Arc::clone(self) + .interrupt_all_child_agents(params.session_id) + .await; + + // Out-of-actor turns (research): actor is free, so we can claim active_turn. + // In-actor turns: finalize already cleared it; fall through to terminal wait. + if let Some(interrupted_turn) = session_handle.interrupt_active_turn().await.flatten() { + if interrupted_turn.turn_id != params.turn_id { + return self.error_response( + request_id, + ProtocolErrorCode::TurnNotFound, + "turn does not exist", + ); + } + return self + .finalize_claimed_interrupted_turn( + request_id, + &session_handle, + params.session_id, + interrupted_turn, + ) + .await; + } + + let snapshot = + match tokio::time::timeout(TURN_INTERRUPT_TERMINAL_TIMEOUT, terminal_rx).await { + Ok(Ok(snapshot)) => snapshot, + Ok(Err(_)) | Err(_) => { + if let Some(snapshot) = self.recent_terminal_turn_status(params.turn_id).await { + snapshot + } else { + return self.error_response( + request_id, + ProtocolErrorCode::TurnNotFound, + "turn is not active", + ); + } + } + }; + + tracing::info!( + session_id = %params.session_id, + turn_id = %params.turn_id, + status = ?snapshot.status, + "interrupted turn" + ); + + self.turn_interrupt_success(request_id, params.turn_id, snapshot.status) + } + + async fn finalize_claimed_interrupted_turn( + self: &Arc, + request_id: serde_json::Value, + session_handle: &crate::runtime::session_actor::SessionHandle, + session_id: SessionId, + interrupted_turn: TurnMetadata, + ) -> serde_json::Value { + self.clear_active_turn_runtime_handles(session_id).await; + + let deferred = session_handle.take_deferred_items().await; + if let Some((item_id, item_seq, text)) = deferred.assistant { + self.complete_item( + session_id, + interrupted_turn.turn_id, + item_id, + item_seq, + ItemKind::AgentMessage, + TurnItem::AgentMessage(TextItem { text: text.clone() }), + serde_json::json!({ "title": "Assistant", "text": text }), + ) + .await; + } + if let Some((item_id, item_seq, text)) = deferred.reasoning { + self.complete_item( + session_id, + interrupted_turn.turn_id, + item_id, + item_seq, + ItemKind::Reasoning, + TurnItem::Reasoning(TextItem { text: text.clone() }), + serde_json::json!({ "title": "Reasoning", "text": text }), + ) + .await; + } + if let Some(persistence) = session_handle.turn_persistence_snapshot().await + && let Some(record) = persistence.record + && let Err(error) = self.rollout_store.append_turn( + &record, + build_turn_record( + &interrupted_turn, + persistence.session_context, + persistence.latest_turn_context, + ), + ) + { + return self.error_response( + request_id, + ProtocolErrorCode::InternalError, + format!("failed to persist interrupted turn: {error}"), + ); + } + + tracing::info!( + session_id = %session_id, + turn_id = %interrupted_turn.turn_id, + status = ?interrupted_turn.status, + "interrupted turn" + ); + self.finalize_turn_workspace_changes(session_id, &interrupted_turn) + .await; + self.broadcast_event(ServerEvent::TurnInterrupted(TurnEventPayload { + session_id, + turn: interrupted_turn.clone(), + })) + .await; + self.broadcast_event(ServerEvent::TurnCompleted(TurnEventPayload { + session_id, + turn: interrupted_turn.clone(), + })) + .await; + self.broadcast_event(ServerEvent::SessionStatusChanged( + SessionStatusChangedPayload { + session_id, + status: SessionRuntimeStatus::Idle, + }, + )) + .await; + self.record_terminal_turn_status( + interrupted_turn.turn_id, + TerminalTurnSnapshot::from_turn(&interrupted_turn), + ) + .await; + + let runtime = Arc::clone(self); + tokio::spawn(async move { + runtime.spawn_next_turn_from_queue(session_id).await; + }); + + self.turn_interrupt_success( + request_id, + interrupted_turn.turn_id, + interrupted_turn.status, + ) + } + + fn turn_interrupt_success( + &self, + request_id: serde_json::Value, + turn_id: TurnId, + status: TurnStatus, + ) -> serde_json::Value { + serde_json::to_value(SuccessResponse { + id: request_id, + result: TurnInterruptResult { turn_id, status }, + }) + .expect("serialize turn/interrupt response") + } +} diff --git a/crates/server/src/runtime/items.rs b/crates/server/src/runtime/items.rs index 49f0e589..7e6758e1 100644 --- a/crates/server/src/runtime/items.rs +++ b/crates/server/src/runtime/items.rs @@ -20,6 +20,20 @@ impl ServerRuntime { self: &Arc, session_id: SessionId, user_input: &str, + ) { + self.maybe_prepare_title_generation_from_user_input(session_id, user_input) + .await; + self.maybe_schedule_final_title_generation(session_id, None) + .await; + } + + /// Assigns a provisional title and records the first user input without + /// calling the title model. Used at turn start while the session actor may + /// soon block on `ExecuteTurn`; final title generation runs post-turn. + pub(super) async fn maybe_prepare_title_generation_from_user_input( + self: &Arc, + session_id: SessionId, + user_input: &str, ) { self.maybe_assign_provisional_title(session_id, user_input) .await; @@ -30,6 +44,16 @@ impl ServerRuntime { let _ = session_handle .set_first_user_input_if_unset(user_input.to_string()) .await; + } + + pub(super) async fn maybe_schedule_final_title_generation( + self: &Arc, + session_id: SessionId, + first_input_override: Option, + ) { + let Some(session_handle) = self.session(session_id).await else { + return; + }; let Some(title_context) = session_handle.title_generation_context().await else { return; }; @@ -37,19 +61,27 @@ impl ServerRuntime { title_context.title_state, SessionTitleState::Unset | SessionTitleState::Provisional ); - if needs_title { - let first_input = session_handle + if !needs_title { + return; + } + let first_input = if let Some(first_input) = first_input_override { + first_input + } else { + session_handle .export_runtime_session() .await .and_then(|session| session.first_user_input) - .unwrap_or_else(|| user_input.to_string()); - let runtime = Arc::clone(self); - tokio::spawn(async move { - runtime - .maybe_generate_final_title(session_id, first_input) - .await; - }); + .unwrap_or_default() + }; + if first_input.is_empty() { + return; } + let runtime = Arc::clone(self); + tokio::spawn(async move { + runtime + .maybe_generate_final_title(session_id, first_input) + .await; + }); } pub(super) async fn maybe_assign_provisional_title( diff --git a/crates/server/src/runtime/research.rs b/crates/server/src/runtime/research.rs index d756fbad..14fde4bd 100644 --- a/crates/server/src/runtime/research.rs +++ b/crates/server/src/runtime/research.rs @@ -290,10 +290,8 @@ impl ServerRuntime { .lock() .await .insert(params.session_id, CancellationToken::new()); - self.active_turn_ids - .lock() - .await - .insert(params.session_id, turn.turn_id); + self.register_runtime_active_turn(params.session_id, turn.clone()) + .await; if let Some(connection_id) = connection_id { self.active_turn_connections .lock() @@ -301,7 +299,7 @@ impl ServerRuntime { .insert(params.session_id, connection_id); } let research_display_input = research_display_input(&display_input); - self.maybe_start_title_generation_from_user_input( + self.maybe_prepare_title_generation_from_user_input( params.session_id, &research_display_input, ) diff --git a/crates/server/src/runtime/session_actor/commands.rs b/crates/server/src/runtime/session_actor/commands.rs index 83562769..df54a65b 100644 --- a/crates/server/src/runtime/session_actor/commands.rs +++ b/crates/server/src/runtime/session_actor/commands.rs @@ -124,9 +124,6 @@ pub(crate) enum SessionCommand { SetSessionIdle { latest_turn: Option, }, - EnqueuePendingTurnInput { - item: PendingInputItem, - }, ActivateQueuedTurn { turn: TurnMetadata, turn_config: TurnConfig, diff --git a/crates/server/src/runtime/session_actor/handle.rs b/crates/server/src/runtime/session_actor/handle.rs index ca26fc59..33b445fe 100644 --- a/crates/server/src/runtime/session_actor/handle.rs +++ b/crates/server/src/runtime/session_actor/handle.rs @@ -1,5 +1,7 @@ +use std::collections::VecDeque; use std::path::PathBuf; use std::sync::Arc; +use std::sync::Mutex as StdMutex; use devo_protocol::ApprovalScopeValue; use devo_protocol::CollaborationMode; @@ -35,6 +37,8 @@ const SESSION_MAILBOX_CAPACITY: usize = 64; pub(crate) struct SessionHandle { session_id: SessionId, tx: mpsc::Sender, + pending_turn_queue: Arc>>, + max_turns: Option, } impl SessionHandle { @@ -42,13 +46,35 @@ impl SessionHandle { self.session_id } + pub(crate) fn pending_turn_queue(&self) -> Arc>> { + Arc::clone(&self.pending_turn_queue) + } + + pub(crate) fn max_turns(&self) -> Option { + self.max_turns + } + + pub(crate) fn push_pending_turn_input(&self, item: PendingInputItem) { + self.pending_turn_queue + .lock() + .expect("pending turn queue mutex should not be poisoned") + .push_back(item); + } + pub(crate) fn spawn( session_id: SessionId, state: SessionActorState, runtime: Arc, ) -> Self { + let pending_turn_queue = Arc::clone(&state.pending_turn_queue); + let max_turns = state.max_turns; let (tx, rx) = mpsc::channel(SESSION_MAILBOX_CAPACITY); - let handle = Self { session_id, tx }; + let handle = Self { + session_id, + tx, + pending_turn_queue, + max_turns, + }; tokio::spawn(super::loop_::run_session_actor(state, rx, runtime)); handle } @@ -130,6 +156,10 @@ impl SessionHandle { let _ = self.send(SessionCommand::SetActiveGoal { goal }).await; } + pub(crate) fn try_set_active_goal(&self, goal: Option) -> bool { + self.try_send(SessionCommand::SetActiveGoal { goal }) + } + pub(crate) async fn runtime_context( &self, ) -> Option> { @@ -420,10 +450,8 @@ impl SessionHandle { .await; } - pub(crate) async fn enqueue_pending_turn_input(&self, item: PendingInputItem) { - let _ = self - .send(SessionCommand::EnqueuePendingTurnInput { item }) - .await; + pub(crate) fn enqueue_pending_turn_input(&self, item: PendingInputItem) { + self.push_pending_turn_input(item); } pub(crate) async fn activate_queued_turn(&self, turn: TurnMetadata, turn_config: TurnConfig) { diff --git a/crates/server/src/runtime/session_actor/loop_.rs b/crates/server/src/runtime/session_actor/loop_.rs index cb462ed0..55385fc7 100644 --- a/crates/server/src/runtime/session_actor/loop_.rs +++ b/crates/server/src/runtime/session_actor/loop_.rs @@ -33,9 +33,26 @@ pub(super) async fn run_session_actor( } => { let session_id = request.session_id; execute_turn_in_actor(&mut state, turn_runtime.clone(), request).await; + // Interrupted turns must not auto-start continuation here: that would + // re-block the actor mailbox before the interrupting handler finishes + // (goal replace/clear/cancel). Failed turns still enter maybe_start so + // `pause_goal_continuation_after_failed_turn` can suppress looping. + // Explicit restarts go through goal handlers' maybe_start calls. + let should_auto_continue_goal = state.latest_turn.as_ref().is_some_and(|turn| { + matches!(turn.status, TurnStatus::Completed | TurnStatus::Failed) + }); let _ = reply.send(()); tokio::spawn(async move { - if !turn_runtime.chain_queued_followup_turn(session_id).await { + turn_runtime + .maybe_schedule_final_title_generation(session_id, None) + .await; + if turn_runtime.chain_queued_followup_turn(session_id).await { + return; + } + if turn_runtime.spawn_next_turn_from_queue(session_id).await { + return; + } + if should_auto_continue_goal { turn_runtime .maybe_start_goal_continuation_turn(session_id) .await; @@ -267,13 +284,6 @@ pub(super) async fn run_session_actor( Some(goal) => state.core.set_active_goal(goal), None => state.core.clear_active_goal(), }, - SessionCommand::EnqueuePendingTurnInput { item } => { - state - .pending_turn_queue - .lock() - .expect("pending turn queue mutex should not be poisoned") - .push_back(item); - } SessionCommand::ActivateQueuedTurn { turn, turn_config } => { let now = Utc::now(); apply_turn_config_to_session_summary(&mut state.summary, &turn_config); diff --git a/crates/server/src/runtime/session_actor/registry.rs b/crates/server/src/runtime/session_actor/registry.rs index 9c43d8a5..8ee7e135 100644 --- a/crates/server/src/runtime/session_actor/registry.rs +++ b/crates/server/src/runtime/session_actor/registry.rs @@ -64,6 +64,38 @@ impl ServerRuntime { summaries } + /// Reads turn reservation state, preferring runtime caches while the session + /// actor is blocked in `ExecuteTurn` (mailbox would deadlock callers). + pub(crate) async fn session_turn_reservation_snapshot( + &self, + session_id: SessionId, + ) -> Option { + if self.runtime_active_turn_id(session_id).await.is_some() + && self.active_stream_state(session_id).await.is_some() + { + let handle = self.session(session_id).await?; + let spawn = self.active_spawn_snapshot_for_session(session_id).await?; + let active_turn = self + .active_turn_metadata + .lock() + .await + .get(&session_id) + .cloned()?; + return Some(super::snapshots::TurnReservationSnapshot { + max_turns: handle.max_turns(), + active_turn: Some(active_turn), + latest_turn: spawn.parent_latest_turn, + pending_turn_queue: handle.pending_turn_queue(), + ephemeral: spawn.parent_summary.ephemeral, + parent_session_id: spawn.parent_summary.parent_session_id, + summary: spawn.parent_summary, + runtime_context: spawn.runtime_context, + }); + } + let handle = self.session(session_id).await?; + handle.turn_reservation_snapshot().await + } + pub(crate) async fn register_turn_spawn_snapshot( &self, session_id: SessionId, diff --git a/crates/server/src/runtime/turn_exec/finalize.rs b/crates/server/src/runtime/turn_exec/finalize.rs index 8a04c26a..ac2b2d7d 100644 --- a/crates/server/src/runtime/turn_exec/finalize.rs +++ b/crates/server/src/runtime/turn_exec/finalize.rs @@ -67,6 +67,7 @@ impl ServerRuntime { .await .remove(&session_id); self.active_turn_ids.lock().await.remove(&session_id); + self.active_turn_metadata.lock().await.remove(&session_id); self.active_turn_connections .lock() .await diff --git a/crates/server/src/runtime/turn_exec/followup.rs b/crates/server/src/runtime/turn_exec/followup.rs index 3765928b..29b954ef 100644 --- a/crates/server/src/runtime/turn_exec/followup.rs +++ b/crates/server/src/runtime/turn_exec/followup.rs @@ -13,17 +13,17 @@ impl ServerRuntime { pub(in crate::runtime) async fn spawn_next_turn_from_queue( self: &Arc, session_id: SessionId, - ) { + ) -> bool { let Some(queued) = self .pop_next_queued_turn_input(session_id, /*require_idle_session*/ false) .await else { - return; + return false; }; self.broadcast_updated_queue(session_id).await; let Some((turn, turn_config)) = self.prepare_queued_turn_start(session_id, &queued).await else { - return; + return false; }; if let Some((parent_session_id, parent_turn_id)) = queued.subagent_usage_owner { self.register_subagent_usage_owner(parent_session_id, session_id, parent_turn_id) @@ -51,6 +51,7 @@ impl ServerRuntime { }) .await; }); + true } /// After a turn completes, chain directly into the next queued input when present. @@ -59,7 +60,7 @@ impl ServerRuntime { session_id: SessionId, ) -> bool { let Some(queued) = self - .pop_next_queued_turn_input(session_id, /*require_idle_session*/ true) + .pop_next_queued_turn_input(session_id, /*require_idle_session*/ false) .await else { return false; diff --git a/crates/server/src/runtime/turn_exec/tests.rs b/crates/server/src/runtime/turn_exec/tests.rs index dc702aca..96a905de 100644 --- a/crates/server/src/runtime/turn_exec/tests.rs +++ b/crates/server/src/runtime/turn_exec/tests.rs @@ -38,7 +38,7 @@ fn command_progress_uses_command_execution_item_id() { ); assert_eq!( command_execution_item_id_for_progress(&pending_tool_calls, "read"), - None + Some(tool_item_id) ); assert_eq!( command_execution_item_id_for_progress(&pending_tool_calls, "missing"), diff --git a/crates/server/src/runtime/turn_reservation.rs b/crates/server/src/runtime/turn_reservation.rs index 7bccb783..586a5b17 100644 --- a/crates/server/src/runtime/turn_reservation.rs +++ b/crates/server/src/runtime/turn_reservation.rs @@ -55,6 +55,21 @@ impl ServerRuntime { self.active_turn_ids.lock().await.get(&session_id).copied() } + pub(super) async fn register_runtime_active_turn( + &self, + session_id: SessionId, + turn: TurnMetadata, + ) { + self.active_turn_ids + .lock() + .await + .insert(session_id, turn.turn_id); + self.active_turn_metadata + .lock() + .await + .insert(session_id, turn); + } + pub(super) async fn clear_active_turn_runtime_handles(&self, session_id: SessionId) { self.active_tasks.lock().await.remove(&session_id); self.active_turn_cancellations @@ -62,6 +77,7 @@ impl ServerRuntime { .await .remove(&session_id); self.active_turn_ids.lock().await.remove(&session_id); + self.active_turn_metadata.lock().await.remove(&session_id); self.active_turn_connections .lock() .await diff --git a/crates/server/tests/deep_research_e2e.rs b/crates/server/tests/deep_research_e2e.rs index 464dabfe..13a33c0f 100644 --- a/crates/server/tests/deep_research_e2e.rs +++ b/crates/server/tests/deep_research_e2e.rs @@ -1,9 +1,7 @@ -use std::io::Write; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::time::Duration; -use std::time::{SystemTime, UNIX_EPOCH}; use anyhow::Context; use anyhow::Result; @@ -901,6 +899,7 @@ async fn interrupted_research_closes_delegated_child_agent() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 33, "method": "_devo/turn/interrupt", "params": { @@ -923,6 +922,7 @@ async fn interrupted_research_closes_delegated_child_agent() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 34, "method": "_devo/agent/list", "params": { @@ -1006,6 +1006,7 @@ async fn interrupted_research_clears_pending_clarification_request() -> Result<( .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 32, "method": "_devo/turn/interrupt", "params": { @@ -1017,8 +1018,10 @@ async fn interrupted_research_clears_pending_clarification_request() -> Result<( ) .await .context("turn/interrupt response")?; + eprintln!("response: {:?}", interrupt_response); let interrupt_response: devo_server::SuccessResponse = serde_json::from_value(interrupt_response)?; + assert_eq!(interrupt_response.result.turn_id.to_string(), turn_id); let stale_response = respond_to_clarification_raw( @@ -1497,21 +1500,6 @@ fn supervisor_stream_events(request: &ModelRequest) -> Vec> let wait_id = supervisor_wait_tool_id(attempt, poll_index); let wait_content = request_tool_result_content(request, &wait_id).unwrap_or_default(); - // #region agent log - agent_debug_log( - "deep_research_e2e.rs:supervisor_stream_events", - "supervisor wait_agent poll result", - serde_json::json!({ - "attempt": attempt, - "pollIndex": poll_index, - "waitId": wait_id, - "terminal": wait_agent_result_is_terminal(&wait_content), - "failed": wait_agent_result_indicates_failure(&wait_content), - "succeeded": wait_agent_result_indicates_success(&wait_content), - }), - "A", - ); - // #endregion if wait_agent_result_indicates_failure(&wait_content) { break; } @@ -1572,10 +1560,6 @@ fn wait_agent_result_indicates_success(content: &str) -> bool { }) || wait_agent_statuses(content).any(|status| status == "completed") } -fn wait_agent_result_is_terminal(content: &str) -> bool { - wait_agent_result_indicates_failure(content) || wait_agent_result_indicates_success(content) -} - fn wait_agent_events(content: &str) -> Vec { serde_json::from_str::(content) .ok() @@ -1600,31 +1584,6 @@ fn wait_agent_statuses(content: &str) -> impl Iterator { }) } -// #region agent log -fn agent_debug_log(location: &str, message: &str, data: serde_json::Value, hypothesis_id: &str) { - let timestamp = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_millis()) - .unwrap_or(0); - let payload = serde_json::json!({ - "sessionId": "ec5e4e", - "location": location, - "message": message, - "data": data, - "hypothesisId": hypothesis_id, - "runId": "post-fix", - "timestamp": timestamp, - }); - if let Ok(mut file) = std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../../debug-ec5e4e.log")) - { - let _ = writeln!(file, "{payload}"); - } -} -// #endregion - fn supervisor_worker_message(attempt: usize) -> String { format!( "You are a delegated DeepResearch worker for attempt {attempt}.\n\n\nUse the parent-provided current date, timezone, and cwd.\n\n\n\nResearch the current official DeepSeek website domain. Use web search, keep the final report short, and include source URLs.\n\n\n\nResearch DeepSeek official website.\n\n\nReturn concise evidence notes with searches/tool calls, key findings, source table, uncertainty, and recommended citations. Do not write report files." @@ -1909,6 +1868,7 @@ async fn start_research_turn( .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 3, "method": "_devo/turn/start", "params": { @@ -1948,6 +1908,7 @@ async fn start_regular_turn_after_research( .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 30, "method": "_devo/turn/start", "params": { @@ -1987,6 +1948,7 @@ async fn queue_regular_turn_during_research( .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 31, "method": "_devo/turn/start", "params": { @@ -2391,6 +2353,7 @@ async fn respond_to_clarification_raw( .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 4, "method": "_devo/request_user_input/respond", "params": { diff --git a/crates/server/tests/end_to_end.rs b/crates/server/tests/end_to_end.rs index 49789a4c..2c4c13a4 100644 --- a/crates/server/tests/end_to_end.rs +++ b/crates/server/tests/end_to_end.rs @@ -245,6 +245,7 @@ async fn stdio_server_process_supports_handshake_and_session_start() -> Result<( format!( "{}\n", serde_json::json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/new", "params": { @@ -377,6 +378,7 @@ async fn second_stdio_server_process_proxies_to_singleton() -> Result<()> { format!( "{}\n", serde_json::json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/new", "params": { @@ -484,6 +486,7 @@ async fn websocket_listener_supports_handshake_subscription_and_turn_lifecycle() socket .send(Message::Text( serde_json::json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/new", "params": { @@ -523,6 +526,7 @@ async fn websocket_listener_supports_handshake_subscription_and_turn_lifecycle() socket .send(Message::Text( serde_json::json!({ + "jsonrpc": "2.0", "id": 3, "method": "_devo/turn/start", "params": { @@ -573,6 +577,7 @@ async fn websocket_listener_supports_handshake_subscription_and_turn_lifecycle() socket .send(Message::Text( serde_json::json!({ + "jsonrpc": "2.0", "id": 4, "method": "_devo/turn/interrupt", "params": { @@ -693,6 +698,7 @@ async fn websocket_turn_streams_final_tool_metadata_for_read_and_glob() -> Resul socket .send(Message::Text( serde_json::json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/new", "params": { @@ -728,6 +734,7 @@ async fn websocket_turn_streams_final_tool_metadata_for_read_and_glob() -> Resul socket .send(Message::Text( serde_json::json!({ + "jsonrpc": "2.0", "id": 3, "method": "_devo/turn/start", "params": { diff --git a/crates/server/tests/goal_continuation.rs b/crates/server/tests/goal_continuation.rs index 10820b2a..fc688767 100644 --- a/crates/server/tests/goal_continuation.rs +++ b/crates/server/tests/goal_continuation.rs @@ -55,6 +55,7 @@ async fn goal_token_budget_reached_after_turn_enters_budget_limited() -> Result< .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 9, "method": "_devo/goal/set", "params": { @@ -75,6 +76,7 @@ async fn goal_token_budget_reached_after_turn_enters_budget_limited() -> Result< .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 10, "method": "_devo/goal/status", "params": { @@ -124,6 +126,7 @@ async fn budget_limited_goal_pause_interrupts_pending_wrapup_turn() -> Result<() .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 110, "method": "_devo/goal/set", "params": { @@ -152,6 +155,7 @@ async fn budget_limited_goal_pause_interrupts_pending_wrapup_turn() -> Result<() .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 111, "method": "_devo/goal/set", "params": { @@ -209,6 +213,7 @@ async fn persisted_paused_goal_replays_without_continuation() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 11, "method": "_devo/goal/set", "params": { @@ -232,6 +237,7 @@ async fn persisted_paused_goal_replays_without_continuation() -> Result<()> { .handle_incoming( replayed_connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 12, "method": "_devo/goal/status", "params": { @@ -266,6 +272,7 @@ async fn persisted_active_goal_pauses_on_restart_without_continuation() -> Resul .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 13, "method": "_devo/turn/start", "params": { @@ -287,6 +294,7 @@ async fn persisted_active_goal_pauses_on_restart_without_continuation() -> Resul .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 14, "method": "_devo/goal/set", "params": { @@ -310,6 +318,7 @@ async fn persisted_active_goal_pauses_on_restart_without_continuation() -> Resul .handle_incoming( replayed_connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 15, "method": "_devo/goal/status", "params": { @@ -344,6 +353,7 @@ async fn goal_pause_interrupts_active_hidden_continuation_turn() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 16, "method": "_devo/goal/set", "params": { @@ -368,6 +378,7 @@ async fn goal_pause_interrupts_active_hidden_continuation_turn() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 17, "method": "_devo/goal/set", "params": { @@ -411,6 +422,7 @@ async fn provider_400_tool_call_adjacency_failure_pauses_goal_without_looping() .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 13, "method": "_devo/goal/set", "params": { @@ -429,6 +441,7 @@ async fn provider_400_tool_call_adjacency_failure_pauses_goal_without_looping() .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 14, "method": "_devo/goal/status", "params": { @@ -462,6 +475,7 @@ async fn goal_set_starts_hidden_continuation_turn() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 19, "method": "_devo/turn/start", "params": { @@ -482,6 +496,7 @@ async fn goal_set_starts_hidden_continuation_turn() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 20, "method": "_devo/goal/set", "params": { @@ -511,6 +526,7 @@ async fn goal_set_starts_hidden_continuation_turn() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 21, "method": "_devo/goal/set", "params": { @@ -561,6 +577,7 @@ async fn goal_set_does_not_start_continuation_while_turn_is_active() -> Result<( .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 30, "method": "_devo/turn/start", "params": { @@ -582,6 +599,7 @@ async fn goal_set_does_not_start_continuation_while_turn_is_active() -> Result<( .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 31, "method": "_devo/goal/set", "params": { @@ -600,6 +618,7 @@ async fn goal_set_does_not_start_continuation_while_turn_is_active() -> Result<( .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 32, "method": "_devo/goal/set", "params": { @@ -617,6 +636,7 @@ async fn goal_set_does_not_start_continuation_while_turn_is_active() -> Result<( .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 33, "method": "_devo/turn/interrupt", "params": { @@ -644,6 +664,7 @@ async fn goal_create_starts_hidden_continuation_turn() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 34, "method": "_devo/goal/create", "params": { @@ -676,6 +697,7 @@ async fn goal_resume_starts_hidden_continuation_turn() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 35, "method": "_devo/goal/set", "params": { @@ -693,6 +715,7 @@ async fn goal_resume_starts_hidden_continuation_turn() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 36, "method": "_devo/goal/resume", "params": { @@ -728,6 +751,7 @@ async fn queued_user_turn_runs_before_goal_continuation() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 40, "method": "_devo/turn/start", "params": { @@ -755,6 +779,7 @@ async fn queued_user_turn_runs_before_goal_continuation() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 41, "method": "_devo/turn/start", "params": { @@ -788,6 +813,7 @@ async fn queued_user_turn_runs_before_goal_continuation() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 42, "method": "_devo/goal/set", "params": { @@ -817,6 +843,7 @@ async fn queued_user_turn_runs_before_goal_continuation() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 43, "method": "_devo/goal/set", "params": { @@ -831,6 +858,7 @@ async fn queued_user_turn_runs_before_goal_continuation() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 44, "method": "_devo/turn/interrupt", "params": { diff --git a/crates/server/tests/goal_lifecycle_interrupts.rs b/crates/server/tests/goal_lifecycle_interrupts.rs index e0bde886..52611d41 100644 --- a/crates/server/tests/goal_lifecycle_interrupts.rs +++ b/crates/server/tests/goal_lifecycle_interrupts.rs @@ -46,6 +46,7 @@ async fn goal_clear_interrupts_active_hidden_continuation_turn() -> Result<()> { .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 120, "method": "_devo/goal/clear", "params": { @@ -86,6 +87,7 @@ async fn goal_complete_interrupts_active_hidden_continuation_turn() -> Result<() .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 121, "method": "_devo/goal/complete", "params": { @@ -128,6 +130,7 @@ async fn goal_cancel_interrupts_active_hidden_continuation_turn() -> Result<()> .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 126, "method": "goal/cancel", "params": { @@ -185,6 +188,7 @@ async fn replacing_goal_interrupts_old_hidden_turn_and_starts_new_goal_cleanly() .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 122, "method": "_devo/goal/status", "params": { @@ -239,6 +243,7 @@ async fn pausing_budget_limited_wrapup_preserves_budget_limited_status() -> Resu .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 123, "method": "_devo/goal/set", "params": { @@ -262,6 +267,7 @@ async fn pausing_budget_limited_wrapup_preserves_budget_limited_status() -> Resu .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 124, "method": "_devo/goal/set", "params": { @@ -295,6 +301,7 @@ async fn start_created_goal( .handle_incoming( connection_id, serde_json::json!({ + "jsonrpc": "2.0", "id": 125, "method": "_devo/goal/create", "params": { diff --git a/crates/server/tests/persistence_resume.rs b/crates/server/tests/persistence_resume.rs index fc9262a2..d7a96f0b 100644 --- a/crates/server/tests/persistence_resume.rs +++ b/crates/server/tests/persistence_resume.rs @@ -386,8 +386,8 @@ async fn runtime_generates_final_title_and_persists_explicit_rename() -> Result< .await .context("turn/start response")?; - wait_for_title_update(&mut notifications_rx, "Generated rollout title").await?; wait_for_turn_completed(&mut notifications_rx).await?; + wait_for_title_update(&mut notifications_rx, "Generated rollout title").await?; let resume_after_completion = runtime .handle_incoming( From 4b8d4ed90663120dc9fb3d65da25f9ad6d412c87 Mon Sep 17 00:00:00 2001 From: wangtsiao Date: Sat, 4 Jul 2026 16:18:55 -1000 Subject: [PATCH 16/16] refactor: extract active turn lifecycle handling into dedicated modules --- AGENTS.md | 2 +- crates/server/AGENTS.md | 41 +++ crates/server/src/runtime.rs | 26 +- crates/server/src/runtime/acp_fs.rs | 2 +- crates/server/src/runtime/acp_terminal.rs | 2 +- crates/server/src/runtime/active_turn.rs | 308 ++++++++++++++++++ crates/server/src/runtime/agents.rs | 58 ++-- .../server/src/runtime/agents/coordinator.rs | 17 +- crates/server/src/runtime/agents/lifecycle.rs | 12 +- crates/server/src/runtime/approval.rs | 45 ++- crates/server/src/runtime/connection.rs | 27 +- .../server/src/runtime/goal_continuation.rs | 39 +-- crates/server/src/runtime/goal_handlers.rs | 13 +- .../server/src/runtime/handlers/acp/prompt.rs | 13 +- .../src/runtime/handlers/acp/session.rs | 36 +- .../src/runtime/handlers/message_edit.rs | 48 ++- crates/server/src/runtime/handlers/turn.rs | 90 ++--- .../src/runtime/handlers/turn_interrupt.rs | 13 +- crates/server/src/runtime/research.rs | 22 +- .../src/runtime/research_tool_runtime.rs | 6 +- .../src/runtime/session_actor/commands.rs | 13 +- .../src/runtime/session_actor/handle.rs | 52 ++- .../server/src/runtime/session_actor/loop_.rs | 43 ++- .../server/src/runtime/session_actor/mod.rs | 8 +- .../src/runtime/session_actor/registry.rs | 50 +-- .../src/runtime/session_actor/snapshots.rs | 3 +- .../server/src/runtime/session_actor/state.rs | 4 + .../src/runtime/session_actor/turn_inline.rs | 4 + crates/server/src/runtime/subagent_usage.rs | 3 + .../server/src/runtime/turn_exec/finalize.rs | 12 +- crates/server/src/runtime/turn_exec/query.rs | 6 +- crates/server/src/runtime/turn_exec/shell.rs | 6 +- crates/server/src/runtime/turn_lifecycle.rs | 68 ++++ crates/server/src/runtime/turn_reservation.rs | 39 +-- 34 files changed, 688 insertions(+), 443 deletions(-) create mode 100644 crates/server/AGENTS.md create mode 100644 crates/server/src/runtime/active_turn.rs create mode 100644 crates/server/src/runtime/turn_lifecycle.rs diff --git a/AGENTS.md b/AGENTS.md index 3aa5b8c2..c415ceb4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -37,4 +37,4 @@ This repository is a Rust-based coding agent, currently called `devo`. - Tests that involve filesystem paths or other platform-dependent behavior MUST be platform-aware: - Use `#[cfg(windows)]` and `#[cfg(unix)]` to define platform-specific test cases when behavior differs. - Never rely on Windows-style paths being interpreted correctly on Unix, or Unix-style paths on Windows. - - Always use platform-native path formats in tests so they align with `std::path::Path` semantics. \ No newline at end of file + - Always use platform-native path formats in tests so they align with `std::path::Path` semantics. diff --git a/crates/server/AGENTS.md b/crates/server/AGENTS.md new file mode 100644 index 00000000..b8960de7 --- /dev/null +++ b/crates/server/AGENTS.md @@ -0,0 +1,41 @@ +## Server runtime concurrency + +The server runtime uses **one session actor per session**. Durable session state lives in `SessionActorState`, owned exclusively by that actor's mailbox loop (`run_session_actor`). Cross-session coordination and in-flight turn execution handles live on `ServerRuntime`. + +### Ownership and actor boundaries + +- **Mutate durable session state only through `SessionHandle` → `SessionCommand`.** Do not reach into `SessionActorState` from handlers, turn tasks, or research code except inside the actor loop or via explicit snapshot/command APIs. +- **`ActiveTurnRegistry` is the single source for in-flight turn execution handles** (cancel tokens, abort handles, connection routing, spawn snapshots, active stream state). Register on turn start. Use `clear_active_turn_interrupt_handles` during in-actor finalization so stream/spawn mirrors stay available until inline state merges; use `clear_active_turn_runtime_handles` for full teardown (interrupt handlers, session stop, out-of-actor turns). +- **Use `turn_lifecycle` helpers** (`register_active_turn_execution`, `spawn_active_turn_task`, `signal_active_turn_interrupt`) instead of touching `ActiveTurnRegistry` fields ad hoc from handlers. +- **Two turn execution paths:** + - **In-actor:** normal turns via `SessionCommand::ExecuteTurn`. + - **Out-of-actor:** research and similar work on a spawned task; use paired `BeginInlineTurn` / `EndInlineTurn` to install and merge `SessionStreamState` / inline mutations. +- **Interactive waits (approval, `request_user_input`) live in `SessionInteractiveLanes`, not the session actor.** The actor must not block the mailbox waiting on client responses. +- **Post-turn scheduling runs outside the actor.** After `ExecuteTurn` replies, continuation (queued follow-ups, goal continuation) is spawned in a background task—never inline in the mailbox handler when interrupts may still be in flight. + +### Lock usage + +- **Never hold `ServerRuntime.sessions` (or other runtime `Mutex` maps) across `.await`.** Look up the `SessionHandle`, drop the lock, then call handle methods. +- **Mutate `pending_turn_queue` only through actor commands** (`EnqueuePendingTurnInput`, `RemoveQueuedTurnInput`, `TakeQueuedTurnInputForSteer`, `PopQueuedTurnInput`). Handlers must not lock the queue directly while a session actor is running. +- **`SessionStreamState` uses `Arc>`** and is shared with the turn event stream task. Prefer actor commands for durable merges; use the stream lock only for streaming-era fields (deferred assistant/reasoning, inline turn scratch state). +- **From turn event streams, use `try_send` on the session mailbox** for fire-and-forget updates (`SetActiveGoal`, `ApplyParentUsageSnapshot`, `TouchLastActivity`). Blocking `send().await` from a stream the actor is waiting on can deadlock. +- **Interrupt/cancel:** call `signal_active_turn_interrupt` before relying on mailbox round-trips—the actor may be blocked in permission wait. + +### Turn lifecycle + +- **Reservation:** use `TryBeginActiveTurn` (idle session + empty pending queue) or turn-reservation snapshots when starting turns from handlers. +- **Terminal status:** in-actor turns finalize via `finalize_executed_turn` when the cancel token fires; out-of-actor turns must claim `active_turn` via `InterruptActiveTurn` and finalize explicitly. +- **Always record terminal turn status** (`record_terminal_turn_status`) and clear runtime handles when a turn ends or is interrupted. +- **Subagent usage:** only root sessions own a parent usage ledger; child turns publish into the parent's ledger. + +### Queues + +- **`pending_turn_queue`:** user-visible queued turns while a session is busy. Enqueue via `SessionHandle::enqueue_pending_turn_input`; pop/remove/steer via actor commands only. +- **`btw_input_queue`:** steer / between-turn input during an active turn. Enqueue only via `EnqueueBtwInput` (mailbox); clear at turn finalize. +- **After dequeuing,** broadcast queue updates and start the next turn from a spawned task (`chain_queued_followup_turn` / `spawn_next_turn_from_queue`). + +### Tests + +- **Runtime concurrency changes need integration coverage** in `crates/server/tests/`: interrupt mid-stream, queued follow-ups, goal lifecycle interrupts, persistence/resume, research. +- **Prefer waiting on observable protocol outcomes** (notifications, terminal status) over sleeping or polling internal maps. +- Follow existing test conventions: `pretty_assertions::assert_eq`, compare whole objects where possible, platform-aware paths when touching filesystem behavior. diff --git a/crates/server/src/runtime.rs b/crates/server/src/runtime.rs index dc569096..2fc1ec7c 100644 --- a/crates/server/src/runtime.rs +++ b/crates/server/src/runtime.rs @@ -148,6 +148,7 @@ use crate::workspace_changes::ActiveWorkspaceBaseline; mod acp_fs; mod acp_terminal; +mod active_turn; mod agents; mod approval; mod command_exec; @@ -182,6 +183,7 @@ mod session_interactive; mod skills; mod subagent_usage; mod turn_exec; +mod turn_lifecycle; mod turn_reservation; mod user_input; mod workspace_baseline; @@ -215,21 +217,9 @@ pub struct ServerRuntime { sessions: Mutex>, /// Interactive approval and user-input waits outside session actors. session_interactive: SessionInteractiveLanes, - /// Spawn snapshots for in-flight parent turns keyed by session and turn id. - active_spawn_snapshots: - Mutex>>>, - /// Stream state shared with the turn event task while a session actor turn runs. - active_stream_states: Mutex< - HashMap>>, - >, + /// In-flight turn execution handles keyed by session id. + active_turns: active_turn::ActiveTurnRegistry, connections: Arc>>, - active_tasks: Mutex>, - active_turn_cancellations: Mutex>, - /// Active turn ids tracked at runtime level so cancel/interrupt can avoid session-actor mailbox round-trips while a turn is blocked in permission wait. - active_turn_ids: Mutex>, - /// Full active-turn metadata for reservation snapshots while the session actor is busy in `ExecuteTurn`. - active_turn_metadata: Mutex>, - active_turn_connections: Mutex>, terminal_turn_statuses: Mutex>, acp_prompt_waiters: Mutex>>>, active_goal_continuation_turns: Mutex>, @@ -344,14 +334,8 @@ impl ServerRuntime { goal_durable_store, sessions: Mutex::new(HashMap::new()), session_interactive: SessionInteractiveLanes::default(), - active_spawn_snapshots: Mutex::new(HashMap::new()), - active_stream_states: Mutex::new(HashMap::new()), + active_turns: active_turn::ActiveTurnRegistry::default(), connections: Arc::new(Mutex::new(HashMap::new())), - active_tasks: Mutex::new(HashMap::new()), - active_turn_cancellations: Mutex::new(HashMap::new()), - active_turn_ids: Mutex::new(HashMap::new()), - active_turn_metadata: Mutex::new(HashMap::new()), - active_turn_connections: Mutex::new(HashMap::new()), terminal_turn_statuses: Mutex::new(VecDeque::new()), acp_prompt_waiters: Mutex::new(HashMap::new()), active_goal_continuation_turns: Mutex::new(HashMap::new()), diff --git a/crates/server/src/runtime/acp_fs.rs b/crates/server/src/runtime/acp_fs.rs index c4be572e..d56c9ec3 100644 --- a/crates/server/src/runtime/acp_fs.rs +++ b/crates/server/src/runtime/acp_fs.rs @@ -109,7 +109,7 @@ impl ServerRuntime { session_id: SessionId, capability: AcpFsCapability, ) -> Option { - let connection_id = *self.active_turn_connections.lock().await.get(&session_id)?; + let connection_id = self.active_turns.active_connection_id(session_id).await?; let connections = self.connections.lock().await; let connection = connections.get(&connection_id)?; client_capabilities_support_fs(&connection.acp_client_capabilities, capability) diff --git a/crates/server/src/runtime/acp_terminal.rs b/crates/server/src/runtime/acp_terminal.rs index bb24f441..300e68e7 100644 --- a/crates/server/src/runtime/acp_terminal.rs +++ b/crates/server/src/runtime/acp_terminal.rs @@ -212,7 +212,7 @@ impl ServerRuntime { &self, session_id: SessionId, ) -> Option { - let connection_id = *self.active_turn_connections.lock().await.get(&session_id)?; + let connection_id = self.active_turns.active_connection_id(session_id).await?; let connections = self.connections.lock().await; let connection = connections.get(&connection_id)?; connection diff --git a/crates/server/src/runtime/active_turn.rs b/crates/server/src/runtime/active_turn.rs new file mode 100644 index 00000000..79cc78de --- /dev/null +++ b/crates/server/src/runtime/active_turn.rs @@ -0,0 +1,308 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use devo_core::SessionId; +use devo_core::TurnId; +use tokio::sync::Mutex; +use tokio_util::sync::CancellationToken; + +use crate::turn::TurnMetadata; + +use super::session_actor::state::{SessionStreamState, SpawnSnapshot}; + +/// Per-session execution state for an in-flight turn. +/// +/// Durable session fields remain on `SessionActorState`; this struct holds only +/// runtime coordination needed while a turn blocks the actor or runs outside it. +pub(crate) struct ActiveTurnExecution { + pub turn: Option, + pub cancel_token: Option, + pub abort_handle: Option, + pub connection_id: Option, + pub spawn_snapshots: HashMap>, + pub stream: Option>>, +} + +/// Unified registry replacing the scattered `active_turn_*`, `active_tasks`, +/// `active_stream_states`, and `active_spawn_snapshots` maps. +#[derive(Default)] +pub(crate) struct ActiveTurnRegistry { + turns: Mutex>, +} + +impl ActiveTurnRegistry { + fn entry(_session_id: SessionId) -> ActiveTurnExecution { + ActiveTurnExecution { + turn: None, + cancel_token: None, + abort_handle: None, + connection_id: None, + spawn_snapshots: HashMap::new(), + stream: None, + } + } + + pub(crate) async fn active_turn_id(&self, session_id: SessionId) -> Option { + self.turns + .lock() + .await + .get(&session_id) + .and_then(|execution| execution.turn.as_ref().map(|turn| turn.turn_id)) + } + + pub(crate) async fn active_turn_metadata(&self, session_id: SessionId) -> Option { + self.turns + .lock() + .await + .get(&session_id) + .and_then(|execution| execution.turn.clone()) + } + + pub(crate) async fn active_connection_id(&self, session_id: SessionId) -> Option { + self.turns + .lock() + .await + .get(&session_id) + .and_then(|execution| execution.connection_id) + } + + pub(crate) async fn connection_map(&self) -> HashMap { + self.turns + .lock() + .await + .iter() + .filter_map(|(session_id, execution)| { + execution + .connection_id + .map(|connection_id| (*session_id, connection_id)) + }) + .collect() + } + + pub(crate) async fn cancel_token(&self, session_id: SessionId) -> Option { + self.turns + .lock() + .await + .get(&session_id) + .and_then(|execution| execution.cancel_token.clone()) + } + + pub(crate) async fn insert_cancel_token( + &self, + session_id: SessionId, + token: CancellationToken, + ) { + self.turns + .lock() + .await + .entry(session_id) + .or_insert_with(|| Self::entry(session_id)) + .cancel_token = Some(token); + } + + pub(crate) async fn register_turn_metadata(&self, session_id: SessionId, turn: TurnMetadata) { + self.turns + .lock() + .await + .entry(session_id) + .or_insert_with(|| Self::entry(session_id)) + .turn = Some(turn); + } + + pub(crate) async fn set_abort_handle( + &self, + session_id: SessionId, + abort_handle: tokio::task::AbortHandle, + ) { + if let Some(execution) = self.turns.lock().await.get_mut(&session_id) { + execution.abort_handle = Some(abort_handle); + } + } + + pub(crate) async fn set_connection_id(&self, session_id: SessionId, connection_id: u64) { + self.turns + .lock() + .await + .entry(session_id) + .or_insert_with(|| Self::entry(session_id)) + .connection_id = Some(connection_id); + } + + pub(crate) async fn abort_task(&self, session_id: SessionId) -> bool { + let abort_handle = self + .turns + .lock() + .await + .get_mut(&session_id) + .and_then(|execution| execution.abort_handle.take()); + if let Some(abort_handle) = abort_handle { + abort_handle.abort(); + true + } else { + false + } + } + + pub(crate) async fn remove_cancel_token(&self, session_id: SessionId) { + let mut turns = self.turns.lock().await; + if let Some(execution) = turns.get_mut(&session_id) { + execution.cancel_token = None; + Self::maybe_remove_empty(&mut turns, session_id); + } + } + + pub(crate) async fn remove_abort_handle(&self, session_id: SessionId) { + let mut turns = self.turns.lock().await; + if let Some(execution) = turns.get_mut(&session_id) { + execution.abort_handle = None; + Self::maybe_remove_empty(&mut turns, session_id); + } + } + + /// Clears cancellation, metadata, and connection routing while a turn ends. + /// + /// Stream state and spawn snapshots remain registered until + /// `execute_turn_in_actor` unregisters them after inline finalization. + pub(crate) async fn clear_interrupt_handles(&self, session_id: SessionId) { + let mut turns = self.turns.lock().await; + if let Some(execution) = turns.get_mut(&session_id) { + execution.turn = None; + execution.cancel_token = None; + execution.abort_handle = None; + execution.connection_id = None; + Self::maybe_remove_empty(&mut turns, session_id); + } + } + + /// Drops all runtime state for a session turn, including stream mirrors. + pub(crate) async fn clear_runtime_handles(&self, session_id: SessionId) { + self.turns.lock().await.remove(&session_id); + } + + pub(crate) async fn register_spawn_snapshot( + &self, + session_id: SessionId, + turn_id: TurnId, + snapshot: Arc, + ) { + self.turns + .lock() + .await + .entry(session_id) + .or_insert_with(|| Self::entry(session_id)) + .spawn_snapshots + .insert(turn_id, snapshot); + } + + pub(crate) async fn clear_spawn_snapshot(&self, session_id: SessionId, turn_id: TurnId) { + let mut turns = self.turns.lock().await; + if let Some(execution) = turns.get_mut(&session_id) { + execution.spawn_snapshots.remove(&turn_id); + Self::maybe_remove_empty(&mut turns, session_id); + } + } + + pub(crate) async fn spawn_snapshot_for_session( + &self, + session_id: SessionId, + ) -> Option { + let turns = self.turns.lock().await; + let execution = turns.get(&session_id)?; + execution + .spawn_snapshots + .values() + .next() + .map(|snapshot| (**snapshot).clone()) + } + + pub(crate) async fn register_stream( + &self, + session_id: SessionId, + stream: Arc>, + ) { + self.turns + .lock() + .await + .entry(session_id) + .or_insert_with(|| Self::entry(session_id)) + .stream = Some(stream); + } + + pub(crate) async fn unregister_stream(&self, session_id: SessionId) { + let mut turns = self.turns.lock().await; + if let Some(execution) = turns.get_mut(&session_id) { + execution.stream = None; + Self::maybe_remove_empty(&mut turns, session_id); + } + } + + pub(crate) async fn stream_state( + &self, + session_id: SessionId, + ) -> Option>> { + self.turns + .lock() + .await + .get(&session_id) + .and_then(|execution| execution.stream.clone()) + } + + pub(crate) async fn remove_session(&self, session_id: SessionId) { + self.turns.lock().await.remove(&session_id); + } + + pub(crate) async fn has_session(&self, session_id: SessionId) -> bool { + self.turns.lock().await.contains_key(&session_id) + } + + pub(crate) async fn cancel_token_for_host_or_session( + &self, + host_session_id: SessionId, + session_id: SessionId, + ) -> CancellationToken { + let turns = self.turns.lock().await; + turns + .get(&host_session_id) + .or_else(|| turns.get(&session_id)) + .and_then(|execution| execution.cancel_token.clone()) + .unwrap_or_else(CancellationToken::new) + } + + pub(crate) async fn drop_connection_id(&self, connection_id: u64) { + let mut turns = self.turns.lock().await; + for execution in turns.values_mut() { + if execution.connection_id == Some(connection_id) { + execution.connection_id = None; + } + } + } + + pub(crate) async fn copy_connection_from_parent( + &self, + child_session_id: SessionId, + parent_session_id: SessionId, + ) { + let connection_id = self.active_connection_id(parent_session_id).await; + if let Some(connection_id) = connection_id { + self.set_connection_id(child_session_id, connection_id) + .await; + } + } + + fn maybe_remove_empty( + turns: &mut HashMap, + session_id: SessionId, + ) { + let should_remove = turns.get(&session_id).is_some_and(|execution| { + execution.turn.is_none() + && execution.cancel_token.is_none() + && execution.abort_handle.is_none() + && execution.connection_id.is_none() + && execution.spawn_snapshots.is_empty() + && execution.stream.is_none() + }); + if should_remove { + turns.remove(&session_id); + } + } +} diff --git a/crates/server/src/runtime/agents.rs b/crates/server/src/runtime/agents.rs index 78cc1885..370d153e 100644 --- a/crates/server/src/runtime/agents.rs +++ b/crates/server/src/runtime/agents.rs @@ -463,11 +463,9 @@ impl ServerRuntime { queued_metadata, Utc::now(), ); - reservation - .pending_turn_queue - .lock() - .expect("pending turn queue mutex should not be poisoned") - .push_back(item.clone()); + session_handle + .enqueue_pending_turn_input(item.clone()) + .await; if !reservation.ephemeral && let Err(error) = self .deps @@ -544,20 +542,12 @@ impl ServerRuntime { let runtime = Arc::clone(self); let turn_for_task = turn.clone(); let turn_config_for_task = turn_config.clone(); - let cancel_token = CancellationToken::new(); if let Some(parent_session_id) = reservation.parent_session_id { - let mut connections = self.active_turn_connections.lock().await; - if let Some(connection_id) = connections.get(&parent_session_id).copied() { - connections.insert(session_id, connection_id); - } + self.active_turns + .copy_connection_from_parent(session_id, parent_session_id) + .await; } - self.active_turn_cancellations - .lock() - .await - .insert(session_id, cancel_token); - self.register_runtime_active_turn(session_id, turn.clone()) - .await; - let task = tokio::spawn(async move { + self.spawn_active_turn_task(session_id, turn.clone(), None, async move { runtime .execute_turn(ExecuteTurnRequest { session_id, @@ -570,11 +560,8 @@ impl ServerRuntime { input_mode: TurnInputMode::VisibleUserMessage, }) .await; - }); - self.active_tasks - .lock() - .await - .insert(session_id, task.abort_handle()); + }) + .await; Ok(turn) } @@ -635,7 +622,26 @@ impl ServerRuntime { Ok(route) } - async fn drain_child_mailbox_into_user_turns( + pub(in crate::runtime) async fn child_can_accept_next_turn( + &self, + session_id: SessionId, + ) -> bool { + let Some(reservation) = self.session_turn_reservation_snapshot(session_id).await else { + return false; + }; + !reservation.max_turns.is_some_and(|max_turns| { + max_turns == 0 + || reservation + .active_turn + .as_ref() + .is_some_and(|turn| turn.sequence >= max_turns) + || reservation + .latest_turn + .as_ref() + .is_some_and(|turn| turn.sequence >= max_turns) + }) + } + pub(in crate::runtime) async fn drain_child_mailbox_into_user_turns( self: &Arc, child_session_id: SessionId, ) -> Result<(), ToolCallError> { @@ -1061,11 +1067,7 @@ impl ServerRuntime { // which independently resolves and records the terminal "closed" status // once it sees `close_requested`. Track whether a turn was actually in // flight so we don't also send a duplicate closed notification below. - let had_active_turn = self - .active_turn_cancellations - .lock() - .await - .contains_key(&child_session_id); + let had_active_turn = self.active_turns.has_session(child_session_id).await; let interrupted_turn = self.interrupt_child_runtime_work(child_session_id).await; if already_terminal && interrupted_turn.is_none() { let status = self diff --git a/crates/server/src/runtime/agents/coordinator.rs b/crates/server/src/runtime/agents/coordinator.rs index e85c3be1..65ad56ad 100644 --- a/crates/server/src/runtime/agents/coordinator.rs +++ b/crates/server/src/runtime/agents/coordinator.rs @@ -38,8 +38,14 @@ impl ServerRuntime { let route = self .queue_agent_message(params.session_id, ¶ms.target, params.message) .await?; - self.drain_child_mailbox_into_user_turns(route.to_session_id) - .await?; + if self + .active_turn_id_for_session(route.to_session_id) + .await + .is_none() + { + self.drain_child_mailbox_into_user_turns(route.to_session_id) + .await?; + } Ok(devo_protocol::AgentMessageResult { delivered: true }) } @@ -59,12 +65,7 @@ impl ServerRuntime { None => self.wait_agent_cursor(params.session_id, &cursor_key).await, }; let output_buffer = self.output_buffer(params.session_id).await; - let cancel = self - .active_turn_cancellations - .lock() - .await - .get(¶ms.session_id) - .cloned(); + let cancel = self.active_turns.cancel_token(params.session_id).await; let (events, next_sequence, timed_out) = output_buffer .wait_after( effective_after_sequence, diff --git a/crates/server/src/runtime/agents/lifecycle.rs b/crates/server/src/runtime/agents/lifecycle.rs index 7cd626dc..f04e3a7e 100644 --- a/crates/server/src/runtime/agents/lifecycle.rs +++ b/crates/server/src/runtime/agents/lifecycle.rs @@ -89,18 +89,10 @@ impl ServerRuntime { // the in-flight query) cannot lose the signal by finding the map entry // already gone and falling back to a fresh, disconnected token. The // entry itself is cleaned up later by `finalize_executed_turn`. - if let Some(cancel_token) = self - .active_turn_cancellations - .lock() - .await - .get(&child_session_id) - .cloned() - { + if let Some(cancel_token) = self.active_turns.cancel_token(child_session_id).await { cancel_token.cancel(); } - if let Some(task) = self.active_tasks.lock().await.remove(&child_session_id) { - task.abort(); - } + self.active_turns.abort_task(child_session_id).await; let session_handle = self.sessions.lock().await.get(&child_session_id).cloned()?; session_handle.interrupt_active_turn().await? } diff --git a/crates/server/src/runtime/approval.rs b/crates/server/src/runtime/approval.rs index 82d62344..a87f2c76 100644 --- a/crates/server/src/runtime/approval.rs +++ b/crates/server/src/runtime/approval.rs @@ -141,18 +141,13 @@ impl ServerRuntime { request: &ToolPermissionRequest, ) -> AutoReviewOutcome { let (model, runtime_context) = { - let Some(session_handle) = self.session(session_id).await else { + let Some(reservation) = self.session_turn_reservation_snapshot(session_id).await else { return AutoReviewOutcome::AskUser; }; - let Some(summary) = session_handle.summary().await else { - return AutoReviewOutcome::AskUser; - }; - let Some(runtime_context) = session_handle.runtime_context().await else { - return AutoReviewOutcome::AskUser; - }; - + let runtime_context = reservation.runtime_context; ( - summary + reservation + .summary .model .clone() .unwrap_or_else(|| runtime_context.default_model.clone()), @@ -319,8 +314,12 @@ impl ServerRuntime { let Some(parent_session_id) = self.parent_session_id(session_id).await else { return session_id; }; - let connections = self.active_turn_connections.lock().await; - if connections.contains_key(&parent_session_id) { + if self + .active_turns + .active_connection_id(parent_session_id) + .await + .is_some() + { parent_session_id } else { session_id @@ -334,13 +333,11 @@ impl ServerRuntime { ) -> Result<(), String> { let host_session_id = self.permission_host_session_id(session_id).await; let available_scopes = approval_scopes_for_request(&request); - let connection_id = { - let connections = self.active_turn_connections.lock().await; - connections - .get(&host_session_id) - .or_else(|| connections.get(&session_id)) - .copied() - }; + let connection_id = self + .active_turns + .active_connection_id(host_session_id) + .await + .or(self.active_turns.active_connection_id(session_id).await); let Some(connection_id) = connection_id else { return Err("no ACP client connection is available for permission request".to_string()); }; @@ -369,14 +366,10 @@ impl ServerRuntime { let request_params = acp_request_permission_params(host_session_id, &request, &available_scopes); - let cancel_token = { - let cancellations = self.active_turn_cancellations.lock().await; - cancellations - .get(&host_session_id) - .or_else(|| cancellations.get(&session_id)) - .cloned() - .unwrap_or_else(CancellationToken::new) - }; + let cancel_token = self + .active_turns + .cancel_token_for_host_or_session(host_session_id, session_id) + .await; let response = match self .send_request_to_connection_cancellable( connection_id, diff --git a/crates/server/src/runtime/connection.rs b/crates/server/src/runtime/connection.rs index 5f543d33..11fb2795 100644 --- a/crates/server/src/runtime/connection.rs +++ b/crates/server/src/runtime/connection.rs @@ -170,10 +170,7 @@ impl ServerRuntime { let _ = pending.send(Err("client connection closed".to_string())); } } - self.active_turn_connections - .lock() - .await - .retain(|_, active_connection_id| *active_connection_id != connection_id); + self.active_turns.drop_connection_id(connection_id).await; self.reference_searches .lock() .await @@ -615,11 +612,7 @@ impl ServerRuntime { return; }; let session_id = payload.context.session_id; - let connection_id = { - let active_turn_connections = self.active_turn_connections.lock().await; - active_turn_connections.get(&session_id).copied() - }; - let Some(connection_id) = connection_id else { + let Some(connection_id) = self.active_turns.active_connection_id(session_id).await else { self.broadcast_event(event.clone()).await; return; }; @@ -689,7 +682,7 @@ impl ServerRuntime { let method = event.method_name(); let session_id = event.session_id(); let child_parent_by_session = self.child_parent_by_session().await; - let active_turn_connections = self.active_turn_connections.lock().await.clone(); + let active_turn_connections = self.active_turns.connection_map().await; let notifications = { let mut connections = self.connections.lock().await; connections @@ -1529,10 +1522,9 @@ mod tests { .subscribe_connection_to_session(observer_connection_id, session_id, None) .await; runtime - .active_turn_connections - .lock() - .await - .insert(session_id, owner_connection_id); + .active_turns + .set_connection_id(session_id, owner_connection_id) + .await; runtime .broadcast_event(ServerEvent::ItemDelta { @@ -1588,10 +1580,9 @@ mod tests { .subscribe_connection_to_session(watcher_connection_id, session_id, None) .await; runtime - .active_turn_connections - .lock() - .await - .insert(session_id, owner_connection_id); + .active_turns + .set_connection_id(session_id, owner_connection_id) + .await; runtime .broadcast_event(ServerEvent::ItemDelta { diff --git a/crates/server/src/runtime/goal_continuation.rs b/crates/server/src/runtime/goal_continuation.rs index 96568a8f..df1cab7f 100644 --- a/crates/server/src/runtime/goal_continuation.rs +++ b/crates/server/src/runtime/goal_continuation.rs @@ -179,8 +179,6 @@ impl ServerRuntime { if !still_reserved { false } else { - let mut cancellations = self.active_turn_cancellations.lock().await; - let mut active_tasks = self.active_tasks.lock().await; let still_active_turn = session_handle .active_turn_id() .await @@ -188,7 +186,9 @@ impl ServerRuntime { if !still_active_turn { false } else { - cancellations.insert(session_id, cancel_token); + self.active_turns + .insert_cancel_token(session_id, cancel_token.clone()) + .await; self.register_runtime_active_turn(session_id, turn.clone()) .await; let task = tokio::spawn(async move { @@ -207,7 +207,9 @@ impl ServerRuntime { }) .await; }); - active_tasks.insert(session_id, task.abort_handle()); + self.active_turns + .set_abort_handle(session_id, task.abort_handle()) + .await; true } } @@ -240,11 +242,10 @@ impl ServerRuntime { { return false; } - reservation - .pending_turn_queue - .lock() - .expect("pending turn queue mutex should not be poisoned") - .is_empty() + let Some(snapshot) = session_handle.pending_queue_snapshot().await else { + return false; + }; + snapshot.pending_count == 0 } async fn goal_continuation_candidate( @@ -427,11 +428,8 @@ impl ServerRuntime { session_id }; if let Some(session_id) = session_id { - self.active_turn_cancellations - .lock() - .await - .remove(&session_id); - self.active_tasks.lock().await.remove(&session_id); + self.active_turns.remove_cancel_token(session_id).await; + self.active_turns.remove_abort_handle(session_id).await; } self.goal_continuation_turn_goals .lock() @@ -489,18 +487,7 @@ impl ServerRuntime { // Cancel via a clone rather than `remove`: see the comment in // `interrupt_child_runtime_work` for why removing here races with // `run_turn_model_query` fetching the same token. - if let Some(cancel_token) = self - .active_turn_cancellations - .lock() - .await - .get(&session_id) - .cloned() - { - cancel_token.cancel(); - } - if let Some(task) = self.active_tasks.lock().await.remove(&session_id) { - task.abort(); - } + self.signal_active_turn_interrupt(session_id).await; let completed = match tokio::time::timeout(std::time::Duration::from_secs(5), terminal_rx).await { diff --git a/crates/server/src/runtime/goal_handlers.rs b/crates/server/src/runtime/goal_handlers.rs index d9fd600e..8c3ed308 100644 --- a/crates/server/src/runtime/goal_handlers.rs +++ b/crates/server/src/runtime/goal_handlers.rs @@ -57,12 +57,8 @@ impl ServerRuntime { .await; } self.sync_core_session_goal(session_id, session_goal).await; - self.schedule_goal_followup_work( - session_id, - Some(title_input), - should_continue, - ) - .await; + self.schedule_goal_followup_work(session_id, Some(title_input), should_continue) + .await; result } Err(e) => self.error_response( @@ -559,10 +555,7 @@ impl ServerRuntime { let runtime = Arc::clone(self); tokio::spawn(async move { runtime - .maybe_prepare_title_generation_from_user_input( - session_id, - &title_input, - ) + .maybe_prepare_title_generation_from_user_input(session_id, &title_input) .await; }); } else { diff --git a/crates/server/src/runtime/handlers/acp/prompt.rs b/crates/server/src/runtime/handlers/acp/prompt.rs index 67ab73bc..b22608af 100644 --- a/crates/server/src/runtime/handlers/acp/prompt.rs +++ b/crates/server/src/runtime/handlers/acp/prompt.rs @@ -141,18 +141,7 @@ impl ServerRuntime { tracing::debug!(session_id = %params.session_id, "session/cancel had no active turn"); return; }; - if let Some(cancel_token) = self - .active_turn_cancellations - .lock() - .await - .get(¶ms.session_id) - .cloned() - { - cancel_token.cancel(); - } - if let Some(task) = self.active_tasks.lock().await.remove(¶ms.session_id) { - task.abort(); - } + self.signal_active_turn_interrupt(params.session_id).await; let runtime = Arc::clone(self); tokio::spawn(async move { let _ = runtime diff --git a/crates/server/src/runtime/handlers/acp/session.rs b/crates/server/src/runtime/handlers/acp/session.rs index a2c74a4f..1d75c45a 100644 --- a/crates/server/src/runtime/handlers/acp/session.rs +++ b/crates/server/src/runtime/handlers/acp/session.rs @@ -485,43 +485,13 @@ impl ServerRuntime { if self.recent_terminal_turn_status(turn_id).await.is_some() { return; } - if let Some(cancel_token) = self - .active_turn_cancellations - .lock() - .await - .get(&session_id) - .cloned() - { - cancel_token.cancel(); - } - if let Some(task) = self.active_tasks.lock().await.remove(&session_id) { - task.abort(); - } + self.signal_active_turn_interrupt(session_id).await; let _ = tokio::time::timeout(std::time::Duration::from_secs(5), receiver).await; } async fn clear_deleted_session_runtime_state(&self, session_id: SessionId) { - if let Some(task) = self.active_tasks.lock().await.remove(&session_id) { - task.abort(); - } - // Cancel via a clone rather than `remove`: see the comment in - // `interrupt_child_runtime_work` for why removing here races with - // `run_turn_model_query` fetching the same token. - if let Some(cancellation) = self - .active_turn_cancellations - .lock() - .await - .get(&session_id) - .cloned() - { - cancellation.cancel(); - } - self.active_turn_ids.lock().await.remove(&session_id); - self.active_turn_metadata.lock().await.remove(&session_id); - self.active_turn_connections - .lock() - .await - .remove(&session_id); + self.signal_active_turn_interrupt(session_id).await; + self.active_turns.clear_runtime_handles(session_id).await; if let Some(turn_id) = self .active_goal_continuation_turns .lock() diff --git a/crates/server/src/runtime/handlers/message_edit.rs b/crates/server/src/runtime/handlers/message_edit.rs index b8a60879..e0fde744 100644 --- a/crates/server/src/runtime/handlers/message_edit.rs +++ b/crates/server/src/runtime/handlers/message_edit.rs @@ -462,12 +462,6 @@ impl ServerRuntime { ) .await; - self.active_turn_cancellations - .lock() - .await - .insert(params.session_id, CancellationToken::new()); - self.register_runtime_active_turn(params.session_id, replacement_turn.clone()) - .await; let runtime = Arc::clone(self); let replacement_turn_for_task = replacement_turn.clone(); let turn_config_for_task = turn_config.clone(); @@ -492,26 +486,28 @@ impl ServerRuntime { updated_at: now.timestamp(), }) }; - let task = tokio::spawn(async move { - runtime - .execute_turn(ExecuteTurnRequest { - session_id, - turn: replacement_turn_for_task, - turn_config: turn_config_for_task, - display_input: display_input_for_task, - input: input_for_task, - input_messages: input_messages_for_task, - collaboration_mode, - input_mode: TurnInputMode::HiddenGoalContinuation { - goal: replacement_goal, - }, - }) - .await; - }); - self.active_tasks - .lock() - .await - .insert(params.session_id, task.abort_handle()); + self.spawn_active_turn_task( + params.session_id, + replacement_turn.clone(), + None, + async move { + runtime + .execute_turn(ExecuteTurnRequest { + session_id, + turn: replacement_turn_for_task, + turn_config: turn_config_for_task, + display_input: display_input_for_task, + input: input_for_task, + input_messages: input_messages_for_task, + collaboration_mode, + input_mode: TurnInputMode::HiddenGoalContinuation { + goal: replacement_goal, + }, + }) + .await; + }, + ) + .await; self.broadcast_event(ServerEvent::MessageEditRecorded( crate::MessageEditRecordedPayload { session_id: params.session_id, diff --git a/crates/server/src/runtime/handlers/turn.rs b/crates/server/src/runtime/handlers/turn.rs index bf932324..b27ade67 100644 --- a/crates/server/src/runtime/handlers/turn.rs +++ b/crates/server/src/runtime/handlers/turn.rs @@ -208,7 +208,9 @@ impl ServerRuntime { now, ); let queued_input_id = item.id; - session_handle.push_pending_turn_input(item.clone()); + session_handle + .enqueue_pending_turn_input(item.clone()) + .await; if !reservation.ephemeral && let Err(err) = self.deps @@ -342,18 +344,6 @@ impl ServerRuntime { }, )) .await; - self.active_turn_cancellations - .lock() - .await - .insert(params.session_id, CancellationToken::new()); - self.register_runtime_active_turn(params.session_id, turn.clone()) - .await; - if let Some(connection_id) = connection_id { - self.active_turn_connections - .lock() - .await - .insert(params.session_id, connection_id); - } let runtime = Arc::clone(self); let turn_for_task = turn.clone(); let display_input_for_task = display_input.clone(); @@ -362,7 +352,7 @@ impl ServerRuntime { let turn_config_for_task = turn_config.clone(); let collaboration_mode = params.collaboration_mode; let session_id = params.session_id; - let task = tokio::spawn(async move { + self.spawn_active_turn_task(params.session_id, turn.clone(), connection_id, async move { runtime .execute_turn(ExecuteTurnRequest { session_id, @@ -375,11 +365,8 @@ impl ServerRuntime { input_mode: TurnInputMode::VisibleUserMessage, }) .await; - }); - self.active_tasks - .lock() - .await - .insert(params.session_id, task.abort_handle()); + }) + .await; tracing::info!( session_id = %params.session_id, @@ -572,27 +559,12 @@ impl ServerRuntime { let runtime = Arc::clone(self); let command_for_task = command.clone(); let turn_for_task = turn.clone(); - self.active_turn_cancellations - .lock() - .await - .insert(params.session_id, CancellationToken::new()); - self.register_runtime_active_turn(params.session_id, turn.clone()) - .await; - if let Some(connection_id) = connection_id { - self.active_turn_connections - .lock() - .await - .insert(params.session_id, connection_id); - } - let task = tokio::spawn(async move { + self.spawn_active_turn_task(params.session_id, turn.clone(), connection_id, async move { runtime .execute_shell_command_turn(params.session_id, turn_for_task, command_for_task, cwd) .await; - }); - self.active_tasks - .lock() - .await - .insert(params.session_id, task.abort_handle()); + }) + .await; serde_json::to_value(SuccessResponse { id: request_id, @@ -635,7 +607,7 @@ impl ServerRuntime { "turn steer input is empty", ); }; - let Some(session_handle) = self.session(params.session_id).await else { + let Some(_session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, @@ -728,7 +700,11 @@ impl ServerRuntime { None, chrono::Utc::now(), ); - session_handle.enqueue_btw_input(item.clone()).await; + reservation + .btw_input_queue + .lock() + .expect("btw input queue mutex should not be poisoned") + .push_back(item.clone()); if !reservation.ephemeral && let Err(err) = self @@ -785,7 +761,7 @@ impl ServerRuntime { ); } }; - let Some(_session_handle) = self.session(params.session_id).await else { + let Some(session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, @@ -802,16 +778,11 @@ impl ServerRuntime { "session does not exist", ); }; - let pending_turn_queue = reservation.pending_turn_queue; let is_ephemeral = reservation.ephemeral; - let removed = { - let mut queue = pending_turn_queue - .lock() - .expect("pending turn queue mutex should not be poisoned"); - let before = queue.len(); - queue.retain(|item| item.id != params.queued_input_id); - queue.len() != before - }; + let removed = session_handle + .remove_queued_turn_input(params.queued_input_id) + .await + .unwrap_or(false); if removed && !is_ephemeral && let Err(error) = self.deps.db.remove_pending_by_id( @@ -853,7 +824,7 @@ impl ServerRuntime { ); } }; - let Some(session_handle) = self.session(params.session_id).await else { + let Some(_session_handle) = self.session(params.session_id).await else { return self.error_response( request_id, ProtocolErrorCode::SessionNotFound, @@ -892,10 +863,10 @@ impl ServerRuntime { ); } let turn_id = active_turn.turn_id; - let pending_turn_queue = reservation.pending_turn_queue; let is_ephemeral = reservation.ephemeral; - let (item, display_input) = { - let mut queue = pending_turn_queue + let queued = { + let mut queue = reservation + .pending_turn_queue .lock() .expect("pending turn queue mutex should not be poisoned"); let Some(index) = queue @@ -905,7 +876,7 @@ impl ServerRuntime { return self.error_response( request_id, ProtocolErrorCode::InvalidParams, - "queued input does not exist", + "queued input does not exist or cannot be steered", ); }; let display_input = match &queue[index].kind { @@ -915,17 +886,22 @@ impl ServerRuntime { return self.error_response( request_id, ProtocolErrorCode::InvalidParams, - "queued input cannot be steered", + "queued input does not exist or cannot be steered", ); } }; let item = queue .remove(index) .expect("queued item index should remain valid"); - (item, display_input) + (display_input, item) }; + let (display_input, item) = queued; - session_handle.enqueue_btw_input(item.clone()).await; + reservation + .btw_input_queue + .lock() + .expect("btw input queue mutex should not be poisoned") + .push_back(item.clone()); if !is_ephemeral { if let Err(error) = self.deps.db.remove_pending_by_id( diff --git a/crates/server/src/runtime/handlers/turn_interrupt.rs b/crates/server/src/runtime/handlers/turn_interrupt.rs index 4325a039..7c455925 100644 --- a/crates/server/src/runtime/handlers/turn_interrupt.rs +++ b/crates/server/src/runtime/handlers/turn_interrupt.rs @@ -58,18 +58,7 @@ impl ServerRuntime { // Cancel via a clone rather than `remove`: see the comment in // `interrupt_child_runtime_work` for why removing here races with // `run_turn_model_query` fetching the same token. - if let Some(cancel_token) = self - .active_turn_cancellations - .lock() - .await - .get(¶ms.session_id) - .cloned() - { - cancel_token.cancel(); - } - if let Some(task) = self.active_tasks.lock().await.remove(¶ms.session_id) { - task.abort(); - } + self.signal_active_turn_interrupt(params.session_id).await; let removed_len = self .session_interactive diff --git a/crates/server/src/runtime/research.rs b/crates/server/src/runtime/research.rs index 14fde4bd..cc5bd45b 100644 --- a/crates/server/src/runtime/research.rs +++ b/crates/server/src/runtime/research.rs @@ -7,7 +7,6 @@ use devo_protocol::RequestUserInputQuestion; use devo_protocol::ServerRequestKind; use serde::Deserialize; use tokio::sync::mpsc; -use tokio_util::sync::CancellationToken; use super::research_capture::{ ClarificationQueryCapture, FinalReportWrite, ResearchArtifactQueryCapture, @@ -286,18 +285,6 @@ impl ServerRuntime { .await; } - self.active_turn_cancellations - .lock() - .await - .insert(params.session_id, CancellationToken::new()); - self.register_runtime_active_turn(params.session_id, turn.clone()) - .await; - if let Some(connection_id) = connection_id { - self.active_turn_connections - .lock() - .await - .insert(params.session_id, connection_id); - } let research_display_input = research_display_input(&display_input); self.maybe_prepare_title_generation_from_user_input( params.session_id, @@ -352,7 +339,7 @@ impl ServerRuntime { let turn_for_task = turn.clone(); let display_input_for_task = research_display_input.clone(); let runtime_context_for_task = Arc::clone(&runtime_context); - let task = tokio::spawn(async move { + self.spawn_active_turn_task(params.session_id, turn.clone(), connection_id, async move { runtime .execute_research_turn(ExecuteResearchTurnInput { session_id: params.session_id, @@ -364,11 +351,8 @@ impl ServerRuntime { cwd: effective_cwd, }) .await; - }); - self.active_tasks - .lock() - .await - .insert(params.session_id, task.abort_handle()); + }) + .await; serde_json::to_value(SuccessResponse { id: request_id, diff --git a/crates/server/src/runtime/research_tool_runtime.rs b/crates/server/src/runtime/research_tool_runtime.rs index ff7e8d91..3ae37b72 100644 --- a/crates/server/src/runtime/research_tool_runtime.rs +++ b/crates/server/src/runtime/research_tool_runtime.rs @@ -47,11 +47,9 @@ impl ServerRuntime { .provider_http .clone(); let turn_cancel_token = self - .active_turn_cancellations - .lock() + .active_turns + .cancel_token(session_id) .await - .get(&session_id) - .cloned() .unwrap_or_else(CancellationToken::new); let tool_execution_start_runtime = Arc::clone(self); Ok(ToolRuntime::new_with_context_and_options( diff --git a/crates/server/src/runtime/session_actor/commands.rs b/crates/server/src/runtime/session_actor/commands.rs index df54a65b..2d59f389 100644 --- a/crates/server/src/runtime/session_actor/commands.rs +++ b/crates/server/src/runtime/session_actor/commands.rs @@ -41,9 +41,6 @@ pub(crate) enum SessionCommand { GetCollaborationMode { reply: oneshot::Sender, }, - GetRuntimeContext { - reply: oneshot::Sender>, - }, GetParentSessionId { reply: oneshot::Sender>, }, @@ -70,6 +67,13 @@ pub(crate) enum SessionCommand { require_idle_session: bool, reply: oneshot::Sender>, }, + EnqueuePendingTurnInput { + item: PendingInputItem, + }, + RemoveQueuedTurnInput { + queued_input_id: devo_core::PendingInputId, + reply: oneshot::Sender, + }, GetActiveTurnId { reply: oneshot::Sender>, }, @@ -156,9 +160,6 @@ pub(crate) enum SessionCommand { cwd: std::path::PathBuf, runtime_context: Arc, }, - EnqueueBtwInput { - item: devo_protocol::PendingInputItem, - }, UpdateSessionMetadata { model: Option, model_binding_id: Option, diff --git a/crates/server/src/runtime/session_actor/handle.rs b/crates/server/src/runtime/session_actor/handle.rs index 33b445fe..2c18320a 100644 --- a/crates/server/src/runtime/session_actor/handle.rs +++ b/crates/server/src/runtime/session_actor/handle.rs @@ -1,7 +1,5 @@ -use std::collections::VecDeque; use std::path::PathBuf; use std::sync::Arc; -use std::sync::Mutex as StdMutex; use devo_protocol::ApprovalScopeValue; use devo_protocol::CollaborationMode; @@ -37,7 +35,6 @@ const SESSION_MAILBOX_CAPACITY: usize = 64; pub(crate) struct SessionHandle { session_id: SessionId, tx: mpsc::Sender, - pending_turn_queue: Arc>>, max_turns: Option, } @@ -46,19 +43,14 @@ impl SessionHandle { self.session_id } - pub(crate) fn pending_turn_queue(&self) -> Arc>> { - Arc::clone(&self.pending_turn_queue) - } - pub(crate) fn max_turns(&self) -> Option { self.max_turns } - pub(crate) fn push_pending_turn_input(&self, item: PendingInputItem) { - self.pending_turn_queue - .lock() - .expect("pending turn queue mutex should not be poisoned") - .push_back(item); + pub(crate) async fn enqueue_pending_turn_input(&self, item: PendingInputItem) { + let _ = self + .send(SessionCommand::EnqueuePendingTurnInput { item }) + .await; } pub(crate) fn spawn( @@ -66,13 +58,11 @@ impl SessionHandle { state: SessionActorState, runtime: Arc, ) -> Self { - let pending_turn_queue = Arc::clone(&state.pending_turn_queue); let max_turns = state.max_turns; let (tx, rx) = mpsc::channel(SESSION_MAILBOX_CAPACITY); let handle = Self { session_id, tx, - pending_turn_queue, max_turns, }; tokio::spawn(super::loop_::run_session_actor(state, rx, runtime)); @@ -160,19 +150,6 @@ impl SessionHandle { self.try_send(SessionCommand::SetActiveGoal { goal }) } - pub(crate) async fn runtime_context( - &self, - ) -> Option> { - let (reply_tx, reply_rx) = oneshot::channel(); - if !self - .send(SessionCommand::GetRuntimeContext { reply: reply_tx }) - .await - { - return None; - } - reply_rx.await.ok() - } - pub(crate) async fn parent_session_id(&self) -> Option> { let (reply_tx, reply_rx) = oneshot::channel(); if !self @@ -450,8 +427,21 @@ impl SessionHandle { .await; } - pub(crate) fn enqueue_pending_turn_input(&self, item: PendingInputItem) { - self.push_pending_turn_input(item); + pub(crate) async fn remove_queued_turn_input( + &self, + queued_input_id: devo_core::PendingInputId, + ) -> Option { + let (reply_tx, reply_rx) = oneshot::channel(); + if !self + .send(SessionCommand::RemoveQueuedTurnInput { + queued_input_id, + reply: reply_tx, + }) + .await + { + return None; + } + reply_rx.await.ok() } pub(crate) async fn activate_queued_turn(&self, turn: TurnMetadata, turn_config: TurnConfig) { @@ -544,10 +534,6 @@ impl SessionHandle { .await; } - pub(crate) async fn enqueue_btw_input(&self, item: devo_protocol::PendingInputItem) { - let _ = self.send(SessionCommand::EnqueueBtwInput { item }).await; - } - pub(crate) async fn update_session_metadata( &self, model: Option, diff --git a/crates/server/src/runtime/session_actor/loop_.rs b/crates/server/src/runtime/session_actor/loop_.rs index 55385fc7..10d3e1f3 100644 --- a/crates/server/src/runtime/session_actor/loop_.rs +++ b/crates/server/src/runtime/session_actor/loop_.rs @@ -52,6 +52,17 @@ pub(super) async fn run_session_actor( if turn_runtime.spawn_next_turn_from_queue(session_id).await { return; } + if turn_runtime + .child_parent_and_path(session_id) + .await + .is_some() + && turn_runtime.child_can_accept_next_turn(session_id).await + { + let _ = turn_runtime + .drain_child_mailbox_into_user_turns(session_id) + .await; + return; + } if should_auto_continue_goal { turn_runtime .maybe_start_goal_continuation_turn(session_id) @@ -72,9 +83,6 @@ pub(super) async fn run_session_actor( SessionCommand::GetCollaborationMode { reply } => { let _ = reply.send(state.core.collaboration_mode); } - SessionCommand::GetRuntimeContext { reply } => { - let _ = reply.send(Arc::clone(&state.runtime_context)); - } SessionCommand::GetParentSessionId { reply } => { let _ = reply.send(state.parent_session_id()); } @@ -83,11 +91,12 @@ pub(super) async fn run_session_actor( max_turns: state.max_turns, active_turn: state.active_turn.clone(), latest_turn: state.latest_turn.clone(), - pending_turn_queue: Arc::clone(&state.pending_turn_queue), ephemeral: state.summary.ephemeral, parent_session_id: state.parent_session_id(), summary: state.summary.clone(), runtime_context: Arc::clone(&state.runtime_context), + pending_turn_queue: Arc::clone(&state.pending_turn_queue), + btw_input_queue: Arc::clone(&state.btw_input_queue), }); } SessionCommand::GetHookContextSnapshot { reply } => { @@ -162,6 +171,25 @@ pub(super) async fn run_session_actor( let popped = queue.pop_front().and_then(pop_queued_turn_input_data); let _ = reply.send(popped); } + SessionCommand::EnqueuePendingTurnInput { item } => { + state + .pending_turn_queue + .lock() + .expect("pending turn queue mutex should not be poisoned") + .push_back(item); + } + SessionCommand::RemoveQueuedTurnInput { + queued_input_id, + reply, + } => { + let mut queue = state + .pending_turn_queue + .lock() + .expect("pending turn queue mutex should not be poisoned"); + let before = queue.len(); + queue.retain(|item| item.id != queued_input_id); + let _ = reply.send(queue.len() != before); + } SessionCommand::GetActiveTurnId { reply } => { let _ = reply.send(state.active_turn.as_ref().map(|turn| turn.turn_id)); } @@ -353,13 +381,6 @@ pub(super) async fn run_session_actor( state.core.cwd = cwd.clone(); state.summary.cwd = cwd; } - SessionCommand::EnqueueBtwInput { item } => { - state - .btw_input_queue - .lock() - .expect("btw input queue mutex should not be poisoned") - .push_back(item); - } SessionCommand::UpdateSessionMetadata { model, model_binding_id, diff --git a/crates/server/src/runtime/session_actor/mod.rs b/crates/server/src/runtime/session_actor/mod.rs index 1526a4bb..fccbda9f 100644 --- a/crates/server/src/runtime/session_actor/mod.rs +++ b/crates/server/src/runtime/session_actor/mod.rs @@ -1,3 +1,9 @@ +// Per-session actor: single-writer for durable session state. +// +// Long-running turns still execute inside or beside this actor today. While a +// turn is in flight, transient execution state lives in ActiveTurnRegistry and +// merges back through actor commands when the turn completes. + mod commands; mod handle; mod loop_; @@ -8,4 +14,4 @@ mod turn; mod turn_inline; pub(crate) use handle::SessionHandle; -pub(crate) use state::{SessionActorState, SpawnSnapshot}; +pub(crate) use state::SessionActorState; diff --git a/crates/server/src/runtime/session_actor/registry.rs b/crates/server/src/runtime/session_actor/registry.rs index 8ee7e135..8286a8d2 100644 --- a/crates/server/src/runtime/session_actor/registry.rs +++ b/crates/server/src/runtime/session_actor/registry.rs @@ -37,7 +37,7 @@ impl ServerRuntime { ) -> Option { let handle = self.sessions.lock().await.remove(&session_id)?; self.session_interactive.clear_session(session_id).await; - self.active_spawn_snapshots.lock().await.remove(&session_id); + self.active_turns.remove_session(session_id).await; Some(handle) } @@ -75,21 +75,17 @@ impl ServerRuntime { { let handle = self.session(session_id).await?; let spawn = self.active_spawn_snapshot_for_session(session_id).await?; - let active_turn = self - .active_turn_metadata - .lock() - .await - .get(&session_id) - .cloned()?; + let active_turn = self.active_turns.active_turn_metadata(session_id).await?; return Some(super::snapshots::TurnReservationSnapshot { max_turns: handle.max_turns(), active_turn: Some(active_turn), latest_turn: spawn.parent_latest_turn, - pending_turn_queue: handle.pending_turn_queue(), ephemeral: spawn.parent_summary.ephemeral, parent_session_id: spawn.parent_summary.parent_session_id, summary: spawn.parent_summary, runtime_context: spawn.runtime_context, + pending_turn_queue: spawn.pending_turn_queue, + btw_input_queue: spawn.btw_input_queue, }); } let handle = self.session(session_id).await?; @@ -102,22 +98,15 @@ impl ServerRuntime { turn_id: TurnId, snapshot: Arc, ) { - self.active_spawn_snapshots - .lock() - .await - .entry(session_id) - .or_default() - .insert(turn_id, snapshot); + self.active_turns + .register_spawn_snapshot(session_id, turn_id, snapshot) + .await; } pub(crate) async fn clear_turn_spawn_snapshot(&self, session_id: SessionId, turn_id: TurnId) { - let mut snapshots = self.active_spawn_snapshots.lock().await; - if let Some(turns) = snapshots.get_mut(&session_id) { - turns.remove(&turn_id); - if turns.is_empty() { - snapshots.remove(&session_id); - } - } + self.active_turns + .clear_spawn_snapshot(session_id, turn_id) + .await; } /// Snapshot registered at turn start while the session actor is busy executing. @@ -125,9 +114,9 @@ impl ServerRuntime { &self, session_id: SessionId, ) -> Option { - let snapshots = self.active_spawn_snapshots.lock().await; - let turns = snapshots.get(&session_id)?; - turns.values().next().map(|snapshot| (**snapshot).clone()) + self.active_turns + .spawn_snapshot_for_session(session_id) + .await } pub(crate) async fn register_active_stream( @@ -135,25 +124,18 @@ impl ServerRuntime { session_id: SessionId, stream: Arc>, ) { - self.active_stream_states - .lock() - .await - .insert(session_id, stream); + self.active_turns.register_stream(session_id, stream).await; } pub(crate) async fn unregister_active_stream(&self, session_id: SessionId) { - self.active_stream_states.lock().await.remove(&session_id); + self.active_turns.unregister_stream(session_id).await; } pub(crate) async fn active_stream_state( &self, session_id: SessionId, ) -> Option>> { - self.active_stream_states - .lock() - .await - .get(&session_id) - .cloned() + self.active_turns.stream_state(session_id).await } pub(crate) async fn session_record_snapshot( diff --git a/crates/server/src/runtime/session_actor/snapshots.rs b/crates/server/src/runtime/session_actor/snapshots.rs index db35ed5b..b8143a86 100644 --- a/crates/server/src/runtime/session_actor/snapshots.rs +++ b/crates/server/src/runtime/session_actor/snapshots.rs @@ -23,11 +23,12 @@ pub(crate) struct TurnReservationSnapshot { pub(crate) max_turns: Option, pub(crate) active_turn: Option, pub(crate) latest_turn: Option, - pub(crate) pending_turn_queue: Arc>>, pub(crate) ephemeral: bool, pub(crate) parent_session_id: Option, pub(crate) summary: SessionMetadata, pub(crate) runtime_context: Arc, + pub(crate) pending_turn_queue: Arc>>, + pub(crate) btw_input_queue: Arc>>, } /// Hook runner inputs derived from session actor state. diff --git a/crates/server/src/runtime/session_actor/state.rs b/crates/server/src/runtime/session_actor/state.rs index 21f9db92..b416e898 100644 --- a/crates/server/src/runtime/session_actor/state.rs +++ b/crates/server/src/runtime/session_actor/state.rs @@ -28,6 +28,8 @@ pub(crate) struct SpawnSnapshot { pub(crate) parent_active_turn_id: Option, pub(crate) parent_tool_registry: Option>, pub(crate) runtime_context: Arc, + pub(crate) pending_turn_queue: Arc>>, + pub(crate) btw_input_queue: Arc>>, } /// Approval caches cloned at turn start for permission checks while the actor @@ -128,6 +130,8 @@ impl SessionActorState { .or_else(|| self.latest_turn.as_ref().map(|turn| turn.turn_id)), parent_tool_registry: self.tool_registry.clone(), runtime_context: Arc::clone(&self.runtime_context), + pending_turn_queue: Arc::clone(&self.pending_turn_queue), + btw_input_queue: Arc::clone(&self.btw_input_queue), } } diff --git a/crates/server/src/runtime/session_actor/turn_inline.rs b/crates/server/src/runtime/session_actor/turn_inline.rs index b07ae318..872ca111 100644 --- a/crates/server/src/runtime/session_actor/turn_inline.rs +++ b/crates/server/src/runtime/session_actor/turn_inline.rs @@ -16,6 +16,10 @@ use super::SessionActorState; use super::snapshots::HookContextSnapshot; /// Mutable session fields updated during an in-actor turn without mailbox round-trips. +/// +/// Transient scratch state registered in `ActiveTurnRegistry` while the actor +/// mailbox is blocked or an out-of-actor turn runs. Merges into durable actor +/// state when the turn completes. pub(crate) struct TurnInlineState { pub(crate) turn_id: TurnId, pub(crate) turn_kind: TurnKind, diff --git a/crates/server/src/runtime/subagent_usage.rs b/crates/server/src/runtime/subagent_usage.rs index 39debba5..c6f29d1b 100644 --- a/crates/server/src/runtime/subagent_usage.rs +++ b/crates/server/src/runtime/subagent_usage.rs @@ -432,6 +432,9 @@ impl ServerRuntime { } pub(super) async fn active_turn_id_for_session(&self, session_id: SessionId) -> Option { + if let Some(turn_id) = self.runtime_active_turn_id(session_id).await { + return Some(turn_id); + } let session_handle = self.session(session_id).await?; session_handle.active_turn_id().await.flatten() } diff --git a/crates/server/src/runtime/turn_exec/finalize.rs b/crates/server/src/runtime/turn_exec/finalize.rs index ac2b2d7d..e8d7a56d 100644 --- a/crates/server/src/runtime/turn_exec/finalize.rs +++ b/crates/server/src/runtime/turn_exec/finalize.rs @@ -61,17 +61,7 @@ impl ServerRuntime { snapshot.session_totals.cache_creation_input_tokens; session_total_cache_read_tokens = snapshot.session_totals.cache_read_input_tokens; } - self.active_tasks.lock().await.remove(&session_id); - self.active_turn_cancellations - .lock() - .await - .remove(&session_id); - self.active_turn_ids.lock().await.remove(&session_id); - self.active_turn_metadata.lock().await.remove(&session_id); - self.active_turn_connections - .lock() - .await - .remove(&session_id); + self.clear_active_turn_interrupt_handles(session_id).await; match &result { Ok(()) => { self.run_session_hook_for_actor_state( diff --git a/crates/server/src/runtime/turn_exec/query.rs b/crates/server/src/runtime/turn_exec/query.rs index e6817ad6..b2a0bc58 100644 --- a/crates/server/src/runtime/turn_exec/query.rs +++ b/crates/server/src/runtime/turn_exec/query.rs @@ -115,11 +115,9 @@ impl ServerRuntime { let permission_profile = state.core.config.permission_profile.clone(); let hook_context = Self::hook_context_from_actor_state(state, session_id); let turn_cancel_token = self - .active_turn_cancellations - .lock() + .active_turns + .cancel_token(session_id) .await - .get(&session_id) - .cloned() .unwrap_or_else(CancellationToken::new); let query_cancel_token = turn_cancel_token.clone(); let provider_http = runtime_context diff --git a/crates/server/src/runtime/turn_exec/shell.rs b/crates/server/src/runtime/turn_exec/shell.rs index 2f54ddd5..6d2b80ac 100644 --- a/crates/server/src/runtime/turn_exec/shell.rs +++ b/crates/server/src/runtime/turn_exec/shell.rs @@ -68,11 +68,9 @@ impl ServerRuntime { .provider_http .clone(); let turn_cancel_token = self - .active_turn_cancellations - .lock() + .active_turns + .cancel_token(session_id) .await - .get(&session_id) - .cloned() .unwrap_or_else(CancellationToken::new); let tool_execution_start_runtime = Arc::clone(&self); let tool_execution_start_session_id = session_id; diff --git a/crates/server/src/runtime/turn_lifecycle.rs b/crates/server/src/runtime/turn_lifecycle.rs new file mode 100644 index 00000000..743aec2a --- /dev/null +++ b/crates/server/src/runtime/turn_lifecycle.rs @@ -0,0 +1,68 @@ +use std::sync::Arc; + +use devo_core::SessionId; +use tokio_util::sync::CancellationToken; + +use crate::turn::TurnMetadata; + +use super::ServerRuntime; + +impl ServerRuntime { + /// Registers cancellation, metadata, and optional connection ownership for a + /// turn that is about to run on a background task. + pub(crate) async fn register_active_turn_execution( + &self, + session_id: SessionId, + turn: TurnMetadata, + connection_id: Option, + ) -> CancellationToken { + let cancel_token = CancellationToken::new(); + self.active_turns + .insert_cancel_token(session_id, cancel_token.clone()) + .await; + self.register_runtime_active_turn(session_id, turn).await; + if let Some(connection_id) = connection_id { + self.active_turns + .set_connection_id(session_id, connection_id) + .await; + } + cancel_token + } + + pub(crate) async fn attach_active_turn_abort_handle( + &self, + session_id: SessionId, + abort_handle: tokio::task::AbortHandle, + ) { + self.active_turns + .set_abort_handle(session_id, abort_handle) + .await; + } + + /// Cancels and aborts the active turn for `session_id` without clearing the + /// full runtime handle (used while waiting for terminal status). + pub(crate) async fn signal_active_turn_interrupt(&self, session_id: SessionId) { + if let Some(cancel_token) = self.active_turns.cancel_token(session_id).await { + cancel_token.cancel(); + } + self.active_turns.abort_task(session_id).await; + } + + pub(crate) async fn spawn_active_turn_task( + self: &Arc, + session_id: SessionId, + turn: TurnMetadata, + connection_id: Option, + task: F, + ) where + F: std::future::Future + Send + 'static, + { + self.register_active_turn_execution(session_id, turn, connection_id) + .await; + let runtime = Arc::clone(self); + let join_handle = tokio::spawn(task); + runtime + .attach_active_turn_abort_handle(session_id, join_handle.abort_handle()) + .await; + } +} diff --git a/crates/server/src/runtime/turn_reservation.rs b/crates/server/src/runtime/turn_reservation.rs index 586a5b17..920e246f 100644 --- a/crates/server/src/runtime/turn_reservation.rs +++ b/crates/server/src/runtime/turn_reservation.rs @@ -52,7 +52,7 @@ impl ServerRuntime { } pub(super) async fn runtime_active_turn_id(&self, session_id: SessionId) -> Option { - self.active_turn_ids.lock().await.get(&session_id).copied() + self.active_turns.active_turn_id(session_id).await } pub(super) async fn register_runtime_active_turn( @@ -60,28 +60,17 @@ impl ServerRuntime { session_id: SessionId, turn: TurnMetadata, ) { - self.active_turn_ids - .lock() - .await - .insert(session_id, turn.turn_id); - self.active_turn_metadata - .lock() - .await - .insert(session_id, turn); + self.active_turns + .register_turn_metadata(session_id, turn) + .await; + } + + pub(super) async fn clear_active_turn_interrupt_handles(&self, session_id: SessionId) { + self.active_turns.clear_interrupt_handles(session_id).await; } pub(super) async fn clear_active_turn_runtime_handles(&self, session_id: SessionId) { - self.active_tasks.lock().await.remove(&session_id); - self.active_turn_cancellations - .lock() - .await - .remove(&session_id); - self.active_turn_ids.lock().await.remove(&session_id); - self.active_turn_metadata.lock().await.remove(&session_id); - self.active_turn_connections - .lock() - .await - .remove(&session_id); + self.active_turns.clear_runtime_handles(session_id).await; } pub(super) async fn clear_active_turn_reservation( @@ -281,11 +270,11 @@ mod tests { ) .await; let response: ErrorResponse = serde_json::from_value(value).expect("error response"); - let queued_len = reservation - .pending_turn_queue - .lock() - .expect("pending turn queue mutex should not be poisoned") - .len(); + let queued_len = session_handle + .pending_queue_snapshot() + .await + .map(|snapshot| snapshot.pending_count) + .unwrap_or(0); assert_eq!(response.error.code, ProtocolErrorCode::TurnAlreadyRunning); assert_eq!(queued_len, 0);