Coverage for src / harness_utils / compaction / pruning.py: 23%

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

2 

3Removes old tool outputs while preserving conversation structure. 

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

5""" 

6 

7from dataclasses import dataclass 

8 

9from harness_utils.config import PruningConfig 

10from harness_utils.models.message import Message 

11from harness_utils.models.parts import ToolPart 

12from harness_utils.tokens.estimator import estimate_tokens 

13 

14 

15@dataclass 

16class PruningResult: 

17 """Result of pruning operation.""" 

18 

19 pruned: int 

20 tokens_saved: int 

21 

22 

23def prune_tool_outputs( 

24 messages: list[Message], 

25 config: PruningConfig, 

26 chars_per_token: int = 4, 

27) -> PruningResult: 

28 """Prune tool outputs from conversation history. 

29 

30 Selectively removes old tool outputs while preserving: 

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

32 - Recent outputs (within protection window) 

33 - Protected tool outputs 

34 - Last N turns 

35 

36 Args: 

37 messages: Conversation messages (newest first recommended) 

38 config: Pruning configuration 

39 chars_per_token: Characters per token ratio for estimation 

40 

41 Returns: 

42 PruningResult with count and tokens saved 

43 """ 

44 total_tokens = 0 

45 prunable_tokens = 0 

46 to_prune: list[tuple[Message, ToolPart]] = [] 

47 turns_skipped = 0 

48 

49 for msg in reversed(messages): 

50 if msg.role == "user": 

51 turns_skipped += 1 

52 

53 if turns_skipped < config.protect_turns: 

54 continue 

55 

56 if msg.summary: 

57 break 

58 

59 for part in msg.parts: 

60 if not isinstance(part, ToolPart): 

61 continue 

62 

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

64 continue 

65 

66 if part.tool in config.protected_tools: 

67 continue 

68 

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

70 continue 

71 

72 token_estimate = estimate_tokens(part.state.output, chars_per_token) 

73 total_tokens += token_estimate 

74 

75 if total_tokens > config.prune_protect: 

76 prunable_tokens += token_estimate 

77 to_prune.append((msg, part)) 

78 

79 if prunable_tokens > config.prune_minimum: 

80 for msg, part in to_prune: 

81 part.state.output = "" 

82 part.state.attachments = [] 

83 if part.state.time: 

84 import time 

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

86 

87 return PruningResult(pruned=len(to_prune), tokens_saved=prunable_tokens) 

88 

89 return PruningResult(pruned=0, tokens_saved=0)