Skip to main content

tea_protocol/
command.rs

1use std::fmt;
2use std::str::FromStr;
3
4use serde::{Deserialize, Deserializer, Serialize, Serializer};
5use serde_json::{Map, Value, json};
6use thiserror::Error;
7
8use crate::envelope::{deserialize_unique_value, validate_read_version};
9use crate::{
10    ApprovalId, BranchId, CURRENT_PROTOCOL_VERSION, CanonicalMessage, CommandId, CorrelationId,
11    MessageId, MessageRole, ProtocolError, ProtocolMetadata, ProtocolTimestamp, ProtocolVersion,
12    SessionId,
13};
14
15/// Maximum UTF-8 bytes in a command text fragment.
16pub const MAX_COMMAND_TEXT_BYTES: usize = 256 * 1024;
17/// Maximum UTF-8 bytes in a model or profile selector.
18pub const MAX_SELECTOR_BYTES: usize = 128;
19
20macro_rules! selector {
21    ($name:ident, $doc:literal, $validate:ident) => {
22        #[doc = $doc]
23        #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
24        pub struct $name(String);
25
26        impl $name {
27            /// Returns the canonical selector text.
28            #[must_use]
29            pub fn as_str(&self) -> &str {
30                &self.0
31            }
32        }
33
34        impl fmt::Display for $name {
35            fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
36                formatter.write_str(&self.0)
37            }
38        }
39
40        impl FromStr for $name {
41            type Err = SelectorParseError;
42
43            fn from_str(value: &str) -> Result<Self, Self::Err> {
44                $validate(value)?;
45                Ok(Self(value.to_owned()))
46            }
47        }
48
49        impl Serialize for $name {
50            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
51            where
52                S: Serializer,
53            {
54                $validate(&self.0).map_err(serde::ser::Error::custom)?;
55                serializer.serialize_str(&self.0)
56            }
57        }
58
59        impl<'de> Deserialize<'de> for $name {
60            fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
61            where
62                D: Deserializer<'de>,
63            {
64                let value = String::deserialize(deserializer)?;
65                value.parse().map_err(serde::de::Error::custom)
66            }
67        }
68    };
69}
70
71selector!(
72    ProfileId,
73    "A bounded product profile selector.",
74    validate_profile_selector
75);
76selector!(
77    ModelId,
78    "A bounded canonical model selector.",
79    validate_model_selector
80);
81selector!(
82    ProviderId,
83    "A bounded canonical model-provider selector.",
84    validate_profile_selector
85);
86
87/// Error returned when parsing a model or profile selector.
88#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
89pub enum SelectorParseError {
90    /// Selector text is empty, oversized, or contains non-canonical characters.
91    #[error(
92        "selector must start with lowercase ASCII and contain only supported canonical characters"
93    )]
94    Invalid,
95}
96
97/// A user decision for a pending approval.
98#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
99#[serde(tag = "type", rename_all = "snake_case")]
100pub enum ApprovalDecision {
101    /// Allow only this tool call.
102    AllowOnce,
103    /// Allow matching operations for the current session.
104    AllowSession,
105    /// Deny this tool call.
106    Deny,
107}
108
109/// Stable initial command discriminators.
110#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
111#[serde(rename_all = "snake_case")]
112pub enum AgentCommandType {
113    /// Create a new session.
114    CreateSession,
115    /// Submit a canonical user prompt.
116    Prompt,
117    /// Inject steering text into the active run.
118    Steer,
119    /// Queue a follow-up user message.
120    FollowUp,
121    /// Abort the active run.
122    Abort,
123    /// Resolve a pending approval.
124    ResolveApproval,
125    /// Select a model.
126    SetModel,
127    /// Select a reasoning effort for subsequent turns.
128    SetReasoningEffort,
129    /// Select a product profile.
130    SetProfile,
131    /// Request session compaction.
132    CompactSession,
133    /// Fork from a durable message or branch point.
134    ForkSession,
135}
136
137impl AgentCommandType {
138    /// All initial protocol 1.0 command types.
139    pub const ALL: [Self; 11] = [
140        Self::CreateSession,
141        Self::Prompt,
142        Self::Steer,
143        Self::FollowUp,
144        Self::Abort,
145        Self::ResolveApproval,
146        Self::SetModel,
147        Self::SetReasoningEffort,
148        Self::SetProfile,
149        Self::CompactSession,
150        Self::ForkSession,
151    ];
152}
153
154/// A provider- and transport-neutral agent command payload.
155#[derive(Debug, Clone, PartialEq)]
156pub enum AgentCommand {
157    /// Create a session using a profile and optional extension metadata.
158    CreateSession {
159        /// Product profile selector.
160        profile_id: ProfileId,
161        /// Bounded session extension metadata.
162        metadata: ProtocolMetadata,
163    },
164    /// Submit a canonical user message.
165    Prompt {
166        /// Canonical user message.
167        message: CanonicalMessage,
168    },
169    /// Steer the currently running turn.
170    Steer {
171        /// Bounded steering text.
172        text: CommandText,
173    },
174    /// Queue a follow-up canonical user message.
175    FollowUp {
176        /// Canonical user message.
177        message: CanonicalMessage,
178    },
179    /// Abort the active run.
180    Abort {},
181    /// Resolve a pending approval request.
182    ResolveApproval {
183        /// Approval request identifier.
184        approval_id: ApprovalId,
185        /// User or policy decision.
186        decision: ApprovalDecision,
187    },
188    /// Select a model for subsequent runs.
189    SetModel {
190        /// Complete provider-qualified model selector.
191        model: crate::ModelRef,
192    },
193    /// Select reasoning effort for subsequent turns.
194    SetReasoningEffort {
195        /// Provider-neutral reasoning effort.
196        reasoning_effort: crate::ReasoningEffort,
197    },
198    /// Select a product profile for subsequent runs.
199    SetProfile {
200        /// Product profile selector.
201        profile_id: ProfileId,
202    },
203    /// Request compaction with an optional bounded instruction.
204    CompactSession {
205        /// Optional compaction instruction.
206        instruction: Option<CommandText>,
207    },
208    /// Fork the session from a message into a new branch.
209    ForkSession {
210        /// Source message for the fork.
211        from_message_id: MessageId,
212        /// Identifier assigned to the new branch.
213        branch_id: BranchId,
214    },
215}
216
217impl AgentCommand {
218    /// Returns the stable command discriminator.
219    #[must_use]
220    pub const fn command_type(&self) -> AgentCommandType {
221        match self {
222            Self::CreateSession { .. } => AgentCommandType::CreateSession,
223            Self::Prompt { .. } => AgentCommandType::Prompt,
224            Self::Steer { .. } => AgentCommandType::Steer,
225            Self::FollowUp { .. } => AgentCommandType::FollowUp,
226            Self::Abort {} => AgentCommandType::Abort,
227            Self::ResolveApproval { .. } => AgentCommandType::ResolveApproval,
228            Self::SetModel { .. } => AgentCommandType::SetModel,
229            Self::SetReasoningEffort { .. } => AgentCommandType::SetReasoningEffort,
230            Self::SetProfile { .. } => AgentCommandType::SetProfile,
231            Self::CompactSession { .. } => AgentCommandType::CompactSession,
232            Self::ForkSession { .. } => AgentCommandType::ForkSession,
233        }
234    }
235
236    fn validate(&self) -> Result<(), CommandValidationError> {
237        match self {
238            Self::Prompt { message } | Self::FollowUp { message }
239                if message.role() != MessageRole::User =>
240            {
241                Err(CommandValidationError::MessageMustBeUser)
242            }
243            _ => Ok(()),
244        }
245    }
246}
247
248#[derive(Serialize, Deserialize)]
249#[serde(
250    remote = "AgentCommand",
251    tag = "type",
252    content = "payload",
253    rename_all = "snake_case"
254)]
255enum AgentCommandDef {
256    CreateSession {
257        #[serde(rename = "profileId")]
258        profile_id: ProfileId,
259        #[serde(default, skip_serializing_if = "ProtocolMetadata::is_empty")]
260        metadata: ProtocolMetadata,
261    },
262    Prompt {
263        message: CanonicalMessage,
264    },
265    Steer {
266        text: CommandText,
267    },
268    FollowUp {
269        message: CanonicalMessage,
270    },
271    Abort {},
272    ResolveApproval {
273        #[serde(rename = "approvalId")]
274        approval_id: ApprovalId,
275        decision: ApprovalDecision,
276    },
277    SetModel {
278        model: crate::ModelRef,
279    },
280    SetReasoningEffort {
281        #[serde(rename = "reasoningEffort")]
282        reasoning_effort: crate::ReasoningEffort,
283    },
284    SetProfile {
285        #[serde(rename = "profileId")]
286        profile_id: ProfileId,
287    },
288    CompactSession {
289        #[serde(skip_serializing_if = "Option::is_none")]
290        instruction: Option<CommandText>,
291    },
292    ForkSession {
293        #[serde(rename = "fromMessageId")]
294        from_message_id: MessageId,
295        #[serde(rename = "branchId")]
296        branch_id: BranchId,
297    },
298}
299
300impl Serialize for AgentCommand {
301    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
302    where
303        S: Serializer,
304    {
305        self.validate().map_err(serde::ser::Error::custom)?;
306        AgentCommandDef::serialize(self, serializer)
307    }
308}
309
310impl<'de> Deserialize<'de> for AgentCommand {
311    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
312    where
313        D: Deserializer<'de>,
314    {
315        let command = AgentCommandDef::deserialize(deserializer)?;
316        command.validate().map_err(serde::de::Error::custom)?;
317        Ok(command)
318    }
319}
320
321/// Bounded command text that rejects controls unsafe for logs and transports.
322#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
323#[serde(try_from = "String", into = "String")]
324pub struct CommandText(String);
325
326impl CommandText {
327    /// Creates validated command text.
328    ///
329    /// # Errors
330    ///
331    /// Returns an error when text is empty, oversized, or contains a null character.
332    pub fn new(value: impl Into<String>) -> Result<Self, CommandValidationError> {
333        let value = value.into();
334        if value.is_empty() || value.len() > MAX_COMMAND_TEXT_BYTES || value.contains('\0') {
335            return Err(CommandValidationError::InvalidText);
336        }
337        Ok(Self(value))
338    }
339
340    /// Returns the bounded command text.
341    #[must_use]
342    pub fn as_str(&self) -> &str {
343        &self.0
344    }
345}
346
347impl TryFrom<String> for CommandText {
348    type Error = CommandValidationError;
349
350    fn try_from(value: String) -> Result<Self, Self::Error> {
351        Self::new(value)
352    }
353}
354
355impl From<CommandText> for String {
356    fn from(value: CommandText) -> Self {
357        value.0
358    }
359}
360
361/// A versioned command transport envelope.
362#[derive(Debug, Clone, PartialEq)]
363pub struct CommandEnvelope {
364    protocol_version: ProtocolVersion,
365    command_id: CommandId,
366    session_id: Option<SessionId>,
367    timestamp: ProtocolTimestamp,
368    command: AgentCommand,
369}
370
371impl CommandEnvelope {
372    /// Creates a validated current-version command envelope.
373    ///
374    /// # Errors
375    ///
376    /// Returns an error when session presence or command payload invariants fail.
377    pub fn new(
378        command_id: CommandId,
379        session_id: Option<SessionId>,
380        timestamp: ProtocolTimestamp,
381        command: AgentCommand,
382    ) -> Result<Self, CommandValidationError> {
383        let envelope = Self {
384            protocol_version: CURRENT_PROTOCOL_VERSION,
385            command_id,
386            session_id,
387            timestamp,
388            command,
389        };
390        envelope.validate()?;
391        Ok(envelope)
392    }
393
394    /// Decodes a JSON value while preserving unsupported-command classification.
395    ///
396    /// # Errors
397    ///
398    /// Returns [`CommandDecodeError::UnsupportedType`] for an unknown safe
399    /// discriminator and [`CommandDecodeError::Invalid`] for malformed input.
400    pub fn decode_value(value: Value) -> Result<Self, CommandDecodeError> {
401        let version = decode_version(&value).map_err(CommandDecodeError::Invalid)?;
402        if validate_read_version(version).is_err() {
403            return Err(CommandDecodeError::UnsupportedVersion { version });
404        }
405        let discriminator = value
406            .as_object()
407            .and_then(|object| object.get("type"))
408            .and_then(Value::as_str)
409            .ok_or_else(|| CommandDecodeError::Invalid("missing command type".to_owned()))?;
410        if discriminator.parse::<AgentCommandTypeText>().is_err() {
411            if valid_discriminator(discriminator) {
412                return Err(CommandDecodeError::UnsupportedType {
413                    command_type: discriminator.to_owned(),
414                });
415            }
416            return Err(CommandDecodeError::Invalid(
417                "invalid command type".to_owned(),
418            ));
419        }
420        serde_json::from_value(value)
421            .map_err(|error| CommandDecodeError::Invalid(error.to_string()))
422    }
423
424    /// Returns the protocol version read from or written to the envelope.
425    #[must_use]
426    pub const fn protocol_version(&self) -> ProtocolVersion {
427        self.protocol_version
428    }
429
430    /// Returns the command identifier.
431    #[must_use]
432    pub const fn command_id(&self) -> CommandId {
433        self.command_id
434    }
435
436    /// Returns the target session, absent only for session creation.
437    #[must_use]
438    pub const fn session_id(&self) -> Option<SessionId> {
439        self.session_id
440    }
441
442    /// Returns the command timestamp.
443    #[must_use]
444    pub const fn timestamp(&self) -> ProtocolTimestamp {
445        self.timestamp
446    }
447
448    /// Returns the command payload.
449    #[must_use]
450    pub const fn command(&self) -> &AgentCommand {
451        &self.command
452    }
453
454    /// Returns the stable command discriminator.
455    #[must_use]
456    pub const fn command_type(&self) -> AgentCommandType {
457        self.command.command_type()
458    }
459
460    fn validate(&self) -> Result<(), CommandValidationError> {
461        let is_create = matches!(self.command, AgentCommand::CreateSession { .. });
462        if is_create == self.session_id.is_some() {
463            return Err(CommandValidationError::InvalidSessionPresence);
464        }
465        self.command.validate()
466    }
467}
468
469impl Serialize for CommandEnvelope {
470    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
471    where
472        S: Serializer,
473    {
474        self.validate().map_err(serde::ser::Error::custom)?;
475        let mut value = serde_json::to_value(&self.command).map_err(serde::ser::Error::custom)?;
476        let object = value
477            .as_object_mut()
478            .ok_or_else(|| serde::ser::Error::custom("command must encode as object"))?;
479        insert_envelope_fields(
480            object,
481            self.protocol_version,
482            self.command_id,
483            self.session_id,
484            self.timestamp,
485        );
486        value.serialize(serializer)
487    }
488}
489
490impl<'de> Deserialize<'de> for CommandEnvelope {
491    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
492    where
493        D: Deserializer<'de>,
494    {
495        let mut value = deserialize_unique_value(deserializer)?;
496        let object = value
497            .as_object_mut()
498            .ok_or_else(|| serde::de::Error::custom("command envelope must be an object"))?;
499        let protocol_version = take(object, "protocolVersion").map_err(serde::de::Error::custom)?;
500        validate_read_version(protocol_version).map_err(serde::de::Error::custom)?;
501        let command_id = take(object, "commandId").map_err(serde::de::Error::custom)?;
502        let session_id = take_optional(object, "sessionId").map_err(serde::de::Error::custom)?;
503        let timestamp = take(object, "timestamp").map_err(serde::de::Error::custom)?;
504        let command = AgentCommand::deserialize(Value::Object(std::mem::take(object)))
505            .map_err(serde::de::Error::custom)?;
506        let envelope = Self {
507            protocol_version,
508            command_id,
509            session_id,
510            timestamp,
511            command,
512        };
513        envelope.validate().map_err(serde::de::Error::custom)?;
514        Ok(envelope)
515    }
516}
517
518/// Failure while decoding an untrusted command envelope.
519#[derive(Debug, Error)]
520pub enum CommandDecodeError {
521    /// The protocol major is unsupported and takes precedence over command type.
522    #[error("unsupported protocol version: {version}")]
523    UnsupportedVersion {
524        /// Received canonical protocol version.
525        version: ProtocolVersion,
526    },
527    /// The discriminator is canonical but unsupported by this host.
528    #[error("unsupported command type: {command_type}")]
529    UnsupportedType {
530        /// Bounded canonical unsupported discriminator.
531        command_type: String,
532    },
533    /// The command envelope or known payload is malformed.
534    #[error("invalid command: {0}")]
535    Invalid(String),
536}
537
538impl CommandDecodeError {
539    /// Converts a decode failure to a safe protocol error.
540    #[must_use]
541    pub fn into_protocol_error(self, correlation_id: CorrelationId) -> ProtocolError {
542        match self {
543            Self::UnsupportedVersion { version } => {
544                ProtocolError::unsupported_protocol_version(correlation_id, version)
545            }
546            Self::UnsupportedType { command_type } => {
547                let details = ProtocolMetadata::protocol_compatibility_details(Some(&command_type));
548                ProtocolError::unsupported_command(correlation_id).with_details(details)
549            }
550            Self::Invalid(_) => ProtocolError::invalid_command(correlation_id),
551        }
552    }
553}
554
555/// Error returned when validating command data.
556#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
557pub enum CommandValidationError {
558    /// Only create-session omits a session identifier.
559    #[error("sessionId must be absent for create_session and present for every other command")]
560    InvalidSessionPresence,
561    /// Prompt and follow-up require canonical user messages.
562    #[error("prompt and follow_up messages must have user role")]
563    MessageMustBeUser,
564    /// Command text is empty, oversized, or contains a null character.
565    #[error("command text is invalid")]
566    InvalidText,
567}
568
569#[derive(Debug, Clone, Copy, PartialEq, Eq)]
570struct AgentCommandTypeText;
571
572impl FromStr for AgentCommandTypeText {
573    type Err = ();
574
575    fn from_str(value: &str) -> Result<Self, Self::Err> {
576        if [
577            "create_session",
578            "prompt",
579            "steer",
580            "follow_up",
581            "abort",
582            "resolve_approval",
583            "set_model",
584            "set_reasoning_effort",
585            "set_profile",
586            "compact_session",
587            "fork_session",
588        ]
589        .contains(&value)
590        {
591            Ok(Self)
592        } else {
593            Err(())
594        }
595    }
596}
597
598fn decode_version(value: &Value) -> Result<ProtocolVersion, String> {
599    let version = value
600        .as_object()
601        .and_then(|object| object.get("protocolVersion"))
602        .cloned()
603        .ok_or_else(|| "missing protocolVersion".to_owned())?;
604    serde_json::from_value(version).map_err(|error| error.to_string())
605}
606
607fn validate_profile_selector(value: &str) -> Result<(), SelectorParseError> {
608    validate_selector(value, false)
609}
610
611fn validate_model_selector(value: &str) -> Result<(), SelectorParseError> {
612    validate_selector(value, true)
613}
614
615fn validate_selector(value: &str, allow_colon: bool) -> Result<(), SelectorParseError> {
616    let mut bytes = value.bytes();
617    if value.len() > MAX_SELECTOR_BYTES
618        || !bytes.next().is_some_and(|byte| byte.is_ascii_lowercase())
619        || !bytes.all(|byte| {
620            byte.is_ascii_lowercase()
621                || byte.is_ascii_digit()
622                || matches!(byte, b'_' | b'-' | b'.' | b'/')
623                || (allow_colon && byte == b':')
624        })
625    {
626        Err(SelectorParseError::Invalid)
627    } else {
628        Ok(())
629    }
630}
631
632fn valid_discriminator(value: &str) -> bool {
633    !value.is_empty()
634        && value.len() <= 128
635        && value
636            .bytes()
637            .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || matches!(byte, b'_'))
638}
639
640fn insert_envelope_fields(
641    object: &mut Map<String, Value>,
642    version: ProtocolVersion,
643    command_id: CommandId,
644    session_id: Option<SessionId>,
645    timestamp: ProtocolTimestamp,
646) {
647    object.insert("protocolVersion".to_owned(), json!(version));
648    object.insert("commandId".to_owned(), json!(command_id));
649    if let Some(session_id) = session_id {
650        object.insert("sessionId".to_owned(), json!(session_id));
651    }
652    object.insert("timestamp".to_owned(), json!(timestamp));
653}
654
655fn take<T>(object: &mut Map<String, Value>, key: &str) -> Result<T, serde_json::Error>
656where
657    T: for<'de> Deserialize<'de>,
658{
659    serde_json::from_value(object.remove(key).unwrap_or(Value::Null))
660}
661
662fn take_optional<T>(
663    object: &mut Map<String, Value>,
664    key: &str,
665) -> Result<Option<T>, serde_json::Error>
666where
667    T: for<'de> Deserialize<'de>,
668{
669    object.remove(key).map_or(Ok(None), serde_json::from_value)
670}