diff --git a/index.ts b/index.ts index 02216b33..b53a0994 100644 --- a/index.ts +++ b/index.ts @@ -69,7 +69,7 @@ const server: Plugin = (async (ctx) => { prompts, hostPermissions, ) as any, - "experimental.text.complete": createTextCompleteHandler(), + "experimental.text.complete": createTextCompleteHandler(state), "command.execute.before": createCommandExecuteHandler( ctx.client, state, diff --git a/lib/hooks.ts b/lib/hooks.ts index 6d6f3862..304d2430 100644 --- a/lib/hooks.ts +++ b/lib/hooks.ts @@ -289,12 +289,15 @@ export function createCommandExecuteHandler( } } -export function createTextCompleteHandler() { +export function createTextCompleteHandler(state: SessionState) { return async ( _input: { sessionID: string; messageID: string; partID: string }, output: { text: string }, ) => { - output.text = stripHallucinationsFromString(output.text) + // The session format decides which tags are recognised: compact sessions + // emit @7@ refs, and cleaning with the default XML rules leaves them in + // the rendered reply (#632). + output.text = stripHallucinationsFromString(output.text, state.idFormat) } } diff --git a/lib/messages/utils.ts b/lib/messages/utils.ts index 004de9a5..715ac854 100644 --- a/lib/messages/utils.ts +++ b/lib/messages/utils.ts @@ -11,6 +11,15 @@ const DCP_UNPAIRED_TAG_REGEX = /<\/?dcp[^>]*>/gi const INJECTED_MESSAGE_ID_SUFFIX_REGEX = /(?<=\n)]*>m\d+<\/dcp-message-id>\s*$/ const HALLUCINATED_PARAMETER_SUFFIX_REGEX = /(?<=\n)m\d+<\/parameter>\s*$/ const COMPACT_TAG_REGEX = /@(?:[1-9]\d*|b[1-9]\d*|blocked)@(?:[ \t]+\[(?:low|medium|high)\])?/gi +// Models that imitate the injected tag often emit broken fragments instead of a +// well formed one: a run of trailing lines made only of bracket debris such as +// "<<", "<" or "?[ \t]*` +const DEGENERATE_TAG_DEBRIS_REGEX = new RegExp( + String.raw`(?:\n${DEGENERATE_TAG_LINE_PATTERN}){2,}\n?$`, +) const generateStableId = (prefix: string, seed: string): string => { const hash = createHash("sha256").update(seed).digest("hex").slice(0, SUMMARY_ID_HASH_LENGTH) @@ -175,6 +184,7 @@ export const stripHallucinationsFromString = (text: string, format: IdFormat = " return withoutKnownSuffixes .replace(DCP_PAIRED_TAG_REGEX, "") .replace(DCP_UNPAIRED_TAG_REGEX, "") + .replace(DEGENERATE_TAG_DEBRIS_REGEX, "") } export const stripHallucinations = (messages: WithParts[], format: IdFormat = "xml"): void => { diff --git a/tests/hooks-permission.test.ts b/tests/hooks-permission.test.ts index 58d26c83..6aa0b4a8 100644 --- a/tests/hooks-permission.test.ts +++ b/tests/hooks-permission.test.ts @@ -269,7 +269,7 @@ test("command execute exits after effective permission resolves to deny", async test("text complete strips hallucinated metadata tags", async () => { const output = { text: "alpha beta omega" } - const handler = createTextCompleteHandler() + const handler = createTextCompleteHandler(createSessionState()) await handler({ sessionID: "session-1", messageID: "message-1", partID: "part-1" }, output) diff --git a/tests/message-priority.test.ts b/tests/message-priority.test.ts index d6105eae..de0eff44 100644 --- a/tests/message-priority.test.ts +++ b/tests/message-priority.test.ts @@ -765,7 +765,7 @@ test("hallucination stripping removes all dcp-prefixed XML tags including varian assert.equal(stripHallucinationsFromString(text), "alphaomega") - const handler = createTextCompleteHandler() + const handler = createTextCompleteHandler(createSessionState()) const output = { text } await handler({ sessionID: "session", messageID: "message", partID: "part" }, output) assert.equal(output.text, "alphaomega") diff --git a/tests/text-cleanup.test.ts b/tests/text-cleanup.test.ts new file mode 100644 index 00000000..c8e9f55d --- /dev/null +++ b/tests/text-cleanup.test.ts @@ -0,0 +1,87 @@ +import assert from "node:assert/strict" +import test from "node:test" +import { mkdtempSync, rmSync } from "node:fs" +import { join } from "node:path" +import { tmpdir } from "node:os" + +const root = mkdtempSync(join(tmpdir(), "dcp-text-cleanup-")) +process.env.XDG_DATA_HOME = join(root, "data") +process.env.XDG_CONFIG_HOME = join(root, "config") +process.env.OPENCODE_CONFIG_DIR = join(root, "config", "opencode") +test.after(() => rmSync(root, { recursive: true, force: true })) + +const { createSessionState } = await import("../lib/state") +const { stripHallucinationsFromString } = await import("../lib/messages/utils") +const { createTextCompleteHandler } = await import("../lib/hooks") + +// Verbatim from #562: the reply a model produced after being shown the +// tag, with every fragment of that tag left unclosed. +const KIMI_REPLY = [ + "Which do you prefer?", + "<<", + "<<", + "<<", + "<<", + "<<", + "<<", + "<<", + "<", + " { + assert.equal(stripHallucinationsFromString(KIMI_REPLY), "Which do you prefer?") +}) + +test("a single stray angle bracket is left alone", () => { + const code = "if (ad) {\n return 1\n}" + assert.equal(stripHallucinationsFromString(code), code) + + const prose = "Compare a < b" + assert.equal(stripHallucinationsFromString(prose), prose) + + const oneLine = "The tag looks like { + // Cleanup is deliberately conservative: only a trailing run is removed, so + // bracket-only lines in the middle of a reply are never rewritten. + const text = ["intro", "<<", " { + const text = ["```ts", "<<", "<", "```", "done"].join("\n") + assert.equal(stripHallucinationsFromString(text), text) +}) + +test("well formed injected tags are still stripped", () => { + assert.equal( + stripHallucinationsFromString("answer\nm0477"), + "answer\n", + ) + assert.equal(stripHallucinationsFromString("m0477"), "m0477") +}) + +test("text.complete strips echoed refs for compact sessions", async () => { + const state = createSessionState("compact") + const handler = createTextCompleteHandler(state) + + // Verbatim from #632: a model copying the compact reference tag into its + // reply. The handler used the XML cleanup rules, which do not match @N@. + const output = { text: "I will use @7@ [high] to compress.\n@b1@ done\n@blocked@" } + await handler({ sessionID: "ses", messageID: "m", partID: "p" }, output) + assert.equal(output.text, "I will use to compress.\n done\n") +}) + +test("text.complete leaves ordinary compact text untouched", async () => { + const state = createSessionState("compact") + const handler = createTextCompleteHandler(state) + const output = { text: "mail me at a@b.com, see for the patch" } + await handler({ sessionID: "ses", messageID: "m", partID: "p" }, output) + assert.equal(output.text, "mail me at a@b.com, see for the patch") +})