From e121b620a9fdfc07ba75b1415a64d54a2cfa40de Mon Sep 17 00:00:00 2001 From: Jacob Fu <141651335+FuJacob@users.noreply.github.com> Date: Sat, 30 May 2026 15:47:48 -0700 Subject: [PATCH] Strip hallucinated prompt-scaffolding labels from suggestions Small instruct models on the local llama path sometimes parrot the prompt's section headers ("App:", "Text before caret:", "Continuation:") as the first thing they generate, sometimes inventing labels that were never in the prompt. Those labels leaked into the ghost text. This adds a leading-label backstop to SuggestionTextNormalizer that peels a run of known scaffolding labels (each on its own line or inline) before the single-line collapse, so the real continuation surfaces. Matching is anchored to a known label set, so legitimate colons ("TODO: buy milk") and mid-text labels are preserved. This is a best-effort catch, not the root-cause fix. The durable fix is feeding instruct models their own chat template so prompt instructions never read as content; that lands separately via the CotabbyInference chat-template work. --- .../Support/SuggestionTextNormalizer.swift | 55 +++++++++++++ .../SuggestionTextNormalizerTests.swift | 82 +++++++++++++++++++ 2 files changed, 137 insertions(+) diff --git a/Cotabby/Support/SuggestionTextNormalizer.swift b/Cotabby/Support/SuggestionTextNormalizer.swift index 01bef67d..b919c57d 100644 --- a/Cotabby/Support/SuggestionTextNormalizer.swift +++ b/Cotabby/Support/SuggestionTextNormalizer.swift @@ -53,6 +53,18 @@ enum SuggestionTextNormalizer { // continuation that followed. normalized = normalized.trimmingCharacters(in: .newlines) + // Backstop for prompt-scaffolding hallucination. Small instruct models sometimes parrot the + // prompt's section headers ("App:", "Text before caret:", "Continuation:") as the first + // thing they emit: sometimes as their own line, sometimes inline before the real text, and + // sometimes as labels the model invents that were never in our prompt at all. None of these + // are valid ghost text. Stripping a leading run of known labels runs before the single-line + // collapse so a model that stacks "Task:\nText before caret:\nreal continuation" still + // surfaces the real continuation instead of collapsing to the first label line. This is a + // best-effort catch, not the fix: the durable fix is feeding instruct models their own chat + // template so instructions never read as content in the first place. + normalized = stripLeadingScaffoldingLabels(normalized) + normalized = normalized.trimmingCharacters(in: .newlines) + if request.isMultiLineEnabled { // Multi-line mode: keep content up to the first blank-line boundary (double newline) // to prevent runaway paragraph generation while still allowing multi-line completions. @@ -144,4 +156,47 @@ enum SuggestionTextNormalizer { let afterLastEchoed = lastEchoedWord.endIndex return String(suggestion[afterLastEchoed...]) } + + /// Section-header labels Cotabby's prompts use, plus close variants small models tend to + /// hallucinate. Matching is anchored to this known set so legitimate user text that merely + /// contains a colon ("Note: buy milk", "TODO: ship it") is never treated as scaffolding. + /// Ordered longest-first at match time so "Text before the caret:" wins over "Text before". + private static let scaffoldingLabels: [String] = [ + "Text before the caret:", + "Text before caret:", + "Text after the caret:", + "Text after caret:", + "User Profile Context:", + "Your style preferences:", + "Final instruction:", + "Screen context:", + "Screen content:", + "User's clipboard:", + "Continuation:", + "Application:", + "Task:", + "App:" + ] + + /// Removes a leading run of known prompt-scaffolding labels (see `scaffoldingLabels`), whether + /// each sits on its own line or inline before the continuation. Only labels at the very start + /// are stripped; a label appearing later in the text is left alone because by then it is far + /// more likely to be real user content than echoed scaffolding. + private static func stripLeadingScaffoldingLabels(_ text: String) -> String { + let labelsByLengthDescending = scaffoldingLabels.sorted { $0.count > $1.count } + var working = text + + while true { + // Look past leading whitespace/newlines to find the first real token. We only commit to + // dropping that whitespace if a label actually matches; otherwise `working` is returned + // untouched so the caller's existing leading-space handling still sees the original. + let leading = String(working.drop(while: { $0.isWhitespace })) + guard let label = labelsByLengthDescending.first(where: { + leading.range(of: $0, options: [.caseInsensitive, .anchored]) != nil + }) else { + return working + } + working = String(leading.dropFirst(label.count)) + } + } } diff --git a/CotabbyTests/SuggestionTextNormalizerTests.swift b/CotabbyTests/SuggestionTextNormalizerTests.swift index b5566b01..c5aa4e14 100644 --- a/CotabbyTests/SuggestionTextNormalizerTests.swift +++ b/CotabbyTests/SuggestionTextNormalizerTests.swift @@ -134,4 +134,86 @@ final class SuggestionTextNormalizerTests: XCTestCase { XCTAssertEqual(normalized, "") } + + func test_normalize_stripsLeadingInlineScaffoldingLabel() { + // Caret sits right after a space, so the exposed leading space is dropped and the + // continuation surfaces cleanly without the echoed "Text before caret:" header. + let request = CotabbyTestFixtures.suggestionRequest( + prefixText: "I am ", + prompt: "PROMPT", + precedingText: "I am " + ) + + let normalized = SuggestionTextNormalizer.normalize( + "Text before caret: going to the store", + for: request + ) + + XCTAssertEqual(normalized, "going to the store") + } + + func test_normalize_stripsHallucinatedAppLabel() { + let request = CotabbyTestFixtures.suggestionRequest( + prefixText: "send the ", + prompt: "PROMPT", + precedingText: "send the " + ) + + let normalized = SuggestionTextNormalizer.normalize( + "App: report by Friday", + for: request + ) + + XCTAssertEqual(normalized, "report by Friday") + } + + func test_normalize_stripsStackedScaffoldingLabelLines() { + // Stacked labels across newlines must be peeled before the single-line collapse, otherwise + // the collapse would keep only the first label line ("Task:") and the real text would be + // lost. + let request = CotabbyTestFixtures.suggestionRequest( + prefixText: "The ", + prompt: "PROMPT", + precedingText: "The " + ) + + let normalized = SuggestionTextNormalizer.normalize( + "Task:\nText before caret:\nquick brown fox", + for: request + ) + + XCTAssertEqual(normalized, "quick brown fox") + } + + func test_normalize_keepsLegitimateNonLabelColon() { + // A colon that is not a known scaffolding label is real user content and must survive. + let request = CotabbyTestFixtures.suggestionRequest( + prefixText: "my list ", + prompt: "PROMPT", + precedingText: "my list " + ) + + let normalized = SuggestionTextNormalizer.normalize( + "TODO: buy milk", + for: request + ) + + XCTAssertEqual(normalized, "TODO: buy milk") + } + + func test_normalize_keepsLabelLikeTextWhenNotLeading() { + // "Task:" appears mid-continuation, not at the start, so it is real text and stays. + let request = CotabbyTestFixtures.suggestionRequest( + prefixText: "finish the ", + prompt: "PROMPT", + precedingText: "finish the " + ) + + let normalized = SuggestionTextNormalizer.normalize( + "first Task: review", + for: request + ) + + XCTAssertEqual(normalized, "first Task: review") + } }