Coverage for src / harnessutils / manager.py: 96%
279 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-03-12 18:52 -0600
1"""Main ConversationManager API for harness-utils."""
3import time
4from pathlib import Path
5from typing import Any
7from harnessutils.compaction.pruning import PruningDecision, prune_tool_outputs
8from harnessutils.compaction.summarization import is_overflow, summarize_conversation
9from harnessutils.compaction.truncation import truncate_output
10from harnessutils.config import HarnessConfig, StorageConfig
11from harnessutils.conversion.to_model import to_model_messages
12from harnessutils.hooks import MessageHooks
13from harnessutils.inspection import ContextInspector
14from harnessutils.memory import SemanticMemoryBackend
15from harnessutils.models.conversation import Conversation
16from harnessutils.models.message import Message
17from harnessutils.models.usage import Usage
18from harnessutils.snapshots import Snapshot, SnapshotDiff, SnapshotManager
19from harnessutils.storage.filesystem import FilesystemStorage
20from harnessutils.storage.memory import MemoryStorage
21from harnessutils.types import LLMClient, StorageBackend
22from harnessutils.utils.ids import generate_id
23from harnessutils.workspace import resolve_workspace
26class ConversationManager:
27 """Main interface for managing conversations with context window management.
29 Provides high-level API for:
30 - Creating and managing conversations
31 - Adding messages
32 - Automatic context compaction (truncation, pruning, summarization)
33 - Message storage and retrieval
34 """
36 def __init__(
37 self,
38 storage: StorageBackend | None = None,
39 config: HarnessConfig | None = None,
40 *,
41 workspace_root: Path | str | None = None,
42 semantic_memory: SemanticMemoryBackend | None = None,
43 message_hooks: MessageHooks | None = None,
44 ):
45 """Initialize conversation manager.
47 Args:
48 storage: Storage backend (uses in-memory if None)
49 config: Configuration (uses defaults if None)
50 workspace_root: Optional workspace root. Creates .harness/ dir with stable
51 project UUID. Defaults storage to FilesystemStorage under .harness/sessions/.
52 semantic_memory: Optional semantic memory backend. Stored as attribute;
53 harness-utils does not call it internally.
54 message_hooks: Optional hooks for the add_message() lifecycle.
55 on_before_add_message fires before storage and can modify or reject.
56 on_after_add_message fires after storage for side effects.
57 """
58 self.config = config or HarnessConfig()
59 self.project_id: str | None = None
60 self.semantic_memory = semantic_memory
62 if workspace_root is not None:
63 harness_dir, self.project_id = resolve_workspace(Path(workspace_root))
64 if storage is None:
65 sessions_path = harness_dir / "sessions"
66 storage_config = StorageConfig(base_path=sessions_path)
67 storage = FilesystemStorage(storage_config)
69 self.storage = storage or MemoryStorage()
70 self.message_hooks = message_hooks
71 self._message_cache: dict[str, list[Message]] = {}
72 # Pass storage (not self.storage) so SnapshotManager receives None when no
73 # explicit backend was given — it falls back to its internal in-memory dict.
74 self.snapshot_manager = SnapshotManager(storage)
76 def create_conversation(
77 self,
78 conversation_id: str | None = None,
79 project_id: str | None = None,
80 ) -> Conversation:
81 """Create a new conversation.
83 Args:
84 conversation_id: Optional conversation ID (generated if None)
85 project_id: Optional project ID for grouping
87 Returns:
88 New conversation object
89 """
90 if conversation_id is None:
91 conversation_id = generate_id("conv")
93 now = int(time.time() * 1000)
94 conversation = Conversation(
95 id=conversation_id,
96 project_id=project_id,
97 created=now,
98 updated=now,
99 )
101 self.storage.save_conversation(conversation_id, conversation.to_dict())
102 self._message_cache[conversation_id] = []
104 return conversation
106 def add_message(self, conversation_id: str, message: Message) -> None:
107 """Add a message to a conversation.
109 If message_hooks.on_before_add_message is set, it is called first.
110 Its return value replaces the message that will be stored. Raising
111 from it rejects the message entirely (nothing is stored).
113 If message_hooks.on_after_add_message is set, it is called after
114 successful storage. Its return value is ignored.
116 Args:
117 conversation_id: Conversation to add message to
118 message: Message to add
119 """
120 # Pre-hook — may modify message or raise to reject
121 if self.message_hooks and self.message_hooks.on_before_add_message:
122 message = self.message_hooks.on_before_add_message(conversation_id, message)
124 # Inject timestamp if not set — required for cleanup_stale_data to work
125 if "timestamp" not in message.metadata:
126 message.metadata["timestamp"] = int(time.time() * 1000)
128 self.storage.save_message(conversation_id, message.id, message.to_dict())
130 if conversation_id not in self._message_cache:
131 self._message_cache[conversation_id] = []
132 self._message_cache[conversation_id].append(message)
134 # Load conversation and update
135 conv_data = self.storage.load_conversation(conversation_id)
136 conv = Conversation.from_dict(conv_data)
137 conv.updated = int(time.time() * 1000)
139 # Track velocity if message has token count
140 if message.tokens:
141 tokens_added = message.tokens.total
142 conv.update_velocity(tokens_added)
144 self.storage.save_conversation(conversation_id, conv.to_dict())
146 # Post-hook — side effects only, message already persisted
147 if self.message_hooks and self.message_hooks.on_after_add_message:
148 self.message_hooks.on_after_add_message(conversation_id, message)
150 def get_messages(self, conversation_id: str) -> list[Message]:
151 """Get all messages for a conversation.
153 Args:
154 conversation_id: Conversation ID
156 Returns:
157 List of messages in chronological order
158 """
159 if conversation_id in self._message_cache:
160 return self._message_cache[conversation_id]
162 message_ids = self.storage.list_messages(conversation_id)
163 messages = [
164 Message.from_dict(self.storage.load_message(conversation_id, msg_id))
165 for msg_id in message_ids
166 ]
168 self._message_cache[conversation_id] = messages
169 return messages
171 def query_messages(
172 self,
173 conversation_id: str,
174 limit: int | None = None,
175 offset: int = 0,
176 order: str = "asc",
177 after: int | None = None,
178 before: int | None = None,
179 filter: dict[str, Any] | None = None,
180 ) -> list[Message]:
181 """Query messages with filtering and pagination.
183 Args:
184 conversation_id: Conversation to query
185 limit: Maximum messages to return
186 offset: Skip first N messages
187 order: "asc" (oldest first) or "desc" (newest first)
188 after: Unix ms timestamp - messages after this time
189 before: Unix ms timestamp - messages before this time
190 filter: Filter criteria dict:
191 - has_errors (bool): Only messages with errors
192 - has_warnings (bool): Only messages with warnings
193 - min_importance (float): Minimum importance score
194 - tools (list[str]): Specific tool types
195 - roles (list[str]): Message roles (user/assistant)
196 - has_tool_outputs (bool): Only messages with tool outputs
197 - is_summary (bool): Only summary messages
199 Returns:
200 Filtered and paginated messages
202 Examples:
203 # Get last 10 messages
204 messages = manager.query_messages(conv_id, limit=10, order="desc")
206 # Get messages with errors
207 messages = manager.query_messages(
208 conv_id,
209 filter={"has_errors": True}
210 )
212 # Get recent high-importance messages
213 messages = manager.query_messages(
214 conv_id,
215 filter={"min_importance": 50.0},
216 after=timestamp_24h_ago,
217 limit=20
218 )
219 """
220 from harnessutils.query import MessageFilter, QueryOptions, query_messages
222 # Build filter
223 msg_filter = None
224 if filter is not None:
225 msg_filter = MessageFilter(
226 has_errors=filter.get("has_errors"),
227 has_warnings=filter.get("has_warnings"),
228 min_importance=filter.get("min_importance"),
229 tools=filter.get("tools"),
230 roles=filter.get("roles"),
231 has_tool_outputs=filter.get("has_tool_outputs"),
232 is_summary=filter.get("is_summary"),
233 )
235 # Build query options
236 options = QueryOptions(
237 limit=limit,
238 offset=offset,
239 order=order, # type: ignore
240 after=after,
241 before=before,
242 filter=msg_filter,
243 )
245 # Get all messages and query
246 messages = self.get_messages(conversation_id)
247 return query_messages(messages, options)
249 def get_context_summary(
250 self, conversation_id: str, recent_limit: int = 5
251 ) -> dict[str, Any]:
252 """Get lightweight context summary without loading full content.
254 Provides overview of conversation state including:
255 - Message count and token usage
256 - Summary messages
257 - Recent activity (last N messages)
258 - Key messages (high importance)
259 - Error messages
261 Args:
262 conversation_id: Conversation to summarize
263 recent_limit: Number of recent messages to include (default: 5)
265 Returns:
266 Dictionary with context summary:
267 - conversation_id (str)
268 - message_count (int)
269 - total_tokens (int)
270 - summaries (list): Summary messages
271 - recent_activity (list): Last N message summaries
272 - key_messages (list): High-importance messages
273 - errors (list): Error messages
275 Examples:
276 # Get quick overview
277 summary = manager.get_context_summary(conv_id)
278 print(f"Messages: {summary['message_count']}")
279 print(f"Tokens: {summary['total_tokens']}")
280 print(f"Errors: {len(summary['errors'])}")
281 """
282 from harnessutils.query import build_context_summary
284 messages = self.get_messages(conversation_id)
285 summary = build_context_summary(conversation_id, messages, recent_limit)
286 return summary.to_dict()
288 def inspect_context(self, conversation_id: str) -> ContextInspector:
289 """Create inspector for querying context state.
291 Provides observability into:
292 - What's currently in context
293 - Token counts and breakdowns
294 - Pruning predictions
295 - Audit trail of decisions
297 Args:
298 conversation_id: Conversation to inspect
300 Returns:
301 ContextInspector instance with full query capabilities
303 Example:
304 >>> inspector = manager.inspect_context(conv_id)
305 >>> summary = inspector.summary()
306 >>> print(f"Total tokens: {summary['total_tokens']}")
307 >>> impact = inspector.predict_impact(5000)
308 >>> if impact['would_trigger_pruning']:
309 ... print(f"Would prune ~{impact['estimated_pruned_count']} outputs")
310 """
311 messages = self.get_messages(conversation_id)
313 # Load conversation metadata
314 conv_data = self.storage.load_conversation(conversation_id)
315 conversation = Conversation.from_dict(conv_data)
317 # Load persisted pruning decisions for audit trail
318 decisions_data = conversation.metadata.get("latest_pruning_decisions", [])
319 decisions = [PruningDecision.from_dict(d) for d in decisions_data]
321 return ContextInspector(messages, self.config, conversation, decisions)
323 def prune_before_turn(
324 self,
325 conversation_id: str,
326 auto_mode: bool = False,
327 ) -> dict[str, Any]:
328 """Proactively prune old tool outputs before processing a turn.
330 This is Tier 2 compaction - removes old tool outputs while
331 preserving conversation structure.
333 Args:
334 conversation_id: Conversation to prune
335 auto_mode: Whether this was auto-triggered
337 Returns:
338 Detailed pruning result with token tracking and breakdown:
339 - pruned: Total outputs removed
340 - tokens_saved: Total tokens saved
341 - tokens_before: Token count before pruning
342 - tokens_after: Token count after pruning
343 - duplicates_pruned: Outputs removed due to duplication
344 - importance_pruned: Outputs removed due to low importance
345 - duplicate_tokens_saved: Tokens saved from deduplication
346 - importance_tokens_saved: Tokens saved from importance pruning
347 - reduction_percent: Percentage reduction in token usage
348 """
349 if not self.config.compaction.prune and auto_mode:
350 return {
351 "pruned": 0,
352 "tokens_saved": 0,
353 "tokens_before": 0,
354 "tokens_after": 0,
355 "duplicates_pruned": 0,
356 "importance_pruned": 0,
357 "duplicate_tokens_saved": 0,
358 "importance_tokens_saved": 0,
359 "reduction_percent": 0,
360 }
362 messages = self.get_messages(conversation_id)
363 result = prune_tool_outputs(
364 messages,
365 self.config.pruning,
366 )
368 for msg in messages:
369 self.storage.save_message(conversation_id, msg.id, msg.to_dict())
371 # Update cache with pruned messages — parts are preserved since
372 # prune_tool_outputs only clears tool output strings in-place
373 self._message_cache[conversation_id] = messages
375 # Persist decisions so audit trail and quality metrics work
376 conv_data = self.storage.load_conversation(conversation_id)
377 conv = Conversation.from_dict(conv_data)
378 conv.metadata["latest_pruning_decisions"] = [d.to_dict() for d in result.decisions]
379 self.storage.save_conversation(conversation_id, conv.to_dict())
381 return result.to_dict()
383 def predict_overflow(
384 self,
385 conversation_id: str,
386 current_usage: Usage,
387 ) -> bool:
388 """Predict if conversation will overflow in next N turns.
390 Args:
391 conversation_id: Conversation to check
392 current_usage: Current token usage
394 Returns:
395 True if overflow predicted within lookahead window
396 """
397 if not self.config.compaction.use_predictive:
398 return False
400 # Load conversation and get velocity
401 conv_data = self.storage.load_conversation(conversation_id)
402 conv = Conversation.from_dict(conv_data)
403 velocity = conv.get_velocity()
405 if velocity is None or not velocity.turn_deltas:
406 return False # No velocity data yet
408 # Project tokens ahead
409 lookahead = self.config.compaction.predictive_lookahead
410 predicted_growth = velocity.predict_tokens_ahead(lookahead)
412 # Calculate current total and projected total
413 current_total = current_usage.input + current_usage.cache.read
414 projected_total = current_total + predicted_growth
416 # Check against safety margin
417 usable_space = (
418 self.config.model_limits.default_context_limit
419 - self.config.model_limits.default_output_limit
420 )
421 safety_threshold = usable_space * self.config.compaction.predictive_safety_margin
423 return projected_total > safety_threshold
425 def needs_compaction(
426 self,
427 conversation_id: str,
428 usage: Usage,
429 ) -> bool:
430 """Check if conversation needs summarization (Tier 3).
432 Uses both reactive (overflow) and predictive checks.
434 Args:
435 conversation_id: Conversation to check
436 usage: Token usage from last turn
438 Returns:
439 True if summarization needed
440 """
441 # Reactive check: already overflowed
442 if is_overflow(
443 usage,
444 self.config.model_limits.default_context_limit,
445 self.config.model_limits.default_output_limit,
446 ):
447 return True
449 # Predictive check: will overflow soon
450 return self.predict_overflow(conversation_id, usage)
452 def compact(
453 self,
454 conversation_id: str,
455 llm_client: LLMClient,
456 parent_message_id: str,
457 model: str | None = None,
458 auto_mode: bool = False,
459 ) -> dict[str, Any]:
460 """Compact conversation using LLM summarization (Tier 3).
462 Args:
463 conversation_id: Conversation to compact
464 llm_client: LLM client for summarization
465 parent_message_id: Message that triggered compaction
466 model: Optional model to use for summarization
467 auto_mode: Whether this was auto-triggered
469 Returns:
470 Compaction result with summary message and metrics
471 """
472 if not self.config.compaction.auto and auto_mode:
473 return {"summarized": False}
475 messages = self.get_messages(conversation_id)
476 summary_id = generate_id("msg")
478 result = summarize_conversation(
479 messages=messages,
480 llm_client=llm_client,
481 parent_message_id=parent_message_id,
482 message_id=summary_id,
483 model=model,
484 auto_mode=auto_mode,
485 config=self.config.summarization,
486 )
488 self.add_message(conversation_id, result.summary_message)
490 return {
491 "summarized": True,
492 "summary_message_id": summary_id,
493 "tokens_used": result.tokens_used.total,
494 "cost": result.cost,
495 }
497 def to_model_format(self, conversation_id: str) -> list[dict[str, Any]]:
498 """Convert conversation messages to model format for LLM requests.
500 Args:
501 conversation_id: Conversation to convert
503 Returns:
504 List of messages in model format
505 """
506 messages = self.get_messages(conversation_id)
507 return to_model_messages(messages)
509 def calculate_context_usage(
510 self,
511 conversation_id: str,
512 model: str | None = None,
513 ) -> int:
514 """Calculate exact token count for conversation using tiktoken.
516 This counts ALL tokens in the conversation (user messages, assistant
517 responses, tool outputs, etc.) that will be sent to the model.
519 Uses cl100k_base tokenizer (GPT-4/Claude) which works for most modern LLMs.
521 Args:
522 conversation_id: Conversation to calculate usage for
523 model: Optional model name (currently unused, defaults to cl100k_base)
525 Returns:
526 Exact token count that will be used in context window
527 """
528 from harnessutils.tokens.exact import count_tokens_exact
530 messages = self.to_model_format(conversation_id)
531 return count_tokens_exact(messages, model)
533 def get_tool_output_tokens(self, conversation_id: str) -> dict[str, Any]:
534 """Get detailed breakdown of token usage for tool outputs.
536 Args:
537 conversation_id: Conversation to analyze
539 Returns:
540 Dictionary with token breakdown:
541 - total: Total tokens in tool outputs
542 - by_tool: Token count per tool type
543 - prunable: Tokens that could be pruned
544 - protected: Tokens in protected outputs
545 """
546 from harnessutils.compaction.pruning import calculate_context_tokens
548 messages = self.get_messages(conversation_id)
550 total = calculate_context_tokens(messages)
551 by_tool: dict[str, int] = {}
552 prunable = 0
553 protected = 0
554 turns_skipped = 0
556 for msg in reversed(messages):
557 if msg.role == "user":
558 turns_skipped += 1
560 for part in msg.parts:
561 from harnessutils.models.parts import ToolPart
562 from harnessutils.tokens.exact import count_tokens_fast
564 if not isinstance(part, ToolPart):
565 continue
567 if part.state.status != "completed":
568 continue
570 if not part.state.output:
571 continue
573 tokens = count_tokens_fast(part.state.output)
575 # Track by tool type
576 by_tool[part.tool] = by_tool.get(part.tool, 0) + tokens
578 # Determine if prunable
579 is_protected = (
580 turns_skipped < self.config.pruning.protect_turns
581 or part.tool in self.config.pruning.protected_tools
582 or (part.state.time and part.state.time.compacted)
583 )
585 if is_protected:
586 protected += tokens
587 else:
588 prunable += tokens
590 return {
591 "total": total,
592 "by_tool": by_tool,
593 "prunable": prunable,
594 "protected": protected,
595 "prunability_percent": round((prunable / total * 100) if total > 0 else 0, 1),
596 }
598 def get_context_quality(self, conversation_id: str) -> dict[str, Any]:
599 """Get current quality metrics for conversation.
601 Calculates all quality metrics, updates history, and returns assessment
602 with health status and actionable recommendations.
604 Args:
605 conversation_id: Conversation to analyze
607 Returns:
608 Dictionary with all metrics, health status, and recommendations.
609 Keys: information_density, redundancy_ratio, staleness_score,
610 error_preservation_rate, protected_ratio, health, recommendations
611 """
612 from harnessutils.quality import assess_quality
614 conv_data = self.storage.load_conversation(conversation_id)
615 conv = Conversation.from_dict(conv_data)
616 messages = self.get_messages(conversation_id)
618 # Get pruning decisions from metadata if available
619 decisions = conv.metadata.get("latest_pruning_decisions")
621 snapshot = assess_quality(
622 messages=messages,
623 config=self.config.pruning,
624 decisions=decisions,
625 )
627 # Update history
628 conv.update_quality_history(snapshot)
629 self.storage.save_conversation(conversation_id, conv.to_dict())
631 return snapshot.to_dict()
633 def track_quality_metric(
634 self,
635 conversation_id: str,
636 metric_name: str,
637 value: float,
638 timestamp: int | None = None,
639 ) -> None:
640 """Track a single quality metric value.
642 Note: Prefer get_context_quality() which calculates all metrics.
643 This is for custom/external metrics.
645 Args:
646 conversation_id: Conversation to track
647 metric_name: Name of metric (e.g., "information_density")
648 value: Metric value
649 timestamp: Unix ms timestamp (defaults to now)
650 """
651 from harnessutils.quality import QualitySnapshot
653 if timestamp is None:
654 timestamp = int(time.time() * 1000)
656 conv_data = self.storage.load_conversation(conversation_id)
657 conv = Conversation.from_dict(conv_data)
659 # Create minimal snapshot with just this metric
660 # (Other metrics set to 0.0, empty recommendations)
661 snapshot = QualitySnapshot(
662 timestamp=timestamp,
663 information_density=value if metric_name == "information_density" else 0.0,
664 redundancy_ratio=value if metric_name == "redundancy_ratio" else 0.0,
665 staleness_score=value if metric_name == "staleness_score" else 0.0,
666 error_preservation_rate=value
667 if metric_name == "error_preservation_rate"
668 else 0.0,
669 protected_ratio=value if metric_name == "protected_ratio" else 0.0,
670 health="unknown",
671 recommendations=[],
672 )
674 conv.update_quality_history(snapshot)
675 self.storage.save_conversation(conversation_id, conv.to_dict())
677 def get_quality_trend(
678 self,
679 conversation_id: str,
680 metric: str,
681 window: int = 20,
682 ) -> list[tuple[int, float]]:
683 """Get trend data for a specific metric.
685 Args:
686 conversation_id: Conversation to query
687 metric: Metric name (e.g., "information_density")
688 window: Number of most recent snapshots to return
690 Returns:
691 List of (timestamp, value) tuples, most recent first
692 """
693 conv_data = self.storage.load_conversation(conversation_id)
694 conv = Conversation.from_dict(conv_data)
696 history = conv.get_quality_history()
697 if history is None:
698 return []
700 return history.get_trend(metric, window)
702 def truncate_tool_output(
703 self,
704 output: str,
705 tool_name: str,
706 ) -> str:
707 """Truncate tool output if it exceeds limits (Tier 1).
709 Args:
710 output: Tool output to truncate
711 tool_name: Name of the tool
713 Returns:
714 Potentially truncated output
715 """
716 output_id = generate_id(f"output_{tool_name}")
718 result = truncate_output(
719 output=output,
720 config=self.config.truncation,
721 output_id=output_id,
722 )
724 if result.truncated and result.output_path:
725 self.storage.save_truncated_output(result.output_path, output)
727 return result.content
729 def create_snapshot(
730 self,
731 conversation_id: str,
732 snapshot_id: str | None = None,
733 metadata: dict[str, Any] | None = None,
734 ) -> Snapshot:
735 """Create snapshot of conversation state for reproducibility.
737 Captures full conversation state including:
738 - All messages with content
739 - Conversation metadata (velocity, etc.)
740 - Current configuration
742 Useful for:
743 - Debugging (save state before/after changes)
744 - A/B testing (compare different strategies)
745 - Reproducibility (restore exact state)
747 Args:
748 conversation_id: Conversation to snapshot
749 snapshot_id: Optional snapshot ID (auto-generated if None)
750 metadata: Optional metadata (e.g., {"reason": "before_refactor"})
752 Returns:
753 Created snapshot
755 Example:
756 >>> snap = manager.create_snapshot(conv_id, metadata={"test": "baseline"})
757 >>> # Make changes...
758 >>> snap2 = manager.create_snapshot(conv_id, metadata={"test": "optimized"})
759 >>> diff = manager.compare_snapshots(snap.snapshot_id, snap2.snapshot_id)
760 """
761 messages = self.get_messages(conversation_id)
763 conv_data = self.storage.load_conversation(conversation_id)
764 conversation = Conversation.from_dict(conv_data)
766 # Serialize config to dict recursively (convert dataclasses, Paths, etc.)
767 from dataclasses import asdict, is_dataclass
768 from pathlib import Path
770 def _serialize_config(obj: Any) -> Any:
771 """Recursively convert config to JSON-safe dict."""
772 if isinstance(obj, Path):
773 return str(obj)
774 elif is_dataclass(obj) and not isinstance(obj, type):
775 # Convert dataclass instance to dict, then recursively serialize values
776 obj_dict = asdict(obj)
777 return {k: _serialize_config(v) for k, v in obj_dict.items()}
778 elif isinstance(obj, dict):
779 return {k: _serialize_config(v) for k, v in obj.items()}
780 elif isinstance(obj, (list, tuple)):
781 return [_serialize_config(item) for item in obj]
782 else:
783 return obj
785 config_dict = _serialize_config(self.config)
787 return self.snapshot_manager.create_snapshot(
788 conversation_id=conversation_id,
789 messages=messages,
790 conversation=conversation,
791 config=config_dict,
792 snapshot_id=snapshot_id,
793 metadata=metadata,
794 )
796 def restore_snapshot(self, snapshot_id: str) -> str:
797 """Restore conversation from snapshot.
799 WARNING: This replaces current conversation state with snapshot state.
800 Consider creating a snapshot of current state first.
802 Args:
803 snapshot_id: Snapshot to restore
805 Returns:
806 Conversation ID of restored conversation
808 Raises:
809 FileNotFoundError: If snapshot not found
810 """
811 snapshot = self.snapshot_manager.get_snapshot(snapshot_id)
812 if not snapshot:
813 raise FileNotFoundError(f"Snapshot {snapshot_id} not found")
815 messages, conversation = self.snapshot_manager.restore_snapshot(snapshot)
817 # Clear existing messages (replace entire conversation state)
818 # For memory storage, clear the messages dict for this conversation
819 if hasattr(self.storage, "messages"):
820 if snapshot.conversation_id in self.storage.messages:
821 self.storage.messages[snapshot.conversation_id] = {}
823 # Save conversation
824 self.storage.save_conversation(snapshot.conversation_id, conversation.to_dict())
826 # Save messages
827 for msg in messages:
828 self.storage.save_message(snapshot.conversation_id, msg.id, msg.to_dict())
830 # Clear cache
831 if snapshot.conversation_id in self._message_cache:
832 del self._message_cache[snapshot.conversation_id]
834 return snapshot.conversation_id
836 def compare_snapshots(
837 self, snapshot1_id: str, snapshot2_id: str
838 ) -> SnapshotDiff | None:
839 """Compare two snapshots to see what changed.
841 Args:
842 snapshot1_id: First snapshot ID (earlier)
843 snapshot2_id: Second snapshot ID (later)
845 Returns:
846 SnapshotDiff describing changes:
847 - messages_added: Number of messages added
848 - messages_removed: Number of messages removed
849 - tokens_delta: Change in token count
850 - message_changes: List of specific changes
851 - config_changes: Configuration modifications
852 - metadata_changes: Additional metadata
854 None if snapshots not found
855 """
856 return self.snapshot_manager.compare_snapshots(snapshot1_id, snapshot2_id)
858 def export_snapshot(self, snapshot_id: str, file_path: str) -> None:
859 """Export snapshot to JSON file for version control.
861 Args:
862 snapshot_id: Snapshot to export
863 file_path: Path to write JSON file
865 Raises:
866 FileNotFoundError: If snapshot not found
868 Example:
869 >>> snap = manager.create_snapshot(conv_id)
870 >>> manager.export_snapshot(snap.snapshot_id, "snapshots/baseline.json")
871 >>> # Commit to git for reproducibility
872 """
873 self.snapshot_manager.export_snapshot(snapshot_id, file_path)
875 def import_snapshot(self, file_path: str) -> Snapshot:
876 """Import snapshot from JSON file.
878 Args:
879 file_path: Path to JSON file
881 Returns:
882 Imported snapshot
883 """
884 return self.snapshot_manager.import_snapshot(file_path)
886 def cleanup_stale_data(
887 self,
888 conversation_id: str,
889 max_age_hours: float = 24,
890 keep_errors: bool = True,
891 execute: bool = False,
892 ) -> dict[str, Any]:
893 """Clean up stale conversation data.
895 By default, identifies old messages and tool outputs for cleanup based on age
896 without making any changes. Set execute=True to actually clear stale outputs.
898 Args:
899 conversation_id: Conversation to clean up
900 max_age_hours: Maximum age in hours before considering stale (default: 24)
901 keep_errors: Whether to preserve error messages (default: True)
902 execute: When True, actually clears stale outputs (default: False)
904 Returns:
905 Dictionary with cleanup statistics:
906 - messages_archived (int): Messages eligible for archival
907 - tokens_freed (int): Tokens that would be freed
908 - duplicates_removed (int): Duplicate outputs found
909 - stale_outputs_pruned (int): Stale outputs identified/cleared
910 - operations (list[str]): Description of operations
912 Example:
913 >>> result = manager.cleanup_stale_data(conv_id, max_age_hours=48)
914 >>> print(f"Can free {result['tokens_freed']} tokens")
915 >>> result = manager.cleanup_stale_data(conv_id, max_age_hours=48, execute=True)
916 >>> print(f"Freed {result['tokens_freed']} tokens")
917 """
918 from harnessutils.maintenance import cleanup_stale_data
920 messages = self.get_messages(conversation_id)
921 result = cleanup_stale_data(
922 messages=messages,
923 config=self.config.pruning,
924 max_age_hours=max_age_hours,
925 keep_errors=keep_errors,
926 execute=execute,
927 )
929 if execute:
930 for msg in messages:
931 self.storage.save_message(conversation_id, msg.id, msg.to_dict())
932 # Invalidate cache — outputs were cleared in-place
933 if conversation_id in self._message_cache:
934 del self._message_cache[conversation_id]
936 return result.to_dict()
938 def scan_and_deduplicate(self, conversation_id: str) -> dict[str, Any]:
939 """Scan for duplicate outputs and return statistics.
941 Identifies duplicate tool outputs without removing them.
942 Use prune_before_turn() to actually remove duplicates.
944 Args:
945 conversation_id: Conversation to scan
947 Returns:
948 Dictionary with deduplication statistics:
949 - duplicates_removed (int): Number of duplicates found
950 - tokens_freed (int): Tokens that would be freed
951 - operations (list[str]): Description of findings
953 Example:
954 >>> result = manager.scan_and_deduplicate(conv_id)
955 >>> if result['duplicates_removed'] > 0:
956 ... print(f"Found {result['duplicates_removed']} duplicates")
957 ... # Run compaction to actually remove them
958 ... manager.prune_before_turn(conv_id)
959 """
960 from harnessutils.maintenance import scan_and_deduplicate
962 messages = self.get_messages(conversation_id)
963 result = scan_and_deduplicate(
964 messages=messages,
965 config=self.config.pruning,
966 )
968 return result.to_dict()
970 def get_memory(self, project_id: str, key: str, default: Any = None) -> Any:
971 """Get a value from project-scoped memory.
973 Args:
974 project_id: Project to retrieve memory from
975 key: Memory key
976 default: Default value if key not found
978 Returns:
979 Stored value or default
980 """
981 data = self._load_project_memory(project_id)
982 return data.get(key, default)
984 def set_memory(self, project_id: str, key: str, value: Any) -> None:
985 """Set a value in project-scoped memory.
987 Args:
988 project_id: Project to store memory in
989 key: Memory key
990 value: Value to store
991 """
992 data = self._load_project_memory(project_id)
993 data[key] = value
994 try:
995 self.storage.save_project_memory(project_id, data)
996 except AttributeError:
997 pass # Backend doesn't support project memory
999 def delete_memory(self, project_id: str, key: str) -> None:
1000 """Delete a key from project-scoped memory.
1002 Args:
1003 project_id: Project to delete memory from
1004 key: Memory key to delete
1005 """
1006 data = self._load_project_memory(project_id)
1007 data.pop(key, None)
1008 try:
1009 self.storage.save_project_memory(project_id, data)
1010 except AttributeError:
1011 pass # Backend doesn't support project memory
1013 def list_memory(self, project_id: str) -> dict[str, Any]:
1014 """List all memory keys and values for a project.
1016 Args:
1017 project_id: Project to list memory for
1019 Returns:
1020 Dictionary of all memory key-value pairs
1021 """
1022 return self._load_project_memory(project_id)
1024 def _load_project_memory(self, project_id: str) -> dict[str, Any]:
1025 """Load project memory, returning empty dict on failure.
1027 Args:
1028 project_id: Project to load memory for
1030 Returns:
1031 Memory dictionary (empty if not found or backend unsupported)
1032 """
1033 try:
1034 return self.storage.load_project_memory(project_id)
1035 except (FileNotFoundError, AttributeError):
1036 return {}
1038 def detect_context_issues(self, conversation_id: str) -> list[dict[str, Any]]:
1039 """Detect quality and drift issues in conversation context.
1041 Analyzes conversation for common problems:
1042 - High redundancy (duplicate content)
1043 - Staleness accumulation (old messages)
1044 - Low information density
1045 - Error preservation status
1046 - Excessive protection (limiting pruning)
1048 Args:
1049 conversation_id: Conversation to analyze
1051 Returns:
1052 List of issue dictionaries, each containing:
1053 - issue_type (str): Type of issue
1054 - severity (str): "info", "warning", or "error"
1055 - description (str): Human-readable description
1056 - affected_count (int): Number of items affected
1057 - suggested_fix (str): Recommended action
1058 - metadata (dict): Additional details
1060 Example:
1061 >>> issues = manager.detect_context_issues(conv_id)
1062 >>> for issue in issues:
1063 ... if issue['severity'] == 'warning':
1064 ... print(f"⚠️ {issue['description']}")
1065 ... print(f" Fix: {issue['suggested_fix']}")
1066 """
1067 from harnessutils.maintenance import detect_context_issues
1069 conv_data = self.storage.load_conversation(conversation_id)
1070 conv = Conversation.from_dict(conv_data)
1071 messages = self.get_messages(conversation_id)
1073 issues = detect_context_issues(
1074 messages=messages,
1075 conversation=conv,
1076 config=self.config.pruning,
1077 )
1079 return [issue.to_dict() for issue in issues]