stravinsky 0.2.38__py3-none-any.whl → 0.2.52__py3-none-any.whl

mcp_bridge/__init__.py

@@ -1 +1 @@
-__version__ = "0.2.32"
+__version__ = "0.2.52"

mcp_bridge/auth/token_refresh.py

@@ -0,0 +1,130 @@
+"""
+Background Token Refresh Scheduler
+
+Proactively refreshes OAuth tokens before they expire:
+- Gemini: Refreshes when 30 minutes remaining (tokens expire in ~1 hour)
+- OpenAI: Refreshes when 12 hours remaining (tokens expire in ~24 hours)
+"""
+
+import asyncio
+import logging
+import time
+from typing import TYPE_CHECKING
+
+from .oauth import refresh_access_token as gemini_refresh
+from .openai_oauth import refresh_access_token as openai_refresh
+
+if TYPE_CHECKING:
+    from .token_store import TokenStore
+
+logger = logging.getLogger(__name__)
+
+# Refresh configuration per provider
+REFRESH_CONFIG = {
+    "gemini": {
+        "interval": 3600,     # Token expires in ~1 hour
+        "buffer": 1800,       # Refresh with 30 min remaining
+    },
+    "openai": {
+        "interval": 86400,    # Token expires in ~24 hours
+        "buffer": 43200,      # Refresh with 12 hours remaining
+    },
+}
+
+
+async def background_token_refresh(token_store: "TokenStore") -> None:
+    """
+    Background task that proactively refreshes tokens before expiry.
+
+    Runs in an infinite loop, checking every 5 minutes if any tokens
+    need refreshing. This prevents tokens from expiring during long
+    sessions.
+
+    Args:
+        token_store: The token store to manage.
+    """
+    logger.info("Starting background token refresh scheduler")
+
+    while True:
+        for provider, config in REFRESH_CONFIG.items():
+            await _refresh_if_needed(token_store, provider, config["buffer"])
+
+        # Check every 5 minutes
+        await asyncio.sleep(300)
+
+
+async def _refresh_if_needed(
+    token_store: "TokenStore",
+    provider: str,
+    buffer_seconds: int,
+) -> None:
+    """
+    Refresh a provider's token if it's close to expiring.
+
+    Args:
+        token_store: Token store instance
+        provider: Provider name (gemini, openai)
+        buffer_seconds: Refresh when this many seconds remain
+    """
+    try:
+        # Check if token needs refresh
+        if not token_store.needs_refresh(provider, buffer_seconds=buffer_seconds):
+            return
+
+        token = token_store.get_token(provider)
+        if not token or not token.get("refresh_token"):
+            return  # No token or no refresh token
+
+        # Get the appropriate refresh function
+        refresh_fn = gemini_refresh if provider == "gemini" else openai_refresh
+
+        # Perform refresh
+        result = refresh_fn(token["refresh_token"])
+
+        # Update stored token
+        token_store.update_access_token(
+            provider,
+            result.access_token,
+            result.expires_in,
+        )
+
+        logger.info(
+            f"Proactively refreshed {provider} token "
+            f"(expires in {result.expires_in}s)"
+        )
+
+    except Exception as e:
+        logger.warning(f"Failed to refresh {provider} token: {e}")
+
+
+def get_token_status(token_store: "TokenStore") -> dict[str, dict]:
+    """
+    Get status of all provider tokens.
+
+    Returns:
+        Dict mapping provider to status info.
+    """
+    status = {}
+
+    for provider in REFRESH_CONFIG:
+        token = token_store.get_token(provider)
+
+        if not token:
+            status[provider] = {"status": "not_authenticated"}
+            continue
+
+        expires_at = token.get("expires_at", 0)
+        if expires_at <= 0:
+            status[provider] = {"status": "authenticated", "expires": "unknown"}
+        else:
+            remaining = expires_at - time.time()
+            if remaining <= 0:
+                status[provider] = {"status": "expired"}
+            else:
+                status[provider] = {
+                    "status": "valid",
+                    "expires_in_seconds": int(remaining),
+                    "expires_in_minutes": int(remaining / 60),
+                }
+
+    return status

mcp_bridge/hooks/__init__.py

@@ -1,6 +1,13 @@
 """
 Hooks initialization.
-Registers all Tier 1-4 hooks into the HookManager.
+Registers all Tier 1-5 hooks into the HookManager.
+
+Hook Tiers:
+- Tier 1: Post-tool-call (immediate response modification)
+- Tier 2: Pre-model-invoke (context management)
+- Tier 3: Pre-model-invoke (performance optimization)
+- Tier 4: Pre-model-invoke (behavior enforcement)
+- Tier 5: Session lifecycle (idle detection, compaction)
 """
 
 from .manager import get_hook_manager
@@ -19,6 +26,11 @@ from .auto_slash_command import auto_slash_command_hook
 from .session_recovery import session_recovery_hook
 from .empty_message_sanitizer import empty_message_sanitizer_hook
 
+# New hooks based on oh-my-opencode patterns
+from .session_idle import session_idle_hook
+from .pre_compact import pre_compact_hook
+from .parallel_enforcer import parallel_enforcer_post_tool_hook
+
 
 def initialize_hooks():
     """Register all available hooks."""
@@ -30,6 +42,7 @@ def initialize_hooks():
     manager.register_post_tool_call(comment_checker_hook)
     manager.register_post_tool_call(agent_reminder_hook)
     manager.register_post_tool_call(session_recovery_hook)
+    manager.register_post_tool_call(parallel_enforcer_post_tool_hook)  # NEW: Enforce parallel spawning
 
     # Tier 2: Pre-model-invoke (context management)
     manager.register_pre_model_invoke(directory_context_hook)
@@ -46,4 +59,8 @@ def initialize_hooks():
     manager.register_pre_model_invoke(todo_continuation_hook)
     manager.register_pre_model_invoke(auto_slash_command_hook)
 
+    # Tier 5: Session lifecycle hooks (NEW)
+    manager.register_session_idle(session_idle_hook)  # Stop hook - idle detection
+    manager.register_pre_compact(pre_compact_hook)    # PreCompact - context preservation
+
     return manager

mcp_bridge/hooks/manager.py

@@ -12,6 +12,13 @@ logger = logging.getLogger(__name__)
 class HookManager:
     """
     Manages the registration and execution of hooks.
+
+    Hook Types:
+    - pre_tool_call: Before tool execution (can modify args or block)
+    - post_tool_call: After tool execution (can modify output)
+    - pre_model_invoke: Before model invocation (can modify prompt/params)
+    - session_idle: When session becomes idle (can inject continuation)
+    - pre_compact: Before context compaction (can preserve critical context)
     """
 
     _instance = None
@@ -26,6 +33,13 @@ class HookManager:
         self.pre_model_invoke_hooks: List[
             Callable[[Dict[str, Any]], Awaitable[Optional[Dict[str, Any]]]]
         ] = []
+        # New hook types based on oh-my-opencode patterns
+        self.session_idle_hooks: List[
+            Callable[[Dict[str, Any]], Awaitable[Optional[Dict[str, Any]]]]
+        ] = []
+        self.pre_compact_hooks: List[
+            Callable[[Dict[str, Any]], Awaitable[Optional[Dict[str, Any]]]]
+        ] = []
 
     @classmethod
     def get_instance(cls):
@@ -51,6 +65,18 @@ class HookManager:
         """Run before model invocation. Can modify prompt or parameters."""
         self.pre_model_invoke_hooks.append(hook)
 
+    def register_session_idle(
+        self, hook: Callable[[Dict[str, Any]], Awaitable[Optional[Dict[str, Any]]]]
+    ):
+        """Run when session becomes idle. Can inject continuation prompts."""
+        self.session_idle_hooks.append(hook)
+
+    def register_pre_compact(
+        self, hook: Callable[[Dict[str, Any]], Awaitable[Optional[Dict[str, Any]]]]
+    ):
+        """Run before context compaction. Can preserve critical context."""
+        self.pre_compact_hooks.append(hook)
+
     async def execute_pre_tool_call(
         self, tool_name: str, arguments: Dict[str, Any]
     ) -> Dict[str, Any]:
@@ -91,6 +117,30 @@ class HookManager:
                 logger.error(f"[HookManager] Error in pre_model_invoke hook {hook.__name__}: {e}")
         return current_params
 
+    async def execute_session_idle(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        """Executes all session idle hooks (Stop hook pattern)."""
+        current_params = params
+        for hook in self.session_idle_hooks:
+            try:
+                modified_params = await hook(current_params)
+                if modified_params is not None:
+                    current_params = modified_params
+            except Exception as e:
+                logger.error(f"[HookManager] Error in session_idle hook {hook.__name__}: {e}")
+        return current_params
+
+    async def execute_pre_compact(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        """Executes all pre-compact hooks (context preservation)."""
+        current_params = params
+        for hook in self.pre_compact_hooks:
+            try:
+                modified_params = await hook(current_params)
+                if modified_params is not None:
+                    current_params = modified_params
+            except Exception as e:
+                logger.error(f"[HookManager] Error in pre_compact hook {hook.__name__}: {e}")
+        return current_params
+
 
 def get_hook_manager() -> HookManager:
     return HookManager.get_instance()

mcp_bridge/hooks/parallel_enforcer.py

@@ -0,0 +1,127 @@
+"""
+Parallel Enforcer Hook - Enforce Parallel Agent Spawning.
+
+Detects when 2+ independent tasks exist and injects reminders
+to spawn agents in parallel rather than working sequentially.
+
+Based on oh-my-opencode's parallel execution enforcement pattern.
+"""
+
+import logging
+import re
+from typing import Any, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+# Parallel enforcement prompt
+PARALLEL_ENFORCEMENT_PROMPT = """
+[PARALLEL EXECUTION REQUIRED]
+
+You have {count} independent pending tasks. You MUST spawn agents for ALL of them simultaneously.
+
+CORRECT (Parallel - DO THIS):
+```
+agent_spawn(prompt="Task 1...", agent_type="explore", description="Task 1")
+agent_spawn(prompt="Task 2...", agent_type="explore", description="Task 2")
+agent_spawn(prompt="Task 3...", agent_type="dewey", description="Task 3")
+// All spawned in ONE response, then wait for results
+```
+
+WRONG (Sequential - DO NOT DO THIS):
+```
+Mark task 1 in_progress -> work on it -> complete
+Mark task 2 in_progress -> work on it -> complete  // TOO SLOW!
+```
+
+RULES:
+1. Spawn ALL independent tasks simultaneously using agent_spawn
+2. Do NOT mark any task as in_progress until agents are spawned
+3. Collect results with agent_output AFTER spawning
+4. Only work sequentially when tasks have dependencies
+"""
+
+# Track if enforcement was already triggered this session
+_enforcement_triggered: Dict[str, bool] = {}
+
+
+async def parallel_enforcer_hook(params: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    """
+    Post-tool-call hook that triggers after TodoWrite.
+
+    When 2+ pending todos are detected, injects parallel execution
+    enforcement prompt to prevent sequential work patterns.
+    """
+    tool_name = params.get("tool_name", "")
+    output = params.get("output", "")
+    session_id = params.get("session_id", "default")
+
+    # Only trigger for TodoWrite calls
+    if tool_name.lower() not in ["todowrite", "todo_write"]:
+        return None
+
+    # Count pending todos
+    pending_count = _count_pending_todos(output)
+
+    if pending_count < 2:
+        return None
+
+    # Check if already triggered recently
+    if _enforcement_triggered.get(session_id, False):
+        return None
+
+    # Mark as triggered
+    _enforcement_triggered[session_id] = True
+
+    logger.info(f"[ParallelEnforcerHook] Detected {pending_count} pending todos, enforcing parallel execution")
+
+    # Inject enforcement prompt
+    enforcement = PARALLEL_ENFORCEMENT_PROMPT.format(count=pending_count)
+    modified_output = output + "\n\n" + enforcement
+
+    return modified_output
+
+
+def _count_pending_todos(output: str) -> int:
+    """Count the number of pending todos in TodoWrite output."""
+    # Pattern matches various pending todo formats
+    patterns = [
+        r'\[pending\]',
+        r'"status":\s*"pending"',
+        r"status:\s*pending",
+        r"'status':\s*'pending'",
+    ]
+
+    total = 0
+    for pattern in patterns:
+        matches = re.findall(pattern, output, re.IGNORECASE)
+        total += len(matches)
+
+    return total
+
+
+def reset_enforcement(session_id: str = "default"):
+    """Reset enforcement state for a session."""
+    _enforcement_triggered[session_id] = False
+
+
+async def parallel_enforcer_post_tool_hook(
+    tool_name: str,
+    arguments: Dict[str, Any],
+    output: str
+) -> Optional[str]:
+    """
+    Post-tool-call hook interface for HookManager.
+
+    Wraps parallel_enforcer_hook for the standard hook signature.
+    """
+    params = {
+        "tool_name": tool_name,
+        "arguments": arguments,
+        "output": output,
+    }
+
+    result = await parallel_enforcer_hook(params)
+
+    if isinstance(result, str):
+        return result
+    return None

mcp_bridge/hooks/pre_compact.py

@@ -0,0 +1,224 @@
+"""
+PreCompact Hook - Context Preservation Before Compaction.
+
+Triggers before session compaction to preserve critical context
+that should survive summarization. Uses Gemini for intelligent
+context extraction and preservation.
+
+Based on oh-my-opencode's pre-compact hook pattern.
+"""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+# Flag to prevent recursive calls during Gemini invocation
+_in_preservation = False
+
+# Critical context patterns to preserve
+PRESERVE_PATTERNS = [
+    # Architecture decisions
+    "ARCHITECTURE:",
+    "DESIGN DECISION:",
+    "## Architecture",
+
+    # Important constraints
+    "CONSTRAINT:",
+    "REQUIREMENT:",
+    "MUST NOT:",
+    "NEVER:",
+
+    # Session state
+    "CURRENT TASK:",
+    "BLOCKED BY:",
+    "WAITING FOR:",
+
+    # Critical errors
+    "CRITICAL ERROR:",
+    "SECURITY ISSUE:",
+    "BREAKING CHANGE:",
+]
+
+# Memory anchors to inject into compaction
+MEMORY_ANCHORS: List[str] = []
+
+
+def register_memory_anchor(anchor: str, priority: str = "normal"):
+    """
+    Register a memory anchor to preserve during compaction.
+
+    Args:
+        anchor: The text to preserve
+        priority: "critical" or "normal"
+    """
+    if priority == "critical":
+        MEMORY_ANCHORS.insert(0, f"[CRITICAL] {anchor}")
+    else:
+        MEMORY_ANCHORS.append(anchor)
+
+    # Limit to 10 anchors to prevent bloat
+    while len(MEMORY_ANCHORS) > 10:
+        MEMORY_ANCHORS.pop()
+
+
+def clear_memory_anchors():
+    """Clear all registered memory anchors."""
+    MEMORY_ANCHORS.clear()
+
+
+async def pre_compact_hook(params: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    """
+    Pre-model-invoke hook that runs before context compaction.
+
+    Uses Gemini to intelligently extract and preserve critical context
+    that should survive summarization.
+    """
+    global _in_preservation
+
+    # Prevent recursive calls
+    if _in_preservation:
+        return None
+
+    prompt = params.get("prompt", "")
+
+    # Only activate for compaction-related prompts
+    if not _is_compaction_prompt(prompt):
+        return None
+
+    # Collect pattern-matched context
+    preserved_context = _extract_preserved_context(prompt)
+    preserved_context.extend(MEMORY_ANCHORS)
+
+    # Use Gemini for intelligent context extraction if prompt is long
+    if len(prompt) > 50000:
+        try:
+            _in_preservation = True
+            gemini_context = await _extract_context_with_gemini(prompt)
+            if gemini_context:
+                preserved_context.extend(gemini_context)
+        except Exception as e:
+            logger.warning(f"[PreCompactHook] Gemini extraction failed: {e}")
+        finally:
+            _in_preservation = False
+
+    if not preserved_context:
+        return None
+
+    # Build preservation section
+    preservation_section = _build_preservation_section(preserved_context)
+
+    logger.info(f"[PreCompactHook] Preserving {len(preserved_context)} context items")
+
+    # Inject into prompt
+    modified_prompt = prompt + "\n\n" + preservation_section
+
+    return {**params, "prompt": modified_prompt}
+
+
+async def _extract_context_with_gemini(prompt: str) -> List[str]:
+    """
+    Use Gemini to intelligently extract critical context to preserve.
+
+    Args:
+        prompt: The full conversation/context to analyze
+
+    Returns:
+        List of critical context items to preserve
+    """
+    try:
+        from ..tools.model_invoke import invoke_gemini_impl
+
+        # Truncate prompt if too long for Gemini
+        max_chars = 100000
+        truncated = prompt[:max_chars] if len(prompt) > max_chars else prompt
+
+        extraction_prompt = f"""Analyze this conversation and extract ONLY the most critical information that MUST be preserved during summarization.
+
+Focus on:
+1. Architecture decisions and their rationale
+2. Critical constraints or requirements
+3. Important error patterns or debugging insights
+4. Key file paths and their purposes
+5. Unfinished tasks or blocking issues
+
+Return a bullet list of critical items (max 10). Be extremely concise.
+
+CONVERSATION:
+{truncated}
+
+CRITICAL ITEMS TO PRESERVE:"""
+
+        result = await invoke_gemini_impl(
+            prompt=extraction_prompt,
+            model="gemini-3-flash",
+            max_tokens=2000,
+            temperature=0.1,
+        )
+
+        if not result:
+            return []
+
+        # Parse bullet points from response
+        lines = result.strip().split("\n")
+        items = []
+        for line in lines:
+            line = line.strip()
+            if line.startswith(("-", "*", "•")) or (len(line) > 1 and line[0].isdigit() and line[1] in ".):"):
+                # Clean up the bullet
+                item = line.lstrip("-*•0123456789.): ").strip()
+                if item and len(item) > 10:
+                    items.append(item)
+
+        return items[:10]  # Max 10 items
+
+    except Exception as e:
+        logger.warning(f"[PreCompactHook] Gemini context extraction error: {e}")
+        return []
+
+
+def _is_compaction_prompt(prompt: str) -> bool:
+    """Detect if this is a compaction/summarization prompt."""
+    compaction_signals = [
+        "summarize the conversation",
+        "compact the context",
+        "reduce context size",
+        "context window",
+        "summarization",
+    ]
+
+    prompt_lower = prompt.lower()
+    return any(signal in prompt_lower for signal in compaction_signals)
+
+
+def _extract_preserved_context(prompt: str) -> List[str]:
+    """Extract context matching preservation patterns."""
+    preserved = []
+    lines = prompt.split("\n")
+
+    for i, line in enumerate(lines):
+        for pattern in PRESERVE_PATTERNS:
+            if pattern in line:
+                # Capture the line and next 2 lines for context
+                context_lines = lines[i:i+3]
+                preserved.append("\n".join(context_lines))
+                break
+
+    return preserved
+
+
+def _build_preservation_section(context_items: List[str]) -> str:
+    """Build the preservation section to inject."""
+    section = """
+## CRITICAL CONTEXT TO PRESERVE
+
+The following information MUST be preserved in any summarization:
+
+"""
+    for i, item in enumerate(context_items, 1):
+        section += f"{i}. {item}\n\n"
+
+    section += """
+When summarizing, ensure these items are included verbatim or with minimal paraphrasing.
+"""
+    return section