fix(acp-runner): apply CodeRabbit inline findings

jramirezhdez02 · jramirezhdez02 · commit 3cdc314419bf · 2026-03-25T05:49:41.000-03:00
- prompt_converter: include cache_creation_tokens and cache_read_tokens in
  UsageUpdate context window calculation
- rpc_server: return error replies and early-return on store failures in
  handle_new_session, handle_set_session_mode, handle_set_session_model,
  and handle_fork_session; move new_id UUID generation after successful load
- connect_integration: use get_host() for container host (supports remote/
  rootless Docker), fix NKey seed to valid 58-char canonical format

Signed-off-by: Jorge &lt;jramirezhdez02@gmail.com&gt;
diff --git a/rsworkspace/crates/trogon-acp-runner/src/prompt_converter.rs b/rsworkspace/crates/trogon-acp-runner/src/prompt_converter.rs
@@ -93,11 +93,15 @@ impl PromptEventConverter {
 
             PromptEvent::UsageUpdate {
                 input_tokens,
+                cache_creation_tokens,
+                cache_read_tokens,
                 output_tokens,
                 context_window,
                 ..
             } => {
-                let used = (input_tokens + output_tokens) as u64;
+                let used =
+                    (input_tokens + cache_creation_tokens + cache_read_tokens + output_tokens)
+                        as u64;
                 let size = context_window.unwrap_or(DEFAULT_CONTEXT_WINDOW);
                 let notif = self.notif(SessionUpdate::UsageUpdate(UsageUpdate::new(used, size)));
                 (vec![notif], None)
diff --git a/rsworkspace/crates/trogon-acp-runner/src/rpc_server.rs b/rsworkspace/crates/trogon-acp-runner/src/rpc_server.rs
@@ -282,6 +282,12 @@ impl RpcServer {
 
         if let Err(e) = self.store.save(&session_id, &state).await {
             warn!(session_id = %session_id, error = %e, "rpc: failed to save new session");
+            self.reply(
+                &msg,
+                &serde_json::json!({ "error": format!("failed to save session: {e}") }),
+            )
+            .await;
+            return;
         }
 
         self.publish_session_ready(&session_id).await;
@@ -322,17 +328,28 @@ impl RpcServer {
         };
 
         let session_id = request.session_id.to_string();
-        match self.store.load(&session_id).await {
-            Ok(mut state) => {
-                state.mode = request.mode_id.to_string();
-                state.updated_at = now_iso8601();
-                if let Err(e) = self.store.save(&session_id, &state).await {
-                    warn!(session_id = %session_id, error = %e, "rpc: failed to persist mode update");
-                }
-            }
+        let mut state = match self.store.load(&session_id).await {
+            Ok(s) => s,
             Err(e) => {
                 warn!(session_id = %session_id, error = %e, "rpc: failed to load session for mode update");
+                self.reply(
+                    &msg,
+                    &serde_json::json!({ "error": format!("failed to load session: {e}") }),
+                )
+                .await;
+                return;
             }
+        };
+        state.mode = request.mode_id.to_string();
+        state.updated_at = now_iso8601();
+        if let Err(e) = self.store.save(&session_id, &state).await {
+            warn!(session_id = %session_id, error = %e, "rpc: failed to persist mode update");
+            self.reply(
+                &msg,
+                &serde_json::json!({ "error": format!("failed to save session: {e}") }),
+            )
+            .await;
+            return;
         }
 
         self.reply(&msg, &SetSessionModeResponse::new()).await;
@@ -349,17 +366,28 @@ impl RpcServer {
         };
 
         let session_id = request.session_id.to_string();
-        match self.store.load(&session_id).await {
-            Ok(mut state) => {
-                state.model = Some(request.model_id.to_string());
-                state.updated_at = now_iso8601();
-                if let Err(e) = self.store.save(&session_id, &state).await {
-                    warn!(session_id = %session_id, error = %e, "rpc: failed to persist model update");
-                }
-            }
+        let mut state = match self.store.load(&session_id).await {
+            Ok(s) => s,
             Err(e) => {
                 warn!(session_id = %session_id, error = %e, "rpc: failed to load session for model update");
+                self.reply(
+                    &msg,
+                    &serde_json::json!({ "error": format!("failed to load session: {e}") }),
+                )
+                .await;
+                return;
             }
+        };
+        state.model = Some(request.model_id.to_string());
+        state.updated_at = now_iso8601();
+        if let Err(e) = self.store.save(&session_id, &state).await {
+            warn!(session_id = %session_id, error = %e, "rpc: failed to persist model update");
+            self.reply(
+                &msg,
+                &serde_json::json!({ "error": format!("failed to save session: {e}") }),
+            )
+            .await;
+            return;
         }
 
         self.reply(&msg, &SetSessionModelResponse::new()).await;
@@ -429,20 +457,26 @@ impl RpcServer {
         };
 
         let source_id = request.session_id.to_string();
-        let new_id = uuid::Uuid::new_v4().to_string();
 
-        match self.store.load(&source_id).await {
-            Ok(mut state) => {
-                let now = now_iso8601();
-                state.created_at = now.clone();
-                state.updated_at = now;
-                if let Err(e) = self.store.save(&new_id, &state).await {
-                    warn!(new_id = %new_id, error = %e, "rpc: failed to save forked session");
-                }
-            }
+        let mut state = match self.store.load(&source_id).await {
+            Ok(s) => s,
             Err(e) => {
                 warn!(source_id = %source_id, error = %e, "rpc: failed to load source session for fork");
+                self.reply(
+                    &msg,
+                    &serde_json::json!({ "error": format!("failed to load source session: {e}") }),
+                )
+                .await;
+                return;
             }
+        };
+
+        let new_id = uuid::Uuid::new_v4().to_string();
+        let now = now_iso8601();
+        state.created_at = now.clone();
+        state.updated_at = now;
+        if let Err(e) = self.store.save(&new_id, &state).await {
+            warn!(new_id = %new_id, error = %e, "rpc: failed to save forked session");
         }
 
         self.reply(&msg, &ForkSessionResponse::new(new_id)).await;
diff --git a/rsworkspace/crates/trogon-nats/tests/connect_integration.rs b/rsworkspace/crates/trogon-nats/tests/connect_integration.rs
@@ -15,20 +15,21 @@ use testcontainers_modules::testcontainers::runners::AsyncRunner;
 use trogon_nats::auth::{NatsAuth, NatsConfig};
 use trogon_nats::connect::{ConnectError, connect};
 
-async fn start_nats() -> Result<(ContainerAsync<Nats>, u16), Box<dyn std::error::Error>> {
+async fn start_nats() -> Result<(ContainerAsync<Nats>, String, u16), Box<dyn std::error::Error>> {
     let container = Nats::default().start().await?;
+    let host = container.get_host().await?.to_string();
     let port = container.get_host_port_ipv4(4222).await?;
-    Ok((container, port))
+    Ok((container, host, port))
 }
 
 /// Covers the `NatsAuth::None` arm (lines 123-128) and the success branch (130-138).
 /// Also exercises `apply_reconnect_options` (lines 69-74) indirectly.
 #[tokio::test]
 #[ignore = "requires Docker"]
 async fn connect_with_no_auth_succeeds() -> Result<(), Box<dyn std::error::Error>> {
-    let (_container, port) = start_nats().await?;
+    let (_container, host, port) = start_nats().await?;
 
-    let config = NatsConfig::new(vec![format!("nats://127.0.0.1:{port}")], NatsAuth::None);
+    let config = NatsConfig::new(vec![format!("nats://{host}:{port}")], NatsAuth::None);
 
     connect(&config, Duration::from_secs(10))
         .await
@@ -42,10 +43,10 @@ async fn connect_with_no_auth_succeeds() -> Result<(), Box<dyn std::error::Error
 async fn connect_with_token_auth_succeeds_on_open_server() -> Result<(), Box<dyn std::error::Error>>
 {
     // An open NATS server accepts any token — the token is just passed through.
-    let (_container, port) = start_nats().await?;
+    let (_container, host, port) = start_nats().await?;
 
     let config = NatsConfig::new(
-        vec![format!("nats://127.0.0.1:{port}")],
+        vec![format!("nats://{host}:{port}")],
         NatsAuth::Token("any-token".to_string()),
     );
 
@@ -60,10 +61,10 @@ async fn connect_with_token_auth_succeeds_on_open_server() -> Result<(), Box<dyn
 #[ignore = "requires Docker"]
 async fn connect_with_user_password_succeeds_on_open_server()
 -> Result<(), Box<dyn std::error::Error>> {
-    let (_container, port) = start_nats().await?;
+    let (_container, host, port) = start_nats().await?;
 
     let config = NatsConfig::new(
-        vec![format!("nats://127.0.0.1:{port}")],
+        vec![format!("nats://{host}:{port}")],
         NatsAuth::UserPassword {
             user: "user".to_string(),
             password: "pass".to_string(),
@@ -84,17 +85,15 @@ async fn connect_with_user_password_succeeds_on_open_server()
 #[tokio::test]
 #[ignore = "requires Docker"]
 async fn connect_with_nkey_auth_on_open_server() -> Result<(), Box<dyn std::error::Error>> {
-    let (_container, port) = start_nats().await?;
+    let (_container, host, port) = start_nats().await?;
 
-    // A valid NKey user seed (base32-encoded, 58-char canonical format).
-    // On an open server the key is not validated — the test simply exercises
-    // the `NatsAuth::NKey` branch in `connect()`.
-    let seed = "SUACSSL3UAHUDXKFSNVUZRF5UHPMWZ6BFDTJ7M6USDRCRBZLYKI4LZPFZFR".to_string();
+    // A valid NKey user seed (base32-encoded, 58-char canonical format,
+    // starts with "SU"). On an open server the key is not validated against
+    // a registered user — the test simply exercises the `NatsAuth::NKey`
+    // branch in `connect()`.
+    let seed = "SUANQDPB2RUOE4ETUA26CNX7FUKE5ZZKFCQIIW63OX225F2CO7UEXTM7ZY".to_string();
 
-    let config = NatsConfig::new(
-        vec![format!("nats://127.0.0.1:{port}")],
-        NatsAuth::NKey(seed),
-    );
+    let config = NatsConfig::new(vec![format!("nats://{host}:{port}")], NatsAuth::NKey(seed));
 
     let result = connect(&config, Duration::from_secs(10)).await;
     assert!(
@@ -134,10 +133,11 @@ async fn connect_with_wrong_token_returns_authorization_violation()
         .with_cmd(["--auth", "correct-token"])
         .start()
         .await?;
+    let host = container.get_host().await?.to_string();
     let port = container.get_host_port_ipv4(4222).await?;
 
     let config = NatsConfig::new(
-        vec![format!("nats://127.0.0.1:{port}")],
+        vec![format!("nats://{host}:{port}")],
         NatsAuth::Token("wrong-token".to_string()),
     );
 
@@ -160,10 +160,11 @@ async fn connect_with_correct_token_succeeds() -> Result<(), Box<dyn std::error:
         .with_cmd(["--auth", "correct-token"])
         .start()
         .await?;
+    let host = container.get_host().await?.to_string();
     let port = container.get_host_port_ipv4(4222).await?;
 
     let config = NatsConfig::new(
-        vec![format!("nats://127.0.0.1:{port}")],
+        vec![format!("nats://{host}:{port}")],
         NatsAuth::Token("correct-token".to_string()),
     );
 
@@ -190,7 +191,7 @@ async fn connect_to_unreachable_server_returns_ok_with_background_retry() {
     let port = listener.local_addr().unwrap().port();
     drop(listener);
 
-    let config = NatsConfig::new(vec![format!("nats://127.0.0.1:{port}")], NatsAuth::None);
+    let config = NatsConfig::new(vec![format!("nats://{host}:{port}")], NatsAuth::None);
 
     // connect() must return within a few seconds (INITIAL_CONNECT_CHECK_SECS + margin).
     let result = tokio::time::timeout(
