Coverage for src / harnessutils / manager.py: 96%
305 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-04-11 08:11 -0600
« prev ^ index » next coverage.py v7.13.2, created at 2026-04-11 08:11 -0600
1"""Main ConversationManager API for harness-utils."""
3import time
4from collections.abc import Callable
5from pathlib import Path
6from typing import Any
8from harnessutils.compaction.pruning import PruningDecision, prune_tool_outputs
9from harnessutils.compaction.summarization import is_overflow, summarize_conversation
10from harnessutils.compaction.truncation import truncate_output
11from harnessutils.config import HarnessConfig, StorageConfig
12from harnessutils.conversion.to_model import to_model_messages
13from harnessutils.hooks import MessageHooks
14from harnessutils.inspection import ContextInspector
15from harnessutils.memory import SemanticMemoryBackend
16from harnessutils.models.conversation import Conversation
17from harnessutils.models.message import Message
18from harnessutils.models.usage import Usage
19from harnessutils.snapshots import Snapshot, SnapshotDiff, SnapshotManager
20from harnessutils.storage.filesystem import FilesystemStorage
21from harnessutils.storage.memory import MemoryStorage
22from harnessutils.types import LLMClient, StorageBackend
23from harnessutils.utils.ids import generate_id
24from harnessutils.workspace import resolve_workspace
27class ConversationManager:
28 """Main interface for managing conversations with context window management.
30 Provides high-level API for:
31 - Creating and managing conversations
32 - Adding messages
33 - Automatic context compaction (truncation, pruning, summarization)
34 - Message storage and retrieval
35 """
37 def __init__(
38 self,
39 storage: StorageBackend | None = None,
40 config: HarnessConfig | None = None,
41 *,
42 workspace_root: Path | str | None = None,
43 harness_name: str = ".harness",
44 semantic_memory: SemanticMemoryBackend | None = None,
45 message_hooks: MessageHooks | None = None,
46 ):
47 """Initialize conversation manager.
49 Args:
50 storage: Storage backend (uses in-memory if None)
51 config: Configuration (uses defaults if None)
52 workspace_root: Optional workspace root. Creates harness dir with stable
53 project UUID. Defaults storage to FilesystemStorage under <harness>/sessions/.
54 harness_name: Name of the harness directory (default: ".harness")
55 semantic_memory: Optional semantic memory backend. Stored as attribute;
56 harness-utils does not call it internally.
57 message_hooks: Optional hooks for the add_message() lifecycle.
58 on_before_add_message fires before storage and can modify or reject.
59 on_after_add_message fires after storage for side effects.
60 """
61 self.config = config or HarnessConfig()
62 self.project_id: str | None = None
63 self.semantic_memory = semantic_memory
65 if workspace_root is not None:
66 harness_dir, self.project_id = resolve_workspace(Path(workspace_root), harness_name)
67 if storage is None:
68 sessions_path = harness_dir / "sessions"
69 storage_config = StorageConfig(base_path=sessions_path)
70 storage = FilesystemStorage(storage_config)
72 self.storage = storage or MemoryStorage()
73 self.message_hooks = message_hooks
74 self._message_cache: dict[str, list[Message]] = {}
75 # Pass storage (not self.storage) so SnapshotManager receives None when no
76 # explicit backend was given — it falls back to its internal in-memory dict.
77 self.snapshot_manager = SnapshotManager(storage)
79 def create_conversation(
80 self,
81 conversation_id: str | None = None,
82 project_id: str | None = None,
83 ) -> Conversation:
84 """Create a new conversation.
86 Args:
87 conversation_id: Optional conversation ID (generated if None)
88 project_id: Optional project ID for grouping
90 Returns:
91 New conversation object
92 """
93 if conversation_id is None:
94 conversation_id = generate_id("conv")
96 now = int(time.time() * 1000)
97 conversation = Conversation(
98 id=conversation_id,
99 project_id=project_id,
100 created=now,
101 updated=now,
102 )
104 self.storage.save_conversation(conversation_id, conversation.to_dict())
105 self._message_cache[conversation_id] = []
107 return conversation
109 def add_message(self, conversation_id: str, message: Message) -> None:
110 """Add a message to a conversation.
112 If message_hooks.on_before_add_message is set, it is called first.
113 Its return value replaces the message that will be stored. Raising
114 from it rejects the message entirely (nothing is stored).
116 If message_hooks.on_after_add_message is set, it is called after
117 successful storage. Its return value is ignored.
119 Args:
120 conversation_id: Conversation to add message to
121 message: Message to add
122 """
123 # Pre-hook — may modify message or raise to reject
124 if self.message_hooks and self.message_hooks.on_before_add_message:
125 message = self.message_hooks.on_before_add_message(conversation_id, message)
127 # Inject timestamp if not set — required for cleanup_stale_data to work
128 if "timestamp" not in message.metadata:
129 message.metadata["timestamp"] = int(time.time() * 1000)
131 self.storage.save_message(conversation_id, message.id, message.to_dict())
133 # Persist parts separately so they survive process restart
134 for idx, part in enumerate(message.parts):
135 self.storage.save_part(message.id, f"{idx:04d}", part.to_dict())
137 if conversation_id not in self._message_cache:
138 self._message_cache[conversation_id] = []
139 self._message_cache[conversation_id].append(message)
141 # Load conversation and update
142 conv_data = self.storage.load_conversation(conversation_id)
143 conv = Conversation.from_dict(conv_data)
144 conv.updated = int(time.time() * 1000)
146 # Track velocity if message has token count
147 if message.tokens:
148 tokens_added = message.tokens.total
149 conv.update_velocity(tokens_added)
151 self.storage.save_conversation(conversation_id, conv.to_dict())
153 # Post-hook — side effects only, message already persisted
154 if self.message_hooks and self.message_hooks.on_after_add_message:
155 self.message_hooks.on_after_add_message(conversation_id, message)
157 def get_messages(self, conversation_id: str) -> list[Message]:
158 """Get all messages for a conversation.
160 Args:
161 conversation_id: Conversation ID
163 Returns:
164 List of messages in chronological order
165 """
166 if conversation_id in self._message_cache:
167 return self._message_cache[conversation_id]
169 message_ids = self.storage.list_messages(conversation_id)
170 messages = []
171 for msg_id in message_ids:
172 msg = Message.from_dict(self.storage.load_message(conversation_id, msg_id))
173 # Reload persisted parts (saved by add_message)
174 part_ids = self.storage.list_parts(msg_id)
175 if part_ids:
176 from harnessutils.models.parts import part_from_dict
177 for part_id in part_ids:
178 try:
179 part_data = self.storage.load_part(msg_id, part_id)
180 msg.parts.append(part_from_dict(part_data))
181 except Exception:
182 pass
183 messages.append(msg)
185 self._message_cache[conversation_id] = messages
186 return messages
188 def query_messages(
189 self,
190 conversation_id: str,
191 limit: int | None = None,
192 offset: int = 0,
193 order: str = "asc",
194 after: int | None = None,
195 before: int | None = None,
196 filter: dict[str, Any] | None = None,
197 ) -> list[Message]:
198 """Query messages with filtering and pagination.
200 Args:
201 conversation_id: Conversation to query
202 limit: Maximum messages to return
203 offset: Skip first N messages
204 order: "asc" (oldest first) or "desc" (newest first)
205 after: Unix ms timestamp - messages after this time
206 before: Unix ms timestamp - messages before this time
207 filter: Filter criteria dict:
208 - has_errors (bool): Only messages with errors
209 - has_warnings (bool): Only messages with warnings
210 - min_importance (float): Minimum importance score
211 - tools (list[str]): Specific tool types
212 - roles (list[str]): Message roles (user/assistant)
213 - has_tool_outputs (bool): Only messages with tool outputs
214 - is_summary (bool): Only summary messages
216 Returns:
217 Filtered and paginated messages
219 Examples:
220 # Get last 10 messages
221 messages = manager.query_messages(conv_id, limit=10, order="desc")
223 # Get messages with errors
224 messages = manager.query_messages(
225 conv_id,
226 filter={"has_errors": True}
227 )
229 # Get recent high-importance messages
230 messages = manager.query_messages(
231 conv_id,
232 filter={"min_importance": 50.0},
233 after=timestamp_24h_ago,
234 limit=20
235 )
236 """
237 from harnessutils.query import MessageFilter, QueryOptions, query_messages
239 # Build filter
240 msg_filter = None
241 if filter is not None:
242 msg_filter = MessageFilter(
243 has_errors=filter.get("has_errors"),
244 has_warnings=filter.get("has_warnings"),
245 min_importance=filter.get("min_importance"),
246 tools=filter.get("tools"),
247 roles=filter.get("roles"),
248 has_tool_outputs=filter.get("has_tool_outputs"),
249 is_summary=filter.get("is_summary"),
250 )
252 # Build query options
253 options = QueryOptions(
254 limit=limit,
255 offset=offset,
256 order=order, # type: ignore
257 after=after,
258 before=before,
259 filter=msg_filter,
260 )
262 # Get all messages and query
263 messages = self.get_messages(conversation_id)
264 return query_messages(messages, options)
266 def get_context_summary(
267 self, conversation_id: str, recent_limit: int = 5
268 ) -> dict[str, Any]:
269 """Get lightweight context summary without loading full content.
271 Provides overview of conversation state including:
272 - Message count and token usage
273 - Summary messages
274 - Recent activity (last N messages)
275 - Key messages (high importance)
276 - Error messages
278 Args:
279 conversation_id: Conversation to summarize
280 recent_limit: Number of recent messages to include (default: 5)
282 Returns:
283 Dictionary with context summary:
284 - conversation_id (str)
285 - message_count (int)
286 - total_tokens (int)
287 - summaries (list): Summary messages
288 - recent_activity (list): Last N message summaries
289 - key_messages (list): High-importance messages
290 - errors (list): Error messages
292 Examples:
293 # Get quick overview
294 summary = manager.get_context_summary(conv_id)
295 print(f"Messages: {summary['message_count']}")
296 print(f"Tokens: {summary['total_tokens']}")
297 print(f"Errors: {len(summary['errors'])}")
298 """
299 from harnessutils.query import build_context_summary
301 messages = self.get_messages(conversation_id)
302 summary = build_context_summary(conversation_id, messages, recent_limit)
303 return summary.to_dict()
305 def inspect_context(self, conversation_id: str) -> ContextInspector:
306 """Create inspector for querying context state.
308 Provides observability into:
309 - What's currently in context
310 - Token counts and breakdowns
311 - Pruning predictions
312 - Audit trail of decisions
314 Args:
315 conversation_id: Conversation to inspect
317 Returns:
318 ContextInspector instance with full query capabilities
320 Example:
321 >>> inspector = manager.inspect_context(conv_id)
322 >>> summary = inspector.summary()
323 >>> print(f"Total tokens: {summary['total_tokens']}")
324 >>> impact = inspector.predict_impact(5000)
325 >>> if impact['would_trigger_pruning']:
326 ... print(f"Would prune ~{impact['estimated_pruned_count']} outputs")
327 """
328 messages = self.get_messages(conversation_id)
330 # Load conversation metadata
331 conv_data = self.storage.load_conversation(conversation_id)
332 conversation = Conversation.from_dict(conv_data)
334 # Load persisted pruning decisions for audit trail
335 decisions_data = conversation.metadata.get("latest_pruning_decisions", [])
336 decisions = [PruningDecision.from_dict(d) for d in decisions_data]
338 return ContextInspector(messages, self.config, conversation, decisions)
340 def prune_before_turn(
341 self,
342 conversation_id: str,
343 auto_mode: bool = False,
344 query: str | None = None,
345 summarizer: Callable[[str, str], str] | None = None,
346 ) -> dict[str, Any]:
347 """Proactively prune old tool outputs before processing a turn.
349 This is Tier 2 compaction - removes old tool outputs while
350 preserving conversation structure.
352 Args:
353 conversation_id: Conversation to prune
354 auto_mode: Whether this was auto-triggered
355 query: Current user query for relevance-based importance boosting
356 summarizer: Optional callable(tool_name, output) -> summary_text.
357 When provided, called before clearing each output so a
358 one-sentence summary is preserved in the model context.
360 Returns:
361 Detailed pruning result with token tracking and breakdown:
362 - pruned: Total outputs removed
363 - tokens_saved: Total tokens saved
364 - tokens_before: Token count before pruning
365 - tokens_after: Token count after pruning
366 - duplicates_pruned: Outputs removed due to duplication
367 - importance_pruned: Outputs removed due to low importance
368 - duplicate_tokens_saved: Tokens saved from deduplication
369 - importance_tokens_saved: Tokens saved from importance pruning
370 - reduction_percent: Percentage reduction in token usage
371 """
372 if not self.config.compaction.prune and auto_mode:
373 return {
374 "pruned": 0,
375 "tokens_saved": 0,
376 "tokens_before": 0,
377 "tokens_after": 0,
378 "duplicates_pruned": 0,
379 "importance_pruned": 0,
380 "duplicate_tokens_saved": 0,
381 "importance_tokens_saved": 0,
382 "reduction_percent": 0,
383 }
385 messages = self.get_messages(conversation_id)
387 conv_data = self.storage.load_conversation(conversation_id)
388 conv = Conversation.from_dict(conv_data)
389 already_excluded = set(conv.excluded_outputs)
391 result = prune_tool_outputs(
392 messages,
393 self.config.pruning,
394 query=query,
395 summarizer=summarizer,
396 already_excluded=already_excluded,
397 )
399 # Accumulate new exclusions into conversation metadata
400 if result.excluded_call_ids:
401 conv.add_excluded_outputs(result.excluded_call_ids)
403 conv.metadata["latest_pruning_decisions"] = [d.to_dict() for d in result.decisions]
404 self.storage.save_conversation(conversation_id, conv.to_dict())
406 # If a summarizer wrote summary_text onto parts, persist those messages
407 if summarizer and result.excluded_call_ids:
408 excluded_set = set(result.excluded_call_ids)
409 for msg in messages:
410 for part in msg.parts:
411 if hasattr(part, "call_id") and part.call_id in excluded_set:
412 self.storage.save_message(conversation_id, msg.id, msg.to_dict())
413 break
415 return result.to_dict()
417 def predict_overflow(
418 self,
419 conversation_id: str,
420 current_usage: Usage,
421 ) -> bool:
422 """Predict if conversation will overflow in next N turns.
424 Args:
425 conversation_id: Conversation to check
426 current_usage: Current token usage
428 Returns:
429 True if overflow predicted within lookahead window
430 """
431 if not self.config.compaction.use_predictive:
432 return False
434 # Load conversation and get velocity
435 conv_data = self.storage.load_conversation(conversation_id)
436 conv = Conversation.from_dict(conv_data)
437 velocity = conv.get_velocity()
439 if velocity is None or not velocity.turn_deltas:
440 return False # No velocity data yet
442 # Project tokens ahead
443 lookahead = self.config.compaction.predictive_lookahead
444 predicted_growth = velocity.predict_tokens_ahead(lookahead)
446 # Calculate current total and projected total
447 current_total = current_usage.input + current_usage.cache.read
448 projected_total = current_total + predicted_growth
450 # Check against safety margin
451 usable_space = (
452 self.config.model_limits.default_context_limit
453 - self.config.model_limits.default_output_limit
454 )
455 safety_threshold = usable_space * self.config.compaction.predictive_safety_margin
457 return projected_total > safety_threshold
459 def needs_compaction(
460 self,
461 conversation_id: str,
462 usage: Usage,
463 ) -> bool:
464 """Check if conversation needs summarization (Tier 3).
466 Uses both reactive (overflow) and predictive checks.
468 Args:
469 conversation_id: Conversation to check
470 usage: Token usage from last turn
472 Returns:
473 True if summarization needed
474 """
475 # Reactive check: already overflowed
476 if is_overflow(
477 usage,
478 self.config.model_limits.default_context_limit,
479 self.config.model_limits.default_output_limit,
480 ):
481 return True
483 # Predictive check: will overflow soon
484 return self.predict_overflow(conversation_id, usage)
486 def compact(
487 self,
488 conversation_id: str,
489 llm_client: LLMClient,
490 parent_message_id: str,
491 model: str | None = None,
492 auto_mode: bool = False,
493 ) -> dict[str, Any]:
494 """Compact conversation using LLM summarization (Tier 3).
496 Args:
497 conversation_id: Conversation to compact
498 llm_client: LLM client for summarization
499 parent_message_id: Message that triggered compaction
500 model: Optional model to use for summarization
501 auto_mode: Whether this was auto-triggered
503 Returns:
504 Compaction result with summary message and metrics
505 """
506 if not self.config.compaction.auto and auto_mode:
507 return {"summarized": False}
509 messages = self.get_messages(conversation_id)
510 summary_id = generate_id("msg")
512 result = summarize_conversation(
513 messages=messages,
514 llm_client=llm_client,
515 parent_message_id=parent_message_id,
516 message_id=summary_id,
517 model=model,
518 auto_mode=auto_mode,
519 config=self.config.summarization,
520 )
522 self.add_message(conversation_id, result.summary_message)
524 return {
525 "summarized": True,
526 "summary_message_id": summary_id,
527 "tokens_used": result.tokens_used.total,
528 "cost": result.cost,
529 }
531 def to_model_format(self, conversation_id: str) -> list[dict[str, Any]]:
532 """Convert conversation messages to model format for LLM requests.
534 Args:
535 conversation_id: Conversation to convert
537 Returns:
538 List of messages in model format
539 """
540 messages = self.get_messages(conversation_id)
541 conv_data = self.storage.load_conversation(conversation_id)
542 conv = Conversation.from_dict(conv_data)
543 return to_model_messages(messages, excluded_outputs=set(conv.excluded_outputs))
545 def calculate_context_usage(
546 self,
547 conversation_id: str,
548 model: str | None = None,
549 ) -> int:
550 """Calculate exact token count for conversation using tiktoken.
552 This counts ALL tokens in the conversation (user messages, assistant
553 responses, tool outputs, etc.) that will be sent to the model.
555 Uses cl100k_base tokenizer (GPT-4/Claude) which works for most modern LLMs.
557 Args:
558 conversation_id: Conversation to calculate usage for
559 model: Optional model name (currently unused, defaults to cl100k_base)
561 Returns:
562 Exact token count that will be used in context window
563 """
564 from harnessutils.tokens.exact import count_tokens_exact
566 messages = self.to_model_format(conversation_id)
567 return count_tokens_exact(messages, model)
569 def get_tool_output_tokens(self, conversation_id: str) -> dict[str, Any]:
570 """Get detailed breakdown of token usage for tool outputs.
572 Args:
573 conversation_id: Conversation to analyze
575 Returns:
576 Dictionary with token breakdown:
577 - total: Total tokens in tool outputs
578 - by_tool: Token count per tool type
579 - prunable: Tokens that could be pruned
580 - protected: Tokens in protected outputs
581 """
582 from harnessutils.compaction.pruning import calculate_context_tokens
584 messages = self.get_messages(conversation_id)
586 total = calculate_context_tokens(messages)
587 by_tool: dict[str, int] = {}
588 prunable = 0
589 protected = 0
590 turns_skipped = 0
592 for msg in reversed(messages):
593 if msg.role == "user":
594 turns_skipped += 1
596 for part in msg.parts:
597 from harnessutils.models.parts import ToolPart
598 from harnessutils.tokens.exact import count_tokens_fast
600 if not isinstance(part, ToolPart):
601 continue
603 if part.state.status != "completed":
604 continue
606 if not part.state.output:
607 continue
609 tokens = count_tokens_fast(part.state.output)
611 # Track by tool type
612 by_tool[part.tool] = by_tool.get(part.tool, 0) + tokens
614 # Determine if prunable
615 is_protected = (
616 turns_skipped < self.config.pruning.protect_turns
617 or part.tool in self.config.pruning.protected_tools
618 or (part.state.time and part.state.time.compacted)
619 )
621 if is_protected:
622 protected += tokens
623 else:
624 prunable += tokens
626 return {
627 "total": total,
628 "by_tool": by_tool,
629 "prunable": prunable,
630 "protected": protected,
631 "prunability_percent": round((prunable / total * 100) if total > 0 else 0, 1),
632 }
634 def get_context_quality(self, conversation_id: str) -> dict[str, Any]:
635 """Get current quality metrics for conversation.
637 Calculates all quality metrics, updates history, and returns assessment
638 with health status and actionable recommendations.
640 Args:
641 conversation_id: Conversation to analyze
643 Returns:
644 Dictionary with all metrics, health status, and recommendations.
645 Keys: information_density, redundancy_ratio, staleness_score,
646 error_preservation_rate, protected_ratio, health, recommendations
647 """
648 from harnessutils.quality import assess_quality
650 conv_data = self.storage.load_conversation(conversation_id)
651 conv = Conversation.from_dict(conv_data)
652 messages = self.get_messages(conversation_id)
654 # Get pruning decisions from metadata if available
655 decisions = conv.metadata.get("latest_pruning_decisions")
657 snapshot = assess_quality(
658 messages=messages,
659 config=self.config.pruning,
660 decisions=decisions,
661 )
663 # Update history
664 conv.update_quality_history(snapshot)
665 self.storage.save_conversation(conversation_id, conv.to_dict())
667 return snapshot.to_dict()
669 def track_quality_metric(
670 self,
671 conversation_id: str,
672 metric_name: str,
673 value: float,
674 timestamp: int | None = None,
675 ) -> None:
676 """Track a single quality metric value.
678 Note: Prefer get_context_quality() which calculates all metrics.
679 This is for custom/external metrics.
681 Args:
682 conversation_id: Conversation to track
683 metric_name: Name of metric (e.g., "information_density")
684 value: Metric value
685 timestamp: Unix ms timestamp (defaults to now)
686 """
687 from harnessutils.quality import QualitySnapshot
689 if timestamp is None:
690 timestamp = int(time.time() * 1000)
692 conv_data = self.storage.load_conversation(conversation_id)
693 conv = Conversation.from_dict(conv_data)
695 # Create minimal snapshot with just this metric
696 # (Other metrics set to 0.0, empty recommendations)
697 snapshot = QualitySnapshot(
698 timestamp=timestamp,
699 information_density=value if metric_name == "information_density" else 0.0,
700 redundancy_ratio=value if metric_name == "redundancy_ratio" else 0.0,
701 staleness_score=value if metric_name == "staleness_score" else 0.0,
702 error_preservation_rate=value
703 if metric_name == "error_preservation_rate"
704 else 0.0,
705 protected_ratio=value if metric_name == "protected_ratio" else 0.0,
706 health="unknown",
707 recommendations=[],
708 )
710 conv.update_quality_history(snapshot)
711 self.storage.save_conversation(conversation_id, conv.to_dict())
713 def get_quality_trend(
714 self,
715 conversation_id: str,
716 metric: str,
717 window: int = 20,
718 ) -> list[tuple[int, float]]:
719 """Get trend data for a specific metric.
721 Args:
722 conversation_id: Conversation to query
723 metric: Metric name (e.g., "information_density")
724 window: Number of most recent snapshots to return
726 Returns:
727 List of (timestamp, value) tuples, most recent first
728 """
729 conv_data = self.storage.load_conversation(conversation_id)
730 conv = Conversation.from_dict(conv_data)
732 history = conv.get_quality_history()
733 if history is None:
734 return []
736 return history.get_trend(metric, window)
738 def truncate_tool_output(
739 self,
740 output: str,
741 tool_name: str,
742 ) -> str:
743 """Truncate tool output if it exceeds limits (Tier 1).
745 Args:
746 output: Tool output to truncate
747 tool_name: Name of the tool
749 Returns:
750 Potentially truncated output
751 """
752 output_id = generate_id(f"output_{tool_name}")
754 result = truncate_output(
755 output=output,
756 config=self.config.truncation,
757 output_id=output_id,
758 )
760 if result.truncated and result.output_path:
761 self.storage.save_truncated_output(result.output_path, output)
763 return result.content
765 def create_snapshot(
766 self,
767 conversation_id: str,
768 snapshot_id: str | None = None,
769 metadata: dict[str, Any] | None = None,
770 ) -> Snapshot:
771 """Create snapshot of conversation state for reproducibility.
773 Captures full conversation state including:
774 - All messages with content
775 - Conversation metadata (velocity, etc.)
776 - Current configuration
778 Useful for:
779 - Debugging (save state before/after changes)
780 - A/B testing (compare different strategies)
781 - Reproducibility (restore exact state)
783 Args:
784 conversation_id: Conversation to snapshot
785 snapshot_id: Optional snapshot ID (auto-generated if None)
786 metadata: Optional metadata (e.g., {"reason": "before_refactor"})
788 Returns:
789 Created snapshot
791 Example:
792 >>> snap = manager.create_snapshot(conv_id, metadata={"test": "baseline"})
793 >>> # Make changes...
794 >>> snap2 = manager.create_snapshot(conv_id, metadata={"test": "optimized"})
795 >>> diff = manager.compare_snapshots(snap.snapshot_id, snap2.snapshot_id)
796 """
797 messages = self.get_messages(conversation_id)
799 conv_data = self.storage.load_conversation(conversation_id)
800 conversation = Conversation.from_dict(conv_data)
802 # Serialize config to dict recursively (convert dataclasses, Paths, etc.)
803 from dataclasses import asdict, is_dataclass
804 from pathlib import Path
806 def _serialize_config(obj: Any) -> Any:
807 """Recursively convert config to JSON-safe dict."""
808 if isinstance(obj, Path):
809 return str(obj)
810 elif is_dataclass(obj) and not isinstance(obj, type):
811 # Convert dataclass instance to dict, then recursively serialize values
812 obj_dict = asdict(obj)
813 return {k: _serialize_config(v) for k, v in obj_dict.items()}
814 elif isinstance(obj, dict):
815 return {k: _serialize_config(v) for k, v in obj.items()}
816 elif isinstance(obj, (list, tuple)):
817 return [_serialize_config(item) for item in obj]
818 else:
819 return obj
821 config_dict = _serialize_config(self.config)
823 return self.snapshot_manager.create_snapshot(
824 conversation_id=conversation_id,
825 messages=messages,
826 conversation=conversation,
827 config=config_dict,
828 snapshot_id=snapshot_id,
829 metadata=metadata,
830 )
832 def restore_snapshot(self, snapshot_id: str) -> str:
833 """Restore conversation from snapshot.
835 WARNING: This replaces current conversation state with snapshot state.
836 Consider creating a snapshot of current state first.
838 Args:
839 snapshot_id: Snapshot to restore
841 Returns:
842 Conversation ID of restored conversation
844 Raises:
845 FileNotFoundError: If snapshot not found
846 """
847 snapshot = self.snapshot_manager.get_snapshot(snapshot_id)
848 if not snapshot:
849 raise FileNotFoundError(f"Snapshot {snapshot_id} not found")
851 messages, conversation = self.snapshot_manager.restore_snapshot(snapshot)
853 # Clear existing messages (replace entire conversation state)
854 # For memory storage, clear the messages dict for this conversation
855 if hasattr(self.storage, "messages"):
856 if snapshot.conversation_id in self.storage.messages:
857 self.storage.messages[snapshot.conversation_id] = {}
859 # Save conversation
860 self.storage.save_conversation(snapshot.conversation_id, conversation.to_dict())
862 # Save messages
863 for msg in messages:
864 self.storage.save_message(snapshot.conversation_id, msg.id, msg.to_dict())
866 # Clear cache
867 if snapshot.conversation_id in self._message_cache:
868 del self._message_cache[snapshot.conversation_id]
870 return snapshot.conversation_id
872 def compare_snapshots(
873 self, snapshot1_id: str, snapshot2_id: str
874 ) -> SnapshotDiff | None:
875 """Compare two snapshots to see what changed.
877 Args:
878 snapshot1_id: First snapshot ID (earlier)
879 snapshot2_id: Second snapshot ID (later)
881 Returns:
882 SnapshotDiff describing changes:
883 - messages_added: Number of messages added
884 - messages_removed: Number of messages removed
885 - tokens_delta: Change in token count
886 - message_changes: List of specific changes
887 - config_changes: Configuration modifications
888 - metadata_changes: Additional metadata
890 None if snapshots not found
891 """
892 return self.snapshot_manager.compare_snapshots(snapshot1_id, snapshot2_id)
894 def export_snapshot(self, snapshot_id: str, file_path: str) -> None:
895 """Export snapshot to JSON file for version control.
897 Args:
898 snapshot_id: Snapshot to export
899 file_path: Path to write JSON file
901 Raises:
902 FileNotFoundError: If snapshot not found
904 Example:
905 >>> snap = manager.create_snapshot(conv_id)
906 >>> manager.export_snapshot(snap.snapshot_id, "snapshots/baseline.json")
907 >>> # Commit to git for reproducibility
908 """
909 self.snapshot_manager.export_snapshot(snapshot_id, file_path)
911 def import_snapshot(self, file_path: str) -> Snapshot:
912 """Import snapshot from JSON file.
914 Args:
915 file_path: Path to JSON file
917 Returns:
918 Imported snapshot
919 """
920 return self.snapshot_manager.import_snapshot(file_path)
922 def cleanup_stale_data(
923 self,
924 conversation_id: str,
925 max_age_hours: float = 24,
926 keep_errors: bool = True,
927 execute: bool = False,
928 ) -> dict[str, Any]:
929 """Clean up stale conversation data.
931 By default, identifies old messages and tool outputs for cleanup based on age
932 without making any changes. Set execute=True to actually clear stale outputs.
934 Args:
935 conversation_id: Conversation to clean up
936 max_age_hours: Maximum age in hours before considering stale (default: 24)
937 keep_errors: Whether to preserve error messages (default: True)
938 execute: When True, actually clears stale outputs (default: False)
940 Returns:
941 Dictionary with cleanup statistics:
942 - messages_archived (int): Messages eligible for archival
943 - tokens_freed (int): Tokens that would be freed
944 - duplicates_removed (int): Duplicate outputs found
945 - stale_outputs_pruned (int): Stale outputs identified/cleared
946 - operations (list[str]): Description of operations
948 Example:
949 >>> result = manager.cleanup_stale_data(conv_id, max_age_hours=48)
950 >>> print(f"Can free {result['tokens_freed']} tokens")
951 >>> result = manager.cleanup_stale_data(conv_id, max_age_hours=48, execute=True)
952 >>> print(f"Freed {result['tokens_freed']} tokens")
953 """
954 from harnessutils.maintenance import cleanup_stale_data
956 messages = self.get_messages(conversation_id)
957 result = cleanup_stale_data(
958 messages=messages,
959 config=self.config.pruning,
960 max_age_hours=max_age_hours,
961 keep_errors=keep_errors,
962 execute=execute,
963 )
965 if execute:
966 for msg in messages:
967 self.storage.save_message(conversation_id, msg.id, msg.to_dict())
968 # Parts are stored separately — persist modified parts back to storage
969 for idx, part in enumerate(msg.parts):
970 self.storage.save_part(msg.id, f"{idx:04d}", part.to_dict())
971 # Invalidate cache — outputs were cleared in-place
972 if conversation_id in self._message_cache:
973 del self._message_cache[conversation_id]
975 return result.to_dict()
977 def scan_and_deduplicate(self, conversation_id: str) -> dict[str, Any]:
978 """Scan for duplicate outputs and return statistics.
980 Identifies duplicate tool outputs without removing them.
981 Use prune_before_turn() to actually remove duplicates.
983 Args:
984 conversation_id: Conversation to scan
986 Returns:
987 Dictionary with deduplication statistics:
988 - duplicates_removed (int): Number of duplicates found
989 - tokens_freed (int): Tokens that would be freed
990 - operations (list[str]): Description of findings
992 Example:
993 >>> result = manager.scan_and_deduplicate(conv_id)
994 >>> if result['duplicates_removed'] > 0:
995 ... print(f"Found {result['duplicates_removed']} duplicates")
996 ... # Run compaction to actually remove them
997 ... manager.prune_before_turn(conv_id)
998 """
999 from harnessutils.maintenance import scan_and_deduplicate
1001 messages = self.get_messages(conversation_id)
1002 result = scan_and_deduplicate(
1003 messages=messages,
1004 config=self.config.pruning,
1005 )
1007 return result.to_dict()
1009 def get_memory(self, project_id: str, key: str, default: Any = None) -> Any:
1010 """Get a value from project-scoped memory.
1012 Args:
1013 project_id: Project to retrieve memory from
1014 key: Memory key
1015 default: Default value if key not found
1017 Returns:
1018 Stored value or default
1019 """
1020 data = self._load_project_memory(project_id)
1021 return data.get(key, default)
1023 def set_memory(self, project_id: str, key: str, value: Any) -> None:
1024 """Set a value in project-scoped memory.
1026 Args:
1027 project_id: Project to store memory in
1028 key: Memory key
1029 value: Value to store
1030 """
1031 data = self._load_project_memory(project_id)
1032 data[key] = value
1033 try:
1034 self.storage.save_project_memory(project_id, data)
1035 except AttributeError:
1036 pass # Backend doesn't support project memory
1038 def delete_memory(self, project_id: str, key: str) -> None:
1039 """Delete a key from project-scoped memory.
1041 Args:
1042 project_id: Project to delete memory from
1043 key: Memory key to delete
1044 """
1045 data = self._load_project_memory(project_id)
1046 data.pop(key, None)
1047 try:
1048 self.storage.save_project_memory(project_id, data)
1049 except AttributeError:
1050 pass # Backend doesn't support project memory
1052 def list_memory(self, project_id: str) -> dict[str, Any]:
1053 """List all memory keys and values for a project.
1055 Args:
1056 project_id: Project to list memory for
1058 Returns:
1059 Dictionary of all memory key-value pairs
1060 """
1061 return self._load_project_memory(project_id)
1063 def _load_project_memory(self, project_id: str) -> dict[str, Any]:
1064 """Load project memory, returning empty dict on failure.
1066 Args:
1067 project_id: Project to load memory for
1069 Returns:
1070 Memory dictionary (empty if not found or backend unsupported)
1071 """
1072 try:
1073 return self.storage.load_project_memory(project_id)
1074 except (FileNotFoundError, AttributeError):
1075 return {}
1077 def detect_context_issues(self, conversation_id: str) -> list[dict[str, Any]]:
1078 """Detect quality and drift issues in conversation context.
1080 Analyzes conversation for common problems:
1081 - High redundancy (duplicate content)
1082 - Staleness accumulation (old messages)
1083 - Low information density
1084 - Error preservation status
1085 - Excessive protection (limiting pruning)
1087 Args:
1088 conversation_id: Conversation to analyze
1090 Returns:
1091 List of issue dictionaries, each containing:
1092 - issue_type (str): Type of issue
1093 - severity (str): "info", "warning", or "error"
1094 - description (str): Human-readable description
1095 - affected_count (int): Number of items affected
1096 - suggested_fix (str): Recommended action
1097 - metadata (dict): Additional details
1099 Example:
1100 >>> issues = manager.detect_context_issues(conv_id)
1101 >>> for issue in issues:
1102 ... if issue['severity'] == 'warning':
1103 ... print(f"⚠️ {issue['description']}")
1104 ... print(f" Fix: {issue['suggested_fix']}")
1105 """
1106 from harnessutils.maintenance import detect_context_issues
1108 conv_data = self.storage.load_conversation(conversation_id)
1109 conv = Conversation.from_dict(conv_data)
1110 messages = self.get_messages(conversation_id)
1112 issues = detect_context_issues(
1113 messages=messages,
1114 conversation=conv,
1115 config=self.config.pruning,
1116 )
1118 return [issue.to_dict() for issue in issues]