Skip to main content

devboy_format_pipeline/
budget.rs

1//! Iterative budget pipeline: encode → check → trim → re-encode → verify.
2//!
3//! The pipeline ensures output fits within a token budget by iteratively:
4//! 1. TOON-encode the full data
5//! 2. If fits budget → return as-is
6//! 3. Calculate B_trim = budget / compression_ratio × (1 - margin)
7//! 4. Build TrimTree → apply strategy → trim to B_trim → re-encode
8//! 5. Repeat up to max_iterations, adjusting B_trim each time
9//! 6. If still over budget → hard truncate
10
11use crate::strategy::{TrimStrategyKind, create_strategy};
12use crate::token_counter::estimate_tokens;
13use crate::toon::{self, TrimLevel};
14use crate::tree::TrimNode;
15use crate::trim;
16use crate::truncation;
17
18#[derive(Debug, Clone)]
19pub struct BudgetConfig {
20    /// Maximum token budget for output (default: 8000).
21    pub budget_tokens: usize,
22    /// Safety margin for token estimation inaccuracy (default: 0.20).
23    pub margin: f64,
24    /// Maximum trim-encode-verify iterations (default: 3).
25    pub max_iterations: usize,
26}
27
28impl Default for BudgetConfig {
29    fn default() -> Self {
30        Self {
31            budget_tokens: 8000,
32            margin: 0.20,
33            max_iterations: 3,
34        }
35    }
36}
37
38/// Result of budget pipeline processing.
39#[derive(Debug, Clone)]
40pub struct BudgetResult {
41    /// TOON-encoded content that fits within budget.
42    pub content: String,
43    /// Estimated token count of the content.
44    pub tokens: usize,
45    /// Whether trimming was applied.
46    pub trimmed: bool,
47    /// Indices of items that were excluded (overflow).
48    pub overflow_indices: Vec<usize>,
49    /// Total number of items before trimming.
50    pub total_items: usize,
51    /// Number of items included after trimming.
52    pub included_items: usize,
53}
54
55/// Run the iterative budget pipeline on a TrimTree.
56///
57/// Arguments:
58/// - `tree` — mutable TrimTree with pre-assigned values from strategy
59/// - `full_encoded` — TOON encoding of the complete data (before trimming)
60/// - `config` — budget configuration
61///
62/// Returns BudgetResult with content fitting within budget.
63pub fn run(tree: &mut TrimNode, full_encoded: &str, config: &BudgetConfig) -> BudgetResult {
64    let full_tokens = estimate_tokens(full_encoded);
65    let total_items = tree.included_items_count();
66
67    // Fast path: fits within budget
68    if full_tokens <= config.budget_tokens {
69        return BudgetResult {
70            content: full_encoded.to_string(),
71            tokens: full_tokens,
72            trimmed: false,
73            overflow_indices: vec![],
74            total_items,
75            included_items: total_items,
76        };
77    }
78
79    // Calculate initial compression ratio
80    let r = if full_tokens > 0 {
81        full_tokens as f64 / tree.total_weight() as f64
82    } else {
83        1.0
84    };
85
86    // Calculate initial trim budget: B_trim = budget / r × (1 - margin)
87    let effective_budget = config.budget_tokens as f64 * (1.0 - config.margin);
88    let mut b_trim = if r > 0.0 {
89        (effective_budget / r) as usize
90    } else {
91        config.budget_tokens
92    };
93
94    // Iterative trim-encode-verify loop
95    for _iteration in 0..config.max_iterations {
96        // Trim tree to b_trim
97        trim::trim(tree, b_trim);
98
99        let included_items = tree.included_items_count();
100        let overflow = tree.excluded_item_indices();
101        let current_weight = tree.total_weight();
102
103        // Estimate tokens after trim (proportional to weight reduction)
104        let estimated_tokens = if tree.total_weight() > 0 {
105            (current_weight as f64 * r) as usize
106        } else {
107            0
108        };
109
110        if estimated_tokens <= config.budget_tokens {
111            // Likely fits — we'll verify with actual encode in the caller
112            return BudgetResult {
113                content: String::new(), // Caller will re-encode from included indices
114                tokens: estimated_tokens,
115                trimmed: true,
116                overflow_indices: overflow,
117                total_items,
118                included_items,
119            };
120        }
121
122        // Adjust b_trim for next iteration
123        if estimated_tokens > 0 {
124            b_trim =
125                (b_trim as f64 * config.budget_tokens as f64 / estimated_tokens as f64) as usize;
126        } else {
127            break;
128        }
129    }
130
131    // Final fallback: hard truncate
132    let overflow = tree.excluded_item_indices();
133    let included_items = tree.included_items_count();
134
135    BudgetResult {
136        content: String::new(),
137        tokens: 0,
138        trimmed: true,
139        overflow_indices: overflow,
140        total_items,
141        included_items,
142    }
143}
144
145/// High-level: apply strategy and run budget pipeline on issues.
146pub fn process_issues(
147    issues: &[devboy_core::Issue],
148    strategy_kind: TrimStrategyKind,
149    config: &BudgetConfig,
150) -> devboy_core::Result<BudgetResult> {
151    // Encode full data
152    let full_encoded = toon::encode_issues(issues, TrimLevel::Full)?;
153    let full_tokens = estimate_tokens(&full_encoded);
154
155    // Fast path
156    if full_tokens <= config.budget_tokens {
157        return Ok(BudgetResult {
158            content: full_encoded,
159            tokens: full_tokens,
160            trimmed: false,
161            overflow_indices: vec![],
162            total_items: issues.len(),
163            included_items: issues.len(),
164        });
165    }
166
167    // Build tree and apply strategy
168    let mut tree = crate::tree::build_issues_tree(issues);
169    let strategy = create_strategy(strategy_kind);
170    strategy.assign_values(&mut tree);
171
172    // Run budget pipeline
173    let mut result = run(&mut tree, &full_encoded, config);
174
175    // Re-encode with only included items
176    if result.trimmed {
177        let included_indices = tree.included_item_indices();
178        let included_issues: Vec<devboy_core::Issue> = included_indices
179            .iter()
180            .map(|&i| issues[i].clone())
181            .collect();
182
183        // Try progressively smaller trim levels until it fits
184        for level in [TrimLevel::Full, TrimLevel::Standard, TrimLevel::Minimal] {
185            let encoded = toon::encode_issues(&included_issues, level)?;
186            let tokens = estimate_tokens(&encoded);
187            if tokens <= config.budget_tokens {
188                result.content = encoded;
189                result.tokens = tokens;
190                return Ok(result);
191            }
192        }
193
194        // Final hard truncate of minimal encoding
195        let encoded = toon::encode_issues(&included_issues, TrimLevel::Minimal)?;
196        let max_chars = crate::token_counter::tokens_to_chars(config.budget_tokens);
197        result.content = truncation::truncate_string(&encoded, max_chars);
198        result.tokens = estimate_tokens(&result.content);
199    }
200
201    Ok(result)
202}
203
204/// High-level: apply strategy and run budget pipeline on merge requests.
205pub fn process_merge_requests(
206    mrs: &[devboy_core::MergeRequest],
207    strategy_kind: TrimStrategyKind,
208    config: &BudgetConfig,
209) -> devboy_core::Result<BudgetResult> {
210    let full_encoded = toon::encode_merge_requests(mrs, TrimLevel::Full)?;
211    let full_tokens = estimate_tokens(&full_encoded);
212
213    if full_tokens <= config.budget_tokens {
214        return Ok(BudgetResult {
215            content: full_encoded,
216            tokens: full_tokens,
217            trimmed: false,
218            overflow_indices: vec![],
219            total_items: mrs.len(),
220            included_items: mrs.len(),
221        });
222    }
223
224    let mut tree = crate::tree::build_merge_requests_tree(mrs);
225    let strategy = create_strategy(strategy_kind);
226    strategy.assign_values(&mut tree);
227
228    let mut result = run(&mut tree, &full_encoded, config);
229
230    if result.trimmed {
231        let included_indices = tree.included_item_indices();
232        let included_mrs: Vec<devboy_core::MergeRequest> =
233            included_indices.iter().map(|&i| mrs[i].clone()).collect();
234
235        for level in [TrimLevel::Full, TrimLevel::Standard, TrimLevel::Minimal] {
236            let encoded = toon::encode_merge_requests(&included_mrs, level)?;
237            let tokens = estimate_tokens(&encoded);
238            if tokens <= config.budget_tokens {
239                result.content = encoded;
240                result.tokens = tokens;
241                return Ok(result);
242            }
243        }
244
245        let encoded = toon::encode_merge_requests(&included_mrs, TrimLevel::Minimal)?;
246        let max_chars = crate::token_counter::tokens_to_chars(config.budget_tokens);
247        result.content = truncation::truncate_string(&encoded, max_chars);
248        result.tokens = estimate_tokens(&result.content);
249    }
250
251    Ok(result)
252}
253
254/// High-level: apply strategy and run budget pipeline on diffs.
255pub fn process_diffs(
256    diffs: &[devboy_core::FileDiff],
257    strategy_kind: TrimStrategyKind,
258    config: &BudgetConfig,
259) -> devboy_core::Result<BudgetResult> {
260    let full_encoded = toon::encode_diffs(diffs)?;
261    let full_tokens = estimate_tokens(&full_encoded);
262
263    if full_tokens <= config.budget_tokens {
264        return Ok(BudgetResult {
265            content: full_encoded,
266            tokens: full_tokens,
267            trimmed: false,
268            overflow_indices: vec![],
269            total_items: diffs.len(),
270            included_items: diffs.len(),
271        });
272    }
273
274    let mut tree = crate::tree::build_diffs_tree(diffs);
275    let strategy = create_strategy(strategy_kind);
276    strategy.assign_values(&mut tree);
277
278    let mut result = run(&mut tree, &full_encoded, config);
279
280    if result.trimmed {
281        let included_indices = tree.included_item_indices();
282        let included_diffs: Vec<devboy_core::FileDiff> =
283            included_indices.iter().map(|&i| diffs[i].clone()).collect();
284
285        let encoded = toon::encode_diffs(&included_diffs)?;
286        let tokens = estimate_tokens(&encoded);
287        if tokens <= config.budget_tokens {
288            result.content = encoded;
289            result.tokens = tokens;
290        } else {
291            let max_chars = crate::token_counter::tokens_to_chars(config.budget_tokens);
292            result.content = truncation::truncate_string(&encoded, max_chars);
293            result.tokens = estimate_tokens(&result.content);
294        }
295    }
296
297    Ok(result)
298}
299
300/// High-level: apply strategy and run budget pipeline on discussions.
301pub fn process_discussions(
302    discussions: &[devboy_core::Discussion],
303    strategy_kind: TrimStrategyKind,
304    config: &BudgetConfig,
305) -> devboy_core::Result<BudgetResult> {
306    let full_encoded = toon::encode_discussions(discussions)?;
307    let full_tokens = estimate_tokens(&full_encoded);
308
309    if full_tokens <= config.budget_tokens {
310        return Ok(BudgetResult {
311            content: full_encoded,
312            tokens: full_tokens,
313            trimmed: false,
314            overflow_indices: vec![],
315            total_items: discussions.len(),
316            included_items: discussions.len(),
317        });
318    }
319
320    let mut tree = crate::tree::build_discussions_tree(discussions);
321    let strategy = create_strategy(strategy_kind);
322    strategy.assign_values(&mut tree);
323
324    let mut result = run(&mut tree, &full_encoded, config);
325
326    if result.trimmed {
327        let included_indices = tree.included_item_indices();
328        let included_discussions: Vec<devboy_core::Discussion> = included_indices
329            .iter()
330            .map(|&i| discussions[i].clone())
331            .collect();
332
333        let encoded = toon::encode_discussions(&included_discussions)?;
334        let tokens = estimate_tokens(&encoded);
335        if tokens <= config.budget_tokens {
336            result.content = encoded;
337            result.tokens = tokens;
338        } else {
339            let max_chars = crate::token_counter::tokens_to_chars(config.budget_tokens);
340            result.content = truncation::truncate_string(&encoded, max_chars);
341            result.tokens = estimate_tokens(&result.content);
342        }
343    }
344
345    Ok(result)
346}
347
348/// High-level: apply strategy and run budget pipeline on comments.
349pub fn process_comments(
350    comments: &[devboy_core::Comment],
351    strategy_kind: TrimStrategyKind,
352    config: &BudgetConfig,
353) -> devboy_core::Result<BudgetResult> {
354    let full_encoded = toon::encode_comments(comments)?;
355    let full_tokens = estimate_tokens(&full_encoded);
356
357    if full_tokens <= config.budget_tokens {
358        return Ok(BudgetResult {
359            content: full_encoded,
360            tokens: full_tokens,
361            trimmed: false,
362            overflow_indices: vec![],
363            total_items: comments.len(),
364            included_items: comments.len(),
365        });
366    }
367
368    let mut tree = crate::tree::build_comments_tree(comments);
369    let strategy = create_strategy(strategy_kind);
370    strategy.assign_values(&mut tree);
371
372    let mut result = run(&mut tree, &full_encoded, config);
373
374    if result.trimmed {
375        let included_indices = tree.included_item_indices();
376        let included_comments: Vec<devboy_core::Comment> = included_indices
377            .iter()
378            .map(|&i| comments[i].clone())
379            .collect();
380
381        let encoded = toon::encode_comments(&included_comments)?;
382        let tokens = estimate_tokens(&encoded);
383        if tokens <= config.budget_tokens {
384            result.content = encoded;
385            result.tokens = tokens;
386        } else {
387            let max_chars = crate::token_counter::tokens_to_chars(config.budget_tokens);
388            result.content = truncation::truncate_string(&encoded, max_chars);
389            result.tokens = estimate_tokens(&result.content);
390        }
391    }
392
393    Ok(result)
394}
395
396#[cfg(test)]
397mod tests {
398    use super::*;
399    use crate::tree::NodeKind;
400
401    fn make_tree(weights: &[usize], values: &[f64]) -> TrimNode {
402        let mut root = TrimNode::new(0, NodeKind::Root, 0);
403        for (i, (&w, &v)) in weights.iter().zip(values.iter()).enumerate() {
404            let mut node = TrimNode::new(i + 1, NodeKind::Item { index: i }, w);
405            node.value = v;
406            root.children.push(node);
407        }
408        root
409    }
410
411    #[test]
412    fn test_budget_fast_path() {
413        let mut tree = make_tree(&[10, 10, 10], &[1.0, 1.0, 1.0]);
414        let content = "short content";
415        let config = BudgetConfig {
416            budget_tokens: 100,
417            ..Default::default()
418        };
419
420        let result = run(&mut tree, content, &config);
421        assert!(!result.trimmed);
422        assert_eq!(result.content, "short content");
423        assert!(result.overflow_indices.is_empty());
424    }
425
426    #[test]
427    fn test_budget_triggers_trimming() {
428        let mut tree = make_tree(&[100, 100, 100], &[1.0, 0.5, 0.2]);
429        // Create content that's over budget
430        let content = "x".repeat(1000); // ~286 tokens
431        let config = BudgetConfig {
432            budget_tokens: 50,
433            margin: 0.20,
434            max_iterations: 3,
435        };
436
437        let result = run(&mut tree, &content, &config);
438        assert!(result.trimmed);
439        assert!(result.included_items < 3);
440        assert!(!result.overflow_indices.is_empty());
441    }
442
443    #[test]
444    fn test_budget_config_defaults() {
445        let config = BudgetConfig::default();
446        assert_eq!(config.budget_tokens, 8000);
447        assert!((config.margin - 0.20).abs() < 0.001);
448        assert_eq!(config.max_iterations, 3);
449    }
450
451    #[test]
452    fn test_budget_result_fields() {
453        let mut tree = make_tree(&[50, 50, 50], &[1.0, 0.8, 0.3]);
454        let content = "x".repeat(600); // ~171 tokens
455        let config = BudgetConfig {
456            budget_tokens: 30,
457            margin: 0.10,
458            max_iterations: 2,
459        };
460
461        let result = run(&mut tree, &content, &config);
462        assert_eq!(result.total_items, 3);
463        assert!(result.included_items <= 3);
464        assert_eq!(
465            result.included_items + result.overflow_indices.len(),
466            result.total_items
467        );
468    }
469
470    // --- process_issues tests ---
471
472    fn sample_issues(n: usize) -> Vec<devboy_core::Issue> {
473        (0..n)
474            .map(|i| devboy_core::Issue {
475                key: format!("gh#{}", i + 1),
476                title: format!("Issue {}", i + 1),
477                description: Some(format!(
478                    "Description for issue {} with enough text to make it non-trivial for token counting purposes",
479                    i + 1
480                )),
481                state: "open".into(),
482                source: "github".into(),
483                priority: None,
484                labels: vec!["bug".into()],
485                author: Some(devboy_core::User {
486                    id: format!("{}", i),
487                    username: format!("user{}", i),
488                    name: None,
489                    email: None,
490                    avatar_url: None,
491                }),
492                assignees: vec![],
493                url: Some(format!("https://github.com/test/repo/issues/{}", i + 1)),
494                created_at: Some("2024-01-01T00:00:00Z".into()),
495                updated_at: Some("2024-01-02T00:00:00Z".into()),
496                attachments_count: None,
497            parent: None,
498                subtasks: vec![],
499            })
500            .collect()
501    }
502
503    #[test]
504    fn test_process_issues_fast_path() {
505        let issues = sample_issues(2);
506        let config = BudgetConfig {
507            budget_tokens: 50000,
508            ..Default::default()
509        };
510        let result = process_issues(&issues, TrimStrategyKind::ElementCount, &config).unwrap();
511        assert!(!result.trimmed);
512        assert!(!result.content.is_empty());
513        assert_eq!(result.total_items, 2);
514        assert_eq!(result.included_items, 2);
515    }
516
517    #[test]
518    fn test_process_issues_with_trimming() {
519        let issues = sample_issues(20);
520        let config = BudgetConfig {
521            budget_tokens: 500,
522            margin: 0.20,
523            max_iterations: 3,
524        };
525        let result = process_issues(&issues, TrimStrategyKind::ElementCount, &config).unwrap();
526        assert!(result.trimmed);
527        assert!(result.included_items < 20);
528        assert!(!result.content.is_empty());
529        assert!(result.tokens <= config.budget_tokens);
530    }
531
532    #[test]
533    fn test_process_issues_with_cascading_strategy() {
534        let issues = sample_issues(10);
535        let config = BudgetConfig {
536            budget_tokens: 300,
537            margin: 0.20,
538            max_iterations: 3,
539        };
540        let result = process_issues(&issues, TrimStrategyKind::Cascading, &config).unwrap();
541        assert!(result.trimmed);
542        assert!(!result.content.is_empty());
543    }
544
545    #[test]
546    fn test_process_issues_very_small_budget() {
547        let issues = sample_issues(10);
548        let config = BudgetConfig {
549            budget_tokens: 50,
550            margin: 0.20,
551            max_iterations: 3,
552        };
553        let result = process_issues(&issues, TrimStrategyKind::Default, &config).unwrap();
554        assert!(result.trimmed);
555        // Content should be truncated but non-empty
556        assert!(!result.content.is_empty());
557    }
558
559    #[test]
560    fn test_process_issues_empty() {
561        let config = BudgetConfig::default();
562        let result = process_issues(&[], TrimStrategyKind::Default, &config).unwrap();
563        assert!(!result.trimmed);
564        assert_eq!(result.total_items, 0);
565    }
566
567    // --- process_merge_requests tests ---
568
569    fn sample_merge_requests(n: usize) -> Vec<devboy_core::MergeRequest> {
570        (0..n)
571            .map(|i| devboy_core::MergeRequest {
572                key: format!("mr#{}", i + 1),
573                title: format!("Merge Request {}", i + 1),
574                description: Some(format!(
575                    "Description for MR {} with enough text to make it non-trivial for token counting purposes",
576                    i + 1
577                )),
578                state: "opened".into(),
579                source: "gitlab".into(),
580                source_branch: format!("feature-{}", i + 1),
581                target_branch: "main".into(),
582                author: Some(devboy_core::User {
583                    id: format!("{}", i),
584                    username: format!("user{}", i),
585                    name: None,
586                    email: None,
587                    avatar_url: None,
588                }),
589                assignees: vec![],
590                reviewers: vec![],
591                labels: vec!["enhancement".into()],
592                draft: false,
593                url: Some(format!("https://gitlab.com/test/repo/-/merge_requests/{}", i + 1)),
594                created_at: Some("2024-01-01T00:00:00Z".into()),
595                updated_at: Some("2024-01-02T00:00:00Z".into()),
596            })
597            .collect()
598    }
599
600    #[test]
601    fn test_process_merge_requests_fast_path() {
602        let mrs = sample_merge_requests(2);
603        let config = BudgetConfig {
604            budget_tokens: 50000,
605            ..Default::default()
606        };
607        let result = process_merge_requests(&mrs, TrimStrategyKind::ElementCount, &config).unwrap();
608        assert!(!result.trimmed);
609        assert!(!result.content.is_empty());
610        assert_eq!(result.total_items, 2);
611        assert_eq!(result.included_items, 2);
612    }
613
614    #[test]
615    fn test_process_merge_requests_with_trimming() {
616        let mrs = sample_merge_requests(20);
617        let config = BudgetConfig {
618            budget_tokens: 500,
619            margin: 0.20,
620            max_iterations: 3,
621        };
622        let result = process_merge_requests(&mrs, TrimStrategyKind::ElementCount, &config).unwrap();
623        assert!(result.trimmed);
624        assert!(result.included_items < 20);
625        assert!(!result.content.is_empty());
626        assert!(result.tokens <= config.budget_tokens);
627    }
628
629    #[test]
630    fn test_process_merge_requests_empty() {
631        let config = BudgetConfig::default();
632        let result = process_merge_requests(&[], TrimStrategyKind::Default, &config).unwrap();
633        assert!(!result.trimmed);
634        assert_eq!(result.total_items, 0);
635    }
636
637    // --- process_diffs tests ---
638
639    fn sample_diffs(n: usize) -> Vec<devboy_core::FileDiff> {
640        (0..n)
641            .map(|i| devboy_core::FileDiff {
642                file_path: format!("src/module_{}/file_{}.rs", i / 3, i + 1),
643                old_path: None,
644                new_file: i == 0,
645                deleted_file: false,
646                renamed_file: false,
647                diff: format!(
648                    "@@ -1,10 +1,15 @@\n-old line {i}\n+new line {i}\n+added context for file {i} with enough diff content to be meaningful",
649                    i = i + 1
650                ),
651                additions: Some(2),
652                deletions: Some(1),
653            })
654            .collect()
655    }
656
657    #[test]
658    fn test_process_diffs_fast_path() {
659        let diffs = sample_diffs(2);
660        let config = BudgetConfig {
661            budget_tokens: 50000,
662            ..Default::default()
663        };
664        let result = process_diffs(&diffs, TrimStrategyKind::ElementCount, &config).unwrap();
665        assert!(!result.trimmed);
666        assert!(!result.content.is_empty());
667        assert_eq!(result.total_items, 2);
668        assert_eq!(result.included_items, 2);
669    }
670
671    #[test]
672    fn test_process_diffs_with_trimming() {
673        let diffs = sample_diffs(20);
674        let config = BudgetConfig {
675            budget_tokens: 200,
676            margin: 0.20,
677            max_iterations: 3,
678        };
679        let result = process_diffs(&diffs, TrimStrategyKind::ElementCount, &config).unwrap();
680        assert!(result.trimmed);
681        assert!(result.included_items < 20);
682        assert!(!result.content.is_empty());
683    }
684
685    #[test]
686    fn test_process_diffs_empty() {
687        let config = BudgetConfig::default();
688        let result = process_diffs(&[], TrimStrategyKind::Default, &config).unwrap();
689        assert!(!result.trimmed);
690        assert_eq!(result.total_items, 0);
691    }
692
693    // --- process_discussions tests ---
694
695    fn sample_discussions(n: usize) -> Vec<devboy_core::Discussion> {
696        (0..n)
697            .map(|i| devboy_core::Discussion {
698                id: format!("disc-{}", i + 1),
699                resolved: i % 3 == 0,
700                resolved_by: None,
701                comments: vec![devboy_core::Comment {
702                    id: format!("comment-{}", i + 1),
703                    body: format!(
704                        "Discussion comment {} with enough body text for token counting purposes in the budget pipeline",
705                        i + 1
706                    ),
707                    author: Some(devboy_core::User {
708                        id: format!("{}", i),
709                        username: format!("reviewer{}", i),
710                        name: None,
711                        email: None,
712                        avatar_url: None,
713                    }),
714                    created_at: Some("2024-01-01T00:00:00Z".into()),
715                    updated_at: None,
716                    position: None,
717                }],
718                position: None,
719            })
720            .collect()
721    }
722
723    #[test]
724    fn test_process_discussions_fast_path() {
725        let discussions = sample_discussions(2);
726        let config = BudgetConfig {
727            budget_tokens: 50000,
728            ..Default::default()
729        };
730        let result =
731            process_discussions(&discussions, TrimStrategyKind::ElementCount, &config).unwrap();
732        assert!(!result.trimmed);
733        assert!(!result.content.is_empty());
734        assert_eq!(result.total_items, 2);
735        assert_eq!(result.included_items, 2);
736    }
737
738    #[test]
739    fn test_process_discussions_with_trimming() {
740        let discussions = sample_discussions(20);
741        let config = BudgetConfig {
742            budget_tokens: 300,
743            margin: 0.20,
744            max_iterations: 3,
745        };
746        let result =
747            process_discussions(&discussions, TrimStrategyKind::ElementCount, &config).unwrap();
748        assert!(result.trimmed);
749        assert!(result.included_items < 20);
750        assert!(!result.content.is_empty());
751    }
752
753    #[test]
754    fn test_process_discussions_empty() {
755        let config = BudgetConfig::default();
756        let result = process_discussions(&[], TrimStrategyKind::Default, &config).unwrap();
757        assert!(!result.trimmed);
758        assert_eq!(result.total_items, 0);
759    }
760
761    // --- process_comments tests ---
762
763    fn sample_comments(n: usize) -> Vec<devboy_core::Comment> {
764        (0..n)
765            .map(|i| devboy_core::Comment {
766                id: format!("c-{}", i + 1),
767                body: format!(
768                    "Comment {} with enough body text to make it non-trivial for budget pipeline token counting",
769                    i + 1
770                ),
771                author: Some(devboy_core::User {
772                    id: format!("{}", i),
773                    username: format!("commenter{}", i),
774                    name: None,
775                    email: None,
776                    avatar_url: None,
777                }),
778                created_at: Some("2024-01-01T00:00:00Z".into()),
779                updated_at: Some("2024-01-02T00:00:00Z".into()),
780                position: None,
781            })
782            .collect()
783    }
784
785    #[test]
786    fn test_process_comments_fast_path() {
787        let comments = sample_comments(2);
788        let config = BudgetConfig {
789            budget_tokens: 50000,
790            ..Default::default()
791        };
792        let result = process_comments(&comments, TrimStrategyKind::ElementCount, &config).unwrap();
793        assert!(!result.trimmed);
794        assert!(!result.content.is_empty());
795        assert_eq!(result.total_items, 2);
796        assert_eq!(result.included_items, 2);
797    }
798
799    #[test]
800    fn test_process_comments_with_trimming() {
801        let comments = sample_comments(20);
802        let config = BudgetConfig {
803            budget_tokens: 300,
804            margin: 0.20,
805            max_iterations: 3,
806        };
807        let result = process_comments(&comments, TrimStrategyKind::ElementCount, &config).unwrap();
808        assert!(result.trimmed);
809        assert!(result.included_items < 20);
810        assert!(!result.content.is_empty());
811    }
812
813    #[test]
814    fn test_process_comments_empty() {
815        let config = BudgetConfig::default();
816        let result = process_comments(&[], TrimStrategyKind::Default, &config).unwrap();
817        assert!(!result.trimmed);
818        assert_eq!(result.total_items, 0);
819    }
820}