diff --git a/Cotabby/Support/SuggestionTextNormalizer.swift b/Cotabby/Support/SuggestionTextNormalizer.swift index 01bef67..b919c57 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 b5566b0..c5aa4e1 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") + } }