1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

// Tests for the spacing of the streamed text. The bug they guard: the model
// ends its sentence with line breaks — more and more of them as the turns go by
// — and the screen hollowed out by as much. Reproducible end to end with
// `SCENARIO=padding python3 ../03-one-agent-one-tool/docs/fake-engine.py`.
package engine

import "testing"

// next is called once per streamed chunk; we check what gets printed for each
// of them, in order.
func TestGapHoldsTrailingBlanks(t *testing.T) {
	cases := []struct {
		name      string
		fragments []string
		want      []string
	}{{
		name:      "plain text",
		fragments: []string{"Hello"},
		want:      []string{"Hello"},
	}, {
		name:      "trailing whitespace is not printed",
		fragments: []string{"Let me look:\n\n"},
		want:      []string{"Let me look:"},
	}, {
		name:      "it is printed when text follows, down to one blank line",
		fragments: []string{"para one\n\n\n\n\n", "para two"},
		want:      []string{"para one", "\n\npara two"},
	}, {
		name:      "leading whitespace is dropped before the first word",
		fragments: []string{"\n\n", "Hello"},
		want:      []string{"", "Hello"},
	}, {
		name:      "all-whitespace chunks: accumulated, then capped",
		fragments: []string{"a", "\n", "\n", "\n", "\n", "b"},
		want:      []string{"a", "", "", "", "", "\n\nb"},
	}, {
		name:      "the indentation of a code block survives",
		fragments: []string{"code:", "\n    fmt.Println()"},
		want:      []string{"code:", "\n    fmt.Println()"},
	}, {
		name:      "a single break between two lines stays single",
		fragments: []string{"line one\n", "line two"},
		want:      []string{"line one", "\nline two"},
	}, {
		name:      "empty chunk",
		fragments: []string{"a", "", "b"},
		want:      []string{"a", "", "b"},
	}}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			var g gap
			for i, frag := range c.fragments {
				if got := g.next(frag); got != c.want[i] {
					t.Errorf("chunk %d (%q): printed %q, want %q",
						i, frag, got, c.want[i])
				}
			}
		})
	}
}

// drop is called when a tool is about to write: the whitespace held back
// belonged to the text before it, and printing it after the 🛠️ line would
// reopen the hole.
func TestGapDropForgetsHeldBlanks(t *testing.T) {
	var g gap
	if got := g.next("Let me look:\n\n\n"); got != "Let me look:" {
		t.Fatalf("printed %q", got)
	}
	g.drop() // 🛠️ bash: … gets printed here
	if got := g.next("Now let me check:"); got != "Now let me check:" {
		t.Errorf("after drop: printed %q, want %q with no whitespace in front",
			got, "Now let me check:")
	}
}

// The full scenario of the bug: four turns whose trailing whitespace grows, a
// tool between each. None of that whitespace may be printed.
func TestGapPaddingScenario(t *testing.T) {
	var g gap
	turns := []string{
		"Let me look at the file:" + nl(2),
		"Now let me check something else:" + nl(4),
		"Let me try another approach:" + nl(7),
		"Let me rebuild the program:" + nl(11),
	}
	for i, turn := range turns {
		out := g.next(turn)
		if hasBlank(out) {
			t.Errorf("turn %d: %q contains trailing whitespace", i, out)
		}
		g.drop() // the tool writes its 🛠️ line
	}
}

func TestSqueeze(t *testing.T) {
	cases := map[string]string{
		"":                 "",
		"\n":               "\n",
		"\n\n":             "\n\n",
		"\n\n\n":           "\n\n",
		"\n\n\n\n\n\n\n\n": "\n\n",
		"\n    ":           "\n    ",   // indentation preserved
		"\n\n\n    ":       "\n\n    ", // capped, indentation kept
		"a\n\n\n\nb":       "a\n\nb",   // in the middle of the text too
		"   ":              "   ",      // spaces are left alone
	}
	for in, want := range cases {
		if got := squeeze(in); got != want {
			t.Errorf("squeeze(%q) = %q, want %q", in, got, want)
		}
	}
}

func nl(n int) string {
	s := ""
	for range n {
		s += "\n"
	}
	return s
}

func hasBlank(s string) bool {
	if s == "" {
		return false
	}
	last := s[len(s)-1]
	return last == '\n' || last == ' ' || last == '\t' || last == '\r'
}