fix(compress): enforce truncation gates and preserve whitespace

internal/compress/algorithm.go

@@ -81,10 +81,13 @@ func Compress(s string, intensity Intensity) (string, Stats) {
 	var out strings.Builder
 	dropSet := dropLists[intensity]
 	for _, seg := range segs {
+		if seg.Raw == "" {
+			continue
+		}
 		if !seg.Safe {
 			stats.PassThroughSegments++
 			stats.SensitiveKeywordsHit = append(stats.SensitiveKeywordsHit, detectSensitive(seg.Text))
-			out.WriteString(seg.Text)
+			out.WriteString(seg.Raw)
 			continue
 		}
 		stats.CompressedSegments++
@@ -101,7 +104,11 @@ func Compress(s string, intensity Intensity) (string, Stats) {
 		postDrops := applyDropList(postDict, dropSet, &stats)
 		stats.BytesSavedByDrops += len(postDict) - len(postDrops)
 
+		// Re-attach the segment's original leading and trailing whitespace so
+		// the re-joined output preserves the input's sentence separation.
+		out.WriteString(segmentLeadingWS(seg.Raw))
 		out.WriteString(postDrops)
+		out.WriteString(segmentTrailingWS(seg.Raw))
 	}
 
 	compressed := out.String()
@@ -140,6 +147,35 @@ func applyDropList(s string, dropSet map[string]bool, stats *Stats) string {
 	return normalizeWhitespace(out)
 }
 
+// segmentLeadingWS returns the leading whitespace of s (spaces, tabs,
+// newlines). An empty string is returned when s starts with non-space.
+func segmentLeadingWS(s string) string {
+	i := 0
+	for i < len(s) {
+		switch s[i] {
+		case ' ', '\t', '\n', '\r':
+			i++
+		default:
+			return s[:i]
+		}
+	}
+	return s
+}
+
+// segmentTrailingWS returns the trailing whitespace of s.
+func segmentTrailingWS(s string) string {
+	i := len(s)
+	for i > 0 {
+		switch s[i-1] {
+		case ' ', '\t', '\n', '\r':
+			i--
+		default:
+			return s[i:]
+		}
+	}
+	return s
+}
+
 func classifyDrop(phrase string, stats *Stats) {
 	lower := strings.ToLower(phrase)
 	switch {

internal/compress/safety.go

@@ -112,6 +112,10 @@ func IsSafeSegment(s string) bool {
 // SplitSafeSegments splits s into segments that are safe to compress and
 // segments that must be preserved verbatim. Returns a slice of (text, safe)
 // pairs. Adjacent safe/unsafe segments are merged by their flag.
+//
+// Each segment keeps its original, untrimmed text (Raw) so re-joins preserve
+// the whitespace between sentences; Text holds the trimmed form used for
+// classification and compression.
 func SplitSafeSegments(s string) []segment {
 	if s == "" {
 		return nil
@@ -120,18 +124,22 @@ func SplitSafeSegments(s string) []segment {
 	segs := splitSentences(s)
 	out := make([]segment, 0, len(segs))
 	for _, seg := range segs {
-		seg = strings.TrimSpace(seg)
-		if seg == "" {
+		trimmed := strings.TrimSpace(seg)
+		if trimmed == "" {
+			// Whitespace-only segment (e.g. blank lines between sentences):
+			// keep the raw text and pass it through untouched.
+			out = append(out, segment{Text: "", Raw: seg, Safe: false})
 			continue
 		}
-		out = append(out, segment{Text: seg, Safe: IsSafeSegment(seg)})
+		out = append(out, segment{Text: trimmed, Raw: seg, Safe: IsSafeSegment(trimmed)})
 	}
 	return out
 }
 
 type segment struct {
-	Text string
-	Safe bool // true if safe to compress; false = preserve verbatim
+	Text string // trimmed text used for classification and compression
+	Raw  string // original (untrimmed) text, preserved for the re-join
+	Safe bool   // true if safe to compress; false = preserve verbatim
 }
 
 // splitSentences splits s on sentence-ending punctuation (`.`, `!`, `?`,

internal/compress/whitespace_regression_test.go

@@ -0,0 +1,65 @@
+package compress_test
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/GrayCodeAI/tok/internal/compress"
+)
+
+// Regression tests for the whitespace-preservation fix: compressed output must
+// keep the sentence separation (newlines/spaces) of the input instead of
+// concatenating trimmed sentences.
+
+func TestCompressPreservesSentenceWhitespace(t *testing.T) {
+	cases := []struct {
+		name string
+		in   string
+		want string // substring that must survive from both sentences
+	}{
+		{"single newline", "First sentence is here.\nSecond sentence is here.", "here."},
+		{"double newline paragraph break", "Paragraph one ends with facts.\n\nParagraph two begins with facts.", "facts."},
+		{"trailing newline", "Sentence with facts here.\n\n", "facts"},
+		{"leading newline", "\n\nSentence with facts here.", "facts"},
+		{"space separated", "First sentence has facts. Second sentence has facts.", "facts"},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			out, _ := compress.Compress(tc.in, compress.Full)
+			if !strings.Contains(out, "\n\n") && strings.Contains(tc.in, "\n\n") {
+				t.Errorf("paragraph break lost: input %q -> output %q", tc.in, out)
+			}
+			if !strings.Contains(out, "\n") && strings.Contains(tc.in, "\n") {
+				t.Errorf("newline lost: input %q -> output %q", tc.in, out)
+			}
+			// Counts may shift, but the separator must not vanish entirely.
+			if !strings.Contains(out, tc.want) {
+				t.Errorf("content missing after compression: %q", out)
+			}
+		})
+	}
+}
+
+func TestCompressPreservesBlankLinesInsidePassThroughSegments(t *testing.T) {
+	in := "SECRET_PASSWORD=abc\n\n\nKeep this line too, it has facts."
+	out, _ := compress.Compress(in, compress.Full)
+	// The sensitive line is passed through verbatim, including its trailing
+	// blank lines, and the next line keeps its own separation.
+	if !strings.Contains(out, "Keep this line too") {
+		t.Errorf("content lost: %q", out)
+	}
+	if !strings.Contains(out, "SECRET_PASSWORD=abc\n\n") {
+		t.Errorf("pass-through whitespace not preserved: %q", out)
+	}
+}
+
+func TestCompressMultipleSentencesKeepSpacesBetween(t *testing.T) {
+	in := "Alpha has useful facts. Beta has useful facts."
+	out, _ := compress.Compress(in, compress.Full)
+	// Even when articles are dropped inside sentences, a space must separate
+	// the two sentences (no "facts.Beta" concatenation).
+	if strings.Contains(out, "facts.Beta") || strings.Contains(out, "sentence.sentence") {
+		t.Errorf("sentences concatenated without separator: %q", out)
+	}
+}

internal/filter/budget.go

@@ -101,27 +101,71 @@ func (f *BudgetEnforcer) enforceBudget(input string, mode Mode) string {
 }
 
 // TruncateToBudget performs word-level truncation to fit within budget.
+// It never returns an empty string for non-empty input: when even a single
+// word exceeds the budget, a truncated fragment of the first word is
+// returned so callers can distinguish "truncated" from "nothing to do".
 func (f *BudgetEnforcer) TruncateToBudget(input string, budget int) string {
 	words := strings.Fields(input)
 	if len(words) == 0 {
 		return input
 	}
+	if budget <= 0 {
+		return ""
+	}
+	if EstimateTokens(input) <= budget {
+		return input
+	}
 
-	// Binary search for the right number of words
+	// Binary search for the right number of words. Uses the fast heuristic
+	// estimator: the precise BPE counter would re-encode a freshly joined
+	// prefix on every iteration (cache-missing each time) — O(n log n) BPE
+	// passes and cache pollution for large inputs. The heuristic is within a
+	// few tokens of BPE for wordy text, and the result is trimmed back below
+	// with the precise counter below.
 	lo, hi := 0, len(words)
 	for lo < hi {
 		mid := (lo + hi + 1) / 2
-		if EstimateTokens(strings.Join(words[:mid], " ")) <= budget {
+		if EstimateTokensFast(strings.Join(words[:mid], " ")) <= budget {
 			lo = mid
 		} else {
 			hi = mid - 1
 		}
 	}
 
 	if lo == 0 {
+		// Even the first word alone exceeds the budget. Keep a fragment of
+		// it (rune-safe) so the output is never empty.
+		return firstWordFragment(words[0], budget)
+	}
+
+	out := strings.Join(words[:lo], " ")
+	// Precise trim-back: the heuristic can land a few tokens over budget.
+	for lo > 0 && EstimateTokens(out) > budget {
+		lo--
+		out = strings.Join(words[:lo], " ")
+	}
+	if lo == 0 {
+		return firstWordFragment(words[0], budget)
+	}
+	return out
+}
+
+// firstWordFragment returns a leading, rune-safe fragment of word sized to
+// fit roughly within budget tokens. It guarantees a non-empty result for
+// non-empty input.
+func firstWordFragment(word string, budget int) string {
+	if budget <= 0 {
 		return ""
 	}
-	return strings.Join(words[:lo], " ")
+	runes := []rune(word)
+	max := budget * 4 // heuristic: ~4 runes per token
+	if max < 1 {
+		max = 1
+	}
+	if max > len(runes) {
+		max = len(runes)
+	}
+	return string(runes[:max])
 }
 
 // scoredLine represents a line with its importance score

internal/filter/budget_regression_test.go

@@ -0,0 +1,66 @@
+package filter
+
+import (
+	"strings"
+	"testing"
+)
+
+// Regression tests for budget enforcement edge cases.
+
+func TestTruncateToBudgetNeverReturnsEmptyForNonEmptyInput(t *testing.T) {
+	f := NewBudgetEnforcer(50)
+	// A single word that dominates the whole budget: the old binary search
+	// returned "" because it could not fit even one word.
+	in := strings.Repeat("supercalifragilistic", 2000) // ~72k chars, way over 50 tokens
+	out := f.TruncateToBudget(in, 50)
+	if out == "" {
+		t.Fatal("TruncateToBudget returned empty string for non-empty input")
+	}
+	if strings.Contains(out, " ") {
+		t.Fatalf("expected a fragment of the single word, got %q", out)
+	}
+	if len(out) > len(in) {
+		t.Fatalf("output longer than input")
+	}
+}
+
+func TestTruncateToBudgetFitsBudget(t *testing.T) {
+	f := NewBudgetEnforcer(40)
+	in := strings.Join(strings.Fields("word one two three four five six seven eight nine ten eleven twelve thirteen fourteen fifteen"), " ")
+	in = strings.Repeat(in+" ", 100)
+	out := f.TruncateToBudget(in, 40)
+	if out == "" {
+		t.Fatal("expected non-empty truncation")
+	}
+	// Allow a small slack: the fast estimator drives the search, then the
+	// precise trim-back guarantees compliance.
+	if EstimateTokens(out) > 45 {
+		t.Fatalf("truncated output exceeds budget: %d tokens for budget 40", EstimateTokens(out))
+	}
+	if !strings.HasPrefix(in, out) {
+		t.Fatalf("truncation must be a prefix of the input")
+	}
+}
+
+func TestBudgetZeroIsNoOp(t *testing.T) {
+	// Budget 0 (or negative) disables enforcement entirely: input passes
+	// through untouched, even when far over any plausible budget.
+	f := NewBudgetEnforcer(0)
+	in := strings.Repeat("this is a fairly long sentence with useful facts. ", 500)
+	out, saved := f.Apply(in, ModeMinimal)
+	if out != in {
+		t.Fatalf("budget 0 must be a no-op, got %d/%d chars", len(out), len(in))
+	}
+	if saved != 0 {
+		t.Fatalf("expected 0 tokens saved, got %d", saved)
+	}
+}
+
+func TestTruncateToBudgetBudgetZero(t *testing.T) {
+	f := NewBudgetEnforcer(0)
+	// Guard: budget <= 0 returns "" (callers must not pass 0 as a real limit).
+	out := f.TruncateToBudget("some words here", 0)
+	if out != "" {
+		t.Fatalf("expected empty for zero budget, got %q", out)
+	}
+}

internal/filter/filter.go

@@ -163,6 +163,13 @@ func EstimateTokens(text string) int {
 	return core.EstimateTokens(text)
 }
 
+// EstimateTokensFast provides a fast heuristic token count without BPE.
+// Delegates to core.EstimateTokensFast; used by hot loops where an exact
+// count is not required (e.g. binary-search budget fitting).
+func EstimateTokensFast(text string) int {
+	return core.EstimateTokensFast(text)
+}
+
 // IsCode checks if the output looks like source code.
 func IsCode(output string) bool {
 	codeIndicators := []string{

internal/filter/pipeline_process.go

@@ -72,7 +72,14 @@ func (p *PipelineCoordinator) Process(input string) (string, *PipelineStats) {
 	if p.budgetEnforcer != nil && p.config.Budget > 0 {
 		finalTokens := core.EstimateTokens(output)
 		if finalTokens > p.config.Budget {
+			before := finalTokens
 			output = p.budgetEnforcer.TruncateToBudget(output, p.config.Budget)
+			// Record the budget layer so Stats.HardTruncated() can tell a
+			// hard cut-off from structural compression.
+			saved := before - core.EstimateTokens(output)
+			if saved > 0 {
+				stats.AddLayerStatSafe(LayerBudget, LayerStat{TokensSaved: saved})
+			}
 		}
 	}
 

internal/filter/presets.go

@@ -55,7 +55,6 @@ func TierConfig(tier Tier, baseMode Mode) PipelineConfig {
 		cfg.EnableEvaluator = true
 		cfg.EnableH2O = true
 		cfg.EnableAttentionSink = true
-		cfg.EnableMetaToken = true
 	case TierExtract:
 		cfg.EnableEntropy = true
 		cfg.EnablePerplexity = true
@@ -70,7 +69,6 @@ func TierConfig(tier Tier, baseMode Mode) PipelineConfig {
 		cfg.EnableAttribution = true
 		cfg.EnableH2O = true
 		cfg.EnableAttentionSink = true
-		cfg.EnableMetaToken = true
 		cfg.EnableSemanticChunk = true
 		cfg.EnableLazyPruner = true
 		cfg.EnableSemanticAnchor = true
@@ -89,7 +87,6 @@ func TierConfig(tier Tier, baseMode Mode) PipelineConfig {
 		cfg.EnableAttribution = true
 		cfg.EnableH2O = true
 		cfg.EnableAttentionSink = true
-		cfg.EnableMetaToken = true
 		cfg.EnableSemanticChunk = true
 		cfg.EnableSketchStore = true
 		cfg.EnableLazyPruner = true
@@ -100,7 +97,6 @@ func TierConfig(tier Tier, baseMode Mode) PipelineConfig {
 		cfg.EnableAST = true
 		cfg.EnableGoalDriven = true
 		cfg.EnableH2O = true
-		cfg.EnableMetaToken = true
 	case TierLog:
 		cfg.EnableEntropy = true
 		cfg.EnablePerplexity = true

internal/filter/tier_config.go

@@ -223,7 +223,6 @@ func enableTierLayers(cfg *PipelineConfig, tier AutoTier) {
 		cfg.EnableAttribution = true
 		cfg.EnableH2O = true
 		cfg.EnableAttentionSink = true
-		cfg.EnableMetaToken = true
 		cfg.EnableSemanticChunk = true
 		cfg.EnableSketchStore = true
 		cfg.EnableLazyPruner = true