From 486fae0f3dee94025d4ecbb9fde46e8a8eef70fe Mon Sep 17 00:00:00 2001 From: Julius Guay Date: Mon, 1 Jun 2026 15:36:02 -0500 Subject: [PATCH] feat: add result metadata and structural list truncation to execute_code Adds result_type and result_count fields to ExecutionResult so the LLM has an authoritative anchor for result shape without recounting from the payload. result_count is computed before sanitize_output runs, so it reflects the true pre-truncation total for list results. Large list/tuple results are now capped at MAX_RESULT_ITEMS (200) at element boundaries rather than mid-JSON, preventing the model from attempting to complete malformed structures with invented data. The truncation notice is appended as the final list element so the cutoff is explicit. Co-Authored-By: Claude Sonnet 4.6 --- .gitignore | 5 +- src/openstaad_mcp/sandbox/const.py | 5 ++ src/openstaad_mcp/sandbox/executor.py | 33 ++++++++++ src/openstaad_mcp/sandbox/stdio_helpers.py | 16 +++-- src/openstaad_mcp/server.py | 13 +++- tests/sandbox/test_executor.py | 73 ++++++++++++++++++++++ 6 files changed, 138 insertions(+), 7 deletions(-) diff --git a/.gitignore b/.gitignore index 0b4a99e..9354e80 100644 --- a/.gitignore +++ b/.gitignore @@ -50,4 +50,7 @@ venv/ .vscode/ # Logs -mcp_server.log \ No newline at end of file +mcp_server.log + +# Private notes / drafts (LinkedIn, analysis, personal docs) +notes/ \ No newline at end of file diff --git a/src/openstaad_mcp/sandbox/const.py b/src/openstaad_mcp/sandbox/const.py index 62104fe..e8a3c13 100644 --- a/src/openstaad_mcp/sandbox/const.py +++ b/src/openstaad_mcp/sandbox/const.py @@ -275,3 +275,8 @@ # Maximum length for a single result value to prevent large injection payloads MAX_RESULT_LENGTH = 100_000 + +# Maximum number of items in a list/tuple result before structural truncation. +# Items beyond this limit are dropped; result_count in the response still reports +# the true pre-truncation total so the model can anchor its narrative to it. +MAX_RESULT_ITEMS = 200 diff --git a/src/openstaad_mcp/sandbox/executor.py b/src/openstaad_mcp/sandbox/executor.py index 38bdf3d..8ea28e7 100644 --- a/src/openstaad_mcp/sandbox/executor.py +++ b/src/openstaad_mcp/sandbox/executor.py @@ -31,6 +31,29 @@ from openstaad_mcp.sandbox.stdio_helpers import LimitedStringIO, sanitize_output, sanitize_traceback +def _classify_result(value: Any) -> tuple[str, int | None]: + """Return (result_type, result_count) for *value* before sanitization. + + result_type is one of: "null", "scalar", "bool", "string", "list", "dict". + result_count is the true item count for list/tuple/dict, None otherwise. + Computing this before sanitize_output runs ensures result_count reflects the + full pre-truncation total, not the capped payload the model receives. + """ + if value is None: + return "null", None + if isinstance(value, bool): + return "bool", None + if isinstance(value, (int, float)): + return "scalar", None + if isinstance(value, str): + return "string", None + if isinstance(value, (list, tuple)): + return "list", len(value) + if isinstance(value, dict): + return "dict", len(value) + return "scalar", None + + @dataclass class ExecutionResult: """Structured result returned from :func:`execute`.""" @@ -41,11 +64,15 @@ class ExecutionResult: stderr: str = "" error: str | None = None duration_seconds: float = 0.0 + result_type: str = "null" + result_count: int | None = None def to_dict(self) -> dict[str, Any]: return { "success": self.success, "result": self.result, + "result_type": self.result_type, + "result_count": self.result_count, "stdout": self.stdout, "stderr": self.stderr, "error": self.error, @@ -156,6 +183,10 @@ def execute( else: result_value = None + # Classify before sanitization so result_count reflects the true total, + # not the truncated payload the model will receive. + res_type, res_count = _classify_result(result_value) + result_value = sanitize_output(result_value) # Attempt JSON-safe serialisation; fall back to repr. @@ -170,4 +201,6 @@ def execute( stdout=stdout_text, stderr=stderr_text, duration_seconds=duration, + result_type=res_type, + result_count=res_count, ) diff --git a/src/openstaad_mcp/sandbox/stdio_helpers.py b/src/openstaad_mcp/sandbox/stdio_helpers.py index 0ace623..ac7217e 100644 --- a/src/openstaad_mcp/sandbox/stdio_helpers.py +++ b/src/openstaad_mcp/sandbox/stdio_helpers.py @@ -4,7 +4,7 @@ import re as _re from typing import Any -from openstaad_mcp.sandbox.const import MAX_EXECUTION_STDOUT, MAX_RESULT_LENGTH +from openstaad_mcp.sandbox.const import MAX_EXECUTION_STDOUT, MAX_RESULT_ITEMS, MAX_RESULT_LENGTH class LimitedStringIO(io.StringIO): @@ -64,15 +64,21 @@ def sanitize_traceback(exc: BaseException) -> str: def sanitize_output(value: Any) -> Any: """Sanitize output values to mitigate indirect prompt injection. - Truncates large strings and adds a data provenance marker so AI agents - know the data comes from an external model file. + Truncates large strings and long lists at structural boundaries so the + LLM never receives a malformed partial payload it will try to "complete." + The true pre-truncation count is reported separately via result_count in + the executor result dict; the LLM must use that field, not len(result). """ if isinstance(value, str): if len(value) > MAX_RESULT_LENGTH: value = value[:MAX_RESULT_LENGTH] + "... (truncated)" return value - if isinstance(value, list): - return [sanitize_output(item) for item in value] + if isinstance(value, (list, tuple)): + truncated = len(value) > MAX_RESULT_ITEMS + items = [sanitize_output(item) for item in value[:MAX_RESULT_ITEMS]] + if truncated: + items.append(f"... ({len(value) - MAX_RESULT_ITEMS} more items not shown — see result_count for total)") + return items if isinstance(value, dict): return {k: sanitize_output(v) for k, v in value.items()} return value diff --git a/src/openstaad_mcp/server.py b/src/openstaad_mcp/server.py index 9223976..b4dee80 100644 --- a/src/openstaad_mcp/server.py +++ b/src/openstaad_mcp/server.py @@ -191,6 +191,13 @@ def execute_code(code: str, instance: str | None = None) -> dict[str, Any]: Pass ``instance`` (alias from ``list_instances``, e.g. ``staadPro1``) to target a specific STAAD instance. Omit it when only one instance is running — it will be selected automatically. + + The response always includes ``result_type`` (one of ``"null"``, + ``"bool"``, ``"scalar"``, ``"string"``, ``"list"``, ``"dict"``) and + ``result_count`` (the true pre-truncation item count for list/dict + results, ``null`` otherwise). Large list results are truncated to + the first 200 items in ``result``; ``result_count`` is the only + authoritative total — do not recount from ``result`` itself. """ try: target = _resolve_target(instance) @@ -251,7 +258,11 @@ async def mcp_lifespan(server: Any) -> AsyncIterator[None]: "instructions. Use `list_instances` to see running STAAD instances, " "`execute_code` to run code against a live STAAD.Pro model, and " "`get_status` to check connection. " - "When a `warning` field appears in any tool response, report it to the user." + "When a `warning` field appears in any tool response, report it to the user. " + "When `execute_code` returns a list result, `result_count` is the " + "authoritative item count — it reflects the true pre-truncation total. " + "Do not recount from `result`; report exactly `result_count` items and " + "note when results were truncated (i.e. when len(result) < result_count)." ), lifespan=mcp_lifespan, **fastmcp_kwargs, diff --git a/tests/sandbox/test_executor.py b/tests/sandbox/test_executor.py index 5d153cc..b8b5f0a 100644 --- a/tests/sandbox/test_executor.py +++ b/tests/sandbox/test_executor.py @@ -217,6 +217,8 @@ def test_to_dict_keys(self, staad, executor): assert set(d.keys()) == { "success", "result", + "result_type", + "result_count", "stdout", "stderr", "error", @@ -224,6 +226,77 @@ def test_to_dict_keys(self, staad, executor): } +class TestResultAnchoring: + """result_type and result_count provide pre-truncation metadata.""" + + def test_int_result_type(self, staad, executor): + r = executor.execute("result = 42", staad) + assert r.result_type == "scalar" + assert r.result_count is None + + def test_float_result_type(self, staad, executor): + r = executor.execute("result = 3.14", staad) + assert r.result_type == "scalar" + assert r.result_count is None + + def test_bool_result_type(self, staad, executor): + r = executor.execute("result = True", staad) + assert r.result_type == "bool" + assert r.result_count is None + + def test_string_result_type(self, staad, executor): + r = executor.execute('result = "hello"', staad) + assert r.result_type == "string" + assert r.result_count is None + + def test_none_result_type(self, staad, executor): + r = executor.execute("x = 1", staad) + assert r.result_type == "null" + assert r.result_count is None + + def test_list_result_type_and_count(self, staad, executor): + r = executor.execute("result = [1, 2, 3, 4, 5]", staad) + assert r.result_type == "list" + assert r.result_count == 5 + + def test_dict_result_type_and_count(self, staad, executor): + r = executor.execute('result = {"a": 1, "b": 2}', staad) + assert r.result_type == "dict" + assert r.result_count == 2 + + def test_com_tuple_classified_as_list(self, staad, executor): + # COM APIs (e.g. GetPrimaryLoadCaseNumbers) return tuples — must be + # classified as "list" so the model treats them uniformly. + r = executor.execute("result = staad.Output.GetBeamEndForces(1, 1)", staad) + assert r.result_type == "list" + assert r.result_count == 6 + + def test_error_result_type_is_null(self, staad, executor): + r = executor.execute("1 / 0", staad) + assert not r.success + assert r.result_type == "null" + assert r.result_count is None + + def test_large_list_count_reflects_pretuncation_total(self, staad, executor): + from openstaad_mcp.sandbox.const import MAX_RESULT_ITEMS + + n = MAX_RESULT_ITEMS + 50 + r = executor.execute(f"result = list(range({n}))", staad) + assert r.success + # result_count is the true pre-truncation total + assert r.result_count == n + # the result payload is capped at MAX_RESULT_ITEMS (+1 for the truncation notice) + assert len(r.result) == MAX_RESULT_ITEMS + 1 + + def test_small_list_not_truncated(self, staad, executor): + from openstaad_mcp.sandbox.const import MAX_RESULT_ITEMS + + r = executor.execute(f"result = list(range({MAX_RESULT_ITEMS}))", staad) + assert r.success + assert r.result_count == MAX_RESULT_ITEMS + assert len(r.result) == MAX_RESULT_ITEMS + + class TestErrorHandling: """Test that errors are properly reported."""