Merge pull request #370 from yao177/aether-rust-pioneer

fix(proxy): record manual proxy node traffic

apps/aether-gateway/src/data/state/auth.rs

@@ -2,11 +2,11 @@ use super::{
     AuthApiKeyLookupKey, CreateManagementTokenRecord, DataLayerError, GatewayAuthApiKeySnapshot,
     GatewayDataState, ManagementTokenListQuery, ProxyNodeHeartbeatMutation,
     ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeRegistrationMutation,
-    ProxyNodeRemoteConfigMutation, ProxyNodeTunnelStatusMutation, RegenerateManagementTokenSecret,
+    ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation,
-    StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot, StoredLdapModuleConfig,
+    RegenerateManagementTokenSecret, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
-    StoredManagementToken, StoredManagementTokenListPage, StoredManagementTokenWithUser,
+    StoredLdapModuleConfig, StoredManagementToken, StoredManagementTokenListPage,
-    StoredOAuthProviderConfig, StoredOAuthProviderModuleConfig, StoredProxyNode,
+    StoredManagementTokenWithUser, StoredOAuthProviderConfig, StoredOAuthProviderModuleConfig,
-    StoredProxyNodeEvent, StoredUserAuthRecord, StoredUserPreferenceRecord,
+    StoredProxyNode, StoredProxyNodeEvent, StoredUserAuthRecord, StoredUserPreferenceRecord,
     StoredUserSessionRecord, StoredWalletSnapshot, UpdateManagementTokenRecord,
     UpsertOAuthProviderConfigRecord,
 };
@@ -2310,6 +2310,16 @@ impl GatewayDataState {
         }
     }
 
+    pub(crate) async fn record_proxy_node_traffic(
+        &self,
+        mutation: &ProxyNodeTrafficMutation,
+    ) -> Result<bool, DataLayerError> {
+        match &self.proxy_node_writer {
+            Some(repository) => repository.record_traffic(mutation).await,
+            None => Ok(false),
+        }
+    }
+
     pub(crate) async fn update_proxy_node_tunnel_status(
         &self,
         mutation: &ProxyNodeTunnelStatusMutation,

apps/aether-gateway/src/data/state/mod.rs

@@ -42,7 +42,8 @@ use aether_data::repository::oauth_providers::{
 use aether_data::repository::proxy_nodes::{
     ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation,
     ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
-    ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
+    ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
+    StoredProxyNode, StoredProxyNodeEvent,
 };
 pub(crate) use aether_data::repository::system::{AdminSystemStats, StoredSystemConfigEntry};
 use aether_data::repository::users::{

apps/aether-gateway/src/execution_runtime/stream/execution.rs

@@ -60,8 +60,9 @@ use crate::execution_runtime::submission::{
     submit_local_core_error_or_sync_finalize,
 };
 use crate::execution_runtime::transport::{
-    execute_stream_plan_via_local_tunnel, DirectSyncExecutionRuntime,
+    execute_stream_plan_via_local_tunnel, record_manual_proxy_request_failure,
-    DirectUpstreamStreamExecution, ExecutionRuntimeTransportError,
+    record_manual_proxy_request_success, record_manual_proxy_stream_error,
+    DirectSyncExecutionRuntime, DirectUpstreamStreamExecution, ExecutionRuntimeTransportError,
 };
 use crate::execution_runtime::{
     local_failover_response_text, resolve_core_stream_direct_finalize_report_kind,
@@ -314,7 +315,16 @@ async fn execute_in_process_stream(
         return Ok(execution);
     }
 
-    DirectSyncExecutionRuntime::new().execute_stream(plan).await
+    match DirectSyncExecutionRuntime::new().execute_stream(plan).await {
+        Ok(execution) => {
+            record_manual_proxy_request_success(state, plan).await;
+            Ok(execution)
+        }
+        Err(error) => {
+            record_manual_proxy_request_failure(state, plan).await;
+            Err(error)
+        }
+    }
 }
 
 async fn execute_in_process_stream_with_oauth_retry(
@@ -2119,6 +2129,7 @@ async fn execute_stream_from_frame_stream(
         }
 
         if let Some(failure) = terminal_failure {
+            record_manual_proxy_stream_error(&state_for_report, &plan_for_report).await;
             submit_midstream_stream_failure(
                 &state_for_report,
                 &trace_id_owned,

apps/aether-gateway/src/execution_runtime/transport.rs

@@ -8,6 +8,7 @@ use aether_contracts::{
     ExecutionPlan, ExecutionResult, ExecutionTelemetry, ProxySnapshot, ResponseBody,
     EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER,
 };
+use aether_data::repository::proxy_nodes::ProxyNodeTrafficMutation;
 use aether_http::{apply_http_client_config, HttpClientConfig};
 use axum::body::Bytes;
 use base64::Engine as _;
@@ -267,12 +268,17 @@ pub(crate) async fn execute_sync_plan(
             .map_err(|err| GatewayError::Internal(err.to_string()));
     }
 
-    let _ = state;
     let _ = trace_id;
-    DirectSyncExecutionRuntime::new()
+    match DirectSyncExecutionRuntime::new().execute_sync(plan).await {
-        .execute_sync(plan)
+        Ok(result) => {
-        .await
+            record_manual_proxy_request_outcome(state, plan, result.status_code).await;
-        .map_err(|err| GatewayError::Internal(err.to_string()))
+            Ok(result)
+        }
+        Err(err) => {
+            record_manual_proxy_request_failure(state, plan).await;
+            Err(GatewayError::Internal(err.to_string()))
+        }
+    }
 }
 
 pub(crate) async fn execute_stream_plan_via_local_tunnel(
@@ -327,6 +333,68 @@ fn build_stream_summary_report_context(plan: &ExecutionPlan) -> Value {
     })
 }
 
+pub(crate) async fn record_manual_proxy_request_success(state: &AppState, plan: &ExecutionPlan) {
+    record_manual_proxy_traffic(state, plan, 1, 0, 0, 0).await;
+}
+
+pub(crate) async fn record_manual_proxy_request_outcome(
+    state: &AppState,
+    plan: &ExecutionPlan,
+    status_code: u16,
+) {
+    let failed_requests_delta = i64::from(status_code >= 400);
+    record_manual_proxy_traffic(state, plan, 1, failed_requests_delta, 0, 0).await;
+}
+
+pub(crate) async fn record_manual_proxy_request_failure(state: &AppState, plan: &ExecutionPlan) {
+    record_manual_proxy_traffic(state, plan, 1, 1, 0, 0).await;
+}
+
+pub(crate) async fn record_manual_proxy_stream_error(state: &AppState, plan: &ExecutionPlan) {
+    record_manual_proxy_traffic(state, plan, 0, 0, 0, 1).await;
+}
+
+async fn record_manual_proxy_traffic(
+    state: &AppState,
+    plan: &ExecutionPlan,
+    total_requests_delta: i64,
+    failed_requests_delta: i64,
+    dns_failures_delta: i64,
+    stream_errors_delta: i64,
+) {
+    let Some(node_id) = manual_proxy_node_id(plan.proxy.as_ref()) else {
+        return;
+    };
+    let mutation = ProxyNodeTrafficMutation {
+        node_id: node_id.clone(),
+        total_requests_delta,
+        failed_requests_delta,
+        dns_failures_delta,
+        stream_errors_delta,
+    };
+
+    if let Err(error) = state.record_proxy_node_traffic(&mutation).await {
+        tracing::warn!(
+            node_id = %node_id,
+            error = ?error,
+            "failed to record manual proxy node traffic"
+        );
+    }
+}
+
+fn manual_proxy_node_id(proxy: Option<&ProxySnapshot>) -> Option<String> {
+    let proxy = proxy?;
+    if proxy.enabled == Some(false) || resolve_tunnel_node_id(Some(proxy)).is_some() {
+        return None;
+    }
+    proxy
+        .node_id
+        .as_deref()
+        .map(str::trim)
+        .filter(|value| !value.is_empty())
+        .map(ToOwned::to_owned)
+}
+
 async fn execute_sync_plan_via_local_tunnel(
     state: &AppState,
     plan: &ExecutionPlan,
@@ -1068,8 +1136,11 @@ mod tests {
     use std::sync::Arc;
 
     use aether_contracts::{
-        ExecutionPlan, ExecutionTimeouts, RequestBody, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
+        ExecutionPlan, ExecutionTimeouts, ProxySnapshot, RequestBody,
-        EXECUTION_REQUEST_HTTP1_ONLY_HEADER,
+        EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER,
+    };
+    use aether_data::repository::proxy_nodes::{
+        InMemoryProxyNodeRepository, ProxyNodeReadRepository, StoredProxyNode,
     };
     use axum::body::Bytes;
     use axum::extract::ws::Message;
@@ -1081,7 +1152,10 @@ mod tests {
     use tokio::sync::watch;
 
     use super::{
-        build_client, execute_sync_plan, DirectSyncExecutionRuntime, ExecutionTransportControls,
+        build_client, execute_sync_plan, record_manual_proxy_request_failure,
+        record_manual_proxy_request_outcome, record_manual_proxy_request_success,
+        record_manual_proxy_stream_error, DirectSyncExecutionRuntime,
+        ExecutionTransportControls,
     };
     use crate::constants::{
         EXECUTION_RUNTIME_LOOP_GUARD_HEADER, EXECUTION_RUNTIME_LOOP_GUARD_VIA_TOKEN,
@@ -1162,8 +1236,8 @@ mod tests {
         }
     }
 
-    fn tunnel_proxy_snapshot(base_url: String) -> aether_contracts::ProxySnapshot {
+    fn tunnel_proxy_snapshot(base_url: String) -> ProxySnapshot {
-        aether_contracts::ProxySnapshot {
+        ProxySnapshot {
             enabled: Some(true),
             mode: Some("tunnel".into()),
             node_id: Some("node-1".into()),
@@ -1173,6 +1247,39 @@ mod tests {
         }
     }
 
+    fn manual_proxy_snapshot(node_id: &str) -> ProxySnapshot {
+        ProxySnapshot {
+            enabled: Some(true),
+            mode: Some("http".into()),
+            node_id: Some(node_id.to_string()),
+            label: Some("manual-proxy".into()),
+            url: Some("http://127.0.0.1:1".into()),
+            extra: None,
+        }
+    }
+
+    fn sample_manual_proxy_node(node_id: &str) -> StoredProxyNode {
+        StoredProxyNode::new(
+            node_id.to_string(),
+            "manual-proxy".to_string(),
+            "127.0.0.1".to_string(),
+            1,
+            true,
+            "online".to_string(),
+            0,
+            0,
+            0,
+            0,
+            0,
+            0,
+            false,
+            false,
+            0,
+        )
+        .expect("manual proxy node should build")
+        .with_manual_proxy_fields(Some("http://127.0.0.1:1".into()), None, None)
+    }
+
     fn decode_relay_envelope(body: &[u8]) -> (serde_json::Value, Vec<u8>) {
         assert!(
             body.len() >= 4,
@@ -1259,6 +1366,273 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn execute_sync_plan_records_manual_proxy_success() {
+        let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
+            sample_manual_proxy_node("manual-node-1"),
+        ]));
+        let data = crate::data::GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(
+            &repository,
+        ));
+        let state = AppState::new()
+            .expect("app state should build")
+            .with_data_state_for_tests(data);
+        let plan = ExecutionPlan {
+            request_id: "req-manual-proxy-success".into(),
+            candidate_id: None,
+            provider_name: None,
+            provider_id: "prov-1".into(),
+            endpoint_id: "ep-1".into(),
+            key_id: "key-1".into(),
+            method: "POST".into(),
+            url: "https://example.com/chat".into(),
+            headers: BTreeMap::new(),
+            content_type: None,
+            content_encoding: None,
+            body: RequestBody::from_json(json!({})),
+            stream: false,
+            client_api_format: "openai:chat".into(),
+            provider_api_format: "openai:chat".into(),
+            model_name: None,
+            proxy: Some(manual_proxy_snapshot("manual-node-1")),
+            tls_profile: None,
+            timeouts: None,
+        };
+
+        record_manual_proxy_request_success(&state, &plan).await;
+
+        let node = repository
+            .find_proxy_node("manual-node-1")
+            .await
+            .expect("proxy node lookup should succeed")
+            .expect("manual proxy node should exist");
+        assert_eq!(node.total_requests, 1);
+        assert_eq!(node.failed_requests, 0);
+    }
+
+    #[tokio::test]
+    async fn execute_sync_plan_records_manual_proxy_failure() {
+        let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
+            sample_manual_proxy_node("manual-node-1"),
+        ]));
+        let data = crate::data::GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(
+            &repository,
+        ));
+        let state = AppState::new()
+            .expect("app state should build")
+            .with_data_state_for_tests(data);
+        let plan = ExecutionPlan {
+            request_id: "req-manual-proxy-failure".into(),
+            candidate_id: None,
+            provider_name: None,
+            provider_id: "prov-1".into(),
+            endpoint_id: "ep-1".into(),
+            key_id: "key-1".into(),
+            method: "POST".into(),
+            url: "https://example.com/chat".into(),
+            headers: BTreeMap::new(),
+            content_type: None,
+            content_encoding: None,
+            body: RequestBody::from_json(json!({})),
+            stream: false,
+            client_api_format: "openai:chat".into(),
+            provider_api_format: "openai:chat".into(),
+            model_name: None,
+            proxy: Some(manual_proxy_snapshot("manual-node-1")),
+            tls_profile: None,
+            timeouts: None,
+        };
+
+        record_manual_proxy_request_failure(&state, &plan).await;
+
+        let node = repository
+            .find_proxy_node("manual-node-1")
+            .await
+            .expect("proxy node lookup should succeed")
+            .expect("manual proxy node should exist");
+        assert_eq!(node.total_requests, 1);
+        assert_eq!(node.failed_requests, 1);
+    }
+
+    #[tokio::test]
+    async fn execute_sync_plan_records_manual_proxy_http_error_as_failure() {
+        let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
+            sample_manual_proxy_node("manual-node-1"),
+        ]));
+        let data = crate::data::GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(
+            &repository,
+        ));
+        let state = AppState::new()
+            .expect("app state should build")
+            .with_data_state_for_tests(data);
+        let plan = ExecutionPlan {
+            request_id: "req-manual-proxy-http-error".into(),
+            candidate_id: None,
+            provider_name: None,
+            provider_id: "prov-1".into(),
+            endpoint_id: "ep-1".into(),
+            key_id: "key-1".into(),
+            method: "POST".into(),
+            url: "https://example.com/chat".into(),
+            headers: BTreeMap::new(),
+            content_type: None,
+            content_encoding: None,
+            body: RequestBody::from_json(json!({})),
+            stream: false,
+            client_api_format: "openai:chat".into(),
+            provider_api_format: "openai:chat".into(),
+            model_name: None,
+            proxy: Some(manual_proxy_snapshot("manual-node-1")),
+            tls_profile: None,
+            timeouts: None,
+        };
+
+        record_manual_proxy_request_outcome(&state, &plan, 429).await;
+
+        let node = repository
+            .find_proxy_node("manual-node-1")
+            .await
+            .expect("proxy node lookup should succeed")
+            .expect("manual proxy node should exist");
+        assert_eq!(node.total_requests, 1);
+        assert_eq!(node.failed_requests, 1);
+    }
+
+    #[tokio::test]
+    async fn execute_sync_plan_records_manual_proxy_http_success_without_failure() {
+        let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
+            sample_manual_proxy_node("manual-node-1"),
+        ]));
+        let data = crate::data::GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(
+            &repository,
+        ));
+        let state = AppState::new()
+            .expect("app state should build")
+            .with_data_state_for_tests(data);
+        let plan = ExecutionPlan {
+            request_id: "req-manual-proxy-http-success".into(),
+            candidate_id: None,
+            provider_name: None,
+            provider_id: "prov-1".into(),
+            endpoint_id: "ep-1".into(),
+            key_id: "key-1".into(),
+            method: "POST".into(),
+            url: "https://example.com/chat".into(),
+            headers: BTreeMap::new(),
+            content_type: None,
+            content_encoding: None,
+            body: RequestBody::from_json(json!({})),
+            stream: false,
+            client_api_format: "openai:chat".into(),
+            provider_api_format: "openai:chat".into(),
+            model_name: None,
+            proxy: Some(manual_proxy_snapshot("manual-node-1")),
+            tls_profile: None,
+            timeouts: None,
+        };
+
+        record_manual_proxy_request_outcome(&state, &plan, 200).await;
+
+        let node = repository
+            .find_proxy_node("manual-node-1")
+            .await
+            .expect("proxy node lookup should succeed")
+            .expect("manual proxy node should exist");
+        assert_eq!(node.total_requests, 1);
+        assert_eq!(node.failed_requests, 0);
+    }
+
+    #[tokio::test]
+    async fn execute_sync_plan_records_manual_proxy_stream_error_without_extra_request_count() {
+        let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
+            sample_manual_proxy_node("manual-node-1"),
+        ]));
+        let data = crate::data::GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(
+            &repository,
+        ));
+        let state = AppState::new()
+            .expect("app state should build")
+            .with_data_state_for_tests(data);
+        let plan = ExecutionPlan {
+            request_id: "req-manual-proxy-stream-error".into(),
+            candidate_id: None,
+            provider_name: None,
+            provider_id: "prov-1".into(),
+            endpoint_id: "ep-1".into(),
+            key_id: "key-1".into(),
+            method: "POST".into(),
+            url: "https://example.com/chat".into(),
+            headers: BTreeMap::new(),
+            content_type: None,
+            content_encoding: None,
+            body: RequestBody::from_json(json!({})),
+            stream: true,
+            client_api_format: "openai:chat".into(),
+            provider_api_format: "openai:chat".into(),
+            model_name: None,
+            proxy: Some(manual_proxy_snapshot("manual-node-1")),
+            tls_profile: None,
+            timeouts: None,
+        };
+
+        record_manual_proxy_request_success(&state, &plan).await;
+        record_manual_proxy_stream_error(&state, &plan).await;
+
+        let node = repository
+            .find_proxy_node("manual-node-1")
+            .await
+            .expect("proxy node lookup should succeed")
+            .expect("manual proxy node should exist");
+        assert_eq!(node.total_requests, 1);
+        assert_eq!(node.failed_requests, 0);
+        assert_eq!(node.stream_errors, 1);
+    }
+
+    #[tokio::test]
+    async fn execute_sync_plan_ignores_stream_error_for_tunnel_proxy() {
+        let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
+            sample_manual_proxy_node("manual-node-1"),
+        ]));
+        let data = crate::data::GatewayDataState::with_proxy_node_repository_for_tests(Arc::clone(
+            &repository,
+        ));
+        let state = AppState::new()
+            .expect("app state should build")
+            .with_data_state_for_tests(data);
+        let plan = ExecutionPlan {
+            request_id: "req-tunnel-proxy-stream-error".into(),
+            candidate_id: None,
+            provider_name: None,
+            provider_id: "prov-1".into(),
+            endpoint_id: "ep-1".into(),
+            key_id: "key-1".into(),
+            method: "POST".into(),
+            url: "https://example.com/chat".into(),
+            headers: BTreeMap::new(),
+            content_type: None,
+            content_encoding: None,
+            body: RequestBody::from_json(json!({})),
+            stream: true,
+            client_api_format: "openai:chat".into(),
+            provider_api_format: "openai:chat".into(),
+            model_name: None,
+            proxy: Some(tunnel_proxy_snapshot("http://127.0.0.1:1".to_string())),
+            tls_profile: None,
+            timeouts: None,
+        };
+
+        record_manual_proxy_stream_error(&state, &plan).await;
+
+        let node = repository
+            .find_proxy_node("manual-node-1")
+            .await
+            .expect("proxy node lookup should succeed")
+            .expect("manual proxy node should exist");
+        assert_eq!(node.total_requests, 0);
+        assert_eq!(node.failed_requests, 0);
+        assert_eq!(node.stream_errors, 0);
+    }
+
     #[tokio::test]
     async fn direct_sync_execution_runtime_supports_tunnel_relay() {
         let listener = crate::test_support::bind_loopback_listener()

apps/aether-gateway/src/state/core.rs

@@ -5,7 +5,7 @@ use std::time::Duration;
 
 use aether_data::repository::proxy_nodes::{
     ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation,
-    ProxyNodeTunnelStatusMutation, StoredProxyNode, StoredProxyNodeEvent,
+    ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, StoredProxyNode, StoredProxyNodeEvent,
 };
 use aether_http::{build_http_client, HttpClientConfig};
 use aether_runtime::{
@@ -551,6 +551,16 @@ impl AppState {
             .map_err(|err| GatewayError::Internal(err.to_string()))
     }
 
+    pub(crate) async fn record_proxy_node_traffic(
+        &self,
+        mutation: &ProxyNodeTrafficMutation,
+    ) -> Result<bool, GatewayError> {
+        self.data
+            .record_proxy_node_traffic(mutation)
+            .await
+            .map_err(|err| GatewayError::Internal(err.to_string()))
+    }
+
     pub(crate) async fn unregister_proxy_node(
         &self,
         node_id: &str,

crates/aether-data/src/repository/proxy_nodes/memory.rs

@@ -9,8 +9,8 @@ use uuid::Uuid;
 use super::types::{
     normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeHeartbeatMutation,
     ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeReadRepository,
-    ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTunnelStatusMutation,
+    ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation,
-    ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
+    ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
 };
 use crate::DataLayerError;
 
@@ -437,6 +437,26 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
         Ok(Some(node.clone()))
     }
 
+    async fn record_traffic(
+        &self,
+        mutation: &ProxyNodeTrafficMutation,
+    ) -> Result<bool, DataLayerError> {
+        let mut nodes = self.nodes.write().expect("proxy node repository lock");
+        let Some(node) = nodes.get_mut(&mutation.node_id) else {
+            return Ok(false);
+        };
+        if !node.is_manual {
+            return Ok(false);
+        }
+
+        node.total_requests += mutation.total_requests_delta.max(0);
+        node.failed_requests += mutation.failed_requests_delta.max(0);
+        node.dns_failures += mutation.dns_failures_delta.max(0);
+        node.stream_errors += mutation.stream_errors_delta.max(0);
+        node.updated_at_unix_secs = Self::now_unix_secs();
+        Ok(true)
+    }
+
     async fn update_tunnel_status(
         &self,
         mutation: &ProxyNodeTunnelStatusMutation,

crates/aether-data/src/repository/proxy_nodes/mod.rs

@@ -9,7 +9,7 @@ pub use types::{
     reconcile_remote_config_after_heartbeat, remote_config_scheduling_state,
     remote_config_upgrade_target, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
     ProxyNodeManualUpdateMutation, ProxyNodeReadRepository, ProxyNodeRegistrationMutation,
-    ProxyNodeRemoteConfigMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
+    ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation,
-    StoredProxyNode, StoredProxyNodeEvent, PROXY_NODE_SCHEDULING_STATE_CORDONED,
+    ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
-    PROXY_NODE_SCHEDULING_STATE_DRAINING,
+    PROXY_NODE_SCHEDULING_STATE_CORDONED, PROXY_NODE_SCHEDULING_STATE_DRAINING,
 };

crates/aether-data/src/repository/proxy_nodes/sql.rs

@@ -6,8 +6,8 @@ use sqlx::{postgres::PgRow, PgPool, Row};
 use super::types::{
     normalize_proxy_metadata, reconcile_remote_config_after_heartbeat, ProxyNodeHeartbeatMutation,
     ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeReadRepository,
-    ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTunnelStatusMutation,
+    ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation,
-    ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
+    ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
 };
 use crate::{
     error::{postgres_error, SqlxResultExt},
@@ -320,6 +320,18 @@ WHERE id = $1
   AND is_manual = TRUE
 "#;
 
+const RECORD_PROXY_NODE_TRAFFIC_SQL: &str = r#"
+UPDATE proxy_nodes
+SET
+  total_requests = total_requests + GREATEST($2, 0),
+  failed_requests = failed_requests + GREATEST($3, 0),
+  dns_failures = dns_failures + GREATEST($4, 0),
+  stream_errors = stream_errors + GREATEST($5, 0),
+  updated_at = NOW()
+WHERE id = $1
+  AND is_manual = TRUE
+"#;
+
 const UNREGISTER_PROXY_NODE_SQL: &str = r#"
 UPDATE proxy_nodes
 SET
@@ -832,6 +844,23 @@ impl ProxyNodeWriteRepository for SqlxProxyNodeRepository {
         Ok(Some(updated))
     }
 
+    async fn record_traffic(
+        &self,
+        mutation: &ProxyNodeTrafficMutation,
+    ) -> Result<bool, DataLayerError> {
+        let result = sqlx::query(RECORD_PROXY_NODE_TRAFFIC_SQL)
+            .bind(&mutation.node_id)
+            .bind(mutation.total_requests_delta)
+            .bind(mutation.failed_requests_delta)
+            .bind(mutation.dns_failures_delta)
+            .bind(mutation.stream_errors_delta)
+            .execute(&self.pool)
+            .await
+            .map_postgres_err()?;
+
+        Ok(result.rows_affected() > 0)
+    }
+
     async fn update_tunnel_status(
         &self,
         mutation: &ProxyNodeTunnelStatusMutation,

crates/aether-data/src/repository/proxy_nodes/types.rs

@@ -162,6 +162,15 @@ pub struct ProxyNodeHeartbeatMutation {
     pub proxy_version: Option<String>,
 }
 
+#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
+pub struct ProxyNodeTrafficMutation {
+    pub node_id: String,
+    pub total_requests_delta: i64,
+    pub failed_requests_delta: i64,
+    pub dns_failures_delta: i64,
+    pub stream_errors_delta: i64,
+}
+
 #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
 pub struct ProxyNodeRegistrationMutation {
     pub name: String,
@@ -392,6 +401,11 @@ pub trait ProxyNodeWriteRepository: Send + Sync {
         mutation: &ProxyNodeHeartbeatMutation,
     ) -> Result<Option<StoredProxyNode>, crate::DataLayerError>;
 
+    async fn record_traffic(
+        &self,
+        mutation: &ProxyNodeTrafficMutation,
+    ) -> Result<bool, crate::DataLayerError>;
+
     async fn update_tunnel_status(
         &self,
         mutation: &ProxyNodeTunnelStatusMutation,