Coverage for src / osiris_cli / session.py: 0%
266 statements
« prev ^ index » next coverage.py v7.13.0, created at 2025-12-31 05:01 +0200
1import json
2from pathlib import Path
3from datetime import datetime
4from typing import List, Dict, Optional, Any
5from .config import settings, CONFIG_DIR
6from .osiris_config import get_osiris_config
8SESSION_DIR = CONFIG_DIR / "sessions"
10class SessionManager:
11 def __init__(self):
12 SESSION_DIR.mkdir(parents=True, exist_ok=True)
13 self.current_session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
14 self.messages: List[Dict[str, Any]] = []
15 self.context_files: List[str] = []
16 self.allowed_tools: List[str] = []
17 self.forbidden_tools: List[str] = []
18 self.total_tokens = 0
19 self.last_command_palette_filter = ""
20 self.tags: List[str] = [] # Session tags for organization
21 self.description: str = "" # Session description
22 self.provider: str = "" # Provider used for this session
23 self.model: str = "" # Model used for this session
24 self.activity_log: List[Dict[str, str]] = []
25 self.plan_history: List[Dict[str, str]] = []
26 self._capabilities: Dict[str, Dict[str, bool]] = {}
27 self._tool_policy_allowed = None
28 self._reset_tool_policy()
30 def _reset_tool_policy(self):
31 config = get_osiris_config()
32 self.allowed_tools = list(config.allowed_tools)
33 self.forbidden_tools = list(config.blocked_tools)
34 self._capabilities.clear()
36 def add_message(self, role: str, content: str, tool_calls=None):
37 msg: Dict[str, Any] = {"role": role, "content": content}
38 if tool_calls:
39 # Serialize tool calls for JSON storage if they are objects
40 # Assuming tool_calls is a list of objects from OpenAI SDK
41 # We need to store them as dicts
42 try:
43 serialized_calls = []
44 for tc in tool_calls:
45 if hasattr(tc, 'id') and hasattr(tc, 'type'):
46 # OpenAI object format
47 call_dict = {
48 "id": tc.id,
49 "type": tc.type,
50 "function": {
51 "name": tc.function.name,
52 "arguments": tc.function.arguments
53 }
54 }
55 else:
56 # Already dict format
57 call_dict = tc
58 serialized_calls.append(call_dict)
59 msg["tool_calls"] = serialized_calls
60 except (AttributeError, TypeError):
61 # Handle any format issues
62 msg["tool_calls"] = []
63 self.messages.append(msg)
64 self.save_session()
66 def add_tool_output(self, tool_call_id: str, name: str, output: str):
67 msg = {
68 "role": "tool",
69 "tool_call_id": tool_call_id,
70 "name": name,
71 "content": str(output)
72 }
73 self.messages.append(msg)
74 self.save_session()
76 def add_activity(self, description: str):
77 self.activity_log.append({
78 "timestamp": datetime.now().isoformat(),
79 "description": description
80 })
81 self.save_session()
83 def add_plan_snapshot(self, plan_text: str):
84 self.plan_history.append({
85 "timestamp": datetime.now().isoformat(),
86 "plan": plan_text
87 })
88 self.save_session()
90 def add_usage(self, tokens: int):
91 if tokens:
92 self.total_tokens += tokens
93 self.save_session()
95 def _generate_preview(self) -> str:
96 """Generate a preview of the session content"""
97 for msg in reversed(self.messages):
98 if msg.get("role") == "user" and msg.get("content"):
99 preview = msg["content"][:60].replace("\n", " ")
100 return f"[{self.total_tokens}t] {preview}"
101 return f"[{self.total_tokens}t] Empty session"
103 def save_session(self, tags: Optional[List[str]] = None, description: str = ""):
104 """Enhanced save session with metadata"""
105 if tags:
106 self.tags = tags
107 if description:
108 self.description = description
110 file_path = SESSION_DIR / f"{self.current_session_id}.json"
111 data = {
112 "id": self.current_session_id,
113 "timestamp": datetime.now().isoformat(),
114 "model": getattr(settings, 'default_model', 'unknown'),
115 "provider": getattr(settings, 'provider', 'unknown'),
116 "context_files": self.context_files,
117 "messages": self.messages,
118 "allowed_tools": self.allowed_tools,
119 "forbidden_tools": self.forbidden_tools,
120 "total_tokens": self.total_tokens,
121 "last_command_palette_filter": self.last_command_palette_filter,
122 "tags": self.tags,
123 "description": self.description,
124 "message_count": len(self.messages),
125 "tool_calls": sum(1 for msg in self.messages if msg.get('role') == 'tool'),
126 "preview": self._generate_preview(),
127 "activity_log": self.activity_log,
128 "plan_history": self.plan_history
129 }
130 with open(file_path, "w") as f:
131 json.dump(data, f, indent=2)
133 def load_last_session(self) -> bool:
134 """Load the most recent session. Returns True if successful."""
135 files = sorted(SESSION_DIR.glob("*.json"), reverse=True)
136 if not files:
137 return False
138 return self.load_session(files[0].stem)
140 def delete_session(self, session_id: str) -> bool:
141 """Delete a session by ID."""
142 file_path = SESSION_DIR / f"{session_id}.json"
143 if file_path.exists():
144 try:
145 file_path.unlink()
146 return True
147 except Exception:
148 return False
149 return False
151 def load_session(self, session_id: str) -> bool:
152 """Load a specific session by ID."""
153 file_path = SESSION_DIR / f"{session_id}.json"
154 if not file_path.exists():
155 return False
157 try:
158 with open(file_path, "r") as f:
159 data = json.load(f)
160 self.current_session_id = data["id"]
161 self.messages = data.get("messages", [])
162 self.context_files = data.get("context_files", [])
163 self.allowed_tools = data.get("allowed_tools", [])
164 self.forbidden_tools = data.get("forbidden_tools", [])
165 self.total_tokens = data.get("total_tokens", 0)
166 self.last_command_palette_filter = data.get("last_command_palette_filter", "")
167 self.activity_log = data.get("activity_log", [])
168 self.plan_history = data.get("plan_history", [])
169 return True
170 except Exception:
171 return False
173 def list_sessions(self) -> List[str]:
174 """List all session IDs"""
175 files = sorted(SESSION_DIR.glob("*.json"), reverse=True)
176 return [f.stem for f in files[:10]] # Return top 10
178 def list_session_meta(self) -> List[Dict]:
179 """Return list of session metadata for menus."""
180 files = sorted(SESSION_DIR.glob("*.json"), reverse=True)
181 meta = []
182 for p in files[:25]:
183 try:
184 data = json.loads(p.read_text())
185 sid = data.get("id", p.stem)
186 ts = data.get("timestamp", "")
187 msgs = data.get("messages", [])
188 toks = data.get("total_tokens", 0)
189 preview = ""
190 for m in msgs:
191 if m.get("role") == "user":
192 preview = m.get("content", "")[:60].replace("\n", " ")
193 break
194 # Format preview to include tokens
195 display_preview = f"[{toks}t] {preview}"
196 meta.append({"id": sid, "timestamp": ts, "preview": display_preview, "total_tokens": toks})
197 except Exception:
198 continue
199 return meta
201 def clear(self):
202 self.messages = []
203 self.context_files = []
204 self._reset_tool_policy()
205 self._capabilities.clear()
206 self.total_tokens = 0
207 self.last_command_palette_filter = ""
208 self.activity_log = []
209 self.plan_history = []
210 # Start new session file
211 self.current_session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
212 self.save_session()
214 def enforce_tool_policy(self, tool_name: str):
215 if not self.is_tool_allowed(tool_name):
216 raise PermissionError(f"Tool '{tool_name}' is blocked by the session policy.")
218 def forbid_tool(self, tool_name: str):
219 if tool_name not in self.forbidden_tools:
220 self.forbidden_tools.append(tool_name)
221 self.save_session()
223 def remove_forbidden_tool(self, tool_name: str):
224 if tool_name in self.forbidden_tools:
225 self.forbidden_tools.remove(tool_name)
226 self.save_session()
228 def remember_tool(self, tool_name: str):
229 if tool_name not in self.allowed_tools:
230 self.allowed_tools.append(tool_name)
231 self.save_session()
233 def is_tool_allowed(self, tool_name: str) -> bool:
234 if self.forbidden_tools and tool_name in self.forbidden_tools:
235 return False
236 if self.allowed_tools:
237 return tool_name in self.allowed_tools
238 return True
240 def apply_tool_policy(self, allowed: Optional[List[str]] = None, forbidden: Optional[List[str]] = None) -> bool:
241 updated = False
243 if allowed is not None:
244 normalized = []
245 for tool in allowed:
246 if tool and tool not in normalized:
247 normalized.append(tool)
248 self.allowed_tools = normalized
249 updated = True
251 if forbidden:
252 for tool in forbidden:
253 if tool and tool not in self.forbidden_tools:
254 self.forbidden_tools.append(tool)
255 updated = True
257 if updated:
258 self.save_session()
259 return updated
261 def provider_capability(self, provider: str, model: str, capability: str) -> bool:
262 if not provider or not model:
263 return True
264 key = f"{provider}:{model}"
265 entry = self._capabilities.get(key)
266 if not entry:
267 return True
268 return entry.get(capability, True)
270 def mark_provider_capability(self, provider: str, model: str, capability: str, value: bool):
271 if not provider or not model:
272 return
273 key = f"{provider}:{model}"
274 entry = self._capabilities.setdefault(key, {})
275 entry[capability] = value
277 def provider_supports_tools(self, provider: str, model: str) -> bool:
278 return self.provider_capability(provider, model, "tools")
280 def mark_provider_tool_support(self, provider: str, model: str, supports: bool):
281 self.mark_provider_capability(provider, model, "tools", supports)
283 def reset_provider_capabilities(self):
284 self._capabilities.clear()
286 def undo_last_step(self) -> bool:
287 """
288 Solid Undo: Removes the last 'step' of conversation history.
289 - If last was Assistant text -> Remove it.
290 - If last was Tool Output -> Remove it AND the preceding Assistant Tool Call (and any sibling tool outputs).
291 """
292 if not self.messages: return False
294 last = self.messages.pop()
296 # If we removed a tool output, we must clean up the entire tool interaction chain
297 if last.get("role") == "tool":
298 # 1. Remove any other adjacent tool outputs (parallel execution case)
299 while self.messages and self.messages[-1].get("role") == "tool":
300 self.messages.pop()
302 # 2. Remove the Assistant message that triggered these tools
303 if self.messages and self.messages[-1].get("role") == "assistant" and self.messages[-1].get("tool_calls"):
304 self.messages.pop()
306 self.save_session()
307 return True
309 def search_sessions(self, query: str = "", tags: Optional[List[str]] = None,
310 provider: str = "", model: str = "") -> List[Dict]:
311 """Advanced session search with filtering"""
312 all_sessions = self.list_session_meta()
314 if not query and not tags and not provider and not model:
315 return all_sessions
317 filtered = []
318 query = query.lower()
320 for session in all_sessions:
321 # Text search in preview, description, and ID
322 text_match = (query in session.get('preview', '').lower() or
323 query in session.get('description', '').lower() or
324 query in session.get('id', '').lower())
326 # Tag filtering
327 tag_match = True
328 if tags:
329 session_tags = session.get('tags', [])
330 tag_match = any(tag in session_tags for tag in tags)
332 # Provider filtering
333 provider_match = True
334 if provider:
335 provider_match = session.get('provider', '').lower() == provider.lower()
337 # Model filtering
338 model_match = True
339 if model:
340 model_match = session.get('model', '').lower() == model.lower()
342 if text_match and tag_match and provider_match and model_match:
343 filtered.append(session)
345 return filtered
347 def add_tags(self, tags: List[str]):
348 """Add tags to current session"""
349 for tag in tags:
350 if tag not in self.tags:
351 self.tags.append(tag)
353 def remove_tags(self, tags: List[str]):
354 """Remove tags from current session"""
355 self.tags = [tag for tag in self.tags if tag not in tags]
357 def set_description(self, description: str):
358 """Set session description"""
359 self.description = description
361 def get_session_stats(self) -> Dict:
362 """Get comprehensive session statistics"""
363 meta = self.list_session_meta()
364 total_sessions = len(meta)
366 if not meta:
367 return {"total_sessions": 0}
369 providers = {}
370 models = {}
371 total_messages = 0
372 total_tokens = 0
373 total_tools = 0
375 for session in meta:
376 provider = session.get('provider', 'unknown')
377 model = session.get('model', 'unknown')
378 messages = session.get('message_count', 0)
379 tokens = session.get('total_tokens', 0)
380 tools = session.get('tool_calls', 0)
382 providers[provider] = providers.get(provider, 0) + 1
383 models[model] = models.get(model, 0) + 1
384 total_messages += messages
385 total_tokens += tokens
386 total_tools += tools
388 return {
389 "total_sessions": total_sessions,
390 "providers": providers,
391 "models": models,
392 "total_messages": total_messages,
393 "total_tokens": total_tokens,
394 "total_tool_calls": total_tools,
395 "avg_messages_per_session": total_messages / total_sessions if total_sessions > 0 else 0,
396 "avg_tokens_per_session": total_tokens / total_sessions if total_sessions > 0 else 0
397 }
399 def dump_context(self, target_path: Path):
400 """Serialize current Memory Map for sub-agent pass-through."""
401 data = {
402 "messages": self.messages,
403 "context_files": self.context_files,
404 "allowed_tools": self.allowed_tools,
405 "total_tokens": self.total_tokens,
406 "parent_session_id": self.current_session_id
407 }
408 with open(target_path, "w") as f:
409 json.dump(data, f, indent=2)
411session = SessionManager()