Skip to content

fix(detection): follow local imports in MCP source retrieval - #114

Open
shengqi-gensee wants to merge 4 commits into
uber:mainfrom
shengqi-gensee:feature/import-following-source-retrieval
Open

fix(detection): follow local imports in MCP source retrieval#114
shengqi-gensee wants to merge 4 commits into
uber:mainfrom
shengqi-gensee:feature/import-following-source-retrieval

Conversation

@shengqi-gensee

Copy link
Copy Markdown
Contributor

Related issue:

Relates to #50

What changed?

get_source_code now returns a bounded, dependency-aware source bundle for
registered Python MCP servers instead of exposing only the registered
entrypoint.

The implementation:

  • statically follows local import and from ... import ... statements;
  • includes package initializers and imported implementation modules;
  • handles relative imports and namespace packages;
  • follows Python package-versus-module and built-in/frozen/native-loader
    precedence;
  • constrains resolved paths to the registered server directory;
  • excludes symlinked source candidates;
  • deduplicates cycles and preserves deterministic discovery order;
  • never imports or executes discovered modules; and
  • preserves the existing source_code entrypoint field and MCP input schema.

The additional source is returned through source_files, with
source_bundle_complete indicating whether the bounded bundle was truncated.

No detector prompts, routing rules, model configuration, or verdict semantics
are changed.

How did you test it?

Deterministic tests cover:

  • local and relative imports;
  • package initializers and from package import submodule;
  • regular-package versus same-named-module precedence;
  • namespace-package behavior;
  • built-in and native-module precedence;
  • cycles and unrelated-file exclusion;
  • resolved-path containment and registry traversal attempts;
  • symlink exclusion;
  • confirmation that discovered modules are never executed; and
  • preservation of the existing MCP schema and response shape.

Local results:

  • Python 3.11 focused provider/baseline suite: 74 passed
  • Python 3.12 focused provider/baseline suite: 74 passed
  • Full Detection suite: 129 passed
  • Registry smoke test: 115/115 registered sources retrieved

We also ran an isolated three-repeat ADR-Bench development comparison with
stock Tier 1 and Tier 2 behavior. The only changed variable was entrypoint-only
versus import-following source context. A second control removed comments and
docstrings while retaining executable code.

Source control Metric Entrypoint only Imports
Raw source Mean recall 77.8% 82.5%
Raw source Mean FPR 13.5% 12.9%
Comments/docstrings stripped Mean recall 77.8% 84.1%
Comments/docstrings stripped Mean FPR 15.1% 15.3%

task_008 was consistently recovered by import-following retrieval in all
three raw-source and all three comments/docstrings-stripped repeats. The
entrypoint-only arm classified it as benign in all six corresponding runs.

The imported implementation exposed executable behavior in
function/ai_insights.py and function/system_analyzer.py, including encoded
session-token handling, decoding, endpoint modification, and silent
exfiltration. The result persisted after comments and docstrings were removed,
indicating that the recovery came from executable implementation behavior
rather than vulnerability-describing prose.

@barisozbas barisozbas left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi, Please see the below findings.

  • High — source_bundle_complete
    (

    "source_bundle_complete": not any(
    )
    incorrectly remains true when the 24-file limit omits queued dependencies. I reproduced 25 files returning only 24
    while reporting complete. Exact size-boundary exhaustion has the same problem.

  • Medium — The collector
    (

    source = path.read_text(encoding="utf-8")
    )
    reads entire files before applying the limit and counts characters, not bytes. A 100,002-character source produced
    400,002 UTF-8 bytes despite the 240,000 budget.

  • Medium — A valid non-UTF-8 Python dependency raises UnicodeDecodeError, causing the entire server response—including ts entrypoint—to become status: error.

@shengqi-gensee

Copy link
Copy Markdown
Contributor Author

Thanks for the review—all three findings were valid and are now fixed:

  • Bundle completeness accounts for dependencies omitted by file or byte limits.
  • Reads are bounded and budgeted in bytes, including failed decoding attempts.
  • Python encoding declarations are supported; invalid codecs or dependency decoding failures preserve the entrypoint and mark the bundle incomplete.

Added regression coverage for these cases and truncation boundaries. All 145 tests pass. Ready for re-review.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants