a3s_code_core/orchestrator/
handle.rs

1//! SubAgent 句柄
2
3use crate::error::Result;
4use crate::orchestrator::{ControlSignal, SubAgentActivity, SubAgentConfig, SubAgentState};
5use std::sync::Arc;
6use tokio::sync::RwLock;
7
8/// SubAgent 句柄
9///
10/// 用于控制和查询 SubAgent 的状态。
11#[derive(Clone)]
12pub struct SubAgentHandle {
13    /// SubAgent ID
14    pub id: String,
15
16    /// SubAgent 配置
17    pub(crate) config: SubAgentConfig,
18
19    /// 创建时间（Unix 时间戳，毫秒）
20    pub(crate) created_at: u64,
21
22    /// 控制信号发送器
23    control_tx: tokio::sync::mpsc::Sender<ControlSignal>,
24
25    /// 状态
26    state: Arc<RwLock<SubAgentState>>,
27
28    /// 当前活动
29    pub(crate) activity: Arc<RwLock<SubAgentActivity>>,
30
31    /// 任务句柄（保持任务存活，不直接访问）
32    #[allow(dead_code)]
33    task_handle: Arc<tokio::task::JoinHandle<Result<String>>>,
34}
35
36impl SubAgentHandle {
37    /// 创建新的句柄
38    pub(crate) fn new(
39        id: String,
40        config: SubAgentConfig,
41        control_tx: tokio::sync::mpsc::Sender<ControlSignal>,
42        state: Arc<RwLock<SubAgentState>>,
43        activity: Arc<RwLock<SubAgentActivity>>,
44        task_handle: tokio::task::JoinHandle<Result<String>>,
45    ) -> Self {
46        Self {
47            id,
48            config,
49            created_at: std::time::SystemTime::now()
50                .duration_since(std::time::UNIX_EPOCH)
51                .unwrap()
52                .as_millis() as u64,
53            control_tx,
54            state,
55            activity,
56            task_handle: Arc::new(task_handle),
57        }
58    }
59
60    /// 获取当前状态
61    ///
62    /// 注意：此方法使用非阻塞读取。如果状态锁当前被持有，返回 Initializing 状态。
63    /// 在异步上下文中，建议使用 `state_async()` 以获得准确的状态。
64    pub fn state(&self) -> SubAgentState {
65        self.state
66            .try_read()
67            .map(|guard| guard.clone())
68            .unwrap_or(SubAgentState::Initializing)
69    }
70
71    /// 异步获取当前状态
72    pub async fn state_async(&self) -> SubAgentState {
73        self.state.read().await.clone()
74    }
75
76    /// 获取当前活动
77    pub async fn activity(&self) -> SubAgentActivity {
78        self.activity.read().await.clone()
79    }
80
81    /// 获取配置
82    pub fn config(&self) -> &SubAgentConfig {
83        &self.config
84    }
85
86    /// 获取创建时间
87    pub fn created_at(&self) -> u64 {
88        self.created_at
89    }
90
91    /// 发送控制信号
92    pub async fn send_control(&self, signal: ControlSignal) -> Result<()> {
93        self.control_tx
94            .send(signal)
95            .await
96            .map_err(|_| anyhow::anyhow!("Failed to send control signal: channel closed"))?;
97        Ok(())
98    }
99
100    /// 暂停执行
101    pub async fn pause(&self) -> Result<()> {
102        self.send_control(ControlSignal::Pause).await
103    }
104
105    /// 恢复执行
106    pub async fn resume(&self) -> Result<()> {
107        self.send_control(ControlSignal::Resume).await
108    }
109
110    /// 取消执行
111    pub async fn cancel(&self) -> Result<()> {
112        self.send_control(ControlSignal::Cancel).await
113    }
114
115    /// 调整参数
116    pub async fn adjust_params(
117        &self,
118        max_steps: Option<usize>,
119        timeout_ms: Option<u64>,
120    ) -> Result<()> {
121        self.send_control(ControlSignal::AdjustParams {
122            max_steps,
123            timeout_ms,
124        })
125        .await
126    }
127
128    /// 注入新提示词
129    pub async fn inject_prompt(&self, prompt: impl Into<String>) -> Result<()> {
130        self.send_control(ControlSignal::InjectPrompt {
131            prompt: prompt.into(),
132        })
133        .await
134    }
135
136    /// 等待完成
137    pub async fn wait(&self) -> Result<String> {
138        // Note: 这里需要克隆 Arc 来避免移动所有权问题
139        // 实际使用中，应该只调用一次 wait()
140        loop {
141            let state = self.state_async().await;
142            if state.is_terminal() {
143                match state {
144                    SubAgentState::Completed { output, .. } => return Ok(output),
145                    SubAgentState::Cancelled => {
146                        return Err(anyhow::anyhow!("SubAgent was cancelled").into())
147                    }
148                    SubAgentState::Error { message } => {
149                        return Err(anyhow::anyhow!("SubAgent error: {}", message).into())
150                    }
151                    _ => unreachable!(),
152                }
153            }
154            tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
155        }
156    }
157
158    /// 是否已完成
159    pub fn is_done(&self) -> bool {
160        self.state().is_terminal()
161    }
162
163    /// 是否正在运行
164    pub fn is_running(&self) -> bool {
165        self.state().is_running()
166    }
167
168    /// 是否已暂停
169    pub fn is_paused(&self) -> bool {
170        self.state().is_paused()
171    }
172}
173
174impl std::fmt::Debug for SubAgentHandle {
175    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
176        f.debug_struct("SubAgentHandle")
177            .field("id", &self.id)
178            .field("state", &self.state())
179            .finish()
180    }
181}
a3s_code_core/orchestrator/handle.rs

a3s_code_core/orchestrator/
handle.rs
