Improve provider reliability and usage visibility

crates/galaxy_agent_core/src/provider_run.rs

@@ -502,6 +502,10 @@ impl ProviderRun {
         self.model_retries
     }
 
+    pub fn max_model_retries_per_turn(&self) -> u32 {
+        self.limits.max_model_retries_per_turn
+    }
+
     pub fn is_terminal(&self) -> bool {
         matches!(
             self.state,
@@ -665,20 +669,20 @@ impl ProviderRun {
             }
             ProviderRunState::Failed { failure } => match failure.kind {
                 ProviderRunFailureKind::ModelCall
-                    if !failure
+                    if failure
                         .source
                         .as_ref()
-                        .is_some_and(|source| !source.recoverable) =>
+                        .is_none_or(|source| source.recoverable) =>
                 {
                     return Err(invalid(
                         "model-call failure lacks a non-recoverable source".to_string(),
                     ));
                 }
                 ProviderRunFailureKind::RetryLimitExceeded
-                    if !failure
+                    if failure
                         .source
                         .as_ref()
-                        .is_some_and(|source| source.recoverable) =>
+                        .is_none_or(|source| !source.recoverable) =>
                 {
                     return Err(invalid(
                         "retry-limit failure lacks a recoverable source".to_string(),

crates/galaxy_agent_core/src/types.rs

@@ -223,15 +223,21 @@ impl ToolEvent {
 
 #[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
 pub struct Usage {
+    /// Input tokens not served by or written to a provider prompt cache.
     pub input_tokens: u64,
     pub output_tokens: u64,
+    /// Input tokens served from the provider prompt cache.
     pub cached_input_tokens: u64,
+    /// Input tokens written to the provider prompt cache.
     pub cache_creation_input_tokens: u64,
 }
 
 impl Usage {
     pub fn total_tokens(&self) -> u64 {
-        self.input_tokens.saturating_add(self.output_tokens)
+        self.input_tokens
+            .saturating_add(self.cached_input_tokens)
+            .saturating_add(self.cache_creation_input_tokens)
+            .saturating_add(self.output_tokens)
     }
 }
 
@@ -268,6 +274,9 @@ pub enum AgentEvent {
     TurnStarted {
         runtime_request_id: String,
     },
+    /// A transport heartbeat proving that the current model stream is still connected.
+    /// Consumers should use this to refresh idle timeouts without rendering output.
+    KeepAlive,
     TextDelta {
         text: String,
     },

crates/galaxy_agent_core/src/types_tests.rs

@@ -26,9 +26,9 @@ fn truncates_large_tool_results_for_provider_request() {
 }
 
 #[test]
-fn usage_total_excludes_cached_breakdown_to_avoid_double_counting() {
+fn usage_total_combines_disjoint_cache_and_miss_buckets() {
     let usage = Usage {
-        input_tokens: 100,
+        input_tokens: 10,
         output_tokens: 25,
         cached_input_tokens: 80,
         cache_creation_input_tokens: 10,

crates/galaxy_agent_rig/src/bedrock_tests.rs

@@ -83,7 +83,7 @@ fn normalizes_bedrock_usage_and_max_token_stop() {
     assert_eq!(
         map_usage((&response).into()),
         Usage {
-            input_tokens: 100,
+            input_tokens: 50,
             output_tokens: 25,
             cached_input_tokens: 40,
             cache_creation_input_tokens: 10,

crates/galaxy_agent_rig/src/chatgpt.rs

@@ -116,6 +116,7 @@ impl ChatGPTSubscriptionRuntime {
                 AgentEvent::ReasoningDelta { .. }
                 | AgentEvent::ReasoningCompleted { .. }
                 | AgentEvent::TurnStarted { .. }
+                | AgentEvent::KeepAlive
                 | AgentEvent::UsageUpdated { .. }
                 | AgentEvent::RuntimeActivityUpdated { .. }
                 | AgentEvent::ContextUsageUpdated { .. }

crates/galaxy_agent_rig/src/openai_compatible_tests.rs

@@ -73,7 +73,7 @@ async fn rig_stream_maps_reasoning_text_usage_and_stop() {
             },
             AgentEvent::UsageUpdated {
                 usage: Usage {
-                    input_tokens: 4,
+                    input_tokens: 2,
                     output_tokens: 6,
                     cached_input_tokens: 2,
                     cache_creation_input_tokens: 0,

crates/galaxy_agent_rig/src/stream.rs

@@ -125,11 +125,15 @@ where
                             });
                         }
                         Ok(StreamedAssistantContent::Unknown(value)) => {
-                            yield Err(AgentError::new(
-                                AgentErrorKind::Protocol,
-                                format!("Rig returned an unsupported provider event: {value}"),
-                            ));
-                            return;
+                            if is_keepalive_event(&value) {
+                                yield Ok(AgentEvent::KeepAlive);
+                            } else {
+                                yield Err(AgentError::new(
+                                    AgentErrorKind::Protocol,
+                                    format!("Rig returned an unsupported provider event: {value}"),
+                                ));
+                                return;
+                            }
                         }
                         Err(error) => {
                             if let Some(reason) = completion_error_stop_reason(&error) {
@@ -255,8 +259,24 @@ fn stopped_with_reason(runtime_request_id: String, reason: StopReason) -> AgentE
 }
 
 pub(crate) fn map_usage(usage: rig_core::completion::Usage) -> Usage {
+    // Rig preserves provider-native input semantics: OpenAI includes cached
+    // tokens in `input_tokens`, while Anthropic reports cache reads/writes
+    // separately. `total_tokens - output_tokens` gives the normalized prompt
+    // size for both shapes, so store only the uncached portion in
+    // `input_tokens` and keep cache reads/writes disjoint.
+    let total_input_tokens = if usage.total_tokens > 0 && usage.total_tokens >= usage.output_tokens
+    {
+        usage.total_tokens - usage.output_tokens
+    } else {
+        usage
+            .input_tokens
+            .saturating_add(usage.cached_input_tokens)
+            .saturating_add(usage.cache_creation_input_tokens)
+    };
     Usage {
-        input_tokens: usage.input_tokens,
+        input_tokens: total_input_tokens
+            .saturating_sub(usage.cached_input_tokens)
+            .saturating_sub(usage.cache_creation_input_tokens),
         output_tokens: usage.output_tokens,
         cached_input_tokens: usage.cached_input_tokens,
         cache_creation_input_tokens: usage.cache_creation_input_tokens,
@@ -307,6 +327,10 @@ fn json_value_indicates_context_window_exceeded(value: &serde_json::Value) -> bo
     }
 }
 
+fn is_keepalive_event(value: &serde_json::Value) -> bool {
+    value.get("type").and_then(serde_json::Value::as_str) == Some("keepalive")
+}
+
 fn text_indicates_context_window_exceeded(text: &str) -> bool {
     let normalized = text.to_ascii_lowercase();
     normalized.contains("modelcontextwindowexceeded")
@@ -352,7 +376,7 @@ fn map_completion_error(error: CompletionError) -> AgentError {
     mapped.recoverable = matches!(
         kind,
         AgentErrorKind::RateLimited | AgentErrorKind::Transport
-    );
+    ) || status.is_some_and(|status| (500..=599).contains(&status));
     mapped
 }
 
@@ -361,7 +385,9 @@ mod tests {
     use galaxy_agent_core::{AgentErrorKind, StopReason};
     use rig_core::completion::CompletionError;
 
-    use super::{completion_error_stop_reason, domain_tool_call, map_completion_error};
+    use super::{
+        completion_error_stop_reason, domain_tool_call, is_keepalive_event, map_completion_error,
+    };
 
     #[test]
     fn domain_tool_call_prefers_responses_call_id() {
@@ -428,4 +454,32 @@ mod tests {
             Some(StopReason::ContextWindowExceeded)
         );
     }
+
+    #[test]
+    fn provider_keepalive_event_is_a_transport_heartbeat() {
+        assert!(is_keepalive_event(&serde_json::json!({
+            "type": "keepalive",
+            "sequence_number": 3,
+        })));
+        assert!(!is_keepalive_event(&serde_json::json!({
+            "type": "unsupported",
+        })));
+    }
+
+    #[test]
+    fn provider_server_error_is_recoverable() {
+        let status = rig_core::http_client::Response::builder()
+            .status(503)
+            .body(())
+            .unwrap()
+            .status();
+        let error = CompletionError::from_http_response(
+            status,
+            r#"{"error":{"message":"Service temporarily unavailable"}}"#,
+        );
+
+        let mapped = map_completion_error(error);
+        assert_eq!(mapped.kind, AgentErrorKind::Provider);
+        assert!(mapped.recoverable);
+    }
 }

crates/galaxy_core/src/paths.rs

@@ -389,9 +389,7 @@ fn migrate_directory_contents(source_dir: &Path, target_dir: &Path) {
         }
     }
 
-    if std::fs::read_dir(source_dir)
-        .is_ok_and(|mut entries| entries.next().is_none())
-    {
+    if std::fs::read_dir(source_dir).is_ok_and(|mut entries| entries.next().is_none()) {
         let _ = std::fs::remove_dir(source_dir);
     }
 }