Mark sync usage active earlier

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

@@ -2701,6 +2701,21 @@ fn persisted_usage_body_ref(
         })
 }
 
+fn merge_usage_status_code(
+    existing: Option<&StoredRequestUsageAudit>,
+    incoming_status: &str,
+    incoming_status_code: Option<u16>,
+) -> Option<u16> {
+    if existing.is_some_and(|existing| {
+        existing.status == "streaming"
+            && incoming_status == "streaming"
+            && incoming_status_code.is_none()
+    }) {
+        return existing.and_then(|existing| existing.status_code);
+    }
+    incoming_status_code
+}
+
 #[async_trait]
 impl UsageWriteRepository for InMemoryUsageReadRepository {
     async fn upsert(
@@ -2867,7 +2882,11 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
                     .map(|existing| existing.actual_total_cost_usd)
                     .unwrap_or_default()
             }),
-            status_code: usage.status_code,
+            status_code: merge_usage_status_code(
+                existing.as_ref(),
+                usage.status.as_str(),
+                usage.status_code,
+            ),
             error_message: usage.error_message,
             error_category: usage.error_category,
             response_time_ms: merge_usage_timing(
@@ -3939,7 +3958,7 @@ mod tests {
         let mut refresh = sample_upsert_usage_record("req-streaming-refresh");
         refresh.is_stream = Some(true);
         refresh.status = "streaming".to_string();
-        refresh.status_code = Some(200);
+        refresh.status_code = None;
         repository
             .upsert(refresh)
             .await
@@ -3951,6 +3970,7 @@ mod tests {
             .expect("usage lookup should succeed")
             .expect("usage should exist");
         assert_eq!(stored.status, "streaming");
+        assert_eq!(stored.status_code, Some(200));
         assert_eq!(stored.response_time_ms, Some(45));
         assert_eq!(stored.first_byte_time_ms, Some(12));
     }

crates/aether-data/src/repository/usage/mysql.rs

@@ -202,6 +202,7 @@ ON DUPLICATE KEY UPDATE
   status_code = CASE
     WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN status_code
     WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN status_code
+    WHEN status = 'streaming' AND VALUES(status) = 'streaming' AND VALUES(status_code) IS NULL THEN status_code
     ELSE VALUES(status_code)
   END,
   error_message = CASE
@@ -1625,6 +1626,9 @@ mod tests {
         assert!(super::UPSERT_USAGE_SQL.contains("updated_at_unix_secs = CASE"));
         assert!(super::UPSERT_USAGE_SQL
             .contains("WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN status"));
+        assert!(super::UPSERT_USAGE_SQL.contains(
+            "WHEN status = 'streaming' AND VALUES(status) = 'streaming' AND VALUES(status_code) IS NULL THEN status_code"
+        ));
     }
 
     #[tokio::test]

crates/aether-data/src/repository/usage/postgres/queries/upsert_sql.sql

@@ -190,6 +190,7 @@ DO UPDATE SET
   status_code = CASE WHEN "usage".billing_status = 'pending' THEN CASE
     WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND EXCLUDED.status IN ('pending', 'streaming') THEN "usage".status_code
     WHEN "usage".status = 'streaming' AND EXCLUDED.status = 'pending' THEN "usage".status_code
+    WHEN "usage".status = 'streaming' AND EXCLUDED.status = 'streaming' AND EXCLUDED.status_code IS NULL THEN "usage".status_code
     WHEN EXCLUDED.status IN ('pending', 'streaming', 'completed', 'cancelled') AND EXCLUDED.status_code IS NULL THEN NULL
     ELSE COALESCE(EXCLUDED.status_code, "usage".status_code)
   END ELSE "usage".status_code END,

crates/aether-data/src/repository/usage/postgres/tests.rs

@@ -981,6 +981,9 @@ fn usage_sql_does_not_allow_streaming_to_regress_back_to_pending() {
     assert!(super::UPSERT_SQL.contains(
             "WHEN \"usage\".status = 'streaming' AND EXCLUDED.status = 'pending' THEN \"usage\".status_code"
         ));
+    assert!(super::UPSERT_SQL.contains(
+            "WHEN \"usage\".status = 'streaming' AND EXCLUDED.status = 'streaming' AND EXCLUDED.status_code IS NULL THEN \"usage\".status_code"
+        ));
     assert!(super::UPSERT_SQL.contains(
             "WHEN \"usage\".status = 'streaming' AND EXCLUDED.status = 'pending' THEN \"usage\".error_message"
         ));

crates/aether-data/src/repository/usage/sqlite.rs

@@ -223,6 +223,7 @@ ON CONFLICT (request_id) DO UPDATE SET
     status_code = CASE
         WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".status_code
         WHEN "usage".status = 'streaming' AND excluded.status = 'pending' THEN "usage".status_code
+        WHEN "usage".status = 'streaming' AND excluded.status = 'streaming' AND excluded.status_code IS NULL THEN "usage".status_code
         ELSE excluded.status_code
     END,
     error_message = CASE
@@ -4612,6 +4613,45 @@ mod tests {
         assert_eq!(current.updated_at_unix_secs, 1_000);
     }
 
+    #[tokio::test]
+    async fn sqlite_usage_write_repository_preserves_streaming_response_start_from_late_active() {
+        let pool = sqlx::sqlite::SqlitePoolOptions::new()
+            .max_connections(1)
+            .connect("sqlite::memory:")
+            .await
+            .expect("sqlite pool should connect");
+        run_sqlite_migrations(&pool)
+            .await
+            .expect("sqlite migrations should run");
+        seed_stats_targets(&pool).await;
+
+        let repository = SqliteUsageWriteRepository::new(pool);
+        repository
+            .upsert(sample_usage(
+                "request-late-active",
+                "streaming",
+                "pending",
+                1_000,
+            ))
+            .await
+            .expect("response-start usage should upsert");
+
+        let mut late_active = sample_usage("request-late-active", "streaming", "pending", 1_001);
+        late_active.status_code = None;
+        late_active.response_time_ms = None;
+        late_active.first_byte_time_ms = None;
+
+        let current = repository
+            .upsert(late_active)
+            .await
+            .expect("late active usage should not clear response-start fields");
+
+        assert_eq!(current.status, "streaming");
+        assert_eq!(current.status_code, Some(200));
+        assert_eq!(current.response_time_ms, Some(42));
+        assert_eq!(current.first_byte_time_ms, Some(12));
+    }
+
     #[tokio::test]
     async fn sqlite_usage_write_repository_cleans_stale_pending_requests() {
         let pool = sqlx::sqlite::SqlitePoolOptions::new()