refactor: extract runtime state backends

crates/aether-usage-runtime/Cargo.toml

@@ -11,6 +11,7 @@ aether-ai-formats.workspace = true
 aether-contracts.workspace = true
 aether-data.workspace = true
 aether-data-contracts.workspace = true
+aether-runtime-state.workspace = true
 async-trait.workspace = true
 base64.workspace = true
 serde.workspace = true

crates/aether-usage-runtime/src/queue.rs

@@ -1,39 +1,33 @@
+use std::sync::Arc;
+
 use serde_json::json;
 
-use aether_data::driver::redis::{
-    RedisConsumerGroup, RedisConsumerName, RedisStreamEntry, RedisStreamName,
-    RedisStreamReclaimConfig, RedisStreamRunner, RedisStreamRunnerConfig,
-};
 use aether_data_contracts::DataLayerError;
+use aether_runtime_state::{RuntimeQueueEntry, RuntimeQueueReclaimConfig, RuntimeQueueStore};
 
 use super::config::UsageRuntimeConfig;
 use super::event::UsageEvent;
 
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub struct UsageQueue {
-    runner: RedisStreamRunner,
+    runner: Arc<dyn RuntimeQueueStore>,
     config: UsageRuntimeConfig,
-    stream: RedisStreamName,
-    group: RedisConsumerGroup,
-    dlq_stream: RedisStreamName,
+    stream: String,
+    group: String,
+    dlq_stream: String,
 }
 
 impl UsageQueue {
     pub fn new(
-        runner: RedisStreamRunner,
+        runner: Arc<dyn RuntimeQueueStore>,
         config: UsageRuntimeConfig,
     ) -> Result<Self, DataLayerError> {
         config.validate()?;
-        let tuned_runner = RedisStreamRunner::new(
-            runner.client().clone(),
-            runner.keyspace().clone(),
-            usage_stream_runner_config(&config),
-        )?;
         Ok(Self {
-            runner: tuned_runner,
-            stream: RedisStreamName(config.stream_key.clone()),
-            group: RedisConsumerGroup(config.consumer_group.clone()),
-            dlq_stream: RedisStreamName(config.dlq_stream_key.clone()),
+            runner,
+            stream: config.stream_key.clone(),
+            group: config.consumer_group.clone(),
+            dlq_stream: config.dlq_stream_key.clone(),
             config,
         })
     }
@@ -53,32 +47,36 @@ impl UsageQueue {
 
     pub async fn read_group(
         &self,
-        consumer: &RedisConsumerName,
-    ) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
+        consumer: &str,
+    ) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
         self.runner
-            .read_group(&self.stream, &self.group, consumer)
+            .read_group(
+                &self.stream,
+                &self.group,
+                consumer,
+                self.config.consumer_batch_size.max(1),
+                Some(self.config.consumer_block_ms.max(1)),
+            )
             .await
     }
 
     pub async fn claim_stale(
         &self,
-        consumer: &RedisConsumerName,
+        consumer: &str,
         start_id: &str,
-    ) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
-        Ok(self
-            .runner
+    ) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
+        self.runner
             .claim_stale(
                 &self.stream,
                 &self.group,
                 consumer,
                 start_id,
-                RedisStreamReclaimConfig {
+                RuntimeQueueReclaimConfig {
                     min_idle_ms: self.config.reclaim_idle_ms,
                     count: self.config.reclaim_count,
                 },
             )
-            .await?
-            .entries)
+            .await
     }
 
     pub async fn ack_and_delete(&self, ids: &[String]) -> Result<(), DataLayerError> {
@@ -89,27 +87,37 @@ impl UsageQueue {
 
     pub async fn push_dead_letter(
         &self,
-        entry: &RedisStreamEntry,
+        entry: &RuntimeQueueEntry,
         error: &str,
     ) -> Result<String, DataLayerError> {
+        let fields = std::collections::BTreeMap::from([(
+            "payload".to_string(),
+            serde_json::to_string(&json!({
+                "entry_id": entry.id,
+                "fields": entry.fields,
+                "error": error,
+            }))
+            .map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?,
+        )]);
         self.runner
-            .append_json(
-                &self.dlq_stream,
-                "payload",
-                &json!({
-                    "entry_id": entry.id,
-                    "fields": entry.fields,
-                    "error": error,
-                }),
-            )
+            .append_fields_with_maxlen(&self.dlq_stream, &fields, None)
             .await
     }
 }
 
-fn usage_stream_runner_config(config: &UsageRuntimeConfig) -> RedisStreamRunnerConfig {
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[cfg(test)]
+struct UsageQueueRuntimeSettings {
+    command_timeout_ms: Option<u64>,
+    read_block_ms: Option<u64>,
+    read_count: usize,
+}
+
+#[cfg(test)]
+fn usage_queue_runtime_settings(config: &UsageRuntimeConfig) -> UsageQueueRuntimeSettings {
     let read_block_ms = config.consumer_block_ms.max(1);
     let command_timeout_ms = read_block_ms.saturating_add(2_000).max(5_000);
-    RedisStreamRunnerConfig {
+    UsageQueueRuntimeSettings {
         command_timeout_ms: Some(command_timeout_ms),
         read_block_ms: Some(read_block_ms),
         read_count: config.consumer_batch_size.max(1),
@@ -118,27 +126,10 @@ fn usage_stream_runner_config(config: &UsageRuntimeConfig) -> RedisStreamRunnerC
 
 #[cfg(test)]
 mod tests {
-    use super::{usage_stream_runner_config, UsageQueue};
+    use super::{usage_queue_runtime_settings, UsageQueue, UsageQueueRuntimeSettings};
     use crate::UsageRuntimeConfig;
-    use aether_data::driver::redis::{RedisClientConfig, RedisClientFactory, RedisStreamRunner};
-
-    fn sample_runner() -> RedisStreamRunner {
-        let config = RedisClientConfig {
-            url: "redis://127.0.0.1/0".to_string(),
-            key_prefix: Some("aether".to_string()),
-        };
-        let client = RedisClientFactory::new(config.clone())
-            .expect("factory should build")
-            .connect_lazy()
-            .expect("client should build");
-
-        RedisStreamRunner::new(
-            client,
-            config.keyspace(),
-            aether_data::driver::redis::RedisStreamRunnerConfig::default(),
-        )
-        .expect("runner should build")
-    }
+    use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeState};
+    use std::sync::Arc;
 
     #[test]
     fn usage_queue_applies_runtime_block_and_batch_settings() {
@@ -148,15 +139,19 @@ mod tests {
             consumer_batch_size: 123,
             ..UsageRuntimeConfig::default()
         };
-        let queue = UsageQueue::new(sample_runner(), config)
-            .expect("usage queue should build from runtime config");
+        let queue = UsageQueue::new(
+            Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())),
+            config,
+        )
+        .expect("usage queue should build from runtime config");
 
         assert_eq!(
-            queue.runner.config(),
-            usage_stream_runner_config(&queue.config)
+            usage_queue_runtime_settings(&queue.config),
+            UsageQueueRuntimeSettings {
+                command_timeout_ms: Some(5_000),
+                read_block_ms: Some(750),
+                read_count: 123,
+            }
         );
-        assert_eq!(queue.runner.config().read_block_ms, Some(750));
-        assert_eq!(queue.runner.config().read_count, 123);
-        assert_eq!(queue.runner.config().command_timeout_ms, Some(5_000));
     }
 }

crates/aether-usage-runtime/src/runtime.rs

@@ -3,9 +3,9 @@ use std::pin::Pin;
 use std::sync::Arc;
 
 use aether_contracts::ExecutionTelemetry;
-use aether_data::driver::redis::RedisStreamRunner;
 use aether_data_contracts::repository::usage::UpsertUsageRecord;
 use aether_data_contracts::DataLayerError;
+use aether_runtime_state::RuntimeQueueStore;
 use async_trait::async_trait;
 use tracing::warn;
 
@@ -61,8 +61,8 @@ pub trait UsageRuntimeAccess:
     + Sync
 {
     fn has_usage_writer(&self) -> bool;
-    fn has_usage_worker_runner(&self) -> bool;
-    fn usage_worker_runner(&self) -> Option<RedisStreamRunner>;
+    fn has_usage_worker_queue(&self) -> bool;
+    fn usage_worker_queue(&self) -> Option<Arc<dyn RuntimeQueueStore>>;
 
     async fn body_capture_policy(&self) -> Result<UsageBodyCapturePolicy, DataLayerError> {
         Ok(UsageBodyCapturePolicy::default())
@@ -104,7 +104,7 @@ impl UsageRuntime {
     where
         T: UsageRuntimeAccess,
     {
-        self.is_enabled() && data.has_usage_writer() && data.has_usage_worker_runner()
+        self.is_enabled() && data.has_usage_writer() && data.has_usage_worker_queue()
     }
 
     pub fn spawn_worker<T>(&self, data: Arc<T>) -> Option<tokio::task::JoinHandle<()>>
@@ -114,7 +114,7 @@ impl UsageRuntime {
         if !self.can_spawn_worker(data.as_ref()) {
             return None;
         }
-        let runner = data.usage_worker_runner()?;
+        let runner = data.usage_worker_queue()?;
         let worker = build_usage_queue_worker(runner, data, self.config.clone()).ok()?;
         Some(worker.spawn())
     }
@@ -352,7 +352,7 @@ impl UsageRuntime {
     where
         T: UsageRuntimeAccess,
     {
-        if let Some(runner) = data.usage_worker_runner() {
+        if let Some(runner) = data.usage_worker_queue() {
             match UsageQueue::new(runner, self.config.clone()) {
                 Ok(queue) => match queue.enqueue(&event).await {
                     Ok(_) => return,
@@ -543,16 +543,14 @@ fn now_unix_secs() -> u64 {
 
 #[cfg(test)]
 mod tests {
-    use std::sync::Mutex;
+    use std::sync::{Arc, Mutex};
 
-    use aether_data::driver::redis::{
-        RedisClientConfig, RedisClientFactory, RedisStreamRunner, RedisStreamRunnerConfig,
-    };
     use aether_data_contracts::repository::settlement::{
         StoredUsageSettlement, UsageSettlementInput,
     };
     use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
     use aether_data_contracts::DataLayerError;
+    use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeQueueStore, RuntimeState};
     use async_trait::async_trait;
     use serde_json::json;
 
@@ -571,27 +569,9 @@ mod tests {
         records: Mutex<Vec<UpsertUsageRecord>>,
     }
 
-    struct RedisConfiguredUsageStore {
+    struct QueueConfiguredUsageStore {
         inner: NoRedisUsageStore,
-        runner: RedisStreamRunner,
-    }
-
-    fn sample_runner() -> RedisStreamRunner {
-        let config = RedisClientConfig {
-            url: "redis://127.0.0.1/0".to_string(),
-            key_prefix: Some("aether".to_string()),
-        };
-        let client = RedisClientFactory::new(config.clone())
-            .expect("factory should build")
-            .connect_lazy()
-            .expect("client should build");
-
-        RedisStreamRunner::new(
-            client,
-            config.keyspace(),
-            RedisStreamRunnerConfig::default(),
-        )
-        .expect("runner should build")
+        queue: Arc<dyn RuntimeQueueStore>,
     }
 
     #[async_trait]
@@ -644,17 +624,17 @@ mod tests {
             true
         }
 
-        fn has_usage_worker_runner(&self) -> bool {
+        fn has_usage_worker_queue(&self) -> bool {
             false
         }
 
-        fn usage_worker_runner(&self) -> Option<RedisStreamRunner> {
+        fn usage_worker_queue(&self) -> Option<Arc<dyn RuntimeQueueStore>> {
             None
         }
     }
 
     #[async_trait]
-    impl UsageRecordWriter for RedisConfiguredUsageStore {
+    impl UsageRecordWriter for QueueConfiguredUsageStore {
         async fn upsert_usage_record(
             &self,
             record: UpsertUsageRecord,
@@ -664,7 +644,7 @@ mod tests {
     }
 
     #[async_trait]
-    impl UsageSettlementWriter for RedisConfiguredUsageStore {
+    impl UsageSettlementWriter for QueueConfiguredUsageStore {
         fn has_usage_settlement_writer(&self) -> bool {
             false
         }
@@ -678,14 +658,14 @@ mod tests {
     }
 
     #[async_trait]
-    impl UsageBillingEventEnricher for RedisConfiguredUsageStore {
+    impl UsageBillingEventEnricher for QueueConfiguredUsageStore {
         async fn enrich_usage_event(&self, _event: &mut UsageEvent) -> Result<(), DataLayerError> {
             Ok(())
         }
     }
 
     #[async_trait]
-    impl ManualProxyNodeCounter for RedisConfiguredUsageStore {
+    impl ManualProxyNodeCounter for QueueConfiguredUsageStore {
         async fn increment_manual_proxy_node_requests(
             &self,
             _node_id: &str,
@@ -697,17 +677,17 @@ mod tests {
         }
     }
 
-    impl UsageRuntimeAccess for RedisConfiguredUsageStore {
+    impl UsageRuntimeAccess for QueueConfiguredUsageStore {
         fn has_usage_writer(&self) -> bool {
             true
         }
 
-        fn has_usage_worker_runner(&self) -> bool {
+        fn has_usage_worker_queue(&self) -> bool {
             true
         }
 
-        fn usage_worker_runner(&self) -> Option<RedisStreamRunner> {
-            Some(self.runner.clone())
+        fn usage_worker_queue(&self) -> Option<Arc<dyn RuntimeQueueStore>> {
+            Some(Arc::clone(&self.queue))
         }
     }
 
@@ -750,10 +730,9 @@ mod tests {
             ..UsageRuntimeConfig::default()
         })
         .expect("usage runtime should build");
-        let runner = sample_runner();
-        let store = RedisConfiguredUsageStore {
+        let store = QueueConfiguredUsageStore {
             inner: NoRedisUsageStore::default(),
-            runner,
+            queue: Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())),
         };
         let event = UsageEvent::new(
             UsageEventType::Failed,

crates/aether-usage-runtime/src/worker.rs

@@ -1,9 +1,9 @@
 use std::sync::Arc;
 use std::time::Duration;
 
-use aether_data::driver::redis::{RedisConsumerName, RedisStreamEntry, RedisStreamRunner};
 use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
 use aether_data_contracts::DataLayerError;
+use aether_runtime_state::{RuntimeQueueEntry, RuntimeQueueStore};
 use async_trait::async_trait;
 use tracing::warn;
 
@@ -60,18 +60,18 @@ where
 pub struct UsageQueueWorker {
     queue: UsageQueue,
     recorder: Arc<dyn UsageEventRecorder>,
-    consumer: RedisConsumerName,
+    consumer: String,
     config: UsageRuntimeConfig,
 }
 
 impl UsageQueueWorker {
     pub fn new(
-        runner: RedisStreamRunner,
+        runner: Arc<dyn RuntimeQueueStore>,
         recorder: Arc<dyn UsageEventRecorder>,
         config: UsageRuntimeConfig,
     ) -> Result<Self, DataLayerError> {
         let queue = UsageQueue::new(runner, config.clone())?;
-        let consumer = RedisConsumerName(consumer_name());
+        let consumer = consumer_name();
         Ok(Self {
             queue,
             recorder,
@@ -89,7 +89,7 @@ impl UsageQueueWorker {
             warn!(
                 event_name = "usage_worker_consumer_group_failed",
                 log_type = "ops",
-                worker_consumer = %self.consumer.0,
+                worker_consumer = %self.consumer,
                 worker_group = %self.config.consumer_group,
                 error = %err,
                 "usage worker failed to ensure consumer group"
@@ -111,7 +111,7 @@ impl UsageQueueWorker {
                                 warn!(
                                     event_name = "usage_worker_reclaim_process_failed",
                                     log_type = "ops",
-                                    worker_consumer = %self.consumer.0,
+                                    worker_consumer = %self.consumer,
                                     worker_group = %self.config.consumer_group,
                                     error = %err,
                                     "usage worker failed while reclaiming stale entries"
@@ -121,7 +121,7 @@ impl UsageQueueWorker {
                         Err(err) => warn!(
                             event_name = "usage_worker_reclaim_failed",
                             log_type = "ops",
-                            worker_consumer = %self.consumer.0,
+                            worker_consumer = %self.consumer,
                             worker_group = %self.config.consumer_group,
                             error = %err,
                             "usage worker failed to reclaim stale entries"
@@ -135,7 +135,7 @@ impl UsageQueueWorker {
                                 warn!(
                                     event_name = "usage_worker_process_failed",
                                     log_type = "ops",
-                                    worker_consumer = %self.consumer.0,
+                                    worker_consumer = %self.consumer,
                                     worker_group = %self.config.consumer_group,
                                     error = %err,
                                     "usage worker failed to process queue entries"
@@ -147,7 +147,7 @@ impl UsageQueueWorker {
                             warn!(
                                 event_name = "usage_worker_read_failed",
                                 log_type = "ops",
-                                worker_consumer = %self.consumer.0,
+                                worker_consumer = %self.consumer,
                                 worker_group = %self.config.consumer_group,
                                 error = %err,
                                 "usage worker failed to read queue"
@@ -160,7 +160,7 @@ impl UsageQueueWorker {
         }
     }
 
-    async fn process_entries(&self, entries: Vec<RedisStreamEntry>) -> Result<(), DataLayerError> {
+    async fn process_entries(&self, entries: Vec<RuntimeQueueEntry>) -> Result<(), DataLayerError> {
         if entries.is_empty() {
             return Ok(());
         }
@@ -189,7 +189,7 @@ impl UsageQueueWorker {
         Ok(())
     }
 
-    async fn process_entry(&self, entry: &RedisStreamEntry) -> Result<bool, DataLayerError> {
+    async fn process_entry(&self, entry: &RuntimeQueueEntry) -> Result<bool, DataLayerError> {
         let event = match UsageEvent::from_stream_fields(&entry.fields) {
             Ok(event) => event,
             Err(err) => {
@@ -204,7 +204,7 @@ impl UsageQueueWorker {
 }
 
 pub fn build_usage_queue_worker<T>(
-    runner: RedisStreamRunner,
+    runner: Arc<dyn RuntimeQueueStore>,
     data: Arc<T>,
     config: UsageRuntimeConfig,
 ) -> Result<UsageQueueWorker, DataLayerError>