Improve provider reliability and usage visibility

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);
+    }
 }