@bendyline/gezel-inference 0.0.0-stage → 0.1.0

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 (967) hide show
  1. package/LICENSE +21 -0
  2. package/README.md +13 -2
  3. package/dist/cache/adapter.d.ts +115 -0
  4. package/dist/cache/adapter.js +23 -0
  5. package/dist/cache/budget.d.ts +15 -0
  6. package/dist/cache/budget.js +31 -0
  7. package/dist/cache/controller.d.ts +205 -0
  8. package/dist/cache/controller.js +377 -0
  9. package/dist/catalog/LICENSE +21 -0
  10. package/dist/catalog/chat-models/bt/btl4-35b-q4/manifest.json +146 -0
  11. package/dist/catalog/chat-models/bt/btl4-35b-q4/versions/1.0.0/manifest.json +15 -0
  12. package/dist/catalog/chat-models/bt/btl4-35b-q4/versions/1.0.1/manifest.json +15 -0
  13. package/dist/catalog/chat-models/bt/btl4-compact-iq2/manifest.json +146 -0
  14. package/dist/catalog/chat-models/bt/btl4-compact-iq2/versions/1.0.0/manifest.json +15 -0
  15. package/dist/catalog/chat-models/bt/btl4-compact-iq2/versions/1.0.1/manifest.json +15 -0
  16. package/dist/catalog/chat-models/bt/btl4-compact-iq2/versions/1.0.2/manifest.json +15 -0
  17. package/dist/catalog/chat-models/de/deepseek-r1-8b-q4/manifest.json +166 -0
  18. package/dist/catalog/chat-models/de/deepseek-r1-8b-q4/versions/1.1.0/manifest.json +62 -0
  19. package/dist/catalog/chat-models/de/deepseek-r1-8b-q4/versions/1.1.1/manifest.json +61 -0
  20. package/dist/catalog/chat-models/de/deepseek-v4-flash-284b-mxfp4/manifest.json +125 -0
  21. package/dist/catalog/chat-models/de/deepseek-v4-flash-284b-mxfp4/versions/1.0.0/manifest.json +27 -0
  22. package/dist/catalog/chat-models/de/deepseek-v4-flash-284b-q2/manifest.json +125 -0
  23. package/dist/catalog/chat-models/de/deepseek-v4-flash-284b-q2/versions/1.0.0/manifest.json +27 -0
  24. package/dist/catalog/chat-models/de/deepseek-v4-flash-284b-q2/versions/1.1.0/manifest.json +27 -0
  25. package/dist/catalog/chat-models/de/deepseek-v4-flash-284b-q2q4/manifest.json +126 -0
  26. package/dist/catalog/chat-models/de/deepseek-v4-flash-284b-q2q4/versions/1.0.0/manifest.json +27 -0
  27. package/dist/catalog/chat-models/de/deepseek-v4-flash-284b-q4/manifest.json +318 -0
  28. package/dist/catalog/chat-models/de/deepseek-v4-flash-284b-q4/versions/1.0.0/manifest.json +225 -0
  29. package/dist/catalog/chat-models/de/deepseek-v4-flash-284b-q4/versions/1.1.0/manifest.json +225 -0
  30. package/dist/catalog/chat-models/ge/gemma4-12b-q4/manifest.json +249 -0
  31. package/dist/catalog/chat-models/ge/gemma4-12b-q4/versions/1.0.0/manifest.json +70 -0
  32. package/dist/catalog/chat-models/ge/gemma4-12b-q4/versions/1.1.0/manifest.json +72 -0
  33. package/dist/catalog/chat-models/ge/gemma4-12b-q4/versions/1.1.2/manifest.json +78 -0
  34. package/dist/catalog/chat-models/ge/gemma4-12b-q4/versions/1.1.3/manifest.json +78 -0
  35. package/dist/catalog/chat-models/ge/gemma4-12b-q4/versions/1.2.0/manifest.json +98 -0
  36. package/dist/catalog/chat-models/ge/gemma4-12b-q8/manifest.json +228 -0
  37. package/dist/catalog/chat-models/ge/gemma4-12b-q8/versions/1.0.0/manifest.json +72 -0
  38. package/dist/catalog/chat-models/ge/gemma4-12b-q8/versions/1.0.2/manifest.json +78 -0
  39. package/dist/catalog/chat-models/ge/gemma4-12b-q8/versions/1.0.3/manifest.json +78 -0
  40. package/dist/catalog/chat-models/ge/gemma4-12b-q8/versions/1.1.0/manifest.json +78 -0
  41. package/dist/catalog/chat-models/ge/gemma4-26b-q4/manifest.json +220 -0
  42. package/dist/catalog/chat-models/ge/gemma4-26b-q4/versions/1.1.0/manifest.json +75 -0
  43. package/dist/catalog/chat-models/ge/gemma4-26b-q4/versions/1.2.0/manifest.json +75 -0
  44. package/dist/catalog/chat-models/ge/gemma4-26b-q4/versions/1.2.1/manifest.json +80 -0
  45. package/dist/catalog/chat-models/ge/gemma4-26b-q4/versions/1.2.2/manifest.json +80 -0
  46. package/dist/catalog/chat-models/ge/gemma4-26b-q4/versions/1.3.0/manifest.json +81 -0
  47. package/dist/catalog/chat-models/ge/gemma4-31b-q4/manifest.json +226 -0
  48. package/dist/catalog/chat-models/ge/gemma4-31b-q4/versions/1.1.0/manifest.json +80 -0
  49. package/dist/catalog/chat-models/ge/gemma4-31b-q4/versions/1.1.1/manifest.json +82 -0
  50. package/dist/catalog/chat-models/ge/gemma4-31b-q4/versions/1.2.0/manifest.json +82 -0
  51. package/dist/catalog/chat-models/ge/gemma4-31b-q4/versions/1.2.1/manifest.json +95 -0
  52. package/dist/catalog/chat-models/ge/gemma4-31b-q4/versions/1.2.2/manifest.json +95 -0
  53. package/dist/catalog/chat-models/ge/gemma4-31b-q4/versions/1.3.0/manifest.json +86 -0
  54. package/dist/catalog/chat-models/ge/gemma4-e2b-q4/manifest.json +213 -0
  55. package/dist/catalog/chat-models/ge/gemma4-e2b-q4/versions/1.1.0/manifest.json +67 -0
  56. package/dist/catalog/chat-models/ge/gemma4-e2b-q4/versions/1.1.2/manifest.json +75 -0
  57. package/dist/catalog/chat-models/ge/gemma4-e2b-q4/versions/1.2.0/manifest.json +71 -0
  58. package/dist/catalog/chat-models/ge/gemma4-e2b-q4/versions/1.2.1/manifest.json +71 -0
  59. package/dist/catalog/chat-models/ge/gemma4-e2b-q4/versions/1.3.0/manifest.json +71 -0
  60. package/dist/catalog/chat-models/ge/gemma4-e4b-q4/manifest.json +205 -0
  61. package/dist/catalog/chat-models/ge/gemma4-e4b-q4/versions/1.1.0/manifest.json +72 -0
  62. package/dist/catalog/chat-models/ge/gemma4-e4b-q4/versions/1.1.2/manifest.json +75 -0
  63. package/dist/catalog/chat-models/ge/gemma4-e4b-q4/versions/1.2.0/manifest.json +76 -0
  64. package/dist/catalog/chat-models/ge/gemma4-e4b-q4/versions/1.2.1/manifest.json +76 -0
  65. package/dist/catalog/chat-models/ge/gemma4-e4b-q4/versions/1.3.0/manifest.json +76 -0
  66. package/dist/catalog/chat-models/gl/glm-5.2-754b-q2/manifest.json +127 -0
  67. package/dist/catalog/chat-models/gl/glm-5.2-754b-q2/versions/1.0.0/manifest.json +20 -0
  68. package/dist/catalog/chat-models/gl/glm-5.3-flash-320b-q2/manifest.json +140 -0
  69. package/dist/catalog/chat-models/gl/glm-5.3-flash-320b-q2/versions/1.0.0/manifest.json +23 -0
  70. package/dist/catalog/chat-models/gl/glm-5.3-flash-320b-q2/versions/1.0.1/manifest.json +26 -0
  71. package/dist/catalog/chat-models/gp/gpt-oss-120b-q4/manifest.json +247 -0
  72. package/dist/catalog/chat-models/gp/gpt-oss-120b-q4/versions/1.0.0/manifest.json +137 -0
  73. package/dist/catalog/chat-models/gp/gpt-oss-120b-q4/versions/1.1.0/manifest.json +135 -0
  74. package/dist/catalog/chat-models/gp/gpt-oss-20b-q4/manifest.json +195 -0
  75. package/dist/catalog/chat-models/gp/gpt-oss-20b-q4/versions/1.1.0/manifest.json +72 -0
  76. package/dist/catalog/chat-models/gp/gpt-oss-20b-q4/versions/1.1.1/manifest.json +71 -0
  77. package/dist/catalog/chat-models/gr/granite4.2-30b-q4/manifest.json +200 -0
  78. package/dist/catalog/chat-models/gr/granite4.2-30b-q4/versions/1.0.0/manifest.json +72 -0
  79. package/dist/catalog/chat-models/gr/granite4.2-3b-q4/manifest.json +196 -0
  80. package/dist/catalog/chat-models/gr/granite4.2-3b-q4/versions/1.0.0/manifest.json +57 -0
  81. package/dist/catalog/chat-models/gr/granite4.2-8b-q4/manifest.json +189 -0
  82. package/dist/catalog/chat-models/gr/granite4.2-8b-q4/versions/1.0.0/manifest.json +57 -0
  83. package/dist/catalog/chat-models/index.json +1 -0
  84. package/dist/catalog/chat-models/la/laguna-s-2.1-118b-q4/manifest.json +235 -0
  85. package/dist/catalog/chat-models/la/laguna-s-2.1-118b-q4/versions/1.0.0/manifest.json +124 -0
  86. package/dist/catalog/chat-models/la/laguna-s-2.1-118b-q4/versions/1.0.1/manifest.json +124 -0
  87. package/dist/catalog/chat-models/la/laguna-s-2.1-118b-q4/versions/1.0.2/manifest.json +123 -0
  88. package/dist/catalog/chat-models/la/laguna-s-2.1-118b-q6/manifest.json +257 -0
  89. package/dist/catalog/chat-models/la/laguna-s-2.1-118b-q6/versions/1.0.0/manifest.json +145 -0
  90. package/dist/catalog/chat-models/la/laguna-s-2.1-118b-q6/versions/1.0.1/manifest.json +145 -0
  91. package/dist/catalog/chat-models/la/laguna-s-2.1-118b-q8/manifest.json +285 -0
  92. package/dist/catalog/chat-models/la/laguna-s-2.1-118b-q8/versions/1.0.0/manifest.json +174 -0
  93. package/dist/catalog/chat-models/la/laguna-s-2.1-118b-q8/versions/1.0.1/manifest.json +174 -0
  94. package/dist/catalog/chat-models/la/laguna-s-2.1-118b-q8/versions/1.0.2/manifest.json +173 -0
  95. package/dist/catalog/chat-models/lf/lfm2.5-2.6b-q4/manifest.json +162 -0
  96. package/dist/catalog/chat-models/lf/lfm2.5-2.6b-q4/versions/1.0.0/manifest.json +59 -0
  97. package/dist/catalog/chat-models/lf/lfm2.5-2.6b-q4/versions/1.0.1/manifest.json +58 -0
  98. package/dist/catalog/chat-models/ll/llama3.2-3b-q4/manifest.json +134 -0
  99. package/dist/catalog/chat-models/ll/llama3.2-3b-q4/versions/1.1.0/manifest.json +57 -0
  100. package/dist/catalog/chat-models/ll/llama3.2-3b-q4/versions/1.1.1/manifest.json +56 -0
  101. package/dist/catalog/chat-models/mi/minicpm5-1b-q4/manifest.json +156 -0
  102. package/dist/catalog/chat-models/mi/minicpm5-1b-q4/versions/1.0.0/manifest.json +57 -0
  103. package/dist/catalog/chat-models/mi/minicpm5-2b-q4/manifest.json +161 -0
  104. package/dist/catalog/chat-models/mi/minicpm5-2b-q4/versions/1.0.0/manifest.json +57 -0
  105. package/dist/catalog/chat-models/mi/mistral-7b-q4/manifest.json +159 -0
  106. package/dist/catalog/chat-models/mi/mistral-7b-q4/versions/1.1.0/manifest.json +62 -0
  107. package/dist/catalog/chat-models/mi/mistral-7b-q4/versions/1.1.1/manifest.json +61 -0
  108. package/dist/catalog/chat-models/mi/mistral-medium-3.5-128b-q4/manifest.json +275 -0
  109. package/dist/catalog/chat-models/mi/mistral-medium-3.5-128b-q4/versions/1.0.0/manifest.json +154 -0
  110. package/dist/catalog/chat-models/mi/mistral-medium-3.5-128b-q4/versions/1.0.1/manifest.json +153 -0
  111. package/dist/catalog/chat-models/mu/muse-glimmer-30b-q4/manifest.json +135 -0
  112. package/dist/catalog/chat-models/mu/muse-glimmer-30b-q4/versions/1.0.0/manifest.json +20 -0
  113. package/dist/catalog/chat-models/mu/muse-glimmer-30b-q4/versions/1.0.1/manifest.json +20 -0
  114. package/dist/catalog/chat-models/mu/muse-glimmer-30b-q4/versions/1.0.2/manifest.json +20 -0
  115. package/dist/catalog/chat-models/na/nanbeige4.2-3b-q4/manifest.json +138 -0
  116. package/dist/catalog/chat-models/na/nanbeige4.2-3b-q4/versions/1.0.0/manifest.json +14 -0
  117. package/dist/catalog/chat-models/ne/nemotron3-nano-30b-q4/manifest.json +212 -0
  118. package/dist/catalog/chat-models/ne/nemotron3-nano-30b-q4/versions/1.0.0/manifest.json +94 -0
  119. package/dist/catalog/chat-models/ne/nemotron3-nano-30b-q4/versions/1.0.1/manifest.json +94 -0
  120. package/dist/catalog/chat-models/ne/nemotron3-nano-30b-q4/versions/1.0.2/manifest.json +93 -0
  121. package/dist/catalog/chat-models/ne/nemotron3-super-120b-q4/manifest.json +135 -0
  122. package/dist/catalog/chat-models/ne/nemotron3-super-120b-q4/versions/1.0.0/manifest.json +30 -0
  123. package/dist/catalog/chat-models/ne/nemotron3-super-120b-q4/versions/1.0.1/manifest.json +30 -0
  124. package/dist/catalog/chat-models/ne/nemotron3.5-lightning-30b-q4/manifest.json +192 -0
  125. package/dist/catalog/chat-models/ne/nemotron3.5-lightning-30b-q4/versions/1.0.0/manifest.json +74 -0
  126. package/dist/catalog/chat-models/ne/nemotron3.5-lightning-30b-q4/versions/1.0.1/manifest.json +73 -0
  127. package/dist/catalog/chat-models/ne/nemotron3.5-lightning-30b-q8/manifest.json +207 -0
  128. package/dist/catalog/chat-models/ne/nemotron3.5-lightning-30b-q8/versions/1.0.0/manifest.json +89 -0
  129. package/dist/catalog/chat-models/ne/nemotron3.5-lightning-30b-q8/versions/1.0.1/manifest.json +88 -0
  130. package/dist/catalog/chat-models/or/ornith-35b-q4/manifest.json +224 -0
  131. package/dist/catalog/chat-models/or/ornith-35b-q4/versions/1.0.0/manifest.json +94 -0
  132. package/dist/catalog/chat-models/or/ornith-35b-q4/versions/1.0.1/manifest.json +93 -0
  133. package/dist/catalog/chat-models/or/ornith-35b-q8/manifest.json +244 -0
  134. package/dist/catalog/chat-models/or/ornith-35b-q8/versions/1.0.0/manifest.json +114 -0
  135. package/dist/catalog/chat-models/or/ornith-35b-q8/versions/1.0.1/manifest.json +113 -0
  136. package/dist/catalog/chat-models/or/ornith-9b-q4/manifest.json +215 -0
  137. package/dist/catalog/chat-models/or/ornith-9b-q4/versions/1.0.0/manifest.json +84 -0
  138. package/dist/catalog/chat-models/or/ornith-9b-q4/versions/1.0.1/manifest.json +83 -0
  139. package/dist/catalog/chat-models/or/ornith1.5-35b-a3b-q4/manifest.json +208 -0
  140. package/dist/catalog/chat-models/or/ornith1.5-35b-a3b-q4/versions/1.0.0/manifest.json +74 -0
  141. package/dist/catalog/chat-models/or/ornith1.5-35b-a3b-q8/manifest.json +223 -0
  142. package/dist/catalog/chat-models/or/ornith1.5-35b-a3b-q8/versions/1.0.0/manifest.json +89 -0
  143. package/dist/catalog/chat-models/or/ornith1.5-9b-q4/manifest.json +187 -0
  144. package/dist/catalog/chat-models/or/ornith1.5-9b-q4/versions/1.0.0/manifest.json +53 -0
  145. package/dist/catalog/chat-models/or/ornith1.5-9b-q4/versions/1.0.1/manifest.json +77 -0
  146. package/dist/catalog/chat-models/or/ornith1.5-9b-q8/manifest.json +192 -0
  147. package/dist/catalog/chat-models/or/ornith1.5-9b-q8/versions/1.0.0/manifest.json +58 -0
  148. package/dist/catalog/chat-models/or/ornith1.5-9b-q8/versions/1.0.1/manifest.json +82 -0
  149. package/dist/catalog/chat-models/qw/qwen3.5-0.8b-q4/manifest.json +192 -0
  150. package/dist/catalog/chat-models/qw/qwen3.5-0.8b-q4/versions/1.0.0/manifest.json +75 -0
  151. package/dist/catalog/chat-models/qw/qwen3.5-122b-a10b-q4/manifest.json +315 -0
  152. package/dist/catalog/chat-models/qw/qwen3.5-122b-a10b-q4/versions/1.0.0/manifest.json +164 -0
  153. package/dist/catalog/chat-models/qw/qwen3.5-122b-a10b-q4/versions/1.0.1/manifest.json +164 -0
  154. package/dist/catalog/chat-models/qw/qwen3.5-122b-a10b-q4/versions/1.0.2/manifest.json +164 -0
  155. package/dist/catalog/chat-models/qw/qwen3.5-122b-a10b-q4/versions/1.0.3/manifest.json +163 -0
  156. package/dist/catalog/chat-models/qw/qwen3.5-27b-q4/manifest.json +217 -0
  157. package/dist/catalog/chat-models/qw/qwen3.5-27b-q4/versions/1.0.0/manifest.json +90 -0
  158. package/dist/catalog/chat-models/qw/qwen3.5-2b-q4/manifest.json +194 -0
  159. package/dist/catalog/chat-models/qw/qwen3.5-2b-q4/versions/1.1.0/manifest.json +77 -0
  160. package/dist/catalog/chat-models/qw/qwen3.5-2b-q4/versions/1.1.2/manifest.json +76 -0
  161. package/dist/catalog/chat-models/qw/qwen3.5-2b-q4/versions/1.1.3/manifest.json +76 -0
  162. package/dist/catalog/chat-models/qw/qwen3.5-4b-q4/manifest.json +188 -0
  163. package/dist/catalog/chat-models/qw/qwen3.5-4b-q4/versions/1.1.0/manifest.json +77 -0
  164. package/dist/catalog/chat-models/qw/qwen3.5-4b-q4/versions/1.1.2/manifest.json +76 -0
  165. package/dist/catalog/chat-models/qw/qwen3.5-4b-q4/versions/1.1.3/manifest.json +76 -0
  166. package/dist/catalog/chat-models/qw/qwen3.5-9b-q4/manifest.json +205 -0
  167. package/dist/catalog/chat-models/qw/qwen3.5-9b-q4/versions/1.1.0/manifest.json +82 -0
  168. package/dist/catalog/chat-models/qw/qwen3.5-9b-q4/versions/1.1.2/manifest.json +81 -0
  169. package/dist/catalog/chat-models/qw/qwen3.5-9b-q4/versions/1.1.3/manifest.json +81 -0
  170. package/dist/catalog/chat-models/qw/qwen3.6-27b-q4/manifest.json +230 -0
  171. package/dist/catalog/chat-models/qw/qwen3.6-27b-q4/versions/1.1.0/manifest.json +92 -0
  172. package/dist/catalog/chat-models/qw/qwen3.6-27b-q4/versions/1.1.4/manifest.json +91 -0
  173. package/dist/catalog/chat-models/qw/qwen3.6-27b-q4/versions/1.1.5/manifest.json +91 -0
  174. package/dist/catalog/chat-models/qw/qwen3.6-27b-q4/versions/1.1.6/manifest.json +115 -0
  175. package/dist/catalog/chat-models/qw/qwen3.6-27b-q8/manifest.json +240 -0
  176. package/dist/catalog/chat-models/qw/qwen3.6-27b-q8/versions/1.0.0/manifest.json +102 -0
  177. package/dist/catalog/chat-models/qw/qwen3.6-27b-q8/versions/1.0.2/manifest.json +103 -0
  178. package/dist/catalog/chat-models/qw/qwen3.6-27b-q8/versions/1.0.3/manifest.json +103 -0
  179. package/dist/catalog/chat-models/qw/qwen3.6-27b-q8/versions/1.0.4/manifest.json +127 -0
  180. package/dist/catalog/chat-models/qw/qwen3.6-35b-a3b-q4/manifest.json +233 -0
  181. package/dist/catalog/chat-models/qw/qwen3.6-35b-a3b-q4/versions/1.0.0/manifest.json +92 -0
  182. package/dist/catalog/chat-models/qw/qwen3.6-35b-a3b-q4/versions/1.0.1/manifest.json +93 -0
  183. package/dist/catalog/chat-models/qw/qwen3.6-35b-a3b-q4/versions/1.0.2/manifest.json +93 -0
  184. package/dist/catalog/chat-models/qw/qwen3.6-35b-a3b-q4/versions/1.0.3/manifest.json +117 -0
  185. package/dist/catalog/chat-models/qw/qwen3.6-35b-a3b-q8/manifest.json +252 -0
  186. package/dist/catalog/chat-models/qw/qwen3.6-35b-a3b-q8/versions/1.0.0/manifest.json +112 -0
  187. package/dist/catalog/chat-models/qw/qwen3.6-35b-a3b-q8/versions/1.0.1/manifest.json +113 -0
  188. package/dist/catalog/chat-models/qw/qwen3.6-35b-a3b-q8/versions/1.0.2/manifest.json +113 -0
  189. package/dist/catalog/chat-models/qw/qwen3.6-35b-a3b-q8/versions/1.0.3/manifest.json +137 -0
  190. package/dist/catalog/chat-models/qw/qwen3.8-27b-iq1-s/manifest.json +195 -0
  191. package/dist/catalog/chat-models/qw/qwen3.8-27b-iq1-s/versions/1.0.0/manifest.json +19 -0
  192. package/dist/catalog/chat-models/qw/qwen3.8-27b-q2/manifest.json +194 -0
  193. package/dist/catalog/chat-models/qw/qwen3.8-27b-q2/versions/1.0.0/manifest.json +19 -0
  194. package/dist/catalog/chat-models/qw/qwen3.8-27b-q3/manifest.json +194 -0
  195. package/dist/catalog/chat-models/qw/qwen3.8-27b-q3/versions/1.0.0/manifest.json +19 -0
  196. package/dist/catalog/chat-models/qw/qwen3.8-27b-q3/versions/1.0.1/manifest.json +19 -0
  197. package/dist/catalog/chat-models/qw/qwen3.8-27b-q4/manifest.json +291 -0
  198. package/dist/catalog/chat-models/qw/qwen3.8-27b-q4/versions/1.0.0/manifest.json +20 -0
  199. package/dist/catalog/chat-models/qw/qwen3.8-27b-q4/versions/1.0.1/manifest.json +93 -0
  200. package/dist/catalog/chat-models/qw/qwen3.8-27b-q4/versions/1.0.2/manifest.json +92 -0
  201. package/dist/catalog/chat-models/qw/qwen3.8-27b-q4/versions/1.0.3/manifest.json +116 -0
  202. package/dist/catalog/chat-models/qw/qwen3.8-27b-q4/versions/1.0.4/manifest.json +116 -0
  203. package/dist/catalog/chat-models/qw/qwen3.8-27b-q4/versions/1.0.5/manifest.json +116 -0
  204. package/dist/catalog/chat-models/qw/qwen3.8-27b-q6/manifest.json +278 -0
  205. package/dist/catalog/chat-models/qw/qwen3.8-27b-q6/versions/1.0.0/manifest.json +20 -0
  206. package/dist/catalog/chat-models/qw/qwen3.8-27b-q6/versions/1.0.1/manifest.json +104 -0
  207. package/dist/catalog/chat-models/qw/qwen3.8-27b-q6/versions/1.0.2/manifest.json +103 -0
  208. package/dist/catalog/chat-models/qw/qwen3.8-27b-q6/versions/1.0.3/manifest.json +127 -0
  209. package/dist/catalog/chat-models/qw/qwen3.8-27b-q8/manifest.json +278 -0
  210. package/dist/catalog/chat-models/qw/qwen3.8-27b-q8/versions/1.0.0/manifest.json +113 -0
  211. package/dist/catalog/chat-models/qw/qwen3.8-27b-q8/versions/1.0.1/manifest.json +113 -0
  212. package/dist/catalog/chat-models/qw/qwen3.8-27b-q8/versions/1.0.2/manifest.json +137 -0
  213. package/dist/catalog/chat-models/qw/qwen3.8-flash-next-iq3/manifest.json +221 -0
  214. package/dist/catalog/chat-models/qw/qwen3.8-flash-next-iq3/versions/1.0.0/manifest.json +36 -0
  215. package/dist/catalog/chat-models/qw/qwen3.8-flash-next-iq4/manifest.json +429 -0
  216. package/dist/catalog/chat-models/qw/qwen3.8-flash-next-iq4/versions/1.0.0/manifest.json +36 -0
  217. package/dist/catalog/chat-models/qw/qwen3.8-flash-next-iq4/versions/1.1.0/manifest.json +205 -0
  218. package/dist/catalog/chat-models/qw/qwen3.8-flash-next-iq4/versions/1.1.1/manifest.json +244 -0
  219. package/dist/catalog/chat-models/qw/qwen3.8-flash-next-q2/manifest.json +203 -0
  220. package/dist/catalog/chat-models/qw/qwen3.8-flash-next-q2/versions/1.0.0/manifest.json +21 -0
  221. package/dist/catalog/chat-models/qw/qwen3.8-flash-next-q2/versions/1.0.1/manifest.json +31 -0
  222. package/dist/catalog/chat-models/qw/qwen3.8-flash-next-q2/versions/1.0.2/manifest.json +30 -0
  223. package/dist/catalog/chat-models/raw-index.json +1 -0
  224. package/dist/catalog/chat-models/ta/talkie-1930-13b-q4/manifest.json +79 -0
  225. package/dist/catalog/chat-models/ta/talkie-1930-13b-q4/versions/1.0.0/manifest.json +15 -0
  226. package/dist/catalog/chat-models/ta/talkie-1930-13b-q4/versions/1.0.1/manifest.json +15 -0
  227. package/dist/catalog/chat-models/te/ternary-bonsai-27b-q2/manifest.json +204 -0
  228. package/dist/catalog/chat-models/te/ternary-bonsai-27b-q2/versions/1.0.0/manifest.json +70 -0
  229. package/dist/catalog/chat-models/te/ternary-bonsai-27b-q2/versions/1.0.1/manifest.json +70 -0
  230. package/dist/catalog/chat-models/un/underdog-saluki-27b-iq2-mix/manifest.json +164 -0
  231. package/dist/catalog/chat-models/un/underdog-saluki-27b-iq2-mix/versions/1.0.0/manifest.json +14 -0
  232. package/dist/catalog/knowledge-catalogs/be/bendyline-how-to-books/manifest.json +26 -0
  233. package/dist/catalog/knowledge-catalogs/be/bendyline-how-to-books/versions/2026.10.8/manifest.json +8950 -0
  234. package/dist/catalog/knowledge-catalogs/ca/caselaw-alaska-part-0001/manifest.json +24 -0
  235. package/dist/catalog/knowledge-catalogs/ca/caselaw-alaska-part-0001/versions/2026.10.1/manifest.json +658 -0
  236. package/dist/catalog/knowledge-catalogs/ca/caselaw-nevada-part-0001/manifest.json +24 -0
  237. package/dist/catalog/knowledge-catalogs/ca/caselaw-nevada-part-0001/versions/2026.10.1/manifest.json +666 -0
  238. package/dist/catalog/knowledge-catalogs/ca/caselaw-us-300-399-part-0001/manifest.json +24 -0
  239. package/dist/catalog/knowledge-catalogs/ca/caselaw-us-300-399-part-0001/versions/2026.10.2/manifest.json +182 -0
  240. package/dist/catalog/knowledge-catalogs/ca/caselaw-us-347-349/manifest.json +25 -0
  241. package/dist/catalog/knowledge-catalogs/ca/caselaw-us-347-349/versions/2026.10.5/manifest.json +38 -0
  242. package/dist/catalog/knowledge-catalogs/ca/caselaw-utah-part-0001/manifest.json +24 -0
  243. package/dist/catalog/knowledge-catalogs/ca/caselaw-utah-part-0001/versions/2026.10.1/manifest.json +558 -0
  244. package/dist/catalog/knowledge-catalogs/ca/caselaw-utah-part-0002/manifest.json +24 -0
  245. package/dist/catalog/knowledge-catalogs/ca/caselaw-utah-part-0002/versions/2026.10.1/manifest.json +258 -0
  246. package/dist/catalog/knowledge-catalogs/ca/caselaw-washington-part-0001/manifest.json +24 -0
  247. package/dist/catalog/knowledge-catalogs/ca/caselaw-washington-part-0001/versions/2026.10.1/manifest.json +318 -0
  248. package/dist/catalog/knowledge-catalogs/ca/caselaw-washington-part-0002/manifest.json +24 -0
  249. package/dist/catalog/knowledge-catalogs/ca/caselaw-washington-part-0002/versions/2026.10.1/manifest.json +266 -0
  250. package/dist/catalog/knowledge-catalogs/ca/caselaw-washington-part-0003/manifest.json +24 -0
  251. package/dist/catalog/knowledge-catalogs/ca/caselaw-washington-part-0003/versions/2026.10.1/manifest.json +186 -0
  252. package/dist/catalog/knowledge-catalogs/ca/caselaw-washington-part-0004/manifest.json +24 -0
  253. package/dist/catalog/knowledge-catalogs/ca/caselaw-washington-part-0004/versions/2026.10.1/manifest.json +138 -0
  254. package/dist/catalog/knowledge-catalogs/ca/caselaw-washington-part-0005/manifest.json +24 -0
  255. package/dist/catalog/knowledge-catalogs/ca/caselaw-washington-part-0005/versions/2026.10.1/manifest.json +138 -0
  256. package/dist/catalog/knowledge-catalogs/ca/caselaw-wyoming-part-0001/manifest.json +24 -0
  257. package/dist/catalog/knowledge-catalogs/ca/caselaw-wyoming-part-0001/versions/2026.10.3/manifest.json +626 -0
  258. package/dist/catalog/knowledge-catalogs/index.json +1 -0
  259. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-aspnet-core-en/manifest.json +22 -0
  260. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-aspnet-core-en/versions/2026.10.16/manifest.json +18270 -0
  261. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-azure-ai-platform-en/manifest.json +23 -0
  262. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-azure-ai-platform-en/versions/2026.10.16/manifest.json +1098 -0
  263. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-azure-apps-compute-en/manifest.json +23 -0
  264. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-azure-apps-compute-en/versions/2026.10.16/manifest.json +722 -0
  265. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-azure-data-aiot-en/manifest.json +23 -0
  266. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-azure-data-aiot-en/versions/2026.10.16/manifest.json +530 -0
  267. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-azure-management-hybrid-en/manifest.json +23 -0
  268. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-azure-management-hybrid-en/versions/2026.10.16/manifest.json +602 -0
  269. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-azure-network-security-en/manifest.json +23 -0
  270. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-azure-network-security-en/versions/2026.10.16/manifest.json +250 -0
  271. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-dotnet-en/manifest.json +22 -0
  272. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-dotnet-en/versions/2026.10.16/manifest.json +16794 -0
  273. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-entity-framework-en/manifest.json +22 -0
  274. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-entity-framework-en/versions/2026.10.16/manifest.json +1106 -0
  275. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-microsoft-365-en/manifest.json +22 -0
  276. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-microsoft-365-en/versions/2026.10.16/manifest.json +234 -0
  277. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-powershell-en/manifest.json +22 -0
  278. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-powershell-en/versions/2026.10.16/manifest.json +522 -0
  279. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-sql-server-en/manifest.json +22 -0
  280. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-sql-server-en/versions/2026.10.16/manifest.json +1850 -0
  281. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-windows-apps-en/manifest.json +22 -0
  282. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-windows-apps-en/versions/2026.10.16/manifest.json +1734 -0
  283. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-windows-server-en/manifest.json +22 -0
  284. package/dist/catalog/knowledge-catalogs/mi/microsoftdocs-windows-server-en/versions/2026.10.16/manifest.json +774 -0
  285. package/dist/catalog/knowledge-catalogs/raw-index.json +1 -0
  286. package/dist/catalog/knowledge-catalogs/re/regional-american-south/art.json +12 -0
  287. package/dist/catalog/knowledge-catalogs/re/regional-american-south/logo.webp +0 -0
  288. package/dist/catalog/knowledge-catalogs/re/regional-american-south/manifest.json +48 -0
  289. package/dist/catalog/knowledge-catalogs/re/regional-american-south/versions/2026.10.3/manifest.json +133 -0
  290. package/dist/catalog/knowledge-catalogs/re/regional-andes-amazon/art.json +12 -0
  291. package/dist/catalog/knowledge-catalogs/re/regional-andes-amazon/logo.webp +0 -0
  292. package/dist/catalog/knowledge-catalogs/re/regional-andes-amazon/manifest.json +48 -0
  293. package/dist/catalog/knowledge-catalogs/re/regional-andes-amazon/versions/2026.10.3/manifest.json +133 -0
  294. package/dist/catalog/knowledge-catalogs/re/regional-australia-nz/art.json +12 -0
  295. package/dist/catalog/knowledge-catalogs/re/regional-australia-nz/logo.webp +0 -0
  296. package/dist/catalog/knowledge-catalogs/re/regional-australia-nz/manifest.json +48 -0
  297. package/dist/catalog/knowledge-catalogs/re/regional-australia-nz/versions/2026.10.3/manifest.json +133 -0
  298. package/dist/catalog/knowledge-catalogs/re/regional-brazil/art.json +12 -0
  299. package/dist/catalog/knowledge-catalogs/re/regional-brazil/logo.webp +0 -0
  300. package/dist/catalog/knowledge-catalogs/re/regional-brazil/manifest.json +48 -0
  301. package/dist/catalog/knowledge-catalogs/re/regional-brazil/versions/2026.10.3/manifest.json +133 -0
  302. package/dist/catalog/knowledge-catalogs/re/regional-britain-ireland/art.json +12 -0
  303. package/dist/catalog/knowledge-catalogs/re/regional-britain-ireland/logo.webp +0 -0
  304. package/dist/catalog/knowledge-catalogs/re/regional-britain-ireland/manifest.json +48 -0
  305. package/dist/catalog/knowledge-catalogs/re/regional-britain-ireland/versions/2026.10.3/manifest.json +133 -0
  306. package/dist/catalog/knowledge-catalogs/re/regional-california-southwest/art.json +12 -0
  307. package/dist/catalog/knowledge-catalogs/re/regional-california-southwest/logo.webp +0 -0
  308. package/dist/catalog/knowledge-catalogs/re/regional-california-southwest/manifest.json +48 -0
  309. package/dist/catalog/knowledge-catalogs/re/regional-california-southwest/versions/2026.10.3/manifest.json +133 -0
  310. package/dist/catalog/knowledge-catalogs/re/regional-canada-alaska/art.json +12 -0
  311. package/dist/catalog/knowledge-catalogs/re/regional-canada-alaska/logo.webp +0 -0
  312. package/dist/catalog/knowledge-catalogs/re/regional-canada-alaska/manifest.json +48 -0
  313. package/dist/catalog/knowledge-catalogs/re/regional-canada-alaska/versions/2026.10.3/manifest.json +133 -0
  314. package/dist/catalog/knowledge-catalogs/re/regional-central-eastern-europe/art.json +12 -0
  315. package/dist/catalog/knowledge-catalogs/re/regional-central-eastern-europe/logo.webp +0 -0
  316. package/dist/catalog/knowledge-catalogs/re/regional-central-eastern-europe/manifest.json +48 -0
  317. package/dist/catalog/knowledge-catalogs/re/regional-central-eastern-europe/versions/2026.10.3/manifest.json +133 -0
  318. package/dist/catalog/knowledge-catalogs/re/regional-east-asia/art.json +12 -0
  319. package/dist/catalog/knowledge-catalogs/re/regional-east-asia/logo.webp +0 -0
  320. package/dist/catalog/knowledge-catalogs/re/regional-east-asia/manifest.json +48 -0
  321. package/dist/catalog/knowledge-catalogs/re/regional-east-asia/versions/2026.10.3/manifest.json +133 -0
  322. package/dist/catalog/knowledge-catalogs/re/regional-east-southern-africa/art.json +12 -0
  323. package/dist/catalog/knowledge-catalogs/re/regional-east-southern-africa/logo.webp +0 -0
  324. package/dist/catalog/knowledge-catalogs/re/regional-east-southern-africa/manifest.json +48 -0
  325. package/dist/catalog/knowledge-catalogs/re/regional-east-southern-africa/versions/2026.10.3/manifest.json +133 -0
  326. package/dist/catalog/knowledge-catalogs/re/regional-mediterranean/art.json +12 -0
  327. package/dist/catalog/knowledge-catalogs/re/regional-mediterranean/logo.webp +0 -0
  328. package/dist/catalog/knowledge-catalogs/re/regional-mediterranean/manifest.json +48 -0
  329. package/dist/catalog/knowledge-catalogs/re/regional-mediterranean/versions/2026.10.3/manifest.json +133 -0
  330. package/dist/catalog/knowledge-catalogs/re/regional-mexico-central-america/art.json +12 -0
  331. package/dist/catalog/knowledge-catalogs/re/regional-mexico-central-america/logo.webp +0 -0
  332. package/dist/catalog/knowledge-catalogs/re/regional-mexico-central-america/manifest.json +48 -0
  333. package/dist/catalog/knowledge-catalogs/re/regional-mexico-central-america/versions/2026.10.3/manifest.json +133 -0
  334. package/dist/catalog/knowledge-catalogs/re/regional-middle-east/art.json +12 -0
  335. package/dist/catalog/knowledge-catalogs/re/regional-middle-east/logo.webp +0 -0
  336. package/dist/catalog/knowledge-catalogs/re/regional-middle-east/manifest.json +48 -0
  337. package/dist/catalog/knowledge-catalogs/re/regional-middle-east/versions/2026.10.3/manifest.json +133 -0
  338. package/dist/catalog/knowledge-catalogs/re/regional-midwest/art.json +12 -0
  339. package/dist/catalog/knowledge-catalogs/re/regional-midwest/logo.webp +0 -0
  340. package/dist/catalog/knowledge-catalogs/re/regional-midwest/manifest.json +48 -0
  341. package/dist/catalog/knowledge-catalogs/re/regional-midwest/versions/2026.10.3/manifest.json +133 -0
  342. package/dist/catalog/knowledge-catalogs/re/regional-north-africa/art.json +12 -0
  343. package/dist/catalog/knowledge-catalogs/re/regional-north-africa/logo.webp +0 -0
  344. package/dist/catalog/knowledge-catalogs/re/regional-north-africa/manifest.json +48 -0
  345. package/dist/catalog/knowledge-catalogs/re/regional-north-africa/versions/2026.10.3/manifest.json +133 -0
  346. package/dist/catalog/knowledge-catalogs/re/regional-northeast/art.json +12 -0
  347. package/dist/catalog/knowledge-catalogs/re/regional-northeast/logo.webp +0 -0
  348. package/dist/catalog/knowledge-catalogs/re/regional-northeast/manifest.json +48 -0
  349. package/dist/catalog/knowledge-catalogs/re/regional-northeast/versions/2026.10.3/manifest.json +133 -0
  350. package/dist/catalog/knowledge-catalogs/re/regional-northern-eurasia/art.json +12 -0
  351. package/dist/catalog/knowledge-catalogs/re/regional-northern-eurasia/logo.webp +0 -0
  352. package/dist/catalog/knowledge-catalogs/re/regional-northern-eurasia/manifest.json +48 -0
  353. package/dist/catalog/knowledge-catalogs/re/regional-northern-eurasia/versions/2026.10.3/manifest.json +133 -0
  354. package/dist/catalog/knowledge-catalogs/re/regional-northern-europe/art.json +12 -0
  355. package/dist/catalog/knowledge-catalogs/re/regional-northern-europe/logo.webp +0 -0
  356. package/dist/catalog/knowledge-catalogs/re/regional-northern-europe/manifest.json +48 -0
  357. package/dist/catalog/knowledge-catalogs/re/regional-northern-europe/versions/2026.10.3/manifest.json +133 -0
  358. package/dist/catalog/knowledge-catalogs/re/regional-pacific-islands/art.json +12 -0
  359. package/dist/catalog/knowledge-catalogs/re/regional-pacific-islands/logo.webp +0 -0
  360. package/dist/catalog/knowledge-catalogs/re/regional-pacific-islands/manifest.json +48 -0
  361. package/dist/catalog/knowledge-catalogs/re/regional-pacific-islands/versions/2026.10.3/manifest.json +133 -0
  362. package/dist/catalog/knowledge-catalogs/re/regional-pacific-northwest/art.json +12 -0
  363. package/dist/catalog/knowledge-catalogs/re/regional-pacific-northwest/logo.webp +0 -0
  364. package/dist/catalog/knowledge-catalogs/re/regional-pacific-northwest/manifest.json +48 -0
  365. package/dist/catalog/knowledge-catalogs/re/regional-pacific-northwest/versions/2026.10.3/manifest.json +133 -0
  366. package/dist/catalog/knowledge-catalogs/re/regional-rockies-plains/art.json +12 -0
  367. package/dist/catalog/knowledge-catalogs/re/regional-rockies-plains/logo.webp +0 -0
  368. package/dist/catalog/knowledge-catalogs/re/regional-rockies-plains/manifest.json +48 -0
  369. package/dist/catalog/knowledge-catalogs/re/regional-rockies-plains/versions/2026.10.3/manifest.json +133 -0
  370. package/dist/catalog/knowledge-catalogs/re/regional-south-asia/art.json +12 -0
  371. package/dist/catalog/knowledge-catalogs/re/regional-south-asia/logo.webp +0 -0
  372. package/dist/catalog/knowledge-catalogs/re/regional-south-asia/manifest.json +48 -0
  373. package/dist/catalog/knowledge-catalogs/re/regional-south-asia/versions/2026.10.3/manifest.json +133 -0
  374. package/dist/catalog/knowledge-catalogs/re/regional-southeast-asia/art.json +12 -0
  375. package/dist/catalog/knowledge-catalogs/re/regional-southeast-asia/logo.webp +0 -0
  376. package/dist/catalog/knowledge-catalogs/re/regional-southeast-asia/manifest.json +48 -0
  377. package/dist/catalog/knowledge-catalogs/re/regional-southeast-asia/versions/2026.10.3/manifest.json +133 -0
  378. package/dist/catalog/knowledge-catalogs/re/regional-southern-cone/art.json +12 -0
  379. package/dist/catalog/knowledge-catalogs/re/regional-southern-cone/logo.webp +0 -0
  380. package/dist/catalog/knowledge-catalogs/re/regional-southern-cone/manifest.json +48 -0
  381. package/dist/catalog/knowledge-catalogs/re/regional-southern-cone/versions/2026.10.3/manifest.json +133 -0
  382. package/dist/catalog/knowledge-catalogs/re/regional-west-central-africa/art.json +12 -0
  383. package/dist/catalog/knowledge-catalogs/re/regional-west-central-africa/logo.webp +0 -0
  384. package/dist/catalog/knowledge-catalogs/re/regional-west-central-africa/manifest.json +48 -0
  385. package/dist/catalog/knowledge-catalogs/re/regional-west-central-africa/versions/2026.10.3/manifest.json +133 -0
  386. package/dist/catalog/knowledge-catalogs/re/regional-western-europe/art.json +12 -0
  387. package/dist/catalog/knowledge-catalogs/re/regional-western-europe/logo.webp +0 -0
  388. package/dist/catalog/knowledge-catalogs/re/regional-western-europe/manifest.json +48 -0
  389. package/dist/catalog/knowledge-catalogs/re/regional-western-europe/versions/2026.10.3/manifest.json +133 -0
  390. package/dist/catalog/knowledge-catalogs/wi/wikipedia-arts/art.json +12 -0
  391. package/dist/catalog/knowledge-catalogs/wi/wikipedia-arts/logo.webp +0 -0
  392. package/dist/catalog/knowledge-catalogs/wi/wikipedia-arts/manifest.json +28 -0
  393. package/dist/catalog/knowledge-catalogs/wi/wikipedia-arts/versions/2026.4.2/manifest.json +53 -0
  394. package/dist/catalog/knowledge-catalogs/wi/wikipedia-arts/versions/2026.4.3/manifest.json +885 -0
  395. package/dist/catalog/knowledge-catalogs/wi/wikipedia-arts/versions/2026.4.5/manifest.json +157 -0
  396. package/dist/catalog/knowledge-catalogs/wi/wikipedia-astronomy/art.json +11 -0
  397. package/dist/catalog/knowledge-catalogs/wi/wikipedia-astronomy/logo.webp +0 -0
  398. package/dist/catalog/knowledge-catalogs/wi/wikipedia-astronomy/manifest.json +32 -0
  399. package/dist/catalog/knowledge-catalogs/wi/wikipedia-astronomy/versions/2026.4.2/manifest.json +69 -0
  400. package/dist/catalog/knowledge-catalogs/wi/wikipedia-astronomy/versions/2026.4.3/manifest.json +741 -0
  401. package/dist/catalog/knowledge-catalogs/wi/wikipedia-astronomy/versions/2026.4.5/manifest.json +173 -0
  402. package/dist/catalog/knowledge-catalogs/wi/wikipedia-biology/art.json +12 -0
  403. package/dist/catalog/knowledge-catalogs/wi/wikipedia-biology/logo.webp +0 -0
  404. package/dist/catalog/knowledge-catalogs/wi/wikipedia-biology/manifest.json +69 -0
  405. package/dist/catalog/knowledge-catalogs/wi/wikipedia-biology/versions/2026.4.4/manifest.json +217 -0
  406. package/dist/catalog/knowledge-catalogs/wi/wikipedia-biology/versions/2026.4.5/manifest.json +217 -0
  407. package/dist/catalog/knowledge-catalogs/wi/wikipedia-business/art.json +12 -0
  408. package/dist/catalog/knowledge-catalogs/wi/wikipedia-business/logo.webp +0 -0
  409. package/dist/catalog/knowledge-catalogs/wi/wikipedia-business/manifest.json +64 -0
  410. package/dist/catalog/knowledge-catalogs/wi/wikipedia-business/versions/2026.4.4/manifest.json +197 -0
  411. package/dist/catalog/knowledge-catalogs/wi/wikipedia-business/versions/2026.4.5/manifest.json +197 -0
  412. package/dist/catalog/knowledge-catalogs/wi/wikipedia-film-tv/art.json +12 -0
  413. package/dist/catalog/knowledge-catalogs/wi/wikipedia-film-tv/logo.webp +0 -0
  414. package/dist/catalog/knowledge-catalogs/wi/wikipedia-film-tv/manifest.json +64 -0
  415. package/dist/catalog/knowledge-catalogs/wi/wikipedia-film-tv/versions/2026.4.5/manifest.json +197 -0
  416. package/dist/catalog/knowledge-catalogs/wi/wikipedia-food-drink/art.json +11 -0
  417. package/dist/catalog/knowledge-catalogs/wi/wikipedia-food-drink/logo.webp +0 -0
  418. package/dist/catalog/knowledge-catalogs/wi/wikipedia-food-drink/manifest.json +29 -0
  419. package/dist/catalog/knowledge-catalogs/wi/wikipedia-food-drink/versions/2026.4.2/manifest.json +57 -0
  420. package/dist/catalog/knowledge-catalogs/wi/wikipedia-food-drink/versions/2026.4.3/manifest.json +925 -0
  421. package/dist/catalog/knowledge-catalogs/wi/wikipedia-food-drink/versions/2026.4.5/manifest.json +185 -0
  422. package/dist/catalog/knowledge-catalogs/wi/wikipedia-games/art.json +12 -0
  423. package/dist/catalog/knowledge-catalogs/wi/wikipedia-games/logo.webp +0 -0
  424. package/dist/catalog/knowledge-catalogs/wi/wikipedia-games/manifest.json +28 -0
  425. package/dist/catalog/knowledge-catalogs/wi/wikipedia-games/versions/2026.4.2/manifest.json +53 -0
  426. package/dist/catalog/knowledge-catalogs/wi/wikipedia-games/versions/2026.4.3/manifest.json +1665 -0
  427. package/dist/catalog/knowledge-catalogs/wi/wikipedia-games/versions/2026.4.5/manifest.json +133 -0
  428. package/dist/catalog/knowledge-catalogs/wi/wikipedia-language/art.json +12 -0
  429. package/dist/catalog/knowledge-catalogs/wi/wikipedia-language/logo.webp +0 -0
  430. package/dist/catalog/knowledge-catalogs/wi/wikipedia-language/manifest.json +29 -0
  431. package/dist/catalog/knowledge-catalogs/wi/wikipedia-language/versions/2026.4.2/manifest.json +57 -0
  432. package/dist/catalog/knowledge-catalogs/wi/wikipedia-language/versions/2026.4.3/manifest.json +1233 -0
  433. package/dist/catalog/knowledge-catalogs/wi/wikipedia-language/versions/2026.4.5/manifest.json +141 -0
  434. package/dist/catalog/knowledge-catalogs/wi/wikipedia-law/art.json +11 -0
  435. package/dist/catalog/knowledge-catalogs/wi/wikipedia-law/logo.webp +0 -0
  436. package/dist/catalog/knowledge-catalogs/wi/wikipedia-law/manifest.json +29 -0
  437. package/dist/catalog/knowledge-catalogs/wi/wikipedia-law/versions/2026.4.2/manifest.json +57 -0
  438. package/dist/catalog/knowledge-catalogs/wi/wikipedia-law/versions/2026.4.3/manifest.json +493 -0
  439. package/dist/catalog/knowledge-catalogs/wi/wikipedia-law/versions/2026.4.5/manifest.json +185 -0
  440. package/dist/catalog/knowledge-catalogs/wi/wikipedia-literature/art.json +12 -0
  441. package/dist/catalog/knowledge-catalogs/wi/wikipedia-literature/logo.webp +0 -0
  442. package/dist/catalog/knowledge-catalogs/wi/wikipedia-literature/manifest.json +63 -0
  443. package/dist/catalog/knowledge-catalogs/wi/wikipedia-literature/versions/2026.4.4/manifest.json +193 -0
  444. package/dist/catalog/knowledge-catalogs/wi/wikipedia-literature/versions/2026.4.5/manifest.json +193 -0
  445. package/dist/catalog/knowledge-catalogs/wi/wikipedia-mathematics/art.json +12 -0
  446. package/dist/catalog/knowledge-catalogs/wi/wikipedia-mathematics/logo.webp +0 -0
  447. package/dist/catalog/knowledge-catalogs/wi/wikipedia-mathematics/manifest.json +30 -0
  448. package/dist/catalog/knowledge-catalogs/wi/wikipedia-mathematics/versions/2026.4.2/manifest.json +61 -0
  449. package/dist/catalog/knowledge-catalogs/wi/wikipedia-mathematics/versions/2026.4.3/manifest.json +665 -0
  450. package/dist/catalog/knowledge-catalogs/wi/wikipedia-mathematics/versions/2026.4.5/manifest.json +177 -0
  451. package/dist/catalog/knowledge-catalogs/wi/wikipedia-medicine/art.json +12 -0
  452. package/dist/catalog/knowledge-catalogs/wi/wikipedia-medicine/logo.webp +0 -0
  453. package/dist/catalog/knowledge-catalogs/wi/wikipedia-medicine/manifest.json +29 -0
  454. package/dist/catalog/knowledge-catalogs/wi/wikipedia-medicine/versions/2026.4.2/manifest.json +57 -0
  455. package/dist/catalog/knowledge-catalogs/wi/wikipedia-medicine/versions/2026.4.3/manifest.json +745 -0
  456. package/dist/catalog/knowledge-catalogs/wi/wikipedia-medicine/versions/2026.4.5/manifest.json +185 -0
  457. package/dist/catalog/knowledge-catalogs/wi/wikipedia-military/art.json +12 -0
  458. package/dist/catalog/knowledge-catalogs/wi/wikipedia-military/logo.webp +0 -0
  459. package/dist/catalog/knowledge-catalogs/wi/wikipedia-military/manifest.json +62 -0
  460. package/dist/catalog/knowledge-catalogs/wi/wikipedia-military/versions/2026.4.4/manifest.json +189 -0
  461. package/dist/catalog/knowledge-catalogs/wi/wikipedia-military/versions/2026.4.5/manifest.json +189 -0
  462. package/dist/catalog/knowledge-catalogs/wi/wikipedia-music/art.json +12 -0
  463. package/dist/catalog/knowledge-catalogs/wi/wikipedia-music/logo.webp +0 -0
  464. package/dist/catalog/knowledge-catalogs/wi/wikipedia-music/manifest.json +64 -0
  465. package/dist/catalog/knowledge-catalogs/wi/wikipedia-music/versions/2026.4.5/manifest.json +197 -0
  466. package/dist/catalog/knowledge-catalogs/wi/wikipedia-politics/art.json +12 -0
  467. package/dist/catalog/knowledge-catalogs/wi/wikipedia-politics/logo.webp +0 -0
  468. package/dist/catalog/knowledge-catalogs/wi/wikipedia-politics/manifest.json +64 -0
  469. package/dist/catalog/knowledge-catalogs/wi/wikipedia-politics/versions/2026.4.4/manifest.json +197 -0
  470. package/dist/catalog/knowledge-catalogs/wi/wikipedia-politics/versions/2026.4.5/manifest.json +197 -0
  471. package/dist/catalog/knowledge-catalogs/wi/wikipedia-religion-philosophy/art.json +12 -0
  472. package/dist/catalog/knowledge-catalogs/wi/wikipedia-religion-philosophy/logo.webp +0 -0
  473. package/dist/catalog/knowledge-catalogs/wi/wikipedia-religion-philosophy/manifest.json +30 -0
  474. package/dist/catalog/knowledge-catalogs/wi/wikipedia-religion-philosophy/versions/2026.4.2/manifest.json +61 -0
  475. package/dist/catalog/knowledge-catalogs/wi/wikipedia-religion-philosophy/versions/2026.4.3/manifest.json +1253 -0
  476. package/dist/catalog/knowledge-catalogs/wi/wikipedia-religion-philosophy/versions/2026.4.5/manifest.json +189 -0
  477. package/dist/catalog/knowledge-catalogs/wi/wikipedia-science/art.json +12 -0
  478. package/dist/catalog/knowledge-catalogs/wi/wikipedia-science/logo.webp +0 -0
  479. package/dist/catalog/knowledge-catalogs/wi/wikipedia-science/manifest.json +29 -0
  480. package/dist/catalog/knowledge-catalogs/wi/wikipedia-science/versions/2026.4.2/manifest.json +57 -0
  481. package/dist/catalog/knowledge-catalogs/wi/wikipedia-science/versions/2026.4.3/manifest.json +1957 -0
  482. package/dist/catalog/knowledge-catalogs/wi/wikipedia-science/versions/2026.4.5/manifest.json +185 -0
  483. package/dist/catalog/knowledge-catalogs/wi/wikipedia-sports/art.json +12 -0
  484. package/dist/catalog/knowledge-catalogs/wi/wikipedia-sports/logo.webp +0 -0
  485. package/dist/catalog/knowledge-catalogs/wi/wikipedia-sports/manifest.json +63 -0
  486. package/dist/catalog/knowledge-catalogs/wi/wikipedia-sports/versions/2026.4.5/manifest.json +193 -0
  487. package/dist/catalog/knowledge-catalogs/wi/wikipedia-technology/art.json +11 -0
  488. package/dist/catalog/knowledge-catalogs/wi/wikipedia-technology/logo.webp +0 -0
  489. package/dist/catalog/knowledge-catalogs/wi/wikipedia-technology/manifest.json +45 -0
  490. package/dist/catalog/knowledge-catalogs/wi/wikipedia-technology/versions/2026.4.4/manifest.json +121 -0
  491. package/dist/catalog/knowledge-catalogs/wi/wikipedia-technology/versions/2026.4.5/manifest.json +173 -0
  492. package/dist/catalog/knowledge-catalogs/wi/wikipedia-transportation/art.json +12 -0
  493. package/dist/catalog/knowledge-catalogs/wi/wikipedia-transportation/logo.webp +0 -0
  494. package/dist/catalog/knowledge-catalogs/wi/wikipedia-transportation/manifest.json +62 -0
  495. package/dist/catalog/knowledge-catalogs/wi/wikipedia-transportation/versions/2026.4.4/manifest.json +189 -0
  496. package/dist/catalog/knowledge-catalogs/wi/wikipedia-transportation/versions/2026.4.5/manifest.json +189 -0
  497. package/dist/catalog/provenance.json +8 -0
  498. package/dist/chat/catalog-model-lookup.d.ts +21 -0
  499. package/dist/chat/catalog-model-lookup.js +48 -0
  500. package/dist/chat/local-model-tier.d.ts +1 -0
  501. package/dist/chat/local-model-tier.js +2 -0
  502. package/dist/chat/vision-capability.d.ts +1 -0
  503. package/dist/chat/vision-capability.js +3 -0
  504. package/dist/engine-components.d.ts +30 -0
  505. package/dist/engine-components.js +144 -0
  506. package/dist/engines/electron-native-reuse.d.ts +65 -0
  507. package/dist/engines/electron-native-reuse.js +305 -0
  508. package/dist/engines/native-file-manifest.json +1320 -0
  509. package/dist/engines/native-manifest.d.ts +70 -0
  510. package/dist/engines/native-manifest.js +98 -0
  511. package/dist/engines/registry.d.ts +88 -0
  512. package/dist/engines/registry.js +217 -0
  513. package/dist/engines/resolver.d.ts +121 -0
  514. package/dist/engines/resolver.js +894 -0
  515. package/dist/engines/signature.d.ts +85 -0
  516. package/dist/engines/signature.js +231 -0
  517. package/dist/fs/atomic.d.ts +53 -0
  518. package/dist/fs/atomic.js +202 -0
  519. package/dist/fs/config-store.d.ts +14 -0
  520. package/dist/fs/config-store.js +101 -0
  521. package/dist/fs/home-permissions.d.ts +18 -0
  522. package/dist/fs/home-permissions.js +27 -0
  523. package/dist/fs/safe-paths.d.ts +107 -0
  524. package/dist/fs/safe-paths.js +264 -0
  525. package/dist/fs/security-json.d.ts +13 -0
  526. package/dist/fs/security-json.js +67 -0
  527. package/dist/http/auth-context.d.ts +23 -0
  528. package/dist/http/auth-context.js +1 -0
  529. package/dist/http/engine-runtime-context.d.ts +30 -0
  530. package/dist/http/engine-runtime-context.js +1 -0
  531. package/dist/http/external-conversation-contract.d.ts +49 -0
  532. package/dist/http/external-conversation-contract.js +1 -0
  533. package/dist/http/external-conversation-headers.d.ts +11 -0
  534. package/dist/http/external-conversation-headers.js +11 -0
  535. package/dist/http/inference-context.d.ts +36 -0
  536. package/dist/http/inference-context.js +1 -0
  537. package/dist/http/model-inventory-cache.d.ts +8 -0
  538. package/dist/http/model-inventory-cache.js +8 -0
  539. package/dist/http/openai-compat/action-ledger.d.ts +15 -0
  540. package/dist/http/openai-compat/action-ledger.js +119 -0
  541. package/dist/http/openai-compat/apple-model.d.ts +3 -0
  542. package/dist/http/openai-compat/apple-model.js +44 -0
  543. package/dist/http/openai-compat/caller-owned-profile.d.ts +7 -0
  544. package/dist/http/openai-compat/caller-owned-profile.js +18 -0
  545. package/dist/http/openai-compat/chat-progress.d.ts +9 -0
  546. package/dist/http/openai-compat/chat-progress.js +69 -0
  547. package/dist/http/openai-compat/chat-target.d.ts +76 -0
  548. package/dist/http/openai-compat/chat-target.js +97 -0
  549. package/dist/http/openai-compat/embeddings.d.ts +37 -0
  550. package/dist/http/openai-compat/embeddings.js +86 -0
  551. package/dist/http/openai-compat/request-tuning.d.ts +35 -0
  552. package/dist/http/openai-compat/request-tuning.js +87 -0
  553. package/dist/http/openai-compat/responses-streaming.d.ts +351 -0
  554. package/dist/http/openai-compat/responses-streaming.js +717 -0
  555. package/dist/http/openai-compat/responses-translate.d.ts +215 -0
  556. package/dist/http/openai-compat/responses-translate.js +377 -0
  557. package/dist/http/openai-compat/streaming-telemetry.d.ts +30 -0
  558. package/dist/http/openai-compat/streaming-telemetry.js +88 -0
  559. package/dist/http/openai-compat/streaming.d.ts +141 -0
  560. package/dist/http/openai-compat/streaming.js +687 -0
  561. package/dist/http/openai-compat/tool-call-stream-filter.d.ts +33 -0
  562. package/dist/http/openai-compat/tool-call-stream-filter.js +110 -0
  563. package/dist/http/openai-compat/translate.d.ts +254 -0
  564. package/dist/http/openai-compat/translate.js +485 -0
  565. package/dist/http/routes/engines.d.ts +13 -0
  566. package/dist/http/routes/engines.js +366 -0
  567. package/dist/http/routes/machine-engine-proxy.d.ts +26 -0
  568. package/dist/http/routes/machine-engine-proxy.js +43 -0
  569. package/dist/http/routes/model-bundles.d.ts +4 -0
  570. package/dist/http/routes/model-bundles.js +190 -0
  571. package/dist/http/routes/v1-chat.d.ts +72 -0
  572. package/dist/http/routes/v1-chat.js +645 -0
  573. package/dist/http/routes/v1-embeddings.d.ts +3 -0
  574. package/dist/http/routes/v1-embeddings.js +82 -0
  575. package/dist/http/routes/v1-knowledge.d.ts +4 -0
  576. package/dist/http/routes/v1-knowledge.js +304 -0
  577. package/dist/http/routes/v1-models-ensure.d.ts +32 -0
  578. package/dist/http/routes/v1-models-ensure.js +122 -0
  579. package/dist/http/routes/v1-models.d.ts +8 -0
  580. package/dist/http/routes/v1-models.js +189 -0
  581. package/dist/http/routes/v1-responses.d.ts +20 -0
  582. package/dist/http/routes/v1-responses.js +327 -0
  583. package/dist/index.d.ts +2 -0
  584. package/dist/index.js +2 -0
  585. package/dist/inference-service.d.ts +28 -0
  586. package/dist/inference-service.js +275 -0
  587. package/dist/knowledge/admission.d.ts +3 -0
  588. package/dist/knowledge/admission.js +22 -0
  589. package/dist/knowledge/bundled-catalog.d.ts +2 -0
  590. package/dist/knowledge/bundled-catalog.js +2 -0
  591. package/dist/knowledge/catalog-host.d.ts +103 -0
  592. package/dist/knowledge/catalog-host.js +180 -0
  593. package/dist/knowledge/catalog-source.d.ts +22 -0
  594. package/dist/knowledge/catalog-source.js +35 -0
  595. package/dist/knowledge/install.d.ts +55 -0
  596. package/dist/knowledge/install.js +274 -0
  597. package/dist/knowledge/manager.d.ts +287 -0
  598. package/dist/knowledge/manager.js +1540 -0
  599. package/dist/knowledge/registry.d.ts +35 -0
  600. package/dist/knowledge/registry.js +123 -0
  601. package/dist/knowledge/search-worker.d.ts +17 -0
  602. package/dist/knowledge/search-worker.js +36 -0
  603. package/dist/knowledge/shared-install-contract.d.ts +19 -0
  604. package/dist/knowledge/shared-install-contract.js +1 -0
  605. package/dist/knowledge/spatial-query.d.ts +11 -0
  606. package/dist/knowledge/spatial-query.js +33 -0
  607. package/dist/knowledge/trust-anchors.d.ts +11 -0
  608. package/dist/knowledge/trust-anchors.js +51 -0
  609. package/dist/knowledge/vector-floors.d.ts +44 -0
  610. package/dist/knowledge/vector-floors.js +92 -0
  611. package/dist/knowledge/worker-host.d.ts +9 -0
  612. package/dist/knowledge/worker-host.js +179 -0
  613. package/dist/kokoro-lexicon/LICENSE +202 -0
  614. package/dist/kokoro-lexicon/lexicon-gb-en.txt.gz +0 -0
  615. package/dist/kokoro-lexicon/lexicon-us-en.txt.gz +0 -0
  616. package/dist/kokoro-lexicon/manifest.json +23 -0
  617. package/dist/machine-engine/inspect-runtime.d.ts +3 -0
  618. package/dist/machine-engine/inspect-runtime.js +48 -0
  619. package/dist/machine-engine/knowledge-assets.d.ts +79 -0
  620. package/dist/machine-engine/knowledge-assets.js +505 -0
  621. package/dist/memory/embed-core.d.ts +131 -0
  622. package/dist/memory/embed-core.js +398 -0
  623. package/dist/memory/embed-worker.d.ts +23 -0
  624. package/dist/memory/embed-worker.js +49 -0
  625. package/dist/memory/embeddings.d.ts +105 -0
  626. package/dist/memory/embeddings.js +467 -0
  627. package/dist/model-profile/behaviors/tools-mlx-template-fix.d.ts +1 -0
  628. package/dist/model-profile/behaviors/tools-mlx-template-fix.js +2 -0
  629. package/dist/model-profile/behaviors/turn-ramble-detection.d.ts +1 -0
  630. package/dist/model-profile/behaviors/turn-ramble-detection.js +2 -0
  631. package/dist/model-profile/defaults.d.ts +1 -0
  632. package/dist/model-profile/defaults.js +2 -0
  633. package/dist/model-profile/index.d.ts +9 -0
  634. package/dist/model-profile/index.js +7 -0
  635. package/dist/model-profile/registry.d.ts +15 -0
  636. package/dist/model-profile/registry.js +18 -0
  637. package/dist/model-profile/runtime.d.ts +48 -0
  638. package/dist/model-profile/runtime.js +131 -0
  639. package/dist/model-profile/tool-grammar.d.ts +1 -0
  640. package/dist/model-profile/tool-grammar.js +2 -0
  641. package/dist/model-profile/tuning.d.ts +19 -0
  642. package/dist/model-profile/tuning.js +21 -0
  643. package/dist/model-profile/types.d.ts +1 -0
  644. package/dist/model-profile/types.js +2 -0
  645. package/dist/models/bundle-storage.d.ts +28 -0
  646. package/dist/models/bundle-storage.js +80 -0
  647. package/dist/models/catalog-drift.d.ts +48 -0
  648. package/dist/models/catalog-drift.js +94 -0
  649. package/dist/models/catalog-payload-identity.d.ts +107 -0
  650. package/dist/models/catalog-payload-identity.js +254 -0
  651. package/dist/models/ensure.d.ts +125 -0
  652. package/dist/models/ensure.js +205 -0
  653. package/dist/models/gezmodel.d.ts +47 -0
  654. package/dist/models/gezmodel.js +742 -0
  655. package/dist/models/install-registry.d.ts +101 -0
  656. package/dist/models/install-registry.js +216 -0
  657. package/dist/models/model-id.d.ts +22 -0
  658. package/dist/models/model-id.js +40 -0
  659. package/dist/models/storage-roots.d.ts +274 -0
  660. package/dist/models/storage-roots.js +865 -0
  661. package/dist/native-release-entry.d.ts +16 -0
  662. package/dist/native-release-entry.js +16 -0
  663. package/dist/perf/redact-path.d.ts +2 -0
  664. package/dist/perf/redact-path.js +27 -0
  665. package/dist/perf/responsiveness.d.ts +70 -0
  666. package/dist/perf/responsiveness.js +560 -0
  667. package/dist/providers/active-install-message.d.ts +42 -0
  668. package/dist/providers/active-install-message.js +49 -0
  669. package/dist/providers/apple-foundation-models/binary.d.ts +4 -0
  670. package/dist/providers/apple-foundation-models/binary.js +18 -0
  671. package/dist/providers/apple-foundation-models/helper.d.ts +66 -0
  672. package/dist/providers/apple-foundation-models/helper.js +262 -0
  673. package/dist/providers/apple-foundation-models/provider.d.ts +82 -0
  674. package/dist/providers/apple-foundation-models/provider.js +389 -0
  675. package/dist/providers/apple-foundation-models/status.d.ts +9 -0
  676. package/dist/providers/apple-foundation-models/status.js +49 -0
  677. package/dist/providers/audio/kokoro-frontend.d.ts +50 -0
  678. package/dist/providers/audio/kokoro-frontend.js +113 -0
  679. package/dist/providers/audio/managed-speech-model.d.ts +9 -0
  680. package/dist/providers/audio/managed-speech-model.js +86 -0
  681. package/dist/providers/audio/types.d.ts +202 -0
  682. package/dist/providers/audio/types.js +17 -0
  683. package/dist/providers/audio/whisper-cpp.d.ts +99 -0
  684. package/dist/providers/audio/whisper-cpp.js +420 -0
  685. package/dist/providers/audio-persister-contract.d.ts +13 -0
  686. package/dist/providers/audio-persister-contract.js +1 -0
  687. package/dist/providers/catalog-model-config.d.ts +63 -0
  688. package/dist/providers/catalog-model-config.js +95 -0
  689. package/dist/providers/code-block-salvage.d.ts +1 -0
  690. package/dist/providers/code-block-salvage.js +2 -0
  691. package/dist/providers/codex-cli/reasoning.d.ts +8 -0
  692. package/dist/providers/codex-cli/reasoning.js +17 -0
  693. package/dist/providers/complete-tool-call.d.ts +16 -0
  694. package/dist/providers/complete-tool-call.js +25 -0
  695. package/dist/providers/constrained-turn.d.ts +1 -0
  696. package/dist/providers/constrained-turn.js +2 -0
  697. package/dist/providers/deliverable-read-pacing.d.ts +1 -0
  698. package/dist/providers/deliverable-read-pacing.js +2 -0
  699. package/dist/providers/direct-file-work-prompt.d.ts +1 -0
  700. package/dist/providers/direct-file-work-prompt.js +2 -0
  701. package/dist/providers/ds4/build-provider.d.ts +100 -0
  702. package/dist/providers/ds4/build-provider.js +539 -0
  703. package/dist/providers/ds4/cache-adapter.d.ts +42 -0
  704. package/dist/providers/ds4/cache-adapter.js +110 -0
  705. package/dist/providers/ds4/dspark.d.ts +83 -0
  706. package/dist/providers/ds4/dspark.js +107 -0
  707. package/dist/providers/ds4/mtp.d.ts +32 -0
  708. package/dist/providers/ds4/mtp.js +46 -0
  709. package/dist/providers/ds4/provider.d.ts +67 -0
  710. package/dist/providers/ds4/provider.js +104 -0
  711. package/dist/providers/ds4/residency.d.ts +143 -0
  712. package/dist/providers/ds4/residency.js +139 -0
  713. package/dist/providers/ds4/stdout-parser.d.ts +47 -0
  714. package/dist/providers/ds4/stdout-parser.js +134 -0
  715. package/dist/providers/duplicate-tool-calls.d.ts +1 -0
  716. package/dist/providers/duplicate-tool-calls.js +2 -0
  717. package/dist/providers/file-repair-policy.d.ts +1 -0
  718. package/dist/providers/file-repair-policy.js +2 -0
  719. package/dist/providers/gpu-arbiter.d.ts +159 -0
  720. package/dist/providers/gpu-arbiter.js +301 -0
  721. package/dist/providers/image-persister-contract.d.ts +15 -0
  722. package/dist/providers/image-persister-contract.js +1 -0
  723. package/dist/providers/llama-cpp/build-metadata.d.ts +13 -0
  724. package/dist/providers/llama-cpp/build-metadata.js +39 -0
  725. package/dist/providers/llama-cpp/build-provider.d.ts +66 -0
  726. package/dist/providers/llama-cpp/build-provider.js +1491 -0
  727. package/dist/providers/llama-cpp/cache-adapter.d.ts +242 -0
  728. package/dist/providers/llama-cpp/cache-adapter.js +689 -0
  729. package/dist/providers/llama-cpp/devices.d.ts +130 -0
  730. package/dist/providers/llama-cpp/devices.js +256 -0
  731. package/dist/providers/llama-cpp/disk-cache.d.ts +8 -0
  732. package/dist/providers/llama-cpp/disk-cache.js +112 -0
  733. package/dist/providers/llama-cpp/engine-flags.d.ts +181 -0
  734. package/dist/providers/llama-cpp/engine-flags.js +279 -0
  735. package/dist/providers/llama-cpp/gguf-metadata-async.d.ts +11 -0
  736. package/dist/providers/llama-cpp/gguf-metadata-async.js +180 -0
  737. package/dist/providers/llama-cpp/gguf-metadata-worker.d.ts +7 -0
  738. package/dist/providers/llama-cpp/gguf-metadata-worker.js +27 -0
  739. package/dist/providers/llama-cpp/gguf-metadata.d.ts +198 -0
  740. package/dist/providers/llama-cpp/gguf-metadata.js +615 -0
  741. package/dist/providers/llama-cpp/gguf-quantization-backfill.d.ts +32 -0
  742. package/dist/providers/llama-cpp/gguf-quantization-backfill.js +75 -0
  743. package/dist/providers/llama-cpp/index.d.ts +10 -0
  744. package/dist/providers/llama-cpp/index.js +13 -0
  745. package/dist/providers/llama-cpp/kv-cache-type.d.ts +118 -0
  746. package/dist/providers/llama-cpp/kv-cache-type.js +140 -0
  747. package/dist/providers/llama-cpp/log.d.ts +92 -0
  748. package/dist/providers/llama-cpp/log.js +280 -0
  749. package/dist/providers/llama-cpp/models.d.ts +389 -0
  750. package/dist/providers/llama-cpp/models.js +1011 -0
  751. package/dist/providers/llama-cpp/offload-planner.d.ts +339 -0
  752. package/dist/providers/llama-cpp/offload-planner.js +524 -0
  753. package/dist/providers/llama-cpp/provider.d.ts +493 -0
  754. package/dist/providers/llama-cpp/provider.js +862 -0
  755. package/dist/providers/llama-cpp/reasoning-launch.d.ts +7 -0
  756. package/dist/providers/llama-cpp/reasoning-launch.js +14 -0
  757. package/dist/providers/llama-cpp/spec-draft.d.ts +44 -0
  758. package/dist/providers/llama-cpp/spec-draft.js +44 -0
  759. package/dist/providers/llama-cpp/stdout-parser.d.ts +70 -0
  760. package/dist/providers/llama-cpp/stdout-parser.js +197 -0
  761. package/dist/providers/local-tool-call-salvage.d.ts +1 -0
  762. package/dist/providers/local-tool-call-salvage.js +2 -0
  763. package/dist/providers/local-turn-policy.d.ts +1 -0
  764. package/dist/providers/local-turn-policy.js +2 -0
  765. package/dist/providers/mlx/async-file-handoff.d.ts +11 -0
  766. package/dist/providers/mlx/async-file-handoff.js +25 -0
  767. package/dist/providers/mlx/build-provider.d.ts +65 -0
  768. package/dist/providers/mlx/build-provider.js +717 -0
  769. package/dist/providers/mlx/cache-adapter.d.ts +161 -0
  770. package/dist/providers/mlx/cache-adapter.js +445 -0
  771. package/dist/providers/mlx/chat-protocol.d.ts +48 -0
  772. package/dist/providers/mlx/chat-protocol.js +8 -0
  773. package/dist/providers/mlx/config-parser.d.ts +29 -0
  774. package/dist/providers/mlx/config-parser.js +81 -0
  775. package/dist/providers/mlx/engine-gate.d.ts +51 -0
  776. package/dist/providers/mlx/engine-gate.js +207 -0
  777. package/dist/providers/mlx/engine-hold.d.ts +41 -0
  778. package/dist/providers/mlx/engine-hold.js +59 -0
  779. package/dist/providers/mlx/engine-log-router.d.ts +84 -0
  780. package/dist/providers/mlx/engine-log-router.js +137 -0
  781. package/dist/providers/mlx/engine-stats-sampler.d.ts +36 -0
  782. package/dist/providers/mlx/engine-stats-sampler.js +59 -0
  783. package/dist/providers/mlx/index.d.ts +10 -0
  784. package/dist/providers/mlx/index.js +10 -0
  785. package/dist/providers/mlx/log.d.ts +55 -0
  786. package/dist/providers/mlx/log.js +157 -0
  787. package/dist/providers/mlx/model-geometry.d.ts +35 -0
  788. package/dist/providers/mlx/model-geometry.js +118 -0
  789. package/dist/providers/mlx/models.d.ts +289 -0
  790. package/dist/providers/mlx/models.js +1079 -0
  791. package/dist/providers/mlx/provider.d.ts +404 -0
  792. package/dist/providers/mlx/provider.js +3017 -0
  793. package/dist/providers/mlx/python/cache_persist.py +348 -0
  794. package/dist/providers/mlx/python/cache_seed.py +402 -0
  795. package/dist/providers/mlx/python/gezel_mlx_server.py +4491 -0
  796. package/dist/providers/mlx/python/lfm2_compat.py +55 -0
  797. package/dist/providers/mlx/python/qwen3_5_text_compat.py +60 -0
  798. package/dist/providers/mlx/python/qwen4_ple.py +159 -0
  799. package/dist/providers/mlx/python/spec_decode.py +775 -0
  800. package/dist/providers/mlx/python/template_stability.py +133 -0
  801. package/dist/providers/mlx/python/think_budget.py +113 -0
  802. package/dist/providers/mlx/python/tool_args_json.py +297 -0
  803. package/dist/providers/mlx/python/tool_call_stream.py +265 -0
  804. package/dist/providers/mlx/python/tool_grammar.py +794 -0
  805. package/dist/providers/mlx/python/vision_inputs.py +64 -0
  806. package/dist/providers/mlx/python/wave_policy.py +147 -0
  807. package/dist/providers/mlx/reasoning-stream.d.ts +85 -0
  808. package/dist/providers/mlx/reasoning-stream.js +196 -0
  809. package/dist/providers/mlx/request-shape.d.ts +34 -0
  810. package/dist/providers/mlx/request-shape.js +42 -0
  811. package/dist/providers/mlx/runtime-diagnostics.d.ts +29 -0
  812. package/dist/providers/mlx/runtime-diagnostics.js +162 -0
  813. package/dist/providers/mlx/spec-drafter.d.ts +25 -0
  814. package/dist/providers/mlx/spec-drafter.js +73 -0
  815. package/dist/providers/mlx/stdout-parser.d.ts +60 -0
  816. package/dist/providers/mlx/stdout-parser.js +254 -0
  817. package/dist/providers/mlx/textual-tool-call-salvage.d.ts +46 -0
  818. package/dist/providers/mlx/textual-tool-call-salvage.js +396 -0
  819. package/dist/providers/mlx/tool-call-protocol.d.ts +65 -0
  820. package/dist/providers/mlx/tool-call-protocol.js +140 -0
  821. package/dist/providers/mlx/tool-call-stripper.d.ts +68 -0
  822. package/dist/providers/mlx/tool-call-stripper.js +153 -0
  823. package/dist/providers/mlx/tool-image-retention.d.ts +23 -0
  824. package/dist/providers/mlx/tool-image-retention.js +43 -0
  825. package/dist/providers/mlx/turn-telemetry.d.ts +22 -0
  826. package/dist/providers/mlx/turn-telemetry.js +45 -0
  827. package/dist/providers/mlx/venv.d.ts +91 -0
  828. package/dist/providers/mlx/venv.js +111 -0
  829. package/dist/providers/mlx/vision-mode.d.ts +92 -0
  830. package/dist/providers/mlx/vision-mode.js +187 -0
  831. package/dist/providers/mlx/vision.d.ts +2 -0
  832. package/dist/providers/mlx/vision.js +12 -0
  833. package/dist/providers/mlx/wire-transcript.d.ts +14 -0
  834. package/dist/providers/mlx/wire-transcript.js +43 -0
  835. package/dist/providers/native/capacity-broker.d.ts +812 -0
  836. package/dist/providers/native/capacity-broker.js +1058 -0
  837. package/dist/providers/native/capacity-budget.d.ts +133 -0
  838. package/dist/providers/native/capacity-budget.js +265 -0
  839. package/dist/providers/native/device-capacity-ledger.d.ts +31 -0
  840. package/dist/providers/native/device-capacity-ledger.js +196 -0
  841. package/dist/providers/native/device-capacity.d.ts +84 -0
  842. package/dist/providers/native/device-capacity.js +431 -0
  843. package/dist/providers/native/engine-api-key.d.ts +57 -0
  844. package/dist/providers/native/engine-api-key.js +72 -0
  845. package/dist/providers/native/engine-key.d.ts +18 -0
  846. package/dist/providers/native/engine-key.js +41 -0
  847. package/dist/providers/native/engine-router.d.ts +145 -0
  848. package/dist/providers/native/engine-router.js +191 -0
  849. package/dist/providers/native/gpu-panic-guard.d.ts +54 -0
  850. package/dist/providers/native/gpu-panic-guard.js +59 -0
  851. package/dist/providers/native/inference.d.ts +33 -0
  852. package/dist/providers/native/inference.js +49 -0
  853. package/dist/providers/native/local-engine-runtime.d.ts +412 -0
  854. package/dist/providers/native/local-engine-runtime.js +1714 -0
  855. package/dist/providers/native/measured-budget.d.ts +32 -0
  856. package/dist/providers/native/measured-budget.js +54 -0
  857. package/dist/providers/native/port.d.ts +11 -0
  858. package/dist/providers/native/port.js +29 -0
  859. package/dist/providers/native/provider-pool.d.ts +474 -0
  860. package/dist/providers/native/provider-pool.js +1166 -0
  861. package/dist/providers/native/supervisor.d.ts +503 -0
  862. package/dist/providers/native/supervisor.js +1466 -0
  863. package/dist/providers/observation-context.d.ts +16 -0
  864. package/dist/providers/observation-context.js +1 -0
  865. package/dist/providers/openai-compatible/sse.d.ts +1 -0
  866. package/dist/providers/openai-compatible/sse.js +2 -0
  867. package/dist/providers/patient-fetch.d.ts +36 -0
  868. package/dist/providers/patient-fetch.js +44 -0
  869. package/dist/providers/project-macro-loop-bail.d.ts +7 -0
  870. package/dist/providers/project-macro-loop-bail.js +21 -0
  871. package/dist/providers/prose-document-salvage.d.ts +1 -0
  872. package/dist/providers/prose-document-salvage.js +2 -0
  873. package/dist/providers/provider-disposal.d.ts +1 -0
  874. package/dist/providers/provider-disposal.js +2 -0
  875. package/dist/providers/queue.d.ts +15 -0
  876. package/dist/providers/queue.js +23 -0
  877. package/dist/providers/ramble-abort-message.d.ts +1 -0
  878. package/dist/providers/ramble-abort-message.js +2 -0
  879. package/dist/providers/ramble-detector.d.ts +1 -0
  880. package/dist/providers/ramble-detector.js +2 -0
  881. package/dist/providers/reasoning-depth.d.ts +1 -0
  882. package/dist/providers/reasoning-depth.js +2 -0
  883. package/dist/providers/required-input-reads.d.ts +1 -0
  884. package/dist/providers/required-input-reads.js +2 -0
  885. package/dist/providers/session-defaults.d.ts +13 -0
  886. package/dist/providers/session-defaults.js +49 -0
  887. package/dist/providers/session-hooks.d.ts +26 -0
  888. package/dist/providers/session-hooks.js +1 -0
  889. package/dist/providers/streaming-session.d.ts +1 -0
  890. package/dist/providers/streaming-session.js +2 -0
  891. package/dist/providers/terminal-tool-policy.d.ts +1 -0
  892. package/dist/providers/terminal-tool-policy.js +2 -0
  893. package/dist/providers/tool-arg-schema-coercion.d.ts +1 -0
  894. package/dist/providers/tool-arg-schema-coercion.js +2 -0
  895. package/dist/providers/tool-executor.d.ts +12 -0
  896. package/dist/providers/tool-executor.js +20 -0
  897. package/dist/providers/tool-failure-tracker.d.ts +1 -0
  898. package/dist/providers/tool-failure-tracker.js +2 -0
  899. package/dist/providers/tool-repeat-tracker.d.ts +1 -0
  900. package/dist/providers/tool-repeat-tracker.js +2 -0
  901. package/dist/providers/turn-timeout.d.ts +1 -0
  902. package/dist/providers/turn-timeout.js +1 -0
  903. package/dist/providers/types.d.ts +834 -0
  904. package/dist/providers/types.js +42 -0
  905. package/dist/providers/usage-builder.d.ts +1 -0
  906. package/dist/providers/usage-builder.js +2 -0
  907. package/dist/providers/vision-capability.d.ts +109 -0
  908. package/dist/providers/vision-capability.js +137 -0
  909. package/dist/python/mlx-runtime-status-bus.d.ts +39 -0
  910. package/dist/python/mlx-runtime-status-bus.js +27 -0
  911. package/dist/python/uv-runtime.d.ts +176 -0
  912. package/dist/python/uv-runtime.js +492 -0
  913. package/dist/relevance/install.d.ts +28 -0
  914. package/dist/relevance/install.js +114 -0
  915. package/dist/relevance/manager.d.ts +79 -0
  916. package/dist/relevance/manager.js +211 -0
  917. package/dist/relevance/registry.d.ts +11 -0
  918. package/dist/relevance/registry.js +158 -0
  919. package/dist/relevance/relevance-core.d.ts +67 -0
  920. package/dist/relevance/relevance-core.js +189 -0
  921. package/dist/relevance/relevance-model.d.ts +42 -0
  922. package/dist/relevance/relevance-model.js +292 -0
  923. package/dist/relevance/relevance-worker.d.ts +20 -0
  924. package/dist/relevance/relevance-worker.js +57 -0
  925. package/dist/relevance/settings.d.ts +30 -0
  926. package/dist/relevance/settings.js +81 -0
  927. package/dist/remotes/identity.d.ts +121 -0
  928. package/dist/remotes/identity.js +348 -0
  929. package/dist/remotes/pinned-fetch.d.ts +45 -0
  930. package/dist/remotes/pinned-fetch.js +193 -0
  931. package/dist/remotes/types.d.ts +54 -0
  932. package/dist/remotes/types.js +1 -0
  933. package/dist/runtime-lock.d.ts +29 -0
  934. package/dist/runtime-lock.js +85 -0
  935. package/dist/search/relevance-stage.d.ts +81 -0
  936. package/dist/search/relevance-stage.js +177 -0
  937. package/dist/secrets/types.d.ts +63 -0
  938. package/dist/secrets/types.js +64 -0
  939. package/dist/speech.d.ts +1 -0
  940. package/dist/speech.js +1 -0
  941. package/dist/system/darwin-memory.d.ts +26 -0
  942. package/dist/system/darwin-memory.js +81 -0
  943. package/dist/system/gezel-process-memory.d.ts +55 -0
  944. package/dist/system/gezel-process-memory.js +156 -0
  945. package/dist/system/memory.d.ts +154 -0
  946. package/dist/system/memory.js +501 -0
  947. package/dist/transformers-cache.d.ts +112 -0
  948. package/dist/transformers-cache.js +392 -0
  949. package/dist/utils/disk-space.d.ts +55 -0
  950. package/dist/utils/disk-space.js +72 -0
  951. package/dist/utils/download-common.d.ts +138 -0
  952. package/dist/utils/download-common.js +196 -0
  953. package/dist/utils/download-with-retry.d.ts +85 -0
  954. package/dist/utils/download-with-retry.js +522 -0
  955. package/dist/utils/extract-zip.d.ts +22 -0
  956. package/dist/utils/extract-zip.js +103 -0
  957. package/dist/utils/retire-model-worker.d.ts +15 -0
  958. package/dist/utils/retire-model-worker.js +24 -0
  959. package/dist/utils/service-worker-entry.d.ts +59 -0
  960. package/dist/utils/service-worker-entry.js +69 -0
  961. package/dist/utils/xet-download.d.ts +40 -0
  962. package/dist/utils/xet-download.js +526 -0
  963. package/dist/utils/xet.d.ts +86 -0
  964. package/dist/utils/xet.js +201 -0
  965. package/package.json +103 -3
  966. package/scripts/embedding-packaging.cjs +264 -0
  967. package/scripts/embedding-packaging.d.cts +16 -0
@@ -0,0 +1,4491 @@
1
+ """
2
+ gezel_mlx_server — thin OpenAI-compatible HTTP server around mlx-vlm
3
+ that owns the prompt-cache lifecycle.
4
+
5
+ Why this exists: mlx_vlm.server (the upstream binary) clears the
6
+ prompt cache after every streamed response (see its `finally`:
7
+ `mx.clear_cache()` + `gc.collect()`). That means every turn re-prefills
8
+ the entire context — on a 25K-token Meester session at ~250 tok/s,
9
+ that's a 2-minute wait per turn. This server preserves the cache
10
+ between requests when the caller supplies a `cache_id`, and reuses any
11
+ matching prefix on the next request.
12
+
13
+ Reuse mechanism: mlx-vlm exposes `PromptCacheState` (in
14
+ mlx_vlm/generate.py) for exactly this purpose — pass it as the
15
+ `prompt_cache_state` kwarg to `stream_generate` and only the new
16
+ tokens beyond the common prefix get prefilled. We hold one
17
+ `PromptCacheState` per `cache_id` in a module-level dict, with bounded
18
+ LRU eviction by total estimated bytes. The batched path plans its own
19
+ reuse via `cache_seed` (longest-common-prefix + protocol trim, plus an
20
+ end-of-prompt KV snapshot for windowed models whose rotated caches
21
+ can't rewind) — see cache_seed.py's module docstring for the why.
22
+
23
+ Compatibility contract: we serve the same OpenAI-shape
24
+ `/v1/chat/completions` endpoint as mlx_vlm.server, with the same SSE
25
+ shape — the gezel `MlxSession` SSE reader and our existing tool-call
26
+ stripper / Gemma-format parser keep working unchanged. We do NOT do
27
+ server-side tool-call extraction (mlx_vlm.server's `process_tool_calls`
28
+ path); the gezel side already handles that better. The model's
29
+ textual tool-call markup flows through as content; the gezel
30
+ `LeakyToolCallStripper` extracts and repairs it.
31
+
32
+ Vision requests use the full installed tower, serialized against text batch
33
+ waves. Their KV is deliberately uncached so image identity cannot be lost in
34
+ text-prefix reuse. Ordinary text retains the existing batched cache path.
35
+ """
36
+
37
+ from __future__ import annotations
38
+
39
+ import argparse
40
+ import asyncio
41
+ import collections
42
+ import hashlib
43
+ import json
44
+ import os
45
+ import re
46
+ import sys
47
+ import threading
48
+ import time
49
+ import traceback
50
+ import types
51
+ import uuid
52
+ from pathlib import Path
53
+ from typing import Any, Dict, List, Optional
54
+
55
+ import mlx.core as mx
56
+ from fastapi import FastAPI, HTTPException, Request
57
+ from fastapi.responses import JSONResponse, StreamingResponse
58
+ from pydantic import BaseModel
59
+
60
+ # mlx_vlm internals we lean on. Importing utils.load + prompt_utils +
61
+ # generate.stream_generate gives us the same model loading, chat
62
+ # templating, and generation primitives as the upstream server. The
63
+ # PromptCacheState class is the single bit of upstream-public API
64
+ # that does the work — see its docstring in mlx_vlm/generate.py.
65
+ from mlx_vlm.generate import PromptCacheState, stream_generate
66
+ from mlx_vlm.prompt_utils import apply_chat_template
67
+ from mlx_vlm.utils import load
68
+
69
+ # Local module — disk persistence for PromptCacheState.
70
+ # Imported by relative path so the file works whether the server is
71
+ # launched as a module or as a free script (Electron's bundle layout
72
+ # extracts both .py files into the same dir).
73
+ sys.path.insert(0, str(Path(__file__).parent))
74
+ import cache_persist # noqa: E402
75
+ import think_budget # noqa: E402
76
+ import cache_seed # noqa: E402
77
+ import tool_grammar # noqa: E402
78
+ import spec_decode # noqa: E402
79
+ import tool_call_stream # noqa: E402
80
+ import template_stability # noqa: E402
81
+ import vision_inputs # noqa: E402
82
+ import wave_policy # noqa: E402
83
+ from lfm2_compat import ensure_lfm2_config_compat # noqa: E402
84
+ from qwen3_5_text_compat import is_text_only_qwen3_5_checkpoint # noqa: E402
85
+ from qwen4_ple import prepare_external_ple_view # noqa: E402
86
+
87
+
88
+ def log_contained_exception(tag: str) -> None:
89
+ """Print the current traceback for a failure the caller has CONTAINED.
90
+
91
+ Tagged line by line, and never `traceback.print_exc()`: the supervisor's
92
+ stdout classifier reads a bare `SomeError: <message>` leaf line as proof
93
+ the engine is dead and tears the process down. That is right during
94
+ startup, which is what it was written for, and wrong at every call site
95
+ here — each one has already answered the affected request and left the
96
+ engine serving everyone else on it. A speculative-decode wave that failed
97
+ one request this way SIGKILLed a 27B three minutes into serving and took
98
+ an unrelated session's turn down with it.
99
+ """
100
+ for line in traceback.format_exc().rstrip().splitlines():
101
+ print(f"[{tag}] | {line}", flush=True)
102
+
103
+
104
+ # ───────── Leaked special-token scrub ─────────
105
+
106
+ # Chat-template framing tokens — turn / sequence / tool-response
107
+ # delimiters that tight quants (e.g. gemma4-e4b-q4) sometimes emit as
108
+ # literal special-token text. `tokenizer.decode()` keeps special tokens,
109
+ # so these render into the SSE content stream verbatim; wild-caught case
110
+ # was gemma4-e4b-q4 streaming `<eos><|tool_response><eos>` after firing
111
+ # its tool calls. We scrub them at every detokenization site so they
112
+ # never reach the client.
113
+ #
114
+ # Deliberately NOT scrubbed: the tool-CALL markers `<|tool_call>` /
115
+ # `<tool_call|>` (and `<|tool>` / `<tool|>`). The gezel-side
116
+ # LeakyToolCallStripper relies on those flowing through as text to
117
+ # salvage and repair tool calls — this server intentionally does no
118
+ # server-side tool-call extraction (see module docstring). A blanket
119
+ # `skip_special_tokens=True` would delete them too, which is why we
120
+ # strip a targeted denylist instead.
121
+ _LEAKED_MARKER_RE = re.compile(
122
+ r"<\|?(?:eos|bos|pad|unk|mask|turn|tool_response|start_of_turn|end_of_turn|im_start|im_end)\|?>",
123
+ re.IGNORECASE,
124
+ )
125
+
126
+
127
+ def _make_tool_stream(request_id: str, tools=None):
128
+ """Build a splitter for this request, or None to keep the old behaviour.
129
+
130
+ Returns None whenever the flag is off OR the model's chat template
131
+ matches no mlx_vlm tool parser — in both cases the caller streams raw
132
+ content and the gezel-side salvage path stays in charge. Never guess a
133
+ marker pair: a wrong boundary would silently eat visible output.
134
+ """
135
+ if not getattr(ARGS, "stream_tool_calls", False):
136
+ return None
137
+ markers = tool_call_stream.resolve_tool_markers(PROCESSOR)
138
+ if markers is None:
139
+ return None
140
+ start, end = markers
141
+ return tool_call_stream.ToolCallStreamState(
142
+ start_marker=start, end_marker=end, request_id=request_id, tools=tools
143
+ )
144
+
145
+
146
+ def _tool_stream_fragments(state, chunk_text: str):
147
+ """Fragments to emit for a chunk; empty when extraction is disabled."""
148
+ if state is None:
149
+ return []
150
+ return state.feed(chunk_text)
151
+
152
+
153
+ def _tool_stream_payload(frag, request_id: str, model: str, chunk) -> Dict[str, Any]:
154
+ """Wrap one fragment in the same chunk envelope the content path uses."""
155
+ if frag.content is not None:
156
+ delta: Dict[str, Any] = {"role": "assistant", "content": frag.content}
157
+ else:
158
+ call: Dict[str, Any] = {
159
+ "index": frag.tool_call_index,
160
+ "function": {"arguments": frag.arguments or ""},
161
+ }
162
+ if frag.tool_call_id is not None:
163
+ call["id"] = frag.tool_call_id
164
+ call["type"] = "function"
165
+ # The grammar fixes the name inside the argument body, so the
166
+ # name field stays empty rather than being guessed here; the
167
+ # client's parser reads it from the arguments as it always has.
168
+ call["function"]["name"] = frag.name or ""
169
+ delta = {"role": "assistant", "tool_calls": [call]}
170
+ return {
171
+ "id": request_id,
172
+ "object": "chat.completion.chunk",
173
+ "created": int(time.time()),
174
+ "model": model,
175
+ "choices": [{"index": 0, "delta": delta, "finish_reason": None}],
176
+ "usage": {
177
+ "input_tokens": getattr(chunk, "prompt_tokens", 0),
178
+ "output_tokens": getattr(chunk, "generation_tokens", 0),
179
+ "total_tokens": (
180
+ getattr(chunk, "prompt_tokens", 0) + getattr(chunk, "generation_tokens", 0)
181
+ ),
182
+ "prompt_tps": getattr(chunk, "prompt_tps", 0.0),
183
+ "generation_tps": getattr(chunk, "generation_tps", 0.0),
184
+ "cached_tokens": getattr(chunk, "cached_tokens", 0),
185
+ },
186
+ }
187
+
188
+
189
+ def _scrub_leaked_markers(text: str) -> str:
190
+ """Drop chat-template framing special-tokens that leaked as text."""
191
+ if not text or "<" not in text:
192
+ return text
193
+ return _LEAKED_MARKER_RE.sub("", text)
194
+
195
+
196
+ def _generation_config_eos_ids(model_dir: str) -> "List[int]":
197
+ """Read the stop-token set a model declares in generation_config.json.
198
+
199
+ mlx's tokenizer wrapper doesn't always surface the full eos set the
200
+ model author declared. gemma4-e4b-q4 is the wild-caught case: it
201
+ declares `eos_token_id: [1, 106, 50]` (`<eos>`, `<turn|>`,
202
+ `<|tool_response>`), but the wrapper resolved a narrower set, so the
203
+ batch engine never stopped on `<eos>` / `<|tool_response>` and
204
+ streamed them as text after the model's tool calls. Merging these in
205
+ halts generation at the real turn boundary instead of running on
206
+ into framing tokens (the upstream fix that complements the scrub).
207
+ """
208
+ try:
209
+ with open(os.path.join(model_dir, "generation_config.json")) as f:
210
+ cfg = json.load(f)
211
+ except (OSError, ValueError):
212
+ return []
213
+ eos = cfg.get("eos_token_id")
214
+ if isinstance(eos, bool): # guard: bools are ints in Python
215
+ return []
216
+ if isinstance(eos, int):
217
+ return [eos]
218
+ if isinstance(eos, list):
219
+ return [int(e) for e in eos if isinstance(e, int) and not isinstance(e, bool)]
220
+ return []
221
+
222
+
223
+ def _tool_call_close_stop_ids(model_dir: str) -> "List[int]":
224
+ """`</tool_call>` special-token ids to halt a turn at the first closed
225
+ tool call — for models that expect a server-side tool-call parser we
226
+ don't run.
227
+
228
+ Poolside Laguna (and kin) declare `tool_call_parser` in
229
+ generation_config.json and emit `<tool_call>…</tool_call>` WITHOUT
230
+ listing `</tool_call>` in `eos_token_id`: their production server (the
231
+ `*_v1` parser) detects the closed call and stops generation. On
232
+ mlx_vlm's raw completion path nothing stops it, so the model runs on
233
+ past the first call and loops out hundreds of calls in a single turn
234
+ until the SSE stalls (wild-caught: laguna-s-118b/MLX — ~450
235
+ `<tool_call>` spans in one turn, then a 300s silent stall → the whole
236
+ turn aborts and no call lands). Making the `</tool_call>` special token
237
+ a stop halts the turn at the first closed call; the gezel GLM salvage
238
+ layer promotes it and the agentic loop advances one call per turn.
239
+
240
+ Gated on the model DECLARING a non-empty `tool_call_parser` so models
241
+ that end tool-call turns on their own eos (Qwen/Gemma) are untouched —
242
+ zero regression risk for the already-tuned cast.
243
+ """
244
+ try:
245
+ with open(os.path.join(model_dir, "generation_config.json")) as f:
246
+ cfg = json.load(f)
247
+ except (OSError, ValueError):
248
+ return []
249
+ parser = cfg.get("tool_call_parser")
250
+ if not isinstance(parser, str) or not parser:
251
+ return []
252
+ try:
253
+ with open(os.path.join(model_dir, "tokenizer_config.json")) as f:
254
+ tcfg = json.load(f)
255
+ except (OSError, ValueError):
256
+ return []
257
+ added = tcfg.get("added_tokens_decoder") or {}
258
+ ids: "List[int]" = []
259
+ for sid, meta in added.items():
260
+ if isinstance(meta, dict) and meta.get("content") == "</tool_call>":
261
+ try:
262
+ ids.append(int(sid))
263
+ except (TypeError, ValueError):
264
+ pass
265
+ return ids
266
+
267
+
268
+ # ───────── Configuration (CLI args) ─────────
269
+
270
+ parser = argparse.ArgumentParser(description="gezel MLX wrapper server")
271
+ parser.add_argument("--model", required=True, help="Path to the MLX model directory")
272
+ parser.add_argument(
273
+ "--external-ple-dir",
274
+ type=str,
275
+ default=None,
276
+ help=(
277
+ "Derived model-view directory for Qwen4's disk-backed PLE table. "
278
+ "Applicable checkpoints are hard-linked into this view and the large "
279
+ "n-gram table is memory-mapped by mlx-vlm 0.7.1. Other model families "
280
+ "keep the source directory; an eligible Qwen4 checkpoint fails safely "
281
+ "when the configured runtime cannot prepare the view."
282
+ ),
283
+ )
284
+ parser.add_argument(
285
+ "--allow-resident-ple-fallback",
286
+ action="store_true",
287
+ help=(
288
+ "Advanced/manual override: if a quantized Qwen4 checkpoint is eligible "
289
+ "for --external-ple-dir but the derived view cannot be prepared, load "
290
+ "the full PLE table resident. Gezel never passes this because its "
291
+ "catalog working-set reservation assumes external PLE is active."
292
+ ),
293
+ )
294
+ parser.add_argument("--host", default="127.0.0.1")
295
+ parser.add_argument("--port", type=int, default=8080)
296
+ parser.add_argument(
297
+ "--cache-budget-mb",
298
+ type=int,
299
+ default=4096,
300
+ help="Bound on total in-memory prompt-cache size before LRU eviction.",
301
+ )
302
+ parser.add_argument(
303
+ "--prefill-step-size",
304
+ type=int,
305
+ default=2048,
306
+ help="Tokens per prefill step. Lower = less peak memory, slower.",
307
+ )
308
+ parser.add_argument(
309
+ "--spec-draft-model",
310
+ type=str,
311
+ default=None,
312
+ help="Path to an MTP drafter directory; enables speculative decoding "
313
+ "for eligible (greedy, processor-free) requests. See spec_decode.py.",
314
+ )
315
+ parser.add_argument(
316
+ "--spec-draft-kind",
317
+ type=str,
318
+ default=None,
319
+ help="Drafter kind override (default: auto-detect from the drafter's model_type).",
320
+ )
321
+ parser.add_argument(
322
+ "--spec-block-size",
323
+ type=int,
324
+ default=None,
325
+ help="Draft block size override (default: the drafter's configured block_size).",
326
+ )
327
+ parser.add_argument(
328
+ "--persist-dir",
329
+ type=str,
330
+ default=None,
331
+ help=(
332
+ "Directory for on-disk PromptCacheState persistence. When set, "
333
+ "in-memory cache misses fall back to a disk lookup, evictions "
334
+ "save before dropping, and graceful shutdown flushes all warm "
335
+ "entries. Disabled when unset (matches Phase 2 behavior)."
336
+ ),
337
+ )
338
+ parser.add_argument(
339
+ "--model-fingerprint",
340
+ type=str,
341
+ default=None,
342
+ help=(
343
+ "Stable identifier for the loaded model — sha256-derived from the "
344
+ "manifest. Disk-cache entries are segmented by fingerprint so a "
345
+ "model swap can never load wrong-shape KV state. Required when "
346
+ "--persist-dir is set."
347
+ ),
348
+ )
349
+ parser.add_argument(
350
+ "--disk-cache-budget-mb",
351
+ type=int,
352
+ default=8192,
353
+ help=(
354
+ "Bound on total on-disk persisted cache size. LRU pruning runs "
355
+ "after every save when over budget. 0 disables pruning."
356
+ ),
357
+ )
358
+ parser.add_argument("--vision", action="store_true", help="Load and enable the installed vision tower")
359
+ parser.add_argument(
360
+ "--max-concurrency",
361
+ type=int,
362
+ default=1,
363
+ help=(
364
+ "Max sequences to generate concurrently. Any positive value routes "
365
+ "generation through a single mlx_lm BatchGenerator; 1 (default) is a "
366
+ "singleton queue that still uses the cache-safe prompt-boundary "
367
+ "snapshot path, while N>1 also enables continuous batching. Set 0 only "
368
+ "for the legacy serial stream_generate fallback. Set from the TS "
369
+ "provider's batchMaxConcurrency."
370
+ ),
371
+ )
372
+ parser.add_argument(
373
+ "--gpu-memory-limit-mb",
374
+ type=int,
375
+ default=0,
376
+ help=(
377
+ "Total GPU/unified-memory budget (MB) this engine may use — weights + "
378
+ "KV + compute. Set from the TS capacity broker (min of the model budget "
379
+ "and Metal's recommended working set). Drives memory-aware wave "
380
+ "admission: when live usage nears this ceiling the BatchGenerator "
381
+ "collapses toward serial instead of co-batching concurrent prefills "
382
+ "into a Metal command-buffer OOM (which SIGABRTs the whole process). "
383
+ "0 (default) = unset → admit purely on --max-concurrency, as before."
384
+ ),
385
+ )
386
+ parser.add_argument(
387
+ "--wired-limit-mb",
388
+ type=int,
389
+ default=0,
390
+ help=(
391
+ "How much this engine may PIN (MB) — the planned resident set: weights "
392
+ "plus the KV its slots will hold. Deliberately NOT the same number as "
393
+ "--gpu-memory-limit-mb: that is a ceiling on what may be USED, while "
394
+ "wired pages are unreclaimable by the OS, so wiring the whole share "
395
+ "starves the machine (a 27B wired 96 GiB for a ~45 GB working set and "
396
+ "left 155 MB free on a 128 GB Mac). 0 (default) = leave the wired "
397
+ "limit alone."
398
+ ),
399
+ )
400
+ parser.add_argument(
401
+ "--stream-tool-calls",
402
+ action="store_true",
403
+ help=(
404
+ "Emit OpenAI-shaped `tool_calls` argument deltas instead of streaming "
405
+ "the model's tool markup as content. OFF by default: the gezel-side "
406
+ "salvage path stays the fallback for models whose chat template "
407
+ "matches no mlx_vlm tool parser. On, the client gets a real tool-call "
408
+ "boundary and can end a turn once the call is usable — the same lever "
409
+ "llama-server gives us, and the reason MLX turns otherwise run to "
410
+ "completion (measured 41 turns at a 136s median, 203 min/10 trials)."
411
+ ),
412
+ )
413
+ parser.add_argument(
414
+ "--kv-bits",
415
+ type=int,
416
+ default=0,
417
+ help=(
418
+ "Quantize the KV cache to this many bits. 0 (default) = f16. Halves "
419
+ "(8-bit) or quarters (4-bit) the largest term in a long-context "
420
+ "session: a 256K window on a 16-full-attention-layer 27B is 16 GiB of "
421
+ "f16 KV. Applies to the SERIAL generation path only — mlx-lm's "
422
+ "BatchGenerator builds BatchKVCache layers, which have no "
423
+ "`to_quantized`, and `maybe_quantize_kv_cache` skips them silently. "
424
+ "The launcher therefore plans memory at f16 whenever batching is on, "
425
+ "so the broker never reserves a discount the engine did not take."
426
+ ),
427
+ )
428
+ parser.add_argument(
429
+ "--kv-group-size",
430
+ type=int,
431
+ default=64,
432
+ help="Group size for KV quantization. mlx-lm's default is 64.",
433
+ )
434
+ parser.add_argument(
435
+ "--kv-quant-scheme",
436
+ default="uniform",
437
+ help=(
438
+ "Quantization scheme passed to mlx-vlm. 'uniform' is right for integer "
439
+ "bit values; mlx-vlm selects TurboQuant for fractional ones."
440
+ ),
441
+ )
442
+ parser.add_argument(
443
+ "--quantized-kv-start",
444
+ type=int,
445
+ default=0,
446
+ help=(
447
+ "Only quantize a layer once its cache passes this many tokens. Leaves "
448
+ "short turns on the f16 path, where quantization costs more than the "
449
+ "memory it saves."
450
+ ),
451
+ )
452
+ ARGS = parser.parse_args()
453
+
454
+
455
+ # Resolve the persistence directory once. None means disk persistence
456
+ # is disabled — every code path that touches disk gates on this.
457
+ PERSIST_DIR: Optional[Path] = (
458
+ Path(ARGS.persist_dir).expanduser() if ARGS.persist_dir else None
459
+ )
460
+ MODEL_FINGERPRINT: Optional[str] = ARGS.model_fingerprint
461
+ if PERSIST_DIR is not None and not MODEL_FINGERPRINT:
462
+ print(
463
+ "[cache] --persist-dir set without --model-fingerprint; refusing to "
464
+ "persist (mismatched fingerprints would silently corrupt outputs).",
465
+ flush=True,
466
+ )
467
+ PERSIST_DIR = None
468
+
469
+
470
+ def _disk_enabled() -> bool:
471
+ return PERSIST_DIR is not None and MODEL_FINGERPRINT is not None
472
+
473
+
474
+ # ───────── Parent-death watchdog ─────────
475
+
476
+
477
+ def _start_parent_death_watchdog() -> None:
478
+ """Self-terminate if our parent process goes away.
479
+
480
+ Why: this server holds 10–20 GB of model weights in unified memory.
481
+ If the gezel daemon crashes, is force-quit, or is SIGKILL'd by
482
+ VSCode's debug "Stop" button, we get reparented to launchd (pid 1)
483
+ and sit on that RAM forever — orphan instances accumulate every
484
+ dev iteration.
485
+
486
+ macOS has no `prctl(PR_SET_PDEATHSIG)` equivalent; polling
487
+ `getppid()` from a daemon thread is the portable substitute. The
488
+ GIL is released during MLX/torch native compute and during uvicorn
489
+ epoll waits, so the 2s tick still fires reliably even under load.
490
+ A `daemon=True` thread doesn't block interpreter shutdown when the
491
+ main thread exits cleanly — the watchdog only matters for the
492
+ crash path.
493
+ """
494
+ initial_ppid = os.getppid()
495
+ if initial_ppid == 1:
496
+ # Already orphaned at startup — there's no parent to serve.
497
+ # Refuse to load weights only to immediately leak them.
498
+ print(
499
+ "[gezel-mlx] parent already pid=1 at startup; refusing to load model",
500
+ flush=True,
501
+ )
502
+ os._exit(1)
503
+
504
+ def watch() -> None:
505
+ while True:
506
+ time.sleep(2.0)
507
+ current = os.getppid()
508
+ if current == initial_ppid:
509
+ continue
510
+ # Reparented — original parent is gone. The new ppid is
511
+ # almost always 1 (launchd) or a session leader; either way
512
+ # nobody who started us is still listening, so release the
513
+ # model RAM. `os._exit` skips atexit/finalizers — uvicorn's
514
+ # graceful-shutdown path would try to drain in-flight
515
+ # streams against a dead client and stall, which defeats
516
+ # the point of the watchdog.
517
+ print(
518
+ f"[gezel-mlx] parent vanished (was {initial_ppid}, now {current}); exiting",
519
+ flush=True,
520
+ )
521
+ os._exit(0)
522
+
523
+ t = threading.Thread(
524
+ target=watch, name="gezel-parent-death-watchdog", daemon=True
525
+ )
526
+ t.start()
527
+
528
+
529
+ _start_parent_death_watchdog()
530
+
531
+
532
+ # ───────── Model loading (one-shot at startup) ─────────
533
+
534
+
535
+ print(f"Started server process [{os.getpid()}]", flush=True)
536
+ print("Waiting for application startup.", flush=True)
537
+ _SOURCE_MODEL_DIR = ARGS.model
538
+ ARGS.model, _ple_status, _ple_applicable = prepare_external_ple_view(
539
+ ARGS.model, ARGS.external_ple_dir
540
+ )
541
+ if ARGS.model != _SOURCE_MODEL_DIR:
542
+ print(f"[ple] external view {_ple_status}: {ARGS.model}", flush=True)
543
+ elif ARGS.external_ple_dir:
544
+ if not _ple_applicable:
545
+ print(f"[ple] not applicable: {_ple_status}", flush=True)
546
+ elif ARGS.allow_resident_ple_fallback:
547
+ print(f"[ple] resident fallback explicitly allowed: {_ple_status}", flush=True)
548
+ else:
549
+ raise RuntimeError(
550
+ "external Qwen4 PLE view is required by the planned memory footprint "
551
+ f"but could not be prepared: {_ple_status}"
552
+ )
553
+ print(f"Loading model from: {ARGS.model}", flush=True)
554
+ _LFM2_FF_DIM = ensure_lfm2_config_compat(ARGS.model)
555
+ if _LFM2_FF_DIM is not None:
556
+ print(f"Applied lfm2 config compat: block_ff_dim={_LFM2_FF_DIM}", flush=True)
557
+ _QWEN3_5_TEXT_ONLY = is_text_only_qwen3_5_checkpoint(ARGS.model)
558
+ if _QWEN3_5_TEXT_ONLY:
559
+ print(
560
+ "Detected text-only Qwen 3.5 checkpoint; allowing absent vision-tower weights.",
561
+ flush=True,
562
+ )
563
+
564
+ # ── Tower selection (2026-08-28) ──
565
+ #
566
+ # Text-only launches can use the faster language tower. --vision retains
567
+ # the full VLM tower for pixel encoding and serialized multimodal requests.
568
+ # Yet mlx_vlm's language towers pay a large context-growing decode tax
569
+ # under the BatchGenerator this server runs on. Measured on
570
+ # qwen3.8-27b-q4 (matched-thermal 2x2, ms per decoded token at
571
+ # 10.9k/54.3k prompt tokens):
572
+ #
573
+ # mlx_lm tower x BatchGenerator : 32.1 / 41.2 (= bare serial)
574
+ # mlx_vlm tower x BatchGenerator : 44.7 / 86.7 (slope 0.97 vs 0.34 us/tok)
575
+ #
576
+ # The vlm tower runs 3-axis multimodal RoPE with per-layer array-offset
577
+ # position ids + left-padding mask machinery per token; mlx_lm's text
578
+ # tower is plain rope(offset) + fused SDPA. Same checkpoint, ~2x decode
579
+ # at working context — so batch mode loads the mlx_lm tower when it can.
580
+ #
581
+ # Scope guards, deliberately conservative:
582
+ # - allow-listed architectures only (verified by the 2x2 above); other
583
+ # archs keep the vlm loader until they get their own measurement,
584
+ # - batch mode only (--max-concurrency >= 1, the default): the legacy
585
+ # serial path is mlx_vlm's stream_generate and keeps the vlm loader,
586
+ # - GEZEL_MLX_TEXT_TOWER=off restores the old behavior outright;
587
+ # =force skips the allowlist (measurement runs).
588
+ # Every branch prints a `[tower]` fingerprint line — an A/B that cannot
589
+ # prove which tower served it is how this class of defect stayed hidden.
590
+ _TEXT_TOWER = None
591
+ _TEXT_TOWER_ALLOWED_ARCHS = {"qwen3_5"}
592
+
593
+
594
+ def _model_architecture(model_dir: str) -> str:
595
+ try:
596
+ with open(os.path.join(model_dir, "config.json"), "r") as fh:
597
+ return str(json.load(fh).get("model_type", ""))
598
+ except Exception:
599
+ return ""
600
+
601
+
602
+ _tower_pref = os.environ.get("GEZEL_MLX_TEXT_TOWER", "auto").strip().lower()
603
+ _arch = _model_architecture(ARGS.model)
604
+ if (
605
+ not ARGS.vision
606
+ and _tower_pref not in ("0", "off", "false")
607
+ and int(getattr(ARGS, "max_concurrency", 1) or 0) >= 1
608
+ and (_tower_pref == "force" or _arch in _TEXT_TOWER_ALLOWED_ARCHS)
609
+ # A configured MTP drafter needs the vlm tower (rollback lives there);
610
+ # an explicit GEZEL_MLX_TEXT_TOWER=force still wins and spec turns off.
611
+ and not (getattr(ARGS, "spec_draft_model", None) and _tower_pref != "force")
612
+ ):
613
+ try:
614
+ from mlx_lm.utils import load as _mlx_lm_load
615
+
616
+ MODEL, PROCESSOR = _mlx_lm_load(ARGS.model)
617
+ _TEXT_TOWER = "mlx_lm"
618
+ print(
619
+ f"[tower] active=mlx_lm arch={_arch or 'unknown'} "
620
+ f"(text tower; vlm loader bypassed for batch serving)",
621
+ flush=True,
622
+ )
623
+ except Exception as exc:
624
+ print(
625
+ f"[tower] mlx_lm load failed for arch={_arch or 'unknown'}: {exc} "
626
+ f"— falling back to mlx_vlm",
627
+ flush=True,
628
+ )
629
+ if _TEXT_TOWER is None:
630
+ MODEL, PROCESSOR = load(ARGS.model, strict=not _QWEN3_5_TEXT_ONLY)
631
+ _reason = (
632
+ "native vision enabled" if ARGS.vision else
633
+ "GEZEL_MLX_TEXT_TOWER=off"
634
+ if _tower_pref in ("0", "off", "false")
635
+ else "serial mode (--max-concurrency 0)"
636
+ if int(getattr(ARGS, "max_concurrency", 1) or 0) < 1
637
+ else "spec drafter configured (MTP runs on the vlm tower)"
638
+ if (getattr(ARGS, "spec_draft_model", None) and _tower_pref != "force")
639
+ else f"arch {_arch or 'unknown'} not in text-tower allowlist"
640
+ if _tower_pref != "force"
641
+ else "mlx_lm load failed"
642
+ )
643
+ print(f"[tower] active=mlx_vlm ({_reason})", flush=True)
644
+ _VISION_ENABLED = bool(ARGS.vision and getattr(PROCESSOR, "image_processor", None) is not None)
645
+ if ARGS.vision and not _VISION_ENABLED:
646
+ raise RuntimeError("Native vision was requested but this checkpoint has no image processor")
647
+ print(f"[vision] enabled={_VISION_ENABLED}", flush=True)
648
+ print("Model and processor loaded successfully.", flush=True)
649
+
650
+ # ── MTP speculative decoding (2026-08-28) ──
651
+ # Armed only when a drafter is configured AND the capability probe passes;
652
+ # every failure logs one `[spec] off` line and the server serves normally.
653
+ _SPEC = spec_decode.resolve_spec(ARGS, MODEL, _TEXT_TOWER)
654
+ if _SPEC is not None:
655
+ print(
656
+ f"[spec] active drafter={_SPEC.drafter_dir} kind={_SPEC.kind} "
657
+ f"block={_SPEC.block_size or 'config-default'} tower=mlx_vlm",
658
+ flush=True,
659
+ )
660
+
661
+
662
+ # Prefix-stability work (template relocation + reasoning-stable snapshot
663
+ # boundaries), OFF by default. Wired, unit-tested, and demonstrably firing —
664
+ # a single turn reused 6,980 tokens against 165 prefilled — but the paired
665
+ # arms measured it as a REGRESSION on its own metric:
666
+ #
667
+ # reuse% 34 (off) -> 26 -> 22 -> 14 (all fixes on)
668
+ # prefill/trial 37,486 (off) -> 43,586 (on)
669
+ #
670
+ # Truncating saves at the stable boundary buys matches but shrinks what each
671
+ # match is worth, and warming with tools raises the prefill denominator. Net
672
+ # negative at n=1-5 with ~3x trial spread, so it is neither shippable nor
673
+ # disproven. Same posture as GEZEL_LAYERED_PREFIX_CACHE: keep it measurable.
674
+ _STABLE_PREFIX_ENABLED = os.environ.get("GEZEL_MLX_STABLE_PREFIX") in ("1", "true")
675
+
676
+
677
+ def _resolve_stable_chat_template() -> Optional[str]:
678
+ """Relocate the reasoning preamble below the cache-stable content.
679
+
680
+ Qwen 3.5+ templates put the reasoning-effort line ABOVE the tool block,
681
+ so flipping reasoning mode between turns — which the constrained-write
682
+ path does deliberately — changed the prompt from token 3 onward and
683
+ re-prefilled the entire tool block every time. Measured on
684
+ qwen3.8-27b-q4: 3 shared tokens before, 2,921 after.
685
+
686
+ Verified before adoption and abandoned on any doubt: a mis-rendered
687
+ prompt is far worse than a cold cache.
688
+ """
689
+ if not _STABLE_PREFIX_ENABLED:
690
+ return None
691
+ tok = getattr(PROCESSOR, "tokenizer", None) or PROCESSOR
692
+ original = getattr(tok, "chat_template", None)
693
+ if not isinstance(original, str) or not original:
694
+ return None
695
+ patched = template_stability.relocate_reasoning_preamble(original)
696
+ if patched is None:
697
+ print("[template] preamble relocation not applicable to this template", flush=True)
698
+ return None
699
+ probe_messages = [
700
+ {"role": "system", "content": "Your role is Builder.\n\n" + ("about text. " * 64)},
701
+ {"role": "user", "content": "hi"},
702
+ ]
703
+ probe_tools = [
704
+ {
705
+ "type": "function",
706
+ "function": {
707
+ "name": "write_file",
708
+ "description": "Write a file to the workspace.",
709
+ "parameters": {
710
+ "type": "object",
711
+ "properties": {"path": {"type": "string"}, "content": {"type": "string"}},
712
+ },
713
+ },
714
+ }
715
+ ]
716
+
717
+ def _render(template: str, kwargs: Dict[str, Any]) -> str:
718
+ return tok.apply_chat_template(
719
+ probe_messages,
720
+ tools=probe_tools,
721
+ add_generation_prompt=True,
722
+ tokenize=False,
723
+ chat_template=template,
724
+ **kwargs,
725
+ )
726
+
727
+ try:
728
+ ok, reason = template_stability.verify_prefix_stability(
729
+ _render, tok.encode, patched, original
730
+ )
731
+ except Exception as exc: # noqa: BLE001 — verification must never block boot
732
+ print(f"[template] preamble relocation verification errored: {exc}", flush=True)
733
+ return None
734
+ if not ok:
735
+ print(f"[template] preamble relocation REJECTED: {reason}", flush=True)
736
+ return None
737
+ print(f"[template] preamble relocated for prefix stability — {reason}", flush=True)
738
+ return patched
739
+
740
+
741
+ STABLE_CHAT_TEMPLATE = _resolve_stable_chat_template()
742
+
743
+
744
+ # ───────── GPU memory ceiling (Metal) ─────────
745
+ # Effective ceiling on total unified-memory this engine may use — weights + KV
746
+ # + compute. The batch admission path (BatchEngine._admit_count) reads live
747
+ # usage against it to decide how many concurrent prefills to start, so a wave
748
+ # never piles KV past the point where a Metal command-buffer allocation aborts
749
+ # the whole python process. 0 = unset (no ceiling → admit on width alone).
750
+ _MEM_LIMIT_BYTES = 0
751
+ # The resident-set target supplied by the TS capacity broker. Keep this
752
+ # separately from the allocation ceiling: mlx_lm's BatchGenerator currently
753
+ # overwrites the process wired limit with Metal's full recommended working set
754
+ # during construction, so the Gezel policy must be reapplied afterwards.
755
+ _WIRED_LIMIT_BYTES = 0
756
+ # MLX keeps freed Metal buffers for reuse in a process-global cache. That is
757
+ # useful, but an unbounded pool made a ~62 GB working set retain another
758
+ # ~45 GB until mlx_lm's coarse 512-step clear, producing a 65 -> 120 GB
759
+ # sawtooth. This cap is intentionally distinct from the prompt/KV cache budget.
760
+ _MLX_BUFFER_CACHE_LIMIT_BYTES = 0
761
+ # Throttle a multi-sub wave down to serial once live allocation reaches this
762
+ # fraction of the ceiling. Leaves headroom for the wave's own prefill/compute
763
+ # activation buffers, which land on top of already-resident weights + KV.
764
+ _MEM_ADMIT_FRACTION = 0.85
765
+
766
+
767
+ def _safe_active_memory() -> int:
768
+ """Live MLX allocation in bytes (weights + resident KV), or 0 if the
769
+ runtime doesn't expose it. Read-only; never raises."""
770
+ try:
771
+ return int(mx.get_active_memory())
772
+ except Exception: # noqa: BLE001 — introspection is best-effort
773
+ return 0
774
+
775
+
776
+ def _metal_recommended_working_set() -> int:
777
+ """Metal's `max_recommended_working_set_size` (bytes), or 0 if unavailable.
778
+ This is the OS's own soft ceiling for GPU-resident memory."""
779
+ try:
780
+ # mlx 0.32.2 (resolved by mlx-vlm 0.7.1) deprecates the old
781
+ # mx.metal.device_info() namespace. Keep the fallback so an operator's
782
+ # older, explicitly configured mlx package still gets a useful cap.
783
+ device_info = getattr(mx, "device_info", None)
784
+ info = device_info() if callable(device_info) else mx.metal.device_info()
785
+ return int(info.get("max_recommended_working_set_size", 0) or 0)
786
+ except Exception: # noqa: BLE001 — API shape varies across mlx versions
787
+ return 0
788
+
789
+
790
+ def _safe_cache_memory() -> int:
791
+ """MLX's retained free-buffer pool in bytes, or 0 if unavailable.
792
+
793
+ Distinct from active memory and invisible to it: MLX holds freed buffers
794
+ for reuse rather than returning them, so a process can sit far above its
795
+ live allocation without `get_active_memory` moving. That gap is most of
796
+ the difference between a model's working set and its physical footprint.
797
+ """
798
+ try:
799
+ return int(mx.get_cache_memory())
800
+ except Exception: # noqa: BLE001 — introspection is best-effort
801
+ return 0
802
+
803
+
804
+ def _safe_peak_memory() -> int:
805
+ """Peak MLX allocation in bytes, or 0 if unavailable."""
806
+ try:
807
+ return int(mx.get_peak_memory())
808
+ except Exception: # noqa: BLE001 — introspection is best-effort
809
+ return 0
810
+
811
+
812
+ def _log_memory(phase: str, cache_before: Optional[int] = None) -> None:
813
+ """One line of memory truth per turn.
814
+
815
+ Nothing observed this engine between startup and the next restart, so a
816
+ footprint drifting toward the wired ceiling was only ever reconstructible
817
+ afterwards from `vm_stat`. `cache` is the number that made the difference
818
+ invisible: it does not move `active`, and no other layer measures it —
819
+ the in-engine prompt-cache budget tracks its own token-count estimate.
820
+ """
821
+ active = _safe_active_memory()
822
+ cached = _safe_cache_memory()
823
+ peak = _safe_peak_memory()
824
+ if active <= 0 and cached <= 0:
825
+ return
826
+ mb = 1024 * 1024
827
+ ceiling = f", ceiling={_MEM_LIMIT_BYTES // mb}MB" if _MEM_LIMIT_BYTES > 0 else ""
828
+ wired = f", wired_limit={_WIRED_LIMIT_BYTES // mb}MB" if _WIRED_LIMIT_BYTES > 0 else ""
829
+ cache_limit = (
830
+ f", cache_limit={_MLX_BUFFER_CACHE_LIMIT_BYTES // mb}MB"
831
+ if _MLX_BUFFER_CACHE_LIMIT_BYTES > 0
832
+ else ""
833
+ )
834
+ peak_text = f" peak={peak // mb}MB" if peak > 0 else ""
835
+ reclaimed_text = (
836
+ f" cache_before={cache_before // mb}MB"
837
+ if cache_before is not None and cache_before > cached
838
+ else ""
839
+ )
840
+ print(
841
+ f"[mem] {phase}: active={active // mb}MB cache={cached // mb}MB "
842
+ f"total={(active + cached) // mb}MB{peak_text}{reclaimed_text}"
843
+ f"{ceiling}{wired}{cache_limit}",
844
+ flush=True,
845
+ )
846
+
847
+
848
+ def _apply_mlx_memory_policy(phase: str) -> None:
849
+ """Apply Gezel's process-wide MLX memory policy.
850
+
851
+ This is deliberately idempotent. In particular it is called immediately
852
+ after BatchGenerator construction because mlx_lm resets the wired limit in
853
+ its constructor. Every MLX call is guarded for older sidecar runtimes.
854
+ """
855
+ if _WIRED_LIMIT_BYTES > 0:
856
+ set_wired = getattr(mx, "set_wired_limit", None)
857
+ if set_wired is not None:
858
+ try:
859
+ set_wired(_WIRED_LIMIT_BYTES)
860
+ except Exception as exc: # noqa: BLE001
861
+ print(f"[mem] {phase} set_wired_limit skipped: {exc}", flush=True)
862
+ if _MLX_BUFFER_CACHE_LIMIT_BYTES > 0:
863
+ set_cache = getattr(mx, "set_cache_limit", None)
864
+ if set_cache is not None:
865
+ try:
866
+ set_cache(_MLX_BUFFER_CACHE_LIMIT_BYTES)
867
+ except Exception as exc: # noqa: BLE001
868
+ print(f"[mem] {phase} set_cache_limit skipped: {exc}", flush=True)
869
+
870
+
871
+ def _reclaim_mlx_buffer_cache(
872
+ phase: str, force: bool = False, log: bool = False
873
+ ) -> bool:
874
+ """Release MLX's free-buffer pool at a safe boundary or under pressure.
875
+
876
+ Active allocations (model weights and live KV state) are untouched.
877
+ `force` is used once a batch turn is fully drained; during generation we
878
+ clear only when total MLX memory crosses the admission threshold or the
879
+ runtime has temporarily exceeded the configured free-buffer cap.
880
+ """
881
+ active = _safe_active_memory()
882
+ cached = _safe_cache_memory()
883
+ under_pressure = (
884
+ _MEM_LIMIT_BYTES > 0
885
+ and active + cached >= _MEM_LIMIT_BYTES * _MEM_ADMIT_FRACTION
886
+ )
887
+ over_cache_limit = (
888
+ _MLX_BUFFER_CACHE_LIMIT_BYTES > 0
889
+ and cached > _MLX_BUFFER_CACHE_LIMIT_BYTES
890
+ )
891
+ reclaimed = False
892
+ if force or under_pressure or over_cache_limit:
893
+ clear_cache = getattr(mx, "clear_cache", None)
894
+ if clear_cache is not None:
895
+ try:
896
+ clear_cache()
897
+ reclaimed = True
898
+ except Exception as exc: # noqa: BLE001
899
+ print(f"[mem] {phase} clear_cache skipped: {exc}", flush=True)
900
+ # Turn boundaries and actual pressure events are useful telemetry. A
901
+ # healthy periodic probe should stay silent instead of writing one line
902
+ # every 32 decode steps.
903
+ if log or force or under_pressure or over_cache_limit:
904
+ _log_memory(phase, cache_before=cached if reclaimed else None)
905
+ return reclaimed
906
+
907
+
908
+ def _setup_gpu_memory_ceiling() -> None:
909
+ """Resolve the effective memory ceiling and the wired limit.
910
+
911
+ The ceiling is the smaller of what the supervisor budgeted for this model
912
+ (`--gpu-memory-limit-mb`, from the TS capacity broker) and what Metal
913
+ recommends as a working set — whichever is tighter is the honest bound.
914
+ Everything here is guarded: a version whose mlx lacks these calls must
915
+ degrade to "no ceiling" (width-only admission), never crash startup.
916
+
917
+ The WIRED limit is a separate, smaller number (`--wired-limit-mb`): the
918
+ planned resident set. Wiring the whole ceiling is what let a 27B hold
919
+ ~103 GB for a ~45 GB working set — wired pages never come back under
920
+ pressure, so this must track what steady-state inference actually needs
921
+ resident, not what the engine is allowed to allocate. Falls back to the
922
+ ceiling only when the supervisor passed no plan (unreadable geometry).
923
+ """
924
+ global _MEM_LIMIT_BYTES, _WIRED_LIMIT_BYTES, _MLX_BUFFER_CACHE_LIMIT_BYTES
925
+ recommended = _metal_recommended_working_set()
926
+ arg_limit = max(0, int(getattr(ARGS, "gpu_memory_limit_mb", 0) or 0)) * 1024 * 1024
927
+ candidates = [b for b in (recommended, arg_limit) if b > 0]
928
+ _MEM_LIMIT_BYTES = min(candidates) if candidates else 0
929
+ planned = max(0, int(getattr(ARGS, "wired_limit_mb", 0) or 0)) * 1024 * 1024
930
+ _WIRED_LIMIT_BYTES = (
931
+ min(planned, _MEM_LIMIT_BYTES)
932
+ if planned > 0 and _MEM_LIMIT_BYTES > 0
933
+ else planned or _MEM_LIMIT_BYTES
934
+ )
935
+ # Retain a small amount of free Metal storage for fast reuse, scaled to
936
+ # the planned working set but never allowed to become a second model-sized
937
+ # pool. Current 27B geometry resolves to just under 4 GB.
938
+ cache_scale = _WIRED_LIMIT_BYTES or _MEM_LIMIT_BYTES
939
+ if cache_scale > 0:
940
+ _MLX_BUFFER_CACHE_LIMIT_BYTES = max(
941
+ 256 * 1024 * 1024,
942
+ min(4 * 1024 * 1024 * 1024, cache_scale // 16),
943
+ )
944
+ else:
945
+ _MLX_BUFFER_CACHE_LIMIT_BYTES = 1024 * 1024 * 1024
946
+ _apply_mlx_memory_policy("startup")
947
+ print(
948
+ f"[mem] gpu ceiling={_MEM_LIMIT_BYTES // (1024 * 1024)}MB "
949
+ f"(budget={arg_limit // (1024 * 1024)}MB, "
950
+ f"metal_recommended={recommended // (1024 * 1024)}MB, "
951
+ f"wired={_WIRED_LIMIT_BYTES // (1024 * 1024)}MB, "
952
+ f"buffer_cache_limit={_MLX_BUFFER_CACHE_LIMIT_BYTES // (1024 * 1024)}MB, "
953
+ f"active={_safe_active_memory() // (1024 * 1024)}MB)",
954
+ flush=True,
955
+ )
956
+ # Model loading happened before this policy was known. Drop any free
957
+ # compilation/loading buffers it left behind so the first turn starts from
958
+ # the same bounded baseline as every later turn.
959
+ _reclaim_mlx_buffer_cache("startup-ready", force=True)
960
+
961
+
962
+ _setup_gpu_memory_ceiling()
963
+
964
+
965
+ def _augment_serial_stop_set(processor, model_dir: str) -> None:
966
+ """Merge generation_config.json's eos ids into the serial-path stop set.
967
+
968
+ The serial (non-batch) stream path stops on
969
+ `tokenizer.stopping_criteria.eos_token_ids`, seeded at load from the
970
+ tokenizer's own (sometimes narrower) eos set — the same gap that let
971
+ gemma4-e4b-q4 stream `<eos>` / `<|tool_response>` as text. We append
972
+ the declared ids directly (not via `add_eos_token_ids`, which expects
973
+ token *strings*). Defensive: mlx internals vary by version, so any
974
+ structural mismatch is logged and skipped rather than crashing
975
+ startup — the leaked-marker scrub still covers this path.
976
+ """
977
+ declared = _generation_config_eos_ids(model_dir) + _tool_call_close_stop_ids(model_dir)
978
+ if not declared:
979
+ return
980
+ try:
981
+ tok = getattr(processor, "tokenizer", None) or processor
982
+ criteria = getattr(tok, "stopping_criteria", None)
983
+ existing = getattr(criteria, "eos_token_ids", None)
984
+ if existing is None:
985
+ return
986
+ added = [e for e in declared if e not in existing]
987
+ existing.extend(added)
988
+ print(
989
+ f"[serial] stop set +{len(added)} from generation_config "
990
+ f"(now {len(existing)})",
991
+ flush=True,
992
+ )
993
+ except Exception as exc: # noqa: BLE001 — never fail startup over this
994
+ print(f"[serial] stop-set augment skipped: {exc}", flush=True)
995
+
996
+
997
+ _augment_serial_stop_set(PROCESSOR, ARGS.model)
998
+ print("Application startup complete.", flush=True)
999
+ print(
1000
+ f"Uvicorn running on http://{ARGS.host}:{ARGS.port} (Press CTRL+C to quit)",
1001
+ flush=True,
1002
+ )
1003
+
1004
+
1005
+ # ───────── Prompt cache layer ─────────
1006
+
1007
+ # Per-cache_id state: (cache_state, est_bytes, last_used_ts).
1008
+ # The cache_state holds both the KV cache list and the token_ids
1009
+ # prefix; mlx-vlm uses these together to decide what to prefill.
1010
+ class CacheEntry:
1011
+ __slots__ = ("state", "est_bytes", "last_used_ts", "dirty")
1012
+
1013
+ def __init__(
1014
+ self,
1015
+ state: PromptCacheState,
1016
+ est_bytes: int,
1017
+ last_used_ts: float,
1018
+ dirty: bool = True,
1019
+ ):
1020
+ self.state = state
1021
+ self.est_bytes = est_bytes
1022
+ self.last_used_ts = last_used_ts
1023
+ # True until this exact state is on disk. A shutdown flush writes only
1024
+ # dirty entries: rewriting a 90k-token session that was loaded from
1025
+ # disk unchanged is seconds of I/O, and the supervisor SIGKILLs three
1026
+ # seconds after SIGTERM — measured, the one entry that mattered was
1027
+ # still being written when the kill landed, and the restart re-prefilled
1028
+ # 55k tokens.
1029
+ self.dirty = dirty
1030
+
1031
+
1032
+ # Module-level cache. Bounded by `--cache-budget-mb` via LRU eviction.
1033
+ _CACHE: Dict[str, CacheEntry] = {}
1034
+
1035
+ # Diagnostic (temporary, Phase-0 churn census): tool-roster shape recorded
1036
+ # at save time per cache_id so a later churn event can print BOTH sides'
1037
+ # shapes. Process-lifetime only — entries loaded from disk report
1038
+ # saved_roster=unknown, and the cached-text preview covers those.
1039
+ _LAST_SAVE_ROSTER: Dict[str, str] = {}
1040
+
1041
+
1042
+ def _roster_fp_of(tools) -> str:
1043
+ """Compact fingerprint of a request's tool roster: "sha8/count", or "0"
1044
+ for an empty roster. Diagnostics only — must never raise."""
1045
+ try:
1046
+ names = []
1047
+ for t in tools or []:
1048
+ if isinstance(t, dict):
1049
+ fn = t.get("function") or {}
1050
+ names.append(str(fn.get("name") or t.get("name") or "?"))
1051
+ else:
1052
+ fn = getattr(t, "function", None)
1053
+ names.append(str(getattr(fn, "name", None) or getattr(t, "name", "?")))
1054
+ if not names:
1055
+ return "0"
1056
+ digest = hashlib.sha256("\n".join(sorted(names)).encode("utf-8")).hexdigest()[:8]
1057
+ return f"{digest}/{len(names)}"
1058
+ except Exception: # noqa: BLE001 — diagnostics never break a turn
1059
+ return "err"
1060
+
1061
+
1062
+ # Cache ids (session + its seed prefix) currently being SERVED — registered
1063
+ # at the top of a chat-completions request and cleared when its stream ends.
1064
+ # Eviction (both in-memory `_enforce_budget` and `_enforce_disk_budget`) skips
1065
+ # these so the very entry a live turn is about to reuse can't be dropped out
1066
+ # from under it mid-flight. Without this pin, a long turn on a busy multi-
1067
+ # session box can have its own (least-recently-used, because it's been
1068
+ # streaming) prefix evicted, forcing a cold full re-prefill on the next step
1069
+ # — exactly the kind of churn that compounds the 27B first-token stall.
1070
+ # Multiset-counted (a value > 0 means pinned) so overlapping requests that
1071
+ # share a prefix id don't unpin each other early.
1072
+ _INFLIGHT_CACHE_IDS: "collections.Counter[str]" = collections.Counter()
1073
+
1074
+
1075
+ def _mark_inflight(*cache_ids: Optional[str]) -> List[str]:
1076
+ """Pin the given (non-empty) cache ids from eviction. Returns the list
1077
+ actually pinned, to hand to {@link _unmark_inflight} in a finally."""
1078
+ pinned: List[str] = []
1079
+ for cid in cache_ids:
1080
+ if cid:
1081
+ _INFLIGHT_CACHE_IDS[cid] += 1
1082
+ pinned.append(cid)
1083
+ return pinned
1084
+
1085
+
1086
+ def _unmark_inflight(cache_ids: List[str]) -> None:
1087
+ # `.get` / `.pop(default)` rather than `[cid]` / `del`: a Counter returns 0
1088
+ # for an absent key but `del` on it raises KeyError, so a stray double-
1089
+ # release (or release of something never marked) must not crash the worker.
1090
+ for cid in cache_ids:
1091
+ count = _INFLIGHT_CACHE_IDS.get(cid, 0)
1092
+ if count <= 1:
1093
+ _INFLIGHT_CACHE_IDS.pop(cid, None)
1094
+ else:
1095
+ _INFLIGHT_CACHE_IDS[cid] = count - 1
1096
+
1097
+
1098
+ # Serializes the actual `stream_generate` GPU work across concurrent
1099
+ # requests. MLX text-only generation isn't safely batched in our wrapper
1100
+ # (no shared-prefix batching path, KV state lives on one Metal stream),
1101
+ # so two simultaneous chat completions would race the shared model
1102
+ # state and the per-cache_id KV save/restore. The TS-side ProviderQueue
1103
+ # now defaults to concurrency=2 to break the `ask_specialist`/`ask_gezel`
1104
+ # deadlock; this lock guarantees the engine still sees one stream at a
1105
+ # time. The lock is only held across the `for chunk in stream_generate(...)`
1106
+ # loop — cache lookup and cache save run outside it so the second request
1107
+ # can start its prefill prep while the first is still streaming.
1108
+ #
1109
+ # NOTE: must be constructed inside an event loop. FastAPI's lifespan
1110
+ # handles that — the lock is lazily allocated on first request.
1111
+ _GENERATION_LOCK: Optional[asyncio.Lock] = None
1112
+
1113
+
1114
+ def _get_generation_lock() -> asyncio.Lock:
1115
+ global _GENERATION_LOCK
1116
+ if _GENERATION_LOCK is None:
1117
+ _GENERATION_LOCK = asyncio.Lock()
1118
+ return _GENERATION_LOCK
1119
+
1120
+
1121
+ def _kv_quant_kwargs() -> Dict[str, Any]:
1122
+ """KV-quantization options for the serial `stream_generate` path.
1123
+
1124
+ Forwarded through `stream_generate(**kwargs)` into mlx-vlm's
1125
+ `generate_step`, which converts each layer via `to_quantized` once it
1126
+ passes `quantized_kv_start`. Layer classes without that method
1127
+ (BatchKVCache, ChunkedKVCache) are skipped by `maybe_quantize_kv_cache`
1128
+ rather than raising — safe, but silent, which is why the launcher plans
1129
+ f16 whenever the batched path is in use instead of trusting this flag.
1130
+
1131
+ Empty when disabled, so the default path is byte-identical to before.
1132
+ """
1133
+ if ARGS.kv_bits <= 0:
1134
+ return {}
1135
+ return {
1136
+ "kv_bits": ARGS.kv_bits,
1137
+ "kv_group_size": ARGS.kv_group_size,
1138
+ "kv_quant_scheme": ARGS.kv_quant_scheme,
1139
+ "quantized_kv_start": ARGS.quantized_kv_start,
1140
+ }
1141
+
1142
+
1143
+ def _budget_bytes() -> int:
1144
+ return ARGS.cache_budget_mb * 1024 * 1024
1145
+
1146
+
1147
+ # Fallback per-token byte estimate, used only when a state's layers don't
1148
+ # expose `nbytes` (every mlx_lm/mlx_vlm cache class does today, via
1149
+ # _BaseCache). The old blanket 100 KB/token figure measured ~55% high for
1150
+ # qwen3.8's 16 full-attention layers (64 KiB/token) and was blind to the
1151
+ # fixed ~147 MiB recurrent state of its 48 linear-attention layers; real
1152
+ # per-layer `nbytes` (shape math, no graph eval) prices both correctly.
1153
+ _BYTES_PER_TOKEN = 100 * 1024
1154
+
1155
+
1156
+ def _estimate_cache_bytes(state: PromptCacheState) -> int:
1157
+ if state.token_ids is None:
1158
+ return 0
1159
+ layers = list(getattr(state, "cache", None) or [])
1160
+ total = 0
1161
+ complete = bool(layers)
1162
+ try:
1163
+ for layer in layers:
1164
+ nb = getattr(layer, "nbytes", None)
1165
+ if nb is None:
1166
+ complete = False
1167
+ break
1168
+ total += int(nb)
1169
+ except Exception: # noqa: BLE001 — budget math must never break a turn
1170
+ complete = False
1171
+ if complete and total > 0:
1172
+ return total
1173
+ return len(state.token_ids) * _BYTES_PER_TOKEN
1174
+
1175
+
1176
+ def _enforce_budget() -> None:
1177
+ """LRU eviction down to <= 80% of budget when over budget.
1178
+
1179
+ Two safety rails on top of the basic LRU:
1180
+ 1. NEVER evict the only entry. Evicting it buys nothing — the
1181
+ next chat-completions request will immediately re-prefill it
1182
+ back to the same size. Single-session users would loop
1183
+ forever paying full prefill cost otherwise.
1184
+ 2. Stop evicting once one entry remains, regardless of budget.
1185
+ An over-budget single entry is a signal that the budget is
1186
+ too tight for this model; keeping the entry preserves cache
1187
+ reuse across turns. Operators see usage exceeding budget
1188
+ in /v1/cache/stats and can tune up.
1189
+
1190
+ When disk persistence is enabled, we save each evicted entry before
1191
+ dropping it. The next session that asks for the same cache_id will
1192
+ load it back from disk (cheap mmap read) instead of paying the
1193
+ full prefill cost again.
1194
+ """
1195
+ total = sum(e.est_bytes for e in _CACHE.values())
1196
+ budget = _budget_bytes()
1197
+ if total < budget:
1198
+ return
1199
+ if len(_CACHE) <= 1:
1200
+ return # don't evict the only entry — cache reuse needs it
1201
+ target = int(budget * 0.8)
1202
+ # Sort by lru ascending; pop oldest until under target OR only one entry left.
1203
+ items = sorted(_CACHE.items(), key=lambda kv: kv[1].last_used_ts)
1204
+ for cache_id, entry in items:
1205
+ if total < target or len(_CACHE) <= 1:
1206
+ break
1207
+ # Never evict an entry a live turn is currently serving — dropping
1208
+ # it forces that turn into a cold re-prefill mid-flight.
1209
+ if cache_id in _INFLIGHT_CACHE_IDS:
1210
+ continue
1211
+ # Persist before dropping. Best-effort — failure here just means
1212
+ # the next session pays the prefill cost, which is the
1213
+ # pre-persistence behavior. A clean entry is already on disk.
1214
+ if _disk_enabled() and entry.dirty:
1215
+ cache_persist.save_cache(
1216
+ PERSIST_DIR, MODEL_FINGERPRINT, cache_id, entry.state # type: ignore[arg-type]
1217
+ )
1218
+ del _CACHE[cache_id]
1219
+ total -= entry.est_bytes
1220
+ print(f"[cache] evicted cache_id={cache_id} (LRU, in-engine budget)", flush=True)
1221
+ # Disk-budget enforcement runs independently of in-memory eviction:
1222
+ # we may have just added new on-disk entries that push us over the
1223
+ # disk cap. LRU prune by mtime keeps total disk usage bounded.
1224
+ _enforce_disk_budget()
1225
+
1226
+
1227
+ def _enforce_disk_budget() -> None:
1228
+ """Prune the oldest persisted cache files when over disk budget.
1229
+
1230
+ Off-by-default-when-zero: operators on a machine with abundant SSD
1231
+ space can set --disk-cache-budget-mb=0 and let the OS reclaim
1232
+ however much it needs. The default budget (8 GB) is generous for
1233
+ the typical 25K-token Meester session (~2.5 GB) plus several
1234
+ smaller sibling sessions.
1235
+ """
1236
+ if not _disk_enabled():
1237
+ return
1238
+ if ARGS.disk_cache_budget_mb <= 0:
1239
+ return
1240
+ budget_bytes = ARGS.disk_cache_budget_mb * 1024 * 1024
1241
+ entries = cache_persist.list_persisted(PERSIST_DIR, MODEL_FINGERPRINT) # type: ignore[arg-type]
1242
+ total = sum(size for (_, size, _) in entries)
1243
+ if total < budget_bytes:
1244
+ return
1245
+ target = int(budget_bytes * 0.8)
1246
+ # Never prune the newest file. Same rail as the in-memory budget's "never
1247
+ # evict the only entry": a single long session can outgrow the cap on its
1248
+ # own (a 120k-token qwen3.8-27b snapshot is ~7.7 GB against the 8 GB
1249
+ # default), and deleting the file a shutdown flush wrote a moment ago
1250
+ # makes that flush — the only thing standing between a restart and a
1251
+ # full re-prefill — pointless.
1252
+ newest = entries[-1][0] if entries else None
1253
+ for cache_id, size, _mtime in entries:
1254
+ if total < target:
1255
+ break
1256
+ if cache_id == newest:
1257
+ continue
1258
+ # Skip the on-disk copy of an entry a live turn is serving — the
1259
+ # in-memory pin above keeps it resident, but a sibling turn that
1260
+ # disk-reloads the same prefix mid-flight must still find the file.
1261
+ if cache_id in _INFLIGHT_CACHE_IDS:
1262
+ continue
1263
+ if cache_persist.delete_cache(PERSIST_DIR, MODEL_FINGERPRINT, cache_id): # type: ignore[arg-type]
1264
+ total -= size
1265
+ print(f"[cache] disk-evicted cache_id={cache_id} (LRU, disk budget)", flush=True)
1266
+
1267
+
1268
+ def _save_cache(cache_id: str, state: PromptCacheState) -> None:
1269
+ bytes_est = _estimate_cache_bytes(state)
1270
+ _CACHE[cache_id] = CacheEntry(state=state, est_bytes=bytes_est, last_used_ts=time.time())
1271
+ _enforce_budget()
1272
+
1273
+
1274
+ def _try_load_from_disk(cache_id: str) -> Optional[PromptCacheState]:
1275
+ """Pull a cache entry off disk into _CACHE under the given id.
1276
+
1277
+ Returns the loaded state on hit (also installs in _CACHE), None on
1278
+ miss / corruption / fingerprint mismatch. The next save will write
1279
+ back through the normal path.
1280
+ """
1281
+ if not _disk_enabled():
1282
+ return None
1283
+ state = cache_persist.load_cache(PERSIST_DIR, MODEL_FINGERPRINT, cache_id) # type: ignore[arg-type]
1284
+ if state is None:
1285
+ return None
1286
+ bytes_est = _estimate_cache_bytes(state)
1287
+ _CACHE[cache_id] = CacheEntry(
1288
+ state=state, est_bytes=bytes_est, last_used_ts=time.time(), dirty=False
1289
+ )
1290
+ print(
1291
+ f"[cache] disk-load cache_id={cache_id} tokens={len(state.token_ids or [])}",
1292
+ flush=True,
1293
+ )
1294
+ return state
1295
+
1296
+
1297
+ def _try_seed_from_prefix(cache_id: str, prefix_cache_id: str) -> Optional[PromptCacheState]:
1298
+ """When a session misses, fall back to its gezel-prefix entry.
1299
+
1300
+ Loads the prefix-shaped entry from disk and registers it under
1301
+ `cache_id` so the session inherits it as its starting state. The
1302
+ prefix file is never overwritten by the session — its own delta
1303
+ saves under its own cache_id. This is how new sessions for an
1304
+ existing gezel skip the system-prompt prefill cost.
1305
+
1306
+ Returns None on any miss; the caller falls through to fresh state.
1307
+ """
1308
+ if not _disk_enabled():
1309
+ return None
1310
+ state = cache_persist.load_cache(
1311
+ PERSIST_DIR, MODEL_FINGERPRINT, prefix_cache_id # type: ignore[arg-type]
1312
+ )
1313
+ if state is None:
1314
+ return None
1315
+ bytes_est = _estimate_cache_bytes(state)
1316
+ # Clean: this is the prefix file's content, already on disk under the
1317
+ # prefix id. Persisting a copy under the session id buys nothing.
1318
+ _CACHE[cache_id] = CacheEntry(
1319
+ state=state, est_bytes=bytes_est, last_used_ts=time.time(), dirty=False
1320
+ )
1321
+ print(
1322
+ f"[cache] prefix-seed cache_id={cache_id} from prefix={prefix_cache_id} "
1323
+ f"tokens={len(state.token_ids or [])}",
1324
+ flush=True,
1325
+ )
1326
+ return state
1327
+
1328
+
1329
+ def _request_prefix_ids(request) -> List[str]:
1330
+ """Ordered, most-specific-first prefix ids for this request. Prefers the
1331
+ layered `prefix_cache_ids`; falls back to the singular `prefix_cache_id`
1332
+ for back-compat. Filters out empties."""
1333
+ if getattr(request, "prefix_cache_ids", None):
1334
+ return [p for p in request.prefix_cache_ids if p]
1335
+ if request.prefix_cache_id:
1336
+ return [request.prefix_cache_id]
1337
+ return []
1338
+
1339
+
1340
+ # Prefix ids we've already seeded from a real session save this process.
1341
+ # Mirrors the llama-cpp adapter's `seededPrefixes` — seed once per process
1342
+ # so back-to-back sessions don't repeatedly rewrite the same prefix file.
1343
+ _SEEDED_PREFIXES: set = set()
1344
+
1345
+
1346
+ def _band_snapshot_target(sub) -> Optional[int]:
1347
+ """Token index of the shared band's end for this request, or None.
1348
+
1349
+ The TS side knows the band as a CHARACTER length inside the system
1350
+ message (`stable_prefix_chars`); only the engine has the tokenizer, and
1351
+ the rendered prompt carries template framing plus the tool block ahead
1352
+ of the system content — so the offset is located by finding the first
1353
+ token index whose decode contains the band's tail. See ADR 0010.
1354
+ """
1355
+ n_chars = getattr(sub.request, "stable_prefix_chars", None)
1356
+ if not n_chars or int(n_chars) <= 0:
1357
+ return None
1358
+ system = next(
1359
+ (m for m in (getattr(sub.request, "messages", None) or [])
1360
+ if getattr(m, "role", None) == "system"),
1361
+ None,
1362
+ )
1363
+ content = getattr(system, "content", None) if system is not None else None
1364
+ if not isinstance(content, str):
1365
+ return None
1366
+ marker = content[: int(n_chars)][-_BAND_MARKER_CHARS:]
1367
+ if len(marker) < _BAND_MARKER_MIN_CHARS:
1368
+ return None
1369
+ tok = getattr(PROCESSOR, "tokenizer", None) or PROCESSOR
1370
+ toks = list(sub.prompt_tokens or [])
1371
+ try:
1372
+ target = cache_seed.token_boundary_for_marker(
1373
+ lambda k: tok.decode(toks[:k]), len(toks), marker, _SNAPSHOT_BOUNDARY_MARGIN
1374
+ )
1375
+ except Exception as exc: # noqa: BLE001 — best-effort; fall back to today
1376
+ print(f"[cache] band-boundary probe failed: {exc}", flush=True)
1377
+ return None
1378
+ # One line naming the decision — the thing to grep when a sibling
1379
+ # re-prefills instead of extending. `target=0` means the band's tail was
1380
+ # not found in the render, which is a marker problem, not a cache one.
1381
+ print(
1382
+ f"[cache] band-boundary chars={int(n_chars)} target={target} "
1383
+ f"prompt_tokens={len(toks)}",
1384
+ flush=True,
1385
+ )
1386
+ return target or None
1387
+
1388
+
1389
+ def _seed_prefix_from_session(request, state: PromptCacheState) -> None:
1390
+ """Seed the most-specific (gp) prefix entry from a real session's saved
1391
+ state so a SIBLING session (same gezel+project, different volatile state)
1392
+ inherits the warm `[system+tools]` prefix on its next turn — the
1393
+ cross-session win the layered flag exists for.
1394
+
1395
+ Mirrors the llama-cpp adapter, which copies a slot save into a
1396
+ `prefix-<hash>.bin` file. The seeded entry holds this session's full
1397
+ state (system+tools+transcript); a sibling's longest-common-prefix trim
1398
+ reuses only the shared `[system+tools]` head and re-prefills its own
1399
+ divergent tail. Best-effort, once per prefix per process; never blocks
1400
+ the request path on failure.
1401
+ """
1402
+ if not _disk_enabled():
1403
+ return
1404
+ ids = _request_prefix_ids(request)
1405
+ if not ids:
1406
+ return
1407
+ gp = ids[0]
1408
+ if gp in _SEEDED_PREFIXES:
1409
+ return
1410
+ try:
1411
+ # INVARIANT: a prefix entry may never be replaced by a LONGER state.
1412
+ # On this model family `seed_from_state` reuses an entry only when it
1413
+ # is a strict prefix (`lcp == n` -> extension); an entry longer than
1414
+ # the shared head cannot be trimmed (48 of 64 layers are
1415
+ # linear-attention `ArraysCache`), so it degrades every sibling to a
1416
+ # full re-prefill. Lengthening is therefore never an improvement, and
1417
+ # a shorter entry that already works must survive a later process
1418
+ # that would otherwise overwrite it with a whole session. Measured:
1419
+ # a 7,202-token entry that delivered 89.8% reuse was overwritten by
1420
+ # the same session's 8,008-token post-turn state moments later.
1421
+ incoming = len(state.token_ids or [])
1422
+ existing = cache_persist.load_cache(PERSIST_DIR, MODEL_FINGERPRINT, gp)
1423
+ prior = len(getattr(existing, "token_ids", None) or []) if existing else 0
1424
+ if prior and incoming > prior:
1425
+ _SEEDED_PREFIXES.add(gp)
1426
+ print(
1427
+ f"[cache] prefix-keep {gp}: existing {prior} tokens is shorter "
1428
+ f"than this session's {incoming} — a longer entry cannot be "
1429
+ f"trimmed by a sibling",
1430
+ flush=True,
1431
+ )
1432
+ return
1433
+ if cache_persist.save_cache(PERSIST_DIR, MODEL_FINGERPRINT, gp, state): # type: ignore[arg-type]
1434
+ _SEEDED_PREFIXES.add(gp)
1435
+ _LAST_SAVE_ROSTER[gp] = _roster_fp_of(getattr(request, "tools", None))
1436
+ _enforce_disk_budget()
1437
+ print(
1438
+ f"[cache] prefix-seeded {gp} from cache_id={request.cache_id} "
1439
+ f"tokens={len(state.token_ids or [])}",
1440
+ flush=True,
1441
+ )
1442
+ except Exception as exc: # noqa: BLE001 — best-effort, never fail the turn
1443
+ print(f"[cache] prefix-seed failed for {gp}: {exc}", flush=True)
1444
+
1445
+
1446
+ # ───────── BatchGenerator engine (--max-concurrency >= 1) ─────────
1447
+ #
1448
+ # Any positive concurrency routes through BatchEngine instead of the legacy
1449
+ # serial `stream_generate` + _GENERATION_LOCK path. At N=1 this is still a
1450
+ # singleton queue, but it matters for correctness/performance: BatchGenerator
1451
+ # exposes prompt-segment boundaries, which lets windowed Qwen/Gemma caches save
1452
+ # a coherent prompt snapshot instead of cold-prefilling after every tool-loop
1453
+ # turn. At N>1 the same engine also continuously batches independent sequences.
1454
+ # One asyncio worker drives the generator on the event-loop thread — MLX GPU
1455
+ # streams are thread-local (see the chat_completions note). Every per-session
1456
+ # concern maps to a per-sequence list in the BatchGenerator API (sampler, logits
1457
+ # processors incl. the tool grammar, seed KV cache, max_tokens); on finish we
1458
+ # extract the sequence's KV and persist it via the SAME _save_cache /
1459
+ # _seed_prefix_from_session path the serial route uses.
1460
+
1461
+
1462
+ def _resolve_cache_state(request):
1463
+ """The cache cascade (memory → disk → prefix-seed → fresh), shared by
1464
+ the serial and batched paths. Returns (state, source, matched_prefix)."""
1465
+ cache_state = None
1466
+ cache_source = "fresh"
1467
+ matched_prefix = None
1468
+ if request.cache_id:
1469
+ existing_entry = _CACHE.get(request.cache_id)
1470
+ if existing_entry is not None:
1471
+ cache_state = existing_entry.state
1472
+ existing_entry.last_used_ts = time.time()
1473
+ cache_source = "memory"
1474
+ else:
1475
+ disk_state = _try_load_from_disk(request.cache_id)
1476
+ if disk_state is not None:
1477
+ cache_state = disk_state
1478
+ cache_source = "disk"
1479
+ else:
1480
+ for pid in _request_prefix_ids(request):
1481
+ seeded = _try_seed_from_prefix(request.cache_id, pid)
1482
+ if seeded is not None:
1483
+ cache_state = seeded
1484
+ cache_source = "prefix"
1485
+ matched_prefix = pid
1486
+ break
1487
+ if cache_state is None:
1488
+ cache_state = PromptCacheState()
1489
+ return cache_state, cache_source, matched_prefix
1490
+
1491
+
1492
+ class _UnwrapLM:
1493
+ """mlx_vlm's LanguageModel returns LanguageModelOutput(.logits); mlx_lm's
1494
+ BatchGenerator wants a raw logits array. Thin shim: unwrap `.logits` and
1495
+ delegate every other attribute (layers, config, make_cache, …) to the
1496
+ wrapped model."""
1497
+
1498
+ def __init__(self, inner):
1499
+ self._inner = inner
1500
+ self.layers = inner.layers
1501
+
1502
+ def __call__(self, *args, **kwargs):
1503
+ out = self._inner(*args, **kwargs)
1504
+ return getattr(out, "logits", out)
1505
+
1506
+ def __getattr__(self, name):
1507
+ return getattr(self._inner, name)
1508
+
1509
+
1510
+ # Batched-prefill liveness tuning. Only emit a marker for waves big enough
1511
+ # that prefill is non-trivial (small prompts return a first token fast, so a
1512
+ # marker would just be noise), and no more often than the interval so a
1513
+ # chunk-iterating worker loop doesn't spam the log.
1514
+ _PREFILL_LIVENESS_MIN_TOKENS = 2048
1515
+
1516
+ # Event-loop breathing between engine steps — see BatchEngine._breathe.
1517
+ _BREATHE_EVERY_S = 0.2
1518
+ _BREATHE_FOR_S = 0.002
1519
+ _PREFILL_LIVENESS_INTERVAL_S = 4.0
1520
+
1521
+ # Keep chunk-edge snapshots this many tokens clear of the prompt's end.
1522
+ # The final tokens (`…<start_of_turn>model\n`) re-merge differently once
1523
+ # the next turn's text continues past them (BPE boundary instability), so
1524
+ # a snapshot cut too close to the end misses by a couple of tokens — and
1525
+ # wrapped rotating layers can't trim even that. 16 is generous; the cost
1526
+ # is at most 16 re-prefilled tokens per turn.
1527
+ _SNAPSHOT_BOUNDARY_MARGIN = 16
1528
+ # Tail of the shared band used to locate its token boundary. Long enough to
1529
+ # be unique in a 60k-char prompt, short enough that one BPE re-merge at the
1530
+ # cut cannot swallow it.
1531
+ _BAND_MARKER_CHARS = 256
1532
+ _BAND_MARKER_MIN_CHARS = 64
1533
+
1534
+
1535
+ def _prefix_seed_state(sub, session_state, trimmable: bool):
1536
+ """The state to publish as this request's prefix entry, or None.
1537
+
1538
+ A band-keyed entry (ADR 0010) must be the band itself. A longer
1539
+ untrimmable entry is not a partial hit for a sibling but a full
1540
+ re-prefill, so when the band snapshot is missing — never planted,
1541
+ or dropped by the boundary proof — nothing is published. A trimmable
1542
+ stack keeps publishing its post-turn state: siblings trim it back.
1543
+ """
1544
+ if sub.band_snapshot is not None:
1545
+ layers, count = sub.band_snapshot
1546
+ state = PromptCacheState()
1547
+ state.cache = layers
1548
+ state.token_ids = list(sub.prompt_tokens[:count])
1549
+ return state
1550
+ if getattr(sub.request, "stable_prefix_chars", None) and not trimmable:
1551
+ return None
1552
+ return session_state
1553
+
1554
+
1555
+ def _store_snapshot(sub, snapshot) -> None:
1556
+ """File a verified boundary snapshot under the slot(s) its boundary is:
1557
+ the session's own save, the shared band, or both when they coincide."""
1558
+ at_tokens = snapshot[1]
1559
+ if at_tokens == sub.snapshot_target:
1560
+ sub.prompt_snapshot = snapshot
1561
+ if at_tokens == sub.band_target:
1562
+ sub.band_snapshot = snapshot
1563
+
1564
+
1565
+ class _Sub:
1566
+ """One in-flight batched request: its decode queue, accumulated tokens,
1567
+ and the bits needed to build its sequence + persist its KV on finish."""
1568
+
1569
+ __slots__ = (
1570
+ "request", "prompt_tokens", "grammar", "seed_state", "request_id",
1571
+ "http_request", "uid", "queue", "cancelled", "token_ids", "emitted",
1572
+ "pinned_cache_ids", "prefill_total", "first_token_seen",
1573
+ "reused_tokens", "seed_mode", "prompt_snapshot", "snapshot_target",
1574
+ "band_target", "band_snapshot", "saved_cache_state",
1575
+ "prefill_started_at", "generation_started_at",
1576
+ "prompt_tps", "generation_tps", "roster_fp",
1577
+ "priority", "queued_at",
1578
+ )
1579
+
1580
+ def __init__(
1581
+ self,
1582
+ request,
1583
+ prompt_tokens,
1584
+ grammar,
1585
+ seed_state,
1586
+ request_id,
1587
+ http_request,
1588
+ snapshot_target=None,
1589
+ roster_fp=None,
1590
+ ):
1591
+ self.request = request
1592
+ self.prompt_tokens = prompt_tokens # List[int], full prompt
1593
+ self.grammar = grammar # SafeToolGrammarProcessor | None
1594
+ self.seed_state = seed_state # PromptCacheState from the cascade
1595
+ self.request_id = request_id
1596
+ self.http_request = http_request
1597
+ # Diagnostic (temporary, Phase-0 churn census): the tool-roster
1598
+ # shape this sub's prompt was rendered with. The warm path passes
1599
+ # it explicitly because its internal ChatRequest omits tools.
1600
+ self.roster_fp = (
1601
+ roster_fp
1602
+ if roster_fp is not None
1603
+ else _roster_fp_of(getattr(request, "tools", None))
1604
+ )
1605
+ self.uid = None
1606
+ self.queue: asyncio.Queue = asyncio.Queue()
1607
+ self.cancelled = False
1608
+ self.token_ids = [] # generated token ids
1609
+ self.emitted = 0 # chars already streamed
1610
+ self.pinned_cache_ids: List[str] = [] # eviction pins to release on stream end
1611
+ self.prefill_total = 0 # new tokens this sub must prefill
1612
+ self.first_token_seen = False # flips once decode emits a token
1613
+ self.reused_tokens = 0 # cache tokens actually served (seed plan)
1614
+ self.seed_mode = "fresh" # seed decision, for logs/diagnosis
1615
+ # End-of-prompt KV snapshot (windowed models only) — captured at
1616
+ # first token, before any generated token enters the cache, so
1617
+ # the saved session entry stays a pure prefix of the next turn's
1618
+ # re-templated prompt. See cache_seed.probe_needs_prompt_snapshot.
1619
+ self.prompt_snapshot = None
1620
+ # Absolute token boundary to capture. Normal turns default to
1621
+ # len(prompt)-margin at admission; prefix warming supplies the stable
1622
+ # boundary found by comparing two different synthetic user turns.
1623
+ self.snapshot_target = snapshot_target
1624
+ # Shared-band boundary (ADR 0010) and its snapshot, captured beside
1625
+ # the session's own so a fresh turn can publish the short band as the
1626
+ # sibling prefix entry without shrinking its own saved cache to it.
1627
+ self.band_target = None
1628
+ self.band_snapshot = None
1629
+ # The exact state committed at finish. Cache warming uses this direct
1630
+ # reference for immediate persistence even if the memory-budget pass
1631
+ # evicts the just-built entry from the global LRU.
1632
+ self.saved_cache_state = None
1633
+ # BatchGenerator does not expose per-sequence timing. Track wall-clock
1634
+ # timing per subscription so the OpenAI-compatible usage frames retain
1635
+ # the same prompt/decode speed fields as the serial stream_generate
1636
+ # path. Values are snapshotted into queue items because the worker can
1637
+ # advance again before the SSE coroutine consumes them.
1638
+ self.prefill_started_at = None
1639
+ self.generation_started_at = None
1640
+ self.prompt_tps = 0.0
1641
+ self.generation_tps = 0.0
1642
+ # Who is waiting on this request: "interactive" (a person) or
1643
+ # "background" (task steps, warms, chores). Decides wave membership
1644
+ # and whether a running wave steps aside — see wave_policy.py.
1645
+ self.priority = wave_policy.normalize_priority(getattr(request, "priority", None))
1646
+ self.queued_at = time.time()
1647
+
1648
+
1649
+ class BatchEngine:
1650
+ """Owns one mlx_lm BatchGenerator + the single worker that drives it."""
1651
+
1652
+ def __init__(self, model, tokenizer, max_concurrency):
1653
+ from mlx_lm.generate import BatchGenerator # local: only when batching
1654
+
1655
+ self._tok = tokenizer
1656
+ self._max = max(1, int(max_concurrency))
1657
+ eos = getattr(tokenizer, "eos_token_ids", None)
1658
+ if eos is None:
1659
+ eos = getattr(tokenizer, "eos_token_id", None)
1660
+ if eos is None:
1661
+ eos_list = []
1662
+ elif isinstance(eos, int):
1663
+ eos_list = [eos]
1664
+ else:
1665
+ eos_list = list(eos)
1666
+ # Union in the stop set the model author declared in
1667
+ # generation_config.json — the tokenizer wrapper sometimes
1668
+ # resolves a narrower eos set than the model actually ends turns
1669
+ # on (see _generation_config_eos_ids).
1670
+ eos_from_cfg = _generation_config_eos_ids(ARGS.model)
1671
+ for e in eos_from_cfg:
1672
+ if e not in eos_list:
1673
+ eos_list.append(e)
1674
+ # Halt at the first closed `<tool_call>` for models that declare a
1675
+ # server-side tool_call_parser but don't list `</tool_call>` in eos
1676
+ # (poolside Laguna) — otherwise the raw completion path never stops
1677
+ # after a call and loops hundreds of them per turn until the SSE
1678
+ # stalls. See _tool_call_close_stop_ids.
1679
+ tool_close = _tool_call_close_stop_ids(ARGS.model)
1680
+ for e in tool_close:
1681
+ if e not in eos_list:
1682
+ eos_list.append(e)
1683
+ self._stop_ids = {int(e) for e in eos_list}
1684
+ stop_tokens = [[int(e)] for e in eos_list]
1685
+ self._model = model
1686
+ self._gen_kwargs = dict(
1687
+ max_tokens=2048,
1688
+ stop_tokens=stop_tokens or None,
1689
+ completion_batch_size=self._max,
1690
+ prefill_batch_size=self._max,
1691
+ prefill_step_size=ARGS.prefill_step_size,
1692
+ )
1693
+ self._gen = self._make_generator("batch-init")
1694
+ self._subs = {} # uid -> _Sub
1695
+ self._pending = [] # _Sub awaiting insert
1696
+ self._wake = asyncio.Event()
1697
+ self._worker = None
1698
+ # Preemption state (wave_policy.py). A wave made only of background
1699
+ # work parks while waiting interactive requests run on a second
1700
+ # generator; the parked wave keeps its KV and resumes where it stopped.
1701
+ self._spare_gen = None # second BatchGenerator, built on first preemption
1702
+ self._nested = False # True while an interactive wave runs over a parked one
1703
+ self._parked = [] # subs of the parked wave (paused markers)
1704
+ self._held_background = [] # background arrivals while nested
1705
+ self._spec_sub = None # the sub a spec wave is serving
1706
+ # Set while a spec wave's rounds own the drafter's per-wave state
1707
+ # (its cache, shared kv, position); a wave run inside the pause must
1708
+ # not draft with the same drafter.
1709
+ self._drafter_busy = False
1710
+ self._liveness = wave_policy.LivenessThrottle()
1711
+ self._preempt_memory_noted = False
1712
+ self._last_breath = 0.0
1713
+ # Windowed / state-cache models can't rewind their caches once
1714
+ # the sequence outgrows the window, so a post-generation save is
1715
+ # dead weight the moment history diverges from the raw tokens
1716
+ # (reasoning stripping guarantees it does). For those models we
1717
+ # capture an end-of-prompt KV snapshot per turn instead — see
1718
+ # _maybe_snapshot_prompt_cache. Full-attention stacks skip the
1719
+ # copy; LCP+trim covers them.
1720
+ self._needs_snapshot = cache_seed.probe_needs_prompt_snapshot(model)
1721
+ self._snapshot_warned = False
1722
+ # Prefill-liveness for the current wave. The batched path prefills the
1723
+ # whole prompt inside BatchGenerator with NO tqdm bar (unlike the
1724
+ # serial mlx_vlm path), so the TS watchdog saw zero prefill events and
1725
+ # the engine looked silent through a long first-token wait (the
1726
+ # qwen3.6-27b 37K-token stall). We emit a `Prefill:`-shaped marker the
1727
+ # stdout-parser already understands so the watchdog knows the engine is
1728
+ # alive + prefilling N tokens, and the UI pill shows it.
1729
+ self._prefill_total = 0 # new tokens the active wave is prefilling
1730
+ self._prefill_done = {} # uid → tokens prefilled so far (this wave)
1731
+ # uid → (cache_id, that sub's own prefill total). The engine has ONE
1732
+ # stdout, so an untagged `Prefill:` line is broadcast by the supervisor
1733
+ # to every in-flight session — a session streaming its reply would draw
1734
+ # a neighbour's prefill bar on its own row, at that neighbour's token
1735
+ # count. This map is what makes each marker addressable.
1736
+ self._prefill_meta = {}
1737
+ self._last_prefill_emit = 0.0
1738
+ self._memory_check_steps = 0
1739
+ print(
1740
+ f"[batch] engine ready max_concurrency={self._max} "
1741
+ f"stop_tokens={len(stop_tokens)} (+{len(eos_from_cfg)} from generation_config"
1742
+ f"{f', +{len(tool_close)} </tool_call>' if tool_close else ''}) "
1743
+ f"prompt_snapshot={'y' if self._needs_snapshot else 'n'}",
1744
+ flush=True,
1745
+ )
1746
+
1747
+ def _make_generator(self, phase):
1748
+ from mlx_lm.generate import BatchGenerator # local: only when batching
1749
+
1750
+ gen = BatchGenerator(self._model, **self._gen_kwargs)
1751
+ # mlx_lm BatchGenerator sets the wired limit to Metal's entire
1752
+ # recommended working set in its constructor. Restore the smaller
1753
+ # resident-set plan and free-buffer cap selected by Gezel.
1754
+ _apply_mlx_memory_policy(phase)
1755
+ return gen
1756
+
1757
+ async def _breathe(self):
1758
+ """Yield to the event loop between engine steps.
1759
+
1760
+ `sleep(0)` yields exactly one loop iteration, and a NEW request needs
1761
+ several (accept, parse, route, read body, render, tokenize, submit)
1762
+ before the engine ever sees it. Each engine step blocks the loop for
1763
+ its whole duration — a 2048-token prefill chunk is seconds on a 27B —
1764
+ so one-iteration yields made an arriving request pay that many steps
1765
+ before it could even queue. Measured: 8.0 s to first token for a
1766
+ 37-token question whose engine-side wait was 0.0 s. A short real
1767
+ sleep every ~200 ms lets the loop drain all ready I/O in one go, at
1768
+ a cost bounded by the ratio (~1%) rather than per token."""
1769
+ now = time.perf_counter()
1770
+ if now - self._last_breath >= _BREATHE_EVERY_S:
1771
+ self._last_breath = now
1772
+ await asyncio.sleep(_BREATHE_FOR_S)
1773
+ else:
1774
+ await asyncio.sleep(0)
1775
+
1776
+ def submit(self, sub):
1777
+ if self._worker is None:
1778
+ self._worker = asyncio.ensure_future(self._run())
1779
+ if self._nested and wave_policy.is_background(sub):
1780
+ # An interactive wave is running over a parked one; background
1781
+ # arrivals wait until the parked wave has resumed.
1782
+ self._held_background.append(sub)
1783
+ else:
1784
+ self._pending.append(sub)
1785
+ self._announce_queued(sub)
1786
+ self._wake.set()
1787
+
1788
+ def is_idle(self):
1789
+ """True once no live or queued batch sub can still use scratch buffers."""
1790
+ return (
1791
+ not self._subs
1792
+ and not self._pending
1793
+ and not self._held_background
1794
+ and not self._parked
1795
+ and self._spec_sub is None
1796
+ )
1797
+
1798
+ def _running_subs(self):
1799
+ running = list(self._subs.values())
1800
+ if self._spec_sub is not None:
1801
+ running.append(self._spec_sub)
1802
+ return running
1803
+
1804
+ def _wave_running(self):
1805
+ return bool(self._subs) or self._spec_sub is not None or self._nested
1806
+
1807
+ def _announce_queued(self, sub):
1808
+ """One line naming a request that arrived mid-wave, plus its first
1809
+ waiting marker right away — the TS side flips the turn from
1810
+ "Processing prompt" to "Waiting for …" without waiting for the next
1811
+ worker step, which on a long prefill chunk can be most of a minute."""
1812
+ if not self._wave_running():
1813
+ return
1814
+ cid = getattr(sub.request, "cache_id", None)
1815
+ behind = wave_policy.cache_ids(self._running_subs())
1816
+ ahead = max(0, len(self._pending) + len(self._held_background) - 1)
1817
+ print(
1818
+ f"[batch] queued request={sub.request_id} cache={cid or '-'} "
1819
+ f"priority={sub.priority} ahead={ahead} behind={','.join(behind) or '-'}",
1820
+ flush=True,
1821
+ )
1822
+ if cid and self._liveness.due(cid):
1823
+ print(wave_policy.waiting_line(cid, ahead, behind), flush=True)
1824
+
1825
+ def _emit_waiting_liveness(self):
1826
+ """Liveness for requests the worker is NOT serving right now: queued
1827
+ ones (`waiting`) and a parked background wave (`paused`). Printed from
1828
+ the worker loop, between steps — so, like the prefill marker, it is
1829
+ proof the engine is alive and busy, and a wedged worker goes quiet.
1830
+
1831
+ Without it a queued request's only signal was silence: the TS
1832
+ pre-first-byte watchdog aborted a turn that was waiting its turn, and
1833
+ a parked wave's open stream would trip the streaming-idle bound."""
1834
+ running = self._running_subs()
1835
+ if not running:
1836
+ return
1837
+ behind = wave_policy.cache_ids(running)
1838
+ queued = [s for s in self._pending + self._held_background if wave_policy.is_live(s)]
1839
+ for idx, sub in enumerate(queued):
1840
+ cid = getattr(sub.request, "cache_id", None)
1841
+ if cid and self._liveness.due(cid):
1842
+ print(wave_policy.waiting_line(cid, idx, behind), flush=True)
1843
+ for sub in self._parked:
1844
+ cid = getattr(sub.request, "cache_id", None)
1845
+ if cid and wave_policy.is_live(sub) and self._liveness.due(cid):
1846
+ print(wave_policy.paused_line(cid, behind), flush=True)
1847
+
1848
+ def _announce_admitted(self, sub):
1849
+ cid = getattr(sub.request, "cache_id", None)
1850
+ if cid and self._liveness.seen(cid):
1851
+ print(wave_policy.admitted_line(cid, time.time() - sub.queued_at), flush=True)
1852
+ self._liveness.forget(cid)
1853
+
1854
+ def _drop_cancelled_pending(self):
1855
+ """A request whose client left while it was queued must not be
1856
+ admitted. Measured: a turn the watchdog had already abandoned was
1857
+ admitted two minutes later and prefilled for nobody."""
1858
+ live = []
1859
+ for sub in self._pending:
1860
+ if sub.cancelled:
1861
+ print(
1862
+ f"[batch] dropped request={sub.request_id} "
1863
+ f"cache={getattr(sub.request, 'cache_id', None) or '-'} "
1864
+ f"(client left while queued)",
1865
+ flush=True,
1866
+ )
1867
+ self._liveness.forget(getattr(sub.request, "cache_id", None))
1868
+ continue
1869
+ live.append(sub)
1870
+ self._pending = live
1871
+
1872
+ def _memory_allows_preemption(self):
1873
+ """A parked wave keeps its KV resident while the interactive wave
1874
+ allocates its own, so preempt only with headroom below the same
1875
+ threshold that collapses a wave to one member."""
1876
+ if _MEM_LIMIT_BYTES <= 0:
1877
+ return True
1878
+ active = _safe_active_memory()
1879
+ if active <= 0:
1880
+ return True
1881
+ ok = active < _MEM_LIMIT_BYTES * _MEM_ADMIT_FRACTION
1882
+ if not ok and not self._preempt_memory_noted:
1883
+ self._preempt_memory_noted = True
1884
+ print(
1885
+ f"[batch] preempt held: active={active // (1024 * 1024)}MB is at the "
1886
+ f"admission threshold of ceiling={_MEM_LIMIT_BYTES // (1024 * 1024)}MB; "
1887
+ f"the waiting request runs after this wave",
1888
+ flush=True,
1889
+ )
1890
+ return ok
1891
+
1892
+ def _preempt_due(self, running):
1893
+ if self._nested:
1894
+ return False
1895
+ if not wave_policy.preemptors(running, self._pending):
1896
+ return False
1897
+ return self._memory_allows_preemption()
1898
+
1899
+ async def _serve_preemptors(self, parked):
1900
+ """Run waiting interactive requests to completion while `parked` (a
1901
+ wave of background work) holds still, then hand the engine back.
1902
+
1903
+ The parked wave keeps everything it owns — KV, generator, logits
1904
+ processor state, its open SSE stream — and resumes at the step it
1905
+ stopped on. The interactive wave runs on a second BatchGenerator so
1906
+ the parked one's batch state is untouched; the shared per-wave fields
1907
+ (prefill liveness, the spec sub) are saved and restored around it.
1908
+ Nothing is cancelled, so nothing is recomputed.
1909
+
1910
+ Safe points only: callers park between engine steps (a batched
1911
+ wave's `next()`, a spec wave's prefill chunk or committed round), so
1912
+ no verification transaction or half-written cache is ever left open.
1913
+ """
1914
+ fg = [s for s in self._pending if wave_policy.is_live(s) and not wave_policy.is_background(s)]
1915
+ if not fg:
1916
+ return
1917
+ if self._spare_gen is None:
1918
+ try:
1919
+ self._spare_gen = self._make_generator("batch-preempt-init")
1920
+ except Exception as exc: # noqa: BLE001 — never kill the worker over a preemption
1921
+ print(f"[batch] preempt unavailable (spare generator failed): {exc}", flush=True)
1922
+ log_contained_exception("batch")
1923
+ return
1924
+ held = [s for s in self._pending if s not in fg]
1925
+ parked_ids = wave_policy.cache_ids(parked)
1926
+ print(
1927
+ f"[batch] preempt: parking {len(parked)} background request(s) "
1928
+ f"[{','.join(parked_ids) or '-'}] for {len(fg)} interactive "
1929
+ f"[{','.join(wave_policy.cache_ids(fg)) or '-'}]",
1930
+ flush=True,
1931
+ )
1932
+ started = time.perf_counter()
1933
+ saved = (
1934
+ self._gen,
1935
+ self._subs,
1936
+ self._prefill_total,
1937
+ self._prefill_done,
1938
+ self._prefill_meta,
1939
+ self._spec_sub,
1940
+ )
1941
+ self._gen = self._spare_gen
1942
+ self._subs = {}
1943
+ self._prefill_total = 0
1944
+ self._prefill_done = {}
1945
+ self._prefill_meta = {}
1946
+ self._last_prefill_emit = 0.0
1947
+ self._spec_sub = None
1948
+ self._pending = fg
1949
+ self._held_background = held
1950
+ self._parked = list(parked)
1951
+ self._nested = True
1952
+ try:
1953
+ while self._pending or self._subs:
1954
+ await self._step_wave()
1955
+ finally:
1956
+ self._spare_gen = self._gen
1957
+ (
1958
+ self._gen,
1959
+ self._subs,
1960
+ self._prefill_total,
1961
+ self._prefill_done,
1962
+ self._prefill_meta,
1963
+ self._spec_sub,
1964
+ ) = saved
1965
+ self._last_prefill_emit = 0.0
1966
+ self._pending = self._pending + self._held_background
1967
+ self._held_background = []
1968
+ self._parked = []
1969
+ self._nested = False
1970
+ self._preempt_memory_noted = False
1971
+ for sub in parked:
1972
+ self._liveness.forget(getattr(sub.request, "cache_id", None))
1973
+ print(
1974
+ f"[batch] resume: {len(parked)} background request(s) "
1975
+ f"[{','.join(parked_ids) or '-'}] after {time.perf_counter() - started:.1f}s",
1976
+ flush=True,
1977
+ )
1978
+
1979
+ def _save_cancelled_snapshot(self, sub, cache_layers=None, wave_tokens=None, done_tokens=None):
1980
+ """Keep what a cancelled request already paid for.
1981
+
1982
+ A cancel — the user's stop, a watchdog, or the service shutting down —
1983
+ used to drop the turn's prompt KV on the floor: the end-of-prompt
1984
+ snapshot is filed only by `_finish`, which a cancelled wave never
1985
+ reaches. The restart that motivated this lost a 98k-token prompt that
1986
+ had been fully prefilled two minutes earlier, and the resumed turn
1987
+ re-prefilled from a 14k-token prefix.
1988
+
1989
+ Saves the planted end-of-prompt snapshot when it was captured, or —
1990
+ for a spec wave stopped mid-prefill — the prefix it reached, proven
1991
+ against the prompt like any other snapshot. Only ever an improvement:
1992
+ it must extend past what this turn was seeded with, so a retry of the
1993
+ same turn resumes from here instead of from the seed."""
1994
+ cid = getattr(sub.request, "cache_id", None)
1995
+ if not cid or not self._needs_snapshot:
1996
+ return
1997
+ try:
1998
+ snapshot = sub.prompt_snapshot
1999
+ if snapshot is None and cache_layers is not None and wave_tokens and done_tokens:
2000
+ boundary = int(sub.reused_tokens) + int(done_tokens)
2001
+ layers = spec_decode.clone_layers(cache_layers)
2002
+ if cache_seed.snapshot_matches_prompt(
2003
+ layers, list(wave_tokens)[:boundary], sub.prompt_tokens, boundary
2004
+ ):
2005
+ snapshot = (layers, boundary)
2006
+ if snapshot is None:
2007
+ return
2008
+ layers, count = snapshot
2009
+ if int(count) <= int(sub.reused_tokens or 0):
2010
+ return
2011
+ state = PromptCacheState()
2012
+ state.cache = layers
2013
+ state.token_ids = list(sub.prompt_tokens[: int(count)])
2014
+ _save_cache(cid, state)
2015
+ _LAST_SAVE_ROSTER[cid] = sub.roster_fp
2016
+ print(
2017
+ f"[batch] saved cancelled cache_id={cid} tokens={int(count)} "
2018
+ f"(was seeded with {int(sub.reused_tokens or 0)})",
2019
+ flush=True,
2020
+ )
2021
+ except Exception as exc: # noqa: BLE001 — best-effort; never break the worker
2022
+ print(f"[batch] cancelled-turn save failed cache_id={cid}: {exc}", flush=True)
2023
+
2024
+ def _sampler_for(self, request):
2025
+ from mlx_lm.sample_utils import make_sampler
2026
+
2027
+ return make_sampler(
2028
+ temp=float(request.temperature) if request.temperature is not None else 0.0,
2029
+ top_p=float(request.top_p) if request.top_p is not None else 0.0,
2030
+ top_k=int(request.top_k) if request.top_k is not None else 0,
2031
+ )
2032
+
2033
+ def _processors_for(self, request, grammar):
2034
+ from mlx_lm.sample_utils import make_logits_processors
2035
+
2036
+ procs = []
2037
+ if request.repetition_penalty is not None:
2038
+ procs = make_logits_processors(
2039
+ repetition_penalty=float(request.repetition_penalty),
2040
+ repetition_context_size=int(request.repetition_context_size or 20),
2041
+ )
2042
+ if grammar is not None:
2043
+ # Accepts a single processor or a list (grammar + think-budget).
2044
+ procs = list(procs) + (grammar if isinstance(grammar, list) else [grammar])
2045
+ # BatchGenerator receives one processor list per sequence. A mixed
2046
+ # wave (for example cache warm + pi tool call) must use [] for the
2047
+ # unconstrained sequence: mlx_lm iterates every per-sequence entry and
2048
+ # crashes on None with "'NoneType' object is not iterable".
2049
+ return list(procs or [])
2050
+
2051
+ def _seed_args(self, sub):
2052
+ """Plan cache reuse for this sub: longest-common-prefix with
2053
+ trim-to-prefix, mirroring the serial path's semantics (the old
2054
+ exact-prefix-only check re-prefilled entire 20K-token sessions
2055
+ whenever one stripped reasoning character diverged the history
2056
+ from the raw tokens). Delegates to cache_seed.seed_from_state;
2057
+ stamps the authoritative prefill/reuse counts on the sub so the
2058
+ liveness marker and usage frames report engine reality, and logs
2059
+ the decision — the line to grep when turns re-prefill."""
2060
+ # Late prefix re-check. The lookup in the request handler runs when the
2061
+ # HTTP request ARRIVES; a band is published when its pioneer's turn
2062
+ # ENDS. Static-wave admission puts minutes between the two, so every
2063
+ # sibling dispatched alongside the pioneer — which is the whole fanout
2064
+ # case this exists for — resolved `fresh` before the entry existed.
2065
+ # Measured: two real sibling sessions agreed on
2066
+ # `prefix-band-8977c4fc5f74e5d9`, the pioneer published 16,576 tokens
2067
+ # to it, and the sibling still cold-prefilled 71,346 because its
2068
+ # lookup had run 55ms too early. Re-checking here, at the moment the
2069
+ # wave actually starts, is the only point that sees the publish.
2070
+ if sub.request.cache_id and not getattr(sub.seed_state, "token_ids", None):
2071
+ for pid in _request_prefix_ids(sub.request):
2072
+ seeded = _try_seed_from_prefix(sub.request.cache_id, pid)
2073
+ if seeded is not None:
2074
+ sub.seed_state = seeded
2075
+ break
2076
+ plan = cache_seed.seed_from_state(sub.seed_state, sub.prompt_tokens)
2077
+ # A session entry that diverged and cannot be trimmed is worth
2078
+ # nothing to this prompt, but the shared prefix it was born from
2079
+ # usually still is. Without this a stale session entry was strictly
2080
+ # WORSE than having none: a miss falls back to the prefix band
2081
+ # (~14-17k tokens reused on qwen3.8-27b), a divergent hit re-prefilled
2082
+ # everything. It matters most right after a restart, when the session
2083
+ # entry the engine persisted no longer matches a rebuilt transcript.
2084
+ if sub.request.cache_id and str(plan.mode).startswith("fresh") and plan.cached_len:
2085
+ for pid in _request_prefix_ids(sub.request):
2086
+ seeded = _try_seed_from_prefix(sub.request.cache_id, pid)
2087
+ if seeded is None:
2088
+ continue
2089
+ fallback = cache_seed.seed_from_state(seeded, sub.prompt_tokens)
2090
+ if fallback.reused > plan.reused:
2091
+ print(
2092
+ f"[batch] seed fallback cache_id={sub.request.cache_id} "
2093
+ f"session-entry mode={plan.mode} lcp={plan.lcp} "
2094
+ f"cached={plan.cached_len} -> prefix={pid} reused={fallback.reused}",
2095
+ flush=True,
2096
+ )
2097
+ # Same divergence preview the full-prefill path prints:
2098
+ # what the session entry and this prompt disagree on is
2099
+ # the only thing that says which side to fix.
2100
+ try:
2101
+ _tk = getattr(PROCESSOR, "tokenizer", None) or PROCESSOR
2102
+ _cached = list(getattr(sub.seed_state, "token_ids", None) or [])
2103
+ _lo = max(0, plan.lcp - 6)
2104
+ print(
2105
+ f"[batch] churn cache_id={sub.request.cache_id} at={plan.lcp} "
2106
+ f"cached={_tk.decode(_cached[_lo:plan.lcp + 24])!r} "
2107
+ f"prompt={_tk.decode(list(sub.prompt_tokens)[_lo:plan.lcp + 24])!r}",
2108
+ flush=True,
2109
+ )
2110
+ except Exception as exc: # noqa: BLE001 — diagnostics never break a turn
2111
+ print(f"[batch] churn preview unavailable: {exc}", flush=True)
2112
+ sub.seed_state = seeded
2113
+ plan = fallback
2114
+ break
2115
+ sub.prefill_total = len(plan.segment)
2116
+ sub.reused_tokens = plan.reused
2117
+ sub.seed_mode = plan.mode
2118
+ # Publish the shared band ONLY from a session that reused nothing —
2119
+ # the "pioneer" of this prefix; every sibling after it inherits the
2120
+ # band and skips it. The band is a SECOND cut beside the session's
2121
+ # own end-minus-margin boundary, not a replacement for it: saving the
2122
+ # pioneer at the band made its next turn re-prefill everything after
2123
+ # it, which on Gemma (tool block rendered after the system text) was
2124
+ # 7,313 of 8,604 tokens on the second message of every chat.
2125
+ if sub.snapshot_target is None and str(plan.mode).startswith("fresh"):
2126
+ sub.band_target = _band_snapshot_target(sub)
2127
+ if sub.request.cache_id:
2128
+ print(
2129
+ f"[batch] seed cache_id={sub.request.cache_id} mode={plan.mode} "
2130
+ f"reused={plan.reused} prefill={len(plan.segment)}",
2131
+ flush=True,
2132
+ )
2133
+ if plan.mode == "fresh-untrimmable" and plan.cached_len:
2134
+ # Name the churn instead of leaving it to inference. Two
2135
+ # rounds of this investigation guessed the cause from cache
2136
+ # SIZES (which say what was available, not what matched) and
2137
+ # both guesses were wrong; decoding the two sides at the
2138
+ # divergence point is the only thing that identifies it.
2139
+ try:
2140
+ _tk = getattr(PROCESSOR, "tokenizer", None) or PROCESSOR
2141
+ _cached = list(getattr(sub.seed_state, "token_ids", None) or [])
2142
+ _lo = max(0, plan.lcp - 6)
2143
+ print(
2144
+ f"[batch] churn cache_id={sub.request.cache_id} "
2145
+ f"at={plan.lcp} "
2146
+ f"cached={_tk.decode(_cached[_lo:plan.lcp + 24])!r} "
2147
+ f"prompt={_tk.decode(list(sub.prompt_tokens)[_lo:plan.lcp + 24])!r}",
2148
+ flush=True,
2149
+ )
2150
+ except Exception as exc: # noqa: BLE001 — diagnostics never break a turn
2151
+ print(f"[batch] churn preview unavailable: {exc}", flush=True)
2152
+ if plan.mode == "fresh-untrimmable":
2153
+ # lcp/cached is the diagnostic pair: lcp near 0 means the
2154
+ # prompt churned early (fix the prompt), lcp near cached
2155
+ # means the snapshot boundary was cut past the stable
2156
+ # region (fix the boundary). Same symptom, opposite fixes.
2157
+ print(
2158
+ f"[batch] WARNING full prefill: cache_id={sub.request.cache_id} "
2159
+ f"had divergent untrimmable state without a reusable prompt "
2160
+ f"snapshot (tokens={len(plan.segment)} lcp={plan.lcp} "
2161
+ f"cached={plan.cached_len} "
2162
+ f"lcp_frac={plan.lcp / plan.cached_len:.2f})"
2163
+ if plan.cached_len
2164
+ else f"[batch] WARNING full prefill: cache_id={sub.request.cache_id} "
2165
+ f"had divergent untrimmable state (tokens={len(plan.segment)})",
2166
+ flush=True,
2167
+ )
2168
+ print(
2169
+ f"[batch] churn-shape cache_id={sub.request.cache_id} "
2170
+ f"probe_roster={sub.roster_fp} "
2171
+ f"saved_roster={_LAST_SAVE_ROSTER.get(sub.request.cache_id, 'unknown')} "
2172
+ f"path={'warm' if str(sub.request_id).startswith('cachewarm-') else 'chat'} "
2173
+ f"prompt_len={len(sub.prompt_tokens)}",
2174
+ flush=True,
2175
+ )
2176
+ return plan.segment, plan.caches, plan.all_tokens
2177
+
2178
+ def _snapshot_segments(self, sub, segment):
2179
+ """Plant the per-sequence segment edges for KV extraction: the
2180
+ session's own boundary, plus the shared band's on a pioneer turn.
2181
+
2182
+ BatchGenerator reports every prefill chunk as a prompt response. The
2183
+ old handler copied the growing cache at *every* chunk and kept only the
2184
+ last copy. Long prompts therefore paid dozens of large, discarded KV
2185
+ copies. Dedicated targets make capture one-shot per boundary and also
2186
+ work for unequal-length multi-sequence waves: mlx_lm finalizes each
2187
+ chunk's right-padding before returning, then extract_cache(uid)
2188
+ removes that sequence's left-padding.
2189
+ """
2190
+ if not self._needs_snapshot or not sub.request.cache_id:
2191
+ sub.band_target = None
2192
+ return [segment]
2193
+ plan = cache_seed.plan_snapshot_segments(
2194
+ segment,
2195
+ sub.reused_tokens,
2196
+ len(sub.prompt_tokens),
2197
+ sub.snapshot_target,
2198
+ _SNAPSHOT_BOUNDARY_MARGIN,
2199
+ sub.band_target,
2200
+ )
2201
+ sub.snapshot_target = plan.target
2202
+ sub.band_target = plan.band
2203
+ return plan.segments
2204
+
2205
+ def _emit_prefill_liveness(self) -> None:
2206
+ """While the active wave is still prefilling (no token emitted yet),
2207
+ print a `Prefill:` line the supervisor's stdout-parser turns into a
2208
+ `prefill` phase event — keeping the TS pre-first-byte watchdog anchored
2209
+ on real engine liveness instead of tripping its silent-cold-start
2210
+ bound. Honest by construction: it only fires while `_prefill_total > 0`
2211
+ (cleared the instant the first token lands), and only from the live
2212
+ worker loop — a genuinely wedged engine emits nothing and still trips.
2213
+
2214
+ Progress is real: the worker consumes `BatchGenerator.next()` per
2215
+ internal step, so each prompt chunk reports its (processed, total)
2216
+ and the marker carries an actual percentage instead of a constant
2217
+ 0%.
2218
+
2219
+ Each marker names the sub it belongs to (`cache=<cache_id>`), and
2220
+ carries THAT sub's own done/total rather than the wave aggregate —
2221
+ see `_prefill_meta`. An untagged marker still parses, and the
2222
+ supervisor still broadcasts it, so an engine without the tag degrades
2223
+ to the old behaviour rather than going dark."""
2224
+ total = self._prefill_total
2225
+ if total < _PREFILL_LIVENESS_MIN_TOKENS:
2226
+ return
2227
+ now = time.time()
2228
+ if now - self._last_prefill_emit < _PREFILL_LIVENESS_INTERVAL_S:
2229
+ return
2230
+ self._last_prefill_emit = now
2231
+ # Format MUST match stdout-parser's Prefill regex:
2232
+ # Prefill: 42%| | 8192/<total> [batched] cache=<id>
2233
+ if not self._prefill_meta:
2234
+ done = min(total, sum(self._prefill_done.values()))
2235
+ pct = min(99, (100 * done) // total) if total > 0 else 0
2236
+ print(f"Prefill: {pct:3d}%| | {done}/{total} [batched]", flush=True)
2237
+ return
2238
+ for key, (cache_id, sub_total) in self._prefill_meta.items():
2239
+ sub_total = max(0, int(sub_total or 0))
2240
+ if sub_total <= 0:
2241
+ continue
2242
+ done = min(sub_total, int(self._prefill_done.get(key, 0)))
2243
+ pct = min(99, (100 * done) // sub_total)
2244
+ tag = f" cache={cache_id}" if cache_id else ""
2245
+ print(
2246
+ f"Prefill: {pct:3d}%| | {done}/{sub_total} [batched]{tag}",
2247
+ flush=True,
2248
+ )
2249
+
2250
+ def _admit_count(self, pending: int) -> int:
2251
+ """How many of the `pending` co-arrived subs to start THIS wave.
2252
+
2253
+ `self._max` is the hard concurrency contract advertised by the engine;
2254
+ pending work beyond it must remain queued for the next static wave.
2255
+ Within that width, live allocation is a second backstop to the TS-side
2256
+ slot estimate: when the device is already tight, collapse the wave to
2257
+ one request rather than risking a Metal abort."""
2258
+ capped = min(pending, self._max)
2259
+ if capped <= 1 or _MEM_LIMIT_BYTES <= 0:
2260
+ return capped
2261
+ threshold = _MEM_LIMIT_BYTES * _MEM_ADMIT_FRACTION
2262
+ active = _safe_active_memory()
2263
+ if active <= 0:
2264
+ return capped # can't read usage → keep the configured hard cap
2265
+ if active >= threshold:
2266
+ # Already tight before this wave even starts — reclaim any freed
2267
+ # buffer cache and re-check once before forcing serial.
2268
+ _reclaim_mlx_buffer_cache("batch-admission", force=True)
2269
+ active = _safe_active_memory()
2270
+ return 1 if active >= threshold else capped
2271
+
2272
+ async def _run(self):
2273
+ while True:
2274
+ if not self._pending and not self._subs:
2275
+ self._wake.clear()
2276
+ await self._wake.wait()
2277
+ async with _get_generation_lock():
2278
+ await self._step_wave()
2279
+ while self._subs:
2280
+ await self._step_wave()
2281
+
2282
+ async def _step_wave(self):
2283
+ # Evict client-disconnected sequences before stepping.
2284
+ gone = [uid for uid, s in self._subs.items() if s.cancelled]
2285
+ if gone:
2286
+ for uid in gone:
2287
+ self._save_cancelled_snapshot(self._subs[uid])
2288
+ try:
2289
+ self._gen.remove(gone)
2290
+ except Exception as exc: # noqa: BLE001
2291
+ print(f"[batch] remove failed: {exc}", flush=True)
2292
+ for uid in gone:
2293
+ self._subs.pop(uid, None)
2294
+
2295
+ # Static-wave admission: only start a new batch when the current
2296
+ # one has fully drained. mlx_lm's continuous merge of a mid-flight
2297
+ # prompt batch into the running generation batch isn't supported by
2298
+ # the mlx_vlm model caches here (raises broadcast (N),(M)), so we
2299
+ # batch co-arriving requests into a wave and run it to completion
2300
+ # before admitting the next. Co-arriving turns (multi-pane,
2301
+ # ask-chains) still share batched decode; a turn that arrives
2302
+ # mid-wave waits for the wave to drain — no worse than the old
2303
+ # serial path, and strictly better when turns arrive together.
2304
+ if self._pending and not self._subs:
2305
+ self._drop_cancelled_pending()
2306
+ if not self._pending:
2307
+ return
2308
+ # Interactive requests first, and never in a wave with background
2309
+ # work — see wave_policy.select_wave.
2310
+ batch, self._pending = wave_policy.select_wave(self._pending, self._admit_count)
2311
+ if self._pending:
2312
+ # Deferred by the configured concurrency ceiling, live
2313
+ # memory pressure, or because the wave is interactive-only.
2314
+ # The remaining subs run in the next wave, after this one
2315
+ # drains and frees its KV. The outer loop re-enters admission
2316
+ # once _subs empties (it never blocks on _wake while _pending
2317
+ # is non-empty).
2318
+ held_bg = sum(1 for s in self._pending if wave_policy.is_background(s))
2319
+ print(
2320
+ f"[batch] admission-throttled: admitting {len(batch)}, "
2321
+ f"deferring {len(self._pending)} to next wave "
2322
+ f"({held_bg} of them background, "
2323
+ f"max_concurrency={self._max}, "
2324
+ f"active={_safe_active_memory() // (1024 * 1024)}MB, "
2325
+ f"ceiling={_MEM_LIMIT_BYTES // (1024 * 1024)}MB)",
2326
+ flush=True,
2327
+ )
2328
+ for sub in batch:
2329
+ self._announce_admitted(sub)
2330
+ if len(batch) == 1 and _SPEC is not None and self._drafter_busy:
2331
+ print(
2332
+ f"[spec] off request={batch[0].request_id} reason=the drafter "
2333
+ f"belongs to a paused wave",
2334
+ flush=True,
2335
+ )
2336
+ elif len(batch) == 1 and _SPEC is not None:
2337
+ spec_sub = batch[0]
2338
+ mode, why = spec_decode.spec_mode(
2339
+ spec_sub.request, spec_sub.grammar
2340
+ )
2341
+ if mode is not None:
2342
+ admit, gate_why = spec_decode.draft_prefill_gate(
2343
+ _SPEC, len(spec_sub.prompt_tokens)
2344
+ )
2345
+ if not admit:
2346
+ mode, why = None, gate_why
2347
+ if mode is not None:
2348
+ await self._run_spec_wave(spec_sub, mode)
2349
+ return
2350
+ print(
2351
+ f"[spec] off request={spec_sub.request_id} reason={why}",
2352
+ flush=True,
2353
+ )
2354
+ segments, caches, alltoks, samplers, lps, maxtoks = [], [], [], [], [], []
2355
+ for sub in batch:
2356
+ seg, c, at = self._seed_args(sub)
2357
+ segments.append(self._snapshot_segments(sub, seg))
2358
+ caches.append(c)
2359
+ alltoks.append(at)
2360
+ samplers.append(self._sampler_for(sub.request))
2361
+ lps.append(self._processors_for(sub.request, sub.grammar))
2362
+ maxtoks.append(int(sub.request.max_tokens) if sub.request.max_tokens else 2048)
2363
+ try:
2364
+ uids = self._gen.insert_segments(
2365
+ segments=segments,
2366
+ max_tokens=maxtoks,
2367
+ caches=caches,
2368
+ all_tokens=alltoks,
2369
+ samplers=samplers,
2370
+ logits_processors=lps,
2371
+ )
2372
+ except Exception as exc: # noqa: BLE001
2373
+ print(f"[batch] insert failed: {exc}", flush=True)
2374
+ log_contained_exception("batch")
2375
+ for sub in batch:
2376
+ sub.queue.put_nowait(("err", str(exc)))
2377
+ return
2378
+ for uid, sub in zip(uids, batch):
2379
+ sub.uid = uid
2380
+ sub.prefill_started_at = time.perf_counter()
2381
+ self._subs[uid] = sub
2382
+ # Arm prefill-liveness for the wave: total new tokens to
2383
+ # prefill across its subs (authoritative — from the seed
2384
+ # plans above). Reset the throttle so the first marker
2385
+ # emits immediately (before the blocking decode below).
2386
+ self._prefill_total = sum(
2387
+ max(0, int(getattr(s, "prefill_total", 0))) for s in batch
2388
+ )
2389
+ self._prefill_done = {}
2390
+ self._prefill_meta = {
2391
+ uid: (
2392
+ getattr(s.request, "cache_id", None),
2393
+ max(0, int(getattr(s, "prefill_total", 0))),
2394
+ )
2395
+ for uid, s in self._subs.items()
2396
+ }
2397
+ self._last_prefill_emit = 0.0
2398
+
2399
+ # One engine step across all active sequences. `next()` (not
2400
+ # `next_generated()`) so control returns after EVERY internal
2401
+ # step — prompt chunks included. That is what makes the
2402
+ # end-of-prompt snapshot capturable at its exact boundary
2403
+ # (see _capture_prompt_snapshot), gives the liveness marker
2404
+ # real per-chunk progress instead of a constant 0%, and lets
2405
+ # the event loop breathe between prefill chunks.
2406
+ # A wave of background work steps aside, between steps, for a person
2407
+ # who has been waiting since it started. It resumes from this exact
2408
+ # step once they are served.
2409
+ if self._subs and self._preempt_due(self._subs.values()):
2410
+ await self._serve_preemptors(list(self._subs.values()))
2411
+ if self._subs:
2412
+ # Emit a prefill marker before the (potentially long, event-
2413
+ # loop-blocking) step so the watchdog sees liveness.
2414
+ # No-op once the first token lands (_prefill_total → 0).
2415
+ self._emit_prefill_liveness()
2416
+ self._emit_waiting_liveness()
2417
+ try:
2418
+ prompt_responses, responses = self._gen.next()
2419
+ step_ended_at = time.perf_counter()
2420
+ except Exception as exc: # noqa: BLE001
2421
+ print(f"[batch] step failed: {exc}", flush=True)
2422
+ log_contained_exception("batch")
2423
+ for sub in list(self._subs.values()):
2424
+ sub.queue.put_nowait(("err", str(exc)))
2425
+ self._subs.clear()
2426
+ return
2427
+ # set_cache_limit bounds future reuse, while this periodic
2428
+ # pressure check also catches temporary overshoot during a
2429
+ # long prefill/decode. It replaces reliance on mlx_lm's much
2430
+ # coarser 512-generation-step cache clear.
2431
+ self._memory_check_steps += 1
2432
+ if self._memory_check_steps >= 32:
2433
+ self._memory_check_steps = 0
2434
+ _reclaim_mlx_buffer_cache("batch-pressure")
2435
+ for pr in prompt_responses:
2436
+ progress = getattr(pr, "progress", None)
2437
+ has_progress = isinstance(progress, tuple) and len(progress) == 2
2438
+ if has_progress:
2439
+ self._prefill_done[pr.uid] = int(progress[0])
2440
+ psub = self._subs.get(pr.uid)
2441
+ if psub is None or not self._needs_snapshot or not psub.request.cache_id:
2442
+ continue
2443
+ if getattr(pr, "end_of_prompt", False):
2444
+ continue
2445
+ if has_progress and getattr(pr, "end_of_segment", False):
2446
+ # Capture only the deliberately planted edge. Chunk
2447
+ # edges are observable too, but copying at each one is
2448
+ # expensive and unnecessary. extract_cache(uid) is
2449
+ # padding-aware after PromptProcessingBatch.prompt()
2450
+ # finalizes this step, so every sub in a multi-request
2451
+ # wave can safely take its own snapshot.
2452
+ boundary = psub.reused_tokens + int(progress[0])
2453
+ if boundary in (psub.snapshot_target, psub.band_target):
2454
+ self._capture_prompt_snapshot(psub, boundary)
2455
+ if responses:
2456
+ # A token came back → prefill for this wave is done; stop
2457
+ # the liveness markers (the SSE first byte now drives the
2458
+ # watchdog's tighter streaming-idle bound).
2459
+ self._prefill_total = 0
2460
+ self._prefill_done = {}
2461
+ self._prefill_meta = {}
2462
+ for r in responses:
2463
+ sub = self._subs.get(r.uid)
2464
+ if sub is None:
2465
+ continue
2466
+ if sub.generation_started_at is None:
2467
+ # The step that yields the first token is the closest
2468
+ # boundary BatchGenerator exposes between prefill and
2469
+ # decode. Count it with prefill/TTFT, then measure decode
2470
+ # from the inter-token intervals that follow; otherwise
2471
+ # a short prompt can produce a near-zero denominator and
2472
+ # a wildly inflated first-token rate.
2473
+ sub.generation_started_at = step_ended_at
2474
+ prefill_started_at = sub.prefill_started_at or step_ended_at
2475
+ prefill_seconds = max(
2476
+ step_ended_at - prefill_started_at,
2477
+ 1e-9,
2478
+ )
2479
+ if sub.prefill_total > 0:
2480
+ sub.prompt_tps = sub.prefill_total / prefill_seconds
2481
+ sub.first_token_seen = True
2482
+ sub.token_ids.append(int(r.token))
2483
+ generated_intervals = len(sub.token_ids) - 1
2484
+ generation_seconds = step_ended_at - sub.generation_started_at
2485
+ if generated_intervals > 0 and generation_seconds > 0:
2486
+ sub.generation_tps = generated_intervals / generation_seconds
2487
+ self._emit_delta(sub)
2488
+ if r.finish_reason is not None:
2489
+ self._finish(sub, r)
2490
+ self._subs.pop(r.uid, None)
2491
+ # Yield so SSE generators flush + new requests get admitted.
2492
+ await self._breathe()
2493
+
2494
+
2495
+ async def _run_spec_wave(self, sub, mode="greedy"):
2496
+ """Serve one eligible sub via MTP speculation (sidecar-owned wave).
2497
+
2498
+ Replaces BatchGenerator for this wave only: seed plan + direct
2499
+ chunked prefill (with the planted boundary snapshot) + upstream's
2500
+ serial _mtp_rounds. Emit, liveness, cancellation, stats, and the
2501
+ finish-time save reuse the engine's existing machinery, so a spec
2502
+ turn is indistinguishable downstream from a batch turn. Measured
2503
+ basis + eligibility rationale live in spec_decode.py."""
2504
+ sub.prefill_started_at = time.perf_counter()
2505
+ self._spec_sub = sub
2506
+ try:
2507
+ # SeedPlan's third field is the tokens ALREADY IN the seeded
2508
+ # cache (the reused prefix — empty on fresh), mirroring what
2509
+ # insert_segments' all_tokens means to mlx_lm.
2510
+ seg, seeded, prefix_toks = self._seed_args(sub)
2511
+ tower = getattr(MODEL, "language_model", MODEL)
2512
+ cache_layers = seeded
2513
+ if cache_layers is not None and not spec_decode.cache_matches_tower(
2514
+ cache_layers, tower
2515
+ ):
2516
+ print(
2517
+ f"[spec] seed-guard: cached layer classes mismatch the "
2518
+ f"active tower for cache_id={sub.request.cache_id}; "
2519
+ f"treating as fresh",
2520
+ flush=True,
2521
+ )
2522
+ cache_layers = None
2523
+ seg = list(sub.prompt_tokens)
2524
+ prefix_toks = []
2525
+ sub.reused_tokens = 0
2526
+ sub.prefill_total = len(seg)
2527
+ sub.seed_mode = "fresh"
2528
+ if cache_layers is None:
2529
+ cache_layers = tower.make_cache()
2530
+ full_prompt = list(prefix_toks or []) + list(seg)
2531
+ max_toks = (
2532
+ int(sub.request.max_tokens) if sub.request.max_tokens else 2048
2533
+ )
2534
+
2535
+ cuts: List[int] = []
2536
+ if self._needs_snapshot and sub.request.cache_id:
2537
+ plan = cache_seed.plan_snapshot_segments(
2538
+ seg,
2539
+ sub.reused_tokens,
2540
+ len(sub.prompt_tokens),
2541
+ sub.snapshot_target,
2542
+ _SNAPSHOT_BOUNDARY_MARGIN,
2543
+ sub.band_target,
2544
+ )
2545
+ sub.snapshot_target = plan.target
2546
+ sub.band_target = plan.band
2547
+ edge = 0
2548
+ for part in plan.segments[:-1]:
2549
+ edge += len(part)
2550
+ cuts.append(edge)
2551
+ else:
2552
+ sub.band_target = None
2553
+
2554
+ self._prefill_total = max(0, int(sub.prefill_total or 0))
2555
+ self._prefill_done = {}
2556
+ self._prefill_meta = {
2557
+ 0: (
2558
+ getattr(sub.request, "cache_id", None),
2559
+ max(0, int(sub.prefill_total or 0)),
2560
+ )
2561
+ }
2562
+ self._last_prefill_emit = 0.0
2563
+
2564
+ final_out = None
2565
+ for kind, done, _total, maybe_out in spec_decode.chunked_prefill_steps(
2566
+ tower, cache_layers, list(seg), ARGS.prefill_step_size, cuts=cuts
2567
+ ):
2568
+ if sub.cancelled:
2569
+ print(
2570
+ f"[spec] wave cancelled during prefill "
2571
+ f"request={sub.request_id}",
2572
+ flush=True,
2573
+ )
2574
+ self._save_cancelled_snapshot(sub, cache_layers, full_prompt, done)
2575
+ return
2576
+ self._prefill_done[0] = int(done)
2577
+ if kind == "cut":
2578
+ self._capture_spec_snapshot(
2579
+ sub, cache_layers, full_prompt, int(sub.reused_tokens) + int(done)
2580
+ )
2581
+ elif kind == "final":
2582
+ final_out = maybe_out
2583
+ # Chunk edges are quiescent: the cache holds exactly the
2584
+ # tokens prefilled so far, and the drafter is not in use yet.
2585
+ if self._spec_preempt_due(sub):
2586
+ await self._serve_preemptors([sub])
2587
+ self._emit_prefill_liveness()
2588
+ self._emit_waiting_liveness()
2589
+ await self._breathe()
2590
+
2591
+ self._prefill_total = 0
2592
+ self._prefill_done = {}
2593
+ self._prefill_meta = {}
2594
+
2595
+ # Assisted mode: positioned target sampling (temperature/top_p/
2596
+ # top_k, per-(seed,position) keys) and/or the processed walk.
2597
+ # The processor list is built by the SAME _processors_for the
2598
+ # normal path uses, so constrained decoding is identical either
2599
+ # route. Seed is logged for replay.
2600
+ sampler = None
2601
+ processors = None
2602
+ history = None
2603
+ wave_seed = None
2604
+ if mode == "assisted":
2605
+ processors = self._processors_for(sub.request, sub.grammar) or None
2606
+ t = sub.request.temperature
2607
+ if t not in (None, 0, 0.0):
2608
+ wave_seed = int.from_bytes(os.urandom(6), "big")
2609
+ sampler = spec_decode.PositionedSampler(
2610
+ temperature=float(t),
2611
+ top_p=float(sub.request.top_p)
2612
+ if sub.request.top_p is not None
2613
+ else 0.0,
2614
+ top_k=int(sub.request.top_k)
2615
+ if sub.request.top_k is not None
2616
+ else 0,
2617
+ seed=wave_seed,
2618
+ )
2619
+ history = list(full_prompt)
2620
+
2621
+ def emit(tok):
2622
+ now = time.perf_counter()
2623
+ if sub.generation_started_at is None:
2624
+ # Mirror the batch worker: the step yielding the first
2625
+ # token counts with prefill/TTFT; decode is measured
2626
+ # from inter-token intervals after it.
2627
+ sub.generation_started_at = now
2628
+ prefill_seconds = max(
2629
+ now - (sub.prefill_started_at or now), 1e-9
2630
+ )
2631
+ if sub.prefill_total > 0:
2632
+ sub.prompt_tps = sub.prefill_total / prefill_seconds
2633
+ sub.first_token_seen = True
2634
+ sub.token_ids.append(int(tok))
2635
+ intervals = len(sub.token_ids) - 1
2636
+ gen_seconds = now - sub.generation_started_at
2637
+ if intervals > 0 and gen_seconds > 0:
2638
+ sub.generation_tps = intervals / gen_seconds
2639
+ self._emit_delta(sub)
2640
+
2641
+ finish_reason = None
2642
+ if mode == "assisted":
2643
+ first = spec_decode.processed_first_token(
2644
+ final_out, sampler, processors, history
2645
+ )
2646
+ history.append(first)
2647
+ else:
2648
+ first = spec_decode.first_token_from(final_out)
2649
+ emit(first)
2650
+ if first in self._stop_ids:
2651
+ finish_reason = "stop"
2652
+ elif len(sub.token_ids) >= max_toks:
2653
+ finish_reason = "length"
2654
+
2655
+ if finish_reason is None:
2656
+ # From here the drafter's per-wave state (its cache, shared
2657
+ # kv, position) belongs to this wave until it ends.
2658
+ self._drafter_busy = True
2659
+ if mode == "assisted":
2660
+ rounds = spec_decode.assisted_rounds(
2661
+ MODEL,
2662
+ _SPEC,
2663
+ cache_layers,
2664
+ list(seg),
2665
+ final_out,
2666
+ first,
2667
+ max_toks,
2668
+ sampler=sampler,
2669
+ processors=processors,
2670
+ history=history,
2671
+ )
2672
+ else:
2673
+ rounds = spec_decode.make_rounds(
2674
+ MODEL,
2675
+ _SPEC,
2676
+ cache_layers,
2677
+ list(seg),
2678
+ final_out,
2679
+ first,
2680
+ max_toks,
2681
+ )
2682
+ mem_steps = 0
2683
+ for tok, _lp in rounds:
2684
+ if sub.cancelled:
2685
+ print(
2686
+ f"[spec] wave cancelled request={sub.request_id}",
2687
+ flush=True,
2688
+ )
2689
+ self._save_cancelled_snapshot(sub)
2690
+ return
2691
+ emit(tok)
2692
+ if tok in self._stop_ids:
2693
+ finish_reason = "stop"
2694
+ break
2695
+ if len(sub.token_ids) >= max_toks:
2696
+ finish_reason = "length"
2697
+ break
2698
+ mem_steps += 1
2699
+ if mem_steps >= 32:
2700
+ mem_steps = 0
2701
+ _reclaim_mlx_buffer_cache("spec-pressure")
2702
+ # Each yield follows a committed verify round (upstream
2703
+ # commits before emitting), so the target cache is
2704
+ # consistent here. A wave run inside the pause cannot
2705
+ # draft: _drafter_busy keeps it on the BatchGenerator.
2706
+ if self._spec_preempt_due(sub):
2707
+ await self._serve_preemptors([sub])
2708
+ self._emit_waiting_liveness()
2709
+ await self._breathe()
2710
+ if finish_reason is None:
2711
+ finish_reason = (
2712
+ "length" if len(sub.token_ids) >= max_toks else "stop"
2713
+ )
2714
+
2715
+ line = spec_decode.stats_line(_SPEC)
2716
+ if line:
2717
+ print(
2718
+ f"[spec] stats request={sub.request_id} mode={mode} "
2719
+ + (f"seed={wave_seed} " if wave_seed is not None else "")
2720
+ + f"ctx={len(sub.prompt_tokens)} gen={len(sub.token_ids)} "
2721
+ f":: {line}",
2722
+ flush=True,
2723
+ )
2724
+ # Rounds can exit mid-round (stop token, budget) with verified-
2725
+ # but-rejected tokens still in the KV. Hybrid stacks are immune —
2726
+ # their save always uses the boundary snapshot — but a trimmable
2727
+ # stack must be trimmed back to exactly the tokens we report, or
2728
+ # the saved entry's offsets outrun its token list.
2729
+ try:
2730
+ reported = len(sub.prompt_tokens) + len(sub.token_ids)
2731
+ offsets = [
2732
+ int(getattr(c, "offset", 0) or 0) for c in cache_layers
2733
+ ]
2734
+ excess = max(offsets) - reported if offsets else 0
2735
+ if excess > 0:
2736
+ cache_seed.trim_layers(cache_layers, excess)
2737
+ except Exception: # noqa: BLE001 — snapshot path covers hybrids
2738
+ pass
2739
+ r = types.SimpleNamespace(
2740
+ finish_reason=finish_reason,
2741
+ prompt_cache=list(cache_layers),
2742
+ all_tokens=full_prompt + [int(t) for t in sub.token_ids],
2743
+ token=sub.token_ids[-1] if sub.token_ids else 0,
2744
+ )
2745
+ self._finish(sub, r)
2746
+ except Exception as exc: # noqa: BLE001
2747
+ print(f"[spec] wave failed request={sub.request_id}: {exc}", flush=True)
2748
+ # Tag every traceback line. The supervisor's stdout classifier
2749
+ # reads a bare `SomeError: ...` leaf line as proof the engine is
2750
+ # dead and tears the process down — right during startup, wrong
2751
+ # here, where this handler has already contained the failure and
2752
+ # the engine is still serving other requests.
2753
+ log_contained_exception("spec")
2754
+ sub.queue.put_nowait(("err", str(exc)))
2755
+ finally:
2756
+ self._drafter_busy = False
2757
+ if self._spec_sub is sub:
2758
+ self._spec_sub = None
2759
+
2760
+ def _spec_preempt_due(self, sub):
2761
+ if self._nested or not wave_policy.is_background(sub):
2762
+ return False
2763
+ return self._preempt_due([sub])
2764
+
2765
+ def _capture_spec_snapshot(self, sub, cache_layers, full_prompt, boundary):
2766
+ """Boundary snapshot for the spec wave's direct prefill. The cache
2767
+ is quiescent at the planted chunk edge; clone_layers reconstructs
2768
+ standalone layers (MLX arrays are immutable, so sharing buffers is
2769
+ safe), and the same token-prefix proof as _capture_prompt_snapshot
2770
+ guards the boundary invariant."""
2771
+ try:
2772
+ boundary = int(boundary)
2773
+ layers = spec_decode.clone_layers(cache_layers)
2774
+ if cache_seed.snapshot_matches_prompt(
2775
+ layers, full_prompt[:boundary], sub.prompt_tokens, boundary
2776
+ ):
2777
+ _store_snapshot(sub, (layers, boundary))
2778
+ elif not self._snapshot_warned:
2779
+ self._snapshot_warned = True
2780
+ print(
2781
+ f"[spec] prompt snapshot skipped: boundary={boundary} "
2782
+ f"(post-gen save fallback)",
2783
+ flush=True,
2784
+ )
2785
+ except Exception as exc: # noqa: BLE001 — snapshot is best-effort
2786
+ if not self._snapshot_warned:
2787
+ self._snapshot_warned = True
2788
+ print(f"[spec] prompt snapshot failed: {exc}", flush=True)
2789
+
2790
+ def _capture_prompt_snapshot(self, sub, at_tokens):
2791
+ """Capture the KV state at a mid-prompt boundary of `at_tokens`.
2792
+
2793
+ Called from the worker's prompt-response handling at segment /
2794
+ chunk edges (the worker consumes `BatchGenerator.next()` per
2795
+ internal step, so the batch is quiescent between calls). A
2796
+ segment edge is a token-index cut deep inside text both turns
2797
+ share, so it is immune to the BPE boundary problem — measured
2798
+ live: the end-of-prompt token sequence `…<start_of_turn>model\\n`
2799
+ re-merges differently once the next turn's text continues past
2800
+ it, so a full-length snapshot misses by ~2 tokens and wrapped
2801
+ rotating layers can't trim even that. The admission loop plants
2802
+ a deliberate cut `_SNAPSHOT_BOUNDARY_MARGIN` tokens before the
2803
+ prompt's end so every turn gets one stable, advancing boundary.
2804
+
2805
+ `extract_cache` materializes contiguous per-sequence copies and
2806
+ returns that sequence's unpadded token list, so the snapshot doesn't
2807
+ alias the live batch buffers or accidentally retain another sub's
2808
+ right-padding. Exact token-prefix + layer-offset checks are the runtime
2809
+ proof of the boundary invariant — on any mismatch the snapshot is
2810
+ dropped (logged once) and _finish falls back to the post-gen save.
2811
+ The same edge may be the session boundary, the band boundary, or both;
2812
+ `_store_snapshot` files it under whichever it is.
2813
+ """
2814
+ try:
2815
+ got = self._gen.extract_cache([sub.uid]).get(sub.uid)
2816
+ layers = got[0] if got else None
2817
+ extracted_tokens = got[1] if got else None
2818
+ if cache_seed.snapshot_matches_prompt(
2819
+ layers, extracted_tokens, sub.prompt_tokens, at_tokens
2820
+ ):
2821
+ _store_snapshot(sub, (list(layers), at_tokens))
2822
+ elif not self._snapshot_warned:
2823
+ self._snapshot_warned = True
2824
+ offsets = [getattr(c, "offset", None) for c in (layers or [])]
2825
+ print(
2826
+ f"[batch] prompt snapshot skipped: boundary={at_tokens} "
2827
+ f"extracted_tokens={len(extracted_tokens or [])} "
2828
+ f"offsets={offsets} (post-gen save fallback)",
2829
+ flush=True,
2830
+ )
2831
+ except Exception as exc: # noqa: BLE001 — snapshot is best-effort
2832
+ if not self._snapshot_warned:
2833
+ self._snapshot_warned = True
2834
+ print(f"[batch] prompt snapshot failed: {exc}", flush=True)
2835
+
2836
+ def _emit_delta(self, sub):
2837
+ # Full-decode-delta detokenization: correct across BPE merges /
2838
+ # multibyte chars (hold on a trailing replacement char until the
2839
+ # next token completes it). Fine for chat-length outputs.
2840
+ full = self._tok.decode(sub.token_ids)
2841
+ if full.endswith("�"):
2842
+ return
2843
+ # Strip leaked framing tokens before computing the delta. The scrub
2844
+ # is deterministic over the cumulative decode, so tracking `emitted`
2845
+ # against the scrubbed string stays consistent across steps.
2846
+ full = _scrub_leaked_markers(full)
2847
+ delta = full[sub.emitted :]
2848
+ if not delta:
2849
+ return
2850
+ sub.emitted = len(full)
2851
+ sub.queue.put_nowait(
2852
+ (
2853
+ "tok",
2854
+ delta,
2855
+ len(sub.token_ids),
2856
+ sub.prompt_tps,
2857
+ sub.generation_tps,
2858
+ )
2859
+ )
2860
+
2861
+ def _finish(self, sub, r):
2862
+ full = _scrub_leaked_markers(self._tok.decode(sub.token_ids))
2863
+ if len(full) > sub.emitted:
2864
+ sub.queue.put_nowait(
2865
+ (
2866
+ "tok",
2867
+ full[sub.emitted :],
2868
+ len(sub.token_ids),
2869
+ sub.prompt_tps,
2870
+ sub.generation_tps,
2871
+ )
2872
+ )
2873
+ sub.emitted = len(full)
2874
+ try:
2875
+ if sub.request.cache_id:
2876
+ # On finish the sequence is already out of the generation
2877
+ # batch, so extract_cache(uid) misses — the Response carries
2878
+ # the final KV (`prompt_cache`) + full token list
2879
+ # (`all_tokens`). Fall back to a token reconstruction.
2880
+ cache_layers = getattr(r, "prompt_cache", None)
2881
+ all_tokens = getattr(r, "all_tokens", None)
2882
+ if cache_layers is not None:
2883
+ post_gen = list(cache_layers)
2884
+ # Save decision: a post-generation cache is only worth
2885
+ # keeping if a future turn can trim it back to wherever
2886
+ # the re-templated history diverges from these raw
2887
+ # tokens. Un-trimmable stacks (wrapped windowed layers)
2888
+ # get the mid-prompt snapshot instead — a stable prefix
2889
+ # of the next prompt, no trim ever needed. A tiny turn
2890
+ # whose segment was too short to capture one keeps the
2891
+ # PRIOR entry (still a valid prefix) rather than
2892
+ # overwriting it with un-trimmable full-length state.
2893
+ trimmable = cache_seed.all_trimmable(post_gen)
2894
+ if not trimmable and sub.prompt_snapshot is None and sub.reused_tokens > 0:
2895
+ prior_entry = _CACHE.get(sub.request.cache_id)
2896
+ if prior_entry is not None:
2897
+ sub.saved_cache_state = prior_entry.state
2898
+ print(
2899
+ f"[batch] kept prior cache entry for "
2900
+ f"cache_id={sub.request.cache_id} (segment too short "
2901
+ f"for a stable snapshot; post-turn cache not trimmable)",
2902
+ flush=True,
2903
+ )
2904
+ else:
2905
+ use_snapshot = sub.prompt_snapshot is not None and not trimmable
2906
+ state = PromptCacheState()
2907
+ if use_snapshot:
2908
+ snap_layers, snap_count = sub.prompt_snapshot
2909
+ state.cache = snap_layers
2910
+ state.token_ids = list(sub.prompt_tokens[:snap_count])
2911
+ else:
2912
+ state.cache = post_gen
2913
+ state.token_ids = (
2914
+ [int(t) for t in all_tokens]
2915
+ if all_tokens is not None
2916
+ else list(sub.prompt_tokens) + list(sub.token_ids)
2917
+ )
2918
+ if state.token_ids:
2919
+ _save_cache(sub.request.cache_id, state)
2920
+ sub.saved_cache_state = state
2921
+ _LAST_SAVE_ROSTER[sub.request.cache_id] = sub.roster_fp
2922
+ prefix_state = _prefix_seed_state(sub, state, trimmable)
2923
+ if prefix_state is not None:
2924
+ _seed_prefix_from_session(sub.request, prefix_state)
2925
+ print(
2926
+ f"[batch] saved cache_id={sub.request.cache_id} "
2927
+ f"tokens={len(state.token_ids)}"
2928
+ + (
2929
+ " (prompt snapshot; post-turn cache not trimmable)"
2930
+ if use_snapshot
2931
+ else ""
2932
+ ),
2933
+ flush=True,
2934
+ )
2935
+ # Reuse parity with the serial path's terminal log — the
2936
+ # line that makes "is the cache actually working?" a one-
2937
+ # line grep instead of a timing investigation.
2938
+ print(
2939
+ f"[cache] reuse cache_id={sub.request.cache_id} "
2940
+ f"cached_tokens={sub.reused_tokens} "
2941
+ f"prefilled_tokens={sub.prefill_total} [batched]",
2942
+ flush=True,
2943
+ )
2944
+ except Exception as exc: # noqa: BLE001
2945
+ print(f"[batch] cache save failed: {exc}", flush=True)
2946
+ sub.queue.put_nowait(
2947
+ ("done", r.finish_reason, sub.prompt_tps, sub.generation_tps)
2948
+ )
2949
+
2950
+
2951
+ def _batched_tool_payload(frag, rid: str, model: str, in_toks: int, out_count: int, sub) -> Dict[str, Any]:
2952
+ """Batched-route envelope for one tool-stream fragment.
2953
+
2954
+ Mirrors `_tool_stream_payload` but sources usage from the sub rather than
2955
+ a generation chunk — the batched worker reports counters separately.
2956
+ """
2957
+ if frag.content is not None:
2958
+ delta: Dict[str, Any] = {"role": "assistant", "content": frag.content}
2959
+ else:
2960
+ call: Dict[str, Any] = {
2961
+ "index": frag.tool_call_index,
2962
+ "function": {"arguments": frag.arguments or ""},
2963
+ }
2964
+ if frag.tool_call_id is not None:
2965
+ call["id"] = frag.tool_call_id
2966
+ call["type"] = "function"
2967
+ call["function"]["name"] = frag.name or ""
2968
+ delta = {"role": "assistant", "tool_calls": [call]}
2969
+ return {
2970
+ "id": rid,
2971
+ "object": "chat.completion.chunk",
2972
+ "created": int(time.time()),
2973
+ "model": model,
2974
+ "choices": [{"index": 0, "delta": delta, "finish_reason": None}],
2975
+ "usage": {
2976
+ "input_tokens": in_toks,
2977
+ "output_tokens": out_count,
2978
+ "total_tokens": in_toks + out_count,
2979
+ "prompt_tps": 0.0,
2980
+ "generation_tps": 0.0,
2981
+ "cached_tokens": sub.reused_tokens,
2982
+ },
2983
+ }
2984
+
2985
+
2986
+ async def _batched_stream_iter(sub: "_Sub"):
2987
+ """SSE generator for the batched path — same wire frames as the serial
2988
+ route, fed from the engine worker's per-sub queue."""
2989
+ request = sub.request
2990
+ rid = sub.request_id
2991
+ in_toks = len(sub.prompt_tokens)
2992
+ # Same structured tool-call split as the serial route. The batched path
2993
+ # is the one the eval harness actually exercises (max-concurrency 1 still
2994
+ # runs through the batch engine), so wiring only the serial route left
2995
+ # the flag inert.
2996
+ tool_stream = _make_tool_stream(rid, getattr(request, 'tools', None))
2997
+ try:
2998
+ while True:
2999
+ item = await sub.queue.get()
3000
+ kind = item[0]
3001
+ if kind == "tok":
3002
+ _, delta, out_count, prompt_tps, generation_tps = item
3003
+ if await sub.http_request.is_disconnected():
3004
+ sub.cancelled = True
3005
+ return
3006
+ if tool_stream is not None:
3007
+ for _frag in tool_stream.feed(delta):
3008
+ yield "data: " + json.dumps(
3009
+ _batched_tool_payload(_frag, rid, request.model, in_toks, out_count, sub)
3010
+ ) + "\n\n"
3011
+ continue
3012
+ yield "data: " + json.dumps(
3013
+ {
3014
+ "id": rid,
3015
+ "object": "chat.completion.chunk",
3016
+ "created": int(time.time()),
3017
+ "model": request.model,
3018
+ "choices": [
3019
+ {
3020
+ "index": 0,
3021
+ "delta": {"role": "assistant", "content": delta},
3022
+ "finish_reason": None,
3023
+ }
3024
+ ],
3025
+ "usage": {
3026
+ "input_tokens": in_toks,
3027
+ "output_tokens": out_count,
3028
+ "total_tokens": in_toks + out_count,
3029
+ "prompt_tps": prompt_tps,
3030
+ "generation_tps": generation_tps,
3031
+ # Cache tokens actually served per the seed
3032
+ # plan — parity with the serial path so the
3033
+ # TS session's cachedInputTokens works on
3034
+ # both. 0 means a genuinely cold prefill.
3035
+ "cached_tokens": sub.reused_tokens,
3036
+ },
3037
+ }
3038
+ ) + "\n\n"
3039
+ elif kind == "done":
3040
+ _, reason, prompt_tps, generation_tps = item
3041
+ if tool_stream is not None:
3042
+ for _frag in tool_stream.flush():
3043
+ yield "data: " + json.dumps(
3044
+ _batched_tool_payload(
3045
+ _frag, rid, request.model, in_toks, len(sub.token_ids), sub
3046
+ )
3047
+ ) + "\n\n"
3048
+ yield "data: " + json.dumps(
3049
+ {
3050
+ "id": rid,
3051
+ "object": "chat.completion.chunk",
3052
+ "created": int(time.time()),
3053
+ "model": request.model,
3054
+ "choices": [
3055
+ {
3056
+ "index": 0,
3057
+ "delta": {"role": "assistant", "content": ""},
3058
+ "finish_reason": "stop" if reason == "stop" else "length",
3059
+ }
3060
+ ],
3061
+ "usage": {
3062
+ "input_tokens": in_toks,
3063
+ "output_tokens": len(sub.token_ids),
3064
+ "total_tokens": in_toks + len(sub.token_ids),
3065
+ "prompt_tps": prompt_tps,
3066
+ "generation_tps": generation_tps,
3067
+ "cached_tokens": sub.reused_tokens,
3068
+ },
3069
+ }
3070
+ ) + "\n\n"
3071
+ yield "data: [DONE]\n\n"
3072
+ return
3073
+ elif kind == "err":
3074
+ yield "data: " + json.dumps({"error": item[1]}) + "\n\n"
3075
+ return
3076
+ finally:
3077
+ # If we exit early (disconnect / error), make sure the worker evicts us.
3078
+ sub.cancelled = True
3079
+ # Release the eviction pins taken for this turn's cache + prefix.
3080
+ _unmark_inflight(sub.pinned_cache_ids)
3081
+ sub.pinned_cache_ids = []
3082
+ engine = _BATCH_ENGINE
3083
+ _reclaim_mlx_buffer_cache(
3084
+ "batch-turn-end",
3085
+ force=engine is None or engine.is_idle(),
3086
+ log=True,
3087
+ )
3088
+
3089
+
3090
+ async def _await_batched_completion(sub: "_Sub") -> None:
3091
+ """Drain a non-streaming BatchEngine sub (currently cache warming)."""
3092
+ completed = False
3093
+ try:
3094
+ while True:
3095
+ item = await sub.queue.get()
3096
+ kind = item[0]
3097
+ if kind == "done":
3098
+ completed = True
3099
+ return
3100
+ if kind == "err":
3101
+ raise RuntimeError(item[1])
3102
+ # Generated token deltas are intentionally discarded: warming
3103
+ # needs one decode step to finalize the cache, not its text.
3104
+ finally:
3105
+ if not completed:
3106
+ sub.cancelled = True
3107
+ _unmark_inflight(sub.pinned_cache_ids)
3108
+ sub.pinned_cache_ids = []
3109
+ engine = _BATCH_ENGINE
3110
+ _reclaim_mlx_buffer_cache(
3111
+ "batch-warm-end",
3112
+ force=engine is None or engine.is_idle(),
3113
+ log=True,
3114
+ )
3115
+
3116
+
3117
+ _BATCH_ENGINE: "Optional[BatchEngine]" = None
3118
+
3119
+
3120
+ def _get_batch_engine() -> "BatchEngine":
3121
+ global _BATCH_ENGINE
3122
+ if _BATCH_ENGINE is None:
3123
+ lm = getattr(MODEL, "language_model", MODEL)
3124
+ tok = getattr(PROCESSOR, "tokenizer", None) or PROCESSOR
3125
+ # mlx_lm models return raw logits — no LanguageModelOutput to unwrap.
3126
+ engine_model = lm if _TEXT_TOWER == "mlx_lm" else _UnwrapLM(lm)
3127
+ _BATCH_ENGINE = BatchEngine(engine_model, tok, ARGS.max_concurrency)
3128
+ return _BATCH_ENGINE
3129
+
3130
+
3131
+ def _flush_all_to_disk() -> int:
3132
+ """Write every in-memory entry to disk. Returns the count saved.
3133
+
3134
+ Called from the SIGTERM handler and the FastAPI shutdown event so a
3135
+ graceful supervisor stop preserves every warm session for the next
3136
+ boot. SIGKILL still loses everything (no handler runs); that's the
3137
+ pre-persistence cost as a worst case.
3138
+ """
3139
+ if not _disk_enabled():
3140
+ return 0
3141
+ saved = 0
3142
+ clean = 0
3143
+ started = time.perf_counter()
3144
+ for cache_id, entry in list(_CACHE.items()):
3145
+ if not entry.dirty:
3146
+ clean += 1
3147
+ continue
3148
+ if cache_persist.save_cache(
3149
+ PERSIST_DIR, MODEL_FINGERPRINT, cache_id, entry.state # type: ignore[arg-type]
3150
+ ):
3151
+ entry.dirty = False
3152
+ saved += 1
3153
+ if saved or clean:
3154
+ print(
3155
+ f"[cache] flushed {saved} entries to disk ({clean} already persisted) "
3156
+ f"in {time.perf_counter() - started:.1f}s",
3157
+ flush=True,
3158
+ )
3159
+ _enforce_disk_budget()
3160
+ return saved
3161
+
3162
+
3163
+ # ───────── HTTP request/response models ─────────
3164
+
3165
+
3166
+ class ChatMessageReq(BaseModel):
3167
+ role: str
3168
+ content: Any
3169
+ images: Optional[List[str]] = None # raw base64, never a file path or URL
3170
+ tool_call_id: Optional[str] = None
3171
+ # Undeclared fields are dropped by pydantic, and a dropped
3172
+ # `tool_calls` silently costs the model its own tool-call history:
3173
+ # Gemma's template renders a `role: "tool"` message only when it can
3174
+ # pair it with the preceding assistant call, so without this the
3175
+ # whole call/result exchange vanishes from the prompt. Wild-caught
3176
+ # on gemma4-26b-q4 — 62% of turns re-sent a prompt just ~5 tokens
3177
+ # longer than the last, the model never saw a single tool result,
3178
+ # and it re-issued the same call until the repeat-guard killed the
3179
+ # turn. Qwen masked it: Hermes templates render a bare tool message.
3180
+ tool_calls: Optional[List[Dict[str, Any]]] = None
3181
+
3182
+
3183
+ class ChatRequest(BaseModel):
3184
+ model: str
3185
+ messages: List[ChatMessageReq]
3186
+ stream: bool = True
3187
+ tools: Optional[List[Dict[str, Any]]] = None
3188
+ # Gezel extension. When set, the server preserves the post-generation
3189
+ # cache state under this key. When unset, behaves like upstream
3190
+ # mlx_vlm.server (cache cleared after the stream).
3191
+ cache_id: Optional[str] = None
3192
+ # Cross-session prefix sharing (Phase 1.3). When the session-specific
3193
+ # `cache_id` misses, the server falls back to loading this entry —
3194
+ # typically a per-gezel hash of the stable system prompt — and
3195
+ # registers it under the session's cache_id as a starting state.
3196
+ # First chat turn for a new session inherits the warm prefix
3197
+ # instead of cold-prefilling 2 KB of system prompt.
3198
+ prefix_cache_id: Optional[str] = None
3199
+ # Layered prefix sharing (flag `layeredPrefixCache`). An ordered,
3200
+ # most-specific-first list of prefix ids — typically
3201
+ # `[prefix-gp-<hash>, prefix-gezel-<hash>]`. On a session miss the
3202
+ # server tries each in order and seeds from the first hit, so a new
3203
+ # session inherits the warmest matching layer (gezel+project > gezel).
3204
+ # When set, takes precedence over the singular `prefix_cache_id`.
3205
+ prefix_cache_ids: Optional[List[str]] = None
3206
+ # Length in CHARACTERS of the shared band inside messages[0].content —
3207
+ # the leading run sibling sessions of the same (gezel, project) render
3208
+ # identically. When set, the engine plants its prompt-snapshot cut at
3209
+ # that boundary and publishes the SHORT snapshot as the prefix entry
3210
+ # instead of the session's full post-turn state.
3211
+ #
3212
+ # Why this matters here and not on llama.cpp: 48 of qwen3.8's 64 layers
3213
+ # are linear-attention `ArraysCache` holding a fixed-size recurrent
3214
+ # state with no `trim()`, so a prefix entry LONGER than the shared head
3215
+ # is not a partial hit — `seed_from_state` falls to `fresh-untrimmable`
3216
+ # and the whole prompt is re-prefilled. Short entries are the only
3217
+ # reusable shape. See ADR 0010.
3218
+ stable_prefix_chars: Optional[int] = None
3219
+ # Generation knobs (subset of mlx-vlm's). Not all are honored today —
3220
+ # we just need enough to run sensible defaults.
3221
+ max_tokens: Optional[int] = None
3222
+ temperature: Optional[float] = None
3223
+ top_p: Optional[float] = None
3224
+ top_k: Optional[int] = None
3225
+ repetition_penalty: Optional[float] = None
3226
+ repetition_context_size: Optional[int] = None
3227
+ # `tuning.reasoning.thinkingBudget` — enforced here because MLX has no
3228
+ # `--reasoning-budget` equivalent; see think_budget.py.
3229
+ max_thinking_tokens: Optional[int] = None
3230
+ # Decode-time tool-call grammar (gezel extension). When set, the
3231
+ # server builds an llguidance logits processor that constrains the
3232
+ # function name of any tool call to the known tools — see
3233
+ # `build_tool_grammar_processor`. Shape: {"format": "...", "mode": "name-only"}.
3234
+ tool_grammar: Optional[Dict[str, Any]] = None
3235
+ # Per-request circuit breaker for a session that observed the assisted
3236
+ # speculative path violate a required-argument grammar. The TS provider
3237
+ # latches this after the first impossible call and leaves the drafter
3238
+ # resident for unaffected sessions; only this request takes the ordinary
3239
+ # BatchGenerator path. This is the request-scoped equivalent of the
3240
+ # operator's GEZEL_MLX_SPEC=greedy-only A/B arm.
3241
+ disable_speculation: bool = False
3242
+ # Who is waiting on this turn: "interactive" (a person typed it) or
3243
+ # "background" (task steps, nudges, warms). A running wave made only of
3244
+ # background work parks for a waiting interactive request; absent reads
3245
+ # as interactive so an older daemon never gets its turns parked.
3246
+ priority: Optional[str] = None
3247
+ # Per-request chat-template override (gezel extension). Swaps the
3248
+ # model's stored Jinja template for a curated one without reinstalling.
3249
+ chat_template_override: Optional[str] = None
3250
+ # Template variables the caller wants bound when rendering the prompt —
3251
+ # `enable_thinking`, `reasoning_effort`, and friends. llama-server has
3252
+ # honored this field for a while and the MLX provider has always SENT it;
3253
+ # until it was declared here pydantic dropped it silently, so every
3254
+ # constrained turn that asked for thinking to be off still rendered a
3255
+ # prompt ending in `<think>`. That one missing field was the bulk of the
3256
+ # measured MLX/llama-cpp divergence — see `_build_prompt`.
3257
+ chat_template_kwargs: Optional[Dict[str, Any]] = None
3258
+
3259
+
3260
+ class CacheStatsEntry(BaseModel):
3261
+ cache_id: str
3262
+ token_count: int
3263
+ est_bytes: int
3264
+ last_used_ts: float
3265
+
3266
+
3267
+ # ───────── FastAPI app ─────────
3268
+
3269
+
3270
+ app = FastAPI(title="gezel-mlx-server")
3271
+
3272
+
3273
+ @app.on_event("shutdown")
3274
+ def _on_shutdown() -> None:
3275
+ """uvicorn fires this on SIGTERM (graceful) before exiting.
3276
+
3277
+ Persisting here means a clean supervisor stop preserves every warm
3278
+ session — the user comes back after lunch, the prefix loads from
3279
+ disk in <200ms instead of paying a full re-prefill.
3280
+ """
3281
+ _flush_all_to_disk()
3282
+
3283
+
3284
+ # SIGTERM handler as a belt-and-braces complement to FastAPI's
3285
+ # `on_event('shutdown')`. uvicorn's signal handling installs its own
3286
+ # trap before ours, so the on_event hook is the primary path; this
3287
+ # fallback runs only when uvicorn's handler doesn't reach the
3288
+ # shutdown event (some kill scenarios race the handler installation).
3289
+ def _sigterm_handler(_signum: int, _frame: Any) -> None:
3290
+ _flush_all_to_disk()
3291
+ # Re-raise default behavior — let uvicorn's own shutdown finish.
3292
+ sys.exit(0)
3293
+
3294
+
3295
+ try:
3296
+ import signal as _signal
3297
+
3298
+ _signal.signal(_signal.SIGTERM, _sigterm_handler)
3299
+ except (ValueError, AttributeError):
3300
+ # `signal.signal` only works on the main thread; if we somehow
3301
+ # aren't there, on_event is still wired and that's the common path.
3302
+ pass
3303
+
3304
+
3305
+ @app.get("/health")
3306
+ async def health() -> JSONResponse:
3307
+ """Same shape mlx_vlm.server emits — keeps our supervisor probe happy."""
3308
+ return JSONResponse({"status": "ok", "vision": _VISION_ENABLED})
3309
+
3310
+
3311
+ @app.post("/admin/flush")
3312
+ async def admin_flush(settle_ms: int = 0) -> JSONResponse:
3313
+ """Persist every in-memory cache entry to disk without dropping it.
3314
+
3315
+ Called by the gezel supervisor's Stage-1 idle freeze: the user has
3316
+ been quiet long enough that we should defensively save state, but
3317
+ not so long that we want to evict the model from unified memory.
3318
+ If the user comes back during the freeze window the in-memory
3319
+ cache is still warm; if they DON'T come back and Stage 2 SIGTERMs
3320
+ the process, the freeze flush already wrote everything out so the
3321
+ next boot resumes cleanly.
3322
+
3323
+ Also called right before a deliberate stop, with `settle_ms`: the
3324
+ daemon cancels its in-flight turns first, and a cancelled turn files
3325
+ its prompt snapshot only when the worker next steps. Waiting (bounded)
3326
+ for the engine to drain is what gets that snapshot into this flush
3327
+ instead of losing it to the SIGKILL that follows SIGTERM.
3328
+ """
3329
+ engine = _BATCH_ENGINE
3330
+ settled = True
3331
+ if settle_ms > 0 and engine is not None:
3332
+ deadline = time.time() + min(int(settle_ms), 60_000) / 1000.0
3333
+ while not engine.is_idle() and time.time() < deadline:
3334
+ await asyncio.sleep(0.05)
3335
+ settled = engine.is_idle()
3336
+ saved = _flush_all_to_disk()
3337
+ return JSONResponse({"flushed": saved, "settled": settled})
3338
+
3339
+
3340
+ @app.get("/v1/cache/stats")
3341
+ async def cache_stats() -> List[CacheStatsEntry]:
3342
+ """In-memory entries only.
3343
+
3344
+ Intentionally excludes disk-only persisted entries — the controller
3345
+ polling this endpoint manages the in-memory budget against the
3346
+ in-engine LRU. Disk persistence is a separate, larger budget
3347
+ enforced internally by _enforce_disk_budget on each save. Surfacing
3348
+ disk entries here would force the controller to evict them as if
3349
+ they were live, defeating the persistence behavior.
3350
+ """
3351
+ return [
3352
+ CacheStatsEntry(
3353
+ cache_id=cid,
3354
+ token_count=len(entry.state.token_ids or []),
3355
+ est_bytes=entry.est_bytes,
3356
+ last_used_ts=entry.last_used_ts,
3357
+ )
3358
+ for cid, entry in _CACHE.items()
3359
+ ]
3360
+
3361
+
3362
+ @app.delete("/v1/cache/{cache_id}")
3363
+ async def cache_evict(cache_id: str) -> JSONResponse:
3364
+ """Drop an entry from both in-memory cache and disk persistence.
3365
+
3366
+ Disk + memory in lockstep is intentional: a controller-driven evict
3367
+ means the operator (or an LRU policy) decided this session's cache
3368
+ is no longer useful. Leaving the disk entry behind would defeat
3369
+ the eviction (next access would re-load it). Deleting both is the
3370
+ "clean slate" the controller intends.
3371
+ """
3372
+ in_memory = cache_id in _CACHE
3373
+ if in_memory:
3374
+ del _CACHE[cache_id]
3375
+ on_disk = (
3376
+ cache_persist.delete_cache(PERSIST_DIR, MODEL_FINGERPRINT, cache_id) # type: ignore[arg-type]
3377
+ if _disk_enabled()
3378
+ else False
3379
+ )
3380
+ if in_memory or on_disk:
3381
+ print(
3382
+ f"[cache] evicted cache_id={cache_id} (memory={in_memory} disk={on_disk})",
3383
+ flush=True,
3384
+ )
3385
+ return JSONResponse({"evicted": in_memory or on_disk, "cache_id": cache_id})
3386
+
3387
+
3388
+ class CacheWarmRequest(BaseModel):
3389
+ cache_id: str
3390
+ messages: List[ChatMessageReq]
3391
+ model: Optional[str] = None
3392
+ # A warmed prefix is only reusable if it is a genuine token-prefix of the
3393
+ # real turns that follow, which means rendering through the SAME template
3394
+ # branch they do. Without `tools` this path fell to the non-tokenizer
3395
+ # render, and on Qwen — which puts the tool block at the top of the system
3396
+ # message — the warmed entry shared 3 tokens with every real prompt. Every
3397
+ # prefix-hit was a false hit: disk load, eviction pin, then a full
3398
+ # prefill anyway.
3399
+ tools: Optional[List[Dict[str, Any]]] = None
3400
+ chat_template_kwargs: Optional[Dict[str, Any]] = None
3401
+ # When True, write the warmed entry to disk immediately (don't wait
3402
+ # for eviction or shutdown). Used by the gezel-prefix warming flow
3403
+ # so subsequent sessions can load the prefix on miss without
3404
+ # depending on supervisor lifecycle. Default False matches the
3405
+ # session-warm semantics (in-memory only until evicted).
3406
+ persist: bool = False
3407
+
3408
+
3409
+ @app.post("/v1/cache/warm")
3410
+ async def cache_warm(req: CacheWarmRequest) -> JSONResponse:
3411
+ """Warm a cache through the same snapshot-capable BatchGenerator as chat.
3412
+
3413
+ The older serial warm saved post-generation state. That is reusable for
3414
+ full-attention caches (they can trim back to the shared prefix), but not for
3415
+ wrapped sliding-window caches: their synthetic user/generated tail cannot
3416
+ be rewound, so the real first turn paid a full prefill anyway.
3417
+
3418
+ System-only warms still need a synthetic user for Qwen-family templates.
3419
+ We render a second, deliberately different synthetic user and take their
3420
+ token LCP minus the normal safety margin. That gives BatchEngine an exact
3421
+ snapshot boundary inside the stable system prefix, before either fake user
3422
+ can contaminate an untrimmable entry.
3423
+ """
3424
+ if any(m.images for m in req.messages):
3425
+ raise HTTPException(status_code=422, detail="Vision requests cannot warm text-prefix caches")
3426
+ synthetic_user = not any(m.role == "user" for m in req.messages)
3427
+ warm_messages = list(req.messages)
3428
+ if synthetic_user:
3429
+ warm_messages.append(ChatMessageReq(role="user", content="hi"))
3430
+
3431
+ prompt_text = _build_prompt(
3432
+ warm_messages, req.tools, None, req.chat_template_kwargs
3433
+ )
3434
+ tok = getattr(PROCESSOR, "tokenizer", None) or PROCESSOR
3435
+ prompt_tokens = [int(t) for t in tok.encode(prompt_text)]
3436
+ snapshot_target = None
3437
+ if synthetic_user:
3438
+ alternate_messages = list(req.messages) + [
3439
+ ChatMessageReq(role="user", content="__gezel_cache_boundary_probe__")
3440
+ ]
3441
+ alternate_text = _build_prompt(
3442
+ alternate_messages, req.tools, None, req.chat_template_kwargs
3443
+ )
3444
+ alternate_tokens = [int(t) for t in tok.encode(alternate_text)]
3445
+ target = cache_seed.stable_snapshot_boundary(
3446
+ prompt_tokens, alternate_tokens, _SNAPSHOT_BOUNDARY_MARGIN
3447
+ )
3448
+ snapshot_target = target
3449
+
3450
+ # A warm entry must be stable on EVERY axis the real turns vary, not just
3451
+ # the synthetic user. Measured: warming stopped at the synthetic-user
3452
+ # boundary (2,509 tokens) while a constrained turn diverges 27 tokens
3453
+ # earlier at the reasoning preamble (2,482) — and an untrimmable cache
3454
+ # cannot rewind 27 tokens any more than it can rewind 2,000, so every
3455
+ # prefix hit was still a full prefill. Take the tighter of the two.
3456
+ reasoning_target = _reasoning_stable_boundary(req, prompt_tokens)
3457
+ if reasoning_target:
3458
+ snapshot_target = (
3459
+ reasoning_target if not snapshot_target else min(snapshot_target, reasoning_target)
3460
+ )
3461
+ if snapshot_target == 0:
3462
+ print(
3463
+ f"[cache] no stable synthetic-user boundary for "
3464
+ f"cache_id={req.cache_id}; using normal finish save",
3465
+ flush=True,
3466
+ )
3467
+
3468
+ warm_request = ChatRequest(
3469
+ model=req.model or str(ARGS.model),
3470
+ messages=warm_messages,
3471
+ stream=True,
3472
+ cache_id=req.cache_id,
3473
+ max_tokens=1,
3474
+ temperature=0.0,
3475
+ # Nobody reads a warm's output; it must never hold a person up.
3476
+ priority=wave_policy.BACKGROUND,
3477
+ )
3478
+ cache_state, _, _ = _resolve_cache_state(warm_request)
3479
+ prior_tokens = len(cache_state.token_ids or [])
3480
+ warm_roster_fp = _roster_fp_of(req.tools)
3481
+
3482
+ # A warm must never race a live turn for the same cache_id. Both write
3483
+ # the same entry, so whichever finishes last wins, and the warm's prompt
3484
+ # is by construction the SHORT one. Observed: a 520-token warm landed on
3485
+ # a session mid-turn and replaced its 121,508-token entry with 521
3486
+ # tokens; the in-flight turn had already seeded so it survived, but the
3487
+ # next turn on that session would have re-prefilled 121k tokens from
3488
+ # scratch. The live turn owns the id — refuse, don't queue.
3489
+ if _INFLIGHT_CACHE_IDS.get(req.cache_id, 0) > 0:
3490
+ print(
3491
+ f"[cache] warm skipped for cache_id={req.cache_id}: a request is already "
3492
+ f"in flight on that cache (prior={prior_tokens})",
3493
+ flush=True,
3494
+ )
3495
+ return JSONResponse(
3496
+ {
3497
+ "warmed": False,
3498
+ "cache_id": req.cache_id,
3499
+ "token_count": prior_tokens,
3500
+ "persisted": False,
3501
+ "skipped": "in-flight",
3502
+ }
3503
+ )
3504
+
3505
+ # The tool roster IS part of the prefix identity — it is rendered into
3506
+ # the prompt ahead of everything else, so warming with a different
3507
+ # roster diverges at token ~1 and the "warm" is a full prefill that
3508
+ # then overwrites a good entry with a useless one. Refuse when the
3509
+ # rosters disagree and there is a real entry to lose. An unknown saved
3510
+ # roster is not a mismatch (nothing recorded yet).
3511
+ saved_roster_fp = _LAST_SAVE_ROSTER.get(req.cache_id)
3512
+ if prior_tokens > 0 and saved_roster_fp and saved_roster_fp != warm_roster_fp:
3513
+ print(
3514
+ f"[cache] warm skipped for cache_id={req.cache_id}: roster mismatch "
3515
+ f"(warm={warm_roster_fp} saved={saved_roster_fp}, prior={prior_tokens}) — "
3516
+ f"warming a different tool shape would churn the whole prefix",
3517
+ flush=True,
3518
+ )
3519
+ return JSONResponse(
3520
+ {
3521
+ "warmed": False,
3522
+ "cache_id": req.cache_id,
3523
+ "token_count": prior_tokens,
3524
+ "persisted": False,
3525
+ "skipped": "roster-mismatch",
3526
+ }
3527
+ )
3528
+
3529
+ sub = _Sub(
3530
+ request=warm_request,
3531
+ prompt_tokens=prompt_tokens,
3532
+ grammar=None,
3533
+ seed_state=cache_state,
3534
+ request_id=f"cachewarm-{uuid.uuid4()}",
3535
+ http_request=None,
3536
+ snapshot_target=snapshot_target,
3537
+ roster_fp=warm_roster_fp,
3538
+ )
3539
+ sub.pinned_cache_ids = _mark_inflight(req.cache_id)
3540
+ sub.prefill_total = len(prompt_tokens)
3541
+ try:
3542
+ _get_batch_engine().submit(sub)
3543
+ except Exception as exc:
3544
+ _unmark_inflight(sub.pinned_cache_ids)
3545
+ sub.pinned_cache_ids = []
3546
+ log_contained_exception("http")
3547
+ raise HTTPException(status_code=500, detail=str(exc))
3548
+ try:
3549
+ await _await_batched_completion(sub)
3550
+ except Exception as exc:
3551
+ log_contained_exception("http")
3552
+ raise HTTPException(status_code=500, detail=str(exc))
3553
+
3554
+ cache_state = sub.saved_cache_state
3555
+ new_tokens = len(cache_state.token_ids or []) if cache_state is not None else 0
3556
+ persisted = False
3557
+ if cache_state is not None and new_tokens > 0:
3558
+ # Prefix warming sets persist=True so the entry survives the next
3559
+ # supervisor stop and brand-new sessions can load it on first
3560
+ # request. Session warming leaves it in-memory only; the next
3561
+ # eviction will write through.
3562
+ if req.persist and _disk_enabled():
3563
+ persisted = cache_persist.save_cache(
3564
+ PERSIST_DIR, MODEL_FINGERPRINT, req.cache_id, cache_state # type: ignore[arg-type]
3565
+ )
3566
+ resident = _CACHE.get(req.cache_id)
3567
+ if persisted and resident is not None and resident.state is cache_state:
3568
+ resident.dirty = False
3569
+ _enforce_disk_budget()
3570
+ print(
3571
+ f"[cache] warmed cache_id={req.cache_id} tokens={new_tokens} "
3572
+ f"(prior={prior_tokens}, persist={persisted})",
3573
+ flush=True,
3574
+ )
3575
+ else:
3576
+ print(
3577
+ f"[cache] warm produced no tokens for cache_id={req.cache_id}; "
3578
+ f"leaving prior entry intact (prior={prior_tokens})",
3579
+ flush=True,
3580
+ )
3581
+ return JSONResponse(
3582
+ {
3583
+ "warmed": new_tokens > 0,
3584
+ "cache_id": req.cache_id,
3585
+ "token_count": new_tokens,
3586
+ "persisted": persisted,
3587
+ }
3588
+ )
3589
+
3590
+
3591
+ def _template_tool_calls(tool_calls: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
3592
+ """Shape OpenAI-wire tool calls for Jinja chat templates.
3593
+
3594
+ The wire format carries `function.arguments` as a JSON *string*;
3595
+ HF chat templates near-universally expect a *mapping* and iterate it
3596
+ (`arguments | items`, `arguments.path`). Qwen's template raises
3597
+ `TypeError: Can only get item pairs from a mapping` on a string,
3598
+ which `apply_chat_template`'s caller catches — and its last-resort
3599
+ recovery is to re-render with `tools` dropped entirely, so a string
3600
+ here costs the model its whole tool surface rather than one message.
3601
+ Gemma's template tolerates both, which is why this only bites some
3602
+ families. Parse when we can; pass through untouched when we can't
3603
+ (a non-JSON string is still better than an exception).
3604
+ """
3605
+ shaped: List[Dict[str, Any]] = []
3606
+ for call in tool_calls:
3607
+ if not isinstance(call, dict):
3608
+ shaped.append(call)
3609
+ continue
3610
+ fn = call.get("function")
3611
+ args = fn.get("arguments") if isinstance(fn, dict) else None
3612
+ if not isinstance(args, str):
3613
+ shaped.append(call)
3614
+ continue
3615
+ try:
3616
+ parsed = json.loads(args)
3617
+ except (ValueError, TypeError):
3618
+ shaped.append(call)
3619
+ continue
3620
+ if not isinstance(parsed, dict):
3621
+ shaped.append(call)
3622
+ continue
3623
+ shaped.append({**call, "function": {**fn, "arguments": parsed}})
3624
+ return shaped
3625
+
3626
+
3627
+ _TOOL_LINKAGE_PROBE: Dict[str, bool] = {}
3628
+ _LINKAGE_PROBE_MARK = "__gezel_tool_linkage_probe__"
3629
+
3630
+
3631
+ def _tool_linkage_probe_messages(linked: bool) -> List[Dict[str, Any]]:
3632
+ call = {
3633
+ "id": "probe1",
3634
+ "type": "function",
3635
+ "function": {"name": "probe_tool", "arguments": {"q": "x"}},
3636
+ }
3637
+ assistant: Dict[str, Any] = {"role": "assistant", "content": ""}
3638
+ result: Dict[str, Any] = {"role": "tool", "content": _LINKAGE_PROBE_MARK}
3639
+ if linked:
3640
+ assistant["tool_calls"] = [call]
3641
+ result["tool_call_id"] = "probe1"
3642
+ return [{"role": "user", "content": "go"}, assistant, result]
3643
+
3644
+
3645
+ def _template_needs_tool_linkage(tokenizer: Any, chat_template_override: Any) -> bool:
3646
+ """Does this model's template DROP a tool result unless the assistant
3647
+ call it answers is linked to it?
3648
+
3649
+ Gemma renders a `role: "tool"` message only as the response half of a
3650
+ call/response pair, so without the linkage the entire exchange vanishes
3651
+ and the model never sees a single tool result. Hermes-style templates
3652
+ (Qwen et al.) render a bare tool message just fine.
3653
+
3654
+ That difference is why this is probed rather than always-on. Feeding
3655
+ Qwen the linkage is *more* correct on paper but measurably worse in
3656
+ practice: it starts seeing its own past `<tool_call>` blocks echoed in
3657
+ the transcript and imitates them — core-suite MLX went 8/11 -> 5/11,
3658
+ tool calls 457 -> 753, three fast passes turning into slow model-stuck
3659
+ retry loops. So we only add the linkage where its absence actually
3660
+ costs the model information, and leave every other template on the byte-
3661
+ identical legacy shape.
3662
+
3663
+ Probed once per template: render a marker-bearing tool result with and
3664
+ without linkage; if the bare render loses the marker, linkage is
3665
+ required. Failures answer "no" — the legacy shape is the safe default.
3666
+ """
3667
+ key = str(chat_template_override) if chat_template_override else "<model>"
3668
+ cached = _TOOL_LINKAGE_PROBE.get(key)
3669
+ if cached is not None:
3670
+ return cached
3671
+ probe_tools = [
3672
+ {
3673
+ "type": "function",
3674
+ "function": {
3675
+ "name": "probe_tool",
3676
+ "description": "probe",
3677
+ "parameters": {
3678
+ "type": "object",
3679
+ "properties": {"q": {"type": "string"}},
3680
+ },
3681
+ },
3682
+ }
3683
+ ]
3684
+ kwargs: Dict[str, Any] = dict(
3685
+ tools=probe_tools, add_generation_prompt=True, tokenize=False
3686
+ )
3687
+ if chat_template_override:
3688
+ kwargs["chat_template"] = chat_template_override
3689
+ needs = False
3690
+ try:
3691
+ bare = tokenizer.apply_chat_template(
3692
+ _tool_linkage_probe_messages(False), **kwargs
3693
+ )
3694
+ linked = tokenizer.apply_chat_template(
3695
+ _tool_linkage_probe_messages(True), **kwargs
3696
+ )
3697
+ needs = (_LINKAGE_PROBE_MARK not in bare) and (_LINKAGE_PROBE_MARK in linked)
3698
+ except Exception as exc:
3699
+ print(
3700
+ f"[tool-prompt] tool-linkage probe failed ({exc}); "
3701
+ f"using legacy message shape",
3702
+ flush=True,
3703
+ )
3704
+ needs = False
3705
+ _TOOL_LINKAGE_PROBE[key] = needs
3706
+ print(
3707
+ f"[tool-prompt] tool-linkage probe: template "
3708
+ f"{'REQUIRES' if needs else 'does not need'} assistant/tool pairing"
3709
+ f"{' — passing tool_calls + tool_call_id through' if needs else ''}",
3710
+ flush=True,
3711
+ )
3712
+ return needs
3713
+
3714
+
3715
+ # Template variables naming reasoning DEPTH rather than the on/off switch.
3716
+ # Mirrors `REASONING_DEPTH_TEMPLATE_KWARGS` in providers/reasoning-depth.ts;
3717
+ # kept in sync by reasoning-template-kwargs-wiring.test.ts.
3718
+ _REASONING_DEPTH_KWARGS = frozenset({"reasoning_effort", "reasoning_strength"})
3719
+
3720
+
3721
+ def _reasoning_stable_boundary(request, prompt_tokens) -> Optional[int]:
3722
+ """Last prompt position that does not depend on this turn's reasoning mode.
3723
+
3724
+ Gezel flips `enable_thinking` between constrained and ordinary turns, and
3725
+ the template renders a reasoning preamble whose presence follows it. With
3726
+ the preamble relocated below the stable content that difference sits at
3727
+ the END of the system block — close enough to the whole prompt to look
3728
+ harmless, and fatal anyway, because a hybrid (linear-attention) cache
3729
+ cannot be trimmed back by even one token.
3730
+
3731
+ Returns an absolute boundary to snapshot at, or None to keep the default
3732
+ end-of-prompt behaviour (correct for models whose cache CAN trim).
3733
+ """
3734
+ if not _needs_reasoning_stable_snapshot():
3735
+ return None
3736
+ kwargs = dict(request.chat_template_kwargs or {})
3737
+ flipped = dict(kwargs)
3738
+ flipped["enable_thinking"] = not bool(kwargs.get("enable_thinking", True))
3739
+ try:
3740
+ alternate = _build_prompt(
3741
+ request.messages, request.tools, request.chat_template_override, flipped
3742
+ )
3743
+ tok = getattr(PROCESSOR, "tokenizer", None) or PROCESSOR
3744
+ alternate_tokens = [int(t) for t in tok.encode(alternate)]
3745
+ except Exception as exc: # noqa: BLE001 — never fail a turn over a boundary hint
3746
+ print(f"[cache] reasoning-stable boundary unavailable: {exc}", flush=True)
3747
+ return None
3748
+ boundary = cache_seed.stable_snapshot_boundary(
3749
+ prompt_tokens, alternate_tokens, _SNAPSHOT_BOUNDARY_MARGIN
3750
+ )
3751
+ # 0 disables capture entirely in plan_snapshot_segments. If the two modes
3752
+ # diverge immediately there is nothing stable to save, and the default
3753
+ # end-of-prompt snapshot would just re-save an unusable entry.
3754
+ return boundary or None
3755
+
3756
+
3757
+ def _needs_reasoning_stable_snapshot() -> bool:
3758
+ """True when a volatile tail in the saved state costs the ENTIRE entry.
3759
+
3760
+ Deliberately the same predicate the engine already uses to decide whether
3761
+ to capture prompt snapshots at all (`prompt_snapshot=y` in the boot log),
3762
+ rather than a second hand-rolled cache probe: an earlier version called
3763
+ `all_trimmable(MODEL.make_cache())` and defaulted to False on any problem,
3764
+ which silently disabled this fix on precisely the untrimmable models it
3765
+ exists for. One predicate, one answer.
3766
+ """
3767
+ if not _STABLE_PREFIX_ENABLED:
3768
+ return False
3769
+ global _REASONING_SNAPSHOT_NEEDED
3770
+ if _REASONING_SNAPSHOT_NEEDED is None:
3771
+ # Probe the SAME object the engine probes — `_get_batch_engine`
3772
+ # builds `BatchEngine(_UnwrapLM(lm))`, and the VLM wrapper's
3773
+ # `make_cache` is not the language model's. Probing the wrapper
3774
+ # reported a trimmable cache for a model that logs
3775
+ # `prompt_snapshot=y` and `fresh-untrimmable` in the same run.
3776
+ _REASONING_SNAPSHOT_NEEDED = cache_seed.probe_needs_prompt_snapshot(
3777
+ getattr(MODEL, "language_model", MODEL)
3778
+ )
3779
+ print(
3780
+ f"[cache] reasoning-stable snapshot "
3781
+ f"{'ENABLED (cache cannot be rewound)' if _REASONING_SNAPSHOT_NEEDED else 'not needed (trimmable cache)'}",
3782
+ flush=True,
3783
+ )
3784
+ return _REASONING_SNAPSHOT_NEEDED
3785
+
3786
+
3787
+ _REASONING_SNAPSHOT_NEEDED = None
3788
+
3789
+
3790
+ def _build_prompt(
3791
+ messages: List[ChatMessageReq],
3792
+ tools: Optional[List[Dict[str, Any]]] = None,
3793
+ chat_template_override: Optional[str] = None,
3794
+ chat_template_kwargs: Optional[Dict[str, Any]] = None,
3795
+ ) -> str:
3796
+ """Apply the model's chat template to OpenAI-shape messages.
3797
+
3798
+ When `tools` is provided, route through the tokenizer's
3799
+ `apply_chat_template(..., tools=...)` directly. Qwen's tokenizer
3800
+ (and most tool-trained models) renders the tools list using the
3801
+ model's *canonical* tool-use template — this is what tells the
3802
+ model "use `<tool_call>{json}</tool_call>` envelopes" reliably.
3803
+ Without this, the tools array gets dropped on the floor and the
3804
+ model improvises a tool-call format from training memory.
3805
+
3806
+ Failure handling peels off optional kwargs one at a time rather than
3807
+ silently abandoning `tools` on the first error: `enable_thinking` is
3808
+ a Qwen3-only kwarg, so a template that rejects it must NOT cost us the
3809
+ tool rendering. Dropping `tools` is the last resort and is logged
3810
+ loudly (it directly drives tool-call mangling, so it should surface in
3811
+ eval mechanism metrics rather than hide behind salvage).
3812
+
3813
+ `chat_template_override`, when set, swaps the model's stored Jinja
3814
+ template for a curated one (per-family template fix) without
3815
+ reinstalling — used for the A/B and as an escape hatch for models
3816
+ with a known-bad embedded template.
3817
+
3818
+ Vision messages keep structured image placeholders in the canonical
3819
+ template; decoded pixels are supplied separately to mlx_vlm.
3820
+ """
3821
+ # Text-only models (e.g. laguna): load_processor returns the bare
3822
+ # tokenizer itself — no `.tokenizer` attribute. Falling back to
3823
+ # PROCESSOR keeps the canonical tools render; without it this
3824
+ # function silently drops `tools` via the text-only path below.
3825
+ _probe_tokenizer = getattr(PROCESSOR, "tokenizer", None) or PROCESSOR
3826
+ link_tools = (
3827
+ any(m.tool_calls or m.tool_call_id for m in messages)
3828
+ and _probe_tokenizer is not None
3829
+ and hasattr(_probe_tokenizer, "apply_chat_template")
3830
+ and _template_needs_tool_linkage(_probe_tokenizer, chat_template_override)
3831
+ )
3832
+ raw = []
3833
+ for m in messages:
3834
+ entry: Dict[str, Any] = {"role": m.role, "content": vision_inputs.message_content(m)}
3835
+ # Carry the tool-call linkage through to the template, but ONLY for
3836
+ # templates that actually need it — see _template_needs_tool_linkage.
3837
+ # Templates that render a bare tool message stay byte-identical to
3838
+ # the legacy shape (which also keeps their prompt caches valid).
3839
+ if link_tools:
3840
+ if m.tool_calls:
3841
+ entry["tool_calls"] = _template_tool_calls(m.tool_calls)
3842
+ if m.tool_call_id:
3843
+ entry["tool_call_id"] = m.tool_call_id
3844
+ raw.append(entry)
3845
+
3846
+ tokenizer = getattr(PROCESSOR, "tokenizer", None) or PROCESSOR
3847
+ use_tokenizer = (
3848
+ tokenizer is not None
3849
+ and hasattr(tokenizer, "apply_chat_template")
3850
+ # Connected apps often send no tools or images. Their reasoning
3851
+ # controls still need the canonical renderer; the text-only fallback
3852
+ # drops these kwargs and silently reopens Qwen's thinking block.
3853
+ and (
3854
+ bool(tools)
3855
+ or bool(chat_template_override)
3856
+ or bool(chat_template_kwargs)
3857
+ or any(m.images for m in messages)
3858
+ )
3859
+ )
3860
+ if use_tokenizer:
3861
+ base_kwargs: Dict[str, Any] = dict(
3862
+ tools=tools,
3863
+ add_generation_prompt=True,
3864
+ tokenize=False,
3865
+ )
3866
+ # A caller-supplied override always wins — it is the escape hatch for
3867
+ # a known-bad embedded template and must not be silently second-
3868
+ # guessed. Otherwise use the prefix-stable rewrite when one verified.
3869
+ if chat_template_override:
3870
+ base_kwargs["chat_template"] = chat_template_override
3871
+ elif STABLE_CHAT_TEMPLATE:
3872
+ base_kwargs["chat_template"] = STABLE_CHAT_TEMPLATE
3873
+ # Reasoning defaults ON (Qwen3 reasoning forced inside <think>), then
3874
+ # the caller's template vars override it. The default has to stay True:
3875
+ # every model whose template reads the switch was tuned here, and
3876
+ # flipping the default would silently disable reasoning for ordinary
3877
+ # turns rather than only the constrained ones the caller names.
3878
+ template_vars: Dict[str, Any] = {"enable_thinking": True}
3879
+ if chat_template_kwargs:
3880
+ template_vars.update(chat_template_kwargs)
3881
+ # Attempt order, most-capable first. The DEPTH dials get their own
3882
+ # rung above the switch because they are the ones that reject values:
3883
+ # Qwen 3.8's jinja `raise_exception`s on any effort outside
3884
+ # {xhigh, medium, low}, and a template raises rather than returning —
3885
+ # so a bad depth value must not cost us `enable_thinking`, which is
3886
+ # worth far more (it decides whether the turn reasons at all).
3887
+ # 1. + all template vars (switch + caller's depth)
3888
+ # 2. - depth dials (template rejected a depth VALUE)
3889
+ # 3. - template vars (template reads none of them)
3890
+ # 4. - tools (template can't render tools at all)
3891
+ #
3892
+ # Each rung catches Exception, not (TypeError, ValueError): a Jinja
3893
+ # TemplateError is neither, and the fallbacks are strictly more
3894
+ # conservative than the rung above, so degrading beats a dead turn.
3895
+ depth_keys = [k for k in template_vars if k in _REASONING_DEPTH_KWARGS]
3896
+ try:
3897
+ return tokenizer.apply_chat_template(raw, **template_vars, **base_kwargs)
3898
+ except Exception as exc_vars:
3899
+ if depth_keys:
3900
+ without_depth = {
3901
+ k: v for k, v in template_vars.items() if k not in _REASONING_DEPTH_KWARGS
3902
+ }
3903
+ try:
3904
+ result = tokenizer.apply_chat_template(
3905
+ raw, **without_depth, **base_kwargs
3906
+ )
3907
+ print(
3908
+ f"[tool-prompt] template rejected reasoning depth "
3909
+ f"{ {k: template_vars[k] for k in depth_keys} }; retried "
3910
+ f"with the switch only: {exc_vars}",
3911
+ flush=True,
3912
+ )
3913
+ return result
3914
+ except Exception:
3915
+ pass
3916
+ try:
3917
+ result = tokenizer.apply_chat_template(raw, **base_kwargs)
3918
+ # Loud, because falling back here silently restores the
3919
+ # template's own reasoning defaults — which is exactly the
3920
+ # divergence this parameter exists to close.
3921
+ print(
3922
+ f"[tool-prompt] template rejected {sorted(template_vars)}; "
3923
+ f"retried without them (reasoning falls back to the "
3924
+ f"template default): {exc_vars}",
3925
+ flush=True,
3926
+ )
3927
+ return result
3928
+ except Exception as exc_tools:
3929
+ if tools:
3930
+ print(
3931
+ f"[tool-prompt] TOOLS DROPPED — apply_chat_template could "
3932
+ f"not render {len(tools)} tool(s); model will not see the "
3933
+ f"canonical tool template and may mangle calls: {exc_tools}",
3934
+ flush=True,
3935
+ )
3936
+
3937
+ # No tools, override or template controls, no tokenizer template, or tools-render
3938
+ # failed: text-only path. apply_chat_template signature varies across
3939
+ # mlx-vlm versions; the (processor, config, messages) form is stable.
3940
+ # The mlx_lm tower has no vlm processor/config pair — its tokenizer
3941
+ # renders directly (same chat_template.jinja on disk either way).
3942
+ if _TEXT_TOWER == "mlx_lm":
3943
+ tok = getattr(PROCESSOR, "tokenizer", None) or PROCESSOR
3944
+ return tok.apply_chat_template(raw, tokenize=False, add_generation_prompt=True)
3945
+ config = getattr(MODEL, "config", None) or {}
3946
+ return apply_chat_template(PROCESSOR, config, raw)
3947
+
3948
+
3949
+ def _serial_prompt_token_variants(prompt_text: str) -> "List[List[int]]":
3950
+ """Best-effort re-encodings of the prompt for the serial-path guard.
3951
+
3952
+ mlx_vlm tokenizes the prompt itself (processor pipeline, possibly
3953
+ BOS-prefixed), so our own `encode` can differ by exactly a leading
3954
+ BOS. The guard takes the max LCP across both variants, absorbing
3955
+ that edge without replicating the processor. Empty list = cannot
3956
+ encode = the guard fails safe (reset, cold prefill).
3957
+ """
3958
+ tok = getattr(PROCESSOR, "tokenizer", None) or PROCESSOR
3959
+ try:
3960
+ base = [int(t) for t in tok.encode(prompt_text)]
3961
+ except Exception: # noqa: BLE001 — no encode, no verification
3962
+ return []
3963
+ variants = [base]
3964
+ bos = getattr(tok, "bos_token_id", None)
3965
+ if isinstance(bos, int) and not isinstance(bos, bool):
3966
+ if base[:1] == [bos]:
3967
+ variants.append(base[1:])
3968
+ else:
3969
+ variants.append([bos] + base)
3970
+ return variants
3971
+
3972
+
3973
+ async def _vision_stream(request, http_request, images):
3974
+ """One full-tower vision request. Never share text-prefix KV or speculate."""
3975
+ request_id = f"chatcmpl-{uuid.uuid4()}"
3976
+ generator = None
3977
+ try:
3978
+ prompt = _build_prompt(request.messages, request.tools,
3979
+ request.chat_template_override, request.chat_template_kwargs)
3980
+ kwargs = {"prefill_step_size": ARGS.prefill_step_size, **_kv_quant_kwargs()}
3981
+ for key in ("max_tokens", "temperature", "top_p", "top_k",
3982
+ "repetition_penalty", "repetition_context_size"):
3983
+ value = getattr(request, key, None)
3984
+ if value is not None:
3985
+ kwargs[key] = value
3986
+ tokenizer = getattr(PROCESSOR, "tokenizer", None) or PROCESSOR
3987
+ processors = []
3988
+ if request.tool_grammar:
3989
+ grammar = tool_grammar.build_tool_grammar_processor(
3990
+ tokenizer, request.tools, request.tool_grammar)
3991
+ if grammar is not None:
3992
+ processors.append(grammar)
3993
+ if request.max_thinking_tokens:
3994
+ processor = think_budget.build_think_budget_processor(
3995
+ tokenizer, request.max_thinking_tokens,
3996
+ opens_in_think=prompt.rstrip().endswith("<think>"))
3997
+ if processor is not None:
3998
+ processors.append(processor)
3999
+ if processors:
4000
+ kwargs["logits_processors"] = processors
4001
+ last = None
4002
+ async with _get_generation_lock():
4003
+ if await http_request.is_disconnected():
4004
+ return
4005
+ print(f"[vision] start request={request_id} images={len(images)} cache=disabled", flush=True)
4006
+ with vision_inputs.isolated_position_state(MODEL):
4007
+ generator = stream_generate(model=MODEL, processor=PROCESSOR,
4008
+ prompt=prompt, image=images, **kwargs)
4009
+ try:
4010
+ for chunk in generator:
4011
+ if await http_request.is_disconnected():
4012
+ return
4013
+ if chunk is None or not hasattr(chunk, "text"):
4014
+ continue
4015
+ last = chunk
4016
+ text = _scrub_leaked_markers(chunk.text)
4017
+ if text:
4018
+ yield "data: " + json.dumps({
4019
+ "id": request_id, "object": "chat.completion.chunk",
4020
+ "model": request.model,
4021
+ "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}],
4022
+ }) + "\n\n"
4023
+ await asyncio.sleep(0)
4024
+ finally:
4025
+ generator.close()
4026
+ generator = None
4027
+ _reclaim_mlx_buffer_cache("vision-end", force=True, log=True)
4028
+ usage = {"input_tokens": getattr(last, "prompt_tokens", 0),
4029
+ "output_tokens": getattr(last, "generation_tokens", 0),
4030
+ "cached_tokens": 0,
4031
+ "prompt_tps": getattr(last, "prompt_tps", 0),
4032
+ "generation_tps": getattr(last, "generation_tps", 0)}
4033
+ finish = "length" if request.max_tokens and usage["output_tokens"] >= request.max_tokens else "stop"
4034
+ yield "data: " + json.dumps({
4035
+ "id": request_id, "object": "chat.completion.chunk", "model": request.model,
4036
+ "choices": [{"index": 0, "delta": {}, "finish_reason": finish}], "usage": usage,
4037
+ }) + "\n\n"
4038
+ yield "data: [DONE]\n\n"
4039
+ except Exception as exc:
4040
+ log_contained_exception("vision")
4041
+ yield "data: " + json.dumps({"error": str(exc)}) + "\n\n"
4042
+ finally:
4043
+ if generator is not None:
4044
+ generator.close()
4045
+ for image in images:
4046
+ image.close()
4047
+
4048
+
4049
+ @app.post("/v1/chat/completions")
4050
+ @app.post("/chat/completions")
4051
+ async def chat_completions(request: ChatRequest, http_request: Request):
4052
+ if request.stream is False:
4053
+ # Non-streaming path is rare from gezel's side; reject explicitly
4054
+ # rather than silently behaving differently.
4055
+ raise HTTPException(
4056
+ status_code=501,
4057
+ detail="gezel-mlx-server only supports streaming completions",
4058
+ )
4059
+
4060
+ if any(message.images for message in request.messages):
4061
+ if not _VISION_ENABLED:
4062
+ raise HTTPException(status_code=422, detail="This MLX engine has no enabled vision tower")
4063
+ try:
4064
+ images = vision_inputs.decode_images(request.messages)
4065
+ except Exception as exc:
4066
+ raise HTTPException(status_code=422, detail=f"Invalid image input: {exc}") from exc
4067
+ return StreamingResponse(_vision_stream(request, http_request, images),
4068
+ media_type="text/event-stream",
4069
+ headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
4070
+
4071
+ prompt_text = _build_prompt(
4072
+ request.messages,
4073
+ request.tools,
4074
+ request.chat_template_override,
4075
+ request.chat_template_kwargs,
4076
+ )
4077
+
4078
+ # Build the tool-call grammar logits processor once per request (the
4079
+ # llguidance tokenizer is cached across requests). Any failure returns
4080
+ # None — the turn then runs unconstrained + TS salvage, so this never
4081
+ # breaks a turn.
4082
+ tool_grammar_processor = None
4083
+ if request.tool_grammar:
4084
+ _grammar_tok = getattr(PROCESSOR, "tokenizer", None) or PROCESSOR
4085
+ tool_grammar_processor = tool_grammar.build_tool_grammar_processor(
4086
+ _grammar_tok, request.tools, request.tool_grammar
4087
+ )
4088
+ think_budget_processor = None
4089
+ if request.max_thinking_tokens:
4090
+ _tb_tok = getattr(PROCESSOR, "tokenizer", None) or PROCESSOR
4091
+ think_budget_processor = think_budget.build_think_budget_processor(
4092
+ _tb_tok,
4093
+ request.max_thinking_tokens,
4094
+ # Thinking templates open the assistant turn already inside the
4095
+ # think block (`...<think>\n` generation prompt) — the open tag
4096
+ # never appears in the generated stream, so seed the state from
4097
+ # the prompt tail.
4098
+ opens_in_think=prompt_text.rstrip().endswith("<think>"),
4099
+ )
4100
+
4101
+ # Look up or initialize cache state. Cascade:
4102
+ # 1. In-memory hit on cache_id — best case, no I/O.
4103
+ # 2. Disk hit on cache_id — supervisor restart or
4104
+ # evicted entry, mmap load.
4105
+ # 3. Disk hit on prefix_cache_id — new session inheriting a
4106
+ # warmed gezel prefix.
4107
+ # 4. Fresh PromptCacheState — full prefill required.
4108
+ cache_state, cache_source, matched_prefix = _resolve_cache_state(request)
4109
+
4110
+ if cache_source == "memory":
4111
+ existing = len(cache_state.token_ids or [])
4112
+ print(
4113
+ f"[cache] hit cache_id={request.cache_id} prior_tokens={existing}",
4114
+ flush=True,
4115
+ )
4116
+ elif cache_source == "disk":
4117
+ existing = len(cache_state.token_ids or [])
4118
+ print(
4119
+ f"[cache] disk-hit cache_id={request.cache_id} prior_tokens={existing}",
4120
+ flush=True,
4121
+ )
4122
+ elif cache_source == "prefix":
4123
+ existing = len(cache_state.token_ids or [])
4124
+ print(
4125
+ f"[cache] prefix-hit cache_id={request.cache_id} "
4126
+ f"prefix={matched_prefix} prior_tokens={existing}",
4127
+ flush=True,
4128
+ )
4129
+ elif request.cache_id:
4130
+ print(f"[cache] miss cache_id={request.cache_id}", flush=True)
4131
+
4132
+ # Pin this turn's session + seed-prefix entries from eviction for the life
4133
+ # of the request, so a concurrent save/evict can't drop the KV state this
4134
+ # turn is about to reuse out from under it (released in each stream
4135
+ # iterator's finally).
4136
+ pinned_cache_ids = _mark_inflight(request.cache_id, matched_prefix)
4137
+ prior_tokens = len(cache_state.token_ids or [])
4138
+
4139
+ request_id = f"chatcmpl-{uuid.uuid4()}"
4140
+
4141
+ # BatchGenerator path: use it even for a singleton queue. Besides
4142
+ # concurrency, this path owns the prompt-boundary snapshot needed to reuse
4143
+ # wrapped/windowed KV caches safely after a tool result makes the next
4144
+ # prompt diverge from the raw generated-token tail. The serial fallback is
4145
+ # retained only for an explicit --max-concurrency 0.
4146
+ if int(getattr(ARGS, "max_concurrency", 1) or 0) >= 1:
4147
+ _engine = _get_batch_engine()
4148
+ _tok = getattr(PROCESSOR, "tokenizer", None) or PROCESSOR
4149
+ prompt_tokens = [int(t) for t in _tok.encode(prompt_text)]
4150
+ sub = _Sub(
4151
+ request=request,
4152
+ prompt_tokens=prompt_tokens,
4153
+ grammar=[
4154
+ p for p in (tool_grammar_processor, think_budget_processor) if p is not None
4155
+ ]
4156
+ or None,
4157
+ seed_state=cache_state,
4158
+ request_id=request_id,
4159
+ http_request=http_request,
4160
+ # Save at the last reasoning-mode-invariant position. Without
4161
+ # this the saved entry carries this turn's reasoning preamble,
4162
+ # and the next turn in the other mode diverges from it by ~27
4163
+ # tokens — which an untrimmable hybrid cache cannot rewind, so
4164
+ # the whole multi-thousand-token entry is discarded.
4165
+ snapshot_target=_reasoning_stable_boundary(request, prompt_tokens),
4166
+ )
4167
+ sub.pinned_cache_ids = pinned_cache_ids
4168
+ # Pessimistic placeholder until admission: the authoritative
4169
+ # prefill count comes from the seed plan in _seed_args (the old
4170
+ # entry-size arithmetic here under-reported by ~40× whenever the
4171
+ # seed was rejected, which suppressed the liveness markers during
4172
+ # the exact full re-prefills they exist to surface).
4173
+ sub.prefill_total = len(prompt_tokens)
4174
+ try:
4175
+ _engine.submit(sub)
4176
+ except Exception:
4177
+ _unmark_inflight(pinned_cache_ids)
4178
+ raise
4179
+ return StreamingResponse(
4180
+ _batched_stream_iter(sub),
4181
+ media_type="text/event-stream",
4182
+ headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
4183
+ )
4184
+
4185
+ # Defense in depth: the tower selector requires --max-concurrency >= 1
4186
+ # for the mlx_lm text tower, so this serial branch (mlx_vlm
4187
+ # stream_generate) should be unreachable in lm mode. If a future edit
4188
+ # breaks that invariant, fail loudly instead of feeding an mlx_lm
4189
+ # model into mlx_vlm's generate loop.
4190
+ if _TEXT_TOWER == "mlx_lm":
4191
+ raise HTTPException(
4192
+ status_code=501,
4193
+ detail=(
4194
+ "serial generation path is unavailable with the mlx_lm text "
4195
+ "tower; restart without GEZEL_MLX_TEXT_TOWER or with "
4196
+ "--max-concurrency >= 1"
4197
+ ),
4198
+ )
4199
+
4200
+ # Serial-path guard: mlx_vlm's divergent-prefix branch trims by
4201
+ # slicing `keys[:, :, :lcp]`, which silently corrupts wrapped
4202
+ # sliding-window layers (the rotated cache is SHORTER than the
4203
+ # prefix, gets skipped, and its offset disagrees with the trimmed
4204
+ # full-attention layers from then on). When the state can't be
4205
+ # trimmed coherently and the prompt genuinely diverges, a cold
4206
+ # prefill is the only correct outcome — reset before handing the
4207
+ # state to stream_generate. Pure extensions pass through untouched.
4208
+ if cache_state.cache is not None and (cache_state.token_ids or []):
4209
+ _variants = _serial_prompt_token_variants(prompt_text)
4210
+ if cache_seed.serial_reset_needed(cache_state.token_ids, cache_state.cache, _variants):
4211
+ print(
4212
+ f"[cache] serial divergence on untrimmable cache; cold prefill "
4213
+ f"cache_id={request.cache_id}",
4214
+ flush=True,
4215
+ )
4216
+ cache_state = PromptCacheState()
4217
+
4218
+ # IMPORTANT: this MUST be `async def`, not `def`. Starlette iterates
4219
+ # sync generators via `iterate_in_threadpool` which can hop threads
4220
+ # between yields — but MLX's GPU streams are thread-local, so any
4221
+ # `mx.synchronize` call from a different thread than the one that
4222
+ # created the stream raises `RuntimeError: There is no Stream(gpu, 0)
4223
+ # in current thread.` async generators get iterated directly in the
4224
+ # event-loop thread, which keeps the stream + generation co-located.
4225
+ # (The body remains synchronous beyond the `yield`s — `for chunk in
4226
+ # stream_generate(...)` blocks the event loop, same as upstream
4227
+ # mlx_vlm.server's pattern.)
4228
+ async def stream_iter():
4229
+ # Tracks whether the gezel client TCP-disconnected mid-stream
4230
+ # (turn timeout fired, user cancelled, etc.). When set, we
4231
+ # `break` out of `stream_generate` so MLX stops generating
4232
+ # tokens nobody will read — without this, uvicorn's underlying
4233
+ # ASGI write keeps failing per token and floods the supervisor
4234
+ # log with `socket.send() raised exception.` for the whole
4235
+ # remainder of the would-be response.
4236
+ client_gone = False
4237
+ try:
4238
+ generation_kwargs: Dict[str, Any] = {
4239
+ "prompt_cache_state": cache_state,
4240
+ "prefill_step_size": ARGS.prefill_step_size,
4241
+ **_kv_quant_kwargs(),
4242
+ }
4243
+ if request.max_tokens is not None:
4244
+ generation_kwargs["max_tokens"] = request.max_tokens
4245
+ if request.temperature is not None:
4246
+ generation_kwargs["temperature"] = request.temperature
4247
+ if request.top_p is not None:
4248
+ generation_kwargs["top_p"] = request.top_p
4249
+ if request.top_k is not None:
4250
+ generation_kwargs["top_k"] = request.top_k
4251
+ if request.repetition_penalty is not None:
4252
+ generation_kwargs["repetition_penalty"] = request.repetition_penalty
4253
+ if request.repetition_context_size is not None:
4254
+ generation_kwargs["repetition_context_size"] = (
4255
+ request.repetition_context_size
4256
+ )
4257
+ _lps = [
4258
+ p for p in (tool_grammar_processor, think_budget_processor) if p is not None
4259
+ ]
4260
+ if _lps:
4261
+ # Forwarded through stream_generate(**kwargs) → generate_step,
4262
+ # which accepts `logits_processors`. Grammar constrains the
4263
+ # tool-call function name; think-budget forces `</think>` at
4264
+ # the configured reasoning budget (llama parity).
4265
+ generation_kwargs["logits_processors"] = _lps
4266
+
4267
+ output_text = ""
4268
+
4269
+ tool_stream = _make_tool_stream(request_id, getattr(request, 'tools', None))
4270
+ last_chunk = None
4271
+ # Serialize generation across concurrent requests — see
4272
+ # `_GENERATION_LOCK` doc. Held only across the
4273
+ # `stream_generate` loop, NOT across cache lookup/save, so
4274
+ # request N+1 can prep its prefill state in parallel with
4275
+ # request N's streaming. The lock is fair (asyncio.Lock is
4276
+ # FIFO since Python 3.10) so a long generation can't starve
4277
+ # a queued short one indefinitely.
4278
+ async with _get_generation_lock():
4279
+ for chunk in stream_generate(
4280
+ model=MODEL,
4281
+ processor=PROCESSOR,
4282
+ prompt=prompt_text,
4283
+ **generation_kwargs,
4284
+ ):
4285
+ if chunk is None or not hasattr(chunk, "text"):
4286
+ continue
4287
+ # Cooperative disconnect check. `is_disconnected()` is
4288
+ # cheap (just peeks at the ASGI receive queue for an
4289
+ # `http.disconnect` message) so doing it per token has
4290
+ # no measurable overhead next to MLX's GPU work.
4291
+ if await http_request.is_disconnected():
4292
+ print(
4293
+ f"[stream] client disconnected mid-generation "
4294
+ f"after {getattr(chunk, 'generation_tokens', 0)} tokens; "
4295
+ f"aborting",
4296
+ flush=True,
4297
+ )
4298
+ client_gone = True
4299
+ break
4300
+ last_chunk = chunk
4301
+ # Scrub leaked framing tokens per chunk. Special tokens
4302
+ # finalize as their own detokenizer segment, so a marker
4303
+ # never straddles two chunks — a per-chunk scrub is safe.
4304
+ chunk_text = _scrub_leaked_markers(chunk.text)
4305
+ if not chunk_text:
4306
+ continue
4307
+ output_text += chunk_text
4308
+
4309
+ # With --stream-tool-calls the raw text is split into
4310
+ # content and tool-call argument deltas, so the client
4311
+ # sees a real call boundary instead of markup embedded in
4312
+ # prose. Without it this is a no-op passthrough and the
4313
+ # gezel-side salvage path behaves exactly as before.
4314
+ for _frag in _tool_stream_fragments(tool_stream, chunk_text):
4315
+ yield "data: " + json.dumps(
4316
+ _tool_stream_payload(_frag, request_id, request.model, chunk)
4317
+ ) + "\n\n"
4318
+ if tool_stream is not None:
4319
+ continue
4320
+
4321
+ payload = {
4322
+ "id": request_id,
4323
+ "object": "chat.completion.chunk",
4324
+ "created": int(time.time()),
4325
+ "model": request.model,
4326
+ "choices": [
4327
+ {
4328
+ "index": 0,
4329
+ "delta": {"role": "assistant", "content": chunk_text},
4330
+ "finish_reason": None,
4331
+ }
4332
+ ],
4333
+ "usage": {
4334
+ "input_tokens": getattr(chunk, "prompt_tokens", 0),
4335
+ "output_tokens": getattr(chunk, "generation_tokens", 0),
4336
+ "total_tokens": (
4337
+ getattr(chunk, "prompt_tokens", 0)
4338
+ + getattr(chunk, "generation_tokens", 0)
4339
+ ),
4340
+ "prompt_tps": getattr(chunk, "prompt_tps", 0.0),
4341
+ "generation_tps": getattr(chunk, "generation_tps", 0.0),
4342
+ # Actual prefix reused by mlx-vlm after token
4343
+ # comparison, not merely the size of the cache
4344
+ # entry we found by id. This distinguishes a real
4345
+ # hot continuation from a token-diverged cache
4346
+ # that still has to re-prefill most of the prompt.
4347
+ "cached_tokens": getattr(chunk, "cached_tokens", 0),
4348
+ },
4349
+ }
4350
+ yield f"data: {json.dumps(payload)}\n\n"
4351
+
4352
+ if tool_stream is not None and last_chunk is not None:
4353
+
4354
+ for _frag in tool_stream.flush():
4355
+
4356
+ yield "data: " + json.dumps(
4357
+
4358
+ _tool_stream_payload(_frag, request_id, request.model, last_chunk)
4359
+
4360
+ ) + "\n\n"
4361
+
4362
+ if client_gone:
4363
+ # Skip terminal + [DONE] frames — there's nobody to
4364
+ # receive them, and yielding through a closed transport
4365
+ # is what we're trying to avoid. Cache save in `finally`
4366
+ # still runs so the prefill we paid for isn't wasted.
4367
+ return
4368
+
4369
+ # Terminal chunk with finish_reason. Mirrors the upstream
4370
+ # server's shape; gezel's MlxSession reads finish_reason
4371
+ # off the last chunk to decide truncation handling.
4372
+ terminal = {
4373
+ "id": request_id,
4374
+ "object": "chat.completion.chunk",
4375
+ "created": int(time.time()),
4376
+ "model": request.model,
4377
+ "choices": [
4378
+ {
4379
+ "index": 0,
4380
+ "delta": {"role": "assistant", "content": ""},
4381
+ "finish_reason": "stop",
4382
+ }
4383
+ ],
4384
+ }
4385
+ if last_chunk is not None:
4386
+ prompt_tokens = getattr(last_chunk, "prompt_tokens", 0)
4387
+ cached_tokens = getattr(last_chunk, "cached_tokens", 0)
4388
+ terminal["usage"] = {
4389
+ "input_tokens": prompt_tokens,
4390
+ "output_tokens": getattr(last_chunk, "generation_tokens", 0),
4391
+ "total_tokens": (
4392
+ prompt_tokens
4393
+ + getattr(last_chunk, "generation_tokens", 0)
4394
+ ),
4395
+ "prompt_tps": getattr(last_chunk, "prompt_tps", 0.0),
4396
+ "generation_tps": getattr(last_chunk, "generation_tps", 0.0),
4397
+ "cached_tokens": cached_tokens,
4398
+ }
4399
+ if request.cache_id:
4400
+ print(
4401
+ f"[cache] reuse cache_id={request.cache_id} "
4402
+ f"cached_tokens={cached_tokens} "
4403
+ f"prefilled_tokens={max(0, prompt_tokens - cached_tokens)}",
4404
+ flush=True,
4405
+ )
4406
+ yield f"data: {json.dumps(terminal)}\n\n"
4407
+ yield "data: [DONE]\n\n"
4408
+
4409
+ except Exception as e:
4410
+ print(f"Error during stream generation: {e}", flush=True)
4411
+ log_contained_exception("stream")
4412
+ yield f"data: {json.dumps({'error': str(e)})}\n\n"
4413
+
4414
+ finally:
4415
+ # The cache `finally` decision:
4416
+ # - cache_id set + state populated → save (the happy path)
4417
+ # - cache_id set + state empty (e.g. mid-stream crash before
4418
+ # stream_generate's update() ran) → DON'T save. Saving an
4419
+ # empty entry would overwrite any prior good cache and
4420
+ # guarantee the next turn misses + re-prefills full
4421
+ # - no cache_id → upstream behavior (drop + clear memory)
4422
+ if request.cache_id:
4423
+ tokens = len(cache_state.token_ids or [])
4424
+ if tokens > 0:
4425
+ _save_cache(request.cache_id, cache_state)
4426
+ print(
4427
+ f"[cache] saved cache_id={request.cache_id} tokens={tokens}",
4428
+ flush=True,
4429
+ )
4430
+ # Seed the shared gezel+project prefix from this real
4431
+ # session save so sibling sessions inherit [system+tools].
4432
+ _seed_prefix_from_session(request, cache_state)
4433
+ else:
4434
+ print(
4435
+ f"[cache] generation produced no tokens for "
4436
+ f"cache_id={request.cache_id}; leaving prior entry intact",
4437
+ flush=True,
4438
+ )
4439
+ else:
4440
+ print("Stream finished without a prompt-cache save.", flush=True)
4441
+ # Release the eviction pins taken for this turn — AFTER the save
4442
+ # above, so this turn's own _save_cache → _enforce_budget can't
4443
+ # evict the entry it just wrote.
4444
+ _unmark_inflight(pinned_cache_ids)
4445
+ # Agentic tool loops issue multiple HTTP sub-turns inside one user
4446
+ # turn. Reclaim scratch buffers at every one of those boundaries,
4447
+ # including the cache-preserving path above; prompt KV remains an
4448
+ # active allocation and is not touched by clear_cache().
4449
+ _reclaim_mlx_buffer_cache("turn-end", force=True, log=True)
4450
+
4451
+ return StreamingResponse(
4452
+ stream_iter(),
4453
+ media_type="text/event-stream",
4454
+ headers={
4455
+ "Cache-Control": "no-cache",
4456
+ # NOTE: mlx_vlm.server sets `Connection: keep-alive` here,
4457
+ # which our gezel daemon would have to strip when running
4458
+ # over HTTP/2. We're a separate process not on HTTP/2 so
4459
+ # the header is fine to send, but explicitly omit to keep
4460
+ # behavior predictable across transports.
4461
+ "X-Accel-Buffering": "no",
4462
+ },
4463
+ )
4464
+
4465
+
4466
+ # ───────── Entry point ─────────
4467
+
4468
+ if __name__ == "__main__":
4469
+ import logging
4470
+
4471
+ import uvicorn
4472
+
4473
+ class _DropPollingAccessLogs(logging.Filter):
4474
+ # Supervisor polls /health and the gezel cache adapter polls
4475
+ # /v1/cache/stats every few seconds — at INFO that floods the
4476
+ # parent's [mlx]-prefixed logs. Drop only the 200s for those two
4477
+ # endpoints; anything else (4xx/5xx, other routes) still logs.
4478
+ _NOISY = ("/health", "/v1/cache/stats")
4479
+
4480
+ def filter(self, record: logging.LogRecord) -> bool:
4481
+ args = record.args
4482
+ if not isinstance(args, tuple) or len(args) < 5:
4483
+ return True
4484
+ method, path, _http, status = args[1], args[2], args[3], args[4]
4485
+ return not (
4486
+ method == "GET" and status == 200 and path in self._NOISY
4487
+ )
4488
+
4489
+ logging.getLogger("uvicorn.access").addFilter(_DropPollingAccessLogs())
4490
+
4491
+ uvicorn.run(app, host=ARGS.host, port=ARGS.port, log_level="info")