vibesurf 0.1.9a6__py3-none-any.whl → 0.1.11__py3-none-any.whl

vibe_surf/backend/api/task.py

@@ -21,10 +21,13 @@ from ..shared_state import (
     execute_task_background,
     is_task_running,
     get_active_task_info,
-    clear_active_task
+    clear_active_task,
+    browser_manager
 )
 
-logger = logging.getLogger(__name__)
+from vibe_surf.logger import get_logger
+
+logger = get_logger(__name__)
 
 router = APIRouter(prefix="/tasks", tags=["tasks"])
 
@@ -75,16 +78,16 @@ async def submit_task(
         task_id = uuid7str()
 
         # Get MCP server config for saving
-        from ..shared_state import controller, active_mcp_server
+        from ..shared_state import vibesurf_tools, active_mcp_server
         mcp_server_config = task_request.mcp_server_config
-        if not mcp_server_config and controller and hasattr(controller, 'mcp_server_config'):
-            mcp_server_config = controller.mcp_server_config
-        
+        if not mcp_server_config and vibesurf_tools and hasattr(vibesurf_tools, 'mcp_server_config'):
+            mcp_server_config = vibesurf_tools.mcp_server_config
+
         # Ensure we have a valid MCP server config (never None)
         if mcp_server_config is None:
             mcp_server_config = {"mcpServers": {}}
             logger.info("Using default empty MCP server configuration")
-        
+
         # DEBUG: Log the type and content of mcp_server_config
         logger.info(f"mcp_server_config type: {type(mcp_server_config)}, value: {mcp_server_config}")
 
@@ -139,16 +142,17 @@ async def _ensure_llm_initialized(llm_profile):
     """Ensure LLM is initialized with the specified profile"""
     from ..utils.llm_factory import create_llm_from_profile
     from ..shared_state import vibesurf_agent
-    
+
     if not vibesurf_agent:
         raise HTTPException(status_code=503, detail="VibeSurf agent not initialized")
-    
+
     # Always create new LLM instance to ensure we're using the right profile
     new_llm = create_llm_from_profile(llm_profile)
-    
-    # Update vibesurf agent's LLM
-    vibesurf_agent.llm = new_llm
-    logger.info(f"LLM updated for profile: {llm_profile['profile_name']}")
+
+    # Update vibesurf agent's LLM and register with token cost service
+    if vibesurf_agent and vibesurf_agent.token_cost_service:
+        vibesurf_agent.llm = vibesurf_agent.token_cost_service.register_llm(new_llm)
+        logger.info(f"LLM updated and registered for token tracking for profile: {llm_profile['profile_name']}")
 
 
 @router.post("/pause")
@@ -259,6 +263,36 @@ async def stop_task(control_request: TaskControlRequest):
         raise HTTPException(status_code=500, detail=f"Failed to stop task: {str(e)}")
 
 
+@router.post("/add-new-task")
+async def add_new_task(control_request: TaskControlRequest):
+    """Add a new task or follow-up instruction during execution"""
+    from ..shared_state import vibesurf_agent
+
+    if not vibesurf_agent:
+        raise HTTPException(status_code=503, detail="VibeSurf agent not initialized")
+
+    if not is_task_running():
+        raise HTTPException(status_code=400, detail="No active task to add new instruction to")
+
+    try:
+        # Use the reason field as the new task content
+        new_task = control_request.reason or "No additional task provided"
+        
+        # Add the new task to the running agent
+        await vibesurf_agent.add_new_task(new_task)
+
+        return {
+            "success": True,
+            "message": "New task added successfully",
+            "operation": "add_new_task",
+            "new_task": new_task
+        }
+
+    except Exception as e:
+        logger.error(f"Failed to add new task: {e}")
+        raise HTTPException(status_code=500, detail=f"Failed to add new task: {str(e)}")
+
+
 @router.get("/detailed-status")
 async def get_detailed_task_status():
     """Get detailed task execution status with vibesurf information"""

vibe_surf/backend/database/manager.py

@@ -13,11 +13,14 @@ from .models import Base
 from typing import AsyncGenerator
 import logging
 
-logger = logging.getLogger(__name__)
+from vibe_surf.logger import get_logger
+
+logger = get_logger(__name__)
+
 
 class DatabaseManager:
     """Database connection and session management"""
-    
+
     def __init__(self, database_url: str = None):
         """Initialize database manager
         
@@ -29,7 +32,7 @@ class DatabaseManager:
             'VIBESURF_DATABASE_URL',
             f'sqlite+aiosqlite:///{os.path.join(shared_state.workspace_dir, "vibe_surf.db")}'
         )
-        
+
         # Configure engine based on database type
         if self.database_url.startswith('sqlite'):
             # SQLite configuration for development
@@ -52,23 +55,23 @@ class DatabaseManager:
                 pool_recycle=3600,
                 echo=False
             )
-        
+
         self.async_session_factory = sessionmaker(
-            self.engine, 
-            class_=AsyncSession, 
+            self.engine,
+            class_=AsyncSession,
             expire_on_commit=False
         )
-    
+
     async def create_tables(self):
         """Create all database tables"""
         async with self.engine.begin() as conn:
             await conn.run_sync(Base.metadata.create_all)
-    
+
     async def drop_tables(self):
         """Drop all database tables"""
         async with self.engine.begin() as conn:
             await conn.run_sync(Base.metadata.drop_all)
-    
+
     async def get_session(self) -> AsyncGenerator[AsyncSession, None]:
         """Get async database session"""
         async with self.async_session_factory() as session:
@@ -80,50 +83,53 @@ class DatabaseManager:
                 raise
             finally:
                 await session.close()
-    
+
     async def close(self):
         """Close database connections"""
         await self.engine.dispose()
 
+
 # Dependency for FastAPI
 async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
     """FastAPI dependency for database sessions"""
     from .. import shared_state
-    
+
     if not shared_state.db_manager:
         raise RuntimeError("Database manager not initialized. Call initialize_vibesurf_components() first.")
-    
+
     async for session in shared_state.db_manager.get_session():
         yield session
 
+
 # Database initialization script
 async def init_database():
     """Initialize database with tables"""
     from .. import shared_state
-    
+
     logger.info("🗄️ Initializing VibeSurf database...")
-    
+
     try:
         if not shared_state.db_manager:
             raise RuntimeError("Database manager not initialized. Call initialize_vibesurf_components() first.")
-            
+
         await shared_state.db_manager.create_tables()
         logger.info("✅ Database tables created successfully")
         logger.info("✅ VibeSurf database ready for single-task execution")
-        
+
     except Exception as e:
         logger.error(f"❌ Database initialization failed: {e}")
         raise
 
+
 if __name__ == "__main__":
     # For standalone execution, initialize a temporary db_manager
     import os
     from .. import shared_state
-    
+
     workspace_dir = os.getenv("VIBESURF_WORKSPACE", os.path.join(os.path.dirname(__file__), "../vibesurf_workspace"))
     database_url = os.getenv(
         'VIBESURF_DATABASE_URL',
         f'sqlite+aiosqlite:///{os.path.join(workspace_dir, "vibe_surf.db")}'
     )
     shared_state.db_manager = DatabaseManager(database_url)
-    asyncio.run(init_database())
+    asyncio.run(init_database())