From afeec3ddc83832d9b7956f5c7dad5490e601e51e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20Bajto=C5=A1?= <oss@bajtos.net>
Date: Thu, 13 Nov 2025 10:24:03 -0300
Subject: [PATCH] fix: charge bots for traffic
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This partially reverts e664eb6.

Signed-off-by: Miroslav Bajtoš <oss@bajtos.net>
---
 piece-retriever/bin/piece-retriever.js |  1 -
 piece-retriever/lib/store.js           | 10 ++-----
 piece-retriever/test/retriever.test.js | 31 +++++++++++++++++++
 piece-retriever/test/store.test.js     | 41 --------------------------
 4 files changed, 33 insertions(+), 50 deletions(-)

diff --git a/piece-retriever/bin/piece-retriever.js b/piece-retriever/bin/piece-retriever.js
index 423a2748..3531363d 100644
--- a/piece-retriever/bin/piece-retriever.js
+++ b/piece-retriever/bin/piece-retriever.js
@@ -230,7 +230,6 @@ export default {
             egressBytes,
             cacheMiss: retrievalResult.cacheMiss,
             enforceEgressQuota: env.ENFORCE_EGRESS_QUOTA,
-            isBotTraffic: !!botName,
           })
         })(),
       )
diff --git a/piece-retriever/lib/store.js b/piece-retriever/lib/store.js
index 4f6377e3..70692a3d 100644
--- a/piece-retriever/lib/store.js
+++ b/piece-retriever/lib/store.js
@@ -253,13 +253,7 @@ export async function getRetrievalCandidatesAndValidatePayer(
  */
 export async function updateDataSetStats(
   env,
-  {
-    dataSetId,
-    egressBytes,
-    cacheMiss,
-    enforceEgressQuota = false,
-    isBotTraffic = false,
-  },
+  { dataSetId, egressBytes, cacheMiss, enforceEgressQuota = false },
 ) {
   await env.DB.prepare(
     `
@@ -271,7 +265,7 @@ export async function updateDataSetStats(
     .bind(egressBytes, dataSetId)
     .run()
 
-  if (enforceEgressQuota && !isBotTraffic) {
+  if (enforceEgressQuota) {
     await env.DB.prepare(
       `
       UPDATE data_set_egress_quotas
diff --git a/piece-retriever/test/retriever.test.js b/piece-retriever/test/retriever.test.js
index 19ccd5c6..d1e5ab96 100644
--- a/piece-retriever/test/retriever.test.js
+++ b/piece-retriever/test/retriever.test.js
@@ -462,6 +462,37 @@ describe('piece-retriever.fetch', () => {
     },
   )
 
+  it('charges bots for egress', async () => {
+    const botToken = Object.keys(botTokens)[0]
+    /** @type {string} */
+    const botName = env.BOT_TOKENS[botToken]
+    console.log({ botToken, botName })
+
+    const mockRetrieveFile = vi.fn().mockResolvedValue({
+      response: new Response('fake'),
+      cacheMiss: true,
+    })
+    const ctx = createExecutionContext()
+    const req = withRequest(defaultPayerAddress, realPieceCid, 'GET', {
+      authorization: `Bearer ${botToken}`,
+    })
+    const res = await worker.fetch(req, env, ctx, {
+      retrieveFile: mockRetrieveFile,
+    })
+    await waitOnExecutionContext(ctx)
+    expect(res.status).toBe(200)
+    const readOutput = await env.DB.prepare(
+      'SELECT egress_bytes FROM retrieval_logs WHERE data_set_id = ?',
+    )
+      .bind(String(realDataSetId))
+      .all()
+    expect(readOutput.results).toStrictEqual([
+      expect.objectContaining({
+        egress_bytes: 4,
+      }),
+    ])
+  })
+
   it('requests payment if withCDN=false', async () => {
     const dataSetId = 'test-data-set-no-cdn'
     const pieceId = 'root-no-cdn'
diff --git a/piece-retriever/test/store.test.js b/piece-retriever/test/store.test.js
index 9b5c545c..2be2474a 100644
--- a/piece-retriever/test/store.test.js
+++ b/piece-retriever/test/store.test.js
@@ -861,45 +861,4 @@ describe('updateDataSetStats', () => {
       initialCacheMissQuota - EGRESS_BYTES,
     )
   })
-
-  it('does not decrement quotas when isBotTraffic is true', async () => {
-    const DATA_SET_ID = 'test-data-set-bot-traffic'
-    const EGRESS_BYTES = 100
-    const initialCdnQuota = 500
-    const initialCacheMissQuota = 300
-
-    await withDataSet(env, {
-      dataSetId: DATA_SET_ID,
-      cdnEgressQuota: initialCdnQuota,
-      cacheMissEgressQuota: initialCacheMissQuota,
-    })
-
-    await updateDataSetStats(env, {
-      dataSetId: DATA_SET_ID,
-      egressBytes: EGRESS_BYTES,
-      cacheMiss: false,
-      enforceEgressQuota: true,
-      isBotTraffic: true,
-    })
-
-    const quotaResult = await env.DB.prepare(
-      'SELECT * FROM data_set_egress_quotas WHERE data_set_id = ?',
-    )
-      .bind(DATA_SET_ID)
-      .first()
-
-    expect(quotaResult).toStrictEqual({
-      data_set_id: DATA_SET_ID,
-      cdn_egress_quota: initialCdnQuota,
-      cache_miss_egress_quota: initialCacheMissQuota,
-    })
-
-    const dataSetResult = await env.DB.prepare(
-      'SELECT total_egress_bytes_used FROM data_sets WHERE id = ?',
-    )
-      .bind(DATA_SET_ID)
-      .first()
-
-    expect(dataSetResult.total_egress_bytes_used).toBe(EGRESS_BYTES)
-  })
 })