1#!/usr/bin/env python3

2"""Optimized Session Management MCP Server.

3

4 This is the refactored, modular version of the session management server.

5 It's organized into focused modules for better maintainability and performance.

6"""

7

8 import sys

9 from collections.abc import AsyncGenerator, Callable

10 from contextlib import asynccontextmanager, suppress

11 from pathlib import Path

12 from typing import Any

13

14# Add project root to Python path

15 project_root = Path(__file__).parent.parent

16 if str(project_root) not in sys.path: 16 ↛ 17 line 16 didn't jump to line 17 because the condition on line 16 was never true

17 sys.path.insert(0, str(project_root))

18

19# Lazy loading for FastMCP

20 try:

21 from fastmcp import FastMCP

22

23 MCP_AVAILABLE = True

24 except 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

34

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

42

43 return decorator

44

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

52

53 return decorator

54

55 def run(self, *args: Any, **kwargs: Any) -> None:

56 pass

57

58 FastMCP = MockFastMCP # type: ignore[no-redef]

59 MCP_AVAILABLE = False

60 else:

61 sys.exit(1)

62

63# Initialize logging

64 from session_buddy.utils.logging import get_session_logger

65

66 logger = get_session_logger()

67

68# Import required modules for automatic lifecycle

69 import os

70

71 from session_buddy.core import SessionLifecycleManager

72 from session_buddy.utils.git_operations import get_git_root, is_git_repository

73

74# Global session manager for lifespan handlers

75 lifecycle_manager = SessionLifecycleManager()

76

77# Global connection info for notification display

78_connection_info = None

79

80

81# Lifespan handler for automatic session management

82@asynccontextmanager

83 async def session_lifecycle(app: Any) -> AsyncGenerator[None]:

84 """Automatic session lifecycle for git repositories only."""

85 current_dir = Path.cwd()

86

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}")

92

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")

97

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")

116

117 yield # Server runs normally

118

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}")

129

130

131# Initialize MCP server with lifespan

132 mcp = FastMCP("session-buddy", lifespan=session_lifecycle)

133

134# Register modularized tools

135 from session_buddy.tools import register_memory_tools, register_session_tools

136

137# Core session management tools

138 register_session_tools(mcp)

139

140# Memory and reflection tools

141 register_memory_tools(mcp)

142

143

144@mcp.tool()

145 async def session_welcome() -> str:

146 """Display session connection information and previous session details."""

147 global _connection_info

148

149 if not _connection_info:

150 return "ℹ️ Session information not available (may not be a git repository)"

151

152 output = []

153 output.append("🚀 Session Management Connected!")

154 output.append("=" * 40)

155

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']}")

160

161 # Previous session info

162 previous = _connection_info.get("previous_session")

163 if previous:

164 output.extend(("\n📋 Previous Session Summary:", "-" * 30))

165

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']}")

172

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 )

181

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}")

188

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 )

197

198 # Clear the connection info after display

199 _connection_info = None

200

201 return "\n".join(output)

202

203

204# Import the real SessionPermissionsManager from core module

205 from session_buddy.core.permissions import SessionPermissionsManager

206 from session_buddy.di.container import depends

207

208

209 def _get_permissions_manager() -> SessionPermissionsManager:

210 import typing as t

211 from contextlib import suppress

212

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

220

221 from session_buddy.di.config import SessionPaths

222

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

229

230 paths = SessionPaths.from_home()

231 paths.ensure_directories()

232 manager = SessionPermissionsManager(paths.claude_dir)

233 depends.set(SessionPermissionsManager, manager)

234 return manager

235

236

237@mcp.tool()

238 async def permissions(action: str = "status", operation: str | None = None) -> str:

239 """Manage session permissions for trusted operations to avoid repeated prompts.

240

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)

244

245 """

246 output = []

247 output.extend(("🔐 Session Permissions Management", "=" * 40))

248

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 )

267

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 )

278

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 )

295

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 )

305

306 else:

307 output.extend(

308 (

309 f"❌ Unknown action: {action}",

310 "💡 Valid actions: status, trust, revoke_all",

311 )

312 )

313

314 return "\n".join(output)

315

316

317# Compaction analysis and auto-execution functions

318 def _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

345

346

347 def _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

350

351 git_dir = current_dir / ".git"

352 if not git_dir.exists():

353 return None

354

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

371

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

391

392 return recent_commits, modified_files

393

394 except (subprocess.TimeoutExpired, Exception):

395 return None

396

397

398 def _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, ""

406

407

408 def _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

414

415 if recent_commits >= 3:

416 return (

417 True,

418 f"High development activity ({recent_commits} commits in 24h) - compaction recommended",

419 )

420

421 if modified_files >= 10:

422 return (

423 True,

424 f"Many modified files ({modified_files}) detected - context optimization beneficial",

425 )

426

427 return False, ""

428

429

430 def _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, ""

438

439

440 def _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"

443

444

445 def _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"

448

449

450 def 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

455

456 try:

457 current_dir = Path(os.environ.get("PWD", Path.cwd()))

458

459 # Count significant files in project as a complexity indicator

460 file_count = _count_significant_files(current_dir)

461

462 # Large project heuristic

463 should_compact, reason = _evaluate_large_project_heuristic(file_count)

464 if should_compact:

465 return should_compact, reason

466

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

472

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

477

478 # Default to not suggesting unless we have clear indicators

479 return False, _get_default_compaction_reason()

480

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()

484

485

486 async 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"

495

496

497# Enhanced tools with auto-compaction

498@mcp.tool()

499 async def auto_compact() -> str:

500 """Automatically trigger conversation compaction with intelligent summary."""

501 output = []

502 output.extend(("🗜️ Auto-Compaction Feature", "=" * 30))

503

504 should_compact, reason = should_suggest_compact()

505 output.append(f"📊 Analysis: {reason}")

506

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")

513

514 return "\n".join(output)

515

516

517@mcp.tool()

518 async 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)

531

532

533# Server startup

534 def run_server() -> None:

535 """Run the optimized MCP server."""

536 try:

537 logger.info("Starting optimized session-buddy server")

538

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 )

547

548 if MCP_AVAILABLE:

549 mcp.run()

550 else:

551 logger.warning("Running in mock mode - FastMCP not available")

552

553 except Exception as e:

554 logger.exception(f"Server startup failed: {e}")

555 sys.exit(1)

556

557

558 if __name__ == "__main__":

559 run_server()