1use std::collections::BTreeMap;
8use std::fmt;
9use std::path::PathBuf;
10
11use serde::{Deserialize, Serialize};
12use serde_json::Value;
13
14use crate::runtime_event::RuntimeEventEnvelope;
15
16pub const WORKER_PROTOCOL_VERSION: u16 = 1;
17
18pub fn validate_worker_protocol_version(protocol_version: u16) -> Result<(), RuntimeError> {
19 if protocol_version == WORKER_PROTOCOL_VERSION {
20 Ok(())
21 } else {
22 Err(RuntimeError {
23 code: RuntimeErrorCode::InvalidRequest,
24 message: format!(
25 "unsupported Worker protocol version {protocol_version}; expected {WORKER_PROTOCOL_VERSION}"
26 ),
27 retryable: false,
28 })
29 }
30}
31
32macro_rules! string_id {
33 ($name:ident) => {
34 #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
35 #[serde(transparent)]
36 pub struct $name(String);
37
38 impl $name {
39 pub fn new(value: impl Into<String>) -> Self {
40 Self(value.into())
41 }
42
43 pub fn as_str(&self) -> &str {
44 &self.0
45 }
46
47 pub fn into_inner(self) -> String {
48 self.0
49 }
50 }
51
52 impl From<String> for $name {
53 fn from(value: String) -> Self {
54 Self(value)
55 }
56 }
57
58 impl From<&str> for $name {
59 fn from(value: &str) -> Self {
60 Self(value.to_owned())
61 }
62 }
63
64 impl AsRef<str> for $name {
65 fn as_ref(&self) -> &str {
66 self.as_str()
67 }
68 }
69
70 impl fmt::Display for $name {
71 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
72 f.write_str(self.as_str())
73 }
74 }
75 };
76}
77
78string_id!(RunId);
79string_id!(SessionId);
80string_id!(AgentId);
81string_id!(WorkerId);
82
83#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
84#[serde(rename_all = "snake_case")]
85pub enum RunStatus {
86 #[default]
87 Accepted,
88 Queued,
89 Running,
90 WaitingForApproval,
91 Completed,
92 Failed,
93 TimedOut,
94 Cancelled,
95}
96
97impl RunStatus {
98 pub fn is_terminal(self) -> bool {
99 matches!(
100 self,
101 Self::Completed | Self::Failed | Self::TimedOut | Self::Cancelled
102 )
103 }
104}
105
106#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
107#[serde(rename_all = "snake_case")]
108pub enum ToolPolicy {
109 ReadOnly,
110 WorkspaceWrite,
111 #[default]
112 RuntimeDefault,
113}
114
115#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
116pub struct AgentProfile {
117 pub id: Option<AgentId>,
118 pub name: String,
119 pub system_prompt: Option<String>,
120 #[serde(default)]
121 pub skills: Vec<String>,
122 #[serde(default)]
123 pub tool_policy: ToolPolicy,
124 #[serde(default)]
125 pub metadata: BTreeMap<String, Value>,
126}
127
128#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
129pub struct CallerContext {
130 pub channel: String,
131 pub principal: Option<String>,
132 pub chat_id: Option<i64>,
133}
134
135#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
136pub struct RunRequest {
137 pub run_id: Option<RunId>,
138 pub prompt: String,
139 pub session_id: Option<SessionId>,
140 pub agent_id: Option<AgentId>,
141 pub parent_run_id: Option<RunId>,
142 pub workspace: Option<PathBuf>,
143 #[serde(default)]
144 pub caller: CallerContext,
145 #[serde(default)]
146 pub metadata: BTreeMap<String, Value>,
147}
148
149impl RunRequest {
150 pub fn new(prompt: impl Into<String>) -> Self {
151 Self {
152 run_id: None,
153 prompt: prompt.into(),
154 session_id: None,
155 agent_id: None,
156 parent_run_id: None,
157 workspace: None,
158 caller: CallerContext::default(),
159 metadata: BTreeMap::new(),
160 }
161 }
162}
163
164#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
165pub struct RunResult {
166 pub run_id: RunId,
167 pub session_id: SessionId,
168 pub status: RunStatus,
169 pub final_text: String,
170 pub error: Option<RuntimeError>,
171 #[serde(default)]
172 pub metadata: BTreeMap<String, Value>,
173}
174
175#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
176#[serde(rename_all = "snake_case")]
177pub enum RuntimeErrorCode {
178 InvalidRequest,
179 Configuration,
180 Provider,
181 Tool,
182 Storage,
183 ApprovalDenied,
184 Cancelled,
185 TimedOut,
186 Unavailable,
187 Internal,
188}
189
190#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
191pub struct RuntimeError {
192 pub code: RuntimeErrorCode,
193 pub message: String,
194 pub retryable: bool,
195}
196
197impl fmt::Display for RuntimeError {
198 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
199 write!(f, "{}", self.message)
200 }
201}
202
203impl std::error::Error for RuntimeError {}
204
205#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
206pub struct RuntimeCapabilities {
207 pub streaming: bool,
208 pub cancellation: bool,
209 pub steering: bool,
210 pub approvals: bool,
211 pub skills: bool,
212 pub subagents: bool,
213 pub remote_workers: bool,
214}
215
216#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
217#[serde(tag = "type", rename_all = "snake_case")]
218pub enum RuntimeControl {
219 Cancel {
220 run_id: RunId,
221 },
222 Steer {
223 run_id: RunId,
224 message: String,
225 },
226 ResolveApproval {
227 run_id: RunId,
228 approval_id: String,
229 decision: String,
230 },
231}
232
233#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
234pub struct WorkerDescriptor {
235 pub id: WorkerId,
236 pub name: String,
237 pub capabilities: RuntimeCapabilities,
238 pub max_concurrent_runs: usize,
239 #[serde(default)]
240 pub labels: BTreeMap<String, String>,
241}
242
243#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
244#[serde(rename_all = "snake_case")]
245pub enum WorkerHealthStatus {
246 Ready,
247 Busy,
248 Draining,
249 Unavailable,
250}
251
252#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
253pub struct WorkerHealth {
254 pub worker_id: WorkerId,
255 pub status: WorkerHealthStatus,
256 pub active_runs: usize,
257 #[serde(default)]
258 pub queued_runs: usize,
259 pub observed_at: String,
260}
261
262#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
264#[serde(tag = "type", rename_all = "snake_case")]
265pub enum WorkerCommand {
266 Describe {
267 protocol_version: u16,
268 },
269 Health {
270 protocol_version: u16,
271 },
272 Submit {
273 protocol_version: u16,
274 profile: AgentProfile,
275 request: Box<RunRequest>,
276 },
277 Control {
278 protocol_version: u16,
279 control: RuntimeControl,
280 },
281 ResumeEvents {
282 protocol_version: u16,
283 run_id: RunId,
284 after_sequence: Option<u64>,
285 },
286}
287
288impl WorkerCommand {
289 pub fn protocol_version(&self) -> u16 {
290 match self {
291 Self::Describe { protocol_version }
292 | Self::Health { protocol_version }
293 | Self::Submit {
294 protocol_version, ..
295 }
296 | Self::Control {
297 protocol_version, ..
298 }
299 | Self::ResumeEvents {
300 protocol_version, ..
301 } => *protocol_version,
302 }
303 }
304
305 pub fn validate_protocol(&self) -> Result<(), RuntimeError> {
306 validate_worker_protocol_version(self.protocol_version())
307 }
308}
309
310#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
313#[serde(tag = "type", rename_all = "snake_case")]
314pub enum WorkerFrame {
315 Descriptor {
316 protocol_version: u16,
317 descriptor: WorkerDescriptor,
318 },
319 Health {
320 protocol_version: u16,
321 health: WorkerHealth,
322 },
323 Accepted {
324 protocol_version: u16,
325 run_id: RunId,
326 },
327 Event {
328 protocol_version: u16,
329 envelope: RuntimeEventEnvelope,
330 },
331 ControlAcknowledged {
332 protocol_version: u16,
333 run_id: RunId,
334 },
335 Result {
336 protocol_version: u16,
337 result: RunResult,
338 },
339 Error {
340 protocol_version: u16,
341 run_id: Option<RunId>,
342 error: RuntimeError,
343 },
344}
345
346impl WorkerFrame {
347 pub fn protocol_version(&self) -> u16 {
348 match self {
349 Self::Descriptor {
350 protocol_version, ..
351 }
352 | Self::Health {
353 protocol_version, ..
354 }
355 | Self::Accepted {
356 protocol_version, ..
357 }
358 | Self::Event {
359 protocol_version, ..
360 }
361 | Self::ControlAcknowledged {
362 protocol_version, ..
363 }
364 | Self::Result {
365 protocol_version, ..
366 }
367 | Self::Error {
368 protocol_version, ..
369 } => *protocol_version,
370 }
371 }
372
373 pub fn validate_protocol(&self) -> Result<(), RuntimeError> {
374 validate_worker_protocol_version(self.protocol_version())
375 }
376}
377
378#[cfg(test)]
379mod tests {
380 use super::*;
381
382 #[test]
383 fn identifiers_are_transport_stable() {
384 let id = RunId::new("run-1");
385 assert_eq!(serde_json::to_string(&id).unwrap(), "\"run-1\"");
386 assert_eq!(serde_json::from_str::<RunId>("\"run-1\"").unwrap(), id);
387 assert_eq!(id.to_string(), "run-1");
388 }
389
390 #[test]
391 fn run_request_round_trips_without_product_details() {
392 let mut request = RunRequest::new("inspect this project");
393 request.session_id = Some(SessionId::new("project-1"));
394 request.agent_id = Some(AgentId::new("coder"));
395 request.workspace = Some(PathBuf::from("/workspace"));
396 request.caller = CallerContext {
397 channel: "work".into(),
398 principal: Some("local-user".into()),
399 chat_id: Some(42),
400 };
401
402 let json = serde_json::to_string(&request).unwrap();
403 let decoded: RunRequest = serde_json::from_str(&json).unwrap();
404 assert_eq!(decoded, request);
405 assert!(!json.contains("provider"));
406 assert!(!json.contains("database"));
407 }
408
409 #[test]
410 fn worker_protocol_round_trips_submission_and_result_frames() {
411 let mut request = RunRequest::new("implement the parser");
412 request.run_id = Some(RunId::new("run-7"));
413 request.parent_run_id = Some(RunId::new("parent-2"));
414 let command = WorkerCommand::Submit {
415 protocol_version: WORKER_PROTOCOL_VERSION,
416 profile: AgentProfile {
417 id: Some(AgentId::new("coder")),
418 skills: vec!["rust-review".into()],
419 ..AgentProfile::default()
420 },
421 request: Box::new(request),
422 };
423 let encoded = serde_json::to_string(&command).unwrap();
424 let decoded: WorkerCommand = serde_json::from_str(&encoded).unwrap();
425 assert_eq!(decoded, command);
426 assert_eq!(decoded.protocol_version(), WORKER_PROTOCOL_VERSION);
427
428 let frame = WorkerFrame::Result {
429 protocol_version: WORKER_PROTOCOL_VERSION,
430 result: RunResult {
431 run_id: RunId::new("run-7"),
432 session_id: SessionId::new("session-3"),
433 status: RunStatus::Completed,
434 final_text: "done".into(),
435 error: None,
436 metadata: BTreeMap::new(),
437 },
438 };
439 let encoded = serde_json::to_string(&frame).unwrap();
440 assert_eq!(
441 serde_json::from_str::<WorkerFrame>(&encoded).unwrap(),
442 frame
443 );
444 assert_eq!(frame.protocol_version(), WORKER_PROTOCOL_VERSION);
445 }
446
447 #[test]
448 fn terminal_statuses_are_explicit() {
449 assert!(!RunStatus::Running.is_terminal());
450 assert!(!RunStatus::WaitingForApproval.is_terminal());
451 assert!(RunStatus::Completed.is_terminal());
452 assert!(RunStatus::Failed.is_terminal());
453 assert!(RunStatus::TimedOut.is_terminal());
454 assert!(RunStatus::Cancelled.is_terminal());
455 }
456
457 #[test]
458 fn worker_protocol_rejects_incompatible_versions() {
459 let command = WorkerCommand::Describe {
460 protocol_version: WORKER_PROTOCOL_VERSION + 1,
461 };
462 let error = command.validate_protocol().unwrap_err();
463 assert_eq!(error.code, RuntimeErrorCode::InvalidRequest);
464 assert!(error
465 .message
466 .contains("unsupported Worker protocol version"));
467
468 let frame = WorkerFrame::Accepted {
469 protocol_version: WORKER_PROTOCOL_VERSION,
470 run_id: RunId::new("run-1"),
471 };
472 frame.validate_protocol().unwrap();
473 }
474
475 #[test]
476 fn worker_contract_round_trips() {
477 let descriptor = WorkerDescriptor {
478 id: WorkerId::new("local-1"),
479 name: "Local worker".into(),
480 capabilities: RuntimeCapabilities {
481 streaming: true,
482 cancellation: true,
483 steering: true,
484 approvals: true,
485 skills: true,
486 subagents: true,
487 remote_workers: false,
488 },
489 max_concurrent_runs: 4,
490 labels: BTreeMap::from([("arch".into(), "arm64".into())]),
491 };
492
493 let json = serde_json::to_string(&descriptor).unwrap();
494 assert_eq!(
495 serde_json::from_str::<WorkerDescriptor>(&json).unwrap(),
496 descriptor
497 );
498 }
499}