a3s_code_core/
queue.rs

1//! Per-session command queue with lane-based priority scheduling
2//!
3//! Provides session-isolated command queues where each session has its own
4//! set of lanes with configurable concurrency limits and priorities.
5//!
6//! ## External Task Handling
7//!
8//! Supports pluggable task handlers allowing SDK users to implement custom
9//! processing logic for different lanes:
10//!
11//! - **Internal**: Default, tasks executed within the runtime
12//! - **External**: Tasks sent to SDK, wait for callback completion
13//! - **Hybrid**: Internal execution with external notification
14//!
15//! ## Implementation
16//!
17//! The actual queue implementation is in `SessionLaneQueue` which is backed
18//! by a3s-lane with features like DLQ, metrics, retry policies, and rate limiting.
19
20pub use a3s_lane::MetricsSnapshot;
21use anyhow::Result;
22use async_trait::async_trait;
23use serde::{Deserialize, Serialize};
24use std::collections::HashMap;
25use std::time::{Duration, Instant};
26
27// ============================================================================
28// Session Lane
29// ============================================================================
30
31/// Session lane for queue priority scheduling and HITL auto-approval
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
33pub enum SessionLane {
34    /// Control operations (P0) - pause, resume, cancel
35    Control,
36    /// Query operations (P1) - workspace reads and Code Intelligence queries
37    Query,
38    /// Execute operations (P2) - bash, write, edit
39    Execute,
40    /// Generate operations (P3) - LLM calls
41    Generate,
42}
43
44impl SessionLane {
45    /// Get the priority level (lower = higher priority)
46    pub fn priority(&self) -> u8 {
47        match self {
48            SessionLane::Control => 0,
49            SessionLane::Query => 1,
50            SessionLane::Execute => 2,
51            SessionLane::Generate => 3,
52        }
53    }
54
55    /// Map a tool name to its lane
56    pub fn from_tool_name(tool_name: &str) -> Self {
57        match tool_name {
58            "read" | "glob" | "ls" | "grep" | "list_files" | "search" | "web_fetch"
59            | "web_search" | "code_symbols" | "code_navigation" | "code_diagnostics" => {
60                SessionLane::Query
61            }
62            "bash" | "write" | "edit" | "delete" | "move" | "copy" | "execute" => {
63                SessionLane::Execute
64            }
65            _ => SessionLane::Execute,
66        }
67    }
68}
69
70// ============================================================================
71// Task Handler Configuration
72// ============================================================================
73
74/// Task handler mode determines how tasks in a lane are processed
75#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
76pub enum TaskHandlerMode {
77    /// Tasks are executed internally within the runtime (default)
78    #[default]
79    Internal,
80    /// Tasks are sent to external handler (SDK), wait for callback
81    External,
82    /// Tasks are executed internally but also notify external handler
83    Hybrid,
84}
85
86/// Configuration for a lane's task handler
87#[derive(Debug, Clone, Serialize, Deserialize)]
88pub struct LaneHandlerConfig {
89    /// Processing mode
90    pub mode: TaskHandlerMode,
91    /// Timeout for external processing (ms), default 60000 (60s)
92    pub timeout_ms: u64,
93}
94
95impl Default for LaneHandlerConfig {
96    fn default() -> Self {
97        Self {
98            mode: TaskHandlerMode::Internal,
99            timeout_ms: 60_000,
100        }
101    }
102}
103
104// ============================================================================
105// External Task Types
106// ============================================================================
107
108/// An external task that needs to be processed by SDK
109#[derive(Debug, Clone, Serialize, Deserialize)]
110pub struct ExternalTask {
111    /// Unique task identifier
112    pub task_id: String,
113    /// Session this task belongs to
114    pub session_id: String,
115    /// Lane the task is in
116    pub lane: SessionLane,
117    /// Type of command (e.g., "bash", "read", "write")
118    pub command_type: String,
119    /// Task payload as JSON
120    pub payload: serde_json::Value,
121    /// Timeout in milliseconds
122    pub timeout_ms: u64,
123    /// When the task was created
124    #[serde(skip)]
125    pub created_at: Option<Instant>,
126}
127
128impl ExternalTask {
129    /// Check if this task has timed out
130    pub fn is_timed_out(&self) -> bool {
131        self.created_at
132            .map(|t| t.elapsed() > Duration::from_millis(self.timeout_ms))
133            .unwrap_or(false)
134    }
135
136    /// Get remaining time until timeout in milliseconds
137    pub fn remaining_ms(&self) -> u64 {
138        self.created_at
139            .map(|t| {
140                let elapsed = t.elapsed().as_millis() as u64;
141                self.timeout_ms.saturating_sub(elapsed)
142            })
143            .unwrap_or(self.timeout_ms)
144    }
145}
146
147/// Result of external task processing
148#[derive(Debug, Clone, Serialize, Deserialize)]
149pub struct ExternalTaskResult {
150    /// Whether the task succeeded
151    pub success: bool,
152    /// Result data (JSON)
153    pub result: serde_json::Value,
154    /// Error message if failed
155    pub error: Option<String>,
156}
157
158// ============================================================================
159// Configuration
160// ============================================================================
161
162/// Configuration for a session command queue
163#[derive(Debug, Clone, Serialize, Deserialize)]
164#[serde(rename_all = "camelCase")]
165pub struct SessionQueueConfig {
166    /// Max concurrency for Control lane (P0)
167    #[serde(default = "default_control_concurrency")]
168    pub control_max_concurrency: usize,
169    /// Max concurrency for Query lane (P1)
170    #[serde(default = "default_query_concurrency")]
171    pub query_max_concurrency: usize,
172    /// Max concurrency for Execute lane (P2)
173    #[serde(default = "default_execute_concurrency")]
174    pub execute_max_concurrency: usize,
175    /// Max concurrency for Generate lane (P3)
176    #[serde(default = "default_generate_concurrency")]
177    pub generate_max_concurrency: usize,
178    /// Handler configurations per lane
179    #[serde(default)]
180    pub lane_handlers: HashMap<SessionLane, LaneHandlerConfig>,
181
182    // ========================================================================
183    // a3s-lane v0.4.0 integration features
184    // ========================================================================
185    /// Enable dead letter queue for failed commands
186    #[serde(default)]
187    pub enable_dlq: bool,
188    /// Max size of dead letter queue (None = use default 1000)
189    #[serde(default)]
190    pub dlq_max_size: Option<usize>,
191    /// Enable metrics collection
192    #[serde(default)]
193    pub enable_metrics: bool,
194    /// Enable queue alerts
195    #[serde(default)]
196    pub enable_alerts: bool,
197    /// Default timeout for commands in milliseconds
198    #[serde(default)]
199    pub default_timeout_ms: Option<u64>,
200    /// Persistent storage path (None = in-memory only)
201    #[serde(default)]
202    pub storage_path: Option<std::path::PathBuf>,
203
204    // ========================================================================
205    // a3s-lane v0.4.0 advanced features
206    // ========================================================================
207    /// Retry policy configuration
208    #[serde(default)]
209    pub retry_policy: Option<RetryPolicyConfig>,
210    /// Rate limit configuration
211    #[serde(default)]
212    pub rate_limit: Option<RateLimitConfig>,
213    /// Priority boost configuration
214    #[serde(default)]
215    pub priority_boost: Option<PriorityBoostConfig>,
216    /// Pressure threshold for emitting pressure/idle events
217    #[serde(default)]
218    pub pressure_threshold: Option<usize>,
219    /// Per-lane timeout overrides in milliseconds
220    #[serde(default)]
221    pub lane_timeouts: HashMap<SessionLane, u64>,
222}
223
224/// Retry policy configuration
225#[derive(Debug, Clone, Serialize, Deserialize)]
226#[serde(rename_all = "camelCase")]
227pub struct RetryPolicyConfig {
228    /// Retry strategy: "exponential", "fixed", or "none"
229    pub strategy: String,
230    /// Maximum number of retries
231    #[serde(default = "default_max_retries")]
232    pub max_retries: u32,
233    /// Initial delay in milliseconds (for exponential)
234    #[serde(default = "default_initial_delay_ms")]
235    pub initial_delay_ms: u64,
236    /// Fixed delay in milliseconds (for fixed strategy)
237    #[serde(default)]
238    pub fixed_delay_ms: Option<u64>,
239}
240
241fn default_max_retries() -> u32 {
242    3
243}
244
245fn default_initial_delay_ms() -> u64 {
246    100
247}
248
249/// Rate limit configuration
250#[derive(Debug, Clone, Serialize, Deserialize)]
251#[serde(rename_all = "camelCase")]
252pub struct RateLimitConfig {
253    /// Rate limit type: "per_second", "per_minute", "per_hour", or "unlimited"
254    pub limit_type: String,
255    /// Maximum number of operations per time period
256    #[serde(default)]
257    pub max_operations: Option<u64>,
258}
259
260/// Priority boost configuration
261#[derive(Debug, Clone, Serialize, Deserialize)]
262#[serde(rename_all = "camelCase")]
263pub struct PriorityBoostConfig {
264    /// Boost strategy: "standard", "aggressive", or "disabled"
265    pub strategy: String,
266    /// Deadline in milliseconds
267    #[serde(default)]
268    pub deadline_ms: Option<u64>,
269}
270
271fn default_control_concurrency() -> usize {
272    4
273}
274
275fn default_query_concurrency() -> usize {
276    12 // Balanced: better stability than 8, good performance (between 8 and 16)
277}
278
279fn default_execute_concurrency() -> usize {
280    4
281}
282
283fn default_generate_concurrency() -> usize {
284    2
285}
286
287impl Default for SessionQueueConfig {
288    fn default() -> Self {
289        Self {
290            control_max_concurrency: 2,
291            query_max_concurrency: 4,
292            execute_max_concurrency: 2,
293            generate_max_concurrency: 1,
294            lane_handlers: HashMap::new(),
295            enable_dlq: false,
296            dlq_max_size: None,
297            enable_metrics: false,
298            enable_alerts: false,
299            default_timeout_ms: None,
300            storage_path: None,
301            retry_policy: None,
302            rate_limit: None,
303            priority_boost: None,
304            pressure_threshold: None,
305            lane_timeouts: HashMap::new(),
306        }
307    }
308}
309
310impl SessionQueueConfig {
311    /// Get max concurrency for a lane
312    pub fn max_concurrency(&self, lane: SessionLane) -> usize {
313        match lane {
314            SessionLane::Control => self.control_max_concurrency,
315            SessionLane::Query => self.query_max_concurrency,
316            SessionLane::Execute => self.execute_max_concurrency,
317            SessionLane::Generate => self.generate_max_concurrency,
318        }
319    }
320
321    /// Get handler config for a lane (returns default if not configured)
322    pub fn handler_config(&self, lane: SessionLane) -> LaneHandlerConfig {
323        self.lane_handlers.get(&lane).cloned().unwrap_or_default()
324    }
325
326    /// Enable dead letter queue with optional max size
327    pub fn with_dlq(mut self, max_size: Option<usize>) -> Self {
328        self.enable_dlq = true;
329        self.dlq_max_size = max_size;
330        self
331    }
332
333    /// Enable metrics collection
334    pub fn with_metrics(mut self) -> Self {
335        self.enable_metrics = true;
336        self
337    }
338
339    /// Enable queue alerts
340    pub fn with_alerts(mut self) -> Self {
341        self.enable_alerts = true;
342        self
343    }
344
345    /// Set default timeout for commands
346    pub fn with_timeout(mut self, timeout_ms: u64) -> Self {
347        self.default_timeout_ms = Some(timeout_ms);
348        self
349    }
350
351    /// Set persistent storage path
352    pub fn with_storage(mut self, path: impl Into<std::path::PathBuf>) -> Self {
353        self.storage_path = Some(path.into());
354        self
355    }
356
357    /// Enable all a3s-lane features with sensible defaults
358    pub fn with_lane_features(mut self) -> Self {
359        self.enable_dlq = true;
360        self.dlq_max_size = Some(1000);
361        self.enable_metrics = true;
362        self.enable_alerts = true;
363        self.default_timeout_ms = Some(60_000);
364        self
365    }
366}
367
368// ============================================================================
369// Session Command Trait
370// ============================================================================
371
372/// Command to be executed in a session queue
373#[async_trait]
374pub trait SessionCommand: Send + Sync {
375    /// Execute the command
376    async fn execute(&self) -> Result<serde_json::Value>;
377
378    /// Get command type (for logging/debugging)
379    fn command_type(&self) -> &str;
380
381    /// Get command payload as JSON (for external handling)
382    fn payload(&self) -> serde_json::Value {
383        serde_json::json!({})
384    }
385}
386
387// ============================================================================
388// Queue Status Types
389// ============================================================================
390
391/// Status of a single lane
392#[derive(Debug, Clone, Serialize, Deserialize)]
393pub struct LaneStatus {
394    pub lane: SessionLane,
395    pub pending: usize,
396    pub active: usize,
397    pub max_concurrency: usize,
398    pub handler_mode: TaskHandlerMode,
399}
400
401/// Statistics for a session queue
402#[derive(Debug, Clone, Default, Serialize, Deserialize)]
403pub struct SessionQueueStats {
404    pub total_pending: usize,
405    pub total_active: usize,
406    pub external_pending: usize,
407    pub lanes: HashMap<String, LaneStatus>,
408}
409
410// ============================================================================
411// Tests
412// ============================================================================
413
414#[cfg(test)]
415mod tests {
416    use super::*;
417
418    #[test]
419    fn test_task_handler_mode_default() {
420        let mode = TaskHandlerMode::default();
421        assert_eq!(mode, TaskHandlerMode::Internal);
422    }
423
424    #[test]
425    fn test_lane_handler_config_default() {
426        let config = LaneHandlerConfig::default();
427        assert_eq!(config.mode, TaskHandlerMode::Internal);
428        assert_eq!(config.timeout_ms, 60_000);
429    }
430
431    #[test]
432    fn test_external_task_timeout() {
433        let task = ExternalTask {
434            task_id: "test".to_string(),
435            session_id: "session".to_string(),
436            lane: SessionLane::Query,
437            command_type: "read".to_string(),
438            payload: serde_json::json!({}),
439            timeout_ms: 100,
440            created_at: Some(Instant::now()),
441        };
442
443        assert!(!task.is_timed_out());
444        assert!(task.remaining_ms() <= 100);
445    }
446
447    #[test]
448    fn test_session_queue_config_default() {
449        let config = SessionQueueConfig::default();
450        assert_eq!(config.control_max_concurrency, 2);
451        assert_eq!(config.query_max_concurrency, 4);
452        assert_eq!(config.execute_max_concurrency, 2);
453        assert_eq!(config.generate_max_concurrency, 1);
454        assert!(!config.enable_dlq);
455        assert!(!config.enable_metrics);
456        assert!(!config.enable_alerts);
457    }
458
459    #[test]
460    fn test_session_queue_config_max_concurrency() {
461        let config = SessionQueueConfig::default();
462        assert_eq!(config.max_concurrency(SessionLane::Control), 2);
463        assert_eq!(config.max_concurrency(SessionLane::Query), 4);
464        assert_eq!(config.max_concurrency(SessionLane::Execute), 2);
465        assert_eq!(config.max_concurrency(SessionLane::Generate), 1);
466    }
467
468    #[test]
469    fn test_session_queue_config_handler_config() {
470        let config = SessionQueueConfig::default();
471        let handler = config.handler_config(SessionLane::Execute);
472        assert_eq!(handler.mode, TaskHandlerMode::Internal);
473        assert_eq!(handler.timeout_ms, 60_000);
474    }
475
476    #[test]
477    fn test_session_queue_config_builders() {
478        let config = SessionQueueConfig::default()
479            .with_dlq(Some(500))
480            .with_metrics()
481            .with_alerts()
482            .with_timeout(30_000);
483
484        assert!(config.enable_dlq);
485        assert_eq!(config.dlq_max_size, Some(500));
486        assert!(config.enable_metrics);
487        assert!(config.enable_alerts);
488        assert_eq!(config.default_timeout_ms, Some(30_000));
489    }
490
491    #[test]
492    fn test_session_queue_config_with_lane_features() {
493        let config = SessionQueueConfig::default().with_lane_features();
494
495        assert!(config.enable_dlq);
496        assert_eq!(config.dlq_max_size, Some(1000));
497        assert!(config.enable_metrics);
498        assert!(config.enable_alerts);
499        assert_eq!(config.default_timeout_ms, Some(60_000));
500    }
501
502    #[test]
503    fn test_external_task_result() {
504        let result = ExternalTaskResult {
505            success: true,
506            result: serde_json::json!({"output": "hello"}),
507            error: None,
508        };
509        assert!(result.success);
510        assert!(result.error.is_none());
511    }
512}
a3s_code_core/queue.rs

a3s_code_core/
queue.rs
