mlx-bun 0.0.10 → 0.0.11

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 (173) hide show
  1. package/README.md +66 -21
  2. package/THIRD_PARTY_LICENSES.md +1 -0
  3. package/docs/archive/mlx-bun-lab-report.html +275 -0
  4. package/docs/archive/wwdc-mlx-bun.html +848 -0
  5. package/docs/design/audio-input-plan.md +356 -0
  6. package/docs/design/batching-perf-path.md +58 -35
  7. package/docs/design/batching-v2-plan.md +23 -9
  8. package/docs/design/bucketing-stage.md +2 -1
  9. package/docs/design/compat-cli-surface-design.md +4 -3
  10. package/docs/design/decode-speed-program.md +172 -0
  11. package/docs/design/docs-reorg-plan.md +1 -1
  12. package/docs/design/dspark-serving-program.md +371 -0
  13. package/docs/design/dspark-speculative-decoding.md +167 -5
  14. package/docs/design/faithful-l1-consolidation.md +310 -0
  15. package/docs/design/generic-model-support.md +8 -5
  16. package/docs/design/grammar-spec-batching-integration.md +301 -0
  17. package/docs/design/minicpm5-decode-megakernel.md +16 -11
  18. package/docs/design/mlx-lm-tool-parity-plan.md +11 -1
  19. package/docs/design/omlx-adoption-map.md +14 -6
  20. package/docs/design/orpo-training.md +34 -30
  21. package/docs/design/paged-kv-cache.md +214 -0
  22. package/docs/design/parallel-slots.md +82 -49
  23. package/docs/design/parity-tier-dag.md +43 -33
  24. package/docs/design/repo-cleanup-plan.md +189 -0
  25. package/docs/design/runtime-isolation.md +19 -3
  26. package/docs/design/segmented-backward-training.md +3 -2
  27. package/docs/design/ssd-kv-cold-tier.md +130 -0
  28. package/docs/design/structured-output.md +484 -0
  29. package/docs/design/the-dreaming-master-plan.md +2 -3
  30. package/docs/design/turboquant-kv.md +235 -0
  31. package/docs/design/unified-engine-frontier-plan.md +774 -0
  32. package/docs/design/web-chat-beat-matrix.md +315 -0
  33. package/docs/design/web-chat-redesign.md +1112 -0
  34. package/docs/investigations/dspark-handoff.md +182 -16
  35. package/docs/investigations/lab-build-journal.md +2 -2
  36. package/docs/investigations/openwiki-evaluation.md +249 -0
  37. package/docs/planning/IDEAS.md +14 -0
  38. package/docs/planning/PRODUCT_ROADMAP.md +2 -3
  39. package/docs/planning/chat-ui-vision-memory-first.md +199 -0
  40. package/docs/planning/chat-ui-vision.md +50 -0
  41. package/docs/planning/memory-docs-and-dag-plan.md +15 -6
  42. package/docs/planning/release-notes-v0.0.11.md +239 -0
  43. package/docs/reference/cli.md +35 -16
  44. package/docs/reference/distribution.md +13 -6
  45. package/docs/reference/embedding.md +35 -14
  46. package/docs/reference/features-matrix.md +97 -0
  47. package/docs/reference/library-api.md +30 -9
  48. package/docs/reference/memory.md +26 -6
  49. package/docs/reference/models.md +11 -1
  50. package/docs/reference/orpo-quickstart.md +3 -4
  51. package/docs/reference/server-api.md +634 -22
  52. package/docs/reference/server-config.md +250 -121
  53. package/docs/reference/training.md +14 -15
  54. package/package.json +5 -3
  55. package/src/anthropic.ts +13 -0
  56. package/src/audio/conformer.ts +647 -0
  57. package/src/audio/decode.ts +133 -0
  58. package/src/audio/features.ts +202 -0
  59. package/src/audio/transcode.ts +67 -0
  60. package/src/cli.ts +356 -105
  61. package/src/config.ts +49 -2
  62. package/src/download.ts +40 -0
  63. package/src/eval/README.md +7 -6
  64. package/src/eval/kl.ts +59 -3
  65. package/src/eval/quality-db.ts +2 -2
  66. package/src/eval/runner.ts +3 -3
  67. package/src/expert-offload.ts +15 -2
  68. package/src/flags.ts +13 -0
  69. package/src/generate.ts +330 -24
  70. package/src/grammar.ts +585 -0
  71. package/src/hub-rest.ts +242 -0
  72. package/src/kv-store.ts +442 -88
  73. package/src/lora.ts +8 -0
  74. package/src/md-modules.d.ts +42 -0
  75. package/src/media-fetch.ts +284 -0
  76. package/src/memory/rest.ts +245 -0
  77. package/src/memory/vault.ts +74 -0
  78. package/src/mlx/array.ts +106 -19
  79. package/src/mlx/autograd.ts +2 -2
  80. package/src/mlx/ffi.ts +27 -3
  81. package/src/mlx/metal-kernel.ts +7 -4
  82. package/src/mlx/ops.ts +153 -23
  83. package/src/mlx/turboquant-ops.ts +831 -0
  84. package/src/mlx/turboquant-tables.ts +237 -0
  85. package/src/model/batched-mask.ts +62 -1
  86. package/src/model/batched-quant.ts +206 -0
  87. package/src/model/batched-rotating-quant.ts +346 -0
  88. package/src/model/batched-rotating.ts +46 -1
  89. package/src/model/compiled-decode.ts +4 -15
  90. package/src/model/factory.ts +10 -4
  91. package/src/model/gemma4-base.ts +363 -39
  92. package/src/model/gemma4.ts +122 -33
  93. package/src/model/generated/gemma4-12b.ts +9 -13
  94. package/src/model/generated/gemma4-26b.ts +10 -5
  95. package/src/model/generated/gemma4-e4b.ts +14 -9
  96. package/src/model/minicpm5.ts +58 -102
  97. package/src/model/paged-kv.ts +253 -0
  98. package/src/model/qwen3-delta.ts +76 -15
  99. package/src/model/qwen3-moe.ts +404 -0
  100. package/src/model/qwen3.ts +9 -8
  101. package/src/model/qwen3_5.ts +88 -23
  102. package/src/model/support.ts +12 -1
  103. package/src/model/universal/dense.ts +40 -5
  104. package/src/model/universal/rope.ts +36 -21
  105. package/src/pi-web.ts +1299 -82
  106. package/src/preflight.ts +27 -8
  107. package/src/prompt-cache.ts +293 -33
  108. package/src/quantize/quantizer.ts +16 -1
  109. package/src/registry.ts +57 -2
  110. package/src/serve/batch-scheduler.ts +965 -72
  111. package/src/serve/generation-gateway.ts +214 -11
  112. package/src/serve/isolate.ts +417 -0
  113. package/src/serve/lane-registry.ts +44 -0
  114. package/src/serve/session-search.ts +326 -0
  115. package/src/server.ts +1380 -135
  116. package/src/spec/assistant-source.ts +150 -0
  117. package/src/spec/deepspec-source.ts +143 -0
  118. package/src/spec/dflash-source.ts +124 -0
  119. package/src/spec/dspark/ab-stats.ts +136 -0
  120. package/src/spec/dspark/calibration.ts +112 -0
  121. package/src/spec/dspark/deepspec-module.ts +801 -0
  122. package/src/spec/dspark/generate-dflash.ts +72 -20
  123. package/src/spec/dspark/loader.ts +23 -0
  124. package/src/spec/dspark/module-dflash.ts +278 -36
  125. package/src/spec/dspark/quantize-drafter.ts +91 -0
  126. package/src/spec/generate.ts +59 -25
  127. package/src/spec/ngram-source.ts +132 -0
  128. package/src/spec/serve-loop.ts +536 -0
  129. package/src/spec/source.ts +93 -0
  130. package/src/spec/two-model.ts +171 -0
  131. package/src/ssd-cache.ts +107 -35
  132. package/src/tokenizer.ts +288 -2
  133. package/src/tool-approvals.ts +131 -0
  134. package/src/tool-call.ts +408 -20
  135. package/src/train/flash-cce.ts +2 -2
  136. package/src/train/job.ts +0 -2
  137. package/src/train/trainer.ts +1 -4
  138. package/src/vision/prompt.ts +217 -48
  139. package/src/web/app.html +1184 -1789
  140. package/src/web/app.js +5667 -0
  141. package/src/web/icon.svg +13 -0
  142. package/src/web/manifest.webmanifest +24 -0
  143. package/src/web/src/adapters-panel.ts +288 -0
  144. package/src/web/src/api.ts +57 -0
  145. package/src/web/src/assistant.ts +485 -0
  146. package/src/web/src/chat.ts +1050 -0
  147. package/src/web/src/composer.ts +959 -0
  148. package/src/web/src/dataset.ts +132 -0
  149. package/src/web/src/finetune.ts +267 -0
  150. package/src/web/src/hub.ts +428 -0
  151. package/src/web/src/main.ts +98 -0
  152. package/src/web/src/markdown.ts +0 -0
  153. package/src/web/src/memory-panel.ts +572 -0
  154. package/src/web/src/model-picker.ts +159 -0
  155. package/src/web/src/palette.ts +346 -0
  156. package/src/web/src/protocol.ts +62 -0
  157. package/src/web/src/quantize.ts +167 -0
  158. package/src/web/src/rag.ts +274 -0
  159. package/src/web/src/sessions.ts +469 -0
  160. package/src/web/src/shell.ts +804 -0
  161. package/src/web/src/status.ts +252 -0
  162. package/src/web/src/ui-catalog.ts +92 -0
  163. package/src/web/sw.js +71 -0
  164. package/src/web/vendor/README.md +61 -0
  165. package/src/web/vendor/hljs-LICENSE +29 -0
  166. package/src/web/vendor/hljs-theme.css +59 -0
  167. package/src/web/vendor/hljs.js +26 -0
  168. package/src/model/fused-decode-kernel.ts +0 -312
  169. package/src/model/fused-geglu-kernel.ts +0 -198
  170. package/src/model/fused-mlp-kernel.ts +0 -169
  171. package/src/model/fused-swiglu-kernel.ts +0 -240
  172. package/src/model/steel-linear-kernel.ts +0 -96
  173. package/src/status-page.html +0 -497
package/src/pi-web.ts CHANGED
@@ -32,6 +32,7 @@ import {
32
32
  createAgentSessionFromServices,
33
33
  createAgentSessionRuntime,
34
34
  createAgentSessionServices,
35
+ defineTool,
35
36
  SessionManager,
36
37
  type AgentSession,
37
38
  type AgentSessionEvent,
@@ -41,20 +42,30 @@ import {
41
42
  type SessionEntry,
42
43
  type SessionInfo,
43
44
  type ToolCallEvent,
45
+ type ToolDefinition,
44
46
  } from "@earendil-works/pi-coding-agent";
45
47
  import type { ImageContent } from "@earendil-works/pi-ai";
48
+ import { Type } from "typebox";
46
49
  import { WEB_TOOL_NAMES } from "./web-tools";
47
50
  import { MEMORY_TOOL_NAMES, REFERENCE_TOOL_NAMES } from "./memory/tools";
48
51
  import { buildPiAgentSurface } from "./pi-session";
49
52
  import { buildPiProvider, DEFAULT_CONTEXT_WINDOW, PI_LOCAL_MODEL_ID } from "./pi-provider";
50
53
  import { downloadsSnapshot } from "./download";
54
+ import { getLane, type Lane } from "./serve/lane-registry";
55
+ import { isToolAlwaysAllowed, setToolAlwaysAllowed, listAlwaysAllowedTools } from "./tool-approvals";
51
56
 
52
57
  /**
53
58
  * Tools that never mutate the user's machine; auto-allowed without a browser
54
59
  * round-trip. The web tools make outbound network requests but change nothing
55
60
  * locally, so they're auto-allowed too (and remain usable in read-only mode).
61
+ * The app-aware assistant's three tools (get_current_app_context/
62
+ * navigate_app/spotlight_ui, plan §6.6) join this set for the same reason:
63
+ * navigation and spotlighting are reversible UI actions on the user's OWN
64
+ * open browser tab, never a machine mutation — see createAppAwareTools'
65
+ * doc comment below for the full classification rationale.
56
66
  */
57
- const READ_ONLY_TOOLS = new Set(["read", "grep", "find", "ls", ...WEB_TOOL_NAMES, ...MEMORY_TOOL_NAMES, ...REFERENCE_TOOL_NAMES]);
67
+ export const APP_AWARE_TOOL_NAMES = ["get_current_app_context", "navigate_app", "spotlight_ui"] as const;
68
+ const READ_ONLY_TOOLS = new Set(["read", "grep", "find", "ls", ...WEB_TOOL_NAMES, ...MEMORY_TOOL_NAMES, ...REFERENCE_TOOL_NAMES, ...APP_AWARE_TOOL_NAMES]);
58
69
  /** Tools that require explicit per-call browser approval. */
59
70
  const GATED_TOOLS = new Set(["bash", "edit", "write"]);
60
71
 
@@ -69,22 +80,214 @@ const WELCOME_TOOLS = ["read", "web_search"] as const;
69
80
 
70
81
  /**
71
82
  * The web chat session's tool allowlist: the two welcome tools, plus the
72
- * memory/reference read tools whenever the surface has memory enabled. The
73
- * system prompt (surface.memoryHint) and the bundled memory skill instruct
74
- * the model to call memory_resolve/memory_read/… — pi treats `tools` as an
83
+ * memory/reference read tools whenever the surface has memory enabled, plus
84
+ * — opt-in only (plan §5.4/§6.5/§9 Phase 2's "let the agent touch files on
85
+ * this machine" toggle, default OFF) — the three GATED_TOOLS (bash/edit/
86
+ * write) and the read-only grep/find/ls that make them useful. The system
87
+ * prompt (surface.memoryHint) and the bundled memory skill instruct the
88
+ * model to call memory_resolve/memory_read/… — pi treats `tools` as an
75
89
  * allowlist, so leaving them out makes every such call fail. All additions
76
- * are in READ_ONLY_TOOLS (no approval round-trip).
90
+ * except the coding tools are in READ_ONLY_TOOLS (no approval round-trip);
91
+ * bash/edit/write ARE gated — this is precisely what makes the
92
+ * already-built approval card fire for the first time once codingTools is
93
+ * on (previously dead code: nothing in the default surface could ever
94
+ * reach GATED_TOOLS). APP_AWARE_TOOL_NAMES ride along unconditionally
95
+ * (like WELCOME_TOOLS) — they're read-only-on-the-machine regardless of
96
+ * memory/codingTools/read-only-server state (see APP_AWARE_TOOL_NAMES'
97
+ * doc comment).
77
98
  */
78
- export function webChatToolAllowlist(memoryEnabled: boolean): string[] {
99
+ export function webChatToolAllowlist(memoryEnabled: boolean, codingTools = false): string[] {
79
100
  return [
80
101
  ...WELCOME_TOOLS,
102
+ ...APP_AWARE_TOOL_NAMES,
81
103
  ...(memoryEnabled ? [...MEMORY_TOOL_NAMES, ...REFERENCE_TOOL_NAMES] : []),
104
+ ...(codingTools ? ["grep", "find", "ls", ...GATED_TOOLS] : []),
82
105
  ];
83
106
  }
84
107
 
85
108
  /** Auto-deny a pending approval after this long with no browser decision. */
86
109
  const APPROVAL_TIMEOUT_MS = 120_000;
87
110
 
111
+ // ---- Loop hygiene (beat matrix Axis 7 "Tool-loop hygiene" row) --------
112
+ //
113
+ // What the pi SDK already provides (investigated in
114
+ // node_modules/@earendil-works/pi-agent-core/dist/agent-loop.js and
115
+ // pi-coding-agent/dist/core/agent-session.js before adding any of this):
116
+ // - Malformed tool-call ARGS: pi-ai's `parseJsonWithRepair`/
117
+ // `parseStreamingJson` (dist/utils/json-parse.js) already repairs
118
+ // invalid backslash escapes/control chars and tolerates truncated
119
+ // streaming JSON; `validateToolArguments` (utils/validation.js)
120
+ // coerces JSON-Schema primitive types. Neither touches OUR OWN
121
+ // format-specific tool-call text parsers (gemma4 `call:name{...}`,
122
+ // the OpenAI/Qwen/MiniCPM5 decoded-text shapes) — that repair layer
123
+ // is added in src/tool-call.ts instead (see RepairKind there).
124
+ // - Repeated identical calls: NOT deduped anywhere in the SDK. Every
125
+ // tool_call the model emits executes, full stop.
126
+ // - Failed-call retries: NOT budgeted. A tool error becomes a
127
+ // toolResult with isError:true and the loop just continues; nothing
128
+ // nudges the model to stop repeating a failing call.
129
+ // - Max turns / loop cap: NOT implemented by AgentSession at all (only
130
+ // agent-loop.js's raw runAgentLoop is hook-driven via
131
+ // config.shouldStopAfterTurn, which AgentSession never wires up or
132
+ // exposes to extensions). There is no cap, so a small model that
133
+ // keeps calling tools can loop until the user manually stops it.
134
+ // All three gaps are closed here using the ONE mechanism AgentSession
135
+ // does expose to extensions: pi.on("tool_call", ...) (before execution,
136
+ // can only `block`) and pi.on("tool_result", ...) (after execution, can
137
+ // reshape the result) — see installLoopHygieneHooks below. Turn-index
138
+ // bookkeeping rides pi.on("turn_start"/"turn_end") (AgentSession's own
139
+ // _turnIndex, which resets to 0 once per agent_start i.e. once per
140
+ // session.prompt() call — see agent-session.js's _emitExtensionEvent).
141
+ //
142
+ // KNOWN SDK LIMITATION (documented, not silently worked around): a
143
+ // `tool_call` handler's block path (ToolCallEventResult) can only
144
+ // produce an ERROR toolResult (agent-loop.js's createErrorToolResult —
145
+ // there is no "block but still succeed" return shape). So "dedup: skip
146
+ // re-execution, return cached result" cannot literally inject a
147
+ // non-error result without running the tool again — the closest honest
148
+ // equivalent is blocking the repeat with the PRIOR call's actual result
149
+ // text inlined into the (error-flagged) block reason, so the model
150
+ // still sees the real answer instead of a bare "blocked" message and
151
+ // can move on, rather than getting a second identical error with no
152
+ // information gain.
153
+
154
+ /** Consecutive-identical-successful-call dedup, retry budget, and
155
+ * tool-turn cap. Mutable per-connection (reset at prompt time) so re-
156
+ * export as a class keeps all three concerns in one bookkeeping object
157
+ * that's trivial to unit test via its pure decision methods. */
158
+ export const LOOP_HYGIENE = {
159
+ /** After this many consecutive FAILED calls with the same signature
160
+ * (name + JSON-stable args), the nudge text asks the model to change
161
+ * approach instead of repeating verbatim. optiq Lab's documented
162
+ * budget (beat matrix Axis 7) is 3. */
163
+ MAX_CONSECUTIVE_FAILURES: 3,
164
+ /** Cap on AgentSession's turnIndex (one assistant-response round, per
165
+ * agent-session.js's _turnIndex) within a single user prompt before
166
+ * every further tool call is blocked with a force-finish nudge.
167
+ * optiq Lab's documented cap (beat matrix Axis 7) is 25. */
168
+ MAX_TOOL_TURNS: 25,
169
+ } as const;
170
+
171
+ /** Stable key for "the same call" — used by both the dedup check and the
172
+ * failure-streak counter. JSON.stringify on a plain object from parsed
173
+ * tool-call args is key-order-dependent; that's fine here since we're
174
+ * comparing the SDK's own re-delivered `event.input` (same object shape
175
+ * each time for a literal repeat, which is the only case that matters —
176
+ * semantically-equal-but-differently-ordered args are not "identical"
177
+ * for dedup purposes, only byte-identical repeats are). Exported for
178
+ * unit testing. */
179
+ export function toolCallSignature(toolName: string, args: unknown): string {
180
+ return `${toolName}:${stableStringify(args)}`;
181
+ }
182
+
183
+ /** Deterministic (key-sorted) JSON.stringify so argument key order never
184
+ * causes two semantically-identical calls to hash differently. */
185
+ function stableStringify(value: unknown): string {
186
+ if (Array.isArray(value)) return `[${value.map(stableStringify).join(",")}]`;
187
+ if (value && typeof value === "object") {
188
+ const keys = Object.keys(value as Record<string, unknown>).sort();
189
+ return `{${keys.map((k) => `${JSON.stringify(k)}:${stableStringify((value as Record<string, unknown>)[k])}`).join(",")}}`;
190
+ }
191
+ return JSON.stringify(value);
192
+ }
193
+
194
+ /** Pure decision for the tool_call (before) hook: given the current loop
195
+ * state and the incoming call's signature, decide whether to block it
196
+ * outright (dedup or turn-cap) and, if so, why. Returns undefined when
197
+ * the call should proceed to normal execution. Exported + pure so the
198
+ * three policies (dedup / turn cap / [failure budget is a tool_result-
199
+ * side concern, see nextFailureNudge]) are unit-testable without a live
200
+ * AgentSession. */
201
+ export interface LoopHygieneState {
202
+ /** Signature + result text of the last call that completed successfully
203
+ * (undefined if none yet this prompt). */
204
+ lastSuccess?: { signature: string; resultText: string };
205
+ /** Consecutive failure count keyed by call signature. */
206
+ failureStreaks: Map<string, number>;
207
+ /** Current AgentSession turnIndex (updated on every turn_start). */
208
+ turnIndex: number;
209
+ }
210
+
211
+ export function initialLoopHygieneState(): LoopHygieneState {
212
+ return { lastSuccess: undefined, failureStreaks: new Map(), turnIndex: 0 };
213
+ }
214
+
215
+ export interface LoopHygieneBlock {
216
+ reason: string;
217
+ }
218
+
219
+ /** Decide whether to block a tool call BEFORE it runs. Checked in this
220
+ * order: turn cap first (a model that's been looping for 25 rounds gets
221
+ * force-finish treatment regardless of what it's calling next), then
222
+ * consecutive-identical-successful-call dedup. */
223
+ export function decideBeforeToolCall(
224
+ state: LoopHygieneState,
225
+ toolName: string,
226
+ args: unknown,
227
+ ): LoopHygieneBlock | undefined {
228
+ if (state.turnIndex >= LOOP_HYGIENE.MAX_TOOL_TURNS) {
229
+ return {
230
+ reason:
231
+ `Tool-turn limit (${LOOP_HYGIENE.MAX_TOOL_TURNS}) reached for this message. ` +
232
+ "Stop calling tools now and answer the user directly with what you have so far.",
233
+ };
234
+ }
235
+ const signature = toolCallSignature(toolName, args);
236
+ if (state.lastSuccess && state.lastSuccess.signature === signature) {
237
+ return {
238
+ reason:
239
+ "This exact tool call already ran and succeeded; skipping re-execution. " +
240
+ `Its result was:\n\n${state.lastSuccess.resultText}\n\n` +
241
+ "Use that result instead of calling it again.",
242
+ };
243
+ }
244
+ return undefined;
245
+ }
246
+
247
+ /** Update loop-hygiene bookkeeping AFTER a call executes (or was
248
+ * immediately blocked/errored) and, on repeated failure, return nudge
249
+ * text to append to the error the model sees. `resultText` is the
250
+ * flattened text content of the tool's result (used to seed
251
+ * `lastSuccess` for the dedup check above). Exported + pure. */
252
+ export function recordToolCallOutcome(
253
+ state: LoopHygieneState,
254
+ toolName: string,
255
+ args: unknown,
256
+ outcome: { isError: boolean; resultText: string },
257
+ ): string | undefined {
258
+ const signature = toolCallSignature(toolName, args);
259
+ if (outcome.isError) {
260
+ const streak = (state.failureStreaks.get(signature) ?? 0) + 1;
261
+ state.failureStreaks.set(signature, streak);
262
+ if (streak >= LOOP_HYGIENE.MAX_CONSECUTIVE_FAILURES) {
263
+ return (
264
+ `\n\n(This exact call has now failed ${streak} times in a row. ` +
265
+ "Stop repeating it verbatim — try a different tool, different arguments, " +
266
+ "or explain to the user what's blocking you.)"
267
+ );
268
+ }
269
+ return undefined;
270
+ }
271
+ state.failureStreaks.delete(signature);
272
+ state.lastSuccess = { signature, resultText: outcome.resultText };
273
+ return undefined;
274
+ }
275
+
276
+ /** Flatten a tool result's content (TextContent[] shape) to plain text for
277
+ * dedup-cache display and diagnostics. Mirrors contentText's tolerance
278
+ * for non-array/string shapes. Exported for unit testing. */
279
+ export function toolResultText(content: unknown): string {
280
+ if (typeof content === "string") return content;
281
+ if (Array.isArray(content)) {
282
+ return content
283
+ .filter((p): p is { type: string; text: string } =>
284
+ !!p && (p as { type?: string }).type === "text" && typeof (p as { text?: unknown }).text === "string")
285
+ .map((p) => p.text)
286
+ .join("");
287
+ }
288
+ return "";
289
+ }
290
+
88
291
  /**
89
292
  * System prompt for the web chat assistant.
90
293
  *
@@ -101,7 +304,7 @@ const APPROVAL_TIMEOUT_MS = 120_000;
101
304
  * with the long one). So: state identity + privacy in one breath, then tell it
102
305
  * plainly to answer what was asked and not to greet or recite capabilities.
103
306
  */
104
- export const WEB_CHAT_PROMPT_VERSION = "2026-06-21-minimal-v1";
307
+ export const WEB_CHAT_PROMPT_VERSION = "2026-07-07-app-aware-v1";
105
308
 
106
309
  export function webChatPromptFingerprint(prompt: string): string {
107
310
  return createHash("sha256").update(prompt).digest("hex").slice(0, 12);
@@ -127,8 +330,17 @@ export function buildWebChatSystemPrompt(
127
330
  // those two — and telling it to answer from knowledge first — is what keeps a
128
331
  // small model from reaching for a tool on math/writing/general questions.
129
332
  const hasTools = opts?.hasTools ?? true;
333
+ // App-aware assistant (plan §6.6): tell the model these tools exist and
334
+ // when to reach for them instead of describing the UI blind — "when the
335
+ // user asks where/how in the app, navigate or spotlight instead of
336
+ // describing blind" is the literal steer the task calls for. Included
337
+ // whenever the session has tools at all (these three are always in the
338
+ // allowlist alongside WELCOME_TOOLS — see webChatToolAllowlist).
339
+ const appAwareLine = hasTools
340
+ ? ` You can see and act on the app: \`navigate_app\` switches the view, \`spotlight_ui\` highlights a control (use \`get_current_app_context\` first for its ref). When asked where/how in the app, navigate or spotlight instead of describing it blind.`
341
+ : "";
130
342
  const toolsLine = hasTools
131
- ? ` Answer directly from your own knowledge whenever you can. Only call a tool when you truly need information you don't have: \`web_search\` for current or external facts (news, current events, prices, latest docs), \`read\` for a specific local file the user points you to. Never use a tool for general questions, explanations, math, or writing — just answer. When a tool returns results, pull the answer out of the result text and state it directly — never tell the user to go open the links or check the sources themselves; reading them is your job.`
343
+ ? ` Answer directly from your own knowledge whenever you can. Only call a tool when you truly need information you don't have: \`web_search\` for current or external facts (news, current events, prices, latest docs), \`read\` for a specific local file the user points you to. Never use a tool for general questions, explanations, math, or writing — just answer. When a tool returns results, pull the answer out of the result text and state it directly — never tell the user to go open the links or check the sources themselves; reading them is your job.${appAwareLine}`
132
344
  : ` You have no tools in this session, so answer from your own knowledge; if something needs current or external data you can't reach, say so briefly instead of pretending to look it up.`;
133
345
 
134
346
  return `You are mlx-bun's built-in assistant, running entirely on the user's own Apple-silicon Mac — nothing they type leaves the machine.${modelLine}${aboutLine}${toolsLine}
@@ -145,48 +357,240 @@ type PiWs = ServerWebSocket<PiWsData>;
145
357
 
146
358
  // ---- WS protocol message shapes --------------------------------------
147
359
 
148
- /** Decision a browser returns for a gated tool call. */
149
- type ApprovalDecision = "allow" | "deny";
360
+ /** Decision a browser returns for a gated tool call. `editedArgs`, when
361
+ * present on an "allow", replaces the tool's proposed input before it runs
362
+ * (LM Studio's editable-args pattern — plan §5.4/§6.5). `alwaysAllow` marks
363
+ * the decision as "and remember this tool" — handled by resolveApproval,
364
+ * which persists it to the durable tool-approvals config (src/tool-
365
+ * approvals.ts) before settling the pending promise. */
366
+ export type ApprovalDecision = "allow" | "deny";
150
367
 
151
368
  /** An image attachment from the browser: base64 (no data: prefix) + mime. */
152
- interface ImageAttachment {
369
+ export interface ImageAttachment {
153
370
  data: string;
154
371
  mimeType: string;
155
372
  }
156
373
 
157
- /** Client -> server frames. */
158
- type ClientMessage =
374
+ /** Client -> server frames. Exported type-only for the web frontend
375
+ * (src/web/src/*.ts import it via `import type` — see tsconfig.web.json
376
+ * and scripts/build-web.ts) so the WS contract is checked at compile time
377
+ * on both ends instead of drifting between server TS and untyped inline
378
+ * browser JS (the observed Phase-0 bug class). */
379
+ export type ClientMessage =
159
380
  | { type: "prompt"; text: string; images?: ImageAttachment[] }
160
381
  | { type: "abort" }
161
- | { type: "approval"; callId: string; decision: ApprovalDecision }
382
+ // `editedArgs`: when the user edited the proposed arguments in the
383
+ // approval card's textarea before approving, the (re-parsed) JSON object
384
+ // to substitute for the tool's original input — mutated onto the pi SDK's
385
+ // ToolCallEvent.input in place (the SDK's documented mutation contract;
386
+ // see installApprovalGate). Ignored on "deny". `alwaysAllow`: persist this
387
+ // tool to the durable always-allow config (src/tool-approvals.ts) so
388
+ // future calls to the SAME TOOL NAME skip the card entirely, on this and
389
+ // every other session.
390
+ | { type: "approval"; callId: string; decision: ApprovalDecision; editedArgs?: Record<string, unknown>; alwaysAllow?: boolean }
162
391
  // Toggle the model's reasoning channel on/off (only meaningful when the
163
392
  // model supports thinking — ready.thinking). Maps to Pi's session thinking
164
393
  // level: "medium" (on) ↔ "off". Pi sends it as enable_thinking to the server.
165
394
  | { type: "set_thinking"; enabled: boolean }
395
+ // Opt-in toggle (plan §5.4/§6.5/§9 Phase 2, default OFF): "let the agent
396
+ // touch files on this machine". Persisted per-browser by the client
397
+ // (localStorage `mlxbun.codingTools`); the ALLOWLIST decision is enforced
398
+ // here, server-side. buildPiAgentSurface's tool list is fixed for the
399
+ // life of an AgentSession (pi's DefaultResourceLoader config is baked in
400
+ // at createAgentSessionServices time), so this can't rewire an
401
+ // in-progress session — it takes effect starting with the next
402
+ // new_session/open_session/fork_session (createRuntimeFactory reads
403
+ // this.codingToolsRequested fresh every time the SDK calls it, including
404
+ // from AgentSessionRuntime.newSession()'s internal re-invocation). The
405
+ // `coding_tools` ServerMessage below reports the honest current-vs-pending
406
+ // state so the UI can say so instead of implying it applies immediately.
407
+ | { type: "set_coding_tools"; enabled: boolean }
166
408
  // Select the active LoRA adapter for subsequent turns (null = none/base).
167
409
  // app.html mounts it (POST /v1/adapters) before sending this; the
168
410
  // before_provider_request hook injects it into the provider payload.
169
411
  | { type: "set_adapter"; id: string | null }
170
- // Per-request sampling overrides for subsequent turns. Each field is
171
- // optional; null/undefined means "leave it to the server default"
172
- // (the mode-aware recommended value resolved in toOptions). A present
173
- // numeric value is injected into the provider payload and always wins.
174
- // The before_provider_request hook injects whatever is set here.
175
- | { type: "set_sampling"; temperature?: number | null; top_p?: number | null; top_k?: number | null }
412
+ // Per-request sampling overrides. Each field is optional; null/undefined
413
+ // means "leave it to the server default" (the mode-aware recommended
414
+ // value resolved in toOptions). A present numeric value is injected into
415
+ // the provider payload and always wins. The before_provider_request hook
416
+ // injects whatever is set here.
417
+ //
418
+ // `scope` (plan §9 Phase 3 / beat matrix Axis 4 "per-message sampling
419
+ // scope"): "session" (default, omitting the field means this too) persists
420
+ // the override on the connection until changed — today's behavior,
421
+ // unchanged. "next_turn" stores a ONE-SHOT override applied to exactly
422
+ // the next prompt's injection, then cleared automatically — it never
423
+ // touches the session-level values (a later plain set_sampling with
424
+ // scope:"session" or omitted still sees whatever was there before this
425
+ // message). Composition when both are present: next_turn's SET fields
426
+ // win over session's for that one request; anything next_turn leaves
427
+ // null/unset falls back to the session-level override for that field
428
+ // (composeSampling — session-level itself falls back to the server's
429
+ // mode-aware default per injectSampling's existing contract).
430
+ | {
431
+ type: "set_sampling";
432
+ scope?: "session" | "next_turn";
433
+ temperature?: number | null;
434
+ top_p?: number | null;
435
+ top_k?: number | null;
436
+ min_p?: number | null;
437
+ xtc_probability?: number | null;
438
+ xtc_threshold?: number | null;
439
+ repetition_penalty?: number | null;
440
+ repetition_context_size?: number | null;
441
+ presence_penalty?: number | null;
442
+ frequency_penalty?: number | null;
443
+ seed?: number | null;
444
+ }
445
+ // Per-session USER system prompt, layered ON TOP OF (not replacing) the
446
+ // built-in surface prompt (buildWebChatSystemPrompt + memoryHint). null
447
+ // clears it. Applied by the before_agent_start hook on the NEXT turn (pi
448
+ // calls it fresh inside every session.prompt(), so this is a true
449
+ // per-turn injection, not a session-creation-only setting — see
450
+ // installSystemPromptHook). Presets v1 (composer.ts) send this alongside
451
+ // set_sampling to apply a saved bundle.
452
+ | { type: "set_system_prompt"; text: string | null }
176
453
  // Session management (recent-chats sidebar + new chat).
177
454
  | { type: "new_session" }
178
455
  | { type: "list_sessions" }
179
456
  | { type: "open_session"; path: string }
180
457
  | { type: "fork_session"; path: string }
181
- | { type: "delete_session"; path: string };
458
+ | { type: "delete_session"; path: string }
459
+ // Message actions (plan §5.2). All three navigate the session's leaf
460
+ // pointer via AgentSession.navigateTree() then re-prompt — a sibling
461
+ // branch in the SAME session file, never a new one (see the helpers above
462
+ // findLastUserMessageEntry for the full rationale). Regenerate re-sends the
463
+ // last user message's original content verbatim; edit_resend re-sends it
464
+ // with edited text (images preserved); switch_sibling moves the active
465
+ // leaf to a previously-created sibling (an earlier edit/regeneration)
466
+ // without sending anything new.
467
+ | { type: "regenerate" }
468
+ | { type: "edit_resend"; text: string }
469
+ | { type: "switch_sibling"; entryId: string }
470
+ // App-aware assistant (plan §6.6, §9 Phase 3, beat matrix Axis 12): pushed
471
+ // by chat.ts on every route change AND on wizard-step change (quantize/
472
+ // finetune/dataset), never on a timer. `context` is opaque here (the
473
+ // browser's AppContext shape, src/web/src/assistant.ts) — stored verbatim
474
+ // as the connection's currentAppContext and returned by the
475
+ // get_current_app_context tool, plus mined for the compact one-line
476
+ // ambient context auto-prepended to the NEXT prompt (see
477
+ // installAppContextHook / ambientContextLine). This is a DIFFERENT frame
478
+ // from the server->client `context` ServerMessage (context-window usage,
479
+ // above) — same name, opposite direction, disambiguated by which type
480
+ // union it's a member of; the two are never in scope at the same call
481
+ // site (ClientMessage here vs. ServerMessage below).
482
+ | { type: "context"; context: AppUiContext };
483
+
484
+ // ---- App-aware assistant (plan §6.6, §9 Phase 3, beat matrix Axis 12) --
485
+ //
486
+ // Mirrors src/web/src/assistant.ts's browser-side types (which are the
487
+ // source of truth for shape — this server-side copy is intentionally loose
488
+ // (mostly `unknown`/optional) since the server never constructs a
489
+ // UiSnapshot itself, only stores and forwards what the browser sends. Kept
490
+ // here (not imported) because src/web/src/assistant.ts is bundled into the
491
+ // browser and importing it here would pull browser code into the server
492
+ // build — the inverse of the type-only ClientMessage/ServerMessage import
493
+ // direction the frontend uses for pi-web.ts.
494
+
495
+ /** One element in a captured DOM snapshot — see assistant.ts's
496
+ * UiSnapshotElement for the authoritative shape/doc. */
497
+ export interface UiSnapshotElement {
498
+ ref: string;
499
+ tag: string;
500
+ label: string;
501
+ kind: "interactive" | "region";
502
+ role?: string;
503
+ selector: string;
504
+ spotlightId?: string;
505
+ }
506
+
507
+ export interface UiSnapshot {
508
+ route: string;
509
+ capturedAt: string;
510
+ elements: UiSnapshotElement[];
511
+ }
512
+
513
+ /** The wizard-step process-state fields (§6.6 "carries process state, not
514
+ * just elements") — present only on routes with a step wizard. */
515
+ export interface WizardStep {
516
+ index: number;
517
+ count: number;
518
+ label: string;
519
+ }
520
+
521
+ /** The full context payload pushed by chat.ts on route/step change — see
522
+ * assistant.ts's buildAppContext. Stored verbatim per-connection. */
523
+ export interface AppUiContext {
524
+ route: string;
525
+ view?: string;
526
+ step?: WizardStep;
527
+ snapshot: UiSnapshot;
528
+ }
529
+
530
+ /** Route catalog mirror (src/web/src/ui-catalog.ts's ROUTE_IDS) for
531
+ * server-side validation of navigate_app's route/page param BEFORE it
532
+ * ever reaches the browser — an unknown route is rejected here as a tool
533
+ * error rather than sent as a `ui_navigate` frame the router can't match.
534
+ * Keep in sync with ui-catalog.ts's ROUTE_IDS by hand (no shared-JSON
535
+ * import between the server and browser bundles — see that file's
536
+ * top-of-file comment for why). */
537
+ export const APP_ROUTE_IDS = ["chat", "quantize", "finetune", "dataset", "status"] as const;
538
+ export type AppRouteId = (typeof APP_ROUTE_IDS)[number];
539
+
540
+ export function isAppRouteId(v: string): v is AppRouteId {
541
+ return (APP_ROUTE_IDS as readonly string[]).includes(v);
542
+ }
543
+
544
+ /** Resolve a free-form route/page string the way navigate_app's tool
545
+ * describes its param to the model: either a bare route id ("quantize")
546
+ * or a "#/quantize"-style hash. Returns null for anything else. */
547
+ export function resolveAppRoute(routeOrPage: string): AppRouteId | null {
548
+ const bare = routeOrPage.trim().replace(/^#\/?/, "");
549
+ return isAppRouteId(bare) ? bare : null;
550
+ }
551
+
552
+ /** Compact one-line ambient context (§6.6 "never answer blind") mined from
553
+ * the stored AppUiContext — e.g. "[user is on: Quantize · step 2/4]".
554
+ * Never a snapshot dump; this is the ONLY thing auto-attached to every
555
+ * turn (the full snapshot stays on-demand via get_current_app_context).
556
+ * Returns null when there's no context yet (nothing pushed this
557
+ * connection — e.g. a non-browser WS client, or the very first frame
558
+ * race before chat.ts's first context push lands). Pure + exported for
559
+ * unit testing; mirrors assistant.ts's ambientLine so the two stay in
560
+ * lockstep even though the server never imports the browser module. */
561
+ const APP_ROUTE_LABELS: Record<AppRouteId, string> = {
562
+ chat: "Chat", quantize: "Quantize", finetune: "Fine-tune", dataset: "Build Dataset", status: "Status",
563
+ };
564
+
565
+ export function ambientContextLine(ctx: AppUiContext | null | undefined): string | null {
566
+ if (!ctx || !ctx.route) return null;
567
+ const place = ctx.view ?? APP_ROUTE_LABELS[ctx.route as AppRouteId] ?? ctx.route;
568
+ const stepPart = ctx.step ? ` · step ${ctx.step.index + 1}/${ctx.step.count}` : "";
569
+ return `[user is on: ${place}${stepPart}]`;
570
+ }
182
571
 
183
- /** Server -> client frames. */
184
- type ServerMessage =
572
+ /** Model-author recommended sampling (generation_config.json, server-CLI
573
+ * overrides applied) — sent once on `ready` so the sampling popover can show
574
+ * per-model defaults instead of one hardcoded shape (web-ui-pass-plan.md
575
+ * #14). A field is null when neither the model nor the server configured one
576
+ * (the built-in fallback in toOptions still applies server-side). */
577
+ export interface ReadyGenDefaults {
578
+ temperature: number | null;
579
+ topP: number | null;
580
+ topK: number | null;
581
+ }
582
+
583
+ /** Server -> client frames. Exported type-only — see ClientMessage above. */
584
+ export type ServerMessage =
185
585
  // `vision`: whether the loaded model can accept images (drives the UI's
186
586
  // image-attach affordance — false on e4b until the SigLIP sidecar lands).
587
+ // `audio`: whether the loaded model can accept `input_audio` parts (audio
588
+ // tower loaded or loadable). Discovery-only for now: pi's Model.input type
589
+ // has no "audio" modality, so unlike `vision` this flag never feeds the
590
+ // provider declaration — clients use it instead of probing for a 400.
187
591
  // `thinking`: whether the model has a switchable reasoning channel (drives
188
592
  // the UI's thinking on/off toggle; false hides it).
189
- | { type: "ready"; model: string; vision: boolean; thinking: boolean }
593
+ | { type: "ready"; model: string; vision: boolean; audio: boolean; thinking: boolean; genDefaults: ReadyGenDefaults }
190
594
  | { type: "turn_start" }
191
595
  | { type: "text_delta"; delta: string }
192
596
  | { type: "thinking_delta"; delta: string }
@@ -194,14 +598,49 @@ type ServerMessage =
194
598
  | { type: "tool_approval_request"; callId: string; tool: string; args: unknown }
195
599
  | { type: "tool_update"; callId: string; chunk: unknown }
196
600
  | { type: "tool_end"; callId: string; ok: boolean; result: unknown }
197
- | { type: "turn_end" }
601
+ // `lane`: the serving lane the just-finished turn ran on (serial /
602
+ // serial+spec / batched), correlated via the lane registry keyed by the
603
+ // AssistantMessage's responseId (the pi SDK's own usage parsing drops
604
+ // custom fields — see src/serve/lane-registry.ts). Absent when the turn
605
+ // produced no assistant message with a responseId (e.g. an aborted turn
606
+ // before any model call) — never guessed client-side (risk #5).
607
+ | { type: "turn_end"; lane?: Lane }
198
608
  | { type: "queue_update"; steering: readonly string[]; followUp: readonly string[] }
199
609
  // Replay a session's transcript (rebuilds the thread); and the sidebar list.
200
610
  | { type: "history"; items: HistoryItem[] }
201
611
  | { type: "sessions"; items: SessionListItem[]; activePath?: string }
612
+ // Sibling group for the LAST user message (edit-and-resend's `< i/n >`
613
+ // toggle, plan §5.2). Sent after every history rebuild and after a
614
+ // completed turn. `count` <= 1 means no toggle to show (the common case:
615
+ // no edits/regenerations yet). `entryId` is the currently-active sibling's
616
+ // id; `siblingIds` is the full ordered group so the client can resolve
617
+ // "previous/next" locally and send switch_sibling with a concrete id
618
+ // (never inferred/guessed — the server is the source of truth for order).
619
+ | { type: "siblings"; entryId?: string; index: number; count: number; siblingIds: string[] }
202
620
  // Context-window usage indicator. tokens/percent are null right after a
203
621
  // compaction until the next assistant reply (pi can't estimate yet).
204
622
  | { type: "context"; tokens: number | null; contextWindow: number; percent: number | null }
623
+ // Effective vs. pending codingTools state (set_coding_tools above).
624
+ // `active` is what the CURRENT session actually built with (gates
625
+ // whether the approval card can ever fire); `pending` is the browser's
626
+ // most recent request, which only takes effect on the next
627
+ // new/opened/forked session — surfaced so the settings copy can say
628
+ // "will apply to your next new chat" instead of lying about immediacy.
629
+ | { type: "coding_tools"; active: boolean; pending: boolean }
630
+ // Full always-allow set for THIS machine (src/tool-approvals.ts), sent on
631
+ // `ready` and after any change, so the settings panel can list/forget
632
+ // durable per-tool approvals without a separate REST round-trip.
633
+ | { type: "tool_approvals"; alwaysAllow: string[] }
634
+ // App-aware assistant (plan §6.6): the navigate_app / spotlight_ui tools'
635
+ // side-channel notifications, mapped 1:1 from the tool call's params —
636
+ // never guessed or re-derived, so chat.ts's router/spotlight get exactly
637
+ // what the model asked for. `route` here is always a validated
638
+ // AppRouteId (isAppRouteId already ran server-side before this frame is
639
+ // ever sent — see navigate_app's execute). ui_spotlight's fields mirror
640
+ // SpotlightRequest 1:1 (assistant.ts resolves ref/label/selector/target
641
+ // client-side, in that precedence order, against its last snapshot).
642
+ | { type: "ui_navigate"; route: AppRouteId }
643
+ | { type: "ui_spotlight"; ref?: string; label?: string; selector?: string; target?: string; message?: string; route?: AppRouteId }
205
644
  | { type: "error"; message: string };
206
645
 
207
646
  // ---- Session serialization (pure, unit-tested) -----------------------
@@ -223,6 +662,11 @@ export interface HistoryItem {
223
662
  /** Model reasoning/thinking, kept separate from the final answer. */
224
663
  thinking?: string;
225
664
  tools: HistoryToolItem[];
665
+ /** Session entry id (user items only). Lets the browser correlate "the
666
+ * last user message" in the replayed thread with the `siblings` frame's
667
+ * entryId, so it knows which DOM node gets the edit/sibling-toggle UI
668
+ * (plan §5.2). Undefined for assistant items — they aren't fork targets. */
669
+ entryId?: string;
226
670
  }
227
671
 
228
672
  /** A row in the recent-chats sidebar, derived from pi's SessionInfo. */
@@ -277,7 +721,7 @@ export function serializeHistory(entries: readonly SessionEntry[]): HistoryItem[
277
721
  if (!m) continue;
278
722
  if (m.role === "user") {
279
723
  const text = contentText(m.content);
280
- if (text.trim()) items.push({ role: "user", text, tools: [] });
724
+ if (text.trim()) items.push({ role: "user", text, tools: [], entryId: entry.id });
281
725
  } else if (m.role === "assistant") {
282
726
  const parts = Array.isArray(m.content) ? (m.content as unknown[]) : [];
283
727
  const text = contentText(parts);
@@ -311,6 +755,102 @@ function toPiImages(images?: ImageAttachment[]): ImageContent[] | undefined {
311
755
  return images.map((i) => ({ type: "image", data: i.data, mimeType: i.mimeType }));
312
756
  }
313
757
 
758
+ // ---- Message actions: regenerate / edit-and-resend (pure helpers) ----
759
+ //
760
+ // Both features reuse ONE pi SDK primitive: AgentSession.navigateTree(entryId)
761
+ // on a user-message entry moves the session's leaf pointer to that entry's
762
+ // PARENT (without deleting anything) and returns the message's own text in
763
+ // `editorText`. Calling session.prompt(...) right after appends a NEW child
764
+ // under that same parent — i.e. a sibling of the original message, in the
765
+ // SAME session file. This is exactly "regenerate" (re-send the identical
766
+ // content) and "edit-and-resend" (re-send edited content) without any
767
+ // file-level fork/session spam (plan §5.2's explicit constraint).
768
+ //
769
+ // navigateTree's own editorText is TEXT-ONLY (pi's editor-reopen use case
770
+ // doesn't carry images). We extract the full original `content` ourselves
771
+ // (text + images) from the entry directly so regenerate/edit-resend don't
772
+ // silently drop image attachments — a small improvement over what the SDK
773
+ // primitive alone would give us.
774
+
775
+ /** One user message entry as found by findLastUserMessageEntry/userMessageSiblings. */
776
+ export interface UserMessageEntryInfo {
777
+ id: string;
778
+ parentId: string | null;
779
+ text: string;
780
+ images: ImageContent[];
781
+ }
782
+
783
+ /** Extract {text, images} from a user message's raw `content` field
784
+ * (string | (TextContent|ImageContent)[]). Pure. */
785
+ function extractUserContent(content: unknown): { text: string; images: ImageContent[] } {
786
+ if (typeof content === "string") return { text: content, images: [] };
787
+ if (Array.isArray(content)) {
788
+ const text = contentText(content);
789
+ const images = (content as unknown[])
790
+ .filter((p): p is { type: string; data: string; mimeType: string } =>
791
+ !!p && (p as { type?: string }).type === "image" &&
792
+ typeof (p as { data?: unknown }).data === "string" && typeof (p as { mimeType?: unknown }).mimeType === "string")
793
+ .map((p) => ({ type: "image" as const, data: p.data, mimeType: p.mimeType }));
794
+ return { text, images };
795
+ }
796
+ return { text: "", images: [] };
797
+ }
798
+
799
+ /** Find the LAST user-message entry in a session's entries (in append order).
800
+ * Used by both `regenerate` (re-send it as-is) and as the sibling group's
801
+ * anchor for `edit_resend` (always edits the last user message — plan §5.2
802
+ * scopes this to the last message, not any earlier one). Pure. */
803
+ export function findLastUserMessageEntry(entries: readonly SessionEntry[]): UserMessageEntryInfo | undefined {
804
+ for (let i = entries.length - 1; i >= 0; i--) {
805
+ const e = entries[i];
806
+ if (!e || e.type !== "message") continue;
807
+ const m = (e as { message?: { role?: string; content?: unknown } }).message;
808
+ if (m?.role !== "user") continue;
809
+ const { text, images } = extractUserContent(m.content);
810
+ return { id: e.id, parentId: e.parentId, text, images };
811
+ }
812
+ return undefined;
813
+ }
814
+
815
+ /** Sibling group (edit history) for a user message: every user-message entry
816
+ * sharing the same parentId, in append order, plus the 1-based index of
817
+ * `entryId` within that group. Drives the browser's `< i/n >` toggle. Pure. */
818
+ export interface SiblingGroup {
819
+ parentId: string | null;
820
+ siblingIds: string[];
821
+ index: number; // 1-based position of the queried entryId within siblingIds
822
+ }
823
+
824
+ export function userMessageSiblings(entries: readonly SessionEntry[], entryId: string): SiblingGroup | undefined {
825
+ const target = entries.find((e) => e.id === entryId);
826
+ if (!target) return undefined;
827
+ const parentId = target.parentId;
828
+ const siblingIds = entries
829
+ .filter((e): e is SessionEntry & { type: "message" } =>
830
+ e.type === "message" && e.parentId === parentId &&
831
+ (e as { message?: { role?: string } }).message?.role === "user")
832
+ .map((e) => e.id);
833
+ const index = siblingIds.indexOf(entryId) + 1;
834
+ if (index === 0) return undefined;
835
+ return { parentId, siblingIds, index };
836
+ }
837
+
838
+ /** Walk DOWN from an entry to the deepest descendant, always following the
839
+ * MOST RECENTLY APPENDED child at each level (last in getChildren's
840
+ * insertion order) — i.e. "what would this branch show if it were the
841
+ * active leaf". Used by `switch_sibling` to resolve a sibling's own leaf
842
+ * (which may itself have further edits/regenerations under it) rather than
843
+ * just landing on the sibling's own root entry. Pure given a children-lookup
844
+ * function so it's unit-testable without a live SessionManager. */
845
+ export function deepestLeafFrom(fromId: string, getChildren: (parentId: string) => SessionEntry[]): string {
846
+ let cur = fromId;
847
+ for (;;) {
848
+ const children = getChildren(cur);
849
+ if (children.length === 0) return cur;
850
+ cur = children[children.length - 1]!.id;
851
+ }
852
+ }
853
+
314
854
  /** Map pi's SessionInfo[] to sidebar rows, newest first. */
315
855
  export function toSessionListItems(infos: readonly SessionInfo[]): SessionListItem[] {
316
856
  return infos
@@ -346,14 +886,20 @@ export function mapEventToFrames(event: AgentSessionEvent): ServerMessage[] {
346
886
  // 400'd) WITHOUT any error being thrown up to the WS message handler, so
347
887
  // it would otherwise vanish — the browser just sees an empty turn ("no
348
888
  // messages"). Surface it as an error frame so the UI can show it.
349
- const msg = (event as { message?: { stopReason?: string; errorMessage?: string } }).message;
889
+ const msg = (event as { message?: { stopReason?: string; errorMessage?: string; responseId?: string } }).message;
890
+ // Lane correlation (docs/design/web-chat-redesign.md §2.3 caveat / risk
891
+ // #5): the pi SDK's AssistantMessage.responseId is the chat-completion
892
+ // response's own `id` (chatcmpl-…), which server.ts records against the
893
+ // lane the moment it decides one. Looked up here — never inferred — so
894
+ // a mismatch or absence just omits the field rather than guessing.
895
+ const lane = msg?.responseId ? getLane(msg.responseId) : undefined;
350
896
  if (msg?.stopReason === "error") {
351
897
  return [
352
898
  { type: "error", message: msg.errorMessage || "the model request failed" },
353
- { type: "turn_end" },
899
+ { type: "turn_end", ...(lane ? { lane } : {}) },
354
900
  ];
355
901
  }
356
- return [{ type: "turn_end" }];
902
+ return [{ type: "turn_end", ...(lane ? { lane } : {}) }];
357
903
  }
358
904
  case "message_update": {
359
905
  const ame = event.assistantMessageEvent;
@@ -419,16 +965,35 @@ export function injectAdapter(
419
965
 
420
966
  /** Per-request sampling overrides carried on the session and injected into the
421
967
  * provider payload by the before_provider_request hook. Each field is either a
422
- * user-set override or null/undefined ("use the server's mode-aware default"). */
968
+ * user-set override or null/undefined ("use the server's mode-aware default").
969
+ * Field names match the server's ChatRequest wire names (not toOptions' camelCase)
970
+ * since these are injected directly into the outgoing chat-completions payload. */
423
971
  export interface SamplingOverrides {
424
972
  temperature?: number | null;
425
973
  top_p?: number | null;
426
974
  top_k?: number | null;
975
+ min_p?: number | null;
976
+ xtc_probability?: number | null;
977
+ xtc_threshold?: number | null;
978
+ repetition_penalty?: number | null;
979
+ repetition_context_size?: number | null;
980
+ presence_penalty?: number | null;
981
+ frequency_penalty?: number | null;
982
+ seed?: number | null;
427
983
  }
428
984
 
985
+ /** Every SamplingOverrides field name, in wire order — the single source of
986
+ * truth `injectSampling` and the `set_sampling` handler both walk, so a new
987
+ * field only needs adding here + the interface above. */
988
+ const SAMPLING_FIELDS = [
989
+ "temperature", "top_p", "top_k", "min_p", "xtc_probability", "xtc_threshold",
990
+ "repetition_penalty", "repetition_context_size", "presence_penalty",
991
+ "frequency_penalty", "seed",
992
+ ] as const satisfies readonly (keyof SamplingOverrides)[];
993
+
429
994
  /** Inject the user's sampling overrides into the outgoing chat-completions
430
995
  * payload (which reaches the server's toOptions, where an explicit
431
- * temperature/top_p/top_k always wins). Only finite numbers are injected; a
996
+ * request field always wins). Only finite numbers are injected; a
432
997
  * null/undefined/unset field is left off so the server falls back to its
433
998
  * mode-aware recommended default. Returns undefined when nothing is set so the
434
999
  * hook can keep the payload unchanged. Pure + exported for unit testing. */
@@ -439,21 +1004,247 @@ export function injectSampling(
439
1004
  if (!s) return undefined;
440
1005
  const out: Record<string, unknown> = { ...payload };
441
1006
  let changed = false;
442
- if (typeof s.temperature === "number" && Number.isFinite(s.temperature)) {
443
- out.temperature = s.temperature;
444
- changed = true;
445
- }
446
- if (typeof s.top_p === "number" && Number.isFinite(s.top_p)) {
447
- out.top_p = s.top_p;
448
- changed = true;
449
- }
450
- if (typeof s.top_k === "number" && Number.isFinite(s.top_k)) {
451
- out.top_k = s.top_k;
452
- changed = true;
1007
+ for (const field of SAMPLING_FIELDS) {
1008
+ const v = s[field];
1009
+ if (typeof v === "number" && Number.isFinite(v)) {
1010
+ out[field] = v;
1011
+ changed = true;
1012
+ }
453
1013
  }
454
1014
  return changed ? out : undefined;
455
1015
  }
456
1016
 
1017
+ // ---- Per-message sampling scope (beat matrix Axis 4 / plan §9 Phase 3) --
1018
+ //
1019
+ // set_sampling's optional `scope` field: "session" (default, today's
1020
+ // behavior — persists on the connection until changed) or "next_turn" (a
1021
+ // one-shot override for exactly the next prompt, then cleared). The
1022
+ // composition rule is "next_turn OVER session": a next_turn field that's
1023
+ // a real number wins over the session-level value for that SAME field;
1024
+ // a next_turn field left null/undefined falls through to the session
1025
+ // override (which itself falls through to the server default per
1026
+ // injectSampling's existing null-means-unset contract). This mirrors
1027
+ // optiq Lab's per-message temp/max-tokens/thinking granularity cited in
1028
+ // the beat matrix, layered onto (not replacing) Phase 1's session-level
1029
+ // overrides.
1030
+ //
1031
+ // Lifecycle: armed by the set_sampling handler when scope is "next_turn"
1032
+ // (stored separately from the session-level `sampling` field so a
1033
+ // one-shot override never permanently clobbers the session default it's
1034
+ // laid over). Consumed at the FIRST before_provider_request of the very
1035
+ // next prompt — before_agent_start fires exactly once per session.prompt()
1036
+ // call (see agent-session.js's prompt() body), which is the natural
1037
+ // "one prompt" boundary, but the actual sampling injection happens in
1038
+ // before_provider_request (fired once per model call, i.e. once per
1039
+ // tool-loop turn) — so the override is consumed (cleared) the first time
1040
+ // before_provider_request fires after being armed, guaranteeing it rides
1041
+ // on exactly the next outgoing provider request and never leaks into a
1042
+ // second turn within the same prompt, let alone the next user message.
1043
+
1044
+ /** Compose the effective sampling overrides for one outgoing provider
1045
+ * request: every field in `nextTurn` that's a real number wins; anything
1046
+ * left null/undefined in `nextTurn` falls back to `session`'s value for
1047
+ * that field. Pure + exported for unit testing. When `nextTurn` is
1048
+ * undefined (no one-shot override armed), returns `session` unchanged. */
1049
+ export function composeSampling(
1050
+ session: SamplingOverrides,
1051
+ nextTurn: SamplingOverrides | undefined,
1052
+ ): SamplingOverrides {
1053
+ if (!nextTurn) return session;
1054
+ const out: SamplingOverrides = { ...session };
1055
+ for (const field of SAMPLING_FIELDS) {
1056
+ const v = nextTurn[field];
1057
+ if (typeof v === "number" && Number.isFinite(v)) out[field] = v;
1058
+ }
1059
+ return out;
1060
+ }
1061
+
1062
+ /** Pure state machine for the arm/consume lifecycle above, factored out of
1063
+ * PiWebSession so "set → one turn → cleared" is unit-testable without a
1064
+ * live session. `applySetSampling` is the set_sampling handler body;
1065
+ * `consumeForRequest` is the before_provider_request handler body — both
1066
+ * mutate and return a NEW state object (no hidden mutation) so tests can
1067
+ * assert on each step's output directly. */
1068
+ export interface SamplingScopeState {
1069
+ session: SamplingOverrides;
1070
+ nextTurn: SamplingOverrides | undefined;
1071
+ }
1072
+
1073
+ export function initialSamplingScopeState(): SamplingScopeState {
1074
+ return { session: {}, nextTurn: undefined };
1075
+ }
1076
+
1077
+ /** Apply a set_sampling message to scope state: scope "next_turn" arms the
1078
+ * one-shot override (leaving `session` untouched); scope "session" or
1079
+ * omitted replaces the session-level override (today's behavior). Pure. */
1080
+ export function applySetSampling(
1081
+ state: SamplingScopeState,
1082
+ overrides: SamplingOverrides,
1083
+ scope: "session" | "next_turn" | undefined,
1084
+ ): SamplingScopeState {
1085
+ return scope === "next_turn"
1086
+ ? { session: state.session, nextTurn: overrides }
1087
+ : { session: overrides, nextTurn: state.nextTurn };
1088
+ }
1089
+
1090
+ /** Compute the effective overrides for ONE outgoing provider request and
1091
+ * return the post-consumption state (nextTurn always cleared, whether or
1092
+ * not one was armed — consuming "nothing armed" is a no-op clear). Pure. */
1093
+ export function consumeForRequest(
1094
+ state: SamplingScopeState,
1095
+ ): { effective: SamplingOverrides; nextState: SamplingScopeState } {
1096
+ return {
1097
+ effective: composeSampling(state.session, state.nextTurn),
1098
+ nextState: { session: state.session, nextTurn: undefined },
1099
+ };
1100
+ }
1101
+
1102
+ /**
1103
+ * Layer the user's custom system prompt onto (not replacing) the built-in
1104
+ * surface prompt for one turn. Called from the `before_agent_start` hook
1105
+ * body with `event.systemPrompt` (the fully-assembled base: identity +
1106
+ * memoryHint) and the connection's stored override. Returns undefined when
1107
+ * there's nothing to layer (null/empty override), so the hook can leave
1108
+ * pi's base prompt alone — mirrors injectAdapter/injectSampling's "undefined
1109
+ * means unchanged" contract. Pure + exported for unit testing.
1110
+ */
1111
+ export function injectSystemPrompt(base: string, custom: string | null | undefined): string | undefined {
1112
+ const trimmed = custom?.trim();
1113
+ if (!trimmed) return undefined;
1114
+ return `${base}\n\n---\n\nThe user has set a custom instruction for this chat. Follow it alongside (not instead of) the guidance above:\n\n${trimmed}`;
1115
+ }
1116
+
1117
+ /**
1118
+ * Apply the browser's edited tool arguments onto pi's mutable
1119
+ * ToolCallEvent.input IN PLACE, per the SDK's documented tool_call
1120
+ * mutation contract (installApprovalGate's own doc comment quotes it: no
1121
+ * return-value channel exists for modified args). A no-op when
1122
+ * `editedArgs` is undefined (the common "approved as proposed" case).
1123
+ * Replaces the input's keys wholesale (delete-then-assign) rather than a
1124
+ * shallow merge, so a key the user removed from the textarea actually
1125
+ * disappears from the call instead of lingering from the original
1126
+ * proposal. Pure enough to unit test the mutation shape without a live
1127
+ * ExtensionAPI/ToolCallEvent — takes any mutable record, not the SDK type.
1128
+ */
1129
+ export function applyEditedArgs(input: Record<string, unknown>, editedArgs: Record<string, unknown> | undefined): void {
1130
+ if (!editedArgs) return;
1131
+ for (const key of Object.keys(input)) delete input[key];
1132
+ Object.assign(input, editedArgs);
1133
+ }
1134
+
1135
+ // ---- App-aware assistant tools (plan §6.6, §9 Phase 3, beat matrix ----
1136
+ // Axis 12): get_current_app_context / navigate_app / spotlight_ui.
1137
+ //
1138
+ // Same shape as PortfolioManager's createPortfolioTools(getAppContext,
1139
+ // clientId, notify) — the reference implementation this ports (read in
1140
+ // full: server/src/agent/portfolio-tools.ts). Built here (not in
1141
+ // pi-session.ts's shared buildPiAgentSurface) because these three tools are
1142
+ // PER-CONNECTION: they close over this session's currentAppContext getter
1143
+ // and its `notify` (send) function, whereas buildPiAgentSurface's tools are
1144
+ // stateless singletons shared across every session. Read-only/navigate
1145
+ // class (no approval gate, see APP_AWARE_TOOL_NAMES above): navigation and
1146
+ // spotlighting are reversible UI actions, not machine mutations — exactly
1147
+ // the same "auto-allow" reasoning READ_ONLY_TOOLS already applies to
1148
+ // memory/web tools. approval-gated ui_act (fill a form, start a job) is
1149
+ // explicitly Phase 5 per §6.6's scope guard, not implemented here.
1150
+
1151
+ function appAwareTextResult(text: string, details: unknown = {}): { content: [{ type: "text"; text: string }]; details: unknown } {
1152
+ return { content: [{ type: "text", text }], details };
1153
+ }
1154
+
1155
+ /** Build the three app-aware tools for one connection. `getContext` reads
1156
+ * the connection's currentAppContext (last `context` ClientMessage, or
1157
+ * null before the first one arrives); `notify` sends a ServerMessage
1158
+ * (ui_navigate/ui_spotlight) to that same browser tab. Pure enough to
1159
+ * unit test the tool -> frame mapping via a fake notify collecting sent
1160
+ * frames (tests/pi-web.test.ts). */
1161
+ export function createAppAwareTools(
1162
+ getContext: () => AppUiContext | null,
1163
+ notify: (msg: ServerMessage) => void,
1164
+ ): ToolDefinition[] {
1165
+ const getCurrentAppContext = defineTool({
1166
+ name: "get_current_app_context",
1167
+ label: "Get Current App Context",
1168
+ description:
1169
+ "Returns the current web-app UI context: route, view, wizard step (if on one), and a uiSnapshot of visible interactive elements (ref, label, selector, spotlightId). Call this before spotlight_ui to pick the correct ref, label, or target — do not guess selectors blind.",
1170
+ parameters: Type.Object({}),
1171
+ execute: async () => {
1172
+ const ctx = getContext();
1173
+ if (!ctx) {
1174
+ return appAwareTextResult(
1175
+ "No app context received yet — the browser hasn't reported its current view. Ask the user what page they're on, or proceed without UI tools.",
1176
+ );
1177
+ }
1178
+ return appAwareTextResult(JSON.stringify(ctx, null, 2), { context: ctx });
1179
+ },
1180
+ });
1181
+
1182
+ const navigateApp = defineTool({
1183
+ name: "navigate_app",
1184
+ label: "Navigate App",
1185
+ description:
1186
+ "Navigate the user's browser to a different page in the app. Use when they ask to be taken somewhere (e.g. \"take me to Quantize\", \"open the fine-tune wizard\"). Pass route or page as one of: chat, quantize, finetune, dataset, status (or \"#/quantize\"-style). Reversible and never needs approval. Does not highlight a specific control — use spotlight_ui separately for that.",
1187
+ parameters: Type.Object({
1188
+ route: Type.Optional(Type.String({ description: "chat | quantize | finetune | dataset | status (or #/quantize)" })),
1189
+ page: Type.Optional(Type.String({ description: "Same as route — alternate param name" })),
1190
+ }),
1191
+ execute: async (_id, params) => {
1192
+ const resolved = resolveAppRoute(params.route ?? params.page ?? "");
1193
+ if (!resolved) {
1194
+ return {
1195
+ ...appAwareTextResult(`Unknown route. Valid routes: ${APP_ROUTE_IDS.join(", ")}.`),
1196
+ isError: true,
1197
+ };
1198
+ }
1199
+ notify({ type: "ui_navigate", route: resolved });
1200
+ return appAwareTextResult(
1201
+ `Navigating the user to ${resolved}. They can see the page now — use spotlight_ui if you need to point at a specific control.`,
1202
+ { route: resolved },
1203
+ );
1204
+ },
1205
+ });
1206
+
1207
+ const spotlightUi = defineTool({
1208
+ name: "spotlight_ui",
1209
+ label: "Spotlight UI",
1210
+ description:
1211
+ "Highlight a visible UI control with a brief, auto-dismissing spotlight (never blocks the user — they can keep working through it). Call get_current_app_context first and prefer ref from its uiSnapshot.elements; or pass label (visible text, e.g. \"Source model\"); or selector; or target (a curated catalog id, e.g. hub-browse, quantize-source). Optional route navigates there first. Optional message is a short hint shown in the popover.",
1212
+ parameters: Type.Object({
1213
+ ref: Type.Optional(Type.String({ description: "Element ref from uiSnapshot.elements[].ref (preferred)" })),
1214
+ label: Type.Optional(Type.String({ description: "Visible label/text to match, e.g. Source model, New chat" })),
1215
+ selector: Type.Optional(Type.String({ description: "CSS selector, when ref/label are unknown" })),
1216
+ target: Type.Optional(Type.String({ description: "Curated catalog id, e.g. hub-browse, quantize-source, sampling-pill" })),
1217
+ route: Type.Optional(Type.String({ description: "Navigate to this route first (chat, quantize, finetune, dataset, status)" })),
1218
+ message: Type.Optional(Type.String({ description: "Short hint shown in the spotlight popover" })),
1219
+ }),
1220
+ execute: async (_id, params) => {
1221
+ const hasLocator = !!params.ref?.trim() || !!params.label?.trim() || !!params.selector?.trim() || !!params.target?.trim();
1222
+ if (!hasLocator) {
1223
+ return { ...appAwareTextResult("Provide at least one of: ref (from uiSnapshot), label, selector, or target."), isError: true };
1224
+ }
1225
+ let route: AppRouteId | undefined;
1226
+ if (params.route) {
1227
+ route = resolveAppRoute(params.route) ?? undefined;
1228
+ if (!route) {
1229
+ return { ...appAwareTextResult(`Unknown route "${params.route}". Valid routes: ${APP_ROUTE_IDS.join(", ")}.`), isError: true };
1230
+ }
1231
+ }
1232
+ notify({
1233
+ type: "ui_spotlight",
1234
+ ref: params.ref, label: params.label, selector: params.selector, target: params.target,
1235
+ route, message: params.message,
1236
+ });
1237
+ const desc = params.ref ?? params.label ?? params.selector ?? params.target ?? "element";
1238
+ return appAwareTextResult(
1239
+ `Highlighting ${desc}${params.message ? `: ${params.message}` : ""}. The spotlight fades on its own in a few seconds.`,
1240
+ { ...params, route },
1241
+ );
1242
+ },
1243
+ });
1244
+
1245
+ return [getCurrentAppContext, navigateApp, spotlightUi];
1246
+ }
1247
+
457
1248
  /**
458
1249
  * One browser connection's pi agent. Owns the AgentSession, the event
459
1250
  * subscription, and the pending tool-approval handshakes.
@@ -468,10 +1259,46 @@ class PiWebSession {
468
1259
  /** Active LoRA adapter id for this connection (null = none/base model).
469
1260
  * Read by the before_provider_request hook; set via the set_adapter msg. */
470
1261
  private selectedAdapter: string | null = null;
471
- /** Per-request sampling overrides for this connection (set via set_sampling).
472
- * Read by the before_provider_request hook; unset fields fall back to the
473
- * server's mode-aware recommended defaults. */
474
- private sampling: SamplingOverrides = {};
1262
+ /** Session-level + one-shot-next-turn sampling overrides for this
1263
+ * connection (set via set_sampling; scope "session"/omitted mutates
1264
+ * .session, scope "next_turn" arms .nextTurn). Read and consumed by
1265
+ * the before_provider_request hook via consumeForRequest — see the
1266
+ * "Per-message sampling scope" block comment above composeSampling for
1267
+ * the full lifecycle. Wrapped in SamplingScopeState (not two loose
1268
+ * fields) so the exact same applySetSampling/consumeForRequest pure
1269
+ * functions run here and in tests/pi-web.test.ts. */
1270
+ private samplingScope: SamplingScopeState = initialSamplingScopeState();
1271
+ /** User's custom system-prompt text for this connection (set via
1272
+ * set_system_prompt), null = none. Read by the before_agent_start hook
1273
+ * and layered onto the built-in surface prompt on the NEXT turn — see
1274
+ * installSystemPromptHook/injectSystemPrompt. */
1275
+ private systemPrompt: string | null = null;
1276
+ /** Requested codingTools state (set via set_coding_tools), read fresh by
1277
+ * createRuntimeFactory every time the SDK invokes it. Starts false — the
1278
+ * toggle in settings defaults OFF, and a connection that never sends
1279
+ * set_coding_tools must never get bash/edit/write. */
1280
+ private codingToolsRequested = false;
1281
+ /** What the CURRENT (already-built) session's tool surface actually
1282
+ * contains — set only inside createRuntimeFactory, so it always reflects
1283
+ * reality rather than the latest request. Read by the `coding_tools`
1284
+ * ServerMessage and by installApprovalGate (whether GATED_TOOLS is even
1285
+ * reachable this session, informational only — the allowlist itself is
1286
+ * the real enforcement). */
1287
+ private codingToolsActive = false;
1288
+
1289
+ /** Last `context` ClientMessage the browser pushed (route/view/step/
1290
+ * snapshot, plan §6.6) — null before the first one arrives (a non-
1291
+ * browser WS client, or the race before chat.ts's first push). Read by
1292
+ * get_current_app_context's tool and by installAppContextHook's
1293
+ * ambient-line injection. Set wholesale (not merged) on every "context"
1294
+ * ClientMessage — the browser always sends the full current picture. */
1295
+ private currentAppContext: AppUiContext | null = null;
1296
+
1297
+ /** Dedup/retry-budget/turn-cap bookkeeping (see LOOP_HYGIENE above).
1298
+ * Reset at the start of every new agent turn (turn_start with
1299
+ * turnIndex 0) so a dedup hit or failure streak from a PREVIOUS user
1300
+ * message never bleeds into the next one. */
1301
+ private loopHygiene: LoopHygieneState = initialLoopHygieneState();
475
1302
 
476
1303
  /** Per-connection invariants, built once in start() and reused across
477
1304
  * session switches (new chat / resume / fork). */
@@ -484,13 +1311,16 @@ class PiWebSession {
484
1311
  private readonly sessionDir = join(homedir(), ".mlx-bun", "sessions");
485
1312
 
486
1313
  /** callId -> resolve(decision). Pending browser approvals in flight. */
487
- private readonly pendingApprovals = new Map<string, (decision: ApprovalDecision) => void>();
1314
+ private readonly pendingApprovals = new Map<string, (decision: ApprovalDecision, editedArgs?: Record<string, unknown>, alwaysAllow?: boolean) => void>();
488
1315
  /** callId -> timer handle, so we can clear on resolve/dispose. */
489
1316
  private readonly approvalTimers = new Map<string, ReturnType<typeof setTimeout>>();
490
1317
 
491
1318
  constructor(
492
1319
  private readonly ws: PiWs,
493
- private readonly opts: { port: number | (() => number); modelId: string; contextWindow: number; readOnly: boolean; vision: boolean; thinking: boolean },
1320
+ private readonly opts: {
1321
+ port: number | (() => number); modelId: string; contextWindow: number; readOnly: boolean;
1322
+ vision: boolean; audio: boolean; thinking: boolean; genDefaults: ReadyGenDefaults;
1323
+ },
494
1324
  ) {}
495
1325
 
496
1326
  /** Build the provider, resume the most recent chat, and start streaming. */
@@ -522,7 +1352,12 @@ class PiWebSession {
522
1352
  await this.replaceRuntime(SessionManager.create(this.cwd, this.sessionDir));
523
1353
  if (this.disposed) return;
524
1354
 
525
- this.send({ type: "ready", model: this.opts.modelId, vision: this.opts.vision, thinking: this.opts.thinking });
1355
+ this.send({
1356
+ type: "ready", model: this.opts.modelId, vision: this.opts.vision, audio: this.opts.audio,
1357
+ thinking: this.opts.thinking, genDefaults: this.opts.genDefaults,
1358
+ });
1359
+ this.sendCodingToolsState();
1360
+ this.sendToolApprovals();
526
1361
  this.sendHistory();
527
1362
  await this.sendSessions();
528
1363
  }
@@ -535,12 +1370,24 @@ class PiWebSession {
535
1370
 
536
1371
  // Welcome-assistant tool surface: exactly two read-only tools —
537
1372
  // `web_search` (current/external facts) and `read` (a local file the user
538
- // points to). Both auto-allow (no approval round-trip). We deliberately
539
- // do NOT expose web_fetch/weather/bash/edit/write/grep/find/ls: a 1B model
540
- // over-calls a big toolset. With thinking ON (the web chat default) it
541
- // uses these two appropriately; build the web tools (for the web_search
542
- // definition) then restrict the allowlist to the two names below.
1373
+ // points to) — UNLESS the user opted into codingTools (settings toggle,
1374
+ // default OFF), in which case bash/edit/write/grep/find/ls join the
1375
+ // allowlist too. We deliberately do NOT expose web_fetch/weather: a 1B
1376
+ // model over-calls a big toolset. With thinking ON (the web chat
1377
+ // default) it uses these appropriately; build the web tools (for the
1378
+ // web_search definition) then restrict the allowlist below.
1379
+ //
1380
+ // codingTools:false here is NOT the coding-tools gate (that lives in
1381
+ // webChatToolAllowlist below, which is the array pi actually receives
1382
+ // as `tools:` — see the comment at that call site). buildPiAgentSurface's
1383
+ // own codingTools flag only affects surface.tools, which pi-web.ts never
1384
+ // reads; kept false so surface.tools (unused here) doesn't silently
1385
+ // imply a wider allowlist than webChatToolAllowlist actually grants.
543
1386
  const surface = await buildPiAgentSurface({ webTools: true, codingTools: false });
1387
+ // Snapshot what THIS session is actually being built with — read fresh
1388
+ // here (not at set_coding_tools time) so `coding_tools`'s `active`
1389
+ // field is always true reality, never a stale request.
1390
+ this.codingToolsActive = this.codingToolsRequested;
544
1391
  const webPrompt = buildWebChatSystemPrompt(this.opts.readOnly, {
545
1392
  modelId: this.opts.modelId,
546
1393
  downloadingModel: downloadsSnapshot().find(
@@ -548,7 +1395,7 @@ class PiWebSession {
548
1395
  )?.repoId ?? null,
549
1396
  }, { hasTools: WELCOME_TOOLS.length > 0 }) + surface.memoryHint;
550
1397
  if (process.env.MLX_BUN_PI_DEBUG) {
551
- console.error(`[pi-web] prompt ${WEB_CHAT_PROMPT_VERSION} sha=${webChatPromptFingerprint(webPrompt)} memory=${surface.memoryEnabled ? "on" : "off"}`);
1398
+ console.error(`[pi-web] prompt ${WEB_CHAT_PROMPT_VERSION} sha=${webChatPromptFingerprint(webPrompt)} memory=${surface.memoryEnabled ? "on" : "off"} codingTools=${this.codingToolsActive ? "on" : "off"}`);
552
1399
  }
553
1400
  const services = await createAgentSessionServices({
554
1401
  cwd,
@@ -564,8 +1411,14 @@ class PiWebSession {
564
1411
  additionalSkillPaths: surface.skillPaths,
565
1412
  systemPrompt: webPrompt,
566
1413
  extensionFactories: [
1414
+ // Loop hygiene runs FIRST: a deduped/turn-capped call is
1415
+ // blocked before the approval gate ever sees it (see
1416
+ // installLoopHygieneHooks's doc comment).
1417
+ (pi) => this.installLoopHygieneHooks(pi),
567
1418
  (pi) => this.installApprovalGate(pi),
568
1419
  (pi) => this.installAdapterHook(pi),
1420
+ (pi) => this.installSystemPromptHook(pi),
1421
+ (pi) => this.installAppContextHook(pi),
569
1422
  ],
570
1423
  },
571
1424
  });
@@ -581,9 +1434,18 @@ class PiWebSession {
581
1434
  // exposes only names in this list, so web_fetch/weather stay defined
582
1435
  // but hidden. Memory/reference tools ride along when the surface has
583
1436
  // memory enabled — the prompt + memory skill advertise them, so they
584
- // must be callable (webChatToolAllowlist).
585
- tools: webChatToolAllowlist(surface.memoryEnabled),
586
- customTools: surface.customTools,
1437
+ // must be callable (webChatToolAllowlist). codingTools is further
1438
+ // ANDed with !readOnly here: a read-only server (--read-only) never
1439
+ // advertises bash/edit/write even if the browser opted in, since
1440
+ // installApprovalGate would deny every one of them anyway — no
1441
+ // point offering (and prompting the model to attempt) a tool that
1442
+ // can only ever fail.
1443
+ tools: webChatToolAllowlist(surface.memoryEnabled, this.codingToolsActive && !this.opts.readOnly),
1444
+ // App-aware assistant tools (plan §6.6) are PER-CONNECTION (close
1445
+ // over this.currentAppContext + this.send), unlike surface.customTools
1446
+ // which is a stateless singleton shared across every session — see
1447
+ // createAppAwareTools' doc comment.
1448
+ customTools: [...surface.customTools, ...createAppAwareTools(() => this.currentAppContext, (m) => this.send(m))],
587
1449
  })),
588
1450
  services,
589
1451
  diagnostics: services.diagnostics,
@@ -641,10 +1503,35 @@ class PiWebSession {
641
1503
  if (previous) await previous.dispose();
642
1504
  }
643
1505
 
644
- /** Replay the active session's transcript to the browser (rebuilds thread). */
1506
+ /** Replay the active session's transcript to the browser (rebuilds thread).
1507
+ * Always followed by sendSiblings(): the sibling group for "the last user
1508
+ * message" is a derived view of the same entries and goes stale under
1509
+ * exactly the same conditions history does (new session, switch, fork). */
645
1510
  private sendHistory(): void {
646
1511
  const entries = this.sessionManager?.getEntries() ?? [];
647
1512
  this.send({ type: "history", items: serializeHistory(entries) });
1513
+ this.sendSiblings();
1514
+ }
1515
+
1516
+ /** Push the sibling-group info for the current last-user-message (drives
1517
+ * the browser's `< i/n >` edit toggle). Called by sendHistory() and after
1518
+ * each completed turn, since either can change which message is "last" or
1519
+ * how many siblings it has. */
1520
+ private sendSiblings(): void {
1521
+ const entries = this.sessionManager?.getEntries() ?? [];
1522
+ const last = findLastUserMessageEntry(entries);
1523
+ if (!last) {
1524
+ this.send({ type: "siblings", index: 0, count: 0, siblingIds: [] });
1525
+ return;
1526
+ }
1527
+ const group = userMessageSiblings(entries, last.id);
1528
+ this.send({
1529
+ type: "siblings",
1530
+ entryId: last.id,
1531
+ index: group?.index ?? 1,
1532
+ count: group?.siblingIds.length ?? 1,
1533
+ siblingIds: group?.siblingIds ?? [last.id],
1534
+ });
648
1535
  }
649
1536
 
650
1537
  /** Push current context-window usage to the browser (for the indicator). */
@@ -677,18 +1564,24 @@ class PiWebSession {
677
1564
  return p === root || p.startsWith(root + "/");
678
1565
  }
679
1566
 
680
- /** Start a fresh chat (old one stays on disk, resumable from the sidebar). */
1567
+ /** Start a fresh chat (old one stays on disk, resumable from the sidebar).
1568
+ * Re-invokes createRuntimeFactory (either directly via replaceRuntime, or
1569
+ * through AgentSessionRuntime.newSession()'s own internal call to the
1570
+ * SAME factory) — so this is the earliest point a codingTools toggle
1571
+ * flipped mid-session actually takes effect (see set_coding_tools). */
681
1572
  private async newSession(): Promise<void> {
682
1573
  if (!this.runtime) {
683
1574
  await this.replaceRuntime(SessionManager.create(this.cwd, this.sessionDir));
684
1575
  } else {
685
1576
  await this.runtime.newSession();
686
1577
  }
1578
+ this.sendCodingToolsState();
687
1579
  this.sendHistory();
688
1580
  await this.sendSessions();
689
1581
  }
690
1582
 
691
- /** Resume an existing chat by file path. */
1583
+ /** Resume an existing chat by file path. Same "factory re-runs, codingTools
1584
+ * becomes active if pending" note as newSession above. */
692
1585
  private async openSession(path: string): Promise<void> {
693
1586
  if (!this.isUnderSessionDir(path)) {
694
1587
  this.send({ type: "error", message: "invalid session path" });
@@ -696,6 +1589,7 @@ class PiWebSession {
696
1589
  }
697
1590
  if (!this.runtime) await this.replaceRuntime(SessionManager.open(path, this.sessionDir));
698
1591
  else await this.runtime.switchSession(path);
1592
+ this.sendCodingToolsState();
699
1593
  this.sendHistory();
700
1594
  await this.sendSessions();
701
1595
  }
@@ -709,6 +1603,7 @@ class PiWebSession {
709
1603
  // File-level fork is not a runtime primitive, so create the target
710
1604
  // SessionManager then replace the runtime through the same SDK factory.
711
1605
  await this.replaceRuntime(SessionManager.forkFrom(path, this.cwd, this.sessionDir));
1606
+ this.sendCodingToolsState();
712
1607
  this.sendHistory();
713
1608
  await this.sendSessions();
714
1609
  }
@@ -730,9 +1625,118 @@ class PiWebSession {
730
1625
  else void this.sendSessions();
731
1626
  }
732
1627
 
1628
+ /** Clear an armed one-shot next_turn sampling override after a prompt()
1629
+ * call throws before ever reaching before_provider_request (the hook
1630
+ * that normally consumes it via consumeForRequest — see the
1631
+ * "Per-message sampling scope" comment above composeSampling). Reuses
1632
+ * that same pure helper so the clearing logic has exactly one
1633
+ * implementation; discards its `effective` composition since there's no
1634
+ * outgoing request to apply it to. A no-op when nothing was armed. */
1635
+ private clearArmedOneShotOnFailedPrompt(): void {
1636
+ this.samplingScope = consumeForRequest(this.samplingScope).nextState;
1637
+ }
1638
+
1639
+ /** Move the session leaf to `entryId`'s parent via navigateTree, then
1640
+ * re-prompt with {text, images}. This is the shared mechanism behind
1641
+ * regenerate and edit_resend: navigateTree does NOT delete the old
1642
+ * message/reply (they stay on disk as a sibling branch), and pi resyncs
1643
+ * agent.state.messages from the new leaf internally, so the next
1644
+ * session.prompt() call appends cleanly. Returns false (and sends an
1645
+ * error frame) if the entry can't be found or a turn is already
1646
+ * streaming (mirrors prompt()'s own idle precondition — regenerate/edit
1647
+ * are not queueable follow-ups, unlike a normal composer submit). */
1648
+ private async resendFrom(entryId: string, text: string, images: ImageContent[]): Promise<boolean> {
1649
+ const session = this.session;
1650
+ if (!session) {
1651
+ this.send({ type: "error", message: "session not ready" });
1652
+ return false;
1653
+ }
1654
+ if (session.isStreaming) {
1655
+ this.send({ type: "error", message: "wait for the current reply to finish first" });
1656
+ return false;
1657
+ }
1658
+ if (!text.trim()) {
1659
+ this.send({ type: "error", message: "message text can't be empty" });
1660
+ return false;
1661
+ }
1662
+ await session.navigateTree(entryId);
1663
+ try {
1664
+ await session.prompt(text, images.length ? { images } : {});
1665
+ } catch (err) {
1666
+ // Same pre-flight-failure guard as the "prompt" case above: prompt()
1667
+ // can throw before before_provider_request ever fires (no model
1668
+ // selected, auth failure), which is the only place a one-shot
1669
+ // next_turn sampling override normally gets consumed. Without this,
1670
+ // an override armed right before a failed regenerate/edit-resend
1671
+ // would silently ride along on a later, unrelated successful turn.
1672
+ this.clearArmedOneShotOnFailedPrompt();
1673
+ throw err;
1674
+ }
1675
+ return true;
1676
+ }
1677
+
1678
+ /** Regenerate: re-send the last user message's ORIGINAL content verbatim
1679
+ * as a new sibling branch. Works after abort/stop too — an aborted turn
1680
+ * still leaves the triggering user message entry in place to find. */
1681
+ private async regenerate(): Promise<void> {
1682
+ const entries = this.sessionManager?.getEntries() ?? [];
1683
+ const last = findLastUserMessageEntry(entries);
1684
+ if (!last) {
1685
+ this.send({ type: "error", message: "no message to regenerate" });
1686
+ return;
1687
+ }
1688
+ await this.resendFrom(last.id, last.text, last.images);
1689
+ }
1690
+
1691
+ /** Edit-and-resend: re-send the LAST user message with edited text (plan
1692
+ * §5.2 scopes editing to the last message only — the `< i/n >` toggle is
1693
+ * a linear sibling list, not a tree view). Original images are preserved
1694
+ * since editing text shouldn't silently drop an attachment. */
1695
+ private async editResend(text: string): Promise<void> {
1696
+ const entries = this.sessionManager?.getEntries() ?? [];
1697
+ const last = findLastUserMessageEntry(entries);
1698
+ if (!last) {
1699
+ this.send({ type: "error", message: "no message to edit" });
1700
+ return;
1701
+ }
1702
+ await this.resendFrom(last.id, text, last.images);
1703
+ }
1704
+
1705
+ /** Switch the active leaf to a previously-created sibling (an earlier
1706
+ * edit/regeneration of the last user message) WITHOUT sending anything
1707
+ * new — pure navigation. Walks down to that sibling's own deepest
1708
+ * descendant (its assistant reply, if any) so the replayed history shows
1709
+ * that branch's actual conversation, not just the bare user message. */
1710
+ private switchSibling(entryId: string): void {
1711
+ const sm = this.sessionManager;
1712
+ const session = this.session;
1713
+ if (!sm || !session) {
1714
+ this.send({ type: "error", message: "session not ready" });
1715
+ return;
1716
+ }
1717
+ if (session.isStreaming) {
1718
+ this.send({ type: "error", message: "wait for the current reply to finish first" });
1719
+ return;
1720
+ }
1721
+ const entry = sm.getEntry(entryId);
1722
+ if (!entry) {
1723
+ this.send({ type: "error", message: "message not found" });
1724
+ return;
1725
+ }
1726
+ const leaf = deepestLeafFrom(entryId, (parentId) => sm.getChildren(parentId));
1727
+ sm.branch(leaf);
1728
+ // Mirror what AgentSession.navigateTree() does internally after moving
1729
+ // the leaf: resync agent.state.messages from the new leaf's path so the
1730
+ // in-memory agent state and the persisted leaf pointer never diverge.
1731
+ session.agent.state.messages = sm.buildSessionContext().messages;
1732
+ this.sendHistory();
1733
+ this.sendContextUsage();
1734
+ }
1735
+
733
1736
  /** Register the before_provider_request hook that injects the selected LoRA
734
- * adapter into every provider request (Pi-native adapter control, mirrors the
735
- * CLI extension). Default none = no injection (base model). */
1737
+ * adapter and the effective sampling overrides into every provider
1738
+ * request (Pi-native adapter control, mirrors the CLI extension).
1739
+ * Default none = no injection (base model). */
736
1740
  private installAdapterHook(pi: ExtensionAPI): void {
737
1741
  pi.on("before_provider_request", (event) => {
738
1742
  let payload = event.payload as Record<string, unknown>;
@@ -740,12 +1744,111 @@ class PiWebSession {
740
1744
  // Each returns undefined when it has nothing to change, so we keep the
741
1745
  // prior payload in that case.
742
1746
  payload = injectAdapter(payload, this.selectedAdapter) ?? payload;
743
- payload = injectSampling(payload, this.sampling) ?? payload;
1747
+ // Compose next_turn OVER session, then CONSUME the one-shot override
1748
+ // immediately — this is the very first outgoing provider request
1749
+ // since it was armed (before_provider_request fires once per model
1750
+ // call), so clearing it here guarantees it rides on exactly one
1751
+ // request and never leaks into a second tool-loop turn of the same
1752
+ // prompt, let alone the next user message.
1753
+ const { effective, nextState } = consumeForRequest(this.samplingScope);
1754
+ this.samplingScope = nextState;
1755
+ payload = injectSampling(payload, effective) ?? payload;
744
1756
  return payload;
745
1757
  });
746
1758
  }
747
1759
 
748
- /** Register the tool_call approval gate on the inline extension. */
1760
+ /** Register the before_agent_start hook that layers the user's custom
1761
+ * system prompt onto the built-in surface prompt. This fires fresh inside
1762
+ * EVERY session.prompt() call (pi passes the freshly-built base prompt
1763
+ * each time, see agent-session.js's emitBeforeAgentStart call site), so a
1764
+ * set_system_prompt sent mid-session takes effect starting with the very
1765
+ * next turn — there is no "only at session creation" limitation here, and
1766
+ * the UI copy should say so plainly (it applies going forward, not
1767
+ * retroactively to already-sent turns). Returning undefined leaves pi's
1768
+ * base prompt (identity + memoryHint) untouched. */
1769
+ private installSystemPromptHook(pi: ExtensionAPI): void {
1770
+ pi.on("before_agent_start", (event) => {
1771
+ const merged = injectSystemPrompt(event.systemPrompt, this.systemPrompt);
1772
+ return merged !== undefined ? { systemPrompt: merged } : undefined;
1773
+ });
1774
+ }
1775
+
1776
+ /** Register the before_agent_start hook that auto-prepends the compact
1777
+ * one-line ambient context (plan §6.6 "never answer blind") to every
1778
+ * turn's system prompt — the SAME layering mechanism
1779
+ * installSystemPromptHook uses for the user's custom prompt (pi's
1780
+ * BeforeAgentStartEventResult can only replace systemPrompt, chained
1781
+ * across multiple extensions — there's no channel to rewrite the raw
1782
+ * user prompt text itself). Fires fresh on every session.prompt() call,
1783
+ * so it always reflects the LATEST context push, never a stale one from
1784
+ * session-creation time. Returns undefined (no-op) when nothing has
1785
+ * been pushed yet — a non-browser WS client, or the race before chat.ts's
1786
+ * first `context` frame lands, must never inject a made-up location. */
1787
+ private installAppContextHook(pi: ExtensionAPI): void {
1788
+ pi.on("before_agent_start", (event) => {
1789
+ const line = ambientContextLine(this.currentAppContext);
1790
+ if (!line) return undefined;
1791
+ return { systemPrompt: `${event.systemPrompt}\n\n${line}` };
1792
+ });
1793
+ }
1794
+
1795
+ /**
1796
+ * Register the tool_call approval gate on the inline extension.
1797
+ *
1798
+ * Approval-fatigue defaults (web-chat-redesign.md §5.4 matrix — verified
1799
+ * here, not just asserted): READ_ONLY_TOOLS NEVER prompt, checked first
1800
+ * and unconditionally — no config, always-allow list, or codingTools
1801
+ * state can make a read-only tool gated, and no combination of state can
1802
+ * make a GATED_TOOLS member skip the card except the durable
1803
+ * tool-approvals config (isToolAlwaysAllowed) or read-only-mode's
1804
+ * outright denial. This keeps the fatigue budget spent only on tools
1805
+ * that actually mutate the machine.
1806
+ */
1807
+ /**
1808
+ * Register loop-hygiene bookkeeping: turnIndex tracking (turn_start),
1809
+ * the before-execution dedup/turn-cap block (tool_call), and the
1810
+ * after-execution failure-streak nudge (tool_result). Registered FIRST
1811
+ * among the tool_call handlers (see extensionFactories order in
1812
+ * createRuntimeFactory) so a deduped or turn-capped call never even
1813
+ * reaches the approval gate — no point prompting the user to approve
1814
+ * a call we're about to skip. See the LOOP_HYGIENE block comment above
1815
+ * for what the SDK already does vs. what this closes.
1816
+ */
1817
+ private installLoopHygieneHooks(pi: ExtensionAPI): void {
1818
+ pi.on("turn_start", (event) => {
1819
+ this.loopHygiene.turnIndex = event.turnIndex;
1820
+ // A fresh user prompt always restarts at turnIndex 0 (AgentSession
1821
+ // resets _turnIndex on agent_start, before the first turn_start) —
1822
+ // treat that as "new tool loop" and clear dedup/failure state so a
1823
+ // previous message's history never suppresses/nudges this one.
1824
+ if (event.turnIndex === 0) this.loopHygiene = initialLoopHygieneState();
1825
+ });
1826
+
1827
+ pi.on("tool_call", (event: ToolCallEvent) => {
1828
+ const block = decideBeforeToolCall(this.loopHygiene, event.toolName, event.input);
1829
+ if (block && process.env.MLX_BUN_PI_DEBUG) {
1830
+ console.error(`[pi-web] loop-hygiene blocked ${event.toolName}: ${block.reason.split("\n")[0]}`);
1831
+ }
1832
+ return block ? { block: true, reason: block.reason } : undefined;
1833
+ });
1834
+
1835
+ pi.on("tool_result", (event) => {
1836
+ const nudge = recordToolCallOutcome(this.loopHygiene, event.toolName, event.input, {
1837
+ isError: event.isError,
1838
+ resultText: toolResultText(event.content),
1839
+ });
1840
+ if (!nudge) return undefined;
1841
+ // Append the course-correction nudge to the existing error content
1842
+ // rather than replacing it, so the model still sees the real
1843
+ // failure reason plus the "stop repeating this" instruction.
1844
+ const existingText = toolResultText(event.content);
1845
+ return {
1846
+ content: [{ type: "text", text: existingText + nudge }],
1847
+ isError: event.isError,
1848
+ };
1849
+ });
1850
+ }
1851
+
749
1852
  private installApprovalGate(pi: ExtensionAPI): void {
750
1853
  pi.on("tool_call", async (event: ToolCallEvent) => {
751
1854
  const tool = event.toolName;
@@ -756,7 +1859,9 @@ class PiWebSession {
756
1859
  // Read-only tools never need approval.
757
1860
  if (READ_ONLY_TOOLS.has(tool)) return undefined;
758
1861
 
759
- // Anything mutating is denied outright in read-only mode.
1862
+ // Anything mutating is denied outright in read-only mode — before the
1863
+ // durable always-allow list is even consulted, since --read-only is a
1864
+ // server operator's hard constraint, not a per-tool user preference.
760
1865
  if (this.opts.readOnly && GATED_TOOLS.has(tool)) {
761
1866
  return { block: true, reason: "Read-only session: mutating tools are disabled." };
762
1867
  }
@@ -765,10 +1870,29 @@ class PiWebSession {
765
1870
  // surface, but be safe): allow.
766
1871
  if (!GATED_TOOLS.has(tool)) return undefined;
767
1872
 
1873
+ // Durable "always allow this tool" (risk #6): a prior approval card on
1874
+ // THIS machine checked the box, persisting to ~/.mlx-bun/tool-
1875
+ // approvals.json (src/tool-approvals.ts). Skip the round-trip entirely
1876
+ // — no card, no wait, matching every session and every browser tab.
1877
+ if (isToolAlwaysAllowed(tool)) return undefined;
1878
+
768
1879
  const decision = await this.requestApproval(event);
769
- if (decision === "deny") {
1880
+ if (decision.decision === "deny") {
770
1881
  return { block: true, reason: "Denied by user." };
771
1882
  }
1883
+ // Editable arguments (LM Studio's pattern): the SDK's documented
1884
+ // mutation contract is "mutate event.input in place before returning"
1885
+ // — there is no return-value channel for modified args (ToolCallEventResult
1886
+ // is `{block?, reason?}` only). event.input is asserted as a mutable
1887
+ // Record here since ToolCallEvent's per-tool variants type it as their
1888
+ // specific input shape, but the mutation contract is untyped-record at
1889
+ // the wire level (this mirrors how pi's own bundled permission-gate
1890
+ // example extension does it).
1891
+ applyEditedArgs(event.input as Record<string, unknown>, decision.editedArgs);
1892
+ if (decision.alwaysAllow) {
1893
+ setToolAlwaysAllowed(tool);
1894
+ this.sendToolApprovals();
1895
+ }
772
1896
  return undefined; // allow
773
1897
  });
774
1898
  }
@@ -777,15 +1901,15 @@ class PiWebSession {
777
1901
  * Ask the browser to approve a gated tool call and await its decision.
778
1902
  * Auto-denies after APPROVAL_TIMEOUT_MS or if the connection drops.
779
1903
  */
780
- private requestApproval(event: ToolCallEvent): Promise<ApprovalDecision> {
1904
+ private requestApproval(event: ToolCallEvent): Promise<{ decision: ApprovalDecision; editedArgs?: Record<string, unknown>; alwaysAllow?: boolean }> {
781
1905
  const callId = event.toolCallId;
782
- return new Promise<ApprovalDecision>((resolve) => {
783
- const settle = (decision: ApprovalDecision) => {
1906
+ return new Promise((resolve) => {
1907
+ const settle = (decision: ApprovalDecision, editedArgs?: Record<string, unknown>, alwaysAllow?: boolean) => {
784
1908
  const timer = this.approvalTimers.get(callId);
785
1909
  if (timer) clearTimeout(timer);
786
1910
  this.approvalTimers.delete(callId);
787
1911
  this.pendingApprovals.delete(callId);
788
- resolve(decision);
1912
+ resolve({ decision, editedArgs, alwaysAllow });
789
1913
  };
790
1914
 
791
1915
  this.pendingApprovals.set(callId, settle);
@@ -802,9 +1926,22 @@ class PiWebSession {
802
1926
  }
803
1927
 
804
1928
  /** Resolve a pending approval from a browser `approval` frame. */
805
- resolveApproval(callId: string, decision: ApprovalDecision): void {
1929
+ resolveApproval(callId: string, decision: ApprovalDecision, editedArgs?: Record<string, unknown>, alwaysAllow?: boolean): void {
806
1930
  const settle = this.pendingApprovals.get(callId);
807
- if (settle) settle(decision);
1931
+ if (settle) settle(decision, editedArgs, alwaysAllow);
1932
+ }
1933
+
1934
+ /** Push the effective-vs-pending codingTools state (sent on ready and
1935
+ * after every set_coding_tools). */
1936
+ private sendCodingToolsState(): void {
1937
+ this.send({ type: "coding_tools", active: this.codingToolsActive, pending: this.codingToolsRequested });
1938
+ }
1939
+
1940
+ /** Push the full durable always-allow set (sent on ready and after any
1941
+ * change), so the settings panel can list/forget approvals without a
1942
+ * separate REST round-trip. */
1943
+ private sendToolApprovals(): void {
1944
+ this.send({ type: "tool_approvals", alwaysAllow: listAlwaysAllowedTools() });
808
1945
  }
809
1946
 
810
1947
  private onSessionEvent(event: AgentSessionEvent): void {
@@ -822,6 +1959,9 @@ class PiWebSession {
822
1959
  if (event.type === "turn_end") {
823
1960
  void this.sendSessions();
824
1961
  this.sendContextUsage();
1962
+ // A completed turn (including regenerate/edit_resend's re-prompt) can
1963
+ // change the last user message's sibling count/index.
1964
+ this.sendSiblings();
825
1965
  }
826
1966
  }
827
1967
 
@@ -845,6 +1985,27 @@ class PiWebSession {
845
1985
  case "delete_session":
846
1986
  this.deleteSession(msg.path);
847
1987
  return;
1988
+ case "set_coding_tools":
1989
+ // Record the request; doesn't touch the ALREADY-BUILT session's tool
1990
+ // surface (pi bakes tools/customTools in at createAgentSessionServices
1991
+ // time — there is no live "add a tool to this running session" hook).
1992
+ // Applied by createRuntimeFactory the next time the SDK calls it:
1993
+ // new_session, open_session, fork_session, or a reconnect's initial
1994
+ // start(). sendCodingToolsState reports both active (this session)
1995
+ // and pending (the request just recorded) so the UI can be honest
1996
+ // about "applies to your next new chat" rather than implying it's
1997
+ // live now.
1998
+ this.codingToolsRequested = !!msg.enabled;
1999
+ this.sendCodingToolsState();
2000
+ return;
2001
+ case "context":
2002
+ // App-aware assistant (plan §6.6): stored wholesale, no session
2003
+ // required — chat.ts pushes this on route/step change, which can
2004
+ // race the very first `ready` frame on initial connect. Read by
2005
+ // get_current_app_context (live) and installAppContextHook (the
2006
+ // ambient one-liner auto-prepended to the NEXT turn).
2007
+ this.currentAppContext = msg.context;
2008
+ return;
848
2009
  }
849
2010
 
850
2011
  const session = this.session;
@@ -868,15 +2029,28 @@ class PiWebSession {
868
2029
  console.error(`[pi-web] prompt text=${JSON.stringify(msg.text.slice(0, 500))} images=${msg.images?.length ?? 0} streaming=${session.isStreaming} adapter=${this.selectedAdapter ?? "base"}`);
869
2030
  }
870
2031
  const images = toPiImages(msg.images);
871
- if (session.isStreaming) await session.prompt(msg.text, { streamingBehavior: "followUp", images });
872
- else await session.prompt(msg.text, { images });
2032
+ // Guard a one-shot next_turn override against prompt() throwing
2033
+ // BEFORE it ever issues a provider request (no model selected, auth
2034
+ // failure, etc.) — the normal consumer, installAdapterHook's
2035
+ // before_provider_request hook, never fires in that case, so
2036
+ // without this the armed override would silently survive to land
2037
+ // on some later, unrelated successful turn instead of this one.
2038
+ // sendError (the WS layer's outer catch) still reports the error;
2039
+ // this only prevents the override from outliving its intended turn.
2040
+ try {
2041
+ if (session.isStreaming) await session.prompt(msg.text, { streamingBehavior: "followUp", images });
2042
+ else await session.prompt(msg.text, { images });
2043
+ } catch (err) {
2044
+ this.clearArmedOneShotOnFailedPrompt();
2045
+ throw err;
2046
+ }
873
2047
  return;
874
2048
  }
875
2049
  case "abort":
876
2050
  await session.abort();
877
2051
  return;
878
2052
  case "approval":
879
- this.resolveApproval(msg.callId, msg.decision);
2053
+ this.resolveApproval(msg.callId, msg.decision, msg.editedArgs, msg.alwaysAllow);
880
2054
  return;
881
2055
  case "set_thinking":
882
2056
  // Pi clamps to the model's available levels; a no-op for models
@@ -889,16 +2063,48 @@ class PiWebSession {
889
2063
  // app.html has already mounted it server-side (POST /v1/adapters).
890
2064
  this.selectedAdapter = msg.id;
891
2065
  return;
892
- case "set_sampling":
893
- // Record the per-request sampling overrides; the
894
- // before_provider_request hook injects any set fields into the
895
- // outgoing payload. A null/undefined field clears the override so the
896
- // server's mode-aware default applies again.
897
- this.sampling = {
2066
+ case "set_sampling": {
2067
+ // The full mlx_lm.server sampler extension set (min_p/XTC/penalty
2068
+ // families/seed), not just temperature/top_p/top_k
2069
+ // (web-ui-pass-plan.md #8), for either scope.
2070
+ const overrides: SamplingOverrides = {
898
2071
  temperature: msg.temperature ?? null,
899
2072
  top_p: msg.top_p ?? null,
900
2073
  top_k: msg.top_k ?? null,
2074
+ min_p: msg.min_p ?? null,
2075
+ xtc_probability: msg.xtc_probability ?? null,
2076
+ xtc_threshold: msg.xtc_threshold ?? null,
2077
+ repetition_penalty: msg.repetition_penalty ?? null,
2078
+ repetition_context_size: msg.repetition_context_size ?? null,
2079
+ presence_penalty: msg.presence_penalty ?? null,
2080
+ frequency_penalty: msg.frequency_penalty ?? null,
2081
+ seed: msg.seed ?? null,
901
2082
  };
2083
+ // scope "next_turn": one-shot, armed via applySetSampling and
2084
+ // consumed by installAdapterHook's before_provider_request handler
2085
+ // on the very next outgoing provider request — never touches the
2086
+ // session-level value, so a later plain set_sampling (scope
2087
+ // "session" or omitted) still sees whatever was there before this
2088
+ // message. scope "session"/omitted: today's behavior, persists on
2089
+ // the connection until changed; a null/undefined field clears the
2090
+ // override so the server's mode-aware default applies again.
2091
+ this.samplingScope = applySetSampling(this.samplingScope, overrides, msg.scope);
2092
+ return;
2093
+ }
2094
+ case "set_system_prompt":
2095
+ // Store verbatim (including whitespace-only/empty -> effectively
2096
+ // cleared by injectSystemPrompt's trim+empty check); applied by the
2097
+ // before_agent_start hook starting with the next turn.
2098
+ this.systemPrompt = msg.text;
2099
+ return;
2100
+ case "regenerate":
2101
+ await this.regenerate();
2102
+ return;
2103
+ case "edit_resend":
2104
+ await this.editResend(msg.text);
2105
+ return;
2106
+ case "switch_sibling":
2107
+ this.switchSibling(msg.entryId);
902
2108
  return;
903
2109
  }
904
2110
  }
@@ -934,6 +2140,7 @@ class PiWebSession {
934
2140
  * @param opts.contextWindow Context window advertised to pi. Default: 32768.
935
2141
  * Source: ctx.model.config.text.maxPositionEmbeddings.
936
2142
  * @param opts.readOnly When true, bash/edit/write are denied outright.
2143
+ * @param opts.genDefaults Model-author sampling defaults, sent once on `ready`.
937
2144
  */
938
2145
  export function makePiWsHandler(opts: {
939
2146
  port: number | (() => number);
@@ -942,9 +2149,17 @@ export function makePiWsHandler(opts: {
942
2149
  readOnly?: boolean;
943
2150
  /** Whether the loaded model can accept images (server's ctx.vision != null). */
944
2151
  vision?: boolean;
2152
+ /** Whether the loaded model can accept `input_audio` parts (server's audio
2153
+ * tower loaded or loadable). Rides only the `ready` frame — pi's
2154
+ * Model.input type has no "audio" modality to declare. */
2155
+ audio?: boolean;
945
2156
  /** Whether the model has a switchable reasoning channel
946
2157
  * (server's ctx.template.supportsThinking). Drives the thinking toggle. */
947
2158
  thinking?: boolean;
2159
+ /** Model-author sampling defaults (generation_config.json + server-CLI
2160
+ * overrides), sent once on the `ready` frame. Default: all-null (the
2161
+ * sampling popover falls back to its own hardcoded shape). */
2162
+ genDefaults?: ReadyGenDefaults;
948
2163
  }): WebSocketHandler<PiWsData> {
949
2164
  const resolved = {
950
2165
  port: opts.port,
@@ -952,7 +2167,9 @@ export function makePiWsHandler(opts: {
952
2167
  contextWindow: opts.contextWindow ?? DEFAULT_CONTEXT_WINDOW,
953
2168
  readOnly: opts.readOnly ?? false,
954
2169
  vision: opts.vision ?? false,
2170
+ audio: opts.audio ?? false,
955
2171
  thinking: opts.thinking ?? false,
2172
+ genDefaults: opts.genDefaults ?? { temperature: null, topP: null, topK: null },
956
2173
  };
957
2174
 
958
2175
  return {