Merge pull request #155 from kortix-ai/fix/billing

v1 Fix/billing

Author: markokraemer

.gitignore

@@ -162,7 +162,6 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
-# AgentPress 
 /threads
 state.json
 /workspace/

backend/.gitignore

@@ -163,7 +163,6 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
-# AgentPress 
 /threads
 state.json
 /workspace/

backend/agent/api.py

@@ -15,9 +15,9 @@
 from services.supabase import DBConnection
 from services import redis
 from agent.run import run_agent
-from utils.auth_utils import get_current_user_id, get_user_id_from_stream_auth, verify_thread_access
+from utils.auth_utils import get_current_user_id_from_jwt, get_user_id_from_stream_auth, verify_thread_access
 from utils.logger import logger
-from utils.billing import check_billing_status, get_account_id_from_thread
+from services.billing import check_billing_status
 from sandbox.sandbox import create_sandbox, get_or_start_sandbox
 from services.llm import make_llm_api_call
 
@@ -348,7 +348,7 @@ async def get_or_create_project_sandbox(client, project_id: str):
 async def start_agent(
     thread_id: str,
     body: AgentStartRequest = Body(...),
-    user_id: str = Depends(get_current_user_id)
+    user_id: str = Depends(get_current_user_id_from_jwt)
 ):
     """Start an agent for a specific thread in the background."""
     global instance_id # Ensure instance_id is accessible
@@ -412,7 +412,7 @@ async def start_agent(
     return {"agent_run_id": agent_run_id, "status": "running"}
 
 @router.post("/agent-run/{agent_run_id}/stop")
-async def stop_agent(agent_run_id: str, user_id: str = Depends(get_current_user_id)):
+async def stop_agent(agent_run_id: str, user_id: str = Depends(get_current_user_id_from_jwt)):
     """Stop a running agent."""
     logger.info(f"Received request to stop agent run: {agent_run_id}")
     client = await db.client
@@ -421,7 +421,7 @@ async def stop_agent(agent_run_id: str, user_id: str = Depends(get_current_user_
     return {"status": "stopped"}
 
 @router.get("/thread/{thread_id}/agent-runs")
-async def get_agent_runs(thread_id: str, user_id: str = Depends(get_current_user_id)):
+async def get_agent_runs(thread_id: str, user_id: str = Depends(get_current_user_id_from_jwt)):
     """Get all agent runs for a thread."""
     logger.info(f"Fetching agent runs for thread: {thread_id}")
     client = await db.client
@@ -431,7 +431,7 @@ async def get_agent_runs(thread_id: str, user_id: str = Depends(get_current_user
     return {"agent_runs": agent_runs.data}
 
 @router.get("/agent-run/{agent_run_id}")
-async def get_agent_run(agent_run_id: str, user_id: str = Depends(get_current_user_id)):
+async def get_agent_run(agent_run_id: str, user_id: str = Depends(get_current_user_id_from_jwt)):
     """Get agent run status and responses."""
     logger.info(f"Fetching agent run details: {agent_run_id}")
     client = await db.client
@@ -859,7 +859,7 @@ async def initiate_agent_with_files(
     stream: Optional[bool] = Form(True),
     enable_context_manager: Optional[bool] = Form(False),
     files: List[UploadFile] = File(default=[]),
-    user_id: str = Depends(get_current_user_id)
+    user_id: str = Depends(get_current_user_id_from_jwt)
 ):
     """Initiate a new agent session with optional file attachments."""
     global instance_id # Ensure instance_id is accessible

backend/agent/run.py

@@ -20,7 +20,8 @@
 from agent.tools.data_providers_tool import DataProvidersTool
 from agent.prompt import get_system_prompt
 from utils import logger
-from utils.billing import check_billing_status, get_account_id_from_thread
+from utils.auth_utils import get_account_id_from_thread
+from services.billing import check_billing_status
 
 load_dotenv()
 

backend/api.py

@@ -16,6 +16,7 @@
 # Import the agent API module
 from agent import api as agent_api
 from sandbox import api as sandbox_api
+from services import billing as billing_api
 
 # Load environment variables (these will be available through config)
 load_dotenv()
@@ -132,6 +133,9 @@ async def log_requests_middleware(request: Request, call_next):
 # Include the sandbox router with a prefix
 app.include_router(sandbox_api.router, prefix="/api")
 
+# Include the billing router with a prefix
+app.include_router(billing_api.router, prefix="/api")
+
 @app.get("/api/health")
 async def health_check():
     """Health check endpoint to verify API is working."""
@@ -152,5 +156,6 @@ async def health_check():
         "api:app", 
         host="0.0.0.0", 
         port=8000,
-        workers=workers
+        workers=workers,
+        reload=True
     )

backend/poetry.lock

@@ -50,14 +50,14 @@ nest-asyncio = "^1.6.0"
 vncdotool = "^1.2.0"
 tavily-python = "^0.5.4"
 pytesseract = "^0.3.13"
+stripe = "^12.0.1"
 
 [tool.poetry.scripts]
 agentpress = "agentpress.cli:main"
 
 [[tool.poetry.packages]]
 include = "agentpress"
 
-
 [tool.poetry.group.dev.dependencies]
 daytona-sdk = "^0.14.0"
 

backend/pyproject.toml

@@ -6,7 +6,7 @@
 from pydantic import BaseModel
 
 from utils.logger import logger
-from utils.auth_utils import get_current_user_id, get_user_id_from_stream_auth, get_optional_user_id
+from utils.auth_utils import get_current_user_id_from_jwt, get_user_id_from_stream_auth, get_optional_user_id
 from sandbox.sandbox import get_or_start_sandbox
 from services.supabase import DBConnection
 from agent.api import get_or_create_project_sandbox