chad-code 2.1.0__tar.gz → 2.3.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (120) hide show
  1. chad_code-2.3.0/PKG-INFO +93 -0
  2. chad_code-2.3.0/README.md +56 -0
  3. {chad_code-2.1.0 → chad_code-2.3.0}/pyproject.toml +11 -5
  4. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/__init__.py +1 -1
  5. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/agent.py +71 -30
  6. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/atif.py +151 -17
  7. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/base_engine.py +20 -0
  8. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/cli.py +105 -57
  9. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/completion_engine.py +1 -0
  10. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/engine.py +190 -22
  11. chad_code-2.3.0/src/chad/gguf_pack.py +795 -0
  12. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/mlx_dflash.py +109 -2
  13. chad_code-2.3.0/src/chad/mlx_fastpath.py +942 -0
  14. chad_code-2.3.0/src/chad/mlx_gguf.py +625 -0
  15. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/mlx_qmm_mma.py +150 -54
  16. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/mlx_qsdpa.py +40 -0
  17. chad_code-2.3.0/src/chad/prism_pack.py +253 -0
  18. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/prove.py +1 -1
  19. chad_code-2.3.0/src/chad_code.egg-info/PKG-INFO +93 -0
  20. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad_code.egg-info/SOURCES.txt +9 -1
  21. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad_code.egg-info/requires.txt +1 -0
  22. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_agent_e2e.py +83 -0
  23. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_atif.py +118 -4
  24. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_cli.py +131 -11
  25. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_completion_engine.py +19 -1
  26. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_engine_dflash.py +140 -0
  27. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_engine_kvquant.py +144 -0
  28. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_feel_pack.py +4 -1
  29. chad_code-2.3.0/tests/test_gguf_kernels.py +103 -0
  30. chad_code-2.3.0/tests/test_gguf_pack.py +490 -0
  31. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_mlx_qmm_mma.py +75 -14
  32. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_mlx_qsdpa.py +23 -0
  33. chad_code-2.3.0/tests/test_polyglot_harness.py +308 -0
  34. chad_code-2.3.0/tests/test_polyglot_kit.py +414 -0
  35. chad_code-2.3.0/tests/test_polyglot_proxy.py +439 -0
  36. chad_code-2.3.0/tests/test_prism_pack.py +374 -0
  37. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_prove.py +8 -2
  38. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_warm_prefix_tiers.py +24 -1
  39. chad_code-2.1.0/PKG-INFO +0 -291
  40. chad_code-2.1.0/README.md +0 -255
  41. chad_code-2.1.0/src/chad/mlx_fastpath.py +0 -433
  42. chad_code-2.1.0/src/chad_code.egg-info/PKG-INFO +0 -291
  43. chad_code-2.1.0/tests/test_matrix_bench.py +0 -484
  44. {chad_code-2.1.0 → chad_code-2.3.0}/LICENSE +0 -0
  45. {chad_code-2.1.0 → chad_code-2.3.0}/setup.cfg +0 -0
  46. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/ambient.py +0 -0
  47. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/bench.py +0 -0
  48. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/checkpoint.py +0 -0
  49. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/compaction.py +0 -0
  50. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/config.py +0 -0
  51. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/diag.py +0 -0
  52. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/guardrails.py +0 -0
  53. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/ignore.py +0 -0
  54. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/levers.py +0 -0
  55. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/mcp.py +0 -0
  56. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/mcp_oauth.py +0 -0
  57. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/parakeet/LICENSE +0 -0
  58. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/parakeet/__init__.py +0 -0
  59. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/parakeet/alignment.py +0 -0
  60. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/parakeet/attention.py +0 -0
  61. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/parakeet/audio.py +0 -0
  62. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/parakeet/cache.py +0 -0
  63. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/parakeet/conformer.py +0 -0
  64. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/parakeet/ctc.py +0 -0
  65. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/parakeet/parakeet.py +0 -0
  66. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/parakeet/rnnt.py +0 -0
  67. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/parakeet/tokenizer.py +0 -0
  68. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/parakeet/utils.py +0 -0
  69. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/prompt.py +0 -0
  70. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/render.py +0 -0
  71. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/repomap.py +0 -0
  72. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/seatbelt.py +0 -0
  73. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/session.py +0 -0
  74. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/skills.py +0 -0
  75. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/speech.py +0 -0
  76. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/spill.py +0 -0
  77. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/syntaxgate.py +0 -0
  78. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/toolcall_parse.py +0 -0
  79. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/tools.py +0 -0
  80. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/tui.py +0 -0
  81. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad/validate.py +0 -0
  82. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad_code.egg-info/dependency_links.txt +0 -0
  83. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad_code.egg-info/entry_points.txt +0 -0
  84. {chad_code-2.1.0 → chad_code-2.3.0}/src/chad_code.egg-info/top_level.txt +0 -0
  85. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_agent.py +0 -0
  86. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_agent_guards.py +0 -0
  87. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_ambient.py +0 -0
  88. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_bench.py +0 -0
  89. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_checkpoint.py +0 -0
  90. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_cli_modes.py +0 -0
  91. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_compaction.py +0 -0
  92. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_config.py +0 -0
  93. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_confirm_preview.py +0 -0
  94. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_edit.py +0 -0
  95. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_engine.py +0 -0
  96. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_engine_pld_hybrid.py +0 -0
  97. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_engine_pld_wide.py +0 -0
  98. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_gate.py +0 -0
  99. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_ignore.py +0 -0
  100. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_intent.py +0 -0
  101. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_levers.py +0 -0
  102. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_log_redaction.py +0 -0
  103. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_mcp.py +0 -0
  104. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_mcp_oauth.py +0 -0
  105. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_mlx_fastpath.py +0 -0
  106. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_render.py +0 -0
  107. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_repomap.py +0 -0
  108. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_repomap_polyglot.py +0 -0
  109. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_seatbelt.py +0 -0
  110. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_session.py +0 -0
  111. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_skills.py +0 -0
  112. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_speech.py +0 -0
  113. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_speech_tui.py +0 -0
  114. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_spill.py +0 -0
  115. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_syntaxgate.py +0 -0
  116. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_toolcall_dialect.py +0 -0
  117. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_toolcall_parse.py +0 -0
  118. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_tools.py +0 -0
  119. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_tui.py +0 -0
  120. {chad_code-2.1.0 → chad_code-2.3.0}/tests/test_validate.py +0 -0
@@ -0,0 +1,93 @@
1
+ Metadata-Version: 2.4
2
+ Name: chad-code
3
+ Version: 2.3.0
4
+ Summary: Claude-Code-style coding agent that runs entirely on a 24 GB Apple Silicon Mac (MLX, Qwen3.8-27B)
5
+ License-Expression: MIT
6
+ Project-URL: Repository, https://github.com/nathansutton/chad
7
+ Project-URL: Issues, https://github.com/nathansutton/chad/issues
8
+ Keywords: mlx,coding-agent,apple-silicon,local-llm,tui
9
+ Classifier: Development Status :: 4 - Beta
10
+ Classifier: Environment :: Console
11
+ Classifier: Intended Audience :: Developers
12
+ Classifier: Operating System :: MacOS
13
+ Classifier: Programming Language :: Python :: 3.11
14
+ Classifier: Programming Language :: Python :: 3.12
15
+ Classifier: Programming Language :: Python :: 3.13
16
+ Classifier: Programming Language :: Python :: 3.14
17
+ Classifier: Topic :: Software Development :: Code Generators
18
+ Requires-Python: >=3.11
19
+ Description-Content-Type: text/markdown
20
+ License-File: LICENSE
21
+ Requires-Dist: mlx!=0.32.1,<=0.32.2,>=0.32.0; sys_platform == "darwin"
22
+ Requires-Dist: mlx-lm<0.32,>=0.31.3; sys_platform == "darwin"
23
+ Requires-Dist: transformers<5.13,>=5.0
24
+ Requires-Dist: jinja2>=3.1
25
+ Requires-Dist: numpy<3,>=1.26
26
+ Requires-Dist: huggingface-hub>=1.22.0
27
+ Requires-Dist: prompt_toolkit>=3.0.43
28
+ Requires-Dist: tree-sitter>=0.26.0
29
+ Requires-Dist: tree-sitter-language-pack>=1.12.2
30
+ Requires-Dist: mcp<3,>=2
31
+ Requires-Dist: gguf<0.20,>=0.19.0
32
+ Provides-Extra: highlight
33
+ Requires-Dist: pygments>=2.17; extra == "highlight"
34
+ Provides-Extra: speech
35
+ Requires-Dist: sounddevice>=0.5; extra == "speech"
36
+ Dynamic: license-file
37
+
38
+ # chad: a coding agent for your macbook pro
39
+
40
+ [![tests](https://github.com/nathansutton/chad/actions/workflows/tests.yml/badge.svg)](https://github.com/nathansutton/chad/actions/workflows/tests.yml)
41
+
42
+ <img src="docs/claude-vs-chad.jpg" width="50%" alt="Two staircase newel posts side by side: Claude is a hand-carved wooden horse head, chad is a scuffed plastic toy horse tied on with twine">
43
+
44
+ > Claude can do anything, for anyone, anywhere. chad does one thing. 🗿
45
+ > *Coding under supervision.*
46
+
47
+ chad is a coding agent that runs entirely on an Apple Silicon Mac via
48
+ [MLX](https://github.com/ml-explore/mlx). One 27B model and no API key.
49
+ (Not affiliated with Anthropic.)
50
+
51
+ ```bash
52
+ uvx chad-code # runs anywhere; the command is still `chad`
53
+ uvx chad-code prove # offline smoke test: 4 tiny fix-it tasks, verified, timed 🗿
54
+ ```
55
+
56
+ The first run asks, then downloads the model once (~14 GB). The PyPI package is `chad-code`.
57
+
58
+ ![chad fixing a failing test end to end: reason, read, edit, run pytest, confirm green, all on a local model](docs/demo.gif)
59
+
60
+ > Real session, unedited.
61
+
62
+ ## Why chad
63
+
64
+ Plenty of harnesses run local models now, and [pi](https://pi.dev) is a fantastic default
65
+ for the same reason llama.cpp is: it works with everything. chad is moving the opposite direction.
66
+
67
+ **1 set of silicon**. This project is focused on making the macbook pro you already have usable. Not a $10K GPU.
68
+
69
+ **1 capable model**. `Qwen 3.8 27B`. This isn't the frontier, but you probably aren't solving frontier problems. Focus on 1 model buys _speed_. You'll experience ~ 20 tokens/second generation in a real session instead of ~ 10 tokens/second for stock llama.cpp implementations. This speed comes from MLX, a couple of targeted custom kernels for this model, and a bundled dflash2 drafter. The weights are Unsloth's `UD-Q3_K_XL` GGUF, read natively in MLX. We shipped a smaller ternary build first and it was measurably worse; [here's why we switched](docs/design.md#the-weights).
70
+
71
+ **1 tightly-coupled agent loop**. Instead of a standard /completions endpoint, the agent loop in chad owns the backend process. This comes with nice advantages that make the KV cache more stable and the coding experience measurably better (no long prefills!).
72
+
73
+ ## Why not
74
+
75
+ You do not have an Apple Silicon with 24 GB RAM. You want to pick your local model. You need a frontier model in a data center. The list goes on.
76
+
77
+ ## Documentation
78
+
79
+ - [Installing & using chad](docs/usage.md) covers install, extras and upgrades, the
80
+ terminal UI, and the command-line flags.
81
+ - [Throughput & performance](docs/benchmarks.md) has every number above, the stock-engine
82
+ comparison, the model, and how to reproduce them with `chad-bench`.
83
+ - [Design](docs/design.md) is the argument: why the agent owns the engine, why there are
84
+ five tools, and what 1.x got wrong.
85
+ - [Architecture](docs/architecture.md) is the module map, the session file format and the
86
+ tool-call wire format.
87
+ - [Configuration reference](docs/configuration.md) documents project instructions, Agent
88
+ Skills, MCP servers, plan mode, the slash commands, the context window, every environment
89
+ variable, and the safety opt-outs.
90
+ - [Troubleshooting](docs/troubleshooting.md) maps symptoms to knobs for when a session
91
+ rambles, loops, or slows.
92
+ - [Contributing](CONTRIBUTING.md) says what lands easily and what needs a conversation
93
+ first.
@@ -0,0 +1,56 @@
1
+ # chad: a coding agent for your macbook pro
2
+
3
+ [![tests](https://github.com/nathansutton/chad/actions/workflows/tests.yml/badge.svg)](https://github.com/nathansutton/chad/actions/workflows/tests.yml)
4
+
5
+ <img src="docs/claude-vs-chad.jpg" width="50%" alt="Two staircase newel posts side by side: Claude is a hand-carved wooden horse head, chad is a scuffed plastic toy horse tied on with twine">
6
+
7
+ > Claude can do anything, for anyone, anywhere. chad does one thing. 🗿
8
+ > *Coding under supervision.*
9
+
10
+ chad is a coding agent that runs entirely on an Apple Silicon Mac via
11
+ [MLX](https://github.com/ml-explore/mlx). One 27B model and no API key.
12
+ (Not affiliated with Anthropic.)
13
+
14
+ ```bash
15
+ uvx chad-code # runs anywhere; the command is still `chad`
16
+ uvx chad-code prove # offline smoke test: 4 tiny fix-it tasks, verified, timed 🗿
17
+ ```
18
+
19
+ The first run asks, then downloads the model once (~14 GB). The PyPI package is `chad-code`.
20
+
21
+ ![chad fixing a failing test end to end: reason, read, edit, run pytest, confirm green, all on a local model](docs/demo.gif)
22
+
23
+ > Real session, unedited.
24
+
25
+ ## Why chad
26
+
27
+ Plenty of harnesses run local models now, and [pi](https://pi.dev) is a fantastic default
28
+ for the same reason llama.cpp is: it works with everything. chad is moving the opposite direction.
29
+
30
+ **1 set of silicon**. This project is focused on making the macbook pro you already have usable. Not a $10K GPU.
31
+
32
+ **1 capable model**. `Qwen 3.8 27B`. This isn't the frontier, but you probably aren't solving frontier problems. Focus on 1 model buys _speed_. You'll experience ~ 20 tokens/second generation in a real session instead of ~ 10 tokens/second for stock llama.cpp implementations. This speed comes from MLX, a couple of targeted custom kernels for this model, and a bundled dflash2 drafter. The weights are Unsloth's `UD-Q3_K_XL` GGUF, read natively in MLX. We shipped a smaller ternary build first and it was measurably worse; [here's why we switched](docs/design.md#the-weights).
33
+
34
+ **1 tightly-coupled agent loop**. Instead of a standard /completions endpoint, the agent loop in chad owns the backend process. This comes with nice advantages that make the KV cache more stable and the coding experience measurably better (no long prefills!).
35
+
36
+ ## Why not
37
+
38
+ You do not have an Apple Silicon with 24 GB RAM. You want to pick your local model. You need a frontier model in a data center. The list goes on.
39
+
40
+ ## Documentation
41
+
42
+ - [Installing & using chad](docs/usage.md) covers install, extras and upgrades, the
43
+ terminal UI, and the command-line flags.
44
+ - [Throughput & performance](docs/benchmarks.md) has every number above, the stock-engine
45
+ comparison, the model, and how to reproduce them with `chad-bench`.
46
+ - [Design](docs/design.md) is the argument: why the agent owns the engine, why there are
47
+ five tools, and what 1.x got wrong.
48
+ - [Architecture](docs/architecture.md) is the module map, the session file format and the
49
+ tool-call wire format.
50
+ - [Configuration reference](docs/configuration.md) documents project instructions, Agent
51
+ Skills, MCP servers, plan mode, the slash commands, the context window, every environment
52
+ variable, and the safety opt-outs.
53
+ - [Troubleshooting](docs/troubleshooting.md) maps symptoms to knobs for when a session
54
+ rambles, loops, or slows.
55
+ - [Contributing](CONTRIBUTING.md) says what lands easily and what needs a conversation
56
+ first.
@@ -4,7 +4,7 @@
4
4
  # import name, and command name are independent. `uvx chad-code` runs the alias
5
5
  # script added under [project.scripts].
6
6
  name = "chad-code"
7
- version = "2.1.0"
7
+ version = "2.3.0"
8
8
  description = "Claude-Code-style coding agent that runs entirely on a 24 GB Apple Silicon Mac (MLX, Qwen3.8-27B)"
9
9
  readme = "README.md"
10
10
  license = "MIT"
@@ -75,6 +75,11 @@ dependencies = [
75
75
  # negotiates protocol 2026-07-28; the 1.x line tops out at 2025-11-25. Capped <3:
76
76
  # the same class of break will happen again.
77
77
  "mcp>=2,<3",
78
+ # GGUF checkpoints (gguf_pack.py): the header reader and the numpy block decoders the
79
+ # Metal kernels are tested against. Capped <0.20: the reader's field layout and the
80
+ # quant-type enum are what the loader routes on, and 0.19.0 is the line it was
81
+ # written and checked against.
82
+ "gguf>=0.19.0,<0.20",
78
83
  ]
79
84
 
80
85
  # Extras ride on the INSTALL SPEC, so the command depends on how chad was installed:
@@ -143,8 +148,7 @@ testpaths = ["tests"]
143
148
 
144
149
  [tool.ruff]
145
150
  line-length = 100
146
- # Vendored Exercism exercises (benchmarks/matrix/tasks/) are upstream files, kept byte-identical.
147
- extend-exclude = ["benchmarks/matrix/tasks", "tools/anti_slop"] # vendored linter: not chad's style
151
+ extend-exclude = ["tools/anti_slop"] # vendored linter: not chad's style
148
152
  src = ["src", "tests"]
149
153
 
150
154
  [tool.ruff.lint]
@@ -169,7 +173,9 @@ ignore = ["E741"] # ambiguous var names (l/I) appear in existing math code
169
173
 
170
174
  [tool.mypy]
171
175
  python_version = "3.11"
172
- files = ["src/chad"]
176
+ files = ["src/chad", "benchmarks/polyglot"]
177
+ # The kit's underscore dirs hold fetched exercises, trial workspaces and model output.
178
+ exclude = ['^benchmarks/polyglot/_']
173
179
  ignore_missing_imports = true # mlx/tree_sitter/mcp ship partial or no stubs
174
180
  warn_unused_ignores = true
175
181
  implicit_optional = true # `x: T = None` defaults are idiomatic throughout; this
@@ -215,7 +221,7 @@ exclude = [
215
221
  ".agents/**",
216
222
  ".claude/**",
217
223
  "src/chad/parakeet/**", # vendored upstream, not held to chad's gate (see ruff/mypy)
218
- "benchmarks/matrix/tasks/**", # vendored Exercism exercises, kept byte-identical
224
+ "benchmarks/polyglot/_*/**", # fetched exercises, trial workspaces, model output
219
225
  "tools/anti_slop/**",
220
226
  ]
221
227
  # Pre-existing findings live in the baseline (regenerate: `make slop-baseline`). A run
@@ -4,7 +4,7 @@ A flat collection of cooperating modules behind one console script (``chad``):
4
4
  the inference engine, the tool layer, the agent loop, and the terminal UI.
5
5
  """
6
6
 
7
- __version__ = "2.1.0"
7
+ __version__ = "2.3.0"
8
8
 
9
9
  # chad sets no MLX_* runtime vars. MLX_METAL_FAST_SYNCH, MLX_MAX_OPS_PER_BUFFER
10
10
  # and MLX_MAX_MB_PER_BUFFER were each measured end-to-end on the 35B and every
@@ -584,7 +584,7 @@ class Agent:
584
584
  # ATIF trajectory capture, off unless CHAD_TRAJECTORY_JSON is set.
585
585
  self._atif = atif.recorder()
586
586
  self._atif_seg = self._atif.new_segment() if self._atif else None
587
- self._atif_stats: list = [] # one entry per successful generate, in step order
587
+ self._atif_stats: list[atif.StepStat] = [] # one per stored assistant turn, in order
588
588
  if self._atif and self._atif.model_name is None:
589
589
  self._atif.model_name = engine.model_id
590
590
  self.interrupted = False
@@ -753,6 +753,12 @@ class Agent:
753
753
  self._split_ok = ok
754
754
  return ok
755
755
 
756
+ def _template_extra(self) -> dict:
757
+ """The template kwargs beyond messages/tools/thinking: the reasoning effort,
758
+ when one is set (the user's, else the engine's loader-carried default)."""
759
+ effort = self.reasoning_effort or self.engine.reasoning_effort_default
760
+ return {"reasoning_effort": effort} if effort else {}
761
+
756
762
  def _render(self, thinking: bool = None):
757
763
  # `thinking` overrides self.thinking for THIS render only (no-think
758
764
  # escalation renders one step with <think> off, then restores). None => self.thinking.
@@ -772,13 +778,14 @@ class Agent:
772
778
  # malformed TypeScript it cannot repair. `low` keeps the reasoning block and
773
779
  # shrinks it. Passed only when set, and via **kwargs, so templates that do not
774
780
  # accept the argument (q3_s6's, which has no reasoning_effort at all) render
775
- # byte-identically to before.
776
- extra = {}
777
- if self.reasoning_effort:
778
- extra["reasoning_effort"] = self.reasoning_effort
781
+ # byte-identically to before. A loader may carry a default of its own
782
+ # (engine.reasoning_effort_default) for a template whose default it disagrees
783
+ # with; the user's setting still wins. `_prefix_ids` renders with the same
784
+ # kwargs: the effort lands in the system block, so a prefix priced without
785
+ # it would not be a prefix of the prompt.
779
786
  ids = self._template_ids(self.engine.tok.apply_chat_template(
780
787
  messages, tools=self._active_schemas(), add_generation_prompt=True,
781
- enable_thinking=thinking, **extra,
788
+ enable_thinking=thinking, **self._template_extra(),
782
789
  ))
783
790
  # Debug hook (env-gated, off by default): dump the first decoded render so a
784
791
  # rendered-prompt difference across environments can be diffed. Best-effort.
@@ -808,7 +815,8 @@ class Agent:
808
815
  def render1(u):
809
816
  return self._template_ids(self.engine.tok.apply_chat_template(
810
817
  [sysm, {"role": "user", "content": u}], tools=schemas,
811
- add_generation_prompt=True, enable_thinking=self.thinking))
818
+ add_generation_prompt=True, enable_thinking=self.thinking,
819
+ **self._template_extra()))
812
820
  a, b = render1("a"), render1("the quick brown fox jumps")
813
821
  n = 0
814
822
  for x, y in zip(a, b):
@@ -835,7 +843,8 @@ class Agent:
835
843
  def render1(m):
836
844
  return self._template_ids(self.engine.tok.apply_chat_template(
837
845
  [m, {"role": "user", "content": "a"}], tools=schemas,
838
- add_generation_prompt=True, enable_thinking=self.thinking))
846
+ add_generation_prompt=True, enable_thinking=self.thinking,
847
+ **self._template_extra()))
839
848
  a, b = render1({"role": "system", "content": static}), render1(sysm)
840
849
  n = 0
841
850
  for x, y in zip(a, b):
@@ -905,6 +914,7 @@ class Agent:
905
914
  if self._atif is None or self._atif_seg is None:
906
915
  return
907
916
  try:
917
+ self._atif.set_pending(self._atif_seg, None) # its message has landed by now
908
918
  self._atif.set_segment(
909
919
  self._atif_seg,
910
920
  atif.steps_from_messages(self.messages, self._atif.model_name,
@@ -913,6 +923,18 @@ class Agent:
913
923
  except Exception as e: # never let telemetry break a turn
914
924
  log.warning("atif: sync failed: %s", e)
915
925
 
926
+ def _atif_pending(self, text: str, thinking: bool, generated: int, started: str) -> None:
927
+ """Show the step still being generated: the one a runaway `<think>` never finishes,
928
+ and so the one a rebuild from `messages` can never contain."""
929
+ if self._atif is None or self._atif_seg is None:
930
+ return
931
+ try:
932
+ self._atif.set_pending(self._atif_seg, atif.pending_step(
933
+ text, thinking, self._atif.model_name, generated, started))
934
+ self._atif.dump()
935
+ except Exception as e: # never let telemetry break a turn
936
+ log.warning("atif: pending dump failed: %s", e)
937
+
916
938
  def run_turn(self, user_text: str, stream=True):
917
939
  """Thin wrapper so the trajectory is flushed on EVERY exit from the turn —
918
940
  normal return, interrupt, or an exception escaping the loop. Harbor SIGKILLs chad
@@ -1176,11 +1198,18 @@ class Agent:
1176
1198
 
1177
1199
  view = _StreamView(self._emit, started_in_think=step_thinking) if stream else None
1178
1200
  gen_count = [0] # decoded chunks this step (~tokens); fed to the live ↓ counter
1201
+ atif_parts: list[str] = [] # this step's text so far, kept only when recording
1202
+ atif_started = atif.now() if self._atif else ""
1179
1203
 
1180
1204
  def on_token(t):
1181
1205
  if view:
1182
1206
  view.feed(t)
1183
1207
  gen_count[0] += 1
1208
+ if self._atif:
1209
+ atif_parts.append(t)
1210
+ if gen_count[0] % atif.PENDING_EVERY == 0:
1211
+ self._atif_pending("".join(atif_parts), step_thinking,
1212
+ gen_count[0], atif_started)
1184
1213
  # Throttle: a status emit per ~16 tokens keeps the queue cheap while the
1185
1214
  # bottom-line ↓ counter still climbs visibly (the refresher renders ~20 Hz).
1186
1215
  if gen_count[0] % 16 == 0:
@@ -1224,12 +1253,6 @@ class Agent:
1224
1253
  # non-trimmable cache its decoded text won't re-tokenize identically, so it
1225
1254
  # also forces the next-step re-prefill). Tracked so the no-call branch can
1226
1255
  # tell "truncated mid-thought" apart from "deliberately answered."
1227
- if self._atif: # one entry per SUCCESSFUL generate — a retried step appends
1228
- # nothing, keeping this list aligned with assistant messages
1229
- self._atif_stats.append(
1230
- {"prompt_tokens": stats.prompt_tokens,
1231
- "cached_tokens": stats.cached_tokens,
1232
- "generated_tokens": stats.generated_tokens})
1233
1256
  hit_cap = stats.generated_tokens >= self.max_gen_tokens
1234
1257
  # Close-and-continue fired: the engine force-closed a runaway <think>
1235
1258
  # and decoded the action in the same step. Counted for autopsy (a salvaged step
@@ -1259,6 +1282,40 @@ class Agent:
1259
1282
  # content.
1260
1283
  text = _SPECIAL_TOKEN_RE.sub("", text).rstrip()
1261
1284
 
1285
+ # Estimate reasoning overhead: the generation opens inside <think> (the
1286
+ # template emits the opening tag), so everything up to </think> is thinking,
1287
+ # and a thinking step that never closed the block — cut by a soft cap, the
1288
+ # token cap or an interrupt — is reasoning end to end. Counted before the
1289
+ # interrupt return below: the step a wall clock cuts is the runaway think,
1290
+ # and dropping it made the longest reasoning in a run read as none at all.
1291
+ # (A repetition stop that fired after </think> takes the fraction path.)
1292
+ if "</think>" in text:
1293
+ frac = len(text.split("</think>", 1)[0]) / len(text)
1294
+ _think_delta = int(stats.generated_tokens * frac)
1295
+ elif step_thinking or stats.stop_condition_fired:
1296
+ _think_delta = stats.generated_tokens
1297
+ else:
1298
+ _think_delta = 0
1299
+ self.think_tokens += _think_delta
1300
+ if self._atif and text: # one entry per STORED assistant turn — an interrupt
1301
+ # with no text stores none, and the list stays aligned
1302
+ finish = ("interrupted" if self._should_stop() else
1303
+ "max_tokens" if hit_cap else
1304
+ "repeat_stop" if rep_fired[0] else
1305
+ "think_cap" if stats.stop_condition_fired else "eos")
1306
+ self._atif_stats.append(atif.StepStat(
1307
+ prompt_tokens=stats.prompt_tokens, cached_tokens=stats.cached_tokens,
1308
+ generated_tokens=stats.generated_tokens, think_tokens=_think_delta,
1309
+ prefill_s=stats.prefill_s, gen_s=stats.gen_s, forwards=stats.forwards,
1310
+ draft_proposed=stats.draft_proposed,
1311
+ draft_accepted=stats.draft_accepted, finish=finish,
1312
+ salvaged=stats.salvaged, compacted=compacted,
1313
+ # "before </think>" is reasoning only when the step opened in a
1314
+ # think block; a no-think step drafted its whole turn as action.
1315
+ draft_hist_think=stats.draft_hist if step_thinking else {},
1316
+ draft_hist_act=stats.draft_hist_acting if step_thinking
1317
+ else stats.draft_hist))
1318
+
1262
1319
  # Interrupted (often mid-prefill, so text is empty): stop cleanly without
1263
1320
  # appending an empty assistant turn.
1264
1321
  if self._should_stop():
@@ -1275,22 +1332,6 @@ class Agent:
1275
1332
  self.messages.append({"role": "assistant",
1276
1333
  "content": close_unclosed_think(text, step_thinking)})
1277
1334
 
1278
- # Estimate reasoning overhead: the generation opens inside <think> (the
1279
- # template emits the opening tag), so everything up to </think> is thinking.
1280
- # A soft-cap stop fires only while still inside <think>, so ALL of
1281
- # this step's tokens are reasoning — count them so think-token telemetry (the
1282
- # metric the budget is measured against) doesn't under-report the capped runs.
1283
- # (A repetition stop can also land inside think — same accounting; one that
1284
- # fired after </think> falls to the fraction path below like any other turn.)
1285
- if stats.stop_condition_fired and "</think>" not in text:
1286
- _think_delta = stats.generated_tokens
1287
- elif "</think>" in text and len(text):
1288
- frac = len(text.split("</think>", 1)[0]) / len(text)
1289
- _think_delta = int(stats.generated_tokens * frac)
1290
- else:
1291
- _think_delta = 0
1292
- self.think_tokens += _think_delta
1293
-
1294
1335
  # Per-step first, turn-cumulative second. The cumulative ratio alone can't
1295
1336
  # show a step where drafting actually paid: a run of span-heavy tool-call
1296
1337
  # args is a few hundred tokens against a whole turn's proposals, so it moves
@@ -1,13 +1,13 @@
1
1
  """ATIF (Agent Trajectory Interchange Format) v1.7 trajectory emitter.
2
2
 
3
- The Terminal-Bench leaderboard requires an ATIF trajectory for **every passing trial**
4
- (the Apr-2026 "Leaderboard Integrity Update"). Harbor's built-in agents get this for free
5
- by subclassing its agent base; chad is driven as a CLI, so it has to emit the document
6
- itself or a submission is rejected.
3
+ A trial's row says whether it passed; its trajectory says where the tokens and the seconds
4
+ went, step by step, and `benchmarks/polyglot/trace.py` reads one back as a table — even
5
+ while the trial is still running. The document is ATIF, Harbor's published schema, rather
6
+ than a private format, so a trajectory another harness wrote reads the same way.
7
7
 
8
- Set `CHAD_TRAJECTORY_JSON=/path/trajectory.json` to turn this on. It is a pure observer:
9
- nothing here changes agent behavior, and any failure is swallowed (a benchmark harness
10
- must never lose a trial because telemetry broke).
8
+ Set `CHAD_TRAJECTORY_JSON=/path/trajectory.json`, or call `start()`, to turn this on. It
9
+ is a pure observer: nothing here changes agent behavior, and any failure is swallowed (a
10
+ benchmark harness must never lose a trial because telemetry broke).
11
11
 
12
12
  **Why derive from `agent.messages` rather than instrument the tool dispatch.**
13
13
  `run_turn` appends `{"role": "tool", ...}` from ~9 different sites (validation rejects,
@@ -26,6 +26,17 @@ reasoning_content, tool_calls[], observation, metrics, model_name}`. Every model
26
26
  `extra: "forbid"`, and a validator requires each `observation.results[].source_call_id`
27
27
  to name a `tool_call_id` **in the same step** — hence one step carries both a turn's tool
28
28
  calls and their results.
29
+
30
+ **What a step cost.** The schema's three token counts cannot tell a step that reasoned for
31
+ twenty minutes from one that decoded a long file, so each agent step's `metrics.extra`
32
+ carries the rest of what the engine measured (`StepStat`): the reasoning slice, prefill and
33
+ decode seconds, speculative yield, and why the generation ended.
34
+
35
+ **The step in flight.** A rebuild from `messages` only ever shows finished steps, and the
36
+ step worth seeing is the one that never finishes — a runaway `<think>` a harness kills at
37
+ its timeout leaves no message behind. So the agent hands over the partial text every
38
+ `PENDING_EVERY` tokens and the document ends in one step marked `extra.in_flight`, replaced
39
+ by the real step once its message lands.
29
40
  """
30
41
 
31
42
  import json
@@ -35,6 +46,8 @@ import re
35
46
  import tempfile
36
47
  import threading
37
48
  import uuid
49
+ from collections.abc import Mapping, Sequence
50
+ from dataclasses import dataclass, field
38
51
  from datetime import datetime, timezone
39
52
  from typing import Callable, Optional
40
53
 
@@ -48,7 +61,7 @@ SCHEMA_VERSION = "ATIF-v1.7"
48
61
  _THINK_BLOCK = re.compile(r"<think>(.*?)</think>", re.S)
49
62
 
50
63
 
51
- def _now() -> str:
64
+ def now() -> str:
52
65
  return datetime.now(timezone.utc).isoformat()
53
66
 
54
67
 
@@ -64,14 +77,66 @@ def split_think(text: str) -> tuple[str, str]:
64
77
  return reasoning, strip_think(text)
65
78
 
66
79
 
67
- def _metrics(stat: dict) -> dict:
80
+ PENDING_EVERY = 256 # tokens between rewrites of the in-flight step (~10 s of decode)
81
+
82
+
83
+ @dataclass(frozen=True)
84
+ class StepStat:
85
+ """What one generation cost, as the engine measured it. `finish` names why it ended:
86
+ `eos`, `max_tokens`, `think_cap`, `repeat_stop` or `interrupted`."""
87
+ prompt_tokens: int = 0 # newly prefilled — excludes the cached prefix
88
+ cached_tokens: int = 0
89
+ generated_tokens: int = 0
90
+ think_tokens: int = 0
91
+ prefill_s: float = 0.0
92
+ gen_s: float = 0.0
93
+ forwards: int = 0
94
+ draft_proposed: int = 0
95
+ draft_accepted: int = 0
96
+ finish: str = "eos"
97
+ salvaged: bool = False
98
+ compacted: bool = False # compaction shrank the render, so this step's prefill
99
+ # is a re-prefill of kept context, not new content
100
+ # Drafted rounds by proposed width k: entry i counts the rounds that accepted exactly
101
+ # i of the k. Split by phase, because reasoning and a tool call accept differently.
102
+ draft_hist_think: Mapping[int, Sequence[int]] = field(default_factory=dict)
103
+ draft_hist_act: Mapping[int, Sequence[int]] = field(default_factory=dict)
104
+
105
+
106
+ def _draft_phase(hist: Mapping[int, Sequence[int]]) -> dict[str, JsonValue]:
107
+ """One phase's drafting: totals, the per-verify yield, and the histogram itself.
108
+ JSON keys are strings, so the width is spelled out."""
109
+ rounds = sum(sum(row) for row in hist.values())
110
+ proposed = sum(k * sum(row) for k, row in hist.items())
111
+ accepted = sum(i * n for row in hist.values() for i, n in enumerate(row))
112
+ return {"rounds": rounds, "proposed": proposed, "accepted": accepted,
113
+ "accepted_per_round": round(accepted / rounds, 2) if rounds else 0.0,
114
+ "first_token_miss": round(sum(row[0] for row in hist.values()) / rounds, 3)
115
+ if rounds else 0.0,
116
+ "by_width": {str(k): list(row) for k, row in sorted(hist.items())}}
117
+
118
+
119
+ def _metrics(stat: StepStat) -> dict[str, JsonValue]:
68
120
  """chad's GenStats counts `prompt_tokens` as tokens *actually prefilled*, excluding the
69
121
  prefix served from cache. ATIF's `prompt_tokens` is documented as including cached
70
122
  tokens, so add them back — otherwise a warm KV cache reads as a shrinking prompt."""
71
- prompt = stat.get("prompt_tokens", 0) + stat.get("cached_tokens", 0)
72
- return {"prompt_tokens": prompt,
73
- "completion_tokens": stat.get("generated_tokens", 0),
74
- "cached_tokens": stat.get("cached_tokens", 0)}
123
+ extra: dict[str, JsonValue] = {
124
+ "think_tokens": stat.think_tokens, "prefill_s": round(stat.prefill_s, 2),
125
+ "gen_s": round(stat.gen_s, 2),
126
+ "tok_per_s": round(stat.generated_tokens / stat.gen_s, 1) if stat.gen_s > 0 else 0.0,
127
+ "forwards": stat.forwards,
128
+ "draft_proposed": stat.draft_proposed, "draft_accepted": stat.draft_accepted,
129
+ "finish": stat.finish}
130
+ if stat.salvaged:
131
+ extra["salvaged"] = True
132
+ if stat.compacted:
133
+ extra["compacted"] = True
134
+ if stat.draft_hist_think or stat.draft_hist_act:
135
+ extra["draft"] = {"think": _draft_phase(stat.draft_hist_think),
136
+ "act": _draft_phase(stat.draft_hist_act)}
137
+ return {"prompt_tokens": stat.prompt_tokens + stat.cached_tokens,
138
+ "completion_tokens": stat.generated_tokens,
139
+ "cached_tokens": stat.cached_tokens, "extra": extra}
75
140
 
76
141
 
77
142
  # First-seen timestamps, keyed by the identity of the message each step describes.
@@ -90,7 +155,7 @@ def _metrics(stat: dict) -> dict:
90
155
  _STAMPS: dict[int, tuple[dict, str]] = {}
91
156
 
92
157
 
93
- def _stamp(m: dict, clock: Callable[[], str] = _now) -> str:
158
+ def _stamp(m: dict, clock: Callable[[], str] = now) -> str:
94
159
  """The time `m` was first seen in a rebuild, minted on first sight."""
95
160
  hit = _STAMPS.get(id(m))
96
161
  if hit is not None and hit[0] is m:
@@ -100,8 +165,27 @@ def _stamp(m: dict, clock: Callable[[], str] = _now) -> str:
100
165
  return ts
101
166
 
102
167
 
168
+ def pending_step(text: str, thinking: bool, model_name: Optional[str], generated: int,
169
+ started: str) -> dict[str, JsonValue]:
170
+ """The step still being generated, from the text decoded so far.
171
+
172
+ On a thinking step the template opened `<think>`, so text with no close tag yet is all
173
+ reasoning — `split_think` alone would file it as the visible message."""
174
+ if thinking and "</think>" not in text:
175
+ reasoning, visible = text, ""
176
+ else:
177
+ reasoning, visible = split_think(text)
178
+ step: dict[str, JsonValue] = {"source": "agent", "message": visible, "timestamp": started,
179
+ "extra": {"in_flight": True, "generated_tokens": generated}}
180
+ if model_name:
181
+ step["model_name"] = model_name
182
+ if reasoning.strip():
183
+ step["reasoning_content"] = reasoning
184
+ return step
185
+
186
+
103
187
  def steps_from_messages(messages: list, model_name: Optional[str],
104
- stats: list, clock: Callable[[], str] = _now) -> list[dict]:
188
+ stats: list[StepStat], clock: Callable[[], str] = now) -> list[dict]:
105
189
  """Convert one Agent's `messages` into ATIF steps (without global `step_id`s).
106
190
 
107
191
  A `role: "tool"` message is not a step — it is an *observation* attached to the
@@ -167,6 +251,33 @@ def steps_from_messages(messages: list, model_name: Optional[str],
167
251
  return steps
168
252
 
169
253
 
254
+ def _totals(mets: list[dict]) -> dict[str, JsonValue]:
255
+ """The trial in one row: what a score table needs without re-reading every step."""
256
+ prompt = sum(m["prompt_tokens"] for m in mets)
257
+ gen_s = sum(m["extra"]["gen_s"] for m in mets)
258
+ out: dict[str, JsonValue] = {
259
+ "total_think_tokens": sum(m["extra"]["think_tokens"] for m in mets),
260
+ "total_gen_s": round(gen_s, 2),
261
+ "total_prefill_s": round(sum(m["extra"]["prefill_s"] for m in mets), 2),
262
+ "tok_per_s": round(sum(m["completion_tokens"] for m in mets) / gen_s, 1)
263
+ if gen_s > 0 else 0.0,
264
+ "cache_hit_rate": round(sum(m["cached_tokens"] for m in mets) / prompt, 3)
265
+ if prompt else 0.0,
266
+ "compactions": sum(1 for m in mets if m["extra"].get("compacted")),
267
+ "finishes": {f: sum(1 for m in mets if m["extra"]["finish"] == f)
268
+ for f in sorted({m["extra"]["finish"] for m in mets})}}
269
+ for phase in ("think", "act"):
270
+ blocks = [m["extra"]["draft"][phase] for m in mets if "draft" in m["extra"]]
271
+ rounds = sum(b["rounds"] for b in blocks)
272
+ if rounds:
273
+ proposed = sum(b["proposed"] for b in blocks)
274
+ accepted = sum(b["accepted"] for b in blocks)
275
+ out[f"draft_{phase}"] = {
276
+ "rounds": rounds, "accept_rate": round(accepted / proposed, 3),
277
+ "accepted_per_round": round(accepted / rounds, 2)}
278
+ return out
279
+
280
+
170
281
  def _pkg_version() -> str:
171
282
  """The installed chad-code version, for the trajectory's agent.version field.
172
283
  Was a hardcoded literal, which silently rots on every release — a leaderboard
@@ -191,6 +302,7 @@ class TrajectoryRecorder:
191
302
  self.session_id = str(uuid.uuid4())
192
303
  self.model_name: Optional[str] = None
193
304
  self._segments: list[list[dict]] = []
305
+ self._pending: dict[int, dict[str, JsonValue]] = {}
194
306
  self._lock = threading.Lock()
195
307
 
196
308
  def new_segment(self) -> int:
@@ -208,9 +320,21 @@ class TrajectoryRecorder:
208
320
  if 0 <= idx < len(self._segments):
209
321
  self._segments[idx] = steps
210
322
 
323
+ def set_pending(self, idx: int, step: Optional[dict[str, JsonValue]]) -> None:
324
+ """Set or clear (None) the in-flight step that trails segment `idx`."""
325
+ with self._lock:
326
+ if step is None:
327
+ self._pending.pop(idx, None)
328
+ else:
329
+ self._pending[idx] = step
330
+
211
331
  def to_dict(self) -> dict:
212
332
  with self._lock:
213
- steps = [s for seg in self._segments for s in seg]
333
+ steps = []
334
+ for idx, seg in enumerate(self._segments):
335
+ steps.extend(seg)
336
+ if idx in self._pending:
337
+ steps.append(dict(self._pending[idx]))
214
338
  for n, s in enumerate(steps, 1): # ATIF: sequential from 1, document-wide
215
339
  s["step_id"] = n
216
340
  agent: dict[str, JsonValue] = {"name": self.agent_name, "version": self.agent_version}
@@ -227,7 +351,8 @@ class TrajectoryRecorder:
227
351
  "total_prompt_tokens": sum(m["prompt_tokens"] for m in mets),
228
352
  "total_completion_tokens": sum(m["completion_tokens"] for m in mets),
229
353
  "total_cached_tokens": sum(m["cached_tokens"] for m in mets),
230
- "total_steps": len(steps)}
354
+ "total_steps": len(steps),
355
+ "extra": _totals(mets)}
231
356
  return doc
232
357
 
233
358
  def dump(self) -> None:
@@ -264,6 +389,15 @@ def recorder() -> Optional[TrajectoryRecorder]:
264
389
  return _RECORDER
265
390
 
266
391
 
392
+ def start(path: str) -> TrajectoryRecorder:
393
+ """Install a fresh process recorder writing to `path`, for a driver that runs many
394
+ trials in one process and wants a document per trial. Agents built afterwards record
395
+ into it; one built earlier keeps the recorder it was constructed with."""
396
+ global _RECORDER, _INIT
397
+ _RECORDER, _INIT = TrajectoryRecorder(path), True
398
+ return _RECORDER
399
+
400
+
267
401
  def _reset_for_tests() -> None:
268
402
  global _RECORDER, _INIT
269
403
  _RECORDER, _INIT = None, False