Merge pull request #42 from softwaremill/fix/duplicate-token-accumulation

Fix/duplicate token accumulation

Author: softberries

crates/tracevault-server/src/repo/events.rs

@@ -92,11 +92,12 @@ impl EventRepo {
     }
 
     /// INSERT INTO transcript_chunks ... ON CONFLICT DO NOTHING.
+    /// Returns true if the row was actually inserted (not a duplicate).
     pub async fn insert_transcript_chunk(
         pool: &PgPool,
         req: &InsertTranscriptChunk,
-    ) -> Result<(), AppError> {
-        sqlx::query(
+    ) -> Result<bool, AppError> {
+        let result = sqlx::query(
             "INSERT INTO transcript_chunks (session_id, chunk_index, data)
              VALUES ($1, $2, $3)
              ON CONFLICT (session_id, chunk_index) DO NOTHING",
@@ -106,7 +107,7 @@ impl EventRepo {
         .bind(&req.data)
         .execute(pool)
         .await?;
-        Ok(())
+        Ok(result.rows_affected() > 0)
     }
 
     /// INSERT INTO user_software_usage ... ON CONFLICT DO UPDATE.

crates/tracevault-server/src/service/stream.rs

@@ -67,7 +67,7 @@ impl StreamService {
 
                 for (i, line) in lines.iter().enumerate() {
                     let chunk_index = offset as i32 + i as i32;
-                    EventRepo::insert_transcript_chunk(
+                    let was_inserted = EventRepo::insert_transcript_chunk(
                         &state.pool,
                         &InsertTranscriptChunk {
                             session_id: session_db_id,
@@ -77,6 +77,12 @@ impl StreamService {
                     )
                     .await?;
 
+                    // Only count tokens from newly inserted chunks to avoid
+                    // double-counting when overlapping batches are sent
+                    if !was_inserted {
+                        continue;
+                    }
+
                     // Extract token usage from assistant messages
                     if let Some(msg) = line.get("message") {
                         if let Some(usage) = msg.get("usage") {
@@ -114,8 +120,11 @@ impl StreamService {
                     // input_tokens from the API includes cache_read and cache_write,
                     // subtract to get fresh (non-cached) input only
                     let fresh_input = (batch_input - batch_cache_read - batch_cache_write).max(0);
-                    let batch_cost = crate::pricing::estimate_cost(
-                        model_name,
+                    let pricing =
+                        crate::pricing::fetch_pricing_for_model(&state.pool, model_name, None)
+                            .await;
+                    let batch_cost = crate::pricing::estimate_cost_with_pricing(
+                        &pricing,
                         fresh_input,
                         batch_output,
                         batch_cache_read,

crates/tracevault-server/tests/repo_events_test.rs

@@ -127,14 +127,16 @@ async fn insert_transcript_chunk_dedup(pool: sqlx::PgPool) {
         data: serde_json::json!({"role": "user", "content": "hello"}),
     };
 
-    EventRepo::insert_transcript_chunk(&pool, &chunk)
+    let inserted = EventRepo::insert_transcript_chunk(&pool, &chunk)
         .await
         .unwrap();
+    assert!(inserted, "first insert should return true");
 
-    // Duplicate should not error (ON CONFLICT DO NOTHING)
-    EventRepo::insert_transcript_chunk(&pool, &chunk)
+    // Duplicate should not error (ON CONFLICT DO NOTHING) and should return false
+    let inserted_again = EventRepo::insert_transcript_chunk(&pool, &chunk)
         .await
         .unwrap();
+    assert!(!inserted_again, "duplicate insert should return false");
 
     let (count,): (i64,) =
         sqlx::query_as("SELECT COUNT(*) FROM transcript_chunks WHERE session_id = $1")