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

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 excluded_call_ids: list[str] = field(default_factory=list) # Call IDs to add to conversation exclusion list 

107 

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

109 """Convert to dictionary for reporting. 

110 

111 Returns: 

112 Dictionary with all pruning metrics 

113 """ 

114 return { 

115 "pruned": self.pruned, 

116 "tokens_saved": self.tokens_saved, 

117 "tokens_before": self.tokens_before, 

118 "tokens_after": self.tokens_after, 

119 "duplicates_pruned": self.duplicates_pruned, 

120 "importance_pruned": self.importance_pruned, 

121 "duplicate_tokens_saved": self.duplicate_tokens_saved, 

122 "importance_tokens_saved": self.importance_tokens_saved, 

123 "reduction_percent": round( 

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

125 1, 

126 ), 

127 } 

128 

129 def __str__(self) -> str: 

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

131 if self.pruned == 0: 

132 return "No pruning needed" 

133 

134 lines = [ 

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

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

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

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

139 ] 

140 

141 if self.duplicates_pruned > 0: 

142 lines.append( 

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

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

145 ) 

146 

147 if self.importance_pruned > 0: 

148 lines.append( 

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

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

151 ) 

152 

153 return "\n".join(lines) 

154 

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

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

157 

158 Returns: 

159 List of decisions where output was pruned 

160 """ 

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

162 

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

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

165 

166 Returns: 

167 List of decisions where output was protected 

168 """ 

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

170 

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

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

173 

174 Returns: 

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

176 """ 

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

178 

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

180 """Find decision for specific tool part. 

181 

182 Args: 

183 part_id: Tool part ID to find 

184 

185 Returns: 

186 Decision for that part, or None if not found 

187 """ 

188 for decision in self.decisions: 

189 if decision.part_id == part_id: 

190 return decision 

191 return None 

192 

193 

194@dataclass 

195class OutputImportance: 

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

197 

198 recency_score: float # Exponential decay based on age 

199 size_penalty: float # Penalty for large outputs 

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

201 tool_priority: float # Tool type importance 

202 token_count: int # Actual token count 

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

204 

205 @property 

206 def total_score(self) -> float: 

207 """Calculate weighted total importance score. 

208 

209 Higher score = more important = keep longer. 

210 Lower score = less important = prune first. 

211 """ 

212 return ( 

213 self.recency_score 

214 + self.size_penalty 

215 + self.semantic_score 

216 + self.tool_priority 

217 + self.query_relevance 

218 ) 

219 

220 

221_STOP_WORDS = frozenset({ 

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

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

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

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

226}) 

227 

228 

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

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

231 

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

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

234 

235 Args: 

236 query: Current user query 

237 output: Tool output text to score 

238 weight: Scaling weight from PruningConfig.query_weight 

239 

240 Returns: 

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

242 """ 

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

244 if not query_words or not output: 

245 return 0.0 

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

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

248 return coverage * 100 * weight 

249 

250 

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

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

253 

254 Args: 

255 text: Text to generate shingles from 

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

257 

258 Returns: 

259 Set of n-gram strings 

260 """ 

261 # Normalize text 

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

263 

264 if len(words) < n: 

265 # If text too short, use the whole thing 

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

267 

268 # Generate n-grams 

269 shingles = set() 

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

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

272 shingles.add(shingle) 

273 

274 return shingles 

275 

276 

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

278 """Calculate Jaccard similarity between two sets. 

279 

280 Jaccard similarity = |intersection| / |union| 

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

282 

283 Args: 

284 set1: First set 

285 set2: Second set 

286 

287 Returns: 

288 Similarity score [0.0, 1.0] 

289 """ 

290 if not set1 or not set2: 

291 return 0.0 

292 

293 intersection = len(set1 & set2) 

294 union = len(set1 | set2) 

295 

296 if union == 0: 

297 return 0.0 

298 

299 return intersection / union 

300 

301 

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

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

304 

305 Args: 

306 text: Text to hash 

307 

308 Returns: 

309 MD5 hash hex string 

310 """ 

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

312 

313 

314def find_duplicate_output( 

315 part: ToolPart, 

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

317 similarity_threshold: float = 0.8, 

318) -> ToolPart | None: 

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

320 

321 Args: 

322 part: Tool part to check for duplication 

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

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

325 

326 Returns: 

327 The duplicate part if found, None otherwise 

328 """ 

329 # Only compare against same tool type 

330 # Different tools are unlikely to produce identical outputs 

331 # This prevents false positives across tool types 

332 same_tool_parts = [ 

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

334 ] 

335 

336 if not same_tool_parts: 

337 return None 

338 

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

340 current_hash = compute_content_hash(part.state.output) 

341 

342 for recent_part, _ in same_tool_parts: 

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

344 return recent_part # Exact duplicate 

345 

346 # Generate shingles for current part 

347 current_shingles = generate_shingles(part.state.output) 

348 

349 # Check similarity against recent parts of same tool 

350 for recent_part, recent_shingles in same_tool_parts: 

351 similarity = jaccard_similarity(current_shingles, recent_shingles) 

352 

353 if similarity >= similarity_threshold: 

354 return recent_part # Found similar output 

355 

356 return None 

357 

358 

359def calculate_context_tokens( 

360 messages: list[Message], 

361 excluded_outputs: set[str] | None = None, 

362) -> int: 

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

364 

365 Args: 

366 messages: All conversation messages 

367 excluded_outputs: Call IDs whose outputs are excluded from the API payload 

368 

369 Returns: 

370 Total token count for outputs that will appear in the API payload 

371 """ 

372 excluded = excluded_outputs or set() 

373 total = 0 

374 for msg in messages: 

375 for part in msg.parts: 

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

377 if part.call_id in excluded: 

378 continue 

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

380 continue # Legacy in-place pruned output 

381 if part.state.output: 

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

383 return total 

384 

385 

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

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

388 

389 Args: 

390 messages: All conversation messages 

391 target_msg: Message to calculate age for 

392 

393 Returns: 

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

395 """ 

396 turn_count = 0 

397 for msg in reversed(messages): 

398 if msg.id == target_msg.id: 

399 return turn_count 

400 if msg.role == "user": 

401 turn_count += 1 

402 return turn_count 

403 

404 

405def score_tool_output( 

406 part: ToolPart, 

407 message: Message, 

408 messages: list[Message], 

409 config: PruningConfig, 

410 query: str | None = None, 

411) -> OutputImportance: 

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

413 

414 Args: 

415 part: Tool part to score 

416 message: Message containing this part 

417 messages: All conversation messages 

418 config: Pruning configuration with scoring weights 

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

420 

421 Returns: 

422 OutputImportance with all score components 

423 """ 

424 token_count = count_tokens_fast(part.state.output) 

425 

426 # 1. Recency score (exponential decay) 

427 age = calculate_turn_age(messages, message) 

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

429 

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

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

432 

433 # 3. Semantic importance 

434 semantic = 0.0 

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

436 

437 # Check for errors 

438 if any( 

439 keyword in output_lower 

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

441 ): 

442 semantic += config.error_boost 

443 

444 # Check for warnings 

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

446 semantic += config.warning_boost 

447 

448 # Check if user explicitly requested this 

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

450 semantic += config.user_requested_boost 

451 

452 semantic *= config.semantic_weight 

453 

454 # 4. Tool type priority 

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

456 

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

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

459 

460 return OutputImportance( 

461 recency_score=recency, 

462 size_penalty=size_penalty, 

463 semantic_score=semantic, 

464 tool_priority=tool_priority, 

465 token_count=token_count, 

466 query_relevance=query_relevance, 

467 ) 

468 

469 

470def prune_tool_outputs( 

471 messages: list[Message], 

472 config: PruningConfig, 

473 query: str | None = None, 

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

475 already_excluded: set[str] | None = None, 

476) -> PruningResult: 

477 """Prune tool outputs from conversation history. 

478 

479 Selectively removes old tool outputs while preserving: 

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

481 - Recent outputs (within protection window) 

482 - Protected tool outputs 

483 - Last N turns 

484 

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

486 

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

488 importance-based or FIFO pruning strategy. 

489 

490 Args: 

491 messages: Conversation messages (newest first recommended) 

492 config: Pruning configuration 

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

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

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

496 

497 Returns: 

498 PruningResult with detailed token tracking and breakdown 

499 """ 

500 excluded = already_excluded or set() 

501 

502 # Calculate token usage before pruning (excluding already-excluded outputs) 

503 tokens_before = calculate_context_tokens(messages, excluded_outputs=excluded) 

504 

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

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

507 if config.detect_duplicates: 

508 duplicate_result = _detect_and_prune_duplicates( 

509 messages, config, summarizer=summarizer, already_excluded=excluded 

510 ) 

511 excluded = excluded | set(duplicate_result.excluded_call_ids) 

512 

513 # Phase 2: Importance-based or FIFO pruning 

514 if config.use_importance_scoring: 

515 pruning_result = _prune_with_importance_scoring_only( 

516 messages, config, query=query, summarizer=summarizer, already_excluded=excluded 

517 ) 

518 else: 

519 pruning_result = _prune_simple(messages, config, summarizer=summarizer, already_excluded=excluded) 

520 

521 # Calculate token usage after pruning 

522 all_new_excluded = duplicate_result.excluded_call_ids + pruning_result.excluded_call_ids 

523 tokens_after = calculate_context_tokens( 

524 messages, excluded_outputs=excluded | set(pruning_result.excluded_call_ids) 

525 ) 

526 

527 # Combine results with detailed tracking 

528 all_decisions = duplicate_result.decisions + pruning_result.decisions 

529 

530 return PruningResult( 

531 pruned=duplicate_result.pruned + pruning_result.pruned, 

532 tokens_saved=duplicate_result.tokens_saved + pruning_result.tokens_saved, 

533 tokens_before=tokens_before, 

534 tokens_after=tokens_after, 

535 duplicates_pruned=duplicate_result.pruned, 

536 importance_pruned=pruning_result.pruned, 

537 duplicate_tokens_saved=duplicate_result.tokens_saved, 

538 importance_tokens_saved=pruning_result.tokens_saved, 

539 decisions=all_decisions, 

540 excluded_call_ids=all_new_excluded, 

541 ) 

542 

543 

544def _prune_simple( 

545 messages: list[Message], 

546 config: PruningConfig, 

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

548 already_excluded: set[str] | None = None, 

549) -> PruningResult: 

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

551 

552 Prunes oldest outputs first when token budget exceeded. 

553 """ 

554 excluded = already_excluded or set() 

555 total_tokens = 0 

556 prunable_tokens = 0 

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

558 decisions: list[PruningDecision] = [] 

559 turns_skipped = 0 

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

561 

562 for msg in reversed(messages): 

563 if msg.role == "user": 

564 turns_skipped += 1 

565 

566 # Check if past summary boundary 

567 past_summary = msg.summary 

568 

569 for part in msg.parts: 

570 if not isinstance(part, ToolPart): 

571 continue 

572 

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

574 continue 

575 

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

577 

578 # Track protection reasons 

579 if turns_skipped < config.protect_turns: 

580 decisions.append( 

581 PruningDecision( 

582 message_id=msg.id, 

583 part_id=part.call_id, 

584 tool=part.tool, 

585 decision="protected_recent", 

586 tokens=token_count, 

587 timestamp=now_ms, 

588 metadata={"turns_back": turns_skipped}, 

589 ) 

590 ) 

591 continue 

592 

593 if past_summary: 

594 decisions.append( 

595 PruningDecision( 

596 message_id=msg.id, 

597 part_id=part.call_id, 

598 tool=part.tool, 

599 decision="protected_summary", 

600 tokens=token_count, 

601 timestamp=now_ms, 

602 ) 

603 ) 

604 break 

605 

606 if part.tool in config.protected_tools: 

607 decisions.append( 

608 PruningDecision( 

609 message_id=msg.id, 

610 part_id=part.call_id, 

611 tool=part.tool, 

612 decision="protected_tool", 

613 tokens=token_count, 

614 timestamp=now_ms, 

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

616 ) 

617 ) 

618 continue 

619 

620 if part.call_id in excluded or (part.state.time and part.state.time.compacted): 

621 # Already excluded or legacy in-place pruned — skip 

622 continue 

623 

624 total_tokens += token_count 

625 

626 if total_tokens > config.prune_protect: 

627 prunable_tokens += token_count 

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

629 else: 

630 # Kept within budget 

631 decisions.append( 

632 PruningDecision( 

633 message_id=msg.id, 

634 part_id=part.call_id, 

635 tool=part.tool, 

636 decision="kept", 

637 tokens=token_count, 

638 timestamp=now_ms, 

639 metadata={"total_tokens_so_far": total_tokens}, 

640 ) 

641 ) 

642 

643 if prunable_tokens > config.prune_minimum: 

644 # Collect call IDs to exclude; optionally populate summary_text 

645 new_excluded: list[str] = [] 

646 for msg, part, tokens in to_prune: 

647 new_excluded.append(part.call_id) 

648 if summarizer and part.state.output and not part.state.summary_text: 

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

650 decisions.append( 

651 PruningDecision( 

652 message_id=msg.id, 

653 part_id=part.call_id, 

654 tool=part.tool, 

655 decision="pruned_fifo", 

656 tokens=tokens, 

657 timestamp=now_ms, 

658 metadata={ 

659 "prunable_tokens": prunable_tokens, 

660 "prune_minimum": config.prune_minimum, 

661 }, 

662 ) 

663 ) 

664 

665 return PruningResult( 

666 pruned=len(to_prune), 

667 tokens_saved=prunable_tokens, 

668 decisions=decisions, 

669 excluded_call_ids=new_excluded, 

670 ) 

671 

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

673 for msg, part, tokens in to_prune: 

674 decisions.append( 

675 PruningDecision( 

676 message_id=msg.id, 

677 part_id=part.call_id, 

678 tool=part.tool, 

679 decision="kept", 

680 tokens=tokens, 

681 timestamp=now_ms, 

682 metadata={ 

683 "reason": "below_minimum_threshold", 

684 "prunable_tokens": prunable_tokens, 

685 "prune_minimum": config.prune_minimum, 

686 }, 

687 ) 

688 ) 

689 

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

691 

692 

693def _prune_with_importance_scoring_only( 

694 messages: list[Message], 

695 config: PruningConfig, 

696 query: str | None = None, 

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

698 already_excluded: set[str] | None = None, 

699) -> PruningResult: 

700 """Smart pruning using importance scoring. 

701 

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

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

704 

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

706 """ 

707 excluded = already_excluded or set() 

708 # Collect all prunable outputs with scores 

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

710 decisions: list[PruningDecision] = [] 

711 turns_skipped = 0 

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

713 

714 for msg in reversed(messages): 

715 if msg.role == "user": 

716 turns_skipped += 1 

717 

718 past_summary = msg.summary 

719 

720 for part in msg.parts: 

721 if not isinstance(part, ToolPart): 

722 continue 

723 

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

725 continue 

726 

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

728 

729 # Track protection reasons (same as FIFO) 

730 if turns_skipped < config.protect_turns: 

731 decisions.append( 

732 PruningDecision( 

733 message_id=msg.id, 

734 part_id=part.call_id, 

735 tool=part.tool, 

736 decision="protected_recent", 

737 tokens=token_count, 

738 timestamp=now_ms, 

739 metadata={"turns_back": turns_skipped}, 

740 ) 

741 ) 

742 continue 

743 

744 if past_summary: 

745 decisions.append( 

746 PruningDecision( 

747 message_id=msg.id, 

748 part_id=part.call_id, 

749 tool=part.tool, 

750 decision="protected_summary", 

751 tokens=token_count, 

752 timestamp=now_ms, 

753 ) 

754 ) 

755 break 

756 

757 if part.tool in config.protected_tools: 

758 decisions.append( 

759 PruningDecision( 

760 message_id=msg.id, 

761 part_id=part.call_id, 

762 tool=part.tool, 

763 decision="protected_tool", 

764 tokens=token_count, 

765 timestamp=now_ms, 

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

767 ) 

768 ) 

769 continue 

770 

771 if part.call_id in excluded or (part.state.time and part.state.time.compacted): 

772 # Already excluded or legacy in-place pruned — skip 

773 continue 

774 

775 # Score this output 

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

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

778 

779 # Calculate total tokens 

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

781 

782 # If under budget, mark all as kept 

783 if total_tokens <= config.prune_protect: 

784 for msg, part, importance in scored_outputs: 

785 decisions.append( 

786 PruningDecision( 

787 message_id=msg.id, 

788 part_id=part.call_id, 

789 tool=part.tool, 

790 decision="kept", 

791 importance_score=importance.total_score, 

792 tokens=importance.token_count, 

793 timestamp=now_ms, 

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

795 ) 

796 ) 

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

798 

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

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

801 

802 # Prune until we're under budget 

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

804 current_tokens = total_tokens 

805 

806 for msg, part, importance in scored_outputs: 

807 if current_tokens <= config.prune_protect: 

808 # Keep this one 

809 decisions.append( 

810 PruningDecision( 

811 message_id=msg.id, 

812 part_id=part.call_id, 

813 tool=part.tool, 

814 decision="kept", 

815 importance_score=importance.total_score, 

816 tokens=importance.token_count, 

817 timestamp=now_ms, 

818 metadata={ 

819 "current_tokens": current_tokens, 

820 "prune_protect": config.prune_protect, 

821 }, 

822 ) 

823 ) 

824 else: 

825 # Mark for pruning 

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

827 current_tokens -= importance.token_count 

828 

829 # Only prune if savings meet minimum threshold 

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

831 

832 if tokens_saved >= config.prune_minimum: 

833 # Collect call IDs to exclude; optionally populate summary_text 

834 new_excluded: list[str] = [] 

835 for msg, part, importance in to_prune: 

836 new_excluded.append(part.call_id) 

837 if summarizer and part.state.output and not part.state.summary_text: 

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

839 decisions.append( 

840 PruningDecision( 

841 message_id=msg.id, 

842 part_id=part.call_id, 

843 tool=part.tool, 

844 decision="pruned_low_importance", 

845 importance_score=importance.total_score, 

846 tokens=importance.token_count, 

847 timestamp=now_ms, 

848 metadata={ 

849 "recency_score": importance.recency_score, 

850 "size_penalty": importance.size_penalty, 

851 "semantic_score": importance.semantic_score, 

852 "tool_priority": importance.tool_priority, 

853 "tokens_saved": tokens_saved, 

854 }, 

855 ) 

856 ) 

857 

858 return PruningResult( 

859 pruned=len(to_prune), 

860 tokens_saved=tokens_saved, 

861 decisions=decisions, 

862 excluded_call_ids=new_excluded, 

863 ) 

864 

865 # Not enough savings - keep all 

866 for msg, part, importance in to_prune: 

867 decisions.append( 

868 PruningDecision( 

869 message_id=msg.id, 

870 part_id=part.call_id, 

871 tool=part.tool, 

872 decision="kept", 

873 importance_score=importance.total_score, 

874 tokens=importance.token_count,