Sanitize streamed local model output

Co-authored-by: Codex <codex@openai.com>

static/engine.js

@@ -33,6 +33,7 @@ export async function loadModel(onProgress) {
 
 export async function generateCode(prompt, language, onToken, onComplete) {
     if (!generator) throw new Error("Model not loaded");
+    let streamedText = "";
 
     const messages = [
         {
@@ -54,20 +55,23 @@ export async function generateCode(prompt, language, onToken, onComplete) {
 
     const streamer = new TextStreamer(generator.tokenizer, {
         skip_prompt: true,
-        callback_function: onToken,
+        callback_function: (token) => {
+            streamedText += token;
+            onToken(token);
+        },
     });
 
     const result = await generator(messages, {
         max_new_tokens: 1024,
-        temperature: 0.2,
-        do_sample: true,
+        do_sample: false,
         streamer,
     });
 
     const generated = result?.[0]?.generated_text;
-    const fullCodeRaw = Array.isArray(generated)
+    const resultText = Array.isArray(generated)
         ? generated.at(-1).content
         : String(generated || "");
+    const fullCodeRaw = resultText && resultText.trim() ? resultText : streamedText;
     const fullCode = stripMarkdownCodeFence(fullCodeRaw);
     onComplete(fullCode);
     return fullCode;

static/ui.js

@@ -7,7 +7,7 @@ export function appendToken(token) {
 
     const display = document.getElementById("stream-display");
     if (display) {
-        display.textContent = streamBuffer;
+        display.textContent = cleanVisibleStream(streamBuffer);
         display.scrollTop = display.scrollHeight;
     }
 }
@@ -67,7 +67,7 @@ export async function rollbackAndReplace(correctedCode, reason, verdict = "FIX")
 }
 
 export function getCurrentCode() {
-    return streamBuffer;
+    return cleanVisibleStream(streamBuffer);
 }
 
 export function setCode(code) {
@@ -96,6 +96,13 @@ function sleep(ms) {
     return new Promise((resolve) => setTimeout(resolve, ms));
 }
 
+function cleanVisibleStream(code) {
+    if (typeof window.stripMarkdownCodeFence === "function") {
+        return window.stripMarkdownCodeFence(code);
+    }
+    return code;
+}
+
 Object.assign(window, {
     appendToken,
     setStatus,