diff --git a/services/gitauto_handler.py b/services/gitauto_handler.py index c33cf85f..b51effae 100644 --- a/services/gitauto_handler.py +++ b/services/gitauto_handler.py @@ -3,6 +3,7 @@ import time from uuid import uuid4 + # Local imports from config import ( PRODUCT_ID, @@ -39,6 +40,8 @@ async def handle_gitauto(payload: GitHubLabeledPayload, trigger_type: str) -> None: """Core functionality to create comments on issue, create PRs, and update progress.""" + current_time = time.time() + # Extract label and validate it if trigger_type == "label" and payload["label"]["name"] != PRODUCT_ID: return @@ -223,9 +226,11 @@ async def handle_gitauto(payload: GitHubLabeledPayload, trigger_type: str) -> No body=pull_request_completed(pull_request_url=pull_request_url), ) + end_time = time.time() supabase_manager.complete_and_update_usage_record( usage_record_id=usage_record_id, token_input=token_input, token_output=token_output, + total_seconds=int(end_time - current_time), ) return diff --git a/services/supabase/gitauto_manager.py b/services/supabase/gitauto_manager.py index 4d7cdbfa..e824d237 100644 --- a/services/supabase/gitauto_manager.py +++ b/services/supabase/gitauto_manager.py @@ -17,6 +17,7 @@ def complete_and_update_usage_record( usage_record_id: int, token_input: int, token_output: int, + total_seconds: int, ) -> None: """Add agent information to usage record and set is_completed to True.""" try: @@ -25,6 +26,7 @@ def complete_and_update_usage_record( "is_completed": True, "token_input": token_input, "token_output": token_output, + "total_seconds": total_seconds, } ).eq(column="id", value=usage_record_id).execute() except Exception as e: diff --git a/tests/services/supabase/test_gitauto_manager.py b/tests/services/supabase/test_gitauto_manager.py index 02f58ae2..dc176378 100644 --- a/tests/services/supabase/test_gitauto_manager.py +++ b/tests/services/supabase/test_gitauto_manager.py @@ -62,6 +62,7 @@ def test_create_update_user_request_works() -> None: usage_record_id=usage_record_id, token_input=1000, token_output=100, + total_seconds=100, ) is None ) @@ -118,6 +119,7 @@ def test_complete_and_update_usage_record_only_updates_one_record() -> None: usage_record_id=usage_record_id, token_input=1000, token_output=100, + total_seconds=100, ) is None ) diff --git a/tests/services/supabase/test_users_manager.py b/tests/services/supabase/test_users_manager.py index 1517c6c8..276b6d09 100644 --- a/tests/services/supabase/test_users_manager.py +++ b/tests/services/supabase/test_users_manager.py @@ -84,6 +84,7 @@ def test_create_and_update_user_request_works() -> None: usage_record_id=usage_record_id, token_input=1000, token_output=100, + total_seconds=100, ) is None )