From a5ff6a1fbd8d031d349c46550157b6887ba43e0d Mon Sep 17 00:00:00 2001
From: pablo <p.lopez.lpz@gmail.com>
Date: Thu, 3 Jul 2025 12:10:15 +0300
Subject: [PATCH] feat: expose available needed by chatsdk

---
 tests/common/test_tokenbucket.nim       | 30 +++++++++++++++++++++++++
 waku/common/rate_limit/token_bucket.nim |  8 +++++++
 2 files changed, 38 insertions(+)

diff --git a/tests/common/test_tokenbucket.nim b/tests/common/test_tokenbucket.nim
index 5bc1a0583..413558e24 100644
--- a/tests/common/test_tokenbucket.nim
+++ b/tests/common/test_tokenbucket.nim
@@ -10,8 +10,13 @@
 
 import testutils/unittests
 import chronos
+import std/math
 import ../../waku/common/rate_limit/token_bucket
 
+# Helper function for approximate float equality
+proc equals(a, b: float, tolerance: float = 1e-6): bool =
+  abs(a - b) <= tolerance
+
 suite "Token Bucket":
   test "TokenBucket Sync test - strict":
     var bucket = TokenBucket.newStrict(1000, 1.milliseconds)
@@ -67,3 +72,28 @@ suite "Token Bucket":
     check bucket.tryConsume(1, start) == false
 
     check bucket.tryConsume(15000, start + 1.milliseconds) == true
+
+  test "TokenBucket getAvailableCapacityRatio":
+    var bucket = TokenBucket.new(1000, 1.minutes, ReplenishMode.Strict)
+    var reqTime = Moment.now()
+    
+    # Test full bucket capacity ratio
+    check equals(bucket.getAvailableCapacityRatio(reqTime), 1.0) # 1000/1000 = 1.0
+    
+    # Consume some tokens and check ratio
+    reqTime += 1.seconds
+    check bucket.tryConsume(400, reqTime) == true
+    check equals(bucket.getAvailableCapacityRatio(reqTime), 0.6)  # 600/1000 = 0.6
+    
+    # Consume more tokens
+    reqTime += 1.seconds
+    check bucket.tryConsume(300, reqTime) == true
+    check equals(bucket.getAvailableCapacityRatio(reqTime), 0.3)  # 300/1000 = 0.3
+    
+    # Test when period has elapsed (should return 1.0)
+    reqTime += 1.minutes
+    check equals(bucket.getAvailableCapacityRatio(reqTime), 1.0) # 1000/1000 = 1.0
+    
+    # Test with empty bucket
+    check bucket.tryConsume(1000, reqTime) == true
+    check equals(bucket.getAvailableCapacityRatio(reqTime), 0.0)  # 0/1000 = 0.0
diff --git a/waku/common/rate_limit/token_bucket.nim b/waku/common/rate_limit/token_bucket.nim
index 799817ebd..265d8847a 100644
--- a/waku/common/rate_limit/token_bucket.nim
+++ b/waku/common/rate_limit/token_bucket.nim
@@ -109,6 +109,14 @@ proc update(bucket: TokenBucket, currentTime: Moment) =
   else:
     updateStrict(bucket, currentTime)
 
+## Returns the available capacity ratio of the bucket.
+## This is a ratio of the bucket capacity that is currently available.
+## It is a float number between 0.0 (empty) and 1.0 (full).
+proc getAvailableCapacityRatio*(bucket: TokenBucket, currentTime: Moment): float =
+  if periodElapsed(bucket, currentTime):
+    return 1.0
+  return bucket.budget.float / bucket.budgetCap.float
+
 proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool =
   ## If `tokens` are available, consume them,
   ## Otherwhise, return false.