"""Tool-runtime utilities — open a DB session, load the authenticated user, write an audit log, commit on writes, surface any error in a shape MCP clients can render. The :func:`tool_runtime` async context manager is the single integration point every tool uses; it ensures we never forget to commit a write, never log a tool call without an audit row, and never leak a database session if a tool raises mid-flight. """ from __future__ import annotations import logging import re as _re from collections.abc import Awaitable, Callable from contextlib import asynccontextmanager from dataclasses import dataclass, field from datetime import UTC from typing import Any from sqlalchemy.ext.asyncio import AsyncSession from app.audit.service import AuditService from app.auth.models import AuthUser # Module-attribute reference (not a bare ``from … import``) so test # monkeypatching of ``app.core.database.async_session_factory`` takes # effect on every tool call. See app/mcp/middleware.py for the same # pattern + rationale. from app.core import database as _db from app.core.dependencies import _load_user_with_roles from app.core.exceptions import AppException from app.mcp.context import ( get_or_init_trace_id, mcp_client_ip, mcp_user_agent, require_mcp_user_id, ) logger = logging.getLogger(__name__) @dataclass(slots=True) class ToolRuntime: """Per-tool-call dependencies handed to the tool body.""" db: AsyncSession user: AuthUser _after_commit: list[Callable[[], Awaitable[None]]] = field( default_factory=list, repr=False, ) def after_commit(self, callback: Callable[[], Awaitable[None]]) -> None: """Run best-effort external work after the business commit lands.""" self._after_commit.append(callback) _AUDIT_STR_CAP = 500 _AUDIT_LIST_CAP = 50 _AUDIT_DICT_KEY_CAP = 50 _AUDIT_TOTAL_CAP = 8_000 # rough char budget per audit row's args dict _AUDIT_DEPTH_CAP = 6 # recursion guard against pathological structures # Exact-match key allowlist. The check is `key.lower() in _SENSITIVE_KEYS` # applied at every depth (round-2 fix). Synonyms enumerated explicitly # so the next tool author who uses an unexpected name (e.g. `api_key`, # `access_key`, `client_secret`) still gets redacted by default. _SENSITIVE_KEYS = { "password", "new_password", "old_password", "passphrase", "hashed_password", "secret", "client_secret", "webhook_secret", "signing_key", "signing_secret", # ``app_secret`` is the actual field name on every Feishu integration # driver's config dict (chat-approval / software / hrbp / notifications). # No current tool returns it, but a future tool that hands back a # decrypted integration config (or a Pydantic config model containing # ``app_secret``) would otherwise leak plaintext to the LLM. (R3 # security finding 2026-05-28.) "app_secret", "token", "access_token", "refresh_token", "id_token", "pat", "personal_access_token", "oauth_token", # PAT-mint plaintext field names — the POST /tokens response uses # ``plaintext``; future MCP tool authors may pick synonyms. Keep # this list in sync with ``_SERIALIZE_DENY_FIELDS`` below. "plaintext", "pat_plaintext", "token_plaintext", "pat_value", "pat_hash", "token_hash", "secret_value", "auth_token", "raw_token", "encrypted_token", "encrypted_value", "api_key", "apikey", "access_key", "secret_key", "private_key", # Real ORM column names that carry credential ciphertext — mirror # ``_SERIALIZE_DENY_FIELDS`` so a tool arg named after the column # gets redacted on the inbound side too. "secrets_encrypted", "secrets_nonce", "key_value", "key_hash", "credential", "credentials", "service_account", "gcp_credentials", "authorization", "bearer", # Session/MFA family — keep specific names. ``session_key`` / # ``session_token`` are unambiguous credentials; bare ``session`` # collides with legitimate domain words ("training session", # "inventory session") and would over-rotate. Same for ``mfa`` — # specific ``mfa_secret`` / ``totp_secret`` only. "session_key", "session_token", "cookie", "csrf", "pin", "otp", "mfa_secret", "totp_secret", "verification_code", "recovery_code", } # Substring patterns to redact inside string values. Catches the case # where a user pastes their own PAT into a free-text arg like ``reason`` # — without this the first 500 chars of the PAT would land verbatim in # audit_logs.detail. Patterns intentionally tight: only well-formed # tokens get redacted, not arbitrary look-alikes. _TOKEN_REDACT_PATTERNS = [ # Helix PAT: helix_pat_<32+ url-safe chars> _re.compile(r"helix_pat_[A-Za-z0-9_\-]{16,}"), # JWT-shape tokens (header.payload.signature, base64url) _re.compile(r"eyJ[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+"), # AWS access key id (AKIA + 16 chars) _re.compile(r"AKIA[0-9A-Z]{16}"), # AWS secret access key shape (40 base64-y chars after = sign or # explicit prefix). Tight to avoid false positives. _re.compile(r"aws_secret_access_key[\"'\s:=]{1,4}[A-Za-z0-9/+]{40}"), # Aliyun access key id (LTAI + 12-20 alphanumeric). Helix runs on # Aliyun for one tenant — without this, an LTAI shape pasted into a # ``reason`` field lands plaintext in audit_logs. _re.compile(r"\bLTAI[A-Za-z0-9]{12,20}\b"), # Aliyun STS / RAM-issued secret (prefix ``STS.`` + 16+ alphanumeric). # Helix's Aliyun integrations issue STS credentials for ECS / OSS # access; a paste into a free-text arg should not land in audit. _re.compile(r"\bSTS\.[A-Za-z0-9]{16,}"), # Tencent Cloud secret id (AKID + 32 alphanumeric) _re.compile(r"\bAKID[A-Za-z0-9]{32}\b"), # Azure SAS token signature — appears inside blob/queue/file URLs. _re.compile(r"sig=[A-Za-z0-9%/+=]{20,}(?:&sv=\d{4}-\d{2}-\d{2})?"), # HuggingFace user / org tokens (``hf_`` + 30+ alphanumeric). Some # internal AI tools may carry these in tool args. _re.compile(r"\bhf_[A-Za-z0-9]{30,}\b"), # GCP service-account JSON marker — redact the whole short-form repr # of a credentials JSON paste. We only target the type-tag + private # key marker; tight bounds avoid false positives on legit user prose. _re.compile(r'"type"\s*:\s*"service_account"'), # GitHub tokens — ghp_ / ghs_ / ghu_ / gho_ / ghr_ (36 base62 chars) _re.compile(r"gh[pousr]_[A-Za-z0-9]{36,}"), # Slack tokens — xoxb / xoxa / xoxp / xoxr / xoxs followed by id-id-secret _re.compile(r"xox[baprs]-[A-Za-z0-9\-]{10,}"), # Google OAuth access tokens — ya29. _re.compile(r"ya29\.[A-Za-z0-9_\-]{20,}"), # Stripe live / test secret keys _re.compile(r"sk_(?:live|test)_[A-Za-z0-9]{16,}"), # Anthropic API keys — sk-ant-<...> _re.compile(r"sk-ant-[A-Za-z0-9_\-]{16,}"), # OpenAI API keys. Two shapes, matched as alternatives: # # 1. Prefixed (``sk-proj-`` / ``sk-svcacct-`` / ``sk-live-`` / # ``sk-test-``) followed by a 31+ char body that CAN include # hyphens / underscores. Real OpenAI project keys are shaped # ``sk-proj---...`` with hyphens INSIDE the # body — so the body alphabet must allow them once the prefix # anchor is satisfied. # 2. Legacy / un-prefixed keys: ``sk-`` + 32+ contiguous # alphanumeric (no hyphens). This tight body avoids false # positives on hyphenated English like # ``sk-learn-model-evaluation-pipeline``. _re.compile( r"\bsk-(?:" r"(?:proj|svcacct|live|test)-[A-Za-z0-9][A-Za-z0-9_\-]{31,}" r"|" r"[A-Za-z0-9]{32,}" r")\b", ), # GitLab PATs — ``glpat-`` + 20-char base62. _re.compile(r"\bglpat-[A-Za-z0-9_\-]{20,}\b"), # Notion integration tokens — ``secret_`` + 43-char base62. _re.compile(r"\bsecret_[A-Za-z0-9]{43}\b"), # Generic Bearer header inline (e.g. "Authorization: Bearer eyJ..."). # The previous pattern matched any 16+ run of [A-Za-z0-9_\-\.], # which over-fired on English-prose words like "bearer responsibilities # of the role". Tighten by requiring the value contain at least one # non-alpha character (digit, dot, underscore, dash, slash, plus, # or equals) — real tokens almost always do; English words don't. _re.compile( r"(?i)bearer\s+(?=[A-Za-z0-9_\-\.+/=]*[\d._\-+/=])" r"[A-Za-z0-9_\-\.+/=]{16,}", ), # PEM private-key blocks (single-line repr from pasted multiline) _re.compile( r"-----BEGIN [A-Z ]*PRIVATE KEY-----[\s\S]+?-----END [A-Z ]*PRIVATE KEY-----", ), ] # Translate-table that strips C0 control chars (except tab/newline) and # the unicode bidi-override characters that auditors otherwise read as # misleading text. Built once at module load. ``127`` is the DEL # control byte — included to keep the strip symmetric with the CLI- # side ``_TERMINAL_STRIP_TABLE``. _CTRL_TABLE = { # range(0, 9): NUL–BS; 9=TAB kept; 10=LF kept; 11=VT; 12=FF; # 13=CR — MUST be stripped to prevent audit-log line-injection via # \r overwrite; range(14,32): SO–US; 127=DEL. c: None for c in [*range(0, 9), 11, 12, 13, *range(14, 32), 127] } | { # C1 controls (U+0080–U+009F). U+009B is a single-character CSI # and U+009D is OSC — xterm-class terminals honor them even in # UTF-8 mode, so "›2J" clears the screen exactly like the # 7-bit "\x1b[2J" the C0 strip blocks. Without this range, the # claimed defense-in-depth (server scrub + CLI scrub, both of # which previously stripped C0 only) was bypassable end-to-end. # (Launch audit 2026-06-10.) c: None for c in range(0x80, 0xA0) } | {c: None for c in (0x202A, 0x202B, 0x202C, 0x202D, 0x202E, 0x2066, 0x2067, 0x2068, 0x2069)} def _redact_str(s: str) -> str: """Truncate + strip dangerous control / RTL-override chars + scrub any inline credential the user may have pasted. Auditors reading ``audit_logs.detail`` in the admin UI shouldn't see a string that visually re-orders itself (RTL override) or carries NUL / CR / LF that could be mis-rendered into the wrapping JSON. Also scrubs Helix PATs, JWTs, and AWS access-key-shaped substrings — a user pasting their own token into ``reason`` shouldn't have it durably stored in the audit table. """ s = s.translate(_CTRL_TABLE) for pat in _TOKEN_REDACT_PATTERNS: s = pat.sub("[redacted-credential]", s) if len(s) > _AUDIT_STR_CAP: return s[:_AUDIT_STR_CAP] + f"...[{len(s) - _AUDIT_STR_CAP} more chars]" return s def _flatten_deny(deny: frozenset | set) -> set: """Pre-compute the flat (no `_` / no `-`) form of every entry in a deny set so a header-style input like ``Hashed-Password`` can match a snake_case entry like ``hashed_password`` via single lookup without falling back to the dash-split component scan (which over-rotates on legitimate field names containing a short deny word). """ return {s.replace("_", "").replace("-", "") for s in deny} _SENSITIVE_KEYS_FLAT = _flatten_deny(_SENSITIVE_KEYS) def _is_sensitive_key( key: str, *, deny_set: frozenset | set | None = None, deny_set_flat: frozenset | set | None = None, strict: bool = False, ) -> bool: """Return True if ``key`` should be redacted. Normalises common header-style variants so an LLM client passing ``X-Api-Key`` / ``X-Refresh-Token`` / ``Set-Cookie`` / ``Proxy-Authorization`` still hits the deny-list. Without this, ``key.lower() in _SENSITIVE_KEYS`` misses any name with hyphens or an ``x-`` prefix. The optional ``deny_set`` lets the outbound serializer share this normalization against ``_SERIALIZE_DENY_FIELDS``. ``deny_set_flat`` is the precomputed flat form (no ``_``/``-``) of the deny set — callers that share the input pre-compute it once at module load. ``strict=True`` skips the dash-split component scan. The component-scan branch is appropriate for INBOUND audit-args from arbitrary clients (HTTP-header names like ``Set-Cookie`` carry secret intent on the ``cookie`` component). On the OUTBOUND serializer path, application dict keys are ordinary snake_case Python identifiers (``last_token_rotation_at``, ``pin_location``, ``otp_required_flag``) — applying the component scan there masks legitimate fields whose name happens to CONTAIN a deny short word. Even in strict mode, the flat-deny lookup still catches ``Hashed-Password`` → ``hashedpassword`` matching the precomputed flat form of ``hashed_password`` — so common header-style credential names stay covered. """ deny = deny_set if deny_set is not None else _SENSITIVE_KEYS deny_flat = ( deny_set_flat if deny_set_flat is not None else _SENSITIVE_KEYS_FLAT if deny is _SENSITIVE_KEYS else _flatten_deny(deny) ) k = key.lower().strip() if k in deny: return True if k.startswith("x-") and k[2:] in deny: return True flat = k.replace("-", "").replace("_", "") if flat in deny_flat: return True if flat.startswith("x") and flat[1:] in deny_flat: return True if strict: return False for part in k.replace("_", "-").split("-"): if part and part in deny: return True return False def _summarise_value(v: Any, *, depth: int = 0, seen: set[int] | None = None) -> Any: """Recursively bound a single audit value. Strings > 500 chars are truncated. Lists / tuples > 50 elements drop the tail and add a marker. Dicts > 50 keys drop excess keys. Nested structures recurse with a depth + cycle guard so a self-referential or pathologically-deep input cannot RecursionError us mid-audit. Sensitive keys (``password``, ``secret``, ``token``, …) are redacted to ``"***"`` AT EVERY DEPTH, not just top-level — so e.g. ``answers={"password": "x"}`` doesn't bleed into ``audit_logs``. """ if seen is None: seen = set() if depth > _AUDIT_DEPTH_CAP: return "...[max depth]" if isinstance(v, str): return _redact_str(v) if isinstance(v, (list, tuple)): if id(v) in seen: return "...[cycle]" seen = seen | {id(v)} bounded = [ _summarise_value(x, depth=depth + 1, seen=seen) for x in v[:_AUDIT_LIST_CAP] ] if len(v) > _AUDIT_LIST_CAP: bounded.append(f"...[{len(v) - _AUDIT_LIST_CAP} more items]") return bounded if isinstance(v, dict): if id(v) in seen: return "...[cycle]" seen = seen | {id(v)} out: dict[str, Any] = {} for i, (k, vv) in enumerate(v.items()): if i >= _AUDIT_DICT_KEY_CAP: out["__truncated__"] = ( f"...[{len(v) - _AUDIT_DICT_KEY_CAP} more keys]" ) break key_s = str(k) if _is_sensitive_key(key_s): out[key_s] = "***" continue out[key_s] = _summarise_value(vv, depth=depth + 1, seen=seen) return out return v def _summarise_args(args: dict[str, Any]) -> dict[str, Any]: """Make tool inputs safe to drop in an audit log. * Truncates strings > 500 chars (auth tokens, pasted secrets). * Caps lists at 50 entries and dicts at 50 keys recursively so a malicious / buggy LLM client can't bloat ``audit_logs.detail`` to multi-MB by submitting a giant nested structure. * Recursive-depth + cycle guard prevents RecursionError on pathological inputs. * Strips C0 control chars + Unicode bidi overrides so audit rows cannot misrender in the admin UI. * Replaces obvious password / secret / token fields (at any depth) with ``***`` so a slip in tool naming doesn't surface plaintext. * Final safety: if the JSON-serialised payload still exceeds ``_AUDIT_TOTAL_CAP`` chars (~8KB), the whole thing is replaced with a marker dict — better to lose detail than DOS the audit table. """ out: dict[str, Any] = {} for k, v in args.items(): if _is_sensitive_key(k): out[k] = "***" continue out[k] = _summarise_value(v) # Rough total-size guard. We don't try to reach the JSON column # max_size; this is a sanity cap to keep individual audit rows # readable and to prevent OOM-on-write under abuse. import json try: serialised = json.dumps(out, default=str) except (TypeError, ValueError): serialised = "" if len(serialised) > _AUDIT_TOTAL_CAP: # Cap each key length so a malicious client can't sneak a 100k- # char key past the size guard via the marker payload itself. # ``_redact_str`` runs ctrl/RTL scrub + truncation. capped_keys = [ _redact_str(str(k))[:80] for k in list(out.keys())[:_AUDIT_DICT_KEY_CAP] ] return { "__oversize__": True, "__original_size__": len(serialised), "__keys__": capped_keys, } return out async def _safe_rollback(db: AsyncSession) -> None: """Best-effort rollback that never replaces the original error. On a degraded DB connection, ``db.rollback()`` itself can raise (the very condition that often produced the original failure). Without this guard, that secondary exception would replace the typed ``AppException`` the caller would otherwise see — turning a clean 4xx into a raw SQLAlchemy 500. """ try: await db.rollback() except Exception: # noqa: BLE001 logger.exception( "MCP db.rollback() failed — original error still propagates", ) async def _write_failure_audit( *, tool_name: str, trace_id, user_id, audit_args, error_code: str, error_message: str | None, client_ip, user_agent, raw_exception: str | None = None, ) -> None: """Best-effort failure-audit writer using a fresh session. Wrapped in a swallowing try/except: audit-DB outage MUST NOT mask the original error code we want to return to the caller. We log the audit failure and continue — the error itself still propagates. ``error_message`` may be None when an upstream exception type doesn't carry one — we coerce defensively so the audit still lands. """ safe_message = (error_message or "(no message)")[:1000] safe_tool = (tool_name or "unknown")[:120] try: async with _db.async_session_factory() as audit_db: detail: dict[str, Any] = { "args": audit_args, "error_code": error_code, "error_message": safe_message, } if raw_exception is not None: detail["raw_exception"] = raw_exception await AuditService(audit_db).log( trace_id=trace_id, actor_type="user", actor_user_id=user_id, action=f"mcp.{safe_tool}", status="failure", resource_type="mcp_tool", resource_id=safe_tool, detail=detail, ip_address=client_ip, user_agent=user_agent, ) await audit_db.commit() except Exception: # noqa: BLE001 logger.exception( "MCP failure-audit write failed for %s — original error " "still propagates to caller", safe_tool, ) @asynccontextmanager async def tool_runtime( tool_name: str, args: dict[str, Any] | None = None, *, mutates: bool = False, leaver_allowed: bool = False, ): """Open a DB session, load the user, write an audit row. Usage:: async with tool_runtime("list_my_assets", {}) as rt: rows = await AssetService(rt.db).get_user_assets(rt.user.id) return {"assets": [...]} On a non-mutating tool, the session never commits — any incidental writes (e.g. caching) roll back at end-of-block. On a mutating tool, the session commits ONCE the body returns successfully; any exception triggers a rollback and a failure audit row. """ user_id = require_mcp_user_id() trace_id = get_or_init_trace_id() client_ip = mcp_client_ip.get() user_agent = mcp_user_agent.get() audit_args = _summarise_args(args or {}) async with _db.async_session_factory() as db: # User-load can throw on transient DB errors. Without this # try/except the raw SQLAlchemy error string would propagate to # the MCP client (FastMCP wraps it in a ToolError that includes # the message). Catch + map to INTERNAL_ERROR so client-visible # error messages stay generic. The full trace still goes to logs. try: user = await _load_user_with_roles(db, user_id) except Exception as exc: # noqa: BLE001 logger.exception("MCP user-load failed for %s", user_id) await _write_failure_audit( tool_name=tool_name, trace_id=trace_id, user_id=user_id, audit_args=audit_args, error_code="INTERNAL_ERROR", error_message="User load failed.", raw_exception=f"{type(exc).__name__}"[:120], client_ip=client_ip, user_agent=user_agent, ) raise AppException( code="INTERNAL_ERROR", message="The tool encountered an unexpected error.", status_code=500, ) from exc # Combined "user gone or offboarded" check — middleware already # rejects offboarded users at PAT-auth time, but a race window # exists where HRIS sync flips employment_status to 'left' # between middleware and runtime. We re-validate here so a # brand-new MCP call from a just-offboarded user fails closed. from app.auth.service import is_effectively_departed from app.core.rbac import is_active_employee if user is None or is_effectively_departed(user): await _write_failure_audit( tool_name=tool_name, trace_id=trace_id, user_id=user_id, audit_args=audit_args, error_code="USER_NOT_FOUND", error_message="User no longer eligible.", client_ip=client_ip, user_agent=user_agent, ) raise AppException( code="USER_NOT_FOUND", message="The token's owner is no longer active.", status_code=401, ) # Leaver Blocked runtime gate for MUTATING tools — parity with # the REST `require_active_employee_for_mutations` dependency # in app.core.dependencies. # # ``leaver_allowed=True`` is the narrow carve-out for actions a # departing employee still MUST be able to take, where the REST # twin deliberately carries no gate either. Today that is only # attendance monthly-report confirmation: the report is the # payroll basis for the period being worked, and an intern whose # last day is booked (``leaver_pending``) still has to sign off # the final one. Without the flag the two surfaces answered the # same state differently — HTTP allowed it, MCP returned # LEAVER_RUNTIME_LOCKED (#2160). Anything added here needs the # same "and the HTTP twin is intentionally ungated, for this # reason" justification written down at BOTH call sites. # # On the two predicates: ``is_effectively_departed`` (used by # ``get_current_user``) only blocks status=='left' and # past-termination leaver_pending, whereas the REST mutation gate # is stricter (employment_status != 'active' OR deleted_at) — so # a leaver_pending user IS read-only via REST but could otherwise # mutate via MCP, which is what this check closes. Read-only # tools intentionally remain accessible so the user can see what # they own while wrapping up. if mutates and not leaver_allowed and not is_active_employee(user): await _write_failure_audit( tool_name=tool_name, trace_id=trace_id, user_id=user_id, audit_args=audit_args, error_code="LEAVER_RUNTIME_LOCKED", error_message=( "Account is in offboarding state — mutating " "actions disabled." ), client_ip=client_ip, user_agent=user_agent, ) raise AppException( code="LEAVER_RUNTIME_LOCKED", message=( "Your account is in offboarding state. Mutating " "actions are disabled while IT completes the " "handover. Read access remains available." ), status_code=403, ) runtime = ToolRuntime(db=db, user=user) # Note: success-audit writes happen in a FRESH session AFTER # the business commit (see mutating path below) — we no longer # share the tool-body session with the audit write, so we don't # instantiate ``AuditService(db)`` here. Failure audits go # through ``_write_failure_audit`` which also opens its own # session. try: yield runtime except AppException as exc: await _safe_rollback(db) await _write_failure_audit( tool_name=tool_name, trace_id=trace_id, user_id=user_id, audit_args=audit_args, error_code=exc.code, error_message=exc.message, client_ip=client_ip, user_agent=user_agent, ) raise except Exception as exc: # noqa: BLE001 # Pydantic ValidationError (subclass of ValueError) is the # standard surface for "the AI client passed garbage" — we # turn that into a VALIDATION_ERROR so the LLM gets a clear, # actionable hint instead of an opaque INTERNAL_ERROR. Real # crashes (typos, KeyError, …) still hit the catch-all and # surface as INTERNAL_ERROR with the full trace logged. from pydantic import ValidationError as PydanticValidationError await _safe_rollback(db) if isinstance(exc, PydanticValidationError): error_code = "VALIDATION_ERROR" error_message = "; ".join( f"{'.'.join(str(p) for p in (e.get('loc') or []))}: {e.get('msg', '')}" for e in exc.errors() )[:1000] or "Invalid input." status_code = 422 raw_exception = None else: logger.exception("MCP tool %s crashed", tool_name) error_code = "INTERNAL_ERROR" error_message = "The tool encountered an unexpected error." status_code = 500 # Apply the same redactor used for audit args so any # JWT / PAT / AKIA shape that leaked into the exception # stringification gets scrubbed before durable write. raw_exception = _redact_str( f"{type(exc).__name__}: {exc}", )[:500] await _write_failure_audit( tool_name=tool_name, trace_id=trace_id, user_id=user_id, audit_args=audit_args, error_code=error_code, error_message=error_message[:1000], raw_exception=raw_exception, client_ip=client_ip, user_agent=user_agent, ) raise AppException( code=error_code, message=error_message, status_code=status_code, ) from exc # Success path — write audit, commit (only commits writes for # mutating tools; reads have nothing to commit and the audit # row gets its own session below to avoid surprising the read # transaction with a write). # # Ordering: commit business writes FIRST, then write the audit # row. The previous order (audit-in-same-session → commit) had # a poisoned-session bug: if ``audit.log`` raised (column-length # / FK / constraint), the local ``except`` swallowed the error # and proceeded to ``await db.commit()`` — but the asyncmy # transaction was already in a failed state, so the commit # itself raised, the tool's business writes were rolled back, # AND we wrote two fallback audit rows (AUDIT_WRITE_DEGRADED + # COMMIT_FAILED) for what was really one event. By committing # business state first we cleanly separate "business write # succeeded" (the durable invariant the user cares about) from # "we successfully logged the event". An audit failure now # never reverts business writes; a commit failure now means # business writes are confirmed lost. (R1 review finding, # 2026-05-28.) if mutates: commit_ok = False try: await db.commit() commit_ok = True except Exception as exc: # noqa: BLE001 # Commit-time failure (deadlock revealed at commit, FK # check, lost connection, etc.) — business writes did # NOT land. Roll back the now-broken session, write a # separate-session failure audit (so forensics has a # trail), and re-raise as a uniform AppException. The # message explicitly says "safe to retry" because the # state is unchanged from the pre-call snapshot. logger.exception( "MCP tool %s commit failed", tool_name, ) await _safe_rollback(db) await _write_failure_audit( tool_name=tool_name, trace_id=trace_id, user_id=user_id, audit_args=audit_args, error_code="COMMIT_FAILED", error_message=( "Tool body completed but the database commit " "failed; no changes were applied." ), raw_exception=_redact_str( f"{type(exc).__name__}: {exc}", )[:500], client_ip=client_ip, user_agent=user_agent, ) raise AppException( code="COMMIT_FAILED", message=( "Tool body completed but the database commit " "failed; no changes were applied. Safe to retry." ), status_code=500, ) from exc # Business writes are now durable. Write the audit row in # a FRESH session so an audit-write failure never poisons # the (already-closed) business txn. This mirrors the # read-only path's audit-write pattern: we treat success # audits as a best-effort observability concern, and # surface lost ones through the failure-audit channel. if commit_ok: for callback in runtime._after_commit: try: await callback() except Exception: # noqa: BLE001 # The business commit is already durable. Post-commit # hooks are for best-effort side effects backed by a # durable intent (for example notification outbox # delivery), so they must never turn success into a # client-visible failure or a misleading safe-retry. logger.exception( "MCP tool %s post-commit callback failed", tool_name, ) try: async with _db.async_session_factory() as audit_db: await AuditService(audit_db).log( trace_id=trace_id, actor_type="user", actor_user_id=user_id, action=f"mcp.{tool_name}", status="success", resource_type="mcp_tool", resource_id=tool_name, detail={"args": audit_args}, ip_address=client_ip, user_agent=user_agent, ) await audit_db.commit() except Exception: # noqa: BLE001 logger.exception( "MCP tool %s success-audit write failed; " "business commit already landed", tool_name, ) await _write_failure_audit( tool_name=tool_name, trace_id=trace_id, user_id=user_id, audit_args=audit_args, error_code="AUDIT_WRITE_DEGRADED", error_message=( "Tool body succeeded and committed; the " "success audit row could not be written." ), client_ip=client_ip, user_agent=user_agent, ) else: # Read-only — audit in a separate session so the read txn # stays read-only from the DB's POV. # # Wrap in try/except: a transient audit-table commit failure # MUST NOT turn a successful read into a 500 to the caller — # the read result is already computed and the user expects # it. Log + best-effort failure-audit, never re-raise. # (Fix from launch-readiness audit 2026-05-05.) try: async with _db.async_session_factory() as audit_db: await AuditService(audit_db).log( trace_id=trace_id, actor_type="user", actor_user_id=user_id, action=f"mcp.{tool_name}", status="success", resource_type="mcp_tool", resource_id=tool_name, detail={"args": audit_args}, ip_address=client_ip, user_agent=user_agent, ) await audit_db.commit() except Exception: # noqa: BLE001 logger.exception( "MCP read-tool %s success-audit write failed; " "writing fallback audit row in fresh session", tool_name, ) await _write_failure_audit( tool_name=tool_name, trace_id=trace_id, user_id=user_id, audit_args=audit_args, error_code="AUDIT_WRITE_DEGRADED", error_message=( "Read tool succeeded; primary audit-row " "insert failed and was retried in a fresh " "session." ), client_ip=client_ip, user_agent=user_agent, ) _SERIALIZE_DEPTH_CAP = 12 # generous; actual nesting in our payloads ≤ 4 def _sanitize_outbound_string(s: str) -> str: """Strip control chars + Unicode bidi-override codepoints from a string returned to the MCP client. Why: tool return values feed directly into the calling LLM's context. Other-user-controlled free text (asset notes, transfer reasons, approval answers, notification bodies) is read by tool A but acted on by the LLM driving tool B. An adversary embedding NUL / CR / LF, or RTL/LRO override marks, can visually-rewrite the string in the LLM's view and steer downstream tool calls. The audit-write redactor (``_redact_str``) does the same on the INBOUND side (args). This is the symmetric outbound treatment — minus the truncation, since legitimate return values can be long (e.g. a full asset list page). We intentionally do NOT scrub credential-shaped substrings on return values: returns are already filtered by ``_SERIALIZE_DENY_FIELDS`` at the field-name level, and a tool that legitimately returns a PAT (only ``reveal_*`` does) MUST be the only path that does so. """ if not s: return s return s.translate(_CTRL_TABLE) # Field names that must NEVER cross the MCP wire even if a tool # accidentally returns a raw ORM row (e.g. an IdentityPerson, an # Integration carrying ``secrets_encrypted``, or a PersonalAccessToken # carrying ``token_hash``). The fallback ``__dict__`` walk below would # otherwise leak the bcrypt / AES-GCM ciphertext / raw PII to the LLM. # Filtered at every depth. # # This list is a SUPERSET of ``_SENSITIVE_KEYS`` (audit-args redactor) # so a future tool that returns a structure containing any audit-redacted # key gets the same protection on the wire. Keep them aligned: when you # add an entry to ``_SENSITIVE_KEYS``, mirror it here. _SERIALIZE_DENY_FIELDS = frozenset({ # Passwords / passphrases "password", "hashed_password", "new_password", "old_password", "passphrase", # Secrets / signing material "secret", "secret_key", "client_secret", "webhook_secret", "signing_secret", "signing_key", "private_key", "encrypted_value", # Mirror ``_SENSITIVE_KEYS`` addition — Feishu driver config dicts # store the tenant's plaintext ``app_secret`` and we must never # serialize that to the LLM. (R3 security finding 2026-05-28.) "app_secret", # Real ORM column names that carry credential ciphertext (do NOT # remove without checking app/integrations/models.py and # app/auth/tokens/models.py — these are the actual DB columns). "secrets_encrypted", "secrets_nonce", "token_hash", "key_value", "key_hash", # Tokens "token", "access_token", "refresh_token", "id_token", "session_token", "session_key", "pat", "pat_value", "pat_hash", "personal_access_token", "oauth_token", # PAT-mint plaintext field names — the POST /tokens response # exposes ``plaintext`` once. If a future MCP tool wraps PAT # creation and returns the row through a Pydantic schema, the # plaintext field name MUST be in this list to keep it from # crossing the wire to the LLM. Mirrors ``_SENSITIVE_KEYS``. "plaintext", "pat_plaintext", "token_plaintext", "secret_value", "auth_token", "raw_token", "encrypted_token", # API / cloud credentials "api_key", "apikey", "access_key", "credential", "credentials", "service_account", "gcp_credentials", # Headers / cookies (rarely surface, but defense-in-depth) "authorization", "bearer", "cookie", "csrf", # MFA / recovery "totp_secret", "mfa_secret", "recovery_code", "pin", "otp", "verification_code", # NOTE: PII column names like ``mobile_phone`` / ``personal_email`` # are intentionally NOT in this deny-list. The legitimate-reveal # surface (``reveal_my_contact``) returns those fields by exact # name and would be broken by masking. Defense against a tool # accidentally returning a raw IdentityPerson row is layered # elsewhere: every Person-returning tool today builds a manual # dict; the serializer's ORM-row fallback (``__dict__`` walk) is # only hit on a future-tool regression and the audit-log redactor # still scrubs token-shaped substrings. Adding PII column names # to the deny-list would over-rotate. }) # Pre-computed flat form (no `_` / no `-`) of the outbound deny set so # ``Hashed-Password`` (header style) matches ``hashed_password`` (ORM # style) via a single set lookup without falling back to the dash-split # component scan — that scan over-rotates on legitimate field names # containing a short deny word like ``token`` / ``cookie`` / ``pin`` # (see ``_is_sensitive_key`` doc for the rationale). _SERIALIZE_DENY_FIELDS_FLAT = _flatten_deny(_SERIALIZE_DENY_FIELDS) def serialize(value: Any, *, _depth: int = 0, _seen: set[int] | None = None) -> Any: """JSON-safe coercion used by tool return values. Pydantic models → ``model_dump(mode="json")``. SQLAlchemy ORM rows fall through to ``__dict__`` minus internal state. Lists/dicts recurse with a depth + cycle guard so a self-referential ORM relationship or a malformed dict can't RecursionError us. Anything else (UUID / datetime / Decimal) becomes a str. """ from datetime import date, datetime from decimal import Decimal from uuid import UUID from pydantic import BaseModel if value is None: return None if isinstance(value, BaseModel): # Recurse through model_dump so the deny-list filter on the # dict branch below applies. Without this, a Pydantic schema # exposing a denylisted field name (e.g. ``hashed_password``) # would short-circuit straight to plaintext output. return serialize( value.model_dump(mode="json"), _depth=_depth + 1, _seen=_seen if _seen is not None else set(), ) if isinstance(value, (UUID, Decimal)): return str(value) if isinstance(value, datetime): # MySQL DATETIME(timezone=True) round-trips to a naive datetime # on fresh load even when the column declares ``DateTime(tz=True)`` # (see project memory ``feedback_mysql_datetime_naive_on_fresh_load``). # ``.isoformat()`` on a naive datetime emits an unsuffixed string # which LLM clients / downstream charts silently treat as local # time. Coerce to UTC at the serialization boundary so every # outbound timestamp is unambiguously tz-tagged. (R1 review # finding, 2026-05-28.) if value.tzinfo is None: value = value.replace(tzinfo=UTC) return value.isoformat() if isinstance(value, date): return value.isoformat() if isinstance(value, str): return _sanitize_outbound_string(value) if isinstance(value, (int, float, bool)): return value # Depth + cycle guard apply to recursive types only. Visit-set # tracks identity (id()) since dicts / ORM rows aren't hashable. if _depth > _SERIALIZE_DEPTH_CAP: return "...[max depth]" if _seen is None: _seen = set() if isinstance(value, (list, tuple)): if id(value) in _seen: return "...[cycle]" next_seen = _seen | {id(value)} return [ serialize(v, _depth=_depth + 1, _seen=next_seen) for v in value ] if isinstance(value, dict): if id(value) in _seen: return "...[cycle]" next_seen = _seen | {id(value)} return { _sanitize_outbound_string(str(k)): ( "***" if _is_sensitive_key( str(k), deny_set=_SERIALIZE_DENY_FIELDS, deny_set_flat=_SERIALIZE_DENY_FIELDS_FLAT, strict=True, ) else serialize(v, _depth=_depth + 1, _seen=next_seen) ) for k, v in value.items() } # SQLAlchemy ORM rows: pull public attributes off ``__dict__`` and # recurse. Internal attributes (anything beginning with ``_``, # notably ``_sa_instance_state``) are filtered out so we don't # accidentally serialize SQLAlchemy bookkeeping into the response. # Sensitive field names (hashed_password, key_value, refresh_token, # …) are masked even if a tool accidentally hands back a raw ORM # row — defense in depth against future contributors who forget # to project a row through a Pydantic schema. state = getattr(value, "__dict__", None) if state: if id(value) in _seen: return "...[cycle]" next_seen = _seen | {id(value)} return { k: ( "***" if _is_sensitive_key( k, deny_set=_SERIALIZE_DENY_FIELDS, deny_set_flat=_SERIALIZE_DENY_FIELDS_FLAT, strict=True, ) else serialize(v, _depth=_depth + 1, _seen=next_seen) ) for k, v in state.items() if not k.startswith("_") } # Anything else (NamedTuple, custom class without __dict__, etc.) # is rendered as its ``str()`` form. We deliberately don't return # ``None`` here because a silent drop would lose data; ``str()`` # at least surfaces a diagnostic representation that the AI # client can show the user. Same control-char / RTL scrub as the # str-branch so unexpected types can't bypass output sanitization. return _sanitize_outbound_string(str(value)) __all__ = ["ToolRuntime", "_safe_rollback", "serialize", "tool_runtime"]