Coverage for session_buddy / server_optimized.py: 18.34%
217 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-04 00:43 -0800
1#!/usr/bin/env python3
2"""Optimized Session Management MCP Server.
4This is the refactored, modular version of the session management server.
5It's organized into focused modules for better maintainability and performance.
6"""
8import sys
9from collections.abc import AsyncGenerator, Callable
10from contextlib import asynccontextmanager, suppress
11from pathlib import Path
12from typing import Any
14# Add project root to Python path
15project_root = Path(__file__).parent.parent
16if str(project_root) not in sys.path: 16 ↛ 17line 16 didn't jump to line 17 because the condition on line 16 was never true
17 sys.path.insert(0, str(project_root))
19# Lazy loading for FastMCP
20try:
21 from fastmcp import FastMCP
23 MCP_AVAILABLE = True
24except ImportError:
25 # Check if we're in a test environment
26 if "pytest" in sys.modules or "test" in sys.argv[0].lower():
27 # Create a minimal mock FastMCP for testing
28 class MockFastMCP:
29 def __init__(self, name: str, lifespan: Any = None, **kwargs: Any) -> None:
30 self.name = name
31 self.tools: dict[str, Any] = {}
32 self.prompts: dict[str, Any] = {}
33 self.lifespan = lifespan
35 def tool(
36 self,
37 *args: Any,
38 **kwargs: Any,
39 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
40 def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
41 return func
43 return decorator
45 def prompt(
46 self,
47 *args: Any,
48 **kwargs: Any,
49 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
50 def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
51 return func
53 return decorator
55 def run(self, *args: Any, **kwargs: Any) -> None:
56 pass
58 FastMCP = MockFastMCP # type: ignore[no-redef]
59 MCP_AVAILABLE = False
60 else:
61 sys.exit(1)
63# Initialize logging
64from session_buddy.utils.logging import get_session_logger
66logger = get_session_logger()
68# Import required modules for automatic lifecycle
69import os
71from session_buddy.core import SessionLifecycleManager
72from session_buddy.utils.git_operations import get_git_root, is_git_repository
74# Global session manager for lifespan handlers
75lifecycle_manager = SessionLifecycleManager()
77# Global connection info for notification display
78_connection_info = None
81# Lifespan handler for automatic session management
82@asynccontextmanager
83async def session_lifecycle(app: Any) -> AsyncGenerator[None]:
84 """Automatic session lifecycle for git repositories only."""
85 current_dir = Path.cwd()
87 # Only auto-initialize for git repositories
88 if is_git_repository(current_dir):
89 try:
90 git_root = get_git_root(current_dir)
91 logger.info(f"Git repository detected at {git_root}")
93 # Run the same logic as the init tool but with connection notification
94 result = await lifecycle_manager.initialize_session(str(current_dir))
95 if result["success"]:
96 logger.info("✅ Auto-initialized session for git repository")
98 # Store connection info for display via tools
99 global _connection_info
100 _connection_info = {
101 "connected_at": "just connected",
102 "project": result["project"],
103 "quality_score": result["quality_score"],
104 "previous_session": result.get("previous_session"),
105 "recommendations": result["quality_data"].get(
106 "recommendations",
107 [],
108 ),
109 }
110 else:
111 logger.warning(f"Auto-init failed: {result['error']}")
112 except Exception as e:
113 logger.warning(f"Auto-init failed (non-critical): {e}")
114 else:
115 logger.debug("Non-git directory - skipping auto-initialization")
117 yield # Server runs normally
119 # On disconnect - cleanup for git repos only
120 if is_git_repository(current_dir):
121 try:
122 result = await lifecycle_manager.end_session()
123 if result["success"]:
124 logger.info("✅ Auto-ended session for git repository")
125 else:
126 logger.warning(f"Auto-cleanup failed: {result['error']}")
127 except Exception as e:
128 logger.warning(f"Auto-cleanup failed (non-critical): {e}")
131# Initialize MCP server with lifespan
132mcp = FastMCP("session-buddy", lifespan=session_lifecycle)
134# Register modularized tools
135from session_buddy.tools import register_memory_tools, register_session_tools
137# Core session management tools
138register_session_tools(mcp)
140# Memory and reflection tools
141register_memory_tools(mcp)
144@mcp.tool()
145async def session_welcome() -> str:
146 """Display session connection information and previous session details."""
147 global _connection_info
149 if not _connection_info:
150 return "ℹ️ Session information not available (may not be a git repository)"
152 output = []
153 output.append("🚀 Session Management Connected!")
154 output.append("=" * 40)
156 # Current session info
157 output.append(f"📁 Project: {_connection_info['project']}")
158 output.append(f"📊 Current quality score: {_connection_info['quality_score']}/100")
159 output.append(f"🔗 Connection status: {_connection_info['connected_at']}")
161 # Previous session info
162 previous = _connection_info.get("previous_session")
163 if previous:
164 output.extend(("\n📋 Previous Session Summary:", "-" * 30))
166 if "ended_at" in previous:
167 output.append(f"⏰ Last session ended: {previous['ended_at']}")
168 if "quality_score" in previous:
169 output.append(f"📈 Final score: {previous['quality_score']}")
170 if "top_recommendation" in previous:
171 output.append(f"💡 Key recommendation: {previous['top_recommendation']}")
173 output.append("\n✨ Session continuity restored - your progress is preserved!")
174 else:
175 output.extend(
176 (
177 "\n🌟 This is your first session in this project!",
178 "💡 Session data will be preserved for future continuity",
179 )
180 )
182 # Current recommendations
183 recommendations = _connection_info.get("recommendations", [])
184 if recommendations:
185 output.append("\n🎯 Current Recommendations:")
186 for i, rec in enumerate(recommendations[:3], 1):
187 output.append(f" {i}. {rec}")
189 output.extend(
190 (
191 "\n🔧 Use other session-buddy tools for:",
192 " • /session-buddy:status - Detailed project health",
193 " • /session-buddy:checkpoint - Mid-session quality check",
194 " • /session-buddy:end - Graceful session cleanup",
195 )
196 )
198 # Clear the connection info after display
199 _connection_info = None
201 return "\n".join(output)
204# Import the real SessionPermissionsManager from core module
205from session_buddy.core.permissions import SessionPermissionsManager
206from session_buddy.di.container import depends
209def _get_permissions_manager() -> SessionPermissionsManager:
210 import typing as t
211 from contextlib import suppress
213 with suppress(Exception):
214 manager = t.cast(
215 "SessionPermissionsManager | None",
216 depends.get_sync(SessionPermissionsManager),
217 )
218 if isinstance(manager, SessionPermissionsManager):
219 return manager
221 from session_buddy.di.config import SessionPaths
223 with suppress(Exception):
224 paths = depends.get_sync(SessionPaths)
225 if isinstance(paths, SessionPaths):
226 manager = SessionPermissionsManager(paths.claude_dir)
227 depends.set(SessionPermissionsManager, manager)
228 return manager
230 paths = SessionPaths.from_home()
231 paths.ensure_directories()
232 manager = SessionPermissionsManager(paths.claude_dir)
233 depends.set(SessionPermissionsManager, manager)
234 return manager
237@mcp.tool()
238async def permissions(action: str = "status", operation: str | None = None) -> str:
239 """Manage session permissions for trusted operations to avoid repeated prompts.
241 Args:
242 action: Action to perform: status (show current), trust (add operation), revoke_all (reset)
243 operation: Operation to trust (required for 'trust' action)
245 """
246 output = []
247 output.extend(("🔐 Session Permissions Management", "=" * 40))
249 permissions_manager = _get_permissions_manager()
250 if action == "status":
251 if permissions_manager.trusted_operations:
252 output.append(
253 f"✅ {len(permissions_manager.trusted_operations)} trusted operations:",
254 )
255 for op in sorted(permissions_manager.trusted_operations):
256 output.append(f" • {op}")
257 output.append(
258 "\n💡 These operations will not prompt for permission in future sessions",
259 )
260 else:
261 output.extend(
262 (
263 "⚠️ No operations are currently trusted",
264 "💡 Operations will be automatically trusted on first successful use",
265 )
266 )
268 output.extend(
269 (
270 "\n🛠️ Common Operations That Can Be Trusted:",
271 " • UV Package Management - uv sync, pip operations",
272 " • Git Repository Access - git status, commit, push",
273 " • Project File Access - reading/writing project files",
274 " • Subprocess Execution - running build tools, tests",
275 " • Claude Directory Access - accessing ~/.claude/",
276 )
277 )
279 elif action == "trust":
280 if not operation:
281 output.extend(
282 (
283 "❌ Error: 'operation' parameter required for 'trust' action",
284 "💡 Example: permissions with action='trust' and operation='uv_package_management'",
285 )
286 )
287 else:
288 permissions_manager.trust_operation(operation)
289 output.extend(
290 (
291 f"✅ Operation '{operation}' has been added to trusted operations",
292 "💡 This operation will no longer require permission prompts",
293 )
294 )
296 elif action == "revoke_all":
297 count = len(permissions_manager.trusted_operations)
298 permissions_manager.trusted_operations.clear()
299 output.extend(
300 (
301 f"🗑️ Revoked {count} trusted operations",
302 "💡 All operations will now require permission prompts",
303 )
304 )
306 else:
307 output.extend(
308 (
309 f"❌ Unknown action: {action}",
310 "💡 Valid actions: status, trust, revoke_all",
311 )
312 )
314 return "\n".join(output)
317# Compaction analysis and auto-execution functions
318def _count_significant_files(current_dir: Path) -> int:
319 """Count significant files in project as a complexity indicator."""
320 file_count = 0
321 with suppress(OSError, PermissionError, FileNotFoundError, ValueError):
322 for file_path in current_dir.rglob("*"):
323 if (
324 file_path.is_file()
325 and not any(part.startswith(".") for part in file_path.parts)
326 and file_path.suffix
327 in {
328 ".py",
329 ".js",
330 ".ts",
331 ".jsx",
332 ".tsx",
333 ".go",
334 ".rs",
335 ".java",
336 ".cpp",
337 ".c",
338 ".h",
339 }
340 ):
341 file_count += 1
342 if file_count > 50: # Stop counting after threshold
343 break
344 return file_count
347def _check_git_activity(current_dir: Path) -> tuple[int, int] | None:
348 """Check for active development via git and return (recent_commits, modified_files)."""
349 import subprocess # nosec B404
351 git_dir = current_dir / ".git"
352 if not git_dir.exists():
353 return None
355 try:
356 # Check number of recent commits as activity indicator
357 result = subprocess.run(
358 ["git", "log", "--oneline", "-20", "--since='24 hours ago'"],
359 check=False,
360 capture_output=True,
361 text=True,
362 cwd=current_dir,
363 timeout=5,
364 )
365 if result.returncode == 0:
366 recent_commits = len(
367 [line for line in result.stdout.strip().split("\n") if line.strip()],
368 )
369 else:
370 recent_commits = 0
372 # Check for large number of modified files
373 status_result = subprocess.run(
374 ["git", "status", "--porcelain"],
375 check=False,
376 capture_output=True,
377 text=True,
378 cwd=current_dir,
379 timeout=5,
380 )
381 if status_result.returncode == 0:
382 modified_files = len(
383 [
384 line
385 for line in status_result.stdout.strip().split("\n")
386 if line.strip()
387 ],
388 )
389 else:
390 modified_files = 0
392 return recent_commits, modified_files
394 except (subprocess.TimeoutExpired, Exception):
395 return None
398def _evaluate_large_project_heuristic(file_count: int) -> tuple[bool, str]:
399 """Evaluate if the project is large enough to benefit from compaction."""
400 if file_count > 50:
401 return (
402 True,
403 "Large codebase with 50+ source files detected - context compaction recommended",
404 )
405 return False, ""
408def _evaluate_git_activity_heuristic(
409 git_activity: tuple[int, int] | None,
410) -> tuple[bool, str]:
411 """Evaluate if git activity suggests compaction would be beneficial."""
412 if git_activity:
413 recent_commits, modified_files = git_activity
415 if recent_commits >= 3:
416 return (
417 True,
418 f"High development activity ({recent_commits} commits in 24h) - compaction recommended",
419 )
421 if modified_files >= 10:
422 return (
423 True,
424 f"Many modified files ({modified_files}) detected - context optimization beneficial",
425 )
427 return False, ""
430def _evaluate_python_project_heuristic(current_dir: Path) -> tuple[bool, str]:
431 """Evaluate if this is a Python project that might benefit from compaction."""
432 if (current_dir / "tests").exists() and (current_dir / "pyproject.toml").exists():
433 return (
434 True,
435 "Python project with tests detected - compaction may improve focus",
436 )
437 return False, ""
440def _get_default_compaction_reason() -> str:
441 """Get the default reason when no strong indicators are found."""
442 return "Context appears manageable - compaction not immediately needed"
445def _get_fallback_compaction_reason() -> str:
446 """Get fallback reason when evaluation fails."""
447 return "Unable to assess context complexity - compaction may be beneficial as a precaution"
450def should_suggest_compact() -> tuple[bool, str]:
451 """Determine if compacting would be beneficial and provide reasoning.
452 Returns (should_compact, reason).
453 """
454 from pathlib import Path
456 try:
457 current_dir = Path(os.environ.get("PWD", Path.cwd()))
459 # Count significant files in project as a complexity indicator
460 file_count = _count_significant_files(current_dir)
462 # Large project heuristic
463 should_compact, reason = _evaluate_large_project_heuristic(file_count)
464 if should_compact:
465 return should_compact, reason
467 # Check for active development via git
468 git_activity = _check_git_activity(current_dir)
469 should_compact, reason = _evaluate_git_activity_heuristic(git_activity)
470 if should_compact:
471 return should_compact, reason
473 # Check for common patterns suggesting complex session
474 should_compact, reason = _evaluate_python_project_heuristic(current_dir)
475 if should_compact:
476 return should_compact, reason
478 # Default to not suggesting unless we have clear indicators
479 return False, _get_default_compaction_reason()
481 except Exception:
482 # If we can't determine, err on the side of suggesting compaction for safety
483 return True, _get_fallback_compaction_reason()
486async def _execute_auto_compact() -> str:
487 """Execute internal compaction instead of recommending /compact command."""
488 try:
489 # This would trigger the same logic as /compact but automatically
490 # For now, we use the memory system's auto-compaction
491 return "✅ Context automatically optimized via intelligent memory management"
492 except Exception as e:
493 logger.warning(f"Auto-compact execution failed: {e}")
494 return f"⚠️ Auto-compact failed: {e!s} - recommend manual /compact"
497# Enhanced tools with auto-compaction
498@mcp.tool()
499async def auto_compact() -> str:
500 """Automatically trigger conversation compaction with intelligent summary."""
501 output = []
502 output.extend(("🗜️ Auto-Compaction Feature", "=" * 30))
504 should_compact, reason = should_suggest_compact()
505 output.append(f"📊 Analysis: {reason}")
507 if should_compact:
508 output.append("\n🔄 Executing automatic compaction...")
509 compact_result = await _execute_auto_compact()
510 output.append(compact_result)
511 else:
512 output.append("✅ Context optimization not needed at this time")
514 return "\n".join(output)
517@mcp.tool()
518async def quality_monitor() -> str:
519 """Phase 3: Proactive quality monitoring with early warning system."""
520 output = []
521 output.extend(
522 (
523 "📊 Quality Monitoring",
524 "=" * 25,
525 "✅ Quality monitoring is integrated into the session management system",
526 "💡 Use the 'status' tool to get current quality metrics",
527 "💡 Use the 'checkpoint' tool for comprehensive quality assessment",
528 )
529 )
530 return "\n".join(output)
533# Server startup
534def run_server() -> None:
535 """Run the optimized MCP server."""
536 try:
537 logger.info("Starting optimized session-buddy server")
539 # Log the modular structure
540 logger.info(
541 "Modular components loaded",
542 session_tools=True,
543 memory_tools=True,
544 git_operations=True,
545 logging_utils=True,
546 )
548 if MCP_AVAILABLE:
549 mcp.run()
550 else:
551 logger.warning("Running in mock mode - FastMCP not available")
553 except Exception as e:
554 logger.exception(f"Server startup failed: {e}")
555 sys.exit(1)
558if __name__ == "__main__":
559 run_server()