Skip to content

ExpressionAssistant: improve fence parsing #7963

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
labkey-nicka merged 3 commits into from
Aug 21, 2026
+236 −38
Merged

ExpressionAssistant: improve fence parsing #7963

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
807b9d9
Revise prompt to remove four backticks
labkey-nicka Aug 20, 2026
ab01c5e
Update fence parsing edge cases
labkey-nicka Aug 20, 2026
b8057e5
Review feedback
labkey-nicka Aug 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
271 changes: 234 additions & 37 deletions query/src/org/labkey/query/controllers/ExpressionAssistantAgentAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -194,48 +195,74 @@ private static JSONArray buildSegments(List<McpService.MessageResponse> 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 — 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
{
// 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 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 its reply:\n{}", text);

return segments;
}

/**
* 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);
String displayed = "html".equals(segment.optString("type"))
? segment.optString("html", "")
: segment.optString("sql", "");
if (displayed.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;
Expand Down Expand Up @@ -277,19 +304,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]);
Expand Down Expand Up @@ -348,6 +394,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()
{
Expand Down Expand Up @@ -469,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()
{
Expand All @@ -490,6 +567,126 @@ 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 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, 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");
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()
{
// 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());
assertExpressionSegment(segments, 0, "Int7 + Int6", "INTEGER");
assertNoPayloadInProse(segments);
}

@Test
public void expressionFenceSplitAcrossResponsesIsRecognized()
{
// 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("```");
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()
{
Expand Down
3 changes: 2 additions & 1 deletion query/src/org/labkey/query/controllers/prompts/ExpressionAssistant.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down