Skip to content

Commit 52e123f

Browse files
chlowellchlowell
authored
Unify token counting between check and tokens count (#146)
1 parent 8910269 commit 52e123f

16 files changed

Lines changed: 288 additions & 71 deletions

File tree

‎cmd/waza/cmd_check_test.go‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"testing"
1111
"unicode/utf8"
1212

13+
"github.com/microsoft/waza/cmd/waza/tokens"
1314
"github.com/microsoft/waza/internal/scaffold"
1415
"github.com/microsoft/waza/internal/scoring"
1516
"github.com/microsoft/waza/internal/validation"
@@ -84,6 +85,55 @@ This is the body of the test skill.
8485
assert.Contains(t, result, "Compliance Score:")
8586
}
8687

88+
func TestCheckTokenBudgetMatchesTokensCountForSameSkill(t *testing.T) {
89+
tmpDir := t.TempDir()
90+
skillContent := strings.Join([]string{
91+
"---",
92+
"name: shared-token-count",
93+
"description: This skill exists to prove that check and tokens count agree on the token count for the same SKILL.md file.",
94+
"---",
95+
"",
96+
"# Shared Token Count",
97+
"",
98+
"1. Normalize the content.",
99+
"2. Count with the default tokenizer.",
100+
"3. Report the same number in both commands.",
101+
"",
102+
"```bash",
103+
"echo consistent",
104+
"```",
105+
"",
106+
}, "\r\n")
107+
require.NoError(t, os.WriteFile(filepath.Join(tmpDir, "SKILL.md"), []byte(skillContent), 0644))
108+
109+
t.Chdir(tmpDir)
110+
111+
checkCmd := newRootCommand()
112+
var checkOut bytes.Buffer
113+
checkCmd.SetOut(&checkOut)
114+
checkCmd.SetErr(&checkOut)
115+
checkCmd.SetArgs([]string{"check", ".", "--format", "json"})
116+
require.NoError(t, checkCmd.Execute(), checkOut.String())
117+
118+
var checkReport checkJSONReport
119+
require.NoError(t, json.Unmarshal(checkOut.Bytes(), &checkReport), checkOut.String())
120+
require.Len(t, checkReport.Skills, 1)
121+
122+
countCmd := newRootCommand()
123+
var countOut bytes.Buffer
124+
countCmd.SetOut(&countOut)
125+
countCmd.SetErr(&countOut)
126+
countCmd.SetArgs([]string{"tokens", "count", ".", "--format", "json"})
127+
require.NoError(t, countCmd.Execute(), countOut.String())
128+
129+
var countReport tokens.CountJSONOutput
130+
require.NoError(t, json.Unmarshal(countOut.Bytes(), &countReport), countOut.String())
131+
132+
entry, ok := countReport.Files["SKILL.md"]
133+
require.True(t, ok, "expected SKILL.md in tokens count output: %s", countOut.String())
134+
require.Equal(t, entry.Tokens, checkReport.Skills[0].TokenBudget.Count)
135+
}
136+
87137
func TestCheckCommandNoSkillMd(t *testing.T) {
88138
tmpDir := t.TempDir()
89139

‎cmd/waza/dev/helpers_test.go‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,34 @@
11
package dev
22

33
import (
4+
"regexp"
5+
"strconv"
46
"strings"
57
"testing"
68

79
"github.com/microsoft/waza/internal/skill"
10+
"github.com/microsoft/waza/internal/testutil"
811
"github.com/microsoft/waza/internal/tokens"
912
"github.com/stretchr/testify/require"
1013
)
1114

15+
// requireOutputMatch verifies that actual matches expected after masking
16+
// token-count numbers. It then separately verifies that the first displayed
17+
// token count equals wantInitialTokens (the BPE count for the initial skill).
18+
func requireOutputMatch(t *testing.T, expected, actual string, wantInitialTokens int) {
19+
t.Helper()
20+
require.Equal(t, testutil.StripTokenCounts(expected), testutil.StripTokenCounts(actual),
21+
"output format mismatch (token counts masked)")
22+
got := -1
23+
var err error
24+
if m := regexp.MustCompile(`Tokens: (\d+)`).FindStringSubmatch(actual); m != nil {
25+
got, err = strconv.Atoi(m[1])
26+
require.NoError(t, err, "parsing token count from output")
27+
}
28+
require.Equal(t, wantInitialTokens, got,
29+
"first displayed token count should match BPE count of initial skill")
30+
}
31+
1232
func makeSkill(name, description string) *skill.Skill {
1333
raw := "---\nname: " + name + "\ndescription: " + description + "\n---\n"
1434
return &skill.Skill{

‎cmd/waza/dev/loop_test.go‎

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,7 @@ Anti-triggers: 0
126126
⚠️ body-structure: Advisory 17: body structure quality — body lacks actionable instructions (no code blocks, numbered steps, or commands); no examples section found; no error handling or troubleshooting section found
127127
✅ progressive-disclosure: Content structure supports progressive disclosure
128128
`
129-
require.Equal(t, expected, buf.String())
129+
requireOutputMatch(t, expected, buf.String(), skill.Tokens)
130130
}
131131

132132
func TestDevLoop_ScoreProgressionPath(t *testing.T) {
@@ -199,6 +199,9 @@ Already meets Medium-High target.
199199
err := runDevLoop(cfg)
200200
require.NoError(t, err)
201201

202+
sk, err := readSkillFile(filepath.Join(skillDir, "SKILL.md"))
203+
require.NoError(t, err)
204+
202205
expected := `Skill: compliant-skill
203206
Score: High
204207
Tokens: 94
@@ -236,7 +239,7 @@ MCP Integration: 1/4
236239
237240
✅ Target adherence level High reached!
238241
`
239-
require.Equal(t, expected, buf.String())
242+
requireOutputMatch(t, expected, buf.String(), sk.Tokens)
240243
}
241244

242245
func TestDevLoop_RunDevLoop_MaxIterationsHit(t *testing.T) {
@@ -254,6 +257,9 @@ description: "Short"
254257
`
255258
require.NoError(t, os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte(skillContent), 0o644))
256259

260+
initialSkill, err := readSkillFile(filepath.Join(skillDir, "SKILL.md"))
261+
require.NoError(t, err)
262+
257263
var buf bytes.Buffer
258264
cfg := &devConfig{
259265
SkillDir: skillDir,
@@ -264,7 +270,7 @@ description: "Short"
264270
In: &bytes.Buffer{},
265271
}
266272

267-
err := runDevLoop(cfg)
273+
err = runDevLoop(cfg)
268274
require.NoError(t, err)
269275

270276
expected := `
@@ -324,7 +330,7 @@ Short. Provides comprehensive support for common use cases and edge cases. Provi
324330
║ TOKEN STATUS: ✅ Under budget (66 < 500) ║
325331
╚══════════════════════════════════════════════════════════════════╝
326332
`
327-
require.Equal(t, expected, buf.String())
333+
requireOutputMatch(t, expected, buf.String(), initialSkill.Tokens)
328334
}
329335

330336
func TestDevConfig_Defaults(t *testing.T) {
@@ -365,6 +371,9 @@ description: "Short"
365371
`
366372
require.NoError(t, os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte(skillContent), 0o644))
367373

374+
initialSkill, err := readSkillFile(filepath.Join(skillDir, "SKILL.md"))
375+
require.NoError(t, err)
376+
368377
var buf bytes.Buffer
369378
cfg := &devConfig{
370379
SkillDir: skillDir,
@@ -375,7 +384,7 @@ description: "Short"
375384
In: &bytes.Buffer{},
376385
}
377386

378-
err := runDevLoop(cfg)
387+
err = runDevLoop(cfg)
379388
require.NoError(t, err)
380389

381390
expected := `
@@ -552,7 +561,7 @@ Anti-triggers: 2
552561
║ TOKEN STATUS: ✅ Under budget (162 < 500) ║
553562
╚══════════════════════════════════════════════════════════════════╝
554563
`
555-
require.Equal(t, expected, buf.String())
564+
requireOutputMatch(t, expected, buf.String(), initialSkill.Tokens)
556565

557566
final, err := readSkillFile(filepath.Join(skillDir, "SKILL.md"))
558567
require.NoError(t, err)
@@ -580,6 +589,9 @@ description: "Short"
580589
`
581590
require.NoError(t, os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte(skillContent), 0o644))
582591

592+
initialSkill, err := readSkillFile(filepath.Join(skillDir, "SKILL.md"))
593+
require.NoError(t, err)
594+
583595
var buf bytes.Buffer
584596
cfg := &devConfig{
585597
SkillDir: skillDir,
@@ -590,7 +602,7 @@ description: "Short"
590602
In: strings.NewReader(""),
591603
}
592604

593-
err := runDevLoop(cfg)
605+
err = runDevLoop(cfg)
594606
require.NoError(t, err)
595607

596608
expected := `
@@ -641,5 +653,5 @@ USE FOR: declined-skill, declined-skill help, use declined-skill, how to decline
641653
642654
No improvements applied.
643655
`
644-
require.Equal(t, expected, buf.String())
656+
requireOutputMatch(t, expected, buf.String(), initialSkill.Tokens)
645657
}

‎cmd/waza/tokens/check_test.go‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"strings"
88
"testing"
99

10+
"github.com/microsoft/waza/internal/testutil"
1011
"github.com/stretchr/testify/require"
1112
)
1213

@@ -37,7 +38,8 @@ references/two.md 6 10000 ✅ OK
3738
3839
4/4 files within limits
3940
`
40-
require.Equal(t, expected, out.String())
41+
require.Equal(t, testutil.StripTokenCounts(expected), testutil.StripTokenCounts(out.String()),
42+
"output format mismatch (token counts masked)")
4143
}
4244

4345
func TestCheck_SomeExceedLimit(t *testing.T) {
@@ -62,7 +64,8 @@ references/two.md 6 100 ✅ OK
6264
⚠️ 1 file(s) exceed their token limits:
6365
SKILL.md: 424 tokens (324 over limit of 100)
6466
`
65-
require.Equal(t, expected, out.String())
67+
require.Equal(t, testutil.StripTokenCounts(expected), testutil.StripTokenCounts(out.String()),
68+
"output format mismatch (token counts masked)")
6669
}
6770

6871
func TestCheck_StrictFails(t *testing.T) {
@@ -88,7 +91,8 @@ references/two.md 6 100 ✅ OK
8891
⚠️ 1 file(s) exceed their token limits:
8992
SKILL.md: 424 tokens (324 over limit of 100)
9093
`
91-
require.Equal(t, expected, err.Error())
94+
require.Equal(t, testutil.StripTokenCounts(expected), testutil.StripTokenCounts(err.Error()),
95+
"output format mismatch (token counts masked)")
9296
}
9397

9498
func TestCheck_StrictPassesWhenWithinLimit(t *testing.T) {
@@ -192,7 +196,8 @@ references/two.md 6 100 ✅ OK
192196
⚠️ 1 file(s) exceed their token limits:
193197
SKILL.md: 424 tokens (324 over limit of 100)
194198
`
195-
require.Equal(t, expected, err.Error())
199+
require.Equal(t, testutil.StripTokenCounts(expected), testutil.StripTokenCounts(err.Error()),
200+
"output format mismatch (token counts masked)")
196201
}
197202

198203
func TestCheck_SpecificFile(t *testing.T) {
@@ -296,9 +301,9 @@ SKILL.md 3 500 ✅ OK
296301
297302
2/2 files within limits
298303
`
299-
require.Equal(t, expected, out.String())
304+
require.Equal(t, testutil.StripTokenCounts(expected), testutil.StripTokenCounts(out.String()),
305+
"output format mismatch (token counts masked)")
300306
}
301-
302307
func TestCheck_ConfigPatternInJSON(t *testing.T) {
303308
td := checkFixture(t, "pattern")
304309
t.Chdir(td)

‎cmd/waza/tokens/compare.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -433,15 +433,15 @@ func compareFilesets(baseRef, headRef, rootDir string, baseFiles, headFiles map[
433433
before = &fileTokens{
434434
Tokens: beforeTokens,
435435
Characters: utf8.RuneCountInString(baseContent),
436-
Lines: countLines(baseContent),
436+
Lines: tokens.CountLines(baseContent),
437437
}
438438
}
439439
var after *fileTokens
440440
if hasHead {
441441
after = &fileTokens{
442442
Tokens: afterTokens,
443443
Characters: utf8.RuneCountInString(headContent),
444-
Lines: countLines(headContent),
444+
Lines: tokens.CountLines(headContent),
445445
}
446446
}
447447

‎cmd/waza/tokens/count.go‎

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -42,14 +42,14 @@ tokenizer provides a faster, less accurate estimate based on character count.`,
4242
return cmd
4343
}
4444

45-
type countJSONOutput struct {
45+
type CountJSONOutput struct {
4646
GeneratedAt string `json:"generatedAt"`
4747
TotalTokens int `json:"totalTokens"`
4848
TotalFiles int `json:"totalFiles"`
49-
Files map[string]countFileEntry `json:"files"`
49+
Files map[string]CountFileEntry `json:"files"`
5050
}
5151

52-
type countFileEntry struct {
52+
type CountFileEntry struct {
5353
Tokens int `json:"tokens"`
5454
Characters int `json:"characters"`
5555
Lines int `json:"lines"`
@@ -143,12 +143,11 @@ func runCount(cmd *cobra.Command, args []string) error {
143143
}
144144

145145
func countTokens(counter tokens.Counter, text, relPath string) *FileResult {
146-
text = strings.ReplaceAll(text, "\r\n", "\n")
147146
return &FileResult{
148147
Path: filepath.ToSlash(filepath.Clean(relPath)),
149148
Tokens: counter.Count(text),
150149
Characters: len(text),
151-
Lines: countLines(text),
150+
Lines: tokens.CountLines(text),
152151
}
153152
}
154153

@@ -202,18 +201,18 @@ func outputCountTable(w io.Writer, results []FileResult, showTotal bool) {
202201
}
203202

204203
func outputCountJSON(w io.Writer, results []FileResult) error {
205-
files := make(map[string]countFileEntry, len(results))
204+
files := make(map[string]CountFileEntry, len(results))
206205
totalTokens := 0
207206
for _, r := range results {
208207
totalTokens += r.Tokens
209-
files[r.Path] = countFileEntry{
208+
files[r.Path] = CountFileEntry{
210209
Tokens: r.Tokens,
211210
Characters: r.Characters,
212211
Lines: r.Lines,
213212
}
214213
}
215214

216-
out := countJSONOutput{
215+
out := CountJSONOutput{
217216
GeneratedAt: nowISO(),
218217
TotalTokens: totalTokens,
219218
TotalFiles: len(results),

0 commit comments

Comments
 (0)