From 807b9d935d8410186d9044cc762b4572e8459577 Mon Sep 17 00:00:00 2001 From: labkey-nicka Date: Thu, 20 Aug 2026 16:51:46 -0700 Subject: [PATCH 1/3] Revise prompt to remove four backticks --- .../labkey/query/controllers/prompts/ExpressionAssistant.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/query/src/org/labkey/query/controllers/prompts/ExpressionAssistant.md b/query/src/org/labkey/query/controllers/prompts/ExpressionAssistant.md index 1eafede8e9e..1f8be652624 100644 --- a/query/src/org/labkey/query/controllers/prompts/ExpressionAssistant.md +++ b/query/src/org/labkey/query/controllers/prompts/ExpressionAssistant.md @@ -96,7 +96,8 @@ ambiguity materially affects the result. 2. **Validate Silently:** When you produce a SQL expression, you must validate it using the `validateCalculatedColumnExpression` tool. You must not mention this tool to the user. 3. **Format Final Expressions:** When presenting a final SQL expression for the user to apply, you must place the tool's - JSON return value verbatim inside a fenced code block tagged `expression` (e.g., ````expression\n{...}\n````). + JSON return value verbatim inside a fenced code block tagged `expression`. Open the block with exactly three + backticks followed by `expression`, and close it with exactly three backticks — see the example below. * Emit this block **ONLY AFTER** a successful validation. * The body of the block must be exactly the JSON string the tool returned. Do not reformat, strip fields, add fields, summarize, or pretty-print it differently than the tool produced. From ab01c5e0df3195c103de5e3daf52e6f2a73d6c88 Mon Sep 17 00:00:00 2001 From: labkey-nicka Date: Thu, 20 Aug 2026 16:52:18 -0700 Subject: [PATCH 2/3] Update fence parsing edge cases --- .../ExpressionAssistantAgentAction.java | 219 +++++++++++++++--- 1 file changed, 182 insertions(+), 37 deletions(-) diff --git a/query/src/org/labkey/query/controllers/ExpressionAssistantAgentAction.java b/query/src/org/labkey/query/controllers/ExpressionAssistantAgentAction.java index 297fbb80825..215b909a300 100644 --- a/query/src/org/labkey/query/controllers/ExpressionAssistantAgentAction.java +++ b/query/src/org/labkey/query/controllers/ExpressionAssistantAgentAction.java @@ -43,6 +43,7 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.stream.Collectors; import static org.apache.commons.lang3.StringUtils.isBlank; import static org.apache.commons.lang3.StringUtils.isNotBlank; @@ -194,48 +195,68 @@ private static JSONArray buildSegments(List response MarkdownService md = MarkdownService.get(); StringBuilder htmlBuf = new StringBuilder(); - for (var response : responses) - { - String text = response.text(); - if (isBlank(text)) - continue; + // Scan the turns as one document: tool calling can split a single assistant turn so that a + // fence opens in one MessageResponse and closes in the next. + String text = responses.stream() + .map(McpService.MessageResponse::text) + .filter(StringUtils::isNotBlank) + .collect(Collectors.joining("\n")); + + LOG.debug("Expression assistant raw response:\n{}", text); - String[] lines = text.split("\n", -1); - int i = 0; - while (i < lines.length) + String[] lines = text.split("\n", -1); + int i = 0; + while (i < lines.length) + { + Fence f = readFence(lines, i); + if (f != null && f.terminated && ("sql".equals(f.tag) || "expression".equals(f.tag))) { - Fence f = readFence(lines, i); - if (f != null && f.terminated && ("sql".equals(f.tag) || "expression".equals(f.tag))) - { - flushHtmlSegment(segments, htmlBuf, md); - segments.put(buildSqlSegment(f.tag, f.body)); - i = f.nextIndex; - } - else if (f != null && !f.terminated) - { - // Unterminated fence — fold the body back into prose so we don't drop content, - // but skip the opening fence line itself so the user doesn't see a stray - // "```expression" rendered as a code marker. - if (!htmlBuf.isEmpty()) htmlBuf.append("\n"); - htmlBuf.append(f.body); - break; - } - else - { - // Either not a fence opener or an unknown tag (e.g., python). In the unknown-tag - // case we leave the fence intact in prose so the Markdown renderer turns it into - // a code block. - if (!htmlBuf.isEmpty()) htmlBuf.append("\n"); - htmlBuf.append(lines[i]); - i++; - } + flushHtmlSegment(segments, htmlBuf, md); + segments.put(buildSqlSegment(f.tag, f.body)); + i = f.nextIndex; + } + else if (f != null && !f.terminated) + { + // Unterminated fence — fold the body back into prose so we don't drop content, + // but skip the opening fence line itself so the user doesn't see a stray + // "```expression" rendered as a code marker. + if (!htmlBuf.isEmpty()) htmlBuf.append("\n"); + htmlBuf.append(f.body); + break; + } + else + { + // Either not a fence opener or an unknown tag (e.g., python). In the unknown-tag + // case we leave the fence intact in prose so the Markdown renderer turns it into + // a code block. + if (!htmlBuf.isEmpty()) htmlBuf.append("\n"); + htmlBuf.append(lines[i]); + i++; } } flushHtmlSegment(segments, htmlBuf, md); + + // A payload in prose means the model broke the fence protocol, so the user is reading raw JSON + // instead of being offered an Apply Expression action. Log the response that caused it. + if (hasLeakedPayload(segments)) + LOG.warn("Expression assistant leaked a validator payload into prose:\n{}", text); + return segments; } + /** True when a validateCalculatedColumnExpression payload reached prose instead of an `expression` fence. */ + private static boolean hasLeakedPayload(JSONArray segments) + { + for (int i = 0; i < segments.length(); i++) + { + JSONObject segment = segments.getJSONObject(i); + if ("html".equals(segment.optString("type")) && segment.optString("html", "").contains("jdbcType")) + return true; + } + return false; + } + /** * Build a segment JSON object for an `sql` or `expression` fenced block. For `expression` * blocks the body is expected to be the JSON returned by validateCalculatedColumnExpression; @@ -277,19 +298,38 @@ private static JSONObject buildSqlSegment(String tag, String body) private record Fence(String tag, String body, int nextIndex, boolean terminated) {} + /** Length of the leading backtick run if {@code trimmed} is long enough to delimit a fence, else 0. */ + private static int backtickRun(String trimmed) + { + int n = 0; + while (n < trimmed.length() && trimmed.charAt(n) == '`') + n++; + return n >= 3 ? n : 0; + } + + /** Per CommonMark a closing fence is backticks only, and at least as long as the opener. */ + private static boolean isFenceClose(String line, int openLength) + { + String trimmed = line.trim(); + int n = backtickRun(trimmed); + return n == trimmed.length() && n >= openLength; + } + private static Fence readFence(String[] lines, int i) { String trimmed = lines[i].trim(); - if (!trimmed.startsWith("```")) + int openLength = backtickRun(trimmed); + if (openLength == 0) return null; - String rest = trimmed.substring(3).trim(); + String rest = trimmed.substring(openLength).trim(); if (rest.isEmpty()) return null; - String tag = rest.toLowerCase(); + // Only the first word of the info string is the tag; models pad it ("expression json"). + String tag = rest.split("\\s+", 2)[0].toLowerCase(); int j = i + 1; StringBuilder body = new StringBuilder(); - while (j < lines.length && !"```".equals(lines[j].trim())) + while (j < lines.length && !isFenceClose(lines[j], openLength)) { if (!body.isEmpty()) body.append("\n"); body.append(lines[j]); @@ -348,6 +388,24 @@ private static String expressionPayload(String sql, String jdbcType) return json.toString(); } + private static void assertExpressionSegment(JSONArray segments, int i, String sql, String jdbcType) + { + assertEquals("expression", segment(segments, i).getString("type")); + assertEquals(sql, segment(segments, i).getString("sql")); + assertEquals(jdbcType, segment(segments, i).getString("jdbcType")); + } + + /** The validator payload is transport between the tool and this action; it must never reach the user as text. */ + private static void assertNoPayloadInProse(JSONArray segments) + { + for (int i = 0; i < segments.length(); i++) + { + JSONObject seg = segments.getJSONObject(i); + if ("html".equals(seg.optString("type"))) + assertFalse("validator payload leaked into prose: " + seg, seg.getString("html").contains("jdbcType")); + } + } + @Test public void emptyResponseList() { @@ -490,6 +548,93 @@ public void fenceTagIsCaseInsensitive() assertEquals("INTEGER", segment(segments, 0).getString("jdbcType")); } + @Test + public void expressionFenceIsRecognizedAtAnyBacktickLength() + { + for (int n = 3; n <= 6; n++) + { + String delim = "`".repeat(n); + String md = delim + "expression\n" + expressionPayload("Int7 + Int6", "INTEGER") + "\n" + delim; + JSONArray segments = buildSegments(List.of(markdownResponse(md))); + assertEquals("fence length " + n, 1, segments.length()); + assertExpressionSegment(segments, 0, "Int7 + Int6", "INTEGER"); + assertNoPayloadInProse(segments); + } + } + + @Test + public void longerClosingFenceTerminatesShorterOpener() + { + String md = "```expression\n" + expressionPayload("Int7 + Int6", "INTEGER") + "\n````"; + JSONArray segments = buildSegments(List.of(markdownResponse(md))); + assertEquals(1, segments.length()); + assertExpressionSegment(segments, 0, "Int7 + Int6", "INTEGER"); + } + + @Test + public void shorterClosingFenceDoesNotTerminateLongerOpener() + { + // Per CommonMark the short line is fence content, not a closer, so the block never terminates. + String md = "````expression\n" + expressionPayload("Int7 + Int6", "INTEGER") + "\n```"; + JSONArray segments = buildSegments(List.of(markdownResponse(md))); + assertEquals(1, segments.length()); + assertEquals("html", segment(segments, 0).getString("type")); + } + + @Test + public void leakedPayloadDetectedWhenFenceIsMissing() + { + // No opening fence, so the payload lands in prose and the user gets no Apply affordance. + String md = "Adding Int7 and Int6.\n\n" + expressionPayload("Int7 + Int6", "INTEGER") + "\n```"; + JSONArray segments = buildSegments(List.of(markdownResponse(md))); + assertEquals(1, segments.length()); + assertEquals("html", segment(segments, 0).getString("type")); + assertTrue("prose payload should be reported as a leak", hasLeakedPayload(segments)); + } + + @Test + public void wellFormedExpressionSegmentIsNotALeak() + { + // "jdbcType" is a key on the expression segment itself, so only prose can leak. + String md = "```expression\n" + expressionPayload("Int7 + Int6", "INTEGER") + "\n```"; + JSONArray segments = buildSegments(List.of(markdownResponse(md))); + assertExpressionSegment(segments, 0, "Int7 + Int6", "INTEGER"); + assertFalse("a parsed expression segment is not a leak", hasLeakedPayload(segments)); + } + + @Test + public void proseWithoutPayloadIsNotALeak() + { + JSONArray segments = buildSegments(List.of(markdownResponse("Which date field should be used?"))); + assertEquals("html", segment(segments, 0).getString("type")); + assertFalse(hasLeakedPayload(segments)); + } + + @Test + public void expressionFenceWithExtraInfoStringTokenIsRecognized() + { + // The tag is matched against the whole info string, so any trailing token defeats it. + String md = "```expression json\n" + expressionPayload("Int7 + Int6", "INTEGER") + "\n```"; + JSONArray segments = buildSegments(List.of(markdownResponse(md))); + assertEquals(1, segments.length()); + assertExpressionSegment(segments, 0, "Int7 + Int6", "INTEGER"); + assertNoPayloadInProse(segments); + } + + @Test + public void expressionFenceSplitAcrossResponsesIsRecognized() + { + // Fence state does not survive the per-response scan, so a turn-split mid-fence loses the payload. + String payload = expressionPayload("Int7 + Int6", "INTEGER"); + var r1 = markdownResponse("Adding Int7 and Int6.\n\n```expression\n" + payload); + var r2 = markdownResponse("```"); + JSONArray segments = buildSegments(List.of(r1, r2)); + assertEquals(2, segments.length()); + assertEquals("html", segment(segments, 0).getString("type")); + assertExpressionSegment(segments, 1, "Int7 + Int6", "INTEGER"); + assertNoPayloadInProse(segments); + } + @Test public void preservesMultilineSqlBody() { From b8057e55e82c141d9baad1d6c1e82b72d21a9534 Mon Sep 17 00:00:00 2001 From: labkey-nicka Date: Thu, 20 Aug 2026 22:36:17 -0700 Subject: [PATCH 3/3] Review feedback --- .../ExpressionAssistantAgentAction.java | 80 +++++++++++++++---- 1 file changed, 66 insertions(+), 14 deletions(-) diff --git a/query/src/org/labkey/query/controllers/ExpressionAssistantAgentAction.java b/query/src/org/labkey/query/controllers/ExpressionAssistantAgentAction.java index 215b909a300..070f73b973b 100644 --- a/query/src/org/labkey/query/controllers/ExpressionAssistantAgentAction.java +++ b/query/src/org/labkey/query/controllers/ExpressionAssistantAgentAction.java @@ -217,12 +217,10 @@ private static JSONArray buildSegments(List response } else if (f != null && !f.terminated) { - // Unterminated fence — fold the body back into prose so we don't drop content, - // but skip the opening fence line itself so the user doesn't see a stray - // "```expression" rendered as a code marker. - if (!htmlBuf.isEmpty()) htmlBuf.append("\n"); - htmlBuf.append(f.body); - break; + // Unterminated fence — skip only the opening line, so the user doesn't see a stray + // "```expression" rendered as a code marker, and keep scanning. The body lands in prose + // line by line via the branch below, and a well-formed fence later in the turn still parses. + i++; } else { @@ -237,21 +235,29 @@ else if (f != null && !f.terminated) flushHtmlSegment(segments, htmlBuf, md); - // A payload in prose means the model broke the fence protocol, so the user is reading raw JSON - // instead of being offered an Apply Expression action. Log the response that caused it. + // A payload the user can read means the model broke the fence protocol, so it renders as raw JSON + // instead of a working Apply Expression action. Log the response that caused it. if (hasLeakedPayload(segments)) - LOG.warn("Expression assistant leaked a validator payload into prose:\n{}", text); + LOG.warn("Expression assistant leaked a validator payload into its reply:\n{}", text); return segments; } - /** True when a validateCalculatedColumnExpression payload reached prose instead of an `expression` fence. */ + /** + * True when a validateCalculatedColumnExpression payload reached the user as text rather than being unpacked + * into an `expression` segment. Prose is one surface; the others are an illustrative `sql` block and an + * `expression` block whose body didn't yield an expression, which shows the raw JSON behind an Apply button. + * A well-formed `expression` segment carries jdbcType as its own key, so only the displayed text is checked. + */ private static boolean hasLeakedPayload(JSONArray segments) { for (int i = 0; i < segments.length(); i++) { JSONObject segment = segments.getJSONObject(i); - if ("html".equals(segment.optString("type")) && segment.optString("html", "").contains("jdbcType")) + String displayed = "html".equals(segment.optString("type")) + ? segment.optString("html", "") + : segment.optString("sql", ""); + if (displayed.contains("jdbcType")) return true; } return false; @@ -527,6 +533,19 @@ public void unterminatedFenceFallsBackToHtml() assertFalse("opening fence must be stripped: " + html, html.contains("```")); } + @Test + public void unterminatedFenceDoesNotDiscardLaterResponses() + { + // An unterminated fence must cost only its own block — the four-backtick opener below is never + // closed (its closer is shorter), and the well-formed block after it still has to parse. + var r1 = markdownResponse("````expression\n" + expressionPayload("SELECT 1", "INTEGER") + "\n```"); + var r2 = markdownResponse("Corrected:\n```expression\n" + expressionPayload("SELECT 2", "BIGINT") + "\n```"); + JSONArray segments = buildSegments(List.of(r1, r2)); + assertEquals(2, segments.length()); + assertEquals("html", segment(segments, 0).getString("type")); + assertExpressionSegment(segments, 1, "SELECT 2", "BIGINT"); + } + @Test public void unknownFenceLanguageIsTreatedAsProse() { @@ -592,10 +611,43 @@ public void leakedPayloadDetectedWhenFenceIsMissing() assertTrue("prose payload should be reported as a leak", hasLeakedPayload(segments)); } + @Test + public void leakedPayloadDetectedInSqlFence() + { + // The model reached for the illustrative `sql` tag, so the payload renders as a read-only code block. + String md = "```sql\n" + expressionPayload("Int7 + Int6", "INTEGER") + "\n```"; + JSONArray segments = buildSegments(List.of(markdownResponse(md))); + assertEquals(1, segments.length()); + assertEquals("sql", segment(segments, 0).getString("type")); + assertTrue("payload in a sql block should be reported as a leak", hasLeakedPayload(segments)); + } + + @Test + public void leakedPayloadDetectedWhenExpressionBodyIsNotJson() + { + // The model narrated inside the fence, so the body doesn't parse and Apply would write JSON to the field. + String body = "Here is the payload: " + expressionPayload("Int7 + Int6", "INTEGER"); + JSONArray segments = buildSegments(List.of(markdownResponse("```expression\n" + body + "\n```"))); + assertEquals(1, segments.length()); + assertEquals("expression", segment(segments, 0).getString("type")); + assertEquals(body, segment(segments, 0).getString("sql")); + assertTrue("unparsable payload should be reported as a leak", hasLeakedPayload(segments)); + } + + @Test + public void leakedPayloadDetectedWhenExpressionKeyIsMissing() + { + // Parses, but nothing to unpack, so buildSqlSegment falls back to the raw body. + String body = "{\"jdbcType\":\"INTEGER\",\"sql\":\"Int7 + Int6\"}"; + JSONArray segments = buildSegments(List.of(markdownResponse("```expression\n" + body + "\n```"))); + assertEquals(body, segment(segments, 0).getString("sql")); + assertTrue("payload without an expression key should be reported as a leak", hasLeakedPayload(segments)); + } + @Test public void wellFormedExpressionSegmentIsNotALeak() { - // "jdbcType" is a key on the expression segment itself, so only prose can leak. + // "jdbcType" is a key on the expression segment itself, not part of the SQL the user sees. String md = "```expression\n" + expressionPayload("Int7 + Int6", "INTEGER") + "\n```"; JSONArray segments = buildSegments(List.of(markdownResponse(md))); assertExpressionSegment(segments, 0, "Int7 + Int6", "INTEGER"); @@ -613,7 +665,7 @@ public void proseWithoutPayloadIsNotALeak() @Test public void expressionFenceWithExtraInfoStringTokenIsRecognized() { - // The tag is matched against the whole info string, so any trailing token defeats it. + // Only the info string's first word is the tag, so "expression json" still resolves to "expression". String md = "```expression json\n" + expressionPayload("Int7 + Int6", "INTEGER") + "\n```"; JSONArray segments = buildSegments(List.of(markdownResponse(md))); assertEquals(1, segments.length()); @@ -624,7 +676,7 @@ public void expressionFenceWithExtraInfoStringTokenIsRecognized() @Test public void expressionFenceSplitAcrossResponsesIsRecognized() { - // Fence state does not survive the per-response scan, so a turn-split mid-fence loses the payload. + // Tool calling can split one assistant turn mid-fence; the scan joins the turns so the block still closes. String payload = expressionPayload("Int7 + Int6", "INTEGER"); var r1 = markdownResponse("Adding Int7 and Int6.\n\n```expression\n" + payload); var r2 = markdownResponse("```");