Coverage for src / harnessutils / compaction / pruning.py: 87%

1"""Tier 2: Selective pruning of tool outputs. 

2 

3Removes old tool outputs while preserving conversation structure. 

4Cost: Cheap (~50ms), Latency: ~50ms. 

5""" 

6 

7import hashlib 

8import math 

9import time 

10from collections.abc import Callable 

11from dataclasses import dataclass, field 

12from typing import Any, Literal 

13 

14from harnessutils.config import PruningConfig 

15from harnessutils.models.message import Message 

16from harnessutils.models.parts import ToolPart 

17from harnessutils.tokens.exact import count_tokens_fast 

18 

19 

20@dataclass 

21class PruningDecision: 

22 """Individual pruning decision for a tool output. 

23 

24 Tracks why a specific output was kept or pruned, enabling full 

25 audit trail and debugging of pruning behavior. 

26 """ 

27 

28 message_id: str # Message containing this output 

29 part_id: str # Tool part ID 

30 tool: str # Tool name 

31 decision: Literal[ 

32 "kept", # Output kept in context 

33 "pruned_duplicate", # Removed as duplicate 

34 "pruned_low_importance", # Removed due to low importance score 

35 "pruned_fifo", # Removed by FIFO (oldest first) 

36 "protected_recent", # Protected by recency (protect_turns) 

37 "protected_tool", # Protected by tool type (protected_tools) 

38 "protected_summary", # Protected (after summary boundary) 

39 ] 

40 importance_score: float | None = None # Importance score (if scored) 

41 duplicate_of: str | None = None # Part ID of duplicate (if duplicate) 

42 tokens: int = 0 # Token count of this output 

43 timestamp: int = 0 # When decision was made (unix ms) 

44 metadata: dict[str, Any] = field(default_factory=dict) # Additional context 

45 

46 def to_dict(self) -> dict[str, Any]: 

47 """Convert to dictionary for serialization.""" 

48 return { 

49 "message_id": self.message_id, 

50 "part_id": self.part_id, 

51 "tool": self.tool, 

52 "decision": self.decision, 

53 "importance_score": self.importance_score, 

54 "duplicate_of": self.duplicate_of, 

55 "tokens": self.tokens, 

56 "timestamp": self.timestamp, 

57 "metadata": self.metadata, 

58 } 

59 

60 @classmethod 

61 def from_dict(cls, data: dict[str, Any]) -> "PruningDecision": 

62 """Create from dictionary. 

63 

64 Args: 

65 data: Dictionary representation 

66 

67 Returns: 

68 PruningDecision instance 

69 """ 

70 return cls( 

71 message_id=data["message_id"], 

72 part_id=data["part_id"], 

73 tool=data["tool"], 

74 decision=data["decision"], 

75 importance_score=data.get("importance_score"), 

76 duplicate_of=data.get("duplicate_of"), 

77 tokens=data.get("tokens", 0), 

78 timestamp=data.get("timestamp", 0), 

79 metadata=data.get("metadata", {}), 

80 ) 

81 

82 @property 

83 def was_pruned(self) -> bool: 

84 """Check if this output was pruned.""" 

85 return self.decision.startswith("pruned_") 

86 

87 @property 

88 def was_protected(self) -> bool: 

89 """Check if this output was protected from pruning.""" 

90 return self.decision.startswith("protected_") 

91 

92 

93@dataclass 

94class PruningResult: 

95 """Result of pruning operation with detailed token tracking.""" 

96 

97 pruned: int # Total outputs pruned 

98 tokens_saved: int # Total tokens saved 

99 tokens_before: int = 0 # Token count before pruning 

100 tokens_after: int = 0 # Token count after pruning 

101 duplicates_pruned: int = 0 # Outputs pruned due to duplication 

102 importance_pruned: int = 0 # Outputs pruned due to low importance 

103 duplicate_tokens_saved: int = 0 # Tokens saved from deduplication 

104 importance_tokens_saved: int = 0 # Tokens saved from importance pruning 

105 decisions: list[PruningDecision] = field(default_factory=list) # Full decision log 

106 

107 def to_dict(self) -> dict[str, int | float]: 

108 """Convert to dictionary for reporting. 

109 

110 Returns: 

111 Dictionary with all pruning metrics 

112 """ 

113 return { 

114 "pruned": self.pruned, 

115 "tokens_saved": self.tokens_saved, 

116 "tokens_before": self.tokens_before, 

117 "tokens_after": self.tokens_after, 

118 "duplicates_pruned": self.duplicates_pruned, 

119 "importance_pruned": self.importance_pruned, 

120 "duplicate_tokens_saved": self.duplicate_tokens_saved, 

121 "importance_tokens_saved": self.importance_tokens_saved, 

122 "reduction_percent": round( 

123 (self.tokens_saved / self.tokens_before * 100) if self.tokens_before > 0 else 0, 

124 1, 

125 ), 

126 } 

127 

128 def __str__(self) -> str: 

129 """Human-readable summary of pruning results.""" 

130 if self.pruned == 0: 

131 return "No pruning needed" 

132 

133 lines = [ 

134 f"Pruned {self.pruned} outputs, saved {self.tokens_saved:,} tokens", 

135 f" Before: {self.tokens_before:,} tokens", 

136 f" After: {self.tokens_after:,} tokens", 

137 f" Reduction: {self.to_dict()['reduction_percent']}%", 

138 ] 

139 

140 if self.duplicates_pruned > 0: 

141 lines.append( 

142 f" Duplicates: {self.duplicates_pruned} removed " 

143 f"({self.duplicate_tokens_saved:,} tokens)" 

144 ) 

145 

146 if self.importance_pruned > 0: 

147 lines.append( 

148 f" Low importance: {self.importance_pruned} removed " 

149 f"({self.importance_tokens_saved:,} tokens)" 

150 ) 

151 

152 return "\n".join(lines) 

153 

154 def get_pruned_decisions(self) -> list[PruningDecision]: 

155 """Get all decisions that resulted in pruning. 

156 

157 Returns: 

158 List of decisions where output was pruned 

159 """ 

160 return [d for d in self.decisions if d.was_pruned] 

161 

162 def get_protected_decisions(self) -> list[PruningDecision]: 

163 """Get all decisions that resulted in protection. 

164 

165 Returns: 

166 List of decisions where output was protected 

167 """ 

168 return [d for d in self.decisions if d.was_protected] 

169 

170 def get_kept_decisions(self) -> list[PruningDecision]: 

171 """Get all decisions where output was kept. 

172 

173 Returns: 

174 List of decisions where output was kept (not pruned or protected) 

175 """ 

176 return [d for d in self.decisions if d.decision == "kept"] 

177 

178 def get_decision_by_part(self, part_id: str) -> PruningDecision | None: 

179 """Find decision for specific tool part. 

180 

181 Args: 

182 part_id: Tool part ID to find 

183 

184 Returns: 

185 Decision for that part, or None if not found 

186 """ 

187 for decision in self.decisions: 

188 if decision.part_id == part_id: 

189 return decision 

190 return None 

191 

192 

193@dataclass 

194class OutputImportance: 

195 """Importance score components for a tool output.""" 

196 

197 recency_score: float # Exponential decay based on age 

198 size_penalty: float # Penalty for large outputs 

199 semantic_score: float # Content-based importance (errors, warnings) 

200 tool_priority: float # Tool type importance 

201 token_count: int # Actual token count 

202 query_relevance: float = 0.0 # Relevance to current user query (0 when no query) 

203 

204 @property 

205 def total_score(self) -> float: 

206 """Calculate weighted total importance score. 

207 

208 Higher score = more important = keep longer. 

209 Lower score = less important = prune first. 

210 """ 

211 return ( 

212 self.recency_score 

213 + self.size_penalty 

214 + self.semantic_score 

215 + self.tool_priority 

216 + self.query_relevance 

217 ) 

218 

219 

220_STOP_WORDS = frozenset({ 

221 "the", "a", "an", "is", "are", "was", "were", "be", "been", "have", "has", 

222 "had", "do", "does", "did", "will", "would", "could", "should", "may", 

223 "might", "in", "on", "at", "to", "for", "of", "and", "or", "but", "it", 

224 "this", "that", "with", "as", "by", "from", "not", "what", "how", "i", 

225}) 

226 

227 

228def _compute_query_relevance(query: str, output: str, weight: float) -> float: 

229 """Score how relevant a tool output is to the current user query. 

230 

231 Uses word-level coverage: what fraction of meaningful query terms appear 

232 in the output. Filters stop words so common words don't inflate scores. 

233 

234 Args: 

235 query: Current user query 

236 output: Tool output text to score 

237 weight: Scaling weight from PruningConfig.query_weight 

238 

239 Returns: 

240 Relevance score (0.0 when no query or no meaningful terms) 

241 """ 

242 query_words = set(query.lower().split()) - _STOP_WORDS 

243 if not query_words or not output: 

244 return 0.0 

245 output_word_set = set(output.lower().split()) 

246 coverage = len(query_words & output_word_set) / len(query_words) 

247 return coverage * 100 * weight 

248 

249 

250def generate_shingles(text: str, n: int = 5) -> set[str]: 

251 """Generate word-level n-grams (shingles) for similarity detection. 

252 

253 Args: 

254 text: Text to generate shingles from 

255 n: Size of n-grams (default: 5 words) 

256 

257 Returns: 

258 Set of n-gram strings 

259 """ 

260 # Normalize text 

261 words = text.lower().split() 

262 

263 if len(words) < n: 

264 # If text too short, use the whole thing 

265 return {" ".join(words)} if words else set() 

266 

267 # Generate n-grams 

268 shingles = set() 

269 for i in range(len(words) - n + 1): 

270 shingle = " ".join(words[i : i + n]) 

271 shingles.add(shingle) 

272 

273 return shingles 

274 

275 

276def jaccard_similarity(set1: set[str], set2: set[str]) -> float: 

277 """Calculate Jaccard similarity between two sets. 

278 

279 Jaccard similarity = |intersection| / |union| 

280 Returns value between 0.0 (no overlap) and 1.0 (identical). 

281 

282 Args: 

283 set1: First set 

284 set2: Second set 

285 

286 Returns: 

287 Similarity score [0.0, 1.0] 

288 """ 

289 if not set1 or not set2: 

290 return 0.0 

291 

292 intersection = len(set1 & set2) 

293 union = len(set1 | set2) 

294 

295 if union == 0: 

296 return 0.0 

297 

298 return intersection / union 

299 

300 

301def compute_content_hash(text: str) -> str: 

302 """Compute fast hash of text content for exact duplicate detection. 

303 

304 Args: 

305 text: Text to hash 

306 

307 Returns: 

308 MD5 hash hex string 

309 """ 

310 return hashlib.md5(text.encode("utf-8")).hexdigest() 

311 

312 

313def find_duplicate_output( 

314 part: ToolPart, 

315 recent_parts: list[tuple[ToolPart, set[str]]], 

316 similarity_threshold: float = 0.8, 

317) -> ToolPart | None: 

318 """Find if output is duplicate of a recent output. 

319 

320 Args: 

321 part: Tool part to check for duplication 

322 recent_parts: List of (part, shingles) tuples to compare against 

323 similarity_threshold: Minimum similarity to consider duplicate (default: 0.8) 

324 

325 Returns: 

326 The duplicate part if found, None otherwise 

327 """ 

328 # Only compare against same tool type 

329 # Different tools are unlikely to produce identical outputs 

330 # This prevents false positives across tool types 

331 same_tool_parts = [ 

332 (p, shingles) for p, shingles in recent_parts if p.tool == part.tool 

333 ] 

334 

335 if not same_tool_parts: 

336 return None 

337 

338 # Fast path: exact duplicate check (within same tool type) 

339 current_hash = compute_content_hash(part.state.output) 

340 

341 for recent_part, _ in same_tool_parts: 

342 if compute_content_hash(recent_part.state.output) == current_hash: 

343 return recent_part # Exact duplicate 

344 

345 # Generate shingles for current part 

346 current_shingles = generate_shingles(part.state.output) 

347 

348 # Check similarity against recent parts of same tool 

349 for recent_part, recent_shingles in same_tool_parts: 

350 similarity = jaccard_similarity(current_shingles, recent_shingles) 

351 

352 if similarity >= similarity_threshold: 

353 return recent_part # Found similar output 

354 

355 return None 

356 

357 

358def calculate_context_tokens(messages: list[Message]) -> int: 

359 """Calculate total token count for all tool outputs in conversation. 

360 

361 Args: 

362 messages: All conversation messages 

363 

364 Returns: 

365 Total token count 

366 """ 

367 total = 0 

368 for msg in messages: 

369 for part in msg.parts: 

370 if isinstance(part, ToolPart) and part.state.status == "completed": 

371 if part.state.output: # Only count non-empty outputs 

372 total += count_tokens_fast(part.state.output) 

373 return total 

374 

375 

376def calculate_turn_age(messages: list[Message], target_msg: Message) -> int: 

377 """Calculate how many turns ago a message was created. 

378 

379 Args: 

380 messages: All conversation messages 

381 target_msg: Message to calculate age for 

382 

383 Returns: 

384 Number of user turns since this message (0 = current turn) 

385 """ 

386 turn_count = 0 

387 for msg in reversed(messages): 

388 if msg.id == target_msg.id: 

389 return turn_count 

390 if msg.role == "user": 

391 turn_count += 1 

392 return turn_count 

393 

394 

395def score_tool_output( 

396 part: ToolPart, 

397 message: Message, 

398 messages: list[Message], 

399 config: PruningConfig, 

400 query: str | None = None, 

401) -> OutputImportance: 

402 """Calculate importance score for a tool output. 

403 

404 Args: 

405 part: Tool part to score 

406 message: Message containing this part 

407 messages: All conversation messages 

408 config: Pruning configuration with scoring weights 

409 query: Current user query for relevance boosting (optional) 

410 

411 Returns: 

412 OutputImportance with all score components 

413 """ 

414 token_count = count_tokens_fast(part.state.output) 

415 

416 # 1. Recency score (exponential decay) 

417 age = calculate_turn_age(messages, message) 

418 recency = math.exp(-age * config.recency_decay) * 100 * config.recency_weight 

419 

420 # 2. Size penalty (prefer removing large outputs) 

421 size_penalty = math.log(token_count + 1) * 10 * config.size_weight 

422 

423 # 3. Semantic importance 

424 semantic = 0.0 

425 output_lower = part.state.output.lower() 

426 

427 # Check for errors 

428 if any( 

429 keyword in output_lower 

430 for keyword in ["error", "exception", "traceback", "failed"] 

431 ): 

432 semantic += config.error_boost 

433 

434 # Check for warnings 

435 if "warning" in output_lower or "warn" in output_lower: 

436 semantic += config.warning_boost 

437 

438 # Check if user explicitly requested this 

439 if part.state.metadata and part.state.metadata.get("user_requested"): 

440 semantic += config.user_requested_boost 

441 

442 semantic *= config.semantic_weight 

443 

444 # 4. Tool type priority 

445 tool_priority = config.tool_importance.get(part.tool, 50.0) * config.tool_priority_weight 

446 

447 # 5. Query relevance (boost outputs semantically related to current query) 

448 query_relevance = _compute_query_relevance(query or "", part.state.output, config.query_weight) 

449 

450 return OutputImportance( 

451 recency_score=recency, 

452 size_penalty=size_penalty, 

453 semantic_score=semantic, 

454 tool_priority=tool_priority, 

455 token_count=token_count, 

456 query_relevance=query_relevance, 

457 ) 

458 

459 

460def prune_tool_outputs( 

461 messages: list[Message], 

462 config: PruningConfig, 

463 query: str | None = None, 

464 summarizer: Callable[[str, str], str] | None = None, 

465) -> PruningResult: 

466 """Prune tool outputs from conversation history. 

467 

468 Selectively removes old tool outputs while preserving: 

469 - Tool call metadata (name, input, title, timing) 

470 - Recent outputs (within protection window) 

471 - Protected tool outputs 

472 - Last N turns 

473 

474 Uses exact token counting via tiktoken for accurate pruning decisions. 

475 

476 First runs duplicate detection (if enabled), then applies either 

477 importance-based or FIFO pruning strategy. 

478 

479 Args: 

480 messages: Conversation messages (newest first recommended) 

481 config: Pruning configuration 

482 query: Current user query for relevance-based importance boosting (optional) 

483 summarizer: Optional callable(tool_name, output) -> summary_text. When provided, 

484 called before clearing each output so a one-sentence summary is preserved. 

485 

486 Returns: 

487 PruningResult with detailed token tracking and breakdown 

488 """ 

489 # Calculate token usage before pruning 

490 tokens_before = calculate_context_tokens(messages) 

491 

492 # Phase 1: Duplicate detection (runs regardless of scoring strategy) 

493 duplicate_result = PruningResult(pruned=0, tokens_saved=0) 

494 if config.detect_duplicates: 

495 duplicate_result = _detect_and_prune_duplicates(messages, config, summarizer=summarizer) 

496 

497 # Phase 2: Importance-based or FIFO pruning 

498 if config.use_importance_scoring: 

499 pruning_result = _prune_with_importance_scoring_only(messages, config, query=query, summarizer=summarizer) 

500 else: 

501 pruning_result = _prune_simple(messages, config, summarizer=summarizer) 

502 

503 # Calculate token usage after pruning 

504 tokens_after = calculate_context_tokens(messages) 

505 

506 # Combine results with detailed tracking 

507 all_decisions = duplicate_result.decisions + pruning_result.decisions 

508 

509 return PruningResult( 

510 pruned=duplicate_result.pruned + pruning_result.pruned, 

511 tokens_saved=duplicate_result.tokens_saved + pruning_result.tokens_saved, 

512 tokens_before=tokens_before, 

513 tokens_after=tokens_after, 

514 duplicates_pruned=duplicate_result.pruned, 

515 importance_pruned=pruning_result.pruned, 

516 duplicate_tokens_saved=duplicate_result.tokens_saved, 

517 importance_tokens_saved=pruning_result.tokens_saved, 

518 decisions=all_decisions, 

519 ) 

520 

521 

522def _prune_simple( 

523 messages: list[Message], 

524 config: PruningConfig, 

525 summarizer: Callable[[str, str], str] | None = None, 

526) -> PruningResult: 

527 """Simple FIFO pruning (original algorithm). 

528 

529 Prunes oldest outputs first when token budget exceeded. 

530 """ 

531 total_tokens = 0 

532 prunable_tokens = 0 

533 to_prune: list[tuple[Message, ToolPart, int]] = [] 

534 decisions: list[PruningDecision] = [] 

535 turns_skipped = 0 

536 now_ms = int(time.time() * 1000) 

537 

538 for msg in reversed(messages): 

539 if msg.role == "user": 

540 turns_skipped += 1 

541 

542 # Check if past summary boundary 

543 past_summary = msg.summary 

544 

545 for part in msg.parts: 

546 if not isinstance(part, ToolPart): 

547 continue 

548 

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

550 continue 

551 

552 token_count = count_tokens_fast(part.state.output) if part.state.output else 0 

553 

554 # Track protection reasons 

555 if turns_skipped < config.protect_turns: 

556 decisions.append( 

557 PruningDecision( 

558 message_id=msg.id, 

559 part_id=part.call_id, 

560 tool=part.tool, 

561 decision="protected_recent", 

562 tokens=token_count, 

563 timestamp=now_ms, 

564 metadata={"turns_back": turns_skipped}, 

565 ) 

566 ) 

567 continue 

568 

569 if past_summary: 

570 decisions.append( 

571 PruningDecision( 

572 message_id=msg.id, 

573 part_id=part.call_id, 

574 tool=part.tool, 

575 decision="protected_summary", 

576 tokens=token_count, 

577 timestamp=now_ms, 

578 ) 

579 ) 

580 break 

581 

582 if part.tool in config.protected_tools: 

583 decisions.append( 

584 PruningDecision( 

585 message_id=msg.id, 

586 part_id=part.call_id, 

587 tool=part.tool, 

588 decision="protected_tool", 

589 tokens=token_count, 

590 timestamp=now_ms, 

591 metadata={"protected_tools": config.protected_tools}, 

592 ) 

593 ) 

594 continue 

595 

596 if part.state.time and part.state.time.compacted: 

597 # Already compacted, skip (don't track as decision) 

598 continue 

599 

600 total_tokens += token_count 

601 

602 if total_tokens > config.prune_protect: 

603 prunable_tokens += token_count 

604 to_prune.append((msg, part, token_count)) 

605 else: 

606 # Kept within budget 

607 decisions.append( 

608 PruningDecision( 

609 message_id=msg.id, 

610 part_id=part.call_id, 

611 tool=part.tool, 

612 decision="kept", 

613 tokens=token_count, 

614 timestamp=now_ms, 

615 metadata={"total_tokens_so_far": total_tokens}, 

616 ) 

617 ) 

618 

619 if prunable_tokens > config.prune_minimum: 

620 # Actually prune 

621 for msg, part, tokens in to_prune: 

622 if summarizer and part.state.output: 

623 part.state.summary_text = summarizer(part.tool, part.state.output) 

624 part.state.output = "" 

625 part.state.attachments = [] 

626 if part.state.time: 

627 part.state.time.compacted = int(time.time() * 1000) 

628 

629 # Record pruning decision 

630 decisions.append( 

631 PruningDecision( 

632 message_id=msg.id, 

633 part_id=part.call_id, 

634 tool=part.tool, 

635 decision="pruned_fifo", 

636 tokens=tokens, 

637 timestamp=now_ms, 

638 metadata={ 

639 "prunable_tokens": prunable_tokens, 

640 "prune_minimum": config.prune_minimum, 

641 }, 

642 ) 

643 ) 

644 

645 return PruningResult( 

646 pruned=len(to_prune), tokens_saved=prunable_tokens, decisions=decisions 

647 ) 

648 

649 # Not enough to prune - mark all candidates as kept 

650 for msg, part, tokens in to_prune: 

651 decisions.append( 

652 PruningDecision( 

653 message_id=msg.id, 

654 part_id=part.call_id, 

655 tool=part.tool, 

656 decision="kept", 

657 tokens=tokens, 

658 timestamp=now_ms, 

659 metadata={ 

660 "reason": "below_minimum_threshold", 

661 "prunable_tokens": prunable_tokens, 

662 "prune_minimum": config.prune_minimum, 

663 }, 

664 ) 

665 ) 

666 

667 return PruningResult(pruned=0, tokens_saved=0, decisions=decisions) 

668 

669 

670def _prune_with_importance_scoring_only( 

671 messages: list[Message], 

672 config: PruningConfig, 

673 query: str | None = None, 

674 summarizer: Callable[[str, str], str] | None = None, 

675) -> PruningResult: 

676 """Smart pruning using importance scoring. 

677 

678 Scores all outputs by importance and prunes lowest-value first. 

679 Preserves critical outputs (errors, warnings, user-requested). 

680 

681 Note: Duplicate detection happens separately before this function is called. 

682 """ 

683 # Collect all prunable outputs with scores 

684 scored_outputs: list[tuple[Message, ToolPart, OutputImportance]] = [] 

685 decisions: list[PruningDecision] = [] 

686 turns_skipped = 0 

687 now_ms = int(time.time() * 1000) 

688 

689 for msg in reversed(messages): 

690 if msg.role == "user": 

691 turns_skipped += 1 

692 

693 past_summary = msg.summary 

694 

695 for part in msg.parts: 

696 if not isinstance(part, ToolPart): 

697 continue 

698 

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

700 continue 

701 

702 token_count = count_tokens_fast(part.state.output) if part.state.output else 0 

703 

704 # Track protection reasons (same as FIFO) 

705 if turns_skipped < config.protect_turns: 

706 decisions.append( 

707 PruningDecision( 

708 message_id=msg.id, 

709 part_id=part.call_id, 

710 tool=part.tool, 

711 decision="protected_recent", 

712 tokens=token_count, 

713 timestamp=now_ms, 

714 metadata={"turns_back": turns_skipped}, 

715 ) 

716 ) 

717 continue 

718 

719 if past_summary: 

720 decisions.append( 

721 PruningDecision( 

722 message_id=msg.id, 

723 part_id=part.call_id, 

724 tool=part.tool, 

725 decision="protected_summary", 

726 tokens=token_count, 

727 timestamp=now_ms, 

728 ) 

729 ) 

730 break 

731 

732 if part.tool in config.protected_tools: 

733 decisions.append( 

734 PruningDecision( 

735 message_id=msg.id, 

736 part_id=part.call_id, 

737 tool=part.tool, 

738 decision="protected_tool", 

739 tokens=token_count, 

740 timestamp=now_ms, 

741 metadata={"protected_tools": config.protected_tools}, 

742 ) 

743 ) 

744 continue 

745 

746 if part.state.time and part.state.time.compacted: 

747 # Already compacted, skip 

748 continue 

749 

750 # Score this output 

751 importance = score_tool_output(part, msg, messages, config, query=query) 

752 scored_outputs.append((msg, part, importance)) 

753 

754 # Calculate total tokens 

755 total_tokens = sum(imp.token_count for _, _, imp in scored_outputs) 

756 

757 # If under budget, mark all as kept 

758 if total_tokens <= config.prune_protect: 

759 for msg, part, importance in scored_outputs: 

760 decisions.append( 

761 PruningDecision( 

762 message_id=msg.id, 

763 part_id=part.call_id, 

764 tool=part.tool, 

765 decision="kept", 

766 importance_score=importance.total_score, 

767 tokens=importance.token_count, 

768 timestamp=now_ms, 

769 metadata={"total_tokens": total_tokens, "prune_protect": config.prune_protect}, 

770 ) 

771 ) 

772 return PruningResult(pruned=0, tokens_saved=0, decisions=decisions) 

773 

774 # Sort by importance (lowest score first = prune first) 

775 scored_outputs.sort(key=lambda x: x[2].total_score) 

776 

777 # Prune until we're under budget 

778 to_prune: list[tuple[Message, ToolPart, OutputImportance]] = [] 

779 current_tokens = total_tokens 

780 

781 for msg, part, importance in scored_outputs: 

782 if current_tokens <= config.prune_protect: 

783 # Keep this one 

784 decisions.append( 

785 PruningDecision( 

786 message_id=msg.id, 

787 part_id=part.call_id, 

788 tool=part.tool, 

789 decision="kept", 

790 importance_score=importance.total_score, 

791 tokens=importance.token_count, 

792 timestamp=now_ms, 

793 metadata={ 

794 "current_tokens": current_tokens, 

795 "prune_protect": config.prune_protect, 

796 }, 

797 ) 

798 ) 

799 else: 

800 # Mark for pruning 

801 to_prune.append((msg, part, importance)) 

802 current_tokens -= importance.token_count 

803 

804 # Only prune if savings meet minimum threshold 

805 tokens_saved = sum(imp.token_count for _, _, imp in to_prune) 

806 

807 if tokens_saved >= config.prune_minimum: 

808 # Actually prune 

809 for msg, part, importance in to_prune: 

810 if summarizer and part.state.output: 

811 part.state.summary_text = summarizer(part.tool, part.state.output) 

812 part.state.output = "" 

813 part.state.attachments = [] 

814 if part.state.time: 

815 part.state.time.compacted = int(time.time() * 1000) 

816 

817 # Record pruning decision with importance score 

818 decisions.append( 

819 PruningDecision( 

820 message_id=msg.id, 

821 part_id=part.call_id, 

822 tool=part.tool, 

823 decision="pruned_low_importance", 

824 importance_score=importance.total_score, 

825 tokens=importance.token_count, 

826 timestamp=now_ms, 

827 metadata={ 

828 "recency_score": importance.recency_score, 

829 "size_penalty": importance.size_penalty, 

830 "semantic_score": importance.semantic_score, 

831 "tool_priority": importance.tool_priority, 

832 "tokens_saved": tokens_saved, 

833 }, 

834 ) 

835 ) 

836 

837 return PruningResult(pruned=len(to_prune), tokens_saved=tokens_saved, decisions=decisions) 

838 

839 # Not enough savings - keep all 

840 for msg, part, importance in to_prune: 

841 decisions.append( 

842 PruningDecision( 

843 message_id=msg.id, 

844 part_id=part.call_id, 

845 tool=part.tool, 

846 decision="kept", 

847 importance_score=importance.total_score, 

848 tokens=importance.token_count, 

849 timestamp=now_ms, 

850 metadata={ 

851 "reason": "below_minimum_threshold", 

852 "tokens_saved": tokens_saved, 

853 "prune_minimum": config.prune_minimum, 

854 }, 

855 ) 

856 ) 

857 

858 return PruningResult(pruned=0, tokens_saved=0, decisions=decisions) 

859 

860 

861def _detect_and_prune_duplicates( 

862 messages: list[Message], 

863 config: PruningConfig, 

864 summarizer: Callable[[str, str], str] | None = None, 

865) -> PruningResult: 

866 """Detect and prune duplicate tool outputs. 

867 

868 Uses similarity hashing (shingles + Jaccard) to find near-duplicates. 

869 Aggressively prunes older duplicates while keeping the most recent. 

870 

871 Args: 

872 messages: Conversation messages 

873 config: Pruning configuration 

874 summarizer: Optional callable(tool_name, output) -> summary_text 

875 

876 Returns: 

877 PruningResult with duplicates pruned 

878 """ 

879 # Build index of recent outputs with their shingles 

880 recent_outputs: list[tuple[ToolPart, set[str], Message]] = [] 

881 # (msg, part, tokens, duplicate_of_id) 

882 duplicates_to_prune: list[tuple[Message, ToolPart, int, str]] = []