Merge remote-tracking branch 'origin/main' into SDK-client-UA-registration

server/crates/arbiter-server/src/actors/bootstrap.rs

@@ -3,12 +3,7 @@ use diesel::QueryDsl;
 use diesel_async::RunQueryDsl;
 use kameo::{Actor, messages};
 use miette::Diagnostic;
-use rand::{
-    RngExt,
-    distr::{Alphanumeric},
-    make_rng,
-    rngs::StdRng,
-};
+use rand::{RngExt, distr::Alphanumeric, make_rng, rngs::StdRng};
 use thiserror::Error;
 
 use crate::db::{self, DatabasePool, schema};
@@ -61,7 +56,6 @@ impl Bootstrapper {
 
         drop(conn);
 
-
         let token = if row_count == 0 {
             let token = generate_token().await?;
             Some(token)

server/crates/arbiter-server/src/actors/client/auth.rs

@@ -1,18 +1,12 @@
 use arbiter_proto::{
     format_challenge,
-    proto::client::{
-        AuthChallenge, AuthChallengeSolution, AuthOk, ClientConnectError, ClientRequest,
-        ClientResponse, client_connect_error::Code as ConnectErrorCode,
-        client_request::Payload as ClientRequestPayload,
-        client_response::Payload as ClientResponsePayload,
-    },
-    transport::expect_message,
+    transport::{Bi, expect_message},
 };
 use diesel::{
     ExpressionMethods as _, OptionalExtension as _, QueryDsl as _, dsl::insert_into, update,
 };
 use diesel_async::RunQueryDsl as _;
-use ed25519_dalek::VerifyingKey;
+use ed25519_dalek::{Signature, VerifyingKey};
 use kameo::error::SendError;
 use tracing::error;
 
@@ -24,35 +18,8 @@ use crate::{
     db::{self, schema::program_client},
 };
 
-use super::session::ClientSession;
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub struct ClientId(i32);
-
-impl ClientId {
-    pub fn new(raw: i32) -> Self {
-        Self(raw)
-    }
-
-    pub fn as_i32(self) -> i32 {
-        self.0
-    }
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-struct ClientNonceState {
-    client_id: ClientId,
-    nonce: i32,
-}
-
 #[derive(thiserror::Error, Debug, Clone, PartialEq, Eq)]
 pub enum Error {
-    #[error("Unexpected message payload")]
-    UnexpectedMessagePayload,
-    #[error("Invalid client public key length")]
-    InvalidClientPubkeyLength,
-    #[error("Invalid client public key encoding")]
-    InvalidAuthPubkeyEncoding,
     #[error("Database pool unavailable")]
     DatabasePoolUnavailable,
     #[error("Database operation failed")]
@@ -61,8 +28,6 @@ pub enum Error {
     InvalidChallengeSolution,
     #[error("Client approval request failed")]
     ApproveError(#[from] ApproveError),
-    #[error("Internal error")]
-    InternalError,
     #[error("Transport error")]
     Transport,
 }
@@ -77,12 +42,21 @@ pub enum ApproveError {
     Upstream(router::ApprovalError),
 }
 
+#[derive(Debug, Clone)]
+pub enum Inbound {
+    AuthChallengeRequest { pubkey: VerifyingKey },
+    AuthChallengeSolution { signature: Signature },
+}
+
+#[derive(Debug, Clone)]
+pub enum Outbound {
+    AuthChallenge { pubkey: VerifyingKey, nonce: i32 },
+    AuthSuccess,
+}
+
 /// Atomically reads and increments the nonce for a known client.
 /// Returns `None` if the pubkey is not registered.
-async fn get_nonce(
-    db: &db::DatabasePool,
-    pubkey: &VerifyingKey,
-) -> Result<Option<ClientNonceState>, Error> {
+async fn get_nonce(db: &db::DatabasePool, pubkey: &VerifyingKey) -> Result<Option<i32>, Error> {
     let pubkey_bytes = pubkey.as_bytes().to_vec();
 
     let mut conn = db.get().await.map_err(|e| {
@@ -109,10 +83,8 @@ async fn get_nonce(
                 .execute(conn)
                 .await?;
 
-            Ok(Some(ClientNonceState {
-                client_id: ClientId::new(client_id),
-                nonce: current_nonce,
-            }))
+            let _ = client_id;
+            Ok(Some(current_nonce))
         })
     })
     .await
@@ -148,7 +120,7 @@ async fn approve_new_client(
 }
 
 enum InsertClientResult {
-    Inserted(ClientId),
+    Inserted,
     AlreadyExists,
 }
 
@@ -198,126 +170,80 @@ async fn insert_client(
             Error::DatabaseOperationFailed
         })?;
 
-    Ok(InsertClientResult::Inserted(ClientId::new(client_id)))
+    let _ = client_id;
+    Ok(InsertClientResult::Inserted)
 }
 
-async fn challenge_client(
-    props: &mut ClientConnection,
+async fn challenge_client<T>(
+    transport: &mut T,
     pubkey: VerifyingKey,
     nonce: i32,
-) -> Result<(), Error> {
-    let challenge = AuthChallenge {
-        pubkey: pubkey.as_bytes().to_vec(),
-        nonce,
-    };
-
-    props
-        .transport
-        .send(Ok(ClientResponse {
-            payload: Some(ClientResponsePayload::AuthChallenge(challenge.clone())),
-        }))
+) -> Result<(), Error>
+where
+    T: Bi<Inbound, Result<Outbound, Error>> + ?Sized,
+{
+    transport
+        .send(Ok(Outbound::AuthChallenge { pubkey, nonce }))
         .await
         .map_err(|e| {
             error!(error = ?e, "Failed to send auth challenge");
             Error::Transport
         })?;
 
-    let AuthChallengeSolution { signature } =
-        expect_message(&mut *props.transport, |req: ClientRequest| {
-            match req.payload? {
-                ClientRequestPayload::AuthChallengeSolution(s) => Some(s),
-                _ => None,
-            }
-        })
-        .await
-        .map_err(|e| {
-            error!(error = ?e, "Failed to receive challenge solution");
-            Error::Transport
-        })?;
-
-    let formatted = format_challenge(nonce, &challenge.pubkey);
-    let sig = signature.as_slice().try_into().map_err(|_| {
-        error!("Invalid signature length");
-        Error::InvalidChallengeSolution
+    let signature = expect_message(transport, |req: Inbound| match req {
+        Inbound::AuthChallengeSolution { signature } => Some(signature),
+        _ => None,
+    })
+    .await
+    .map_err(|e| {
+        error!(error = ?e, "Failed to receive challenge solution");
+        Error::Transport
     })?;
 
-    pubkey.verify_strict(&formatted, &sig).map_err(|_| {
+    let formatted = format_challenge(nonce, pubkey.as_bytes());
+
+    pubkey.verify_strict(&formatted, &signature).map_err(|_| {
         error!("Challenge solution verification failed");
         Error::InvalidChallengeSolution
     })?;
 
-    props
-        .transport
-        .send(Ok(ClientResponse {
-            payload: Some(ClientResponsePayload::AuthOk(AuthOk {})),
-        }))
-        .await
-        .map_err(|e| {
-            error!(error = ?e, "Failed to send auth ok");
-            Error::Transport
-        })?;
-
     Ok(())
 }
 
-fn connect_error_code(err: &Error) -> ConnectErrorCode {
-    match err {
-        Error::ApproveError(ApproveError::Denied) => ConnectErrorCode::ApprovalDenied,
-        Error::ApproveError(ApproveError::Upstream(
-            router::ApprovalError::NoUserAgentsConnected,
-        )) => ConnectErrorCode::NoUserAgentsOnline,
-        _ => ConnectErrorCode::Unknown,
-    }
-}
-
-async fn authenticate(props: &mut ClientConnection) -> Result<(VerifyingKey, ClientId), Error> {
-    let Some(ClientRequest {
-        payload: Some(ClientRequestPayload::AuthChallengeRequest(challenge)),
-    }) = props.transport.recv().await
+pub async fn authenticate<T>(
+    props: &mut ClientConnection,
+    transport: &mut T,
+) -> Result<VerifyingKey, Error>
+where
+    T: Bi<Inbound, Result<Outbound, Error>> + Send + ?Sized,
+{
+    let Some(Inbound::AuthChallengeRequest { pubkey }) = transport.recv().await
     else {
         return Err(Error::Transport);
     };
 
-    let pubkey_bytes = challenge
-        .pubkey
-        .as_array()
-        .ok_or(Error::InvalidClientPubkeyLength)?;
-    let pubkey =
-        VerifyingKey::from_bytes(pubkey_bytes).map_err(|_| Error::InvalidAuthPubkeyEncoding)?;
-
-    let (client_id, nonce) = match get_nonce(&props.db, &pubkey).await? {
-        Some(state) => (state.client_id, state.nonce),
+    let nonce = match get_nonce(&props.db, &pubkey).await? {
+        Some(nonce) => nonce,
         None => {
             approve_new_client(&props.actors, pubkey).await?;
             match insert_client(&props.db, &pubkey).await? {
-                InsertClientResult::Inserted(client_id) => (client_id, 0),
+                InsertClientResult::Inserted => 0,
                 InsertClientResult::AlreadyExists => match get_nonce(&props.db, &pubkey).await? {
-                    Some(state) => (state.client_id, state.nonce),
-                    None => return Err(Error::InternalError),
+                    Some(nonce) => nonce,
+                    None => return Err(Error::DatabaseOperationFailed),
                 },
             }
         }
     };
 
-    challenge_client(props, pubkey, nonce).await?;
+    challenge_client(transport, pubkey, nonce).await?;
+    transport
+        .send(Ok(Outbound::AuthSuccess))
+        .await
+        .map_err(|e| {
+            error!(error = ?e, "Failed to send auth success");
+            Error::Transport
+        })?;
 
-    Ok((pubkey, client_id))
-}
-
-pub async fn authenticate_and_create(mut props: ClientConnection) -> Result<ClientSession, Error> {
-    match authenticate(&mut props).await {
-        Ok((_pubkey, client_id)) => Ok(ClientSession::new(props, client_id)),
-        Err(err) => {
-            let code = connect_error_code(&err);
-            let _ = props
-                .transport
-                .send(Ok(ClientResponse {
-                    payload: Some(ClientResponsePayload::ClientConnectError(
-                        ClientConnectError { code: code.into() },
-                    )),
-                }))
-                .await;
-            Err(err)
-        }
-    }
+    Ok(pubkey)
 }

server/crates/arbiter-server/src/actors/client/mod.rs

@@ -1,7 +1,4 @@
-use arbiter_proto::{
-    proto::client::{ClientRequest, ClientResponse},
-    transport::Bi,
-};
+use arbiter_proto::transport::Bi;
 use kameo::actor::Spawn;
 use tracing::{error, info};
 
@@ -10,48 +7,31 @@ use crate::{
     db,
 };
 
-#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
-pub enum ClientError {
-    #[error("Expected message with payload")]
-    MissingRequestPayload,
-    #[error("Unexpected request payload")]
-    UnexpectedRequestPayload,
-    #[error("State machine error")]
-    StateTransitionFailed,
-    #[error("Connection registration failed")]
-    ConnectionRegistrationFailed,
-    #[error(transparent)]
-    Auth(#[from] auth::Error),
-}
-
-pub type Transport = Box<dyn Bi<ClientRequest, Result<ClientResponse, ClientError>> + Send>;
-
 pub struct ClientConnection {
     pub(crate) db: db::DatabasePool,
-    pub(crate) transport: Transport,
     pub(crate) actors: GlobalActors,
 }
 
 impl ClientConnection {
-    pub fn new(db: db::DatabasePool, transport: Transport, actors: GlobalActors) -> Self {
-        Self {
-            db,
-            transport,
-            actors,
-        }
+    pub fn new(db: db::DatabasePool, actors: GlobalActors) -> Self {
+        Self { db, actors }
     }
 }
 
 pub mod auth;
 pub mod session;
 
-pub async fn connect_client(props: ClientConnection) {
-    match auth::authenticate_and_create(props).await {
-        Ok(session) => {
-            ClientSession::spawn(session);
+pub async fn connect_client<T>(mut props: ClientConnection, transport: &mut T)
+where
+    T: Bi<auth::Inbound, Result<auth::Outbound, auth::Error>> + Send + ?Sized,
+{
+    match auth::authenticate(&mut props, transport).await {
+        Ok(_pubkey) => {
+            ClientSession::spawn(ClientSession::new(props));
             info!("Client authenticated, session started");
         }
         Err(err) => {
+            let _ = transport.send(Err(err.clone())).await;
             error!(?err, "Authentication failed, closing connection");
         }
     }

server/crates/arbiter-server/src/actors/client/session.rs

@@ -1,92 +1,45 @@
-use alloy::{consensus::TxEip1559, primitives::Address, rlp::Decodable};
-use arbiter_proto::proto::{
-    client::{
-        ClientRequest, ClientResponse, client_request::Payload as ClientRequestPayload,
-        client_response::Payload as ClientResponsePayload,
-    },
-    evm::{
-        EvmError, EvmSignTransactionResponse, evm_sign_transaction_response::Result as SignResult,
-    },
-};
-use kameo::Actor;
-use tokio::select;
-use tracing::{error, info};
+use kameo::{Actor, messages};
+use tracing::error;
 
 use crate::{
     actors::{
-        GlobalActors,
-        client::{ClientConnection, ClientError, auth::ClientId},
-        evm::ClientSignTransaction,
-        router::RegisterClient,
+        GlobalActors, client::ClientConnection, keyholder::KeyHolderState, router::RegisterClient,
     },
     db,
 };
 
 pub struct ClientSession {
     props: ClientConnection,
-    client_id: ClientId,
 }
 
 impl ClientSession {
-    pub(crate) fn new(props: ClientConnection, client_id: ClientId) -> Self {
-        Self { props, client_id }
-    }
-
-    pub async fn process_transport_inbound(&mut self, req: ClientRequest) -> Output {
-        let msg = req.payload.ok_or_else(|| {
-            error!(actor = "client", "Received message with no payload");
-            ClientError::MissingRequestPayload
-        })?;
-
-        match msg {
-            ClientRequestPayload::EvmSignTransaction(sign_req) => {
-                let wallet_address: [u8; 20] = sign_req
-                    .wallet_address
-                    .try_into()
-                    .map_err(|_| ClientError::UnexpectedRequestPayload)?;
-
-                let mut rlp_bytes: &[u8] = &sign_req.rlp_transaction;
-                let tx = TxEip1559::decode(&mut rlp_bytes)
-                    .map_err(|_| ClientError::UnexpectedRequestPayload)?;
-
-                let result = self
-                    .props
-                    .actors
-                    .evm
-                    .ask(ClientSignTransaction {
-                        client_id: self.client_id.as_i32(),
-                        wallet_address: Address::from_slice(&wallet_address),
-                        transaction: tx,
-                    })
-                    .await;
-
-                let response_result = match result {
-                    Ok(signature) => SignResult::Signature(signature.as_bytes().to_vec()),
-                    Err(err) => {
-                        error!(?err, "client sign transaction failed");
-                        SignResult::Error(EvmError::Internal.into())
-                    }
-                };
-
-                Ok(ClientResponse {
-                    payload: Some(ClientResponsePayload::EvmSignTransaction(
-                        EvmSignTransactionResponse {
-                            result: Some(response_result),
-                        },
-                    )),
-                })
-            }
-            _ => Err(ClientError::UnexpectedRequestPayload),
-        }
+    pub(crate) fn new(props: ClientConnection) -> Self {
+        Self { props }
     }
 }
 
-type Output = Result<ClientResponse, ClientError>;
+#[messages]
+impl ClientSession {
+    #[message]
+    pub(crate) async fn handle_query_vault_state(&mut self) -> Result<KeyHolderState, Error> {
+        use crate::actors::keyholder::GetState;
+
+        let vault_state = match self.props.actors.key_holder.ask(GetState {}).await {
+            Ok(state) => state,
+            Err(err) => {
+                error!(?err, actor = "client", "keyholder.query.failed");
+                return Err(Error::Internal);
+            }
+        };
+
+        Ok(vault_state)
+    }
+}
 
 impl Actor for ClientSession {
     type Args = Self;
 
-    type Error = ClientError;
+    type Error = Error;
 
     async fn on_start(
         args: Self::Args,
@@ -97,55 +50,22 @@ impl Actor for ClientSession {
             .router
             .ask(RegisterClient { actor: this })
             .await
-            .map_err(|_| ClientError::ConnectionRegistrationFailed)?;
+            .map_err(|_| Error::ConnectionRegistrationFailed)?;
         Ok(args)
     }
-
-    async fn next(
-        &mut self,
-        _actor_ref: kameo::prelude::WeakActorRef<Self>,
-        mailbox_rx: &mut kameo::prelude::MailboxReceiver<Self>,
-    ) -> Option<kameo::mailbox::Signal<Self>> {
-        loop {
-            select! {
-                signal = mailbox_rx.recv() => {
-                    return signal;
-                }
-                msg = self.props.transport.recv() => {
-                    match msg {
-                        Some(request) => {
-                            match self.process_transport_inbound(request).await {
-                                Ok(resp) => {
-                                    if self.props.transport.send(Ok(resp)).await.is_err() {
-                                        error!(actor = "client", reason = "channel closed", "send.failed");
-                                        return Some(kameo::mailbox::Signal::Stop);
-                                    }
-                                }
-                                Err(err) => {
-                                    let _ = self.props.transport.send(Err(err)).await;
-                                    return Some(kameo::mailbox::Signal::Stop);
-                                }
-                            }
-                        }
-                        None => {
-                            info!(actor = "client", "transport.closed");
-                            return Some(kameo::mailbox::Signal::Stop);
-                        }
-                    }
-                }
-            }
-        }
-    }
 }
 
 impl ClientSession {
     pub fn new_test(db: db::DatabasePool, actors: GlobalActors) -> Self {
-        use arbiter_proto::transport::DummyTransport;
-        let transport: super::Transport = Box::new(DummyTransport::new());
-        let props = ClientConnection::new(db, transport, actors);
-        Self {
-            props,
-            client_id: ClientId::new(0),
-        }
+        let props = ClientConnection::new(db, actors);
+        Self { props }
     }
 }
+
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("Connection registration failed")]
+    ConnectionRegistrationFailed,
+    #[error("Internal error")]
+    Internal,
+}

server/crates/arbiter-server/src/actors/evm/mod.rs

@@ -1,21 +1,26 @@
 use alloy::{consensus::TxEip1559, primitives::Address, signers::Signature};
-use diesel::{ExpressionMethods, OptionalExtension as _, QueryDsl, SelectableHelper as _, dsl::insert_into};
+use diesel::{
+    ExpressionMethods, OptionalExtension as _, QueryDsl, SelectableHelper as _, dsl::insert_into,
+};
 use diesel_async::RunQueryDsl;
 use kameo::{Actor, actor::ActorRef, messages};
-use memsafe::MemSafe;
 use rand::{SeedableRng, rng, rngs::StdRng};
 
 use crate::{
     actors::keyholder::{CreateNew, Decrypt, KeyHolder},
-    db::{self, DatabasePool, models::{self, EvmBasicGrant, SqliteTimestamp}, schema},
+    db::{
+        self, DatabasePool,
+        models::{self, SqliteTimestamp},
+        schema,
+    },
     evm::{
-        self, RunKind,
+        self, ListGrantsError, RunKind,
         policies::{
-            FullGrant, SharedGrantSettings, SpecificGrant, SpecificMeaning,
-            ether_transfer::EtherTransfer,
-            token_transfers::TokenTransfer,
+            FullGrant, Grant, SharedGrantSettings, SpecificGrant, SpecificMeaning,
+            ether_transfer::EtherTransfer, token_transfers::TokenTransfer,
         },
     },
+    safe_cell::{SafeCell, SafeCellHandle as _},
 };
 
 pub use crate::evm::safe_signer;
@@ -88,7 +93,12 @@ impl EvmActor {
         // todo: audit
         let rng = StdRng::from_rng(&mut rng());
         let engine = evm::Engine::new(db.clone());
-        Self { keyholder, db, rng, engine }
+        Self {
+            keyholder,
+            db,
+            rng,
+            engine,
+        }
     }
 }
 
@@ -98,11 +108,7 @@ impl EvmActor {
     pub async fn generate(&mut self) -> Result<Address, Error> {
         let (mut key_cell, address) = safe_signer::generate(&mut self.rng);
 
-        // Move raw key bytes into a Vec<u8> MemSafe for KeyHolder
-        let plaintext = {
-            let reader = key_cell.read().expect("MemSafe read");
-            MemSafe::new(reader.to_vec()).expect("MemSafe allocation")
-        };
+        let plaintext = key_cell.read_inline(|reader| SafeCell::new(reader.to_vec()));
 
         let aead_id: i32 = self
             .keyholder
@@ -149,12 +155,24 @@ impl EvmActor {
         match grant {
             SpecificGrant::EtherTransfer(settings) => {
                 self.engine
-                    .create_grant::<EtherTransfer>(client_id, FullGrant { basic, specific: settings })
+                    .create_grant::<EtherTransfer>(
+                        client_id,
+                        FullGrant {
+                            basic,
+                            specific: settings,
+                        },
+                    )
                     .await
             }
             SpecificGrant::TokenTransfer(settings) => {
                 self.engine
-                    .create_grant::<TokenTransfer>(client_id, FullGrant { basic, specific: settings })
+                    .create_grant::<TokenTransfer>(
+                        client_id,
+                        FullGrant {
+                            basic,
+                            specific: settings,
+                        },
+                    )
                     .await
             }
         }
@@ -172,19 +190,12 @@ impl EvmActor {
     }
 
     #[message]
-    pub async fn useragent_list_grants(
-        &mut self,
-        wallet_id: Option<i32>,
-    ) -> Result<Vec<EvmBasicGrant>, Error> {
-        let mut conn = self.db.get().await?;
-        let mut query = schema::evm_basic_grant::table
-            .select(EvmBasicGrant::as_select())
-            .filter(schema::evm_basic_grant::revoked_at.is_null())
-            .into_boxed();
-        if let Some(wid) = wallet_id {
-            query = query.filter(schema::evm_basic_grant::wallet_id.eq(wid));
+    pub async fn useragent_list_grants(&mut self) -> Result<Vec<Grant<SpecificGrant>>, Error> {
+        match self.engine.list_all_grants().await {
+            Ok(grants) => Ok(grants),
+            Err(ListGrantsError::Database(db)) => Err(Error::Database(db)),
+            Err(ListGrantsError::Pool(pool)) => Err(Error::DatabasePool(pool)),
         }
-        Ok(query.load(&mut conn).await?)
     }
 
     #[message]
@@ -204,8 +215,14 @@ impl EvmActor {
             .ok_or(SignTransactionError::WalletNotFound)?;
         drop(conn);
 
-        let meaning = self.engine
-            .evaluate_transaction(wallet.id, client_id, transaction.clone(), RunKind::Execution)
+        let meaning = self
+            .engine
+            .evaluate_transaction(
+                wallet.id,
+                client_id,
+                transaction.clone(),
+                RunKind::Execution,
+            )
             .await?;
 
         Ok(meaning)
@@ -228,16 +245,23 @@ impl EvmActor {
             .ok_or(SignTransactionError::WalletNotFound)?;
         drop(conn);
 
-        let raw_key: MemSafe<Vec<u8>> = self
+        let raw_key: SafeCell<Vec<u8>> = self
             .keyholder
-            .ask(Decrypt { aead_id: wallet.aead_encrypted_id })
+            .ask(Decrypt {
+                aead_id: wallet.aead_encrypted_id,
+            })
             .await
             .map_err(|_| SignTransactionError::KeyholderSend)?;
 
-        let signer = safe_signer::SafeSigner::from_memsafe(raw_key)?;
+        let signer = safe_signer::SafeSigner::from_cell(raw_key)?;
 
         self.engine
-            .evaluate_transaction(wallet.id, client_id, transaction.clone(), RunKind::Execution)
+            .evaluate_transaction(
+                wallet.id,
+                client_id,
+                transaction.clone(),
+                RunKind::Execution,
+            )
             .await?;
 
         use alloy::network::TxSignerSync as _;

server/crates/arbiter-server/src/actors/keyholder/encryption/v1.rs

@@ -5,12 +5,13 @@ use chacha20poly1305::{
     AeadInPlace, Key, KeyInit as _, XChaCha20Poly1305, XNonce,
     aead::{AeadMut, Error, Payload},
 };
-use memsafe::MemSafe;
 use rand::{
     Rng as _, SeedableRng,
     rngs::{StdRng, SysRng},
 };
 
+use crate::safe_cell::{SafeCell, SafeCellHandle as _};
+
 pub const ROOT_KEY_TAG: &[u8] = "arbiter/seal/v1".as_bytes();
 pub const TAG: &[u8] = "arbiter/private-key/v1".as_bytes();
 
@@ -47,40 +48,37 @@ impl<'a> TryFrom<&'a [u8]> for Nonce {
     }
 }
 
-pub struct KeyCell(pub MemSafe<Key>);
-impl From<MemSafe<Key>> for KeyCell {
-    fn from(value: MemSafe<Key>) -> Self {
+pub struct KeyCell(pub SafeCell<Key>);
+impl From<SafeCell<Key>> for KeyCell {
+    fn from(value: SafeCell<Key>) -> Self {
         Self(value)
     }
 }
-impl TryFrom<MemSafe<Vec<u8>>> for KeyCell {
+impl TryFrom<SafeCell<Vec<u8>>> for KeyCell {
     type Error = ();
 
-    fn try_from(mut value: MemSafe<Vec<u8>>) -> Result<Self, Self::Error> {
-        let value = value.read().unwrap();
+    fn try_from(mut value: SafeCell<Vec<u8>>) -> Result<Self, Self::Error> {
+        let value = value.read();
         if value.len() != size_of::<Key>() {
             return Err(());
         }
-        let mut cell = MemSafe::new(Key::default()).unwrap();
-        {
-            let mut cell_write = cell.write().unwrap();
-            let cell_slice: &mut [u8] = cell_write.as_mut();
-            cell_slice.copy_from_slice(&value);
-        }
+        let cell = SafeCell::new_inline(|cell_write: &mut Key| {
+            cell_write.copy_from_slice(&value);
+        });
         Ok(Self(cell))
     }
 }
 
 impl KeyCell {
     pub fn new_secure_random() -> Self {
-        let mut key = MemSafe::new(Key::default()).unwrap();
-        {
-            let mut key_buffer = key.write().unwrap();
-            let key_buffer: &mut [u8] = key_buffer.as_mut();
-
+        let key = SafeCell::new_inline(|key_buffer: &mut Key| {
+            #[allow(
+                clippy::unwrap_used,
+                reason = "Rng failure is unrecoverable and should panic"
+            )]
             let mut rng = StdRng::try_from_rng(&mut SysRng).unwrap();
             rng.fill_bytes(key_buffer);
-        }
+        });
 
         key.into()
     }
@@ -91,7 +89,7 @@ impl KeyCell {
         associated_data: &[u8],
         mut buffer: impl AsMut<Vec<u8>>,
     ) -> Result<(), Error> {
-        let key_reader = self.0.read().unwrap();
+        let key_reader = self.0.read();
         let key_ref = key_reader.deref();
         let cipher = XChaCha20Poly1305::new(key_ref);
         let nonce = XNonce::from_slice(nonce.0.as_ref());
@@ -102,13 +100,13 @@ impl KeyCell {
         &mut self,
         nonce: &Nonce,
         associated_data: &[u8],
-        buffer: &mut MemSafe<Vec<u8>>,
+        buffer: &mut SafeCell<Vec<u8>>,
     ) -> Result<(), Error> {
-        let key_reader = self.0.read().unwrap();
+        let key_reader = self.0.read();
         let key_ref = key_reader.deref();
         let cipher = XChaCha20Poly1305::new(key_ref);
         let nonce = XNonce::from_slice(nonce.0.as_ref());
-        let mut buffer = buffer.write().unwrap();
+        let mut buffer = buffer.write();
         let buffer: &mut Vec<u8> = buffer.as_mut();
         cipher.decrypt_in_place(nonce, associated_data, buffer)
     }
@@ -119,7 +117,7 @@ impl KeyCell {
         associated_data: &[u8],
         plaintext: impl AsRef<[u8]>,
     ) -> Result<Vec<u8>, Error> {
-        let key_reader = self.0.read().unwrap();
+        let key_reader = self.0.read();
         let key_ref = key_reader.deref();
         let mut cipher = XChaCha20Poly1305::new(key_ref);
         let nonce = XNonce::from_slice(nonce.0.as_ref());
@@ -139,6 +137,10 @@ pub type Salt = [u8; ArgonSalt::RECOMMENDED_LENGTH];
 
 pub fn generate_salt() -> Salt {
     let mut salt = Salt::default();
+    #[allow(
+        clippy::unwrap_used,
+        reason = "Rng failure is unrecoverable and should panic"
+    )]
     let mut rng = StdRng::try_from_rng(&mut SysRng).unwrap();
     rng.fill_bytes(&mut salt);
     salt
@@ -146,19 +148,23 @@ pub fn generate_salt() -> Salt {
 
 /// User password might be of different length, have not enough entropy, etc...
 /// Derive a fixed-length key from the password using Argon2id, which is designed for password hashing and key derivation.
-pub fn derive_seal_key(mut password: MemSafe<Vec<u8>>, salt: &Salt) -> KeyCell {
+pub fn derive_seal_key(mut password: SafeCell<Vec<u8>>, salt: &Salt) -> KeyCell {
+    #[allow(clippy::unwrap_used)]
     let params = argon2::Params::new(262_144, 3, 4, None).unwrap();
     let hasher = Argon2::new(Algorithm::Argon2id, argon2::Version::V0x13, params);
-    let mut key = MemSafe::new(Key::default()).unwrap();
-    {
-        let password_source = password.read().unwrap();
-        let mut key_buffer = key.write().unwrap();
+    let mut key = SafeCell::new(Key::default());
+    password.read_inline(|password_source| {
+        let mut key_buffer = key.write();
         let key_buffer: &mut [u8] = key_buffer.as_mut();
 
+        #[allow(
+            clippy::unwrap_used,
+            reason = "Better fail completely than return a weak key"
+        )]
         hasher
             .hash_password_into(password_source.deref(), salt, key_buffer)
             .unwrap();
-    }
+    });
 
     key.into()
 }
@@ -166,20 +172,20 @@ pub fn derive_seal_key(mut password: MemSafe<Vec<u8>>, salt: &Salt) -> KeyCell {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use memsafe::MemSafe;
+    use crate::safe_cell::SafeCell;
 
     #[test]
     pub fn derive_seal_key_deterministic() {
         static PASSWORD: &[u8] = b"password";
-        let password = MemSafe::new(PASSWORD.to_vec()).unwrap();
-        let password2 = MemSafe::new(PASSWORD.to_vec()).unwrap();
+        let password = SafeCell::new(PASSWORD.to_vec());
+        let password2 = SafeCell::new(PASSWORD.to_vec());
         let salt = generate_salt();
 
         let mut key1 = derive_seal_key(password, &salt);
         let mut key2 = derive_seal_key(password2, &salt);
 
-        let key1_reader = key1.0.read().unwrap();
-        let key2_reader = key2.0.read().unwrap();
+        let key1_reader = key1.0.read();
+        let key2_reader = key2.0.read();
 
         assert_eq!(key1_reader.deref(), key2_reader.deref());
     }
@@ -187,11 +193,11 @@ mod tests {
     #[test]
     pub fn successful_derive() {
         static PASSWORD: &[u8] = b"password";
-        let password = MemSafe::new(PASSWORD.to_vec()).unwrap();
+        let password = SafeCell::new(PASSWORD.to_vec());
         let salt = generate_salt();
 
         let mut key = derive_seal_key(password, &salt);
-        let key_reader = key.0.read().unwrap();
+        let key_reader = key.0.read();
         let key_ref = key_reader.deref();
 
         assert_ne!(key_ref.as_slice(), &[0u8; 32][..]);
@@ -200,7 +206,7 @@ mod tests {
     #[test]
     pub fn encrypt_decrypt() {
         static PASSWORD: &[u8] = b"password";
-        let password = MemSafe::new(PASSWORD.to_vec()).unwrap();
+        let password = SafeCell::new(PASSWORD.to_vec());
         let salt = generate_salt();
 
         let mut key = derive_seal_key(password, &salt);
@@ -212,12 +218,12 @@ mod tests {
             .unwrap();
         assert_ne!(buffer, b"secret data");
 
-        let mut buffer = MemSafe::new(buffer).unwrap();
+        let mut buffer = SafeCell::new(buffer);
 
         key.decrypt_in_place(&nonce, associated_data, &mut buffer)
             .unwrap();
 
-        let buffer = buffer.read().unwrap();
+        let buffer = buffer.read();
         assert_eq!(*buffer, b"secret data");
     }
 

server/crates/arbiter-server/src/actors/keyholder/mod.rs

@@ -5,21 +5,24 @@ use diesel::{
 };
 use diesel_async::{AsyncConnection, RunQueryDsl};
 use kameo::{Actor, Reply, messages};
-use memsafe::MemSafe;
 use strum::{EnumDiscriminants, IntoDiscriminant};
 use tracing::{error, info};
 
-use crate::db::{
-    self,
-    models::{self, RootKeyHistory},
-    schema::{self},
+use crate::safe_cell::SafeCell;
+use crate::{
+    db::{
+        self,
+        models::{self, RootKeyHistory},
+        schema::{self},
+    },
+    safe_cell::SafeCellHandle as _,
 };
 use encryption::v1::{self, KeyCell, Nonce};
 
 pub mod encryption;
 
 #[derive(Default, EnumDiscriminants)]
-#[strum_discriminants(derive(Reply), vis(pub))]
+#[strum_discriminants(derive(Reply), vis(pub), name(KeyHolderState))]
 enum State {
     #[default]
     Unbootstrapped,
@@ -136,7 +139,7 @@ impl KeyHolder {
     }
 
     #[message]
-    pub async fn bootstrap(&mut self, seal_key_raw: MemSafe<Vec<u8>>) -> Result<(), Error> {
+    pub async fn bootstrap(&mut self, seal_key_raw: SafeCell<Vec<u8>>) -> Result<(), Error> {
         if !matches!(self.state, State::Unbootstrapped) {
             return Err(Error::AlreadyBootstrapped);
         }
@@ -148,16 +151,15 @@ impl KeyHolder {
         let root_key_nonce = v1::Nonce::default();
         let data_encryption_nonce = v1::Nonce::default();
 
-        let root_key_ciphertext: Vec<u8> = {
-            let root_key_reader = root_key.0.read().unwrap();
-            let root_key_reader = root_key_reader.as_slice();
+        let root_key_ciphertext: Vec<u8> = root_key.0.read_inline(|reader| {
+            let root_key_reader = reader.as_slice();
             seal_key
                 .encrypt(&root_key_nonce, v1::ROOT_KEY_TAG, root_key_reader)
                 .map_err(|err| {
                     error!(?err, "Fatal bootstrap error");
                     Error::Encryption(err)
-                })?
-        };
+                })
+        })?;
 
         let mut conn = self.db.get().await?;
 
@@ -199,7 +201,7 @@ impl KeyHolder {
     }
 
     #[message]
-    pub async fn try_unseal(&mut self, seal_key_raw: MemSafe<Vec<u8>>) -> Result<(), Error> {
+    pub async fn try_unseal(&mut self, seal_key_raw: SafeCell<Vec<u8>>) -> Result<(), Error> {
         let State::Sealed {
             root_key_history_id,
         } = &self.state
@@ -225,7 +227,7 @@ impl KeyHolder {
         })?;
         let mut seal_key = v1::derive_seal_key(seal_key_raw, &salt);
 
-        let mut root_key = MemSafe::new(current_key.ciphertext.clone()).unwrap();
+        let mut root_key = SafeCell::new(current_key.ciphertext.clone());
 
         let nonce = v1::Nonce::try_from(current_key.root_key_encryption_nonce.as_slice()).map_err(
             |_| {
@@ -256,7 +258,7 @@ impl KeyHolder {
 
     // Decrypts the `aead_encrypted` entry with the given ID and returns the plaintext
     #[message]
-    pub async fn decrypt(&mut self, aead_id: i32) -> Result<MemSafe<Vec<u8>>, Error> {
+    pub async fn decrypt(&mut self, aead_id: i32) -> Result<SafeCell<Vec<u8>>, Error> {
         let State::Unsealed { root_key, .. } = &mut self.state else {
             return Err(Error::NotBootstrapped);
         };
@@ -279,14 +281,14 @@ impl KeyHolder {
             );
             Error::BrokenDatabase
         })?;
-        let mut output = MemSafe::new(row.ciphertext).unwrap();
+        let mut output = SafeCell::new(row.ciphertext);
         root_key.decrypt_in_place(&nonce, v1::TAG, &mut output)?;
         Ok(output)
     }
 
     // Creates new `aead_encrypted` entry in the database and returns it's ID
     #[message]
-    pub async fn create_new(&mut self, mut plaintext: MemSafe<Vec<u8>>) -> Result<i32, Error> {
+    pub async fn create_new(&mut self, mut plaintext: SafeCell<Vec<u8>>) -> Result<i32, Error> {
         let State::Unsealed {
             root_key,
             root_key_history_id,
@@ -299,7 +301,7 @@ impl KeyHolder {
         // Borrow checker note: &mut borrow a few lines above is disjoint from this field
         let nonce = Self::get_new_nonce(&self.db, *root_key_history_id).await?;
 
-        let mut ciphertext_buffer = plaintext.write().unwrap();
+        let mut ciphertext_buffer = plaintext.write();
         let ciphertext_buffer: &mut Vec<u8> = ciphertext_buffer.as_mut();
         root_key.encrypt_in_place(&nonce, v1::TAG, &mut *ciphertext_buffer)?;
 
@@ -313,7 +315,7 @@ impl KeyHolder {
                 current_nonce: nonce.to_vec(),
                 schema_version: 1,
                 associated_root_key_id: *root_key_history_id,
-                created_at: Utc::now().into()
+                created_at: Utc::now().into(),
             })
             .returning(schema::aead_encrypted::id)
             .get_result(&mut conn)
@@ -323,7 +325,7 @@ impl KeyHolder {
     }
 
     #[message]
-    pub fn get_state(&self) -> StateDiscriminants {
+    pub fn get_state(&self) -> KeyHolderState {
         self.state.discriminant()
     }
 
@@ -346,17 +348,19 @@ impl KeyHolder {
 #[cfg(test)]
 mod tests {
     use diesel::SelectableHelper;
-    
-    use diesel_async::RunQueryDsl;
-    use memsafe::MemSafe;
 
-    use crate::db::{self};
+    use diesel_async::RunQueryDsl;
+
+    use crate::{
+        db::{self},
+        safe_cell::SafeCell,
+    };
 
     use super::*;
 
     async fn bootstrapped_actor(db: &db::DatabasePool) -> KeyHolder {
         let mut actor = KeyHolder::new(db.clone()).await.unwrap();
-        let seal_key = MemSafe::new(b"test-seal-key".to_vec()).unwrap();
+        let seal_key = SafeCell::new(b"test-seal-key".to_vec());
         actor.bootstrap(seal_key).await.unwrap();
         actor
     }
@@ -391,7 +395,7 @@ mod tests {
         assert_eq!(root_row.data_encryption_nonce, n2.to_vec());
 
         let id = actor
-            .create_new(MemSafe::new(b"post-interleave".to_vec()).unwrap())
+            .create_new(SafeCell::new(b"post-interleave".to_vec()))
             .await
             .unwrap();
         let row: models::AeadEncrypted = schema::aead_encrypted::table

server/crates/arbiter-server/src/actors/router/mod.rs

@@ -154,7 +154,7 @@ impl MessageRouter {
         ctx: &mut Context<Self, DelegatedReply<Result<bool, ApprovalError>>>,
     ) -> DelegatedReply<Result<bool, ApprovalError>> {
         let (reply, Some(reply_sender)) = ctx.reply_sender() else {
-            panic!("Expected `request_client_approval` to have callback channel");
+            unreachable!("Expected `request_client_approval` to have callback channel");
         };
 
         let weak_refs = self

server/crates/arbiter-server/src/actors/user_agent/auth.rs

@@ -1,114 +1,82 @@
-use arbiter_proto::proto::user_agent::{
-    AuthChallengeRequest, AuthChallengeSolution, KeyType as ProtoKeyType, UserAgentRequest,
-    user_agent_request::Payload as UserAgentRequestPayload,
-};
+use arbiter_proto::transport::Bi;
 use tracing::error;
 
 use crate::actors::user_agent::{
-    UserAgentConnection,
-    auth::state::{AuthContext, AuthPublicKey, AuthStateMachine},
-    session::UserAgentSession,
+    AuthPublicKey, UserAgentConnection,
+    auth::state::{AuthContext, AuthStateMachine},
 };
 
-#[derive(thiserror::Error, Debug, PartialEq)]
-pub enum Error {
-    #[error("Unexpected message payload")]
-    UnexpectedMessagePayload,
-    #[error("Invalid client public key length")]
-    InvalidClientPubkeyLength,
-    #[error("Invalid client public key encoding")]
-    InvalidAuthPubkeyEncoding,
-    #[error("Database pool unavailable")]
-    DatabasePoolUnavailable,
-    #[error("Database operation failed")]
-    DatabaseOperationFailed,
-    #[error("Public key not registered")]
-    PublicKeyNotRegistered,
-    #[error("Transport error")]
-    Transport,
-    #[error("Invalid bootstrap token")]
-    InvalidBootstrapToken,
-    #[error("Bootstrapper actor unreachable")]
-    BootstrapperActorUnreachable,
-    #[error("Invalid challenge solution")]
-    InvalidChallengeSolution,
-}
-
 mod state;
 use state::*;
 
-fn parse_pubkey(key_type: ProtoKeyType, pubkey: Vec<u8>) -> Result<AuthPublicKey, Error> {
-    match key_type {
-        // UNSPECIFIED treated as Ed25519 for backward compatibility
-        ProtoKeyType::Unspecified | ProtoKeyType::Ed25519 => {
-            let pubkey_bytes = pubkey.as_array().ok_or(Error::InvalidClientPubkeyLength)?;
-            let key = ed25519_dalek::VerifyingKey::from_bytes(pubkey_bytes)
-                .map_err(|_| Error::InvalidAuthPubkeyEncoding)?;
-            Ok(AuthPublicKey::Ed25519(key))
-        }
-        ProtoKeyType::EcdsaSecp256k1 => {
-            // Public key is sent as 33-byte SEC1 compressed point
-            let key = k256::ecdsa::VerifyingKey::from_sec1_bytes(&pubkey)
-                .map_err(|_| Error::InvalidAuthPubkeyEncoding)?;
-            Ok(AuthPublicKey::EcdsaSecp256k1(key))
-        }
-        ProtoKeyType::Rsa => {
-            use rsa::pkcs8::DecodePublicKey as _;
-            let key = rsa::RsaPublicKey::from_public_key_der(&pubkey)
-                .map_err(|_| Error::InvalidAuthPubkeyEncoding)?;
-            Ok(AuthPublicKey::Rsa(key))
+#[derive(Debug, Clone)]
+pub enum Inbound {
+    AuthChallengeRequest {
+        pubkey: AuthPublicKey,
+        bootstrap_token: Option<String>,
+    },
+    AuthChallengeSolution {
+        signature: Vec<u8>,
+    },
+}
+
+#[derive(Debug)]
+pub enum Error {
+    UnregisteredPublicKey,
+    InvalidChallengeSolution,
+    InvalidBootstrapToken,
+    Internal { details: String },
+    Transport,
+}
+
+impl Error {
+    fn internal(details: impl Into<String>) -> Self {
+        Self::Internal {
+            details: details.into(),
         }
     }
 }
 
-fn parse_auth_event(payload: UserAgentRequestPayload) -> Result<AuthEvents, Error> {
+#[derive(Debug, Clone)]
+pub enum Outbound {
+    AuthChallenge { nonce: i32 },
+    AuthSuccess,
+}
+
+fn parse_auth_event(payload: Inbound) -> AuthEvents {
     match payload {
-        UserAgentRequestPayload::AuthChallengeRequest(AuthChallengeRequest {
+        Inbound::AuthChallengeRequest {
             pubkey,
             bootstrap_token: None,
-            key_type,
-        }) => {
-            let kt = ProtoKeyType::try_from(key_type).unwrap_or(ProtoKeyType::Unspecified);
-            Ok(AuthEvents::AuthRequest(ChallengeRequest {
-                pubkey: parse_pubkey(kt, pubkey)?,
-            }))
-        }
-        UserAgentRequestPayload::AuthChallengeRequest(AuthChallengeRequest {
+        } => AuthEvents::AuthRequest(ChallengeRequest { pubkey }),
+        Inbound::AuthChallengeRequest {
             pubkey,
             bootstrap_token: Some(token),
-            key_type,
-        }) => {
-            let kt = ProtoKeyType::try_from(key_type).unwrap_or(ProtoKeyType::Unspecified);
-            Ok(AuthEvents::BootstrapAuthRequest(BootstrapAuthRequest {
-                pubkey: parse_pubkey(kt, pubkey)?,
-                token,
-            }))
-        }
-        UserAgentRequestPayload::AuthChallengeSolution(AuthChallengeSolution { signature }) => {
-            Ok(AuthEvents::ReceivedSolution(ChallengeSolution {
+        } => AuthEvents::BootstrapAuthRequest(BootstrapAuthRequest { pubkey, token }),
+        Inbound::AuthChallengeSolution { signature } => {
+            AuthEvents::ReceivedSolution(ChallengeSolution {
                 solution: signature,
-            }))
+            })
         }
-        _ => Err(Error::UnexpectedMessagePayload),
     }
 }
 
-pub async fn authenticate(props: &mut UserAgentConnection) -> Result<AuthPublicKey, Error> {
-    let mut state = AuthStateMachine::new(AuthContext::new(props));
+pub async fn authenticate<T>(
+    props: &mut UserAgentConnection,
+    transport: T,
+) -> Result<AuthPublicKey, Error>
+where
+    T: Bi<Inbound, Result<Outbound, Error>> + Send,
+{
+    let mut state = AuthStateMachine::new(AuthContext::new(props, transport));
 
     loop {
         // `state` holds a mutable reference to `props` so we can't access it directly here
-        let transport = state.context_mut().conn.transport.as_mut();
-        let Some(UserAgentRequest {
-            payload: Some(payload),
-        }) = transport.recv().await
-        else {
+        let Some(payload) = state.context_mut().transport.recv().await else {
             return Err(Error::Transport);
         };
 
-        let event = parse_auth_event(payload)?;
-
-        match state.process_event(event).await {
+        match state.process_event(parse_auth_event(payload)).await {
             Ok(AuthStates::AuthOk(key)) => return Ok(key.clone()),
             Err(AuthError::ActionFailed(err)) => {
                 error!(?err, "State machine action failed");
@@ -131,11 +99,3 @@ pub async fn authenticate(props: &mut UserAgentConnection) -> Result<AuthPublicK
         }
     }
 }
-
-pub async fn authenticate_and_create(
-    mut props: UserAgentConnection,
-) -> Result<UserAgentSession, Error> {
-    let _key = authenticate(&mut props).await?;
-    let session = UserAgentSession::new(props);
-    Ok(session)
-}

server/crates/arbiter-server/src/actors/user_agent/auth/state.rs

@@ -1,52 +1,17 @@
-use arbiter_proto::proto::user_agent::{
-    AuthChallenge, UserAgentResponse, user_agent_response::Payload as UserAgentResponsePayload,
-};
+use arbiter_proto::transport::Bi;
 use diesel::{ExpressionMethods as _, OptionalExtension as _, QueryDsl, update};
 use diesel_async::RunQueryDsl;
 use tracing::error;
 
 use super::Error;
 use crate::{
-    actors::{bootstrap::ConsumeToken, user_agent::UserAgentConnection},
-    db::{models::KeyType, schema},
+    actors::{
+        bootstrap::ConsumeToken,
+        user_agent::{AuthPublicKey, UserAgentConnection, auth::Outbound},
+    },
+    db::schema,
 };
 
-/// Abstraction over Ed25519 / ECDSA-secp256k1 / RSA public keys used during the auth handshake.
-#[derive(Clone)]
-pub enum AuthPublicKey {
-    Ed25519(ed25519_dalek::VerifyingKey),
-    /// Compressed SEC1 public key; signature bytes are raw 64-byte (r||s).
-    EcdsaSecp256k1(k256::ecdsa::VerifyingKey),
-    /// RSA-2048+ public key (Windows Hello / KeyCredentialManager); signature bytes are PSS+SHA-256.
-    Rsa(rsa::RsaPublicKey),
-}
-
-impl AuthPublicKey {
-    /// Canonical bytes stored in DB and echoed back in the challenge.
-    /// Ed25519: raw 32 bytes. ECDSA: SEC1 compressed 33 bytes. RSA: DER-encoded SPKI.
-    pub fn to_stored_bytes(&self) -> Vec<u8> {
-        match self {
-            AuthPublicKey::Ed25519(k) => k.to_bytes().to_vec(),
-            // SEC1 compressed (33 bytes) is the natural compact format for secp256k1
-            AuthPublicKey::EcdsaSecp256k1(k) => k.to_encoded_point(true).as_bytes().to_vec(),
-            AuthPublicKey::Rsa(k) => {
-                use rsa::pkcs8::EncodePublicKey as _;
-                k.to_public_key_der()
-                    .expect("rsa SPKI encoding is infallible")
-                    .to_vec()
-            }
-        }
-    }
-
-    pub fn key_type(&self) -> KeyType {
-        match self {
-            AuthPublicKey::Ed25519(_) => KeyType::Ed25519,
-            AuthPublicKey::EcdsaSecp256k1(_) => KeyType::EcdsaSecp256k1,
-            AuthPublicKey::Rsa(_) => KeyType::Rsa,
-        }
-    }
-}
-
 pub struct ChallengeRequest {
     pub pubkey: AuthPublicKey,
 }
@@ -57,7 +22,7 @@ pub struct BootstrapAuthRequest {
 }
 
 pub struct ChallengeContext {
-    pub challenge: AuthChallenge,
+    pub challenge_nonce: i32,
     pub key: AuthPublicKey,
 }
 
@@ -70,15 +35,15 @@ smlang::statemachine!(
     custom_error: true,
     transitions: {
         *Init + AuthRequest(ChallengeRequest) / async prepare_challenge = SentChallenge(ChallengeContext),
-        Init + BootstrapAuthRequest(BootstrapAuthRequest) [async verify_bootstrap_token] / provide_key_bootstrap = AuthOk(AuthPublicKey),
-        SentChallenge(ChallengeContext) + ReceivedSolution(ChallengeSolution) [async verify_solution] / provide_key = AuthOk(AuthPublicKey),
+        Init + BootstrapAuthRequest(BootstrapAuthRequest) / async verify_bootstrap_token = AuthOk(AuthPublicKey),
+        SentChallenge(ChallengeContext) + ReceivedSolution(ChallengeSolution) / async verify_solution = AuthOk(AuthPublicKey),
     }
 );
 
 async fn create_nonce(db: &crate::db::DatabasePool, pubkey_bytes: &[u8]) -> Result<i32, Error> {
     let mut db_conn = db.get().await.map_err(|e| {
         error!(error = ?e, "Database pool error");
-        Error::DatabasePoolUnavailable
+        Error::internal("Database unavailable")
     })?;
     db_conn
         .exclusive_transaction(|conn| {
@@ -102,11 +67,11 @@ async fn create_nonce(db: &crate::db::DatabasePool, pubkey_bytes: &[u8]) -> Resu
         .optional()
         .map_err(|e| {
             error!(error = ?e, "Database error");
-            Error::DatabaseOperationFailed
+            Error::internal("Database operation failed")
         })?
         .ok_or_else(|| {
             error!(?pubkey_bytes, "Public key not found in database");
-            Error::PublicKeyNotRegistered
+            Error::UnregisteredPublicKey
         })
 }
 
@@ -115,7 +80,7 @@ async fn register_key(db: &crate::db::DatabasePool, pubkey: &AuthPublicKey) -> R
     let key_type = pubkey.key_type();
     let mut conn = db.get().await.map_err(|e| {
         error!(error = ?e, "Database pool error");
-        Error::DatabasePoolUnavailable
+        Error::internal("Database unavailable")
     })?;
 
     diesel::insert_into(schema::useragent_client::table)
@@ -128,31 +93,95 @@ async fn register_key(db: &crate::db::DatabasePool, pubkey: &AuthPublicKey) -> R
         .await
         .map_err(|e| {
             error!(error = ?e, "Database error");
-            Error::DatabaseOperationFailed
+            Error::internal("Database operation failed")
         })?;
 
     Ok(())
 }
 
-pub struct AuthContext<'a> {
+pub struct AuthContext<'a, T> {
     pub(super) conn: &'a mut UserAgentConnection,
+    pub(super) transport: T,
 }
 
-impl<'a> AuthContext<'a> {
-    pub fn new(conn: &'a mut UserAgentConnection) -> Self {
-        Self { conn }
+impl<'a, T> AuthContext<'a, T> {
+    pub fn new(conn: &'a mut UserAgentConnection, transport: T) -> Self {
+        Self { conn, transport }
     }
 }
 
-impl AuthStateMachineContext for AuthContext<'_> {
+impl<T> AuthStateMachineContext for AuthContext<'_, T>
+where
+    T: Bi<super::Inbound, Result<super::Outbound, Error>> + Send,
+{
     type Error = Error;
 
+    async fn prepare_challenge(
+        &mut self,
+        ChallengeRequest { pubkey }: ChallengeRequest,
+    ) -> Result<ChallengeContext, Self::Error> {
+        let stored_bytes = pubkey.to_stored_bytes();
+        let nonce = create_nonce(&self.conn.db, &stored_bytes).await?;
+
+        self.transport
+            .send(Ok(Outbound::AuthChallenge { nonce }))
+            .await
+            .map_err(|e| {
+                error!(?e, "Failed to send auth challenge");
+                Error::Transport
+            })?;
+
+        Ok(ChallengeContext {
+            challenge_nonce: nonce,
+            key: pubkey,
+        })
+    }
+
+    #[allow(missing_docs)]
+    #[allow(clippy::result_unit_err)]
+    async fn verify_bootstrap_token(
+        &mut self,
+        BootstrapAuthRequest { pubkey, token }: BootstrapAuthRequest,
+    ) -> Result<AuthPublicKey, Self::Error> {
+        let token_ok: bool = self
+            .conn
+            .actors
+            .bootstrapper
+            .ask(ConsumeToken {
+                token: token.clone(),
+            })
+            .await
+            .map_err(|e| {
+                error!(?e, "Failed to consume bootstrap token");
+                Error::internal("Failed to consume bootstrap token")
+            })?;
+
+        if !token_ok {
+            error!("Invalid bootstrap token provided");
+            return Err(Error::InvalidBootstrapToken);
+        }
+
+        register_key(&self.conn.db, &pubkey).await?;
+
+        self.transport
+            .send(Ok(Outbound::AuthSuccess))
+            .await
+            .map_err(|_| Error::Transport)?;
+
+        Ok(pubkey)
+    }
+
+    #[allow(missing_docs)]
+    #[allow(clippy::unused_unit)]
     async fn verify_solution(
-        &self,
-        ChallengeContext { challenge, key }: &ChallengeContext,
-        ChallengeSolution { solution }: &ChallengeSolution,
-    ) -> Result<bool, Self::Error> {
-        let formatted = arbiter_proto::format_challenge(challenge.nonce, &challenge.pubkey);
+        &mut self,
+        ChallengeContext {
+            challenge_nonce,
+            key,
+        }: &ChallengeContext,
+        ChallengeSolution { solution }: ChallengeSolution,
+    ) -> Result<AuthPublicKey, Self::Error> {
+        let formatted = arbiter_proto::format_challenge(*challenge_nonce, &key.to_stored_bytes());
 
         let valid = match key {
             AuthPublicKey::Ed25519(vk) => {
@@ -181,117 +210,13 @@ impl AuthStateMachineContext for AuthContext<'_> {
             }
         };
 
-        Ok(valid)
-    }
-
-    async fn prepare_challenge(
-        &mut self,
-        ChallengeRequest { pubkey }: ChallengeRequest,
-    ) -> Result<ChallengeContext, Self::Error> {
-        let stored_bytes = pubkey.to_stored_bytes();
-        let nonce = create_nonce(&self.conn.db, &stored_bytes).await?;
-
-        let challenge = AuthChallenge {
-            pubkey: stored_bytes,
-            nonce,
-        };
-
-        self.conn
-            .transport
-            .send(Ok(UserAgentResponse {
-                payload: Some(UserAgentResponsePayload::AuthChallenge(challenge.clone())),
-            }))
-            .await
-            .map_err(|e| {
-                error!(?e, "Failed to send auth challenge");
-                Error::Transport
-            })?;
-
-        Ok(ChallengeContext {
-            challenge,
-            key: pubkey,
-        })
-    }
-
-    #[allow(missing_docs)]
-    #[allow(clippy::result_unit_err)]
-    async fn verify_bootstrap_token(
-        &self,
-        BootstrapAuthRequest { pubkey, token }: &BootstrapAuthRequest,
-    ) -> Result<bool, Self::Error> {
-        let token_ok: bool = self
-            .conn
-            .actors
-            .bootstrapper
-            .ask(ConsumeToken {
-                token: token.clone(),
-            })
-            .await
-            .map_err(|e| {
-                error!(?e, "Failed to consume bootstrap token");
-                Error::BootstrapperActorUnreachable
-            })?;
-
-        if !token_ok {
-            error!("Invalid bootstrap token provided");
-            return Err(Error::InvalidBootstrapToken);
+        if valid {
+            self.transport
+                .send(Ok(Outbound::AuthSuccess))
+                .await
+                .map_err(|_| Error::Transport)?;
         }
 
-        register_key(&self.conn.db, pubkey).await?;
-
-        Ok(true)
-    }
-
-    fn provide_key_bootstrap(
-        &mut self,
-        event_data: BootstrapAuthRequest,
-    ) -> Result<AuthPublicKey, Self::Error> {
-        Ok(event_data.pubkey)
-    }
-
-    fn provide_key(
-        &mut self,
-        state_data: &ChallengeContext,
-        _: ChallengeSolution,
-    ) -> Result<AuthPublicKey, Self::Error> {
-        // ChallengeContext.key cannot be taken by value because smlang passes it by ref;
-        // we reconstruct stored bytes and return them wrapped in Ed25519 placeholder.
-        // Session uses only the raw bytes, so we carry them via a Vec<u8>.
-        // IMPORTANT: do NOT simplify this by storing the key type separately — the
-        // `AuthPublicKey` enum IS the source of truth for key bytes and type.
-        //
-        // smlang state-machine trait requires returning an owned value from `provide_key`,
-        // but `state_data` is only available by shared reference here.  We extract the
-        // stored bytes and re-wrap as the correct variant so the caller can call
-        // `to_stored_bytes()` / `key_type()` without losing information.
-        let bytes = state_data.challenge.pubkey.clone();
-        let key_type = state_data.key.key_type();
-        let rebuilt = match key_type {
-            crate::db::models::KeyType::Ed25519 => {
-                let arr: &[u8; 32] = bytes
-                    .as_slice()
-                    .try_into()
-                    .expect("ed25519 pubkey must be 32 bytes in challenge");
-                AuthPublicKey::Ed25519(
-                    ed25519_dalek::VerifyingKey::from_bytes(arr)
-                        .expect("key was already validated in parse_auth_event"),
-                )
-            }
-            crate::db::models::KeyType::EcdsaSecp256k1 => {
-                // bytes are SEC1 compressed (33 bytes produced by to_encoded_point(true))
-                AuthPublicKey::EcdsaSecp256k1(
-                    k256::ecdsa::VerifyingKey::from_sec1_bytes(&bytes)
-                        .expect("ecdsa key was already validated in parse_auth_event"),
-                )
-            }
-            crate::db::models::KeyType::Rsa => {
-                use rsa::pkcs8::DecodePublicKey as _;
-                AuthPublicKey::Rsa(
-                    rsa::RsaPublicKey::from_public_key_der(&bytes)
-                        .expect("rsa key was already validated in parse_auth_event"),
-                )
-            }
-        };
-        Ok(rebuilt)
+        Ok(key.clone())
     }
 }

server/crates/arbiter-server/src/actors/user_agent/mod.rs

@@ -1,83 +1,94 @@
-use arbiter_proto::{
-    proto::user_agent::{
-        SdkClientError as ProtoSdkClientError, UserAgentRequest, UserAgentResponse,
-    },
-    transport::Bi,
-};
-use fatality::Fatality;
-use kameo::actor::Spawn as _;
-use tracing::{error, info};
-
 use crate::{
-    actors::{GlobalActors, user_agent::session::UserAgentSession},
-    db::{self},
+    actors::GlobalActors,
+    db::{self, models::KeyType},
 };
 
-#[derive(Debug, thiserror::Error, PartialEq)]
-pub enum TransportResponseError {
-    #[error("Expected message with payload")]
-    MissingRequestPayload,
-    #[error("Unexpected request payload")]
-    UnexpectedRequestPayload,
-    #[error("Invalid state for unseal encrypted key")]
-    InvalidStateForUnsealEncryptedKey,
-    #[error("client_pubkey must be 32 bytes")]
-    InvalidClientPubkeyLength,
-    #[error("State machine error")]
-    StateTransitionFailed,
-    #[error("Vault is not available")]
-    KeyHolderActorUnreachable,
-    #[error("SDK client approve failed: {0:?}")]
-    SdkClientApprove(ProtoSdkClientError),
-    #[error("SDK client list failed: {0:?}")]
-    SdkClientList(ProtoSdkClientError),
-    #[error("SDK client revoke failed: {0:?}")]
-    SdkClientRevoke(ProtoSdkClientError),
-    #[error(transparent)]
-    Auth(#[from] auth::Error),
-    #[error("Failed registering connection")]
-    ConnectionRegistrationFailed,
+/// Abstraction over Ed25519 / ECDSA-secp256k1 / RSA public keys used during the auth handshake.
+#[derive(Clone, Debug)]
+pub enum AuthPublicKey {
+    Ed25519(ed25519_dalek::VerifyingKey),
+    /// Compressed SEC1 public key; signature bytes are raw 64-byte (r||s).
+    EcdsaSecp256k1(k256::ecdsa::VerifyingKey),
+    /// RSA-2048+ public key (Windows Hello / KeyCredentialManager); signature bytes are PSS+SHA-256.
+    Rsa(rsa::RsaPublicKey),
 }
 
-impl Fatality for TransportResponseError {
-    fn is_fatal(&self) -> bool {
-        !matches!(
-            self,
-            Self::SdkClientApprove(_) | Self::SdkClientList(_) | Self::SdkClientRevoke(_)
-        )
+impl AuthPublicKey {
+    /// Canonical bytes stored in DB and echoed back in the challenge.
+    /// Ed25519: raw 32 bytes. ECDSA: SEC1 compressed 33 bytes. RSA: DER-encoded SPKI.
+    pub fn to_stored_bytes(&self) -> Vec<u8> {
+        match self {
+            AuthPublicKey::Ed25519(k) => k.to_bytes().to_vec(),
+            // SEC1 compressed (33 bytes) is the natural compact format for secp256k1
+            AuthPublicKey::EcdsaSecp256k1(k) => k.to_encoded_point(true).as_bytes().to_vec(),
+            AuthPublicKey::Rsa(k) => {
+                use rsa::pkcs8::EncodePublicKey as _;
+                #[allow(clippy::expect_used)]
+                k.to_public_key_der()
+                    .expect("rsa SPKI encoding is infallible")
+                    .to_vec()
+            }
+        }
+    }
+
+    pub fn key_type(&self) -> KeyType {
+        match self {
+            AuthPublicKey::Ed25519(_) => KeyType::Ed25519,
+            AuthPublicKey::EcdsaSecp256k1(_) => KeyType::EcdsaSecp256k1,
+            AuthPublicKey::Rsa(_) => KeyType::Rsa,
+        }
     }
 }
 
-pub type Transport =
-    Box<dyn Bi<UserAgentRequest, Result<UserAgentResponse, TransportResponseError>> + Send>;
+impl TryFrom<(KeyType, Vec<u8>)> for AuthPublicKey {
+    type Error = &'static str;
+
+    fn try_from(value: (KeyType, Vec<u8>)) -> Result<Self, Self::Error> {
+        let (key_type, bytes) = value;
+        match key_type {
+            KeyType::Ed25519 => {
+                let bytes: [u8; 32] = bytes.try_into().map_err(|_| "invalid Ed25519 key length")?;
+                let key = ed25519_dalek::VerifyingKey::from_bytes(&bytes)
+                    .map_err(|_e| "invalid Ed25519 key")?;
+                Ok(AuthPublicKey::Ed25519(key))
+            }
+            KeyType::EcdsaSecp256k1 => {
+                let point =
+                    k256::EncodedPoint::from_bytes(&bytes).map_err(|_e| "invalid ECDSA key")?;
+                let key = k256::ecdsa::VerifyingKey::from_encoded_point(&point)
+                    .map_err(|_e| "invalid ECDSA key")?;
+                Ok(AuthPublicKey::EcdsaSecp256k1(key))
+            }
+            KeyType::Rsa => {
+                use rsa::pkcs8::DecodePublicKey as _;
+                let key = rsa::RsaPublicKey::from_public_key_der(&bytes)
+                    .map_err(|_e| "invalid RSA key")?;
+                Ok(AuthPublicKey::Rsa(key))
+            }
+        }
+    }
+}
+
+// Messages, sent by user agent to connection client without having a request
+#[derive(Debug)]
+pub enum OutOfBand {
+    ClientConnectionRequest { pubkey: ed25519_dalek::VerifyingKey },
+    ClientConnectionCancel,
+}
 
 pub struct UserAgentConnection {
-    db: db::DatabasePool,
-    actors: GlobalActors,
-    transport: Transport,
+    pub(crate) db: db::DatabasePool,
+    pub(crate) actors: GlobalActors,
 }
 
 impl UserAgentConnection {
-    pub fn new(db: db::DatabasePool, actors: GlobalActors, transport: Transport) -> Self {
-        Self {
-            db,
-            actors,
-            transport,
-        }
+    pub fn new(db: db::DatabasePool, actors: GlobalActors) -> Self {
+        Self { db, actors }
     }
 }
 
 pub mod auth;
 pub mod session;
 
-pub async fn connect_user_agent(props: UserAgentConnection) {
-    match auth::authenticate_and_create(props).await {
-        Ok(session) => {
-            UserAgentSession::spawn(session);
-            info!("User authenticated, session started");
-        }
-        Err(err) => {
-            error!(?err, "Authentication failed, closing connection");
-        }
-    }
-}
+pub use auth::authenticate;
+pub use session::UserAgentSession;

server/crates/arbiter-server/src/actors/user_agent/session.rs

@@ -1,525 +1,116 @@
-use std::{ops::DerefMut, sync::Mutex};
+use std::borrow::Cow;
 
-use arbiter_proto::proto::{
-    evm as evm_proto,
-    user_agent::{
-        SdkClientApproveRequest, SdkClientApproveResponse, SdkClientConnectionCancel,
-        SdkClientConnectionRequest, SdkClientEntry, SdkClientError as ProtoSdkClientError,
-        SdkClientList, SdkClientListResponse, SdkClientRevokeRequest, SdkClientRevokeResponse,
-        UnsealEncryptedKey, UnsealResult, UnsealStart, UnsealStartResponse, UserAgentRequest,
-        UserAgentResponse, sdk_client_approve_response, sdk_client_list_response,
-        sdk_client_revoke_response, user_agent_request::Payload as UserAgentRequestPayload,
-        user_agent_response::Payload as UserAgentResponsePayload,
-    },
-};
-use chacha20poly1305::{AeadInPlace, XChaCha20Poly1305, XNonce, aead::KeyInit};
-use diesel::{ExpressionMethods as _, QueryDsl as _, dsl::insert_into};
-use diesel_async::RunQueryDsl as _;
+use arbiter_proto::transport::Sender;
+use async_trait::async_trait;
 use ed25519_dalek::VerifyingKey;
-use fatality::Fatality;
-use kameo::{Actor, error::SendError, messages, prelude::Context};
-use memsafe::MemSafe;
-use tokio::{select, sync::watch};
-use tracing::{error, info};
-use x25519_dalek::{EphemeralSecret, PublicKey};
+use kameo::{Actor, messages};
+use thiserror::Error;
+use tokio::sync::watch;
+use tracing::error;
 
-use crate::{
-    actors::{
-        evm::{Generate, ListWallets},
-        keyholder::{self, TryUnseal},
-        router::RegisterUserAgent,
-        user_agent::{TransportResponseError, UserAgentConnection},
-    },
-    db::schema::program_client,
+use crate::actors::{
+    router::RegisterUserAgent,
+    user_agent::{OutOfBand, UserAgentConnection},
 };
 
 mod state;
-use state::{DummyContext, UnsealContext, UserAgentEvents, UserAgentStateMachine, UserAgentStates};
+use state::{DummyContext, UserAgentEvents, UserAgentStateMachine};
 
-// Error for consumption by other actors
-#[derive(Debug, thiserror::Error, PartialEq)]
+#[derive(Debug, Error)]
 pub enum Error {
-    #[error("User agent session ended due to connection loss")]
-    ConnectionLost,
+    #[error("State transition failed")]
+    State,
 
-    #[error("User agent session ended due to unexpected message")]
-    UnexpectedMessage,
+    #[error("Internal error: {message}")]
+    Internal { message: Cow<'static, str> },
+}
+
+impl Error {
+    pub fn internal(message: impl Into<Cow<'static, str>>) -> Self {
+        Self::Internal {
+            message: message.into(),
+        }
+    }
 }
 
 pub struct UserAgentSession {
     props: UserAgentConnection,
     state: UserAgentStateMachine<DummyContext>,
+    #[allow(dead_code, reason = "The session keeps ownership of the outbound transport even before the state-machine flow starts using it directly")]
+    sender: Box<dyn Sender<OutOfBand>>,
 }
 
+mod connection;
+pub(crate) use connection::{
+    BootstrapError, HandleBootstrapEncryptedKey, HandleEvmWalletCreate, HandleEvmWalletList,
+    HandleGrantCreate, HandleGrantDelete, HandleGrantList, HandleQueryVaultState,
+};
+pub use connection::{HandleUnsealEncryptedKey, HandleUnsealRequest, UnsealError};
+
 impl UserAgentSession {
-    pub(crate) fn new(props: UserAgentConnection) -> Self {
+    pub(crate) fn new(props: UserAgentConnection, sender: Box<dyn Sender<OutOfBand>>) -> Self {
         Self {
             props,
             state: UserAgentStateMachine::new(DummyContext),
+            sender,
         }
     }
 
-    fn transition(&mut self, event: UserAgentEvents) -> Result<(), TransportResponseError> {
+    pub fn new_test(db: crate::db::DatabasePool, actors: crate::actors::GlobalActors) -> Self {
+        struct DummySender;
+
+        #[async_trait]
+        impl Sender<OutOfBand> for DummySender {
+            async fn send(
+                &mut self,
+                _item: OutOfBand,
+            ) -> Result<(), arbiter_proto::transport::Error> {
+                Ok(())
+            }
+        }
+
+        Self::new(UserAgentConnection::new(db, actors), Box::new(DummySender))
+    }
+
+    fn transition(&mut self, event: UserAgentEvents) -> Result<(), Error> {
         self.state.process_event(event).map_err(|e| {
             error!(?e, "State transition failed");
-            TransportResponseError::StateTransitionFailed
+            Error::State
         })?;
         Ok(())
     }
-
-    async fn send_msg<Reply: kameo::Reply>(
-        &mut self,
-        msg: UserAgentResponsePayload,
-        _ctx: &mut Context<Self, Reply>,
-    ) -> Result<(), Error> {
-        self.props
-            .transport
-            .send(Ok(response(msg)))
-            .await
-            .map_err(|_| {
-                error!(
-                    actor = "useragent",
-                    reason = "channel closed",
-                    "send.failed"
-                );
-                Error::ConnectionLost
-            })
-    }
-
-    async fn expect_msg<Extractor, Msg, Reply>(
-        &mut self,
-        extractor: Extractor,
-        ctx: &mut Context<Self, Reply>,
-    ) -> Result<Msg, Error>
-    where
-        Extractor: FnOnce(UserAgentRequestPayload) -> Option<Msg>,
-        Reply: kameo::Reply,
-    {
-        let msg = self.props.transport.recv().await.ok_or_else(|| {
-            error!(
-                actor = "useragent",
-                reason = "channel closed",
-                "recv.failed"
-            );
-            ctx.stop();
-            Error::ConnectionLost
-        })?;
-
-        msg.payload.and_then(extractor).ok_or_else(|| {
-            error!(
-                actor = "useragent",
-                reason = "unexpected message",
-                "recv.failed"
-            );
-            ctx.stop();
-            Error::UnexpectedMessage
-        })
-    }
 }
 
 #[messages]
 impl UserAgentSession {
-    // TODO: Think about refactoring it to state-machine based flow, as we already have one
-    #[message(ctx)]
+    #[message]
     pub async fn request_new_client_approval(
         &mut self,
         client_pubkey: VerifyingKey,
         mut cancel_flag: watch::Receiver<()>,
-        ctx: &mut Context<Self, Result<bool, Error>>,
-    ) -> Result<bool, Error> {
-        self.send_msg(
-            UserAgentResponsePayload::SdkClientConnectionRequest(SdkClientConnectionRequest {
-                pubkey: client_pubkey.as_bytes().to_vec(),
-            }),
-            ctx,
-        )
-        .await?;
-
-        let extractor = |msg| {
-            if let UserAgentRequestPayload::SdkClientConnectionResponse(
-                client_connection_response,
-            ) = msg
-            {
-                Some(client_connection_response)
-            } else {
-                None
-            }
-        };
-
-        tokio::select! {
-            _ = cancel_flag.changed() => {
-                info!(actor = "useragent", "client connection approval cancelled");
-                self.send_msg(
-                    UserAgentResponsePayload::SdkClientConnectionCancel(SdkClientConnectionCancel {}),
-                    ctx,
-                ).await?;
-                Ok(false)
-            }
-            result = self.expect_msg(extractor, ctx) => {
-                let result = result?;
-                info!(actor = "useragent", "received client connection approval result: approved={}", result.approved);
-                Ok(result.approved)
-            }
-        }
-    }
-}
-
-impl UserAgentSession {
-    pub async fn process_transport_inbound(&mut self, req: UserAgentRequest) -> Output {
-        let msg = req.payload.ok_or_else(|| {
-            error!(actor = "useragent", "Received message with no payload");
-            TransportResponseError::MissingRequestPayload
-        })?;
-
-        match msg {
-            UserAgentRequestPayload::UnsealStart(unseal_start) => {
-                self.handle_unseal_request(unseal_start).await
-            }
-            UserAgentRequestPayload::UnsealEncryptedKey(unseal_encrypted_key) => {
-                self.handle_unseal_encrypted_key(unseal_encrypted_key).await
-            }
-            UserAgentRequestPayload::EvmWalletCreate(_) => self.handle_evm_wallet_create().await,
-            UserAgentRequestPayload::EvmWalletList(_) => self.handle_evm_wallet_list().await,
-            UserAgentRequestPayload::SdkClientApprove(req) => {
-                self.handle_sdk_client_approve(req).await
-            }
-            UserAgentRequestPayload::SdkClientRevoke(req) => {
-                self.handle_sdk_client_revoke(req).await
-            }
-            UserAgentRequestPayload::SdkClientList(_) => self.handle_sdk_client_list().await,
-            _ => Err(TransportResponseError::UnexpectedRequestPayload),
-        }
-    }
-}
-
-type Output = Result<UserAgentResponse, TransportResponseError>;
-
-fn response(payload: UserAgentResponsePayload) -> UserAgentResponse {
-    UserAgentResponse {
-        payload: Some(payload),
-    }
-}
-
-impl UserAgentSession {
-    async fn handle_unseal_request(&mut self, req: UnsealStart) -> Output {
-        let secret = EphemeralSecret::random();
-        let public_key = PublicKey::from(&secret);
-
-        let client_pubkey_bytes: [u8; 32] = req
-            .client_pubkey
-            .try_into()
-            .map_err(|_| TransportResponseError::InvalidClientPubkeyLength)?;
-
-        let client_public_key = PublicKey::from(client_pubkey_bytes);
-
-        self.transition(UserAgentEvents::UnsealRequest(UnsealContext {
-            secret: Mutex::new(Some(secret)),
-            client_public_key,
-        }))?;
-
-        Ok(response(UserAgentResponsePayload::UnsealStartResponse(
-            UnsealStartResponse {
-                server_pubkey: public_key.as_bytes().to_vec(),
-            },
-        )))
-    }
-
-    async fn handle_unseal_encrypted_key(&mut self, req: UnsealEncryptedKey) -> Output {
-        let UserAgentStates::WaitingForUnsealKey(unseal_context) = self.state.state() else {
-            error!("Received unseal encrypted key in invalid state");
-            return Err(TransportResponseError::InvalidStateForUnsealEncryptedKey);
-        };
-        let ephemeral_secret = {
-            let mut secret_lock = unseal_context.secret.lock().unwrap();
-            let secret = secret_lock.take();
-            match secret {
-                Some(secret) => secret,
-                None => {
-                    drop(secret_lock);
-                    error!("Ephemeral secret already taken");
-                    self.transition(UserAgentEvents::ReceivedInvalidKey)?;
-                    return Ok(response(UserAgentResponsePayload::UnsealResult(
-                        UnsealResult::InvalidKey.into(),
-                    )));
-                }
-            }
-        };
-
-        let nonce = XNonce::from_slice(&req.nonce);
-
-        let shared_secret = ephemeral_secret.diffie_hellman(&unseal_context.client_public_key);
-        let cipher = XChaCha20Poly1305::new(shared_secret.as_bytes().into());
-
-        let mut seal_key_buffer = MemSafe::new(req.ciphertext.clone()).unwrap();
-
-        let decryption_result = {
-            let mut write_handle = seal_key_buffer.write().unwrap();
-            let write_handle = write_handle.deref_mut();
-            cipher.decrypt_in_place(nonce, &req.associated_data, write_handle)
-        };
-
-        match decryption_result {
-            Ok(_) => {
-                match self
-                    .props
-                    .actors
-                    .key_holder
-                    .ask(TryUnseal {
-                        seal_key_raw: seal_key_buffer,
-                    })
-                    .await
-                {
-                    Ok(_) => {
-                        info!("Successfully unsealed key with client-provided key");
-                        self.transition(UserAgentEvents::ReceivedValidKey)?;
-                        Ok(response(UserAgentResponsePayload::UnsealResult(
-                            UnsealResult::Success.into(),
-                        )))
-                    }
-                    Err(SendError::HandlerError(keyholder::Error::InvalidKey)) => {
-                        self.transition(UserAgentEvents::ReceivedInvalidKey)?;
-                        Ok(response(UserAgentResponsePayload::UnsealResult(
-                            UnsealResult::InvalidKey.into(),
-                        )))
-                    }
-                    Err(SendError::HandlerError(err)) => {
-                        error!(?err, "Keyholder failed to unseal key");
-                        self.transition(UserAgentEvents::ReceivedInvalidKey)?;
-                        Ok(response(UserAgentResponsePayload::UnsealResult(
-                            UnsealResult::InvalidKey.into(),
-                        )))
-                    }
-                    Err(err) => {
-                        error!(?err, "Failed to send unseal request to keyholder");
-                        self.transition(UserAgentEvents::ReceivedInvalidKey)?;
-                        Err(TransportResponseError::KeyHolderActorUnreachable)
-                    }
-                }
-            }
-            Err(err) => {
-                error!(?err, "Failed to decrypt unseal key");
-                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
-                Ok(response(UserAgentResponsePayload::UnsealResult(
-                    UnsealResult::InvalidKey.into(),
-                )))
-            }
-        }
-    }
-}
-
-impl UserAgentSession {
-    async fn handle_evm_wallet_create(&mut self) -> Output {
-        use evm_proto::wallet_create_response::Result as CreateResult;
-
-        let result = match self.props.actors.evm.ask(Generate {}).await {
-            Ok(address) => CreateResult::Wallet(evm_proto::WalletEntry {
-                address: address.as_slice().to_vec(),
-            }),
-            Err(err) => CreateResult::Error(map_evm_error("wallet create", err).into()),
-        };
-
-        Ok(response(UserAgentResponsePayload::EvmWalletCreate(
-            evm_proto::WalletCreateResponse {
-                result: Some(result),
-            },
-        )))
-    }
-
-    async fn handle_evm_wallet_list(&mut self) -> Output {
-        use evm_proto::wallet_list_response::Result as ListResult;
-
-        let result = match self.props.actors.evm.ask(ListWallets {}).await {
-            Ok(wallets) => ListResult::Wallets(evm_proto::WalletList {
-                wallets: wallets
-                    .into_iter()
-                    .map(|addr| evm_proto::WalletEntry {
-                        address: addr.as_slice().to_vec(),
-                    })
-                    .collect(),
-            }),
-            Err(err) => ListResult::Error(map_evm_error("wallet list", err).into()),
-        };
-
-        Ok(response(UserAgentResponsePayload::EvmWalletList(
-            evm_proto::WalletListResponse {
-                result: Some(result),
-            },
-        )))
-    }
-}
-
-impl UserAgentSession {
-    async fn handle_sdk_client_approve(&mut self, req: SdkClientApproveRequest) -> Output {
-        use sdk_client_approve_response::Result as ApproveResult;
-
-        if req.pubkey.len() != 32 {
-            return Err(TransportResponseError::SdkClientApprove(
-                ProtoSdkClientError::Internal,
-            ));
-        }
-
-        let now = std::time::SystemTime::now()
-            .duration_since(std::time::UNIX_EPOCH)
-            .unwrap_or_default()
-            .as_secs() as i32;
-
-        let mut conn = match self.props.db.get().await {
-            Ok(c) => c,
-            Err(e) => {
-                error!(?e, "Failed to get DB connection for sdk_client_approve");
-                return Err(TransportResponseError::SdkClientApprove(
-                    ProtoSdkClientError::Internal,
-                ));
-            }
-        };
-
-        let pubkey_bytes = req.pubkey.clone();
-        let insert_result = insert_into(program_client::table)
-            .values((
-                program_client::public_key.eq(&pubkey_bytes),
-                program_client::nonce.eq(1), // pre-incremented; challenge will use nonce=0
-                program_client::created_at.eq(now),
-                program_client::updated_at.eq(now),
-            ))
-            .returning((
-                program_client::id,
-                program_client::public_key,
-                program_client::created_at,
-            ))
-            .get_result::<(i32, Vec<u8>, i32)>(&mut conn)
-            .await;
-
-        match insert_result {
-            Ok((id, pubkey, created_at)) => Ok(response(
-                UserAgentResponsePayload::SdkClientApprove(SdkClientApproveResponse {
-                    result: Some(ApproveResult::Client(SdkClientEntry {
-                        id,
-                        pubkey,
-                        created_at,
-                    })),
-                }),
-            )),
-            Err(diesel::result::Error::DatabaseError(
-                diesel::result::DatabaseErrorKind::UniqueViolation,
-                _,
-            )) => Err(TransportResponseError::SdkClientApprove(
-                ProtoSdkClientError::AlreadyExists,
-            )),
-            Err(e) => {
-                error!(?e, "Failed to insert SDK client");
-                Err(TransportResponseError::SdkClientApprove(
-                    ProtoSdkClientError::Internal,
-                ))
-            }
-        }
-    }
-
-    async fn handle_sdk_client_list(&mut self) -> Output {
-        let mut conn = match self.props.db.get().await {
-            Ok(c) => c,
-            Err(e) => {
-                error!(?e, "Failed to get DB connection for sdk_client_list");
-                return Err(TransportResponseError::SdkClientList(
-                    ProtoSdkClientError::Internal,
-                ));
-            }
-        };
-
-        match program_client::table
-            .select((
-                program_client::id,
-                program_client::public_key,
-                program_client::created_at,
-            ))
-            .load::<(i32, Vec<u8>, i32)>(&mut conn)
+    ) -> Result<bool, ()> {
+        if self
+            .sender
+            .send(OutOfBand::ClientConnectionRequest {
+                pubkey: client_pubkey,
+            })
             .await
+            .is_err()
         {
-            Ok(rows) => Ok(response(UserAgentResponsePayload::SdkClientList(
-                SdkClientListResponse {
-                    result: Some(sdk_client_list_response::Result::Clients(SdkClientList {
-                        clients: rows
-                            .into_iter()
-                            .map(|(id, pubkey, created_at)| SdkClientEntry {
-                                id,
-                                pubkey,
-                                created_at,
-                            })
-                            .collect(),
-                    })),
-                },
-            ))),
-            Err(e) => {
-                error!(?e, "Failed to list SDK clients");
-                Err(TransportResponseError::SdkClientList(
-                    ProtoSdkClientError::Internal,
-                ))
-            }
+            return Err(());
         }
-    }
 
-    async fn handle_sdk_client_revoke(&mut self, req: SdkClientRevokeRequest) -> Output {
-        use sdk_client_revoke_response::Result as RevokeResult;
+        let _ = cancel_flag.changed().await;
 
-        let mut conn = match self.props.db.get().await {
-            Ok(c) => c,
-            Err(e) => {
-                error!(?e, "Failed to get DB connection for sdk_client_revoke");
-                return Err(TransportResponseError::SdkClientRevoke(
-                    ProtoSdkClientError::Internal,
-                ));
-            }
-        };
-
-        match diesel::delete(program_client::table)
-            .filter(program_client::id.eq(req.client_id))
-            .execute(&mut conn)
-            .await
-        {
-            Ok(0) => Err(TransportResponseError::SdkClientRevoke(
-                ProtoSdkClientError::NotFound,
-            )),
-            Ok(_) => Ok(response(UserAgentResponsePayload::SdkClientRevoke(
-                SdkClientRevokeResponse {
-                    result: Some(RevokeResult::Ok(())),
-                },
-            ))),
-            Err(diesel::result::Error::DatabaseError(
-                diesel::result::DatabaseErrorKind::ForeignKeyViolation,
-                _,
-            )) => Err(TransportResponseError::SdkClientRevoke(
-                ProtoSdkClientError::HasRelatedData,
-            )),
-            Err(e) => {
-                error!(?e, "Failed to delete SDK client");
-                Err(TransportResponseError::SdkClientRevoke(
-                    ProtoSdkClientError::Internal,
-                ))
-            }
-        }
-    }
-}
-
-fn map_evm_error<M>(op: &str, err: SendError<M, crate::actors::evm::Error>) -> evm_proto::EvmError {
-    use crate::actors::{evm::Error as EvmError, keyholder::Error as KhError};
-    match err {
-        SendError::HandlerError(EvmError::Keyholder(KhError::NotBootstrapped)) => {
-            evm_proto::EvmError::VaultSealed
-        }
-        SendError::HandlerError(err) => {
-            error!(?err, "EVM {op} failed");
-            evm_proto::EvmError::Internal
-        }
-        _ => {
-            error!("EVM actor unreachable during {op}");
-            evm_proto::EvmError::Internal
-        }
+        let _ = self.sender.send(OutOfBand::ClientConnectionCancel).await;
+        Ok(false)
     }
 }
 
 impl Actor for UserAgentSession {
     type Args = Self;
 
-    type Error = TransportResponseError;
+    type Error = Error;
 
     async fn on_start(
         args: Self::Args,
@@ -534,63 +125,8 @@ impl Actor for UserAgentSession {
             .await
             .map_err(|err| {
                 error!(?err, "Failed to register user agent connection with router");
-                TransportResponseError::ConnectionRegistrationFailed
+                Error::internal("Failed to register user agent connection with router")
             })?;
         Ok(args)
     }
-
-    async fn next(
-        &mut self,
-        _actor_ref: kameo::prelude::WeakActorRef<Self>,
-        mailbox_rx: &mut kameo::prelude::MailboxReceiver<Self>,
-    ) -> Option<kameo::mailbox::Signal<Self>> {
-        loop {
-            select! {
-                signal = mailbox_rx.recv() => {
-                    return signal;
-                }
-                msg = self.props.transport.recv() => {
-                    match msg {
-                        Some(request) => {
-                            match self.process_transport_inbound(request).await {
-                                Ok(response) => {
-                                    if self.props.transport.send(Ok(response)).await.is_err() {
-                                        error!(actor = "useragent", reason = "channel closed", "send.failed");
-                                        return Some(kameo::mailbox::Signal::Stop);
-                                    }
-                                }
-                                Err(err) => {
-                                    let should_stop = err.is_fatal();
-                                    if self.props.transport.send(Err(err)).await.is_err() {
-                                        error!(actor = "useragent", reason = "channel closed", "send.failed");
-                                        return Some(kameo::mailbox::Signal::Stop);
-                                    }
-
-                                    if should_stop {
-                                        return Some(kameo::mailbox::Signal::Stop);
-                                    }
-                                }
-                            }
-                        }
-                        None => {
-                            info!(actor = "useragent", "transport.closed");
-                            return Some(kameo::mailbox::Signal::Stop);
-                        }
-                    }
-                }
-            }
-        }
-    }
-}
-
-impl UserAgentSession {
-    pub fn new_test(db: crate::db::DatabasePool, actors: crate::actors::GlobalActors) -> Self {
-        use arbiter_proto::transport::DummyTransport;
-        let transport: super::Transport = Box::new(DummyTransport::new());
-        let props = UserAgentConnection::new(db, actors, transport);
-        Self {
-            props,
-            state: UserAgentStateMachine::new(DummyContext),
-        }
-    }
 }

server/crates/arbiter-server/src/actors/user_agent/session/connection.rs

@@ -0,0 +1,354 @@
+use std::sync::Mutex;
+
+use alloy::primitives::Address;
+use chacha20poly1305::{AeadInPlace, XChaCha20Poly1305, XNonce, aead::KeyInit};
+use kameo::error::SendError;
+use kameo::messages;
+use tracing::{error, info};
+use x25519_dalek::{EphemeralSecret, PublicKey};
+
+use crate::actors::keyholder::KeyHolderState;
+use crate::actors::user_agent::session::Error;
+use crate::evm::policies::{Grant, SpecificGrant};
+use crate::safe_cell::SafeCell;
+use crate::{
+    actors::{
+        evm::{
+            Generate, ListWallets, UseragentCreateGrant, UseragentDeleteGrant, UseragentListGrants,
+        },
+        keyholder::{self, Bootstrap, TryUnseal},
+        user_agent::session::{
+                UserAgentSession,
+                state::{UnsealContext, UserAgentEvents, UserAgentStates},
+            },
+    },
+    safe_cell::SafeCellHandle as _,
+};
+
+impl UserAgentSession {
+    fn take_unseal_secret(&mut self) -> Result<(EphemeralSecret, PublicKey), Error> {
+        let UserAgentStates::WaitingForUnsealKey(unseal_context) = self.state.state() else {
+            error!("Received encrypted key in invalid state");
+            return Err(Error::internal("Invalid state for unseal encrypted key"));
+        };
+
+        let ephemeral_secret = {
+            #[allow(
+                clippy::unwrap_used,
+                reason = "Mutex poison is unrecoverable and should panic"
+            )]
+            let mut secret_lock = unseal_context.secret.lock().unwrap();
+            let secret = secret_lock.take();
+            match secret {
+                Some(secret) => secret,
+                None => {
+                    drop(secret_lock);
+                    error!("Ephemeral secret already taken");
+                    return Err(Error::internal("Ephemeral secret already taken"));
+                }
+            }
+        };
+
+        Ok((ephemeral_secret, unseal_context.client_public_key))
+    }
+
+    fn decrypt_client_key_material(
+        ephemeral_secret: EphemeralSecret,
+        client_public_key: PublicKey,
+        nonce: &[u8],
+        ciphertext: &[u8],
+        associated_data: &[u8],
+    ) -> Result<SafeCell<Vec<u8>>, ()> {
+        let nonce = XNonce::from_slice(nonce);
+
+        let shared_secret = ephemeral_secret.diffie_hellman(&client_public_key);
+        let cipher = XChaCha20Poly1305::new(shared_secret.as_bytes().into());
+
+        let mut key_buffer = SafeCell::new(ciphertext.to_vec());
+
+        let decryption_result = key_buffer.write_inline(|write_handle| {
+            cipher.decrypt_in_place(nonce, associated_data, write_handle)
+        });
+
+        match decryption_result {
+            Ok(_) => Ok(key_buffer),
+            Err(err) => {
+                error!(?err, "Failed to decrypt encrypted key material");
+                Err(())
+            }
+        }
+    }
+}
+
+pub struct UnsealStartResponse {
+    pub server_pubkey: PublicKey,
+}
+
+#[derive(Debug, Error)]
+pub enum UnsealError {
+    #[error("Invalid key provided for unsealing")]
+    InvalidKey,
+    #[error("Internal error during unsealing process")]
+    General(#[from] super::Error),
+}
+
+#[derive(Debug, Error)]
+pub enum BootstrapError {
+    #[error("Invalid key provided for bootstrapping")]
+    InvalidKey,
+    #[error("Vault is already bootstrapped")]
+    AlreadyBootstrapped,
+
+    #[error("Internal error during bootstrapping process")]
+    General(#[from] super::Error),
+}
+
+#[messages]
+impl UserAgentSession {
+    #[message]
+    pub async fn handle_unseal_request(
+        &mut self,
+        client_pubkey: x25519_dalek::PublicKey,
+    ) -> Result<UnsealStartResponse, Error> {
+        let secret = EphemeralSecret::random();
+        let public_key = PublicKey::from(&secret);
+
+        self.transition(UserAgentEvents::UnsealRequest(UnsealContext {
+            secret: Mutex::new(Some(secret)),
+            client_public_key: client_pubkey,
+        }))?;
+
+        Ok(UnsealStartResponse {
+            server_pubkey: public_key,
+        })
+    }
+
+    #[message]
+    pub async fn handle_unseal_encrypted_key(
+        &mut self,
+        nonce: Vec<u8>,
+        ciphertext: Vec<u8>,
+        associated_data: Vec<u8>,
+    ) -> Result<(), UnsealError> {
+        let (ephemeral_secret, client_public_key) = match self.take_unseal_secret() {
+            Ok(values) => values,
+            Err(Error::State) => {
+                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
+                return Err(UnsealError::InvalidKey);
+            }
+            Err(_err) => {
+                return Err(Error::internal("Failed to take unseal secret").into());
+            }
+        };
+
+        let seal_key_buffer = match Self::decrypt_client_key_material(
+            ephemeral_secret,
+            client_public_key,
+            &nonce,
+            &ciphertext,
+            &associated_data,
+        ) {
+            Ok(buffer) => buffer,
+            Err(()) => {
+                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
+                return Err(UnsealError::InvalidKey);
+            }
+        };
+
+        match self
+            .props
+            .actors
+            .key_holder
+            .ask(TryUnseal {
+                seal_key_raw: seal_key_buffer,
+            })
+            .await
+        {
+            Ok(_) => {
+                info!("Successfully unsealed key with client-provided key");
+                self.transition(UserAgentEvents::ReceivedValidKey)?;
+                Ok(())
+            }
+            Err(SendError::HandlerError(keyholder::Error::InvalidKey)) => {
+                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
+                Err(UnsealError::InvalidKey)
+            }
+            Err(SendError::HandlerError(err)) => {
+                error!(?err, "Keyholder failed to unseal key");
+                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
+                Err(UnsealError::InvalidKey)
+            }
+            Err(err) => {
+                error!(?err, "Failed to send unseal request to keyholder");
+                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
+                Err(Error::internal("Vault actor error").into())
+            }
+        }
+    }
+
+    #[message]
+    pub(crate) async fn handle_bootstrap_encrypted_key(
+        &mut self,
+        nonce: Vec<u8>,
+        ciphertext: Vec<u8>,
+        associated_data: Vec<u8>,
+    ) -> Result<(), BootstrapError> {
+        let (ephemeral_secret, client_public_key) = match self.take_unseal_secret() {
+            Ok(values) => values,
+            Err(Error::State) => {
+                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
+                return Err(BootstrapError::InvalidKey);
+            }
+            Err(err) => return Err(err.into()),
+        };
+
+        let seal_key_buffer = match Self::decrypt_client_key_material(
+            ephemeral_secret,
+            client_public_key,
+            &nonce,
+            &ciphertext,
+            &associated_data,
+        ) {
+            Ok(buffer) => buffer,
+            Err(()) => {
+                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
+                return Err(BootstrapError::InvalidKey);
+            }
+        };
+
+        match self
+            .props
+            .actors
+            .key_holder
+            .ask(Bootstrap {
+                seal_key_raw: seal_key_buffer,
+            })
+            .await
+        {
+            Ok(_) => {
+                info!("Successfully bootstrapped vault with client-provided key");
+                self.transition(UserAgentEvents::ReceivedValidKey)?;
+                Ok(())
+            }
+            Err(SendError::HandlerError(keyholder::Error::AlreadyBootstrapped)) => {
+                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
+                Err(BootstrapError::AlreadyBootstrapped)
+            }
+            Err(SendError::HandlerError(err)) => {
+                error!(?err, "Keyholder failed to bootstrap vault");
+                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
+                Err(BootstrapError::InvalidKey)
+            }
+            Err(err) => {
+                error!(?err, "Failed to send bootstrap request to keyholder");
+                self.transition(UserAgentEvents::ReceivedInvalidKey)?;
+                Err(BootstrapError::General(Error::internal(
+                    "Vault actor error",
+                )))
+            }
+        }
+    }
+}
+
+#[messages]
+impl UserAgentSession {
+    #[message]
+    pub(crate) async fn handle_query_vault_state(&mut self) -> Result<KeyHolderState, Error> {
+        use crate::actors::keyholder::GetState;
+
+        let vault_state = match self.props.actors.key_holder.ask(GetState {}).await {
+            Ok(state) => state,
+            Err(err) => {
+                error!(?err, actor = "useragent", "keyholder.query.failed");
+                return Err(Error::internal("Vault is in broken state"));
+            }
+        };
+
+        Ok(vault_state)
+    }
+}
+
+#[messages]
+impl UserAgentSession {
+    #[message]
+    pub(crate) async fn handle_evm_wallet_create(&mut self) -> Result<Address, Error> {
+        match self.props.actors.evm.ask(Generate {}).await {
+            Ok(address) => Ok(address),
+            Err(SendError::HandlerError(err)) => Err(Error::internal(format!(
+                "EVM wallet generation failed: {err}"
+            ))),
+            Err(err) => {
+                error!(?err, "EVM actor unreachable during wallet create");
+                Err(Error::internal("EVM actor unreachable"))
+            }
+        }
+    }
+
+    #[message]
+    pub(crate) async fn handle_evm_wallet_list(&mut self) -> Result<Vec<Address>, Error> {
+        match self.props.actors.evm.ask(ListWallets {}).await {
+            Ok(wallets) => Ok(wallets),
+            Err(err) => {
+                error!(?err, "EVM wallet list failed");
+                Err(Error::internal("Failed to list EVM wallets"))
+            }
+        }
+    }
+}
+
+#[messages]
+impl UserAgentSession {
+    #[message]
+    pub(crate) async fn handle_grant_list(&mut self) -> Result<Vec<Grant<SpecificGrant>>, Error> {
+        match self.props.actors.evm.ask(UseragentListGrants {}).await {
+            Ok(grants) => Ok(grants),
+            Err(err) => {
+                error!(?err, "EVM grant list failed");
+                Err(Error::internal("Failed to list EVM grants"))
+            }
+        }
+    }
+
+    #[message]
+    pub(crate) async fn handle_grant_create(
+        &mut self,
+        client_id: i32,
+        basic: crate::evm::policies::SharedGrantSettings,
+        grant: crate::evm::policies::SpecificGrant,
+    ) -> Result<i32, Error> {
+        match self
+            .props
+            .actors
+            .evm
+            .ask(UseragentCreateGrant {
+                client_id,
+                basic,
+                grant,
+            })
+            .await
+        {
+            Ok(grant_id) => Ok(grant_id),
+            Err(err) => {
+                error!(?err, "EVM grant create failed");
+                Err(Error::internal("Failed to create EVM grant"))
+            }
+        }
+    }
+
+    #[message]
+    pub(crate) async fn handle_grant_delete(&mut self, grant_id: i32) -> Result<(), Error> {
+        match self
+            .props
+            .actors
+            .evm
+            .ask(UseragentDeleteGrant { grant_id })
+            .await
+        {
+            Ok(()) => Ok(()),
+            Err(err) => {
+                error!(?err, "EVM grant delete failed");
+                Err(Error::internal("Failed to delete EVM grant"))
+            }
+        }
+    }
+}