aion-proto 0.31.0

//! Worker protocol serde/prost wire types.

use crate::{ProtoActivityId, ProtoPayload, ProtoRunId, ProtoWorkflowId, WireError};

/// Proto representation of `ActivityErrorKind`. Zero is invalid on decode.
#[derive(
    Clone,
    Copy,
    Debug,
    PartialEq,
    Eq,
    Hash,
    serde::Serialize,
    serde::Deserialize,
    prost::Enumeration,
)]
#[repr(i32)]
pub enum ProtoActivityErrorKind {
    /// Missing/invalid kind.
    Unspecified = 0,
    /// Activity failure may be retried by the engine.
    Retryable = 1,
    /// Provider policy refusal to be routed before ordinary retry handling.
    PolicyRefused = 3,
    /// Activity failure is terminal.
    Terminal = 2,
}

/// Proto representation of `ActivityError`.
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, prost::Message)]
pub struct ProtoActivityError {
    /// Explicit retryability classification.
    #[prost(enumeration = "ProtoActivityErrorKind", tag = "1")]
    pub kind: i32,
    /// Human-readable error message.
    #[prost(string, tag = "2")]
    pub message: String,
    /// Optional structured failure details.
    #[prost(message, optional, tag = "3")]
    pub details: Option<ProtoPayload>,
}

/// Typed activity surface advertised during worker registration.
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, prost::Message)]
pub struct ProtoActivityDescriptor {
    /// Activity type.
    #[prost(string, tag = "1")]
    pub name: String,
    /// JSON-encoded input schema.
    #[prost(string, tag = "2")]
    pub input_schema_json: String,
    /// JSON-encoded output schema.
    #[prost(string, tag = "3")]
    pub output_schema_json: String,
}

/// Stable identity of one worker instance within an operator deployment.
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, prost::Message)]
pub struct ProtoWorkerInstanceIdentity {
    /// Deployment record name.
    #[prost(string, tag = "1")]
    pub deployment: String,
    /// Operator/launcher-assigned instance identifier.
    #[prost(string, tag = "2")]
    pub instance_id: String,
}

/// Worker registration advertisement.
///
/// A registration scopes the worker to a SET of `namespaces` correctness
/// boundaries, one `task_queue` pool/flavour lane within each, and an optional
/// `node` locality. The dimensions are disjoint: `namespaces` is the set of
/// correctness/isolation boundaries the worker is authorized for, `task_queue`
/// is the pool/flavour selector within each, and `node` is an OPTIONAL locality
/// used to filter within a pool. `activity_type` (carried in `activity_types`)
/// is matched inside a pool, not used to select it.
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, prost::Message)]
pub struct ProtoRegisterWorker {
    /// Set of correctness/isolation boundaries this worker is authorized for.
    /// The server indexes the worker under each, preserving wire order.
    #[prost(string, repeated, tag = "1")]
    pub namespaces: Vec<String>,
    /// Activity types implemented by the worker, preserving wire order.
    #[prost(string, repeated, tag = "2")]
    pub activity_types: Vec<String>,
    /// Pool/flavour selector within each namespace. The worker-pool address is
    /// `(namespace, task_queue)`; the server normalizes an empty value to the
    /// literal `"default"` pool.
    #[prost(string, tag = "3")]
    pub task_queue: String,
    /// Optional locality affinity this worker advertises. A dispatch pinned to a
    /// node reaches only workers whose node equals it (round-robin among them);
    /// an empty value means the worker carries no locality. Default = the
    /// worker's machine hostname, resolved by the SDK at registration.
    #[prost(string, tag = "4")]
    pub node: String,
    /// Typed activity descriptors. Empty marks a pre-contract worker.
    #[prost(message, repeated, tag = "5")]
    pub activities: Vec<ProtoActivityDescriptor>,
    /// Stable worker-build identity.
    #[prost(string, tag = "6")]
    pub identity: String,
    /// Optional deployment/instance association, distinct from build identity.
    #[prost(message, optional, tag = "7")]
    pub instance: Option<ProtoWorkerInstanceIdentity>,
    /// How many activities this worker executes AT ONCE. The server counts a
    /// worker's in-flight dispatches against this number and never selects a
    /// worker that has reached it.
    ///
    /// `Option` because presence is load-bearing: a bare proto3 scalar cannot
    /// tell `0` from unset, and reading `0` as "unstated" would be an assumed
    /// default for the one number deciding how much work a process is handed.
    /// `None` and `Some(0)` are both REFUSED at registration — the same
    /// treatment `activities` gets for an empty set.
    #[prost(uint32, optional, tag = "8")]
    pub max_concurrency: Option<u32>,
}

/// A worker declining ONE pushed task because it has no free execution slot.
///
/// Not a failure and not an attempt: the server re-parks the dispatch with no
/// clock started, corrects its per-worker in-flight count, and re-selects.
/// Nothing is recorded against the activity and its attempt number is untouched.
///
/// With capacity on the wire the server should never push past a worker's
/// advertised concurrency, so this frame is the honest handling of a DRIFTED
/// count (a reconnect race, a redelivery joining an outstanding generation) —
/// never the primary backpressure mechanism.
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, prost::Message)]
pub struct ProtoTaskRefused {
    /// Workflow whose activity is being declined.
    #[prost(message, optional, tag = "1")]
    pub workflow_id: Option<ProtoWorkflowId>,
    /// Activity within that workflow, echoed from the dispatched task.
    #[prost(message, optional, tag = "2")]
    pub activity_id: Option<ProtoActivityId>,
    /// Operator-facing reason the worker declined, in the worker's own words.
    #[prost(string, tag = "3")]
    pub reason: String,
    /// Verbatim execution-generation token from the dispatched task.
    ///
    /// Load-bearing rather than decorative: a refusal that arrived from a
    /// SUPERSEDED generation must not re-park the attempt that superseded it,
    /// and this token is what the server's completion fence checks that with.
    #[prost(string, tag = "4")]
    pub completion_token: String,
}

/// Activity invocation pushed to a worker.
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, prost::Message)]
pub struct ProtoActivityTask {
    /// Owning workflow id.
    #[prost(message, optional, tag = "1")]
    pub workflow_id: Option<ProtoWorkflowId>,
    /// Correlating activity id.
    #[prost(message, optional, tag = "2")]
    pub activity_id: Option<ProtoActivityId>,
    /// Activity type name.
    #[prost(string, tag = "3")]
    pub activity_type: String,
    /// Serialized activity input.
    #[prost(message, optional, tag = "4")]
    pub input: Option<ProtoPayload>,
    /// One-based delivery attempt stamped by the dispatching engine seam.
    /// Zero is malformed: consumers reject a task whose attempt is 0 (the
    /// proto3 default means the producer failed to stamp it).
    #[prost(uint32, tag = "5")]
    pub attempt: u32,
    /// Human-meaningful display labels the workflow attached to the activity
    /// (for example `brief=IP-001`). Display metadata only — the engine never
    /// interprets them; they ride to the worker so its logs and the ops console
    /// can show what a dispatch is working on. Empty when none were set.
    #[prost(map = "string, string", tag = "6")]
    pub labels: ::std::collections::HashMap<String, String>,
    /// Concrete run that owns this dispatch. The protobuf field remains optional
    /// so old bytes can be decoded, but current servers and workers refuse its
    /// absence rather than execute an effect without run scope.
    #[prost(message, optional, tag = "7")]
    pub run_id: Option<ProtoRunId>,
    /// Opaque execution-generation proof assigned by the server. Workers echo
    /// this verbatim and never interpret it. Empty is an incompatible task.
    #[prost(string, tag = "8")]
    pub completion_token: String,
    /// Stable external-effect key for this run and action site. Attempts and
    /// execution generations do not participate in its derivation. Empty is an
    /// incompatible task.
    #[prost(string, tag = "9")]
    pub idempotency_key: String,
}

/// Server-initiated drain: the server is going away (restart, deploy,
/// rebalance). The worker finishes already-assigned work, stops expecting
/// new tasks, and reconnects after the schedule's initial backoff. A drain
/// frame re-classifies the session's eventual stream end (clean or abrupt)
/// as a drain-class drop that consumes no drop budget — distinct from denial
/// (gRPC error status, terminal) and from an unannounced close (budgeted
/// retryable drop).
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, prost::Message)]
pub struct ProtoDrainRequest {}

/// Positive registration acknowledgement — always the first frame on the
/// response stream. There is no negative counterpart: a denied or invalid
/// registration fails the RPC with a gRPC error status.
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, prost::Message)]
pub struct ProtoRegisterAck {
    /// Server-assigned stream identifier for this registration.
    #[prost(uint64, tag = "1")]
    pub worker_id: u64,
    /// The namespace the registration was authorized against.
    #[prost(string, tag = "2")]
    pub namespace: String,
    /// Operator-configured liveness window on this server, in milliseconds.
    #[prost(uint64, tag = "3")]
    pub heartbeat_window_ms: u64,
    /// The concurrency the server RECORDED for this worker, echoed from the
    /// registration. The server's accounting is the one thing the worker cannot
    /// see, so echoing it makes a drift between the two diagnosable at the
    /// single place both sides agree.
    #[prost(uint32, tag = "4")]
    pub accepted_max_concurrency: u32,
}

/// Per-result acknowledgement: the server has consumed the identified
/// `ActivityResult` frame and the worker may stop re-reporting it. Not a
/// durability receipt — the durable truth is the workflow's event history.
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, prost::Message)]
pub struct ProtoResultAck {
    /// Owning workflow id.
    #[prost(message, optional, tag = "1")]
    pub workflow_id: Option<ProtoWorkflowId>,
    /// Correlating activity id.
    #[prost(message, optional, tag = "2")]
    pub activity_id: Option<ProtoActivityId>,
}

/// Activity result or failure reported by a worker.
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, prost::Message)]
pub struct ProtoActivityResult {
    /// Owning workflow id.
    #[prost(message, optional, tag = "1")]
    pub workflow_id: Option<ProtoWorkflowId>,
    /// Correlating activity id.
    #[prost(message, optional, tag = "2")]
    pub activity_id: Option<ProtoActivityId>,
    /// Successful result payload or explicit activity error.
    #[prost(oneof = "proto_activity_result::Outcome", tags = "3, 4")]
    pub outcome: Option<proto_activity_result::Outcome>,
    /// Concrete run echoed from the dispatched task. The protobuf field remains
    /// optional so old bytes can be decoded, but current servers reject its
    /// absence rather than bypass the run gate.
    #[prost(message, optional, tag = "5")]
    pub run_id: Option<ProtoRunId>,
    /// Verbatim execution-generation token from the dispatched task. Empty is
    /// never treated as current; it denotes an incompatible pre-fencing worker.
    #[prost(string, tag = "6")]
    pub completion_token: String,
}

/// Types nested under [`ProtoActivityResult`].
pub mod proto_activity_result {
    use super::{ProtoActivityError, ProtoPayload};

    /// Proto oneof for activity success or failure.
    #[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, prost::Oneof)]
    pub enum Outcome {
        /// Successful activity output.
        #[prost(message, tag = "3")]
        Result(ProtoPayload),
        /// Activity failure preserving retryability classification.
        #[prost(message, tag = "4")]
        Error(ProtoActivityError),
    }
}

/// Transport-level liveness ping pushed by the server down a worker's task
/// stream — the gRPC counterpart of the liminal transport's `LivenessPing`.
///
/// Answered by the SDK RUNTIME, never by action code. It rides the SAME stream
/// a dispatch travels, because the only fact worth measuring is whether the
/// server can reach this worker's DISPATCH path; a parallel channel could be
/// healthy while the real one is not, and an open `mpsc` sender is not an
/// answering worker.
///
/// Distinct from [`ProtoHeartbeat`], which travels worker-to-server and proves
/// the worker PROCESS is alive. Only the server-to-worker-and-back round trip
/// this pair performs is a dispatch precondition.
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, prost::Message)]
pub struct ProtoLivenessPing {
    /// Monotonic ping sequence within this server's probe, echoed verbatim on
    /// the answer so a stale or reordered answer cannot be counted as fresh.
    #[prost(uint64, tag = "1")]
    pub liveness_ping: u64,
    /// How long the worker may hear NOTHING on this stream before it should
    /// treat the link as dead — this server's `worker.heartbeat_window`,
    /// carried on the wire so the worker holds no second copy of it.
    #[prost(uint64, tag = "2")]
    pub silence_window_ms: u64,
}

/// Server-initiated cancellation of ONE in-flight activity (#233).
///
/// A REQUEST, never a guarantee. The server asks the worker holding this
/// activity to stop it; the worker fires that activity's cancellation handle,
/// which cooperative action code observes and which terminates a contained
/// child process group. A worker that no longer holds the named activity
/// ignores it — the `(workflow_id, activity_id)` pair is globally unique, so
/// this can never stop somebody else's work.
///
/// Distinct from [`ProtoDrainRequest`]: drain is about the WORKER (finish what
/// you hold, expect nothing new, reconnect), this is about ONE ACTIVITY (stop
/// this piece of work, keep serving). A drained worker still finishes the very
/// work this message exists to interrupt.
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, prost::Message)]
pub struct ProtoCancelActivity {
    /// Workflow whose activity is to be cancelled.
    #[prost(message, optional, tag = "1")]
    pub workflow_id: Option<ProtoWorkflowId>,
    /// Activity within that workflow. With `workflow_id` it forms the key the
    /// worker's in-flight registry is addressed by.
    #[prost(message, optional, tag = "2")]
    pub activity_id: Option<ProtoActivityId>,
}

/// The worker runtime's answer to a [`ProtoLivenessPing`], echoing its sequence.
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, prost::Message)]
pub struct ProtoLivenessAnswer {
    /// The sequence of the ping being answered, echoed verbatim.
    #[prost(uint64, tag = "1")]
    pub liveness_ping: u64,
}

/// Worker heartbeat for an in-flight activity.
#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize, prost::Message)]
pub struct ProtoHeartbeat {
    /// Owning workflow id.
    #[prost(message, optional, tag = "1")]
    pub workflow_id: Option<ProtoWorkflowId>,
    /// Correlating activity id.
    #[prost(message, optional, tag = "2")]
    pub activity_id: Option<ProtoActivityId>,
    /// Optional opaque progress payload.
    #[prost(message, optional, tag = "3")]
    pub progress: Option<ProtoPayload>,
}

impl From<aion_core::ActivityErrorKind> for ProtoActivityErrorKind {
    fn from(value: aion_core::ActivityErrorKind) -> Self {
        match value {
            aion_core::ActivityErrorKind::Retryable => Self::Retryable,
            aion_core::ActivityErrorKind::PolicyRefused => Self::PolicyRefused,
            aion_core::ActivityErrorKind::Terminal => Self::Terminal,
        }
    }
}

impl TryFrom<ProtoActivityErrorKind> for aion_core::ActivityErrorKind {
    type Error = WireError;

    fn try_from(value: ProtoActivityErrorKind) -> Result<Self, Self::Error> {
        match value {
            ProtoActivityErrorKind::Unspecified => {
                Err(WireError::backend("activity error kind is missing"))
            }
            ProtoActivityErrorKind::Retryable => Ok(Self::Retryable),
            ProtoActivityErrorKind::PolicyRefused => Ok(Self::PolicyRefused),
            ProtoActivityErrorKind::Terminal => Ok(Self::Terminal),
        }
    }
}

impl From<aion_core::ActivityError> for ProtoActivityError {
    fn from(value: aion_core::ActivityError) -> Self {
        Self {
            kind: ProtoActivityErrorKind::from(value.kind) as i32,
            message: value.message,
            details: value.details.map(ProtoPayload::from),
        }
    }
}

impl TryFrom<ProtoActivityError> for aion_core::ActivityError {
    type Error = WireError;

    fn try_from(value: ProtoActivityError) -> Result<Self, Self::Error> {
        let kind = ProtoActivityErrorKind::try_from(value.kind)
            .map_err(|_| WireError::backend("activity error kind is unknown"))?;
        Ok(Self {
            kind: aion_core::ActivityErrorKind::try_from(kind)?,
            message: value.message,
            details: value
                .details
                .map(aion_core::Payload::try_from)
                .transpose()?,
        })
    }
}

#[cfg(test)]
mod tests {
    use prost::Message;
    use serde_json::json;

    use super::{
        ProtoActivityError, ProtoActivityErrorKind, ProtoActivityResult, ProtoActivityTask,
        ProtoDrainRequest, ProtoHeartbeat, ProtoLivenessAnswer, ProtoLivenessPing,
        ProtoRegisterAck, ProtoRegisterWorker, ProtoResultAck, ProtoTaskRefused,
        proto_activity_result,
    };
    use crate::{ProtoActivityId, ProtoPayload, ProtoWorkflowId, WireError};

    fn workflow_id() -> aion_core::WorkflowId {
        aion_core::WorkflowId::new(uuid::Uuid::nil())
    }

    #[test]
    fn activity_error_round_trips_preserving_classification() -> Result<(), WireError> {
        let core = aion_core::ActivityError {
            kind: aion_core::ActivityErrorKind::Retryable,
            message: String::from("connection reset"),
            details: Some(
                aion_core::Payload::from_json(&json!({"retry_after_ms": 500}))
                    .map_err(|_| WireError::backend("test payload could not be created"))?,
            ),
        };

        let proto = ProtoActivityError::from(core.clone());
        assert_eq!(aion_core::ActivityError::try_from(proto.clone())?, core);
        assert!(aion_core::ActivityError::try_from(proto)?.is_retryable());

        let policy_refused = ProtoActivityError {
            kind: ProtoActivityErrorKind::PolicyRefused as i32,
            message: String::from("provider safety policy"),
            details: None,
        };
        let policy_refused = aion_core::ActivityError::try_from(policy_refused)?;
        assert_eq!(
            policy_refused.kind,
            aion_core::ActivityErrorKind::PolicyRefused
        );
        assert!(!policy_refused.is_retryable());

        let terminal = ProtoActivityError {
            kind: ProtoActivityErrorKind::Terminal as i32,
            message: String::from("invalid request"),
            details: None,
        };
        assert!(!aion_core::ActivityError::try_from(terminal)?.is_retryable());

        Ok(())
    }

    #[test]
    fn worker_registration_round_trips_through_serde_and_proto()
    -> Result<(), Box<dyn std::error::Error>> {
        let registration = ProtoRegisterWorker {
            namespaces: vec![String::from("tenant-a"), String::from("tenant-b")],
            activity_types: vec![String::from("charge-card"), String::from("send-email")],
            task_queue: String::from("claude"),
            node: String::from("host-7"),
            activities: Vec::new(),
            identity: String::new(),
            instance: Some(super::ProtoWorkerInstanceIdentity {
                deployment: String::from("shells"),
                instance_id: String::from("instance-7"),
            }),
            max_concurrency: Some(4),
        };

        assert_json_and_proto_round_trip(&registration)
    }

    #[test]
    fn worker_registration_repeated_namespaces_round_trip() -> Result<(), Box<dyn std::error::Error>>
    {
        // A worker advertising a SET of namespaces survives a proto and serde
        // round trip preserving every namespace in wire order.
        let registration = ProtoRegisterWorker {
            namespaces: vec![String::from("a"), String::from("b"), String::from("c")],
            activity_types: vec![String::from("dev")],
            task_queue: String::from("default"),
            node: String::new(),
            activities: Vec::new(),
            identity: String::new(),
            instance: None,
            max_concurrency: None,
        };
        let decoded =
            ProtoRegisterWorker::decode(prost_round_trip_bytes(&registration)?.as_slice())?;
        assert_eq!(decoded.namespaces, vec!["a", "b", "c"]);
        assert_eq!(decoded.node, "");
        assert_json_and_proto_round_trip(&registration)
    }

    #[test]
    fn worker_registration_task_queue_and_node_use_wire_tags_three_and_four()
    -> Result<(), Box<dyn std::error::Error>> {
        // Pins task_queue to proto tag 3 and node to tag 4 (field keys 0x1A and
        // 0x22 = tags 3/4, length-delimited) so the hand-written stubs cannot
        // drift, and confirms an old-shape registration that omits node decodes
        // to the proto3 default `""` (= no locality affinity).
        let registration = ProtoRegisterWorker {
            namespaces: Vec::new(),
            activity_types: Vec::new(),
            task_queue: String::from("gpu"),
            node: String::from("n1"),
            activities: Vec::new(),
            identity: String::new(),
            instance: None,
            max_concurrency: None,
        };
        let mut bytes = Vec::new();
        registration.encode(&mut bytes)?;
        assert_eq!(
            bytes,
            vec![0x1A, 0x03, b'g', b'p', b'u', 0x22, 0x02, b'n', b'1']
        );

        // An encoded registration with no tag-4 field decodes node to "".
        let no_node = ProtoRegisterWorker {
            namespaces: vec![String::from("tenant-a")],
            activity_types: vec![String::from("charge-card")],
            task_queue: String::new(),
            node: String::new(),
            activities: Vec::new(),
            identity: String::new(),
            instance: None,
            max_concurrency: None,
        };
        let mut bytes = Vec::new();
        no_node.encode(&mut bytes)?;
        let decoded = ProtoRegisterWorker::decode(bytes.as_slice())?;
        assert_eq!(decoded.task_queue, "");
        assert_eq!(decoded.node, "");
        Ok(())
    }

    #[test]
    fn worker_instance_identity_uses_wire_tag_seven() -> Result<(), Box<dyn std::error::Error>> {
        let registration = ProtoRegisterWorker {
            namespaces: Vec::new(),
            activity_types: Vec::new(),
            task_queue: String::new(),
            node: String::new(),
            activities: Vec::new(),
            identity: String::new(),
            instance: Some(super::ProtoWorkerInstanceIdentity {
                deployment: "d".to_owned(),
                instance_id: "i".to_owned(),
            }),
            max_concurrency: None,
        };
        let mut bytes = Vec::new();
        registration.encode(&mut bytes)?;
        assert_eq!(bytes, vec![0x3a, 0x06, 0x0a, 0x01, b'd', 0x12, 0x01, b'i']);

        let absent = ProtoRegisterWorker {
            namespaces: Vec::new(),
            activity_types: Vec::new(),
            task_queue: String::new(),
            node: String::new(),
            activities: Vec::new(),
            identity: String::new(),
            instance: None,
            max_concurrency: None,
        };
        let mut absent_bytes = Vec::new();
        absent.encode(&mut absent_bytes)?;
        assert!(absent_bytes.is_empty());
        Ok(())
    }

    /// Pins advertised capacity to proto tag 8 AND pins its PRESENCE.
    ///
    /// Two facts, and the second is the one that matters: `optional` means an
    /// absent value encodes to nothing at all, while `Some(0)` encodes a real
    /// tag-8 varint of zero. A bare `uint32` would collapse those into the same
    /// bytes, and the server would have to invent a meaning for the collapse —
    /// the assumed default this field is `optional` to avoid. Nothing else in
    /// the tree would catch a mis-numbered or presence-losing capacity field:
    /// proto tags are not guarded by exhaustiveness.
    #[test]
    fn worker_registration_max_concurrency_uses_wire_tag_eight_and_keeps_presence()
    -> Result<(), Box<dyn std::error::Error>> {
        let advertised = ProtoRegisterWorker {
            namespaces: Vec::new(),
            activity_types: Vec::new(),
            task_queue: String::new(),
            node: String::new(),
            activities: Vec::new(),
            identity: String::new(),
            instance: None,
            max_concurrency: Some(8),
        };
        let mut bytes = Vec::new();
        advertised.encode(&mut bytes)?;
        // Field key 0x40 = (tag 8 << 3) | varint wire type 0.
        assert_eq!(bytes, vec![0x40, 8]);
        assert_eq!(
            ProtoRegisterWorker::decode(bytes.as_slice())?.max_concurrency,
            Some(8)
        );

        // A ZERO capacity is a value the wire carries, not an absence: the
        // server refuses it by NAME rather than by failing to see it.
        let zero = ProtoRegisterWorker {
            max_concurrency: Some(0),
            ..advertised.clone()
        };
        let mut zero_bytes = Vec::new();
        zero.encode(&mut zero_bytes)?;
        assert_eq!(zero_bytes, vec![0x40, 0]);
        assert_eq!(
            ProtoRegisterWorker::decode(zero_bytes.as_slice())?.max_concurrency,
            Some(0)
        );

        // And an absent one encodes to nothing, so the two are distinguishable
        // on the wire rather than only in this crate's types.
        let absent = ProtoRegisterWorker {
            max_concurrency: None,
            ..advertised
        };
        let mut absent_bytes = Vec::new();
        absent.encode(&mut absent_bytes)?;
        assert!(absent_bytes.is_empty());
        assert_eq!(
            ProtoRegisterWorker::decode(absent_bytes.as_slice())?.max_concurrency,
            None
        );
        Ok(())
    }

    /// The ack's echo of the recorded capacity sits at tag 4, beside the
    /// heartbeat window the worker already reads off the same frame.
    #[test]
    fn register_ack_accepted_max_concurrency_uses_wire_tag_four()
    -> Result<(), Box<dyn std::error::Error>> {
        let ack = ProtoRegisterAck {
            worker_id: 0,
            namespace: String::new(),
            heartbeat_window_ms: 0,
            accepted_max_concurrency: 6,
        };
        let mut bytes = Vec::new();
        ack.encode(&mut bytes)?;
        // Field key 0x20 = (tag 4 << 3) | varint wire type 0.
        assert_eq!(bytes, vec![0x20, 6]);
        Ok(())
    }

    /// The refusal round-trips and pins its own field tags. Its ids are the
    /// dispatched task's, echoed verbatim, so the server can key the re-park
    /// against the same `(workflow, activity)` it pushed.
    #[test]
    fn task_refused_round_trips_and_pins_its_wire_tags() -> Result<(), Box<dyn std::error::Error>> {
        let refusal = ProtoTaskRefused {
            workflow_id: Some(ProtoWorkflowId::from(workflow_id())),
            activity_id: Some(ProtoActivityId::from(
                aion_core::ActivityId::from_sequence_position(4),
            )),
            reason: String::from("worker is at its advertised concurrency"),
            completion_token: String::from("generation-2"),
        };
        let reason_only = ProtoTaskRefused {
            workflow_id: None,
            activity_id: None,
            reason: String::from("no"),
            completion_token: String::new(),
        };
        let mut bytes = Vec::new();
        reason_only.encode(&mut bytes)?;
        // Field key 0x1A = (tag 3 << 3) | length-delimited wire type 2.
        assert_eq!(bytes, vec![0x1A, 0x02, b'n', b'o']);
        assert_json_and_proto_round_trip(&refusal)
    }

    /// The refusal is the FIFTH `WorkerToServer` arm, at oneof tag 5, and it
    /// must not decode as any of the four that came before it — the demux the
    /// server's inbound loop relies on to keep a refusal out of the result,
    /// heartbeat and liveness paths.
    #[cfg(feature = "generated")]
    #[test]
    fn task_refused_is_the_fifth_worker_to_server_arm() -> Result<(), Box<dyn std::error::Error>> {
        let refused = crate::generated::WorkerToServer {
            message: Some(crate::generated::worker_to_server::Message::TaskRefused(
                crate::generated::TaskRefused {
                    workflow_id: None,
                    activity_id: None,
                    reason: String::from("at capacity"),
                    completion_token: String::from("generation-2"),
                },
            )),
        };
        let mut bytes = Vec::new();
        refused.encode(&mut bytes)?;
        // Field key 0x2A = (tag 5 << 3) | length-delimited wire type 2.
        assert_eq!(bytes.first(), Some(&0x2A));
        assert_eq!(
            crate::generated::WorkerToServer::decode(bytes.as_slice())?,
            refused
        );
        Ok(())
    }

    #[test]
    fn activity_task_round_trips_through_serde_and_proto() -> Result<(), Box<dyn std::error::Error>>
    {
        let task = ProtoActivityTask {
            workflow_id: Some(ProtoWorkflowId::from(workflow_id())),
            activity_id: Some(ProtoActivityId::from(
                aion_core::ActivityId::from_sequence_position(7),
            )),
            activity_type: String::from("charge-card"),
            input: Some(ProtoPayload::from(aion_core::Payload::from_json(
                &json!({"amount": 42}),
            )?)),
            attempt: 3,
            labels: [
                (String::from("brief"), String::from("IP-001")),
                (String::from("repo"), String::from("ablative-io/yggdrasil")),
            ]
            .into_iter()
            .collect(),
            run_id: None,
            completion_token: String::from("generation-3"),
            idempotency_key: String::from("effect-key"),
        };

        assert_json_and_proto_round_trip(&task)
    }

    #[test]
    fn drain_request_round_trips_through_serde_and_proto() -> Result<(), Box<dyn std::error::Error>>
    {
        assert_json_and_proto_round_trip(&ProtoDrainRequest {})
    }

    #[test]
    fn register_ack_round_trips_through_serde_and_proto() -> Result<(), Box<dyn std::error::Error>>
    {
        let ack = ProtoRegisterAck {
            worker_id: 7,
            namespace: String::from("tenant-a"),
            heartbeat_window_ms: 30_000,
            accepted_max_concurrency: 8,
        };

        assert_json_and_proto_round_trip(&ack)
    }

    #[test]
    fn result_ack_round_trips_through_serde_and_proto() -> Result<(), Box<dyn std::error::Error>> {
        let ack = ProtoResultAck {
            workflow_id: Some(ProtoWorkflowId::from(workflow_id())),
            activity_id: Some(ProtoActivityId::from(
                aion_core::ActivityId::from_sequence_position(11),
            )),
        };

        assert_json_and_proto_round_trip(&ack)
    }

    #[cfg(feature = "generated")]
    #[test]
    fn server_to_worker_ack_arms_pin_oneof_tags_three_and_four()
    -> Result<(), Box<dyn std::error::Error>> {
        // Pins the new ServerToWorker oneof arms to wire tags 3 (register_ack)
        // and 4 (result_ack): field key = (tag << 3) | 2 (length-delimited).
        let register_ack = crate::generated::ServerToWorker {
            message: Some(crate::generated::server_to_worker::Message::RegisterAck(
                crate::generated::RegisterAck {
                    worker_id: 1,
                    namespace: String::from("tenant-a"),
                    heartbeat_window_ms: 1_000,
                    accepted_max_concurrency: 3,
                },
            )),
        };
        let mut bytes = Vec::new();
        register_ack.encode(&mut bytes)?;
        assert_eq!(bytes.first(), Some(&0x1A));
        assert_eq!(
            crate::generated::ServerToWorker::decode(bytes.as_slice())?,
            register_ack
        );

        let result_ack = crate::generated::ServerToWorker {
            message: Some(crate::generated::server_to_worker::Message::ResultAck(
                crate::generated::ResultAck {
                    workflow_id: None,
                    activity_id: None,
                },
            )),
        };
        let mut bytes = Vec::new();
        result_ack.encode(&mut bytes)?;
        assert_eq!(bytes.first(), Some(&0x22));
        assert_eq!(
            crate::generated::ServerToWorker::decode(bytes.as_slice())?,
            result_ack
        );
        Ok(())
    }

    #[test]
    fn activity_task_attempt_uses_wire_tag_five() -> Result<(), Box<dyn std::error::Error>> {
        // Pins the attempt field to proto tag 5 (field key 0x28 = tag 5,
        // varint wire type) so the hand-written SDK stubs cannot drift.
        let task = ProtoActivityTask {
            workflow_id: None,
            activity_id: None,
            activity_type: String::new(),
            input: None,
            attempt: 9,
            labels: ::std::collections::HashMap::new(),
            run_id: None,
            completion_token: String::new(),
            idempotency_key: String::new(),
        };
        let mut bytes = Vec::new();
        task.encode(&mut bytes)?;
        assert_eq!(bytes, vec![0x28, 9]);
        Ok(())
    }

    #[test]
    fn activity_success_result_round_trips_through_serde_and_proto()
    -> Result<(), Box<dyn std::error::Error>> {
        let result = ProtoActivityResult {
            workflow_id: Some(ProtoWorkflowId::from(workflow_id())),
            activity_id: Some(ProtoActivityId::from(
                aion_core::ActivityId::from_sequence_position(8),
            )),
            outcome: Some(proto_activity_result::Outcome::Result(ProtoPayload::from(
                aion_core::Payload::from_json(&json!({"authorization": "ok"}))?,
            ))),
            run_id: None,
            completion_token: String::from("generation-1"),
        };

        assert_json_and_proto_round_trip(&result)
    }

    #[test]
    fn activity_error_result_round_trips_through_serde_and_proto()
    -> Result<(), Box<dyn std::error::Error>> {
        let result = ProtoActivityResult {
            workflow_id: Some(ProtoWorkflowId::from(workflow_id())),
            activity_id: Some(ProtoActivityId::from(
                aion_core::ActivityId::from_sequence_position(9),
            )),
            outcome: Some(proto_activity_result::Outcome::Error(
                ProtoActivityError::from(aion_core::ActivityError {
                    kind: aion_core::ActivityErrorKind::Terminal,
                    message: String::from("card declined"),
                    details: Some(aion_core::Payload::from_json(&json!({"code": "declined"}))?),
                }),
            )),
            run_id: None,
            completion_token: String::from("generation-1"),
        };

        assert_json_and_proto_round_trip(&result)
    }

    /// The liveness pair round-trips and pins its wire tags: `liveness_ping`
    /// at tag 1 (field key 0x08, varint) and `silence_window_ms` at tag 2
    /// (0x10). The Python/TS stubs are regenerated from the same `.proto`
    /// (#201), so a silent tag drift here would land as a worker that answers
    /// with the wrong sequence and never clears probation.
    #[test]
    fn liveness_ping_round_trips_and_pins_its_wire_tags() -> Result<(), Box<dyn std::error::Error>>
    {
        let ping = ProtoLivenessPing {
            liveness_ping: 7,
            silence_window_ms: 30_000,
        };
        let mut bytes = Vec::new();
        ping.encode(&mut bytes)?;
        assert_eq!(bytes, vec![0x08, 7, 0x10, 0xB0, 0xEA, 0x01]);
        assert_json_and_proto_round_trip(&ping)
    }

    /// The answer echoes the ping's sequence on the SAME field number, so an
    /// SDK that echoes the frame it received is correct by construction.
    #[test]
    fn liveness_answer_round_trips_and_echoes_the_ping_sequence_field()
    -> Result<(), Box<dyn std::error::Error>> {
        let answer = ProtoLivenessAnswer { liveness_ping: 7 };
        let mut bytes = Vec::new();
        answer.encode(&mut bytes)?;
        assert_eq!(bytes, vec![0x08, 7]);
        // The answer's sequence field shares the ping's tag AND its varint
        // encoding: the first two bytes of an encoded ping with the same
        // sequence are byte-identical.
        let mut ping_bytes = Vec::new();
        ProtoLivenessPing {
            liveness_ping: 7,
            silence_window_ms: 0,
        }
        .encode(&mut ping_bytes)?;
        assert_eq!(ping_bytes, bytes);
        assert_json_and_proto_round_trip(&answer)
    }

    /// A liveness ping must never decode as any other pushed frame, and no
    /// other pushed frame may decode as a ping — the demux contract the worker
    /// runtime's stream arm relies on. Asserted on the ONEOF wrapper, which is
    /// what the wire actually carries: bare prost messages are permissive
    /// enough that a same-shaped payload could decode either way, and the
    /// discriminating fact is the oneof tag.
    #[cfg(feature = "generated")]
    #[test]
    fn liveness_frames_are_disjoint_from_the_other_stream_frames()
    -> Result<(), Box<dyn std::error::Error>> {
        // Pins the ServerToWorker liveness arm to tag 5 (field key 0x2A) and
        // the WorkerToServer answer arm to tag 4 (0x22).
        let ping = crate::generated::ServerToWorker {
            message: Some(crate::generated::server_to_worker::Message::LivenessPing(
                crate::generated::LivenessPing {
                    liveness_ping: 3,
                    silence_window_ms: 1_000,
                },
            )),
        };
        let mut bytes = Vec::new();
        ping.encode(&mut bytes)?;
        assert_eq!(bytes.first(), Some(&0x2A));
        assert_eq!(
            crate::generated::ServerToWorker::decode(bytes.as_slice())?,
            ping
        );

        let answer = crate::generated::WorkerToServer {
            message: Some(crate::generated::worker_to_server::Message::LivenessAnswer(
                crate::generated::LivenessAnswer { liveness_ping: 3 },
            )),
        };
        let mut bytes = Vec::new();
        answer.encode(&mut bytes)?;
        assert_eq!(bytes.first(), Some(&0x22));
        assert_eq!(
            crate::generated::WorkerToServer::decode(bytes.as_slice())?,
            answer
        );
        Ok(())
    }

    #[test]
    fn heartbeat_round_trips_through_serde_and_proto() -> Result<(), Box<dyn std::error::Error>> {
        let heartbeat = ProtoHeartbeat {
            workflow_id: Some(ProtoWorkflowId::from(workflow_id())),
            activity_id: Some(ProtoActivityId::from(
                aion_core::ActivityId::from_sequence_position(10),
            )),
            progress: Some(ProtoPayload::from(aion_core::Payload::from_json(
                &json!({"percent": 50}),
            )?)),
        };

        assert_json_and_proto_round_trip(&heartbeat)
    }

    fn assert_json_and_proto_round_trip<T>(value: &T) -> Result<(), Box<dyn std::error::Error>>
    where
        T: Message
            + Default
            + serde::Serialize
            + serde::de::DeserializeOwned
            + PartialEq
            + std::fmt::Debug,
    {
        assert_eq!(
            serde_json::from_str::<T>(&serde_json::to_string(value)?)?,
            *value
        );
        assert_eq!(prost_round_trip(value)?, *value);
        Ok(())
    }

    fn prost_round_trip<T>(value: &T) -> Result<T, Box<dyn std::error::Error>>
    where
        T: Message + Default,
    {
        Ok(T::decode(prost_round_trip_bytes(value)?.as_slice())?)
    }

    fn prost_round_trip_bytes<T>(value: &T) -> Result<Vec<u8>, Box<dyn std::error::Error>>
    where
        T: Message,
    {
        let mut bytes = Vec::new();
        value.encode(&mut bytes)?;
        Ok(bytes)
    }
}