Skip to content
Merged
Show file tree
Hide file tree
Changes from 8 commits
Commits
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
517 changes: 353 additions & 164 deletions packages/agami-core/src/execute_sql.py

Large diffs are not rendered by default.

4 changes: 4 additions & 0 deletions packages/agami-core/src/mcp_http.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@
bootstrap_paths,
record_tool_call,
server_version,
set_injected_executor,
)

_log = logging.getLogger(__name__)
Expand Down Expand Up @@ -367,6 +368,9 @@ def create_app(extra_tools: dict | None = None, adapters: Adapters | None = None
bootstrap_paths()
adapters = adapters or default_adapters()
auth_provider = adapters.auth_provider
# AH-012: register the composition-root executor (None = the default subprocess path). Behind the
# shared guard in `tool_execute_sql`; a hosted consumer injects a pooled/RBAC/tunnel executor here.
set_injected_executor(adapters.executor)
# Validate consumer-supplied tools up front so a malformed entry fails at construction with a
# clear error, not later as a KeyError/500 inside tools/list or tools/call.
for tool_name, meta in (extra_tools or {}).items():
Expand Down
47 changes: 38 additions & 9 deletions packages/agami-core/src/ports.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
"""The four port Protocols — the seams adapters plug into.
"""The five port Protocols — the seams adapters plug into.

agami-core keeps one MCP implementation across deployments; deployment-specific behavior is
swapped at the composition root through these ports, never by forking a tool:
Expand All @@ -7,6 +7,8 @@
- ``OrgResolver`` — single vs multi tenancy as a config flag, not a schema fork
- ``AuthProvider`` — bearer token → principal (presence by default)
- ``GovernancePolicy`` — warn-only by default; enforcement is a paid concern
- ``Executor`` — the connect-and-run step, *behind* the shared guard (built-in by default;
a consumer injects a pooled/RBAC/tunnel executor without forking the guard)

These are **interfaces only** — `typing.Protocol`, so an adapter satisfies a port by shape, with
no import coupling back to core. The OSS default adapters live in ``oss_adapters`` (so the local
Expand All @@ -31,6 +33,12 @@
# @runtime_checkable only checks method *names*, so isinstance() works without these.
from contracts import QueryExecutionRecord

# ``ExecResult`` is defined in ``execute_sql`` (not here): it is the executor's result type and
# ``execute_sql`` ships in the stdlib-lean plugin mirror that does NOT include this module, so it
# cannot import ``ports`` at runtime. Referencing it under TYPE_CHECKING keeps the ``Executor``
# annotation resolvable for type-checkers without a runtime import cycle.
from execute_sql import ExecResult

# ---------------------------------------------------------------------------
# Seam value types (minimal — a consumer extends them when it needs more)
# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -63,7 +71,7 @@ class GovernanceVerdict:


# ---------------------------------------------------------------------------
# The four ports
# The five ports
# ---------------------------------------------------------------------------


Expand Down Expand Up @@ -104,22 +112,43 @@ class GovernancePolicy(Protocol):
def evaluate(self, ctx: object | None = None) -> GovernanceVerdict: ...


@runtime_checkable
class Executor(Protocol):
"""Connect to a datasource and run **already-vetted** SQL — the only swappable part of the
execution path. It runs *inside* the guarded envelope (guard → executor → shape/log), so it
**only ever receives SQL the guard already passed**, never raw user input; it does no guarding,
logging, or governance itself. This is the seam a hosted consumer overrides to supply a
pooled / per-user-RBAC / SSH-tunnel executor **behind agami-core's one guard** — no fork.

``profile`` is the datasource identity a pooling executor keys its reused connection on. The
built-in OSS default (subprocess/direct connect-per-query) implements this same shape, so a
plain deploy is unchanged."""

def execute(self, vetted_sql: str, creds: dict[str, str], *, profile: str) -> ExecResult: ...


# ---------------------------------------------------------------------------
# Composition-root container — the four adapters passed as one argument
# Composition-root container — the adapters passed as one argument
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class Adapters:
"""The four port adapters, bundled so ``mcp_http.create_app`` takes them as one argument.
"""The port adapters, bundled so ``mcp_http.create_app`` takes them as one argument.

A consumer builds this with its own implementations of the ports (its own ``OrgResolver``,
``AuthProvider``, ``ActivitySink``, ``GovernancePolicy``, and optionally an ``Executor``);
passing ``adapters=None`` to ``create_app`` uses the OSS defaults (``mcp_http.default_adapters``).
Today ``create_app`` wires ``auth_provider`` + ``org_resolver`` into the request path;
``activity_sink`` + ``governance`` are carried here for consumers and not yet referenced by a
core call site.

A consumer builds this with its own implementations of the four ports (its own ``OrgResolver``,
``AuthProvider``, ``ActivitySink``, and ``GovernancePolicy``); passing ``adapters=None`` to
``create_app`` uses the OSS defaults (``mcp_http.default_adapters``). Today ``create_app`` wires
``auth_provider`` + ``org_resolver`` into the request path; ``activity_sink`` + ``governance``
are carried here for consumers and not yet referenced by a core call site."""
``executor`` is optional and defaults to ``None`` — meaning "use the built-in executor" (the
subprocess/direct connect-per-query path, byte-identical to today). A consumer sets it to run
execution in-process behind the shared guard (see ``tools.tool_execute_sql``)."""

activity_sink: ActivitySink
org_resolver: OrgResolver
auth_provider: AuthProvider
governance: GovernancePolicy
executor: Executor | None = None
186 changes: 147 additions & 39 deletions packages/agami-core/src/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -871,12 +871,139 @@ def _executor_truncated(stderr: str | None) -> bool:
return False


# The composition-root executor (AH-012). ``None`` (the default) means "fork the execute_sql
# subprocess" — the byte-identical local/single-user path. A consumer injects a ``ports.Executor``
# via ``create_app(adapters=…)`` to run execution IN-PROCESS behind the same guard (no fork, native
# rows). Process-global on purpose: the executor is a composition-root singleton, not per-request.
_INJECTED_EXECUTOR: Any | None = None


def set_injected_executor(executor: Any | None) -> None:
"""Register (or clear) the composition-root executor. Called once by ``mcp_http.create_app`` from
``adapters.executor``; ``None`` keeps the default subprocess path. Validates the shape at
registration so a malformed adapter fails fast at app construction, not as an ``AttributeError``
at query time."""
global _INJECTED_EXECUTOR
if executor is not None:
import ports

if not isinstance(executor, ports.Executor): # runtime_checkable: has execute(...)
raise TypeError(
"injected executor must satisfy ports.Executor "
"(an execute(vetted_sql, creds, *, profile) method)"
)
_INJECTED_EXECUTOR = executor
Comment on lines +881 to +895

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Done in 817de66set_injected_executor now isinstance-checks against ports.Executor (runtime_checkable) at registration, so a malformed adapter fails fast at create_app construction instead of an AttributeError at query time.



def _finalize_execution(
columns: list, data_rows: list, truncated: bool, *, profile: str, sql: str,
execution_ms: int, args: dict[str, Any],
) -> str:
"""Shape a successful result (units + exact-render markdown + trust receipt), log the execution
through the single sink, and return the tool JSON. Shared by both execution paths — the subprocess
fork and the in-process executor — so a query returns the identical envelope whichever ran it."""
# Deterministic, exact rendering — so the numbers a user verifies don't depend on
# how the host LLM chooses to format them. `markdown` is the table to display
# verbatim; `rows` stays raw (exact CSV values) for charting / programmatic use.
unit_map = _resolve_units(profile, sql)
try:
from semantic_model import units # stdlib-only; safe even without model deps

markdown = units.format_table(columns, data_rows, unit_map)
except Exception:
markdown = None

result = {
"columns": columns,
"rows": data_rows,
"row_count": len(data_rows),
"truncated": truncated,
"units": unit_map,
"markdown": markdown, # exact, full numbers (currency symbol + grouping) — render as-is
"sql": sql,
"execution_ms": execution_ms,
# Trust receipt — provenance + anything unapproved this answer used. Same assembler
# the agami-query skill renders, so Desktop gets the same trust panel. Clients should
# surface receipt.warnings and any receipt.metrics whose review_state != "approved"
# (offer to approve/correct via the save_correction tool).
"receipt": _resolve_receipt(profile, sql),
}

# Log the execution through the single chokepoint: the DB sink when AGAMI_DB_URL is set (one
# query_executions row), else the local jsonl the skills use. Best-effort either way.
_record_query(
{
"ts": _now_iso(),
"profile": profile,
"question": args.get("raw_query"),
"sql": sql,
"row_count": len(data_rows),
"source": "mcp_server",
}
)
return json.dumps(result, indent=2, default=str)


def _run_in_process(
sql: str, profile: str, area: str | None, max_rows: int | None, executor: Any
) -> tuple[list, list, bool] | dict:
"""Run through the in-process executor behind the shared guarded envelope (no subprocess, no CSV
round-trip). Returns ``(columns, data_rows, truncated)`` on success, or an error dict on a guard
refusal / execution failure — the same error shape the subprocess branch produces.

Rows are textualized to match the subprocess CSV wire (``None`` → ``""``, else ``str``) so the
two paths return observably identical JSON. Native-typed rows are a deliberately deferred decision
(see the AH-012 spec); flipping this one coercion is the follow-up once that's settled."""
import execute_sql

# The per-call cap rides execute_sql's module global (ACE-044). ACE-028, which makes in-process
# the real serving path, must thread the cap per-call instead — this global is not safe under
# concurrent in-process queries with different caps. Save/restore keeps it inert by default.
prev_cap = execute_sql._max_rows_override
execute_sql._max_rows_override = max_rows
try:
result = execute_sql.execute_guarded(sql, profile, area, executor=executor)
except execute_sql.GuardRefused as refusal:
# A read-only refusal (envelope present) is already caught by tool_execute_sql's upstream
# check_read_only fast-fail, so in practice only the model-safety branch (envelope None) is
# reached here; both are handled for defence-in-depth.
if refusal.envelope is not None:
return {"error": refusal.envelope["error"]}
# A model-safety refusal wrote its detail to the server log (stderr); surface a clean refusal.
return {"error": {"kind": "permission",
"remediation": "Query refused by the semantic-model safety pass."}}

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Correct — and I've fixed the overclaim in eb83acc, but I'm deliberately not folding the full parity fix into this PR. Here's the reasoning:

Successful results are identical across paths (the valuable guarantee). The divergence is only for guard refusals, and it's structural: _model_safety writes its structured envelope to sys.stderr and returns (sql, rc). On the subprocess path tool_execute_sql reads that stderr and surfaces it (kind _classify_exit(1)); in-process, that JSON went to the server log (not captured — capturing it means a thread-unsafe global sys.stderr swap), so I return a clean generic refusal.

True structured-refusal parity requires _model_safety to return its envelope (so execute_guarded can carry it in GuardRefused for both callers), which ripples into:

  1. test_ace051 — the fail-closed guard tests pin both the stderr JSON and the (sql, rc) 2-tuple across ~8 sites, and
  2. the default subprocess path's refusal output (parsing the stderr envelope) — i.e. no longer byte-identical.

That's a change to the security guard's error-reporting contract, separable from the executor seam. Folding it into this PR would expand a security-sensitive, byte-identical-by-design change into the guard internals. So I've corrected the docstrings to scope the "identical envelope" claim to successful results, clarified the in-process remediation points to the server log, and I'm tracking structured-refusal parity as a follow-up (natural fit with the _model_safety-returns-envelope refactor). Both paths still correctly fail closed — only the refusal's presentation differs.

except execute_sql.ExecutorError as exc:
return {"error": {"kind": _classify_exit(exc.code), "remediation": exc.msg}}
Comment on lines +969 to +984

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Addressed in 817de66 with the fail-closed net you suggested: _run_in_process now catches SystemExit, so a deep sys.exit(2) from _load_credentials/_parse_dsn (bad profile/DSN) can no longer escape in-process and take down the host — it becomes a tool error envelope instead.

The deeper refactor (convert those helpers' sys.exit(2)ExecutorError) is a deliberate follow-up: _parse_dsn/_load_credentials have existing tests asserting SystemExit (test_unsupported_scheme_exits, test_random_string_exits) and sit on the byte-identical subprocess/CLI path, so that change ripples beyond this seam — tracking it with ACE-028 (which makes in-process the real serving path). The net makes the seam safe today regardless.

except SystemExit as exc:
# Defence-in-depth. The known credential/DSN failures now raise ExecutorError (handled above,
# carrying their detailed message), so this net catches only a residual/future sys.exit deep
# in a driver — ensuring an in-process query can never take down the host; it becomes a
# fail-closed tool error instead.
code = exc.code if isinstance(exc.code, int) else 2
return {"error": {"kind": _classify_exit(code),
"remediation": "Datasource configuration error."}}
finally:
execute_sql._max_rows_override = prev_cap

Comment on lines +969 to +995

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Good catch — fixed properly in 89f8c55, and I went a different way than stderr-capture on purpose.

Capturing stderr around execute_guarded would require swapping the process-global sys.stderr, which is not thread-safe in the server (handlers run in worker threads via run_blocking; concurrent requests would interleave/steal each other's captured output). So instead I removed the drift at the source: _load_credentials and _parse_dsn now raise ExecutorError(code=2) instead of sys.exit(2). The detailed message rides the exception, so:

  • In-process: _run_in_process catches ExecutorErrorremediation = exc.msg — the exact detailed message, parity with the subprocess path.
  • Subprocess/CLI: main() catches ExecutorError → writes msg + "\n" and returns code 2 — byte-identical stderr + exit to the old sys.exit(2).

Ripple handled: setup_pgauth.py (the one external _parse_dsn caller) translates ExecutorError back to its stderr+exit-2 CLI UX; the SystemExit-asserting tests were updated to assert ExecutorError; and the SystemExit net stays as pure defence-in-depth. New test test_injected_executor_credential_error_surfaces_detailed_remediation pins the parity. This also makes the executor library sys.exit-free (only the __main__ guard remains), which is the cleaner end state you flagged in the other threads.

columns = list(result.columns)
data_rows = [["" if v is None else str(v) for v in row] for row in result.rows]
truncated = result.truncated
if max_rows is not None and len(data_rows) > max_rows: # backstop, matches the subprocess branch
data_rows = data_rows[:max_rows]
truncated = True
return columns, data_rows, truncated


def tool_execute_sql(args: dict[str, Any]) -> str:
"""Local analog of Ask Agami `execute_sql`: run a read-only SELECT locally.

Routes through the sibling execute_sql.py (Tier-3 Python executor) so all
DB types are handled identically and nothing but the rows leaves the
process. Enforces the same read-only guarantee as the hosted connector.

Two execution paths behind the same guard: the default forks the execute_sql subprocess
(isolation, byte-identical local/single-user); an injected executor (AH-012) runs in-process with
native rows. Both funnel through `_finalize_execution` so the returned envelope is identical.
"""
sql = args.get("sql")
if not isinstance(sql, str) or not sql.strip():
Expand All @@ -903,6 +1030,23 @@ def tool_execute_sql(args: dict[str, Any]) -> str:
if max_rows is not None:
max_rows = max(1, min(max_rows, 10_000))

area = str(args["area"]) if args.get("area") else None

# In-process path (AH-012): a consumer injected an executor, so run behind the shared guarded
# envelope with no subprocess and no CSV round-trip. Falls through to the subprocess fork below
# when no executor is injected (the default) — that path stays byte-identical.
if _INJECTED_EXECUTOR is not None:
started = time.monotonic()
outcome = _run_in_process(sql, profile, area, max_rows, _INJECTED_EXECUTOR)
execution_ms = int((time.monotonic() - started) * 1000)
if isinstance(outcome, dict): # guard refusal / execution error
return json.dumps({**outcome, "sql": sql, "execution_ms": execution_ms}, indent=2)
columns, data_rows, truncated = outcome
return _finalize_execution(
columns, data_rows, truncated,
profile=profile, sql=sql, execution_ms=execution_ms, args=args,
)

# The model safety pass (fan/chasm pre-flight + default_filters) runs inside
# execute_sql.py; pass the subject area so default_filters scope correctly.
# Route through the unified executor as a module (the package is installed alongside
Expand Down Expand Up @@ -954,46 +1098,10 @@ def tool_execute_sql(args: dict[str, Any]) -> str:
data_rows = data_rows[:max_rows]
truncated = True

# Deterministic, exact rendering — so the numbers a user verifies don't depend on
# how the host LLM chooses to format them. `markdown` is the table to display
# verbatim; `rows` stays raw (exact CSV values) for charting / programmatic use.
unit_map = _resolve_units(profile, sql)
try:
from semantic_model import units # stdlib-only; safe even without model deps

markdown = units.format_table(columns, data_rows, unit_map)
except Exception:
markdown = None

result = {
"columns": columns,
"rows": data_rows,
"row_count": len(data_rows),
"truncated": truncated,
"units": unit_map,
"markdown": markdown, # exact, full numbers (currency symbol + grouping) — render as-is
"sql": sql,
"execution_ms": execution_ms,
# Trust receipt — provenance + anything unapproved this answer used. Same assembler
# the agami-query skill renders, so Desktop gets the same trust panel. Clients should
# surface receipt.warnings and any receipt.metrics whose review_state != "approved"
# (offer to approve/correct via the save_correction tool).
"receipt": _resolve_receipt(profile, sql),
}

# Log the execution through the single chokepoint: the DB sink when AGAMI_DB_URL is set (one
# query_executions row), else the local jsonl the skills use. Best-effort either way.
_record_query(
{
"ts": _now_iso(),
"profile": profile,
"question": args.get("raw_query"),
"sql": sql,
"row_count": len(data_rows),
"source": "mcp_server",
}
return _finalize_execution(
columns, data_rows, truncated,
profile=profile, sql=sql, execution_ms=execution_ms, args=args,
)
return json.dumps(result, indent=2, default=str)


def _now_iso() -> str:
Expand Down
Loading
Loading