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

1"""Main ConversationManager API for harness-utils.""" 

2 

3import time 

4from collections.abc import Callable 

5from pathlib import Path 

6from typing import Any 

7 

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 

25 

26 

27class ConversationManager: 

28 """Main interface for managing conversations with context window management. 

29 

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

36 

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. 

48 

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 

64 

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) 

71 

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) 

78 

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. 

85 

86 Args: 

87 conversation_id: Optional conversation ID (generated if None) 

88 project_id: Optional project ID for grouping 

89 

90 Returns: 

91 New conversation object 

92 """ 

93 if conversation_id is None: 

94 conversation_id = generate_id("conv") 

95 

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 ) 

103 

104 self.storage.save_conversation(conversation_id, conversation.to_dict()) 

105 self._message_cache[conversation_id] = [] 

106 

107 return conversation 

108 

109 def add_message(self, conversation_id: str, message: Message) -> None: 

110 """Add a message to a conversation. 

111 

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

115 

116 If message_hooks.on_after_add_message is set, it is called after 

117 successful storage. Its return value is ignored. 

118 

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) 

126 

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) 

130 

131 self.storage.save_message(conversation_id, message.id, message.to_dict()) 

132 

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

136 

137 if conversation_id not in self._message_cache: 

138 self._message_cache[conversation_id] = [] 

139 self._message_cache[conversation_id].append(message) 

140 

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) 

145 

146 # Track velocity if message has token count 

147 if message.tokens: 

148 tokens_added = message.tokens.total 

149 conv.update_velocity(tokens_added) 

150 

151 self.storage.save_conversation(conversation_id, conv.to_dict()) 

152 

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) 

156 

157 def get_messages(self, conversation_id: str) -> list[Message]: 

158 """Get all messages for a conversation. 

159 

160 Args: 

161 conversation_id: Conversation ID 

162 

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] 

168 

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) 

184 

185 self._message_cache[conversation_id] = messages 

186 return messages 

187 

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. 

199 

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 

215 

216 Returns: 

217 Filtered and paginated messages 

218 

219 Examples: 

220 # Get last 10 messages 

221 messages = manager.query_messages(conv_id, limit=10, order="desc") 

222 

223 # Get messages with errors 

224 messages = manager.query_messages( 

225 conv_id, 

226 filter={"has_errors": True} 

227 ) 

228 

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 

238 

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 ) 

251 

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 ) 

261 

262 # Get all messages and query 

263 messages = self.get_messages(conversation_id) 

264 return query_messages(messages, options) 

265 

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. 

270 

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 

277 

278 Args: 

279 conversation_id: Conversation to summarize 

280 recent_limit: Number of recent messages to include (default: 5) 

281 

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 

291 

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 

300 

301 messages = self.get_messages(conversation_id) 

302 summary = build_context_summary(conversation_id, messages, recent_limit) 

303 return summary.to_dict() 

304 

305 def inspect_context(self, conversation_id: str) -> ContextInspector: 

306 """Create inspector for querying context state. 

307 

308 Provides observability into: 

309 - What's currently in context 

310 - Token counts and breakdowns 

311 - Pruning predictions 

312 - Audit trail of decisions 

313 

314 Args: 

315 conversation_id: Conversation to inspect 

316 

317 Returns: 

318 ContextInspector instance with full query capabilities 

319 

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) 

329 

330 # Load conversation metadata 

331 conv_data = self.storage.load_conversation(conversation_id) 

332 conversation = Conversation.from_dict(conv_data) 

333 

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] 

337 

338 return ContextInspector(messages, self.config, conversation, decisions) 

339 

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. 

348 

349 This is Tier 2 compaction - removes old tool outputs while 

350 preserving conversation structure. 

351 

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. 

359 

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 } 

384 

385 messages = self.get_messages(conversation_id) 

386 

387 conv_data = self.storage.load_conversation(conversation_id) 

388 conv = Conversation.from_dict(conv_data) 

389 already_excluded = set(conv.excluded_outputs) 

390 

391 result = prune_tool_outputs( 

392 messages, 

393 self.config.pruning, 

394 query=query, 

395 summarizer=summarizer, 

396 already_excluded=already_excluded, 

397 ) 

398 

399 # Accumulate new exclusions into conversation metadata 

400 if result.excluded_call_ids: 

401 conv.add_excluded_outputs(result.excluded_call_ids) 

402 

403 conv.metadata["latest_pruning_decisions"] = [d.to_dict() for d in result.decisions] 

404 self.storage.save_conversation(conversation_id, conv.to_dict()) 

405 

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 

414 

415 return result.to_dict() 

416 

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. 

423 

424 Args: 

425 conversation_id: Conversation to check 

426 current_usage: Current token usage 

427 

428 Returns: 

429 True if overflow predicted within lookahead window 

430 """ 

431 if not self.config.compaction.use_predictive: 

432 return False 

433 

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

438 

439 if velocity is None or not velocity.turn_deltas: 

440 return False # No velocity data yet 

441 

442 # Project tokens ahead 

443 lookahead = self.config.compaction.predictive_lookahead 

444 predicted_growth = velocity.predict_tokens_ahead(lookahead) 

445 

446 # Calculate current total and projected total 

447 current_total = current_usage.input + current_usage.cache.read 

448 projected_total = current_total + predicted_growth 

449 

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 

456 

457 return projected_total > safety_threshold 

458 

459 def needs_compaction( 

460 self, 

461 conversation_id: str, 

462 usage: Usage, 

463 ) -> bool: 

464 """Check if conversation needs summarization (Tier 3). 

465 

466 Uses both reactive (overflow) and predictive checks. 

467 

468 Args: 

469 conversation_id: Conversation to check 

470 usage: Token usage from last turn 

471 

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 

482 

483 # Predictive check: will overflow soon 

484 return self.predict_overflow(conversation_id, usage) 

485 

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

495 

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 

502 

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} 

508 

509 messages = self.get_messages(conversation_id) 

510 summary_id = generate_id("msg") 

511 

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 ) 

521 

522 self.add_message(conversation_id, result.summary_message) 

523 

524 return { 

525 "summarized": True, 

526 "summary_message_id": summary_id, 

527 "tokens_used": result.tokens_used.total, 

528 "cost": result.cost, 

529 } 

530 

531 def to_model_format(self, conversation_id: str) -> list[dict[str, Any]]: 

532 """Convert conversation messages to model format for LLM requests. 

533 

534 Args: 

535 conversation_id: Conversation to convert 

536 

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

544 

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. 

551 

552 This counts ALL tokens in the conversation (user messages, assistant 

553 responses, tool outputs, etc.) that will be sent to the model. 

554 

555 Uses cl100k_base tokenizer (GPT-4/Claude) which works for most modern LLMs. 

556 

557 Args: 

558 conversation_id: Conversation to calculate usage for 

559 model: Optional model name (currently unused, defaults to cl100k_base) 

560 

561 Returns: 

562 Exact token count that will be used in context window 

563 """ 

564 from harnessutils.tokens.exact import count_tokens_exact 

565 

566 messages = self.to_model_format(conversation_id) 

567 return count_tokens_exact(messages, model) 

568 

569 def get_tool_output_tokens(self, conversation_id: str) -> dict[str, Any]: 

570 """Get detailed breakdown of token usage for tool outputs. 

571 

572 Args: 

573 conversation_id: Conversation to analyze 

574 

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 

583 

584 messages = self.get_messages(conversation_id) 

585 

586 total = calculate_context_tokens(messages) 

587 by_tool: dict[str, int] = {} 

588 prunable = 0 

589 protected = 0 

590 turns_skipped = 0 

591 

592 for msg in reversed(messages): 

593 if msg.role == "user": 

594 turns_skipped += 1 

595 

596 for part in msg.parts: 

597 from harnessutils.models.parts import ToolPart 

598 from harnessutils.tokens.exact import count_tokens_fast 

599 

600 if not isinstance(part, ToolPart): 

601 continue 

602 

603 if part.state.status != "completed": 

604 continue 

605 

606 if not part.state.output: 

607 continue 

608 

609 tokens = count_tokens_fast(part.state.output) 

610 

611 # Track by tool type 

612 by_tool[part.tool] = by_tool.get(part.tool, 0) + tokens 

613 

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 ) 

620 

621 if is_protected: 

622 protected += tokens 

623 else: 

624 prunable += tokens 

625 

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 } 

633 

634 def get_context_quality(self, conversation_id: str) -> dict[str, Any]: 

635 """Get current quality metrics for conversation. 

636 

637 Calculates all quality metrics, updates history, and returns assessment 

638 with health status and actionable recommendations. 

639 

640 Args: 

641 conversation_id: Conversation to analyze 

642 

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 

649 

650 conv_data = self.storage.load_conversation(conversation_id) 

651 conv = Conversation.from_dict(conv_data) 

652 messages = self.get_messages(conversation_id) 

653 

654 # Get pruning decisions from metadata if available 

655 decisions = conv.metadata.get("latest_pruning_decisions") 

656 

657 snapshot = assess_quality( 

658 messages=messages, 

659 config=self.config.pruning, 

660 decisions=decisions, 

661 ) 

662 

663 # Update history 

664 conv.update_quality_history(snapshot) 

665 self.storage.save_conversation(conversation_id, conv.to_dict()) 

666 

667 return snapshot.to_dict() 

668 

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. 

677 

678 Note: Prefer get_context_quality() which calculates all metrics. 

679 This is for custom/external metrics. 

680 

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 

688 

689 if timestamp is None: 

690 timestamp = int(time.time() * 1000) 

691 

692 conv_data = self.storage.load_conversation(conversation_id) 

693 conv = Conversation.from_dict(conv_data) 

694 

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 ) 

709 

710 conv.update_quality_history(snapshot) 

711 self.storage.save_conversation(conversation_id, conv.to_dict()) 

712 

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. 

720 

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 

725 

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) 

731 

732 history = conv.get_quality_history() 

733 if history is None: 

734 return [] 

735 

736 return history.get_trend(metric, window) 

737 

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

744 

745 Args: 

746 output: Tool output to truncate 

747 tool_name: Name of the tool 

748 

749 Returns: 

750 Potentially truncated output 

751 """ 

752 output_id = generate_id(f"output_{tool_name}") 

753 

754 result = truncate_output( 

755 output=output, 

756 config=self.config.truncation, 

757 output_id=output_id, 

758 ) 

759 

760 if result.truncated and result.output_path: 

761 self.storage.save_truncated_output(result.output_path, output) 

762 

763 return result.content 

764 

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. 

772 

773 Captures full conversation state including: 

774 - All messages with content 

775 - Conversation metadata (velocity, etc.) 

776 - Current configuration 

777 

778 Useful for: 

779 - Debugging (save state before/after changes) 

780 - A/B testing (compare different strategies) 

781 - Reproducibility (restore exact state) 

782 

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

787 

788 Returns: 

789 Created snapshot 

790 

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) 

798 

799 conv_data = self.storage.load_conversation(conversation_id) 

800 conversation = Conversation.from_dict(conv_data) 

801 

802 # Serialize config to dict recursively (convert dataclasses, Paths, etc.) 

803 from dataclasses import asdict, is_dataclass 

804 from pathlib import Path 

805 

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 

820 

821 config_dict = _serialize_config(self.config) 

822 

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 ) 

831 

832 def restore_snapshot(self, snapshot_id: str) -> str: 

833 """Restore conversation from snapshot. 

834 

835 WARNING: This replaces current conversation state with snapshot state. 

836 Consider creating a snapshot of current state first. 

837 

838 Args: 

839 snapshot_id: Snapshot to restore 

840 

841 Returns: 

842 Conversation ID of restored conversation 

843 

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

850 

851 messages, conversation = self.snapshot_manager.restore_snapshot(snapshot) 

852 

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] = {} 

858 

859 # Save conversation 

860 self.storage.save_conversation(snapshot.conversation_id, conversation.to_dict()) 

861 

862 # Save messages 

863 for msg in messages: 

864 self.storage.save_message(snapshot.conversation_id, msg.id, msg.to_dict()) 

865 

866 # Clear cache 

867 if snapshot.conversation_id in self._message_cache: 

868 del self._message_cache[snapshot.conversation_id] 

869 

870 return snapshot.conversation_id 

871 

872 def compare_snapshots( 

873 self, snapshot1_id: str, snapshot2_id: str 

874 ) -> SnapshotDiff | None: 

875 """Compare two snapshots to see what changed. 

876 

877 Args: 

878 snapshot1_id: First snapshot ID (earlier) 

879 snapshot2_id: Second snapshot ID (later) 

880 

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 

889 

890 None if snapshots not found 

891 """ 

892 return self.snapshot_manager.compare_snapshots(snapshot1_id, snapshot2_id) 

893 

894 def export_snapshot(self, snapshot_id: str, file_path: str) -> None: 

895 """Export snapshot to JSON file for version control. 

896 

897 Args: 

898 snapshot_id: Snapshot to export 

899 file_path: Path to write JSON file 

900 

901 Raises: 

902 FileNotFoundError: If snapshot not found 

903 

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) 

910 

911 def import_snapshot(self, file_path: str) -> Snapshot: 

912 """Import snapshot from JSON file. 

913 

914 Args: 

915 file_path: Path to JSON file 

916 

917 Returns: 

918 Imported snapshot 

919 """ 

920 return self.snapshot_manager.import_snapshot(file_path) 

921 

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. 

930 

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. 

933 

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) 

939 

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 

947 

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 

955 

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 ) 

964 

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] 

974 

975 return result.to_dict() 

976 

977 def scan_and_deduplicate(self, conversation_id: str) -> dict[str, Any]: 

978 """Scan for duplicate outputs and return statistics. 

979 

980 Identifies duplicate tool outputs without removing them. 

981 Use prune_before_turn() to actually remove duplicates. 

982 

983 Args: 

984 conversation_id: Conversation to scan 

985 

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 

991 

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 

1000 

1001 messages = self.get_messages(conversation_id) 

1002 result = scan_and_deduplicate( 

1003 messages=messages, 

1004 config=self.config.pruning, 

1005 ) 

1006 

1007 return result.to_dict() 

1008 

1009 def get_memory(self, project_id: str, key: str, default: Any = None) -> Any: 

1010 """Get a value from project-scoped memory. 

1011 

1012 Args: 

1013 project_id: Project to retrieve memory from 

1014 key: Memory key 

1015 default: Default value if key not found 

1016 

1017 Returns: 

1018 Stored value or default 

1019 """ 

1020 data = self._load_project_memory(project_id) 

1021 return data.get(key, default) 

1022 

1023 def set_memory(self, project_id: str, key: str, value: Any) -> None: 

1024 """Set a value in project-scoped memory. 

1025 

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 

1037 

1038 def delete_memory(self, project_id: str, key: str) -> None: 

1039 """Delete a key from project-scoped memory. 

1040 

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 

1051 

1052 def list_memory(self, project_id: str) -> dict[str, Any]: 

1053 """List all memory keys and values for a project. 

1054 

1055 Args: 

1056 project_id: Project to list memory for 

1057 

1058 Returns: 

1059 Dictionary of all memory key-value pairs 

1060 """ 

1061 return self._load_project_memory(project_id) 

1062 

1063 def _load_project_memory(self, project_id: str) -> dict[str, Any]: 

1064 """Load project memory, returning empty dict on failure. 

1065 

1066 Args: 

1067 project_id: Project to load memory for 

1068 

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 {} 

1076 

1077 def detect_context_issues(self, conversation_id: str) -> list[dict[str, Any]]: 

1078 """Detect quality and drift issues in conversation context. 

1079 

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) 

1086 

1087 Args: 

1088 conversation_id: Conversation to analyze 

1089 

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 

1098 

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 

1107 

1108 conv_data = self.storage.load_conversation(conversation_id) 

1109 conv = Conversation.from_dict(conv_data) 

1110 messages = self.get_messages(conversation_id) 

1111 

1112 issues = detect_context_issues( 

1113 messages=messages, 

1114 conversation=conv, 

1115 config=self.config.pruning, 

1116 ) 

1117 

1118 return [issue.to_dict() for issue in issues]