Coverage for src / moai_adk / core / session_manager.py: 0.00%
193 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-11-20 20:52 +0900
1"""
2Session Manager for MoAI-ADK Agent Orchestration
4Manages sub-agent session IDs and resume logic based on official Claude Code documentation.
5Provides session tracking, result storage, and resume decision making.
7Official Documentation Reference:
8https://code.claude.com/docs/en/sub-agents
10Key Principles:
11- Sub-agents operate in isolated context windows
12- No direct agent-to-agent communication
13- Results flow through main conversation thread (Alfred)
14- Resume preserves full conversation history
15- Each execution gets unique agentId
16"""
18import json
19import logging
20from datetime import datetime
21from pathlib import Path
22from typing import Any, Dict, List, Optional
24logger = logging.getLogger(__name__)
27class SessionManager:
28 """
29 Manages sub-agent session IDs and resume logic.
31 Based on official Claude Code sub-agent pattern:
32 - Each agent execution gets unique agentId
33 - Resume parameter inherits full conversation history
34 - Session transcripts stored in agent-{agentId}.jsonl
36 Attributes:
37 _sessions: Mapping of agent_name to current agentId
38 _results: Storage of agent execution results (agentId → result data)
39 _chains: Workflow chains tracking (chain_name → [agentIds])
40 _session_file: Persistent storage location
41 _transcript_dir: Directory for conversation transcripts
42 """
44 def __init__(
45 self,
46 session_file: Optional[Path] = None,
47 transcript_dir: Optional[Path] = None,
48 ):
49 """
50 Initialize SessionManager.
52 Args:
53 session_file: Path to session storage JSON file
54 (default: .moai/memory/agent-sessions.json)
55 transcript_dir: Directory for agent transcripts
56 (default: .moai/logs/agent-transcripts/)
57 """
58 # Default paths
59 project_root = Path.cwd()
60 self._session_file = (
61 session_file or project_root / ".moai" / "memory" / "agent-sessions.json"
62 )
63 self._transcript_dir = (
64 transcript_dir or project_root / ".moai" / "logs" / "agent-transcripts"
65 )
67 # Ensure directories exist
68 self._session_file.parent.mkdir(parents=True, exist_ok=True)
69 self._transcript_dir.mkdir(parents=True, exist_ok=True)
71 # In-memory storage
72 self._sessions: Dict[str, str] = {} # agent_name → current agentId
73 self._results: Dict[str, Any] = {} # agentId → result data
74 self._chains: Dict[str, List[str]] = {} # chain_name → [agentIds]
75 self._metadata: Dict[str, Dict[str, Any]] = {} # agentId → metadata
77 # Load existing sessions
78 self._load_sessions()
80 def _load_sessions(self) -> None:
81 """Load session data from persistent storage."""
82 if self._session_file.exists():
83 try:
84 with open(self._session_file, "r", encoding="utf-8") as f:
85 data = json.load(f)
86 self._sessions = data.get("sessions", {})
87 self._chains = data.get("chains", {})
88 self._metadata = data.get("metadata", {})
89 logger.info(
90 f"Loaded {len(self._sessions)} sessions from {self._session_file}"
91 )
92 except json.JSONDecodeError as e:
93 logger.warning(f"Failed to load sessions: {e}")
94 self._sessions = {}
95 self._chains = {}
96 self._metadata = {}
98 def _save_sessions(self) -> None:
99 """Save session data to persistent storage."""
100 data = {
101 "sessions": self._sessions,
102 "chains": self._chains,
103 "metadata": self._metadata,
104 "last_updated": datetime.now().isoformat(),
105 }
107 try:
108 with open(self._session_file, "w", encoding="utf-8") as f:
109 json.dump(data, f, indent=2, ensure_ascii=False)
110 logger.debug(f"Saved sessions to {self._session_file}")
111 except IOError as e:
112 logger.error(f"Failed to save sessions: {e}")
114 def register_agent_result(
115 self,
116 agent_name: str,
117 agent_id: str,
118 result: Any,
119 chain_id: Optional[str] = None,
120 ) -> None:
121 """
122 Register agent execution result in main context.
124 This method implements the official pattern:
125 "Results flow through main conversation thread"
127 Args:
128 agent_name: Name of the agent (e.g., "tdd-implementer")
129 agent_id: Unique agentId returned from Task() execution
130 result: Result data from agent execution
131 chain_id: Optional workflow chain identifier (e.g., "SPEC-AUTH-001-implementation")
132 """
133 # Store agent ID mapping
134 self._sessions[agent_name] = agent_id
136 # Store result data
137 self._results[agent_id] = {
138 "agent_name": agent_name,
139 "result": result,
140 "timestamp": datetime.now().isoformat(),
141 "chain_id": chain_id,
142 }
144 # Track in workflow chain
145 if chain_id:
146 if chain_id not in self._chains:
147 self._chains[chain_id] = []
148 self._chains[chain_id].append(agent_id)
150 # Store metadata
151 self._metadata[agent_id] = {
152 "agent_name": agent_name,
153 "created_at": datetime.now().isoformat(),
154 "chain_id": chain_id,
155 "resume_count": 0,
156 }
158 # Persist to disk
159 self._save_sessions()
161 logger.info(
162 f"Registered agent result: {agent_name} (agentId: {agent_id[:8]}..., chain: {chain_id})"
163 )
165 def get_resume_id(
166 self,
167 agent_name: str,
168 chain_id: Optional[str] = None,
169 ) -> Optional[str]:
170 """
171 Get agentId to resume if continuing same work.
173 Official pattern:
174 - resume parameter preserves full conversation history
175 - Same agent can continue work with context
177 Args:
178 agent_name: Name of the agent to resume
179 chain_id: Optional workflow chain to resume
181 Returns:
182 agentId to resume, or None if should start new session
183 """
184 # Check if agent has previous session
185 if agent_name not in self._sessions:
186 logger.debug(f"No previous session for {agent_name}")
187 return None
189 agent_id = self._sessions[agent_name]
191 # Validate chain_id if provided
192 if chain_id:
193 metadata = self._metadata.get(agent_id, {})
194 if metadata.get("chain_id") != chain_id:
195 logger.debug(
196 f"Chain mismatch: {agent_name} was in {metadata.get('chain_id')}, "
197 f"requested {chain_id}"
198 )
199 return None
201 logger.info(f"Resume ID for {agent_name}: {agent_id[:8]}...")
202 return agent_id
204 def should_resume(
205 self,
206 agent_name: str,
207 current_task: str,
208 previous_task: Optional[str] = None,
209 ) -> bool:
210 """
211 Determine if resume or new invocation is appropriate.
213 Decision logic based on official best practices:
214 - Resume: Same agent, continuing previous task, context continuity needed
215 - New: Different agent, independent task, context switch
217 Args:
218 agent_name: Name of the agent
219 current_task: Description of current task
220 previous_task: Description of previous task (if any)
222 Returns:
223 True if should resume, False if should start new session
224 """
225 # No previous session → new
226 if agent_name not in self._sessions:
227 return False
229 # No previous task information → new
230 if not previous_task:
231 return False
233 # Check resume count (prevent infinite loops)
234 agent_id = self._sessions[agent_name]
235 metadata = self._metadata.get(agent_id, {})
236 resume_count = metadata.get("resume_count", 0)
238 if resume_count >= 5: # Max resume depth from config
239 logger.warning(
240 f"{agent_name} has been resumed {resume_count} times, starting new session"
241 )
242 return False
244 # Heuristic: Check if tasks are related
245 # (This can be enhanced with semantic similarity)
246 task_keywords_match = any(
247 keyword in current_task.lower()
248 for keyword in previous_task.lower().split()
249 if len(keyword) > 4
250 )
252 if task_keywords_match:
253 logger.info(f"Tasks appear related, resuming {agent_name}")
254 return True
256 logger.info(f"Tasks appear independent, starting new session for {agent_name}")
257 return False
259 def increment_resume_count(self, agent_id: str) -> None:
260 """
261 Increment resume count for an agent session.
263 Args:
264 agent_id: Agent session ID
265 """
266 if agent_id in self._metadata:
267 self._metadata[agent_id]["resume_count"] += 1
268 self._metadata[agent_id]["last_resumed_at"] = datetime.now().isoformat()
269 self._save_sessions()
271 def get_agent_result(self, agent_id: str) -> Optional[Any]:
272 """
273 Retrieve stored result for an agent execution.
275 Args:
276 agent_id: Agent session ID
278 Returns:
279 Stored result data, or None if not found
280 """
281 result_data = self._results.get(agent_id)
282 if result_data:
283 return result_data["result"]
284 return None
286 def get_chain_results(self, chain_id: str) -> List[Dict[str, Any]]:
287 """
288 Get all agent results in a workflow chain.
290 Args:
291 chain_id: Workflow chain identifier
293 Returns:
294 List of result dictionaries in execution order
295 """
296 if chain_id not in self._chains:
297 return []
299 agent_ids = self._chains[chain_id]
300 results = []
302 for agent_id in agent_ids:
303 if agent_id in self._results:
304 results.append(self._results[agent_id])
306 return results
308 def get_chain_summary(self, chain_id: str) -> Dict[str, Any]:
309 """
310 Get summary of a workflow chain.
312 Args:
313 chain_id: Workflow chain identifier
315 Returns:
316 Summary dictionary with agent names, timestamps, etc.
317 """
318 results = self.get_chain_results(chain_id)
320 if not results:
321 return {"chain_id": chain_id, "status": "not_found"}
323 return {
324 "chain_id": chain_id,
325 "agent_count": len(results),
326 "agents": [r["agent_name"] for r in results],
327 "started_at": results[0]["timestamp"] if results else None,
328 "completed_at": results[-1]["timestamp"] if results else None,
329 "status": "completed",
330 }
332 def clear_agent_session(self, agent_name: str) -> None:
333 """
334 Clear session data for a specific agent.
336 Use when you want to force a new session for an agent.
338 Args:
339 agent_name: Name of the agent
340 """
341 if agent_name in self._sessions:
342 agent_id = self._sessions[agent_name]
343 del self._sessions[agent_name]
345 if agent_id in self._results:
346 del self._results[agent_id]
348 if agent_id in self._metadata:
349 del self._metadata[agent_id]
351 self._save_sessions()
352 logger.info(f"Cleared session for {agent_name}")
354 def clear_chain(self, chain_id: str) -> None:
355 """
356 Clear all sessions in a workflow chain.
358 Args:
359 chain_id: Workflow chain identifier
360 """
361 if chain_id in self._chains:
362 agent_ids = self._chains[chain_id]
364 for agent_id in agent_ids:
365 if agent_id in self._results:
366 del self._results[agent_id]
367 if agent_id in self._metadata:
368 del self._metadata[agent_id]
370 del self._chains[chain_id]
371 self._save_sessions()
372 logger.info(f"Cleared chain: {chain_id}")
374 def get_all_sessions(self) -> Dict[str, Any]:
375 """
376 Get all active sessions.
378 Returns:
379 Dictionary with all session data
380 """
381 return {
382 "sessions": self._sessions,
383 "chains": list(self._chains.keys()),
384 "total_results": len(self._results),
385 }
387 def export_transcript(self, agent_id: str) -> Optional[Path]:
388 """
389 Get path to agent conversation transcript.
391 Official pattern:
392 - Transcripts stored in agent-{agentId}.jsonl
393 - Contains full conversation history
395 Args:
396 agent_id: Agent session ID
398 Returns:
399 Path to transcript file, or None if not found
400 """
401 transcript_file = self._transcript_dir / f"agent-{agent_id}.jsonl"
403 if transcript_file.exists():
404 return transcript_file
406 logger.warning(f"Transcript not found for agentId: {agent_id}")
407 return None
409 def create_chain(
410 self,
411 chain_id: str,
412 agent_sequence: List[str],
413 metadata: Optional[Dict[str, Any]] = None,
414 ) -> None:
415 """
416 Create a new workflow chain.
418 Args:
419 chain_id: Unique chain identifier (e.g., "SPEC-AUTH-001-implementation")
420 agent_sequence: Expected agent execution order
421 metadata: Optional metadata for the chain
422 """
423 self._chains[chain_id] = []
425 chain_metadata = {
426 "created_at": datetime.now().isoformat(),
427 "expected_sequence": agent_sequence,
428 "metadata": metadata or {},
429 }
431 # Store in a separate chains metadata file
432 chains_file = self._session_file.parent / "workflow-chains.json"
434 if chains_file.exists():
435 with open(chains_file, "r", encoding="utf-8") as f:
436 chains_data = json.load(f)
437 else:
438 chains_data = {}
440 chains_data[chain_id] = chain_metadata
442 with open(chains_file, "w", encoding="utf-8") as f:
443 json.dump(chains_data, f, indent=2, ensure_ascii=False)
445 logger.info(
446 f"Created workflow chain: {chain_id} with {len(agent_sequence)} agents"
447 )
450# Global instance (singleton pattern)
451_session_manager_instance: Optional[SessionManager] = None
454def get_session_manager() -> SessionManager:
455 """
456 Get global SessionManager instance (singleton).
458 Returns:
459 SessionManager instance
460 """
461 global _session_manager_instance
463 if _session_manager_instance is None:
464 _session_manager_instance = SessionManager()
466 return _session_manager_instance
469# Convenience functions for direct use
472def register_agent(
473 agent_name: str,
474 agent_id: str,
475 result: Any,
476 chain_id: Optional[str] = None,
477) -> None:
478 """
479 Convenience function to register agent result.
481 Args:
482 agent_name: Name of the agent
483 agent_id: Unique agentId from Task() execution
484 result: Result data
485 chain_id: Optional workflow chain identifier
486 """
487 manager = get_session_manager()
488 manager.register_agent_result(agent_name, agent_id, result, chain_id)
491def get_resume_id(agent_name: str, chain_id: Optional[str] = None) -> Optional[str]:
492 """
493 Convenience function to get resume ID.
495 Args:
496 agent_name: Name of the agent
497 chain_id: Optional workflow chain
499 Returns:
500 agentId to resume, or None
501 """
502 manager = get_session_manager()
503 return manager.get_resume_id(agent_name, chain_id)
506def should_resume(
507 agent_name: str,
508 current_task: str,
509 previous_task: Optional[str] = None,
510) -> bool:
511 """
512 Convenience function to check if should resume.
514 Args:
515 agent_name: Name of the agent
516 current_task: Description of current task
517 previous_task: Description of previous task
519 Returns:
520 True if should resume
521 """
522 manager = get_session_manager()
523 return manager.should_resume(agent_name, current_task, previous_task)
526# Example usage for documentation
527if __name__ == "__main__":
528 """
529 Example usage of SessionManager.
531 This demonstrates the official Claude Code sub-agent patterns.
532 """
534 # Initialize manager
535 manager = SessionManager()
537 # Example 1: Linear Chain (spec-builder → implementation-planner)
538 print("=== Example 1: Linear Chain ===")
540 # Create workflow chain
541 manager.create_chain(
542 chain_id="SPEC-AUTH-001-planning",
543 agent_sequence=["spec-builder", "implementation-planner"],
544 metadata={"spec_id": "SPEC-AUTH-001", "feature": "User Authentication"},
545 )
547 # Simulate spec-builder execution
548 spec_result = {
549 "spec_id": "SPEC-AUTH-001",
550 "files_created": [".moai/specs/SPEC-AUTH-001/spec.md"],
551 "status": "success",
552 }
554 manager.register_agent_result(
555 agent_name="spec-builder",
556 agent_id="spec-abc123",
557 result=spec_result,
558 chain_id="SPEC-AUTH-001-planning",
559 )
561 # Simulate implementation-planner execution
562 plan_result = {
563 "dependencies": {"fastapi": ">=0.118.3"},
564 "status": "success",
565 }
567 manager.register_agent_result(
568 agent_name="implementation-planner",
569 agent_id="plan-def456",
570 result=plan_result,
571 chain_id="SPEC-AUTH-001-planning",
572 )
574 # Get chain summary
575 summary = manager.get_chain_summary("SPEC-AUTH-001-planning")
576 print(f"Chain summary: {json.dumps(summary, indent=2)}")
578 # Example 2: Resume Pattern (tdd-implementer continues work)
579 print("\n=== Example 2: Resume Pattern ===")
581 manager.create_chain(
582 chain_id="SPEC-AUTH-001-implementation",
583 agent_sequence=["tdd-implementer"],
584 )
586 # First execution: Implementation phase 1
587 implementation_001_result = {
588 "phase": "phase_1",
589 "tests_created": ["tests/test_registration.py"],
590 "code_created": ["src/auth/registration.py"],
591 "status": "success",
592 }
594 manager.register_agent_result(
595 agent_name="tdd-implementer",
596 agent_id="tdd-ghi789",
597 result=implementation_001_result,
598 chain_id="SPEC-AUTH-001-implementation",
599 )
601 # Get resume ID for continuing work
602 resume_id = manager.get_resume_id(
603 agent_name="tdd-implementer",
604 chain_id="SPEC-AUTH-001-implementation",
605 )
607 print(f"Resume ID for tdd-implementer: {resume_id}")
609 # Should resume? (continuing user auth flow)
610 should_resume_decision = manager.should_resume(
611 agent_name="tdd-implementer",
612 current_task="Implement user login endpoint",
613 previous_task="Implement user registration endpoint",
614 )
616 print(f"Should resume? {should_resume_decision}")
618 if should_resume_decision and resume_id:
619 print(f"✅ Resume with agentId: {resume_id}")
620 manager.increment_resume_count(resume_id)
621 else:
622 print("❌ Start new session")
624 # Example 3: Parallel Analysis
625 print("\n=== Example 3: Parallel Analysis ===")
627 manager.create_chain(
628 chain_id="SPEC-AUTH-001-review",
629 agent_sequence=["backend-expert", "security-expert", "frontend-expert"],
630 metadata={"review_type": "expert_consultation"},
631 )
633 # All experts run independently (no resume)
634 experts_results = {
635 "backend-expert": {
636 "recommendations": ["Use JWT for auth"],
637 "agent_id": "backend-jkl012",
638 },
639 "security-expert": {
640 "vulnerabilities": ["Rate limiting needed"],
641 "agent_id": "security-mno345",
642 },
643 "frontend-expert": {
644 "ui_concerns": ["Token refresh flow"],
645 "agent_id": "frontend-pqr678",
646 },
647 }
649 for expert_name, data in experts_results.items():
650 manager.register_agent_result(
651 agent_name=expert_name,
652 agent_id=data["agent_id"],
653 result={k: v for k, v in data.items() if k != "agent_id"},
654 chain_id="SPEC-AUTH-001-review",
655 )
657 # Get all review results
658 review_results = manager.get_chain_results("SPEC-AUTH-001-review")
659 print(f"Expert reviews: {len(review_results)} experts")
661 for result in review_results:
662 print(f" - {result['agent_name']}: {list(result['result'].keys())}")
664 # Get all sessions
665 print("\n=== All Sessions ===")
666 all_sessions = manager.get_all_sessions()
667 print(json.dumps(all_sessions, indent=2))