fix: use actual token usage from Kiro API instead of hardcoded 0

- Added usage tracking to AwsEventStreamParser
- Parser now stores usage_credits and context_usage_percentage
- Added get_usage() method to retrieve usage data
- Updated handler to pass usage data to response builder
- Non-streaming: Uses actual usage_credits as total_tokens
- Streaming: Includes usage_credits in final chunk
- Falls back to 0 if Kiro API doesn't provide usage data
- Prevents over-estimation when Kiro returns 0 tokens

aisbf/providers/kiro/handler.py

@@ -252,17 +252,19 @@ class KiroProviderHandler(BaseProviderHandler):
             # Extract content and tool calls
             content = parser.get_content()
             tool_calls = parser.get_tool_calls()
+            usage_data = parser.get_usage()
 
             if AISBF_DEBUG:
                 logging.info(f"KiroProviderHandler: Parsed content length: {len(content)}")
                 logging.info(f"KiroProviderHandler: Parsed tool calls: {len(tool_calls)}")
+                logging.info(f"KiroProviderHandler: Usage credits: {usage_data.get('usage_credits', 0)}")
                 if tool_calls:
                     logging.info(f"KiroProviderHandler: Tool calls: {json.dumps(tool_calls, indent=2)}")
 
             logging.info(f"KiroProviderHandler: Response parsed successfully")
 
             # Build OpenAI-format response
-            openai_response = self._build_openai_response(model, content, tool_calls)
+            openai_response = self._build_openai_response(model, content, tool_calls, usage_data)
 
             if AISBF_DEBUG:
                 logging.info(f"=== FINAL KIRO RESPONSE DICT ===")
@@ -282,9 +284,13 @@ class KiroProviderHandler(BaseProviderHandler):
             self.record_failure()
             raise e
 
-    def _build_openai_response(self, model: str, content: str, tool_calls: List[Dict]) -> Dict:
+    def _build_openai_response(self, model: str, content: str, tool_calls: List[Dict], usage_data: Dict = None) -> Dict:
         """Build OpenAI-format response from parsed Kiro data."""
         finish_reason = "tool_calls" if tool_calls else "stop"
+        
+        # Convert Kiro usage credits to approximate token count
+        # Kiro usage is in credits, roughly 1 credit ≈ 1 token
+        usage_credits = usage_data.get('usage_credits', 0) if usage_data else 0
 
         openai_response = {
             "id": f"kiro-{int(time.time())}",
@@ -302,7 +308,7 @@ class KiroProviderHandler(BaseProviderHandler):
             "usage": {
                 "prompt_tokens": 0,
                 "completion_tokens": 0,
-                "total_tokens": 0
+                "total_tokens": usage_credits  # Use actual usage from Kiro API
             }
         }
 
@@ -371,9 +377,14 @@ class KiroProviderHandler(BaseProviderHandler):
                 logger.info(f"KiroProviderHandler: Streaming completed")
 
                 final_tool_calls = parser.get_tool_calls()
+                usage_data = parser.get_usage()
                 finish_reason = "tool_calls" if final_tool_calls else "stop"
+                
+                # Convert Kiro usage credits to approximate token count
+                usage_credits = usage_data.get('usage_credits', 0)
 
                 logger.info(f"KiroProviderHandler: Final tool calls count: {len(final_tool_calls)}")
+                logger.info(f"KiroProviderHandler: Usage credits: {usage_credits}")
 
                 if final_tool_calls:
                     indexed_tool_calls = []
@@ -423,7 +434,7 @@ class KiroProviderHandler(BaseProviderHandler):
                     "usage": {
                         "prompt_tokens": 0,
                         "completion_tokens": 0,
-                        "total_tokens": 0
+                        "total_tokens": usage_credits  # Use actual usage from Kiro API
                     }
                 }
 

aisbf/providers/kiro/parsers.py

@@ -265,6 +265,8 @@ class AwsEventStreamParser:
         self.current_tool_call: Optional[Dict[str, Any]] = None
         self.tool_calls: List[Dict[str, Any]] = []
         self.content_chunks: List[str] = []
+        self.usage_credits: int = 0  # Track usage credits from Kiro API
+        self.context_usage_percentage: float = 0.0  # Track context usage percentage
     
     def feed(self, chunk: bytes) -> List[Dict[str, Any]]:
         """
@@ -336,9 +338,11 @@ class AwsEventStreamParser:
         elif event_type == 'tool_stop':
             return self._process_tool_stop_event(data)
         elif event_type == 'usage':
-            return {"type": "usage", "data": data.get('usage', 0)}
+            self.usage_credits = data.get('usage', 0)
+            return {"type": "usage", "data": self.usage_credits}
         elif event_type == 'context_usage':
-            return {"type": "context_usage", "data": data.get('contextUsagePercentage', 0)}
+            self.context_usage_percentage = data.get('contextUsagePercentage', 0)
+            return {"type": "context_usage", "data": self.context_usage_percentage}
         
         return None
     
@@ -572,6 +576,18 @@ class AwsEventStreamParser:
             self._finalize_tool_call()
         return deduplicate_tool_calls(self.tool_calls)
     
+    def get_usage(self) -> Dict[str, int]:
+        """
+        Returns usage information from Kiro API.
+        
+        Returns:
+            Dictionary with usage_credits and context_usage_percentage
+        """
+        return {
+            "usage_credits": self.usage_credits,
+            "context_usage_percentage": self.context_usage_percentage
+        }
+    
     def reset(self) -> None:
         """Resets parser state."""
         self.buffer = ""
@@ -579,3 +595,5 @@ class AwsEventStreamParser:
         self.current_tool_call = None
         self.tool_calls = []
         self.content_chunks = []
+        self.usage_credits = 0
+        self.context_usage_percentage = 0.0