Address Polly review: confidential_files is containment-only, not a write grant

Revert the write-scope widening that let any file listed in confidential_files be written/deleted even if the agent never created it and it isn't in write_files. confidential_files is now purely a containment declaration: writing to a confidential file still requires it to be created this session or in write_files, matching the pre-existing write boundary. The demo CUJ is unaffected (it writes to a doc the agent created this session). Also document that the read-latch engages only on reads that name a confidential file by id — content-returning reads that don't target a specific file (drive_search, listing, exports) can surface confidential text without engaging containment. Update tests to the corrected semantics and add a guard that declaring a file confidential does not by itself grant write access.
Add Bell-LaPadula "no write-down" to gdrive_policy; fix MCP field/tool gaps
2026-07-01 07:58:09 +00:00 · 2026-07-01 07:29:30 +00:00
8 changed files with 697 additions and 62 deletions
@@ -102,9 +102,22 @@ _FILE_ID_ARG_KEYS: tuple[str, ...] = (
    "presentation_id",
 )

-# Result keys whose string value is a newly-created file ID.
+# Result keys whose string value is a newly-created file ID. Covers both the
+# raw Google API camelCase (``documentId``) and the snake_case a wrapping MCP
+# may re-emit (``document_id``) — the Databricks Google MCP filters create
+# results down to snake_case fields, so both spellings must be recognized.
 _FILE_RESULT_ID_KEYS: frozenset[str] = frozenset(
-    {"id", "documentId", "spreadsheetId", "presentationId", "fileId"}
+    {
+        "id",
+        "documentId",
+        "spreadsheetId",
+        "presentationId",
+        "fileId",
+        "document_id",
+        "spreadsheet_id",
+        "presentation_id",
+        "file_id",
+    }
 )

 # Max recursion depth when scanning a tool-result payload for created IDs.
@@ -161,6 +174,7 @@ _DRIVE_COMMENT_TOOLS: frozenset[str] = frozenset(
 _DRIVE_WRITE_TOOLS: frozenset[str] = frozenset(
    {
        "docs_document_batch_update",
+        "docs_document_edit_section",
        "docs_document_apply_template_styles",
        "drive_file_update",
        "drive_file_delete",
@@ -174,6 +188,29 @@ _DRIVE_WRITE_TOOLS: frozenset[str] = frozenset(
 # Canonical-name prefixes the Drive policy owns (unrecognized → fail closed).
 _DRIVE_OWNED_PREFIXES: tuple[str, ...] = ("docs_", "drive_", "sheets_", "slides_")

+# ── Bell-LaPadula "no write-down" constants ────────────────────────────────────
+#
+# Bell-LaPadula is the classic lattice-based *confidentiality* model. Its
+# "*-property" (star property) forbids a **write down**: a subject cleared for a
+# sensitive level must not write to a less-classified object, because that moves
+# classified data into a less-protected place (a leak).
+#
+# We model a simple two-level lattice — *confidential* vs. everything else — from
+# a caller-supplied allowlist of confidential file IDs (``confidential_files``).
+# There is intentionally no dependence on any per-document classification label:
+# the set of confidential documents is declared explicitly, so the policy works
+# on any Drive tenant. Once the session has read a confidential file, its writes
+# are confined to the confidential set (a write elsewhere would be a write-down).
+#
+# Scope: this enforces only the write-down rule. It deliberately does NOT enforce
+# the "simple security property" (no read-up) — the agent must be able to read
+# confidential files for the containment to engage. It is a one-directional
+# write-side guard, not a full multilevel-security kernel.
+
+# Session-state key: has the session read any confidential file this session?
+# Public: it surfaces in the conversation's persisted ``session_state``.
+READ_CONFIDENTIAL_STATE_KEY = "gdrive_read_confidential"
+
 # ── Gmail tool sets ───────────────────────────────────────────────────────────

 # Arg keys carrying a target draft ID on a draft update / delete call.
@@ -454,6 +491,11 @@ class _DriveCfg:
    :param write_ids: Normalized file IDs writable regardless of creation.
    :param comment_ids: Normalized file IDs the agent may comment on.
    :param allow_create: Whether new-file creation is permitted.
+    :param confidential_ids: Normalized file IDs that form the confidential
+        compartment. Non-empty enables Bell-LaPadula "no write-down": once the
+        session reads any of these, writes are confined to this set.
+    :param write_down_action: Verdict on a write-down violation — ``"DENY"``
+        (default) or ``"ASK"`` (human approval).
    :param deny_reason: Reason prefix for DENY decisions.
    """

@@ -462,9 +504,88 @@ class _DriveCfg:
    write_ids: set[str]
    comment_ids: set[str]
    allow_create: bool
+    confidential_ids: set[str]
+    write_down_action: str
    deny_reason: str


+def _escalate(cfg: _DriveCfg, reason: str) -> PolicyResponse:
+    """
+    Build the configured write-down escalation response (DENY or ASK).
+
+    :param cfg: Resolved Drive configuration (carries ``write_down_action``).
+    :param reason: Human-readable explanation of the violated rule.
+    :returns: A :class:`PolicyResponse` with the configured verdict.
+    """
+    action = cfg.write_down_action
+    return {"result": action, "reason": f"{cfg.deny_reason} {reason}"}  # type: ignore[typeddict-item]
+
+
+def _read_confidential_update(event: PolicyEvent, cfg: _DriveCfg) -> list[StateUpdateEntry]:
+    """
+    Flag the session as having read a confidential file, if this read is one.
+
+    Inspects the originating read's target file IDs (from ``request_data``) and,
+    when any is in the configured confidential compartment, emits a state-update
+    setting :data:`READ_CONFIDENTIAL_STATE_KEY`. This is the "clearance rises on
+    read" step of Bell-LaPadula, reduced to a single latch (the compartment is
+    two-level: confidential vs. not).
+
+    :param event: A ``tool_result`` event for a Drive read tool.
+    :param cfg: Resolved Drive configuration (holds ``confidential_ids``).
+    :returns: A one-element ``state_updates`` list when a confidential file was
+        read (and the latch isn't already set), else empty.
+    """
+    if not cfg.confidential_ids:
+        return []
+    session_state = event.get("session_state") or {}
+    if session_state.get(READ_CONFIDENTIAL_STATE_KEY):
+        return []  # already latched — nothing to update
+    request_data = event.get("request_data")
+    args = request_data.get("arguments") if isinstance(request_data, dict) else None
+    read_ids = _extract_ids_from_args(args if isinstance(args, dict) else {}, _FILE_ID_ARG_KEYS)
+    if read_ids & cfg.confidential_ids:
+        return [{"key": READ_CONFIDENTIAL_STATE_KEY, "action": "set", "value": True}]
+    return []
+
+
+def _check_no_write_down(
+    target_ids: set[str],
+    event: PolicyEvent,
+    cfg: _DriveCfg,
+) -> PolicyResponse | None:
+    """
+    Apply Bell-LaPadula's "no write-down" rule to a write / comment / create.
+
+    Only meaningful once ``confidential_ids`` is configured. When the session
+    has read a confidential file, a write is allowed only if its target is
+    itself in the confidential compartment; any other target (a less-protected
+    file, or a brand-new one) is a write-down and is DENYed (or ASKed). Returns
+    ``None`` to abstain (session hasn't read confidential material, or the write
+    stays inside the compartment).
+
+    :param target_ids: Normalized target file IDs from the call (empty for a
+        create, which has no pre-existing target).
+    :param event: The ``tool_call`` event.
+    :param cfg: Resolved Drive configuration.
+    :returns: A DENY/ASK :class:`PolicyResponse` on violation, else ``None``.
+    """
+    if not cfg.confidential_ids:
+        return None
+    session_state = event.get("session_state") or {}
+    if not session_state.get(READ_CONFIDENTIAL_STATE_KEY):
+        return None  # no confidential data read yet — no containment
+    # A write stays legal only when every target is inside the compartment.
+    if target_ids and target_ids <= cfg.confidential_ids:
+        return None
+    return _escalate(
+        cfg,
+        "Bell-LaPadula (no write-down): this session has read a confidential "
+        "document, so writes are confined to the confidential set; this target "
+        "is outside it.",
+    )
+
+
 def _decide_drive_tool_call(
    event: PolicyEvent, prefixes: tuple[str, ...], cfg: _DriveCfg
 ) -> PolicyResponse | None:
@@ -493,6 +614,19 @@ def _decide_drive_tool_call(
            f"{cfg.deny_reason} Read restricted to the configured allowlist; "
            f"this call targets a file outside it (or cannot be scoped)."
        )
+    # Bell-LaPadula "no write-down" runs *before* the access-scope rules on any
+    # tool that emits data into a file: once the session has read confidential
+    # material, a write-down leak is blocked even if the file is otherwise
+    # writable (e.g. one the agent created this session).
+    if cfg.confidential_ids and (
+        canonical in _DRIVE_WRITE_TOOLS
+        or canonical in _DRIVE_COMMENT_TOOLS
+        or canonical in _DRIVE_CREATE_TOOLS
+    ):
+        violation = _check_no_write_down(target_ids, event, cfg)
+        if violation is not None:
+            return violation
+
    if canonical in _DRIVE_CREATE_TOOLS:
        return (
            None
@@ -506,6 +640,11 @@ def _decide_drive_tool_call(
    if canonical in _DRIVE_WRITE_TOOLS:
        if not target_ids:
            return _deny(f"{cfg.deny_reason} Write call carries no identifiable target file.")
+        # ``confidential_files`` is purely a containment declaration: it does NOT
+        # by itself grant write access. Writing to a confidential file still
+        # requires it to be created this session or in ``write_files`` — so a
+        # confidential doc the agent created stays writable (until it reads
+        # another confidential file, which the no-write-down check above gates).
        if target_ids <= (cfg.write_ids | created_ids):
            return None
        extra = " or the configured write allowlist" if cfg.write_ids else ""
@@ -543,6 +682,40 @@ def _record_created_drive(event: PolicyEvent, prefixes: tuple[str, ...]) -> Poli
    return {"result": "ALLOW", "state_updates": updates}


+def _record_drive_result(
+    event: PolicyEvent, prefixes: tuple[str, ...], cfg: _DriveCfg
+) -> PolicyResponse | None:
+    """
+    Handle a Drive ``tool_result``: track created IDs and confidential reads.
+
+    Merges two independent state effects into one response so both survive:
+
+    - Created-file tracking (:func:`_record_created_drive`) — always on.
+    - Confidential-read latch (:func:`_read_confidential_update`) — only when
+      ``confidential_ids`` is configured; flags the session once it reads a file
+      from the confidential compartment.
+
+    :param event: A ``tool_result`` event.
+    :param prefixes: Server prefixes for canonicalization.
+    :param cfg: Resolved Drive configuration.
+    :returns: ALLOW with the combined ``state_updates``, or ``None`` when there
+        is nothing to record.
+    """
+    created = _record_created_drive(event, prefixes)
+    updates: list[StateUpdateEntry] = list(created["state_updates"]) if created else []
+
+    if cfg.confidential_ids:
+        raw_tool = event.get("target")
+        if isinstance(raw_tool, str) and _canonical_tool_name(raw_tool, prefixes) in (
+            _DRIVE_READ_TOOLS
+        ):
+            updates.extend(_read_confidential_update(event, cfg))
+
+    if not updates:
+        return None
+    return {"result": "ALLOW", "state_updates": updates}
+
+
 def gdrive_policy(
    *,
    read_all: bool = True,
@@ -550,6 +723,8 @@ def gdrive_policy(
    allow_create: bool = False,
    write_files: list[str] | None = None,
    comment_files: list[str] | None = None,
+    confidential_files: list[str] | None = None,
+    write_down_action: str = "DENY",
    tool_prefixes: list[str] | None = None,
    deny_reason: str = "Google Drive operation blocked by policy.",
 ) -> Callable[[PolicyEvent], PolicyResponse | None]:
@@ -566,31 +741,59 @@ def gdrive_policy(
        means none.
    :param comment_files: File IDs / URLs the agent may comment on (in addition
        to files it created). ``None`` means none.
+    :param confidential_files: File IDs / Google URLs that form the confidential
+        compartment. When non-empty, this layers Bell-LaPadula's classic
+        confidentiality rule — the "*-property", i.e. **no write down** — on top
+        of the access rules: once the session reads any file in this set, its
+        writes are confined to the same set. A write to any other file (or a
+        brand-new one) would move confidential data into a less-protected place
+        and is blocked. The compartment is declared explicitly (rather than
+        inferred from a per-document label), so it works on any Drive tenant.
+        ``None`` / empty means the rule is off and the base access policy behaves
+        exactly as before. This enforces only the write-down rule; it does NOT
+        restrict reads (the agent must read a confidential file for the
+        containment to engage) and does NOT by itself grant write access to the
+        listed files — writing to a confidential file still requires it to be
+        created this session or in ``write_files``. Note the latch engages only
+        on reads that target a confidential file *by id*; content-returning
+        reads that don't name a specific file (e.g. ``drive_search``, listing,
+        or exports) can surface confidential text without engaging containment.
+    :param write_down_action: Verdict on a write-down violation — ``"DENY"``
+        (default, hard block) or ``"ASK"`` (human approval). Ignored when
+        ``confidential_files`` is empty.
    :param tool_prefixes: Server prefixes to strip when canonicalizing tool
        names. ``None`` uses the standard + Databricks defaults.
    :param deny_reason: Reason text attached to DENY decisions.
    :returns: A one-argument policy callable.
+    :raises ValueError: If ``write_down_action`` is not ``"DENY"`` or ``"ASK"``.
    """
+    normalized_action = write_down_action.strip().upper()
+    if normalized_action not in {"DENY", "ASK"}:
+        raise ValueError(
+            f"gdrive_policy: write_down_action must be 'DENY' or 'ASK', got {write_down_action!r}"
+        )
    cfg = _DriveCfg(
        read_all=read_all,
        read_ids=_normalize_file_refs(read_files),
        write_ids=_normalize_file_refs(write_files),
        comment_ids=_normalize_file_refs(comment_files),
        allow_create=allow_create,
+        confidential_ids=_normalize_file_refs(confidential_files),
+        write_down_action=normalized_action,
        deny_reason=deny_reason,
    )
    prefixes = _resolve_prefixes(tool_prefixes)

    def _evaluate(event: PolicyEvent) -> PolicyResponse | None:
        """
-        Route a Drive event: record created files on results, gate tool calls.
+        Route a Drive event: track created files / confidential reads, gate calls.

        :param event: The policy event.
        :returns: A :class:`PolicyResponse`, or ``None`` to abstain.
        """
        phase = event.get("type")
        if phase == "tool_result":
-            return _record_created_drive(event, prefixes)
+            return _record_drive_result(event, prefixes, cfg)
        if phase == "tool_call":
            return _decide_drive_tool_call(event, prefixes, cfg)
        return None
@@ -864,7 +1067,10 @@ POLICY_REGISTRY: list[dict[str, Any]] = [  # type: ignore[explicit-any]
            "Controls access to Google Drive files, Docs, Sheets, and Slides through "
            "any Google MCP server. Restricts reads to an allowlist and restricts "
            "writes/comments to files the agent created this session plus explicitly "
-            "allowed files."
+            "allowed files. Optionally enforces Bell-LaPadula's 'no write-down' rule "
+            "via a confidential-file compartment: once the session reads a confidential "
+            "file, its writes are confined to that set so classified data can't leak "
+            "into a less-protected file."
        ),
        "params_schema": {
            "type": "object",
@@ -894,6 +1100,21 @@ POLICY_REGISTRY: list[dict[str, Any]] = [  # type: ignore[explicit-any]
                    "items": {"type": "string"},
                    "description": "File IDs or URLs the agent may comment on.",
                },
+                "confidential_files": {
+                    "type": "array",
+                    "items": {"type": "string"},
+                    "description": "File IDs or Google URLs forming the confidential "
+                    "compartment. When set, Bell-LaPadula 'no write-down' engages: "
+                    "after the session reads one of these, writes are confined to the "
+                    "set. Empty (default) disables the rule.",
+                },
+                "write_down_action": {
+                    "type": "string",
+                    "enum": ["DENY", "ASK"],
+                    "description": "Verdict on a write-down violation. "
+                    "Ignored unless confidential_files is set.",
+                    "default": "DENY",
+                },
                "tool_prefixes": {
                    "type": "array",
                    "items": {"type": "string"},
@@ -5,12 +5,16 @@ implements the "session risk score" pattern from the sample-policies wishlist:

 - **Accrue risk** as the agent takes risky actions. Two sources:
  - *Tool calls* — each call to a configured tool adds points
-    (``tool_points``), e.g. a ``web_search`` adds 10.
+    (``tool_points``), e.g. a ``web_search`` adds 10. This source needs no
+    special support from any tool, so it works out of the box.
  - *Tool results* — a result carrying a sensitive data-classification label
-    adds points (``sensitive_labels``), e.g. reading a doc tagged
+    adds points (``sensitive_labels``), e.g. a result tagged
    ``"Highly Confidential"`` adds 30. The label is read out of the tool's
-    result payload, so this works for *any* MCP server whose results carry a
-    classification field (Google Drive returns ``label_classification``).
+    result payload (scanning the keys in ``label_keys``), so this works with
+    *any* MCP server that annotates its results with a classification field —
+    configure ``sensitive_labels`` to match whatever values that server emits.
+    (Not every MCP labels its output; when none do, leave ``sensitive_labels``
+    empty and rely on ``tool_points`` alone.)
 - **Gate sensitive actions** once the accrued score crosses ``threshold``:
  configured ``guarded_tools`` (e.g. ``gmail_message_send``) escalate from
  ALLOW to **ASK** (default) or **DENY**, forcing human oversight on a session
@@ -44,6 +48,7 @@ YAML usage::
          arguments:
            threshold: 50
            tool_points: {web_search: 10, fetch: 5}
+            # Only if your MCP annotates results with these exact label values.
            sensitive_labels: {"Highly Confidential": 30, RESTRICTED: 30}
            guarded_tools: [gmail_message_send, drive_permission_create]
            escalate_action: ASK
@@ -169,6 +169,10 @@ _ALT_COVERED: frozenset[str] = frozenset(
        # risk_score_agent: the built-in session-risk-score policy is
        # exercised in tests/runtime/policies/test_example_omnigent_yamls.py.
        "risk_score_agent",
+        # info_flow_agent: the built-in gdrive information-flow policy
+        # (Bell-LaPadula + Biba) is exercised in
+        # tests/runtime/policies/test_example_omnigent_yamls.py.
+        "info_flow_agent",
        # qwen_perm_test: qwen-harness permission fixture exercised by
        # tests/inner/test_qwen_agent_integration.py against a mocked ACP
        # subprocess. The live qwen round-trip lives in
@@ -78,6 +78,7 @@ def tool_result_event(
    tool: str,
    result: str,
    session_state: dict[str, Any] | None = None,  # type: ignore[explicit-any]
+    request_arguments: dict[str, Any] | None = None,  # type: ignore[explicit-any]
 ) -> PolicyEvent:
    """
    Build a ``tool_result`` :class:`PolicyEvent` (server-side shape).
@@ -87,12 +88,21 @@ def tool_result_event(
    :param result: Stringified tool output under ``data.result``, e.g.
        ``'{"documentId": "1New"}'``.
    :param session_state: Optional persisted state. ``None`` means empty.
+    :param request_arguments: Optional arguments of the originating tool call,
+        surfaced under ``request_data`` (the ``{"name", "arguments"}`` shape the
+        server passes on ``tool_result``). Lets a result policy correlate the
+        response with the request — e.g. the info-flow policy learning which
+        file a read returned a classification for. ``None`` omits
+        ``request_data``.
    :returns: A ``tool_result`` event dict.
    """
-    return {
+    event: PolicyEvent = {
        "type": "tool_result",
        "target": tool,
        "data": {"result": result},
        "context": {"actor": {}, "usage": {}},
        "session_state": session_state or {},
    }
+    if request_arguments is not None:
+        event["request_data"] = {"name": tool, "arguments": request_arguments}
+    return event
@@ -171,6 +171,47 @@ def test_drive_write_to_created_file_allowed() -> None:
    assert result is None


+def test_drive_create_result_snake_case_id_recorded() -> None:
+    """A create result using snake_case ``document_id`` is recorded.
+
+    Regression guard: the Databricks Google MCP filters create results down to
+    snake_case id fields (``document_id``), not the raw camelCase
+    ``documentId``. If only camelCase were recognized, the created file would
+    not be tracked and a later write to it would be wrongly denied.
+    """
+    result = gdrive_policy()(
+        tr(
+            "mcp__google__docs_document_create",
+            '{"service": "docs", "operation": "create", "document_id": "1New"}',
+        )
+    )
+    assert result is not None
+    assert result["state_updates"] == [
+        {"key": CREATED_FILES_STATE_KEY, "action": "append", "value": "1New"}
+    ]
+
+
+def test_drive_edit_section_is_a_write_tool() -> None:
+    """``docs_document_edit_section`` is treated as a write (scoped like others).
+
+    Regression guard: the Docs content-edit tool must be recognized as a write,
+    not fall through to the "unknown tool" fail-closed branch. It is scoped by
+    its ``document_id`` arg, so a write to a created file is allowed and one to
+    an unowned file is denied.
+    """
+    policy = gdrive_policy()
+    allowed = policy(
+        tc(
+            "mcp__google__docs_document_edit_section",
+            {"document_id": "1New"},
+            {CREATED_FILES_STATE_KEY: ["1New"]},
+        )
+    )
+    assert allowed is None
+    denied = policy(tc("mcp__google__docs_document_edit_section", {"document_id": "1Foreign"}))
+    assert denied is not None and denied["result"] == "DENY"
+
+
 def test_drive_write_to_uncreated_file_denied() -> None:
    """A write to a file the agent did not create (nor allowlisted) is denied."""
    policy = gdrive_policy()
@@ -350,6 +391,208 @@ async def test_drive_created_file_roundtrip(
    assert write_foreign.action == PolicyAction.DENY


+# ── gdrive_policy: Bell-LaPadula "no write-down" (confidential compartment) ────
+
+_CONF_ID = "1ConfidentialDocABCDEFGHIJKLMNOPQRSTUV"
+_OTHER_ID = "1OtherDocABCDEFGHIJKLMNOPQRSTUVWXYZ0123"
+
+
+def test_no_write_down_off_by_default() -> None:
+    """With no ``confidential_files``, reads and writes are unconstrained.
+
+    Guards backward-compat: the compartment rule is inert unless
+    ``confidential_files`` is configured.
+    """
+    policy = gdrive_policy()  # confidential_files defaults to empty
+    # A read of any file records nothing.
+    assert (
+        policy(
+            tr(
+                "mcp__google__docs_document_get",
+                "{}",
+                request_arguments={"document_id": _CONF_ID},
+            )
+        )
+        is None
+    )
+
+
+def test_no_write_down_reading_confidential_latches_state() -> None:
+    """Reading a confidential file flags the session's confidential-read latch."""
+    policy = gdrive_policy(confidential_files=[_CONF_ID])
+    result = policy(
+        tr(
+            "mcp__google__docs_document_get",
+            "{}",
+            request_arguments={"document_id": _CONF_ID},
+        )
+    )
+    assert result is not None and result["result"] == "ALLOW"
+    assert result["state_updates"] == [
+        {"key": "gdrive_read_confidential", "action": "set", "value": True}
+    ]
+
+
+def test_no_write_down_reading_non_confidential_does_not_latch() -> None:
+    """Reading a file outside the compartment leaves the latch unset."""
+    policy = gdrive_policy(confidential_files=[_CONF_ID])
+    result = policy(
+        tr(
+            "mcp__google__docs_document_get",
+            "{}",
+            request_arguments={"document_id": _OTHER_ID},
+        )
+    )
+    assert result is None  # nothing recorded
+
+
+def test_no_write_down_latch_is_idempotent() -> None:
+    """A second confidential read does not re-emit the latch update."""
+    policy = gdrive_policy(confidential_files=[_CONF_ID])
+    result = policy(
+        tr(
+            "mcp__google__docs_document_get",
+            "{}",
+            session_state={"gdrive_read_confidential": True},
+            request_arguments={"document_id": _CONF_ID},
+        )
+    )
+    assert result is None  # already latched — no duplicate update
+
+
+def test_no_write_down_write_outside_compartment_denied() -> None:
+    """After reading confidential, a write to an outside file is denied."""
+    policy = gdrive_policy(confidential_files=[_CONF_ID])
+    state = {"gdrive_read_confidential": True}
+    result = policy(tc("mcp__google__drive_file_update", {"file_id": _OTHER_ID}, state))
+    assert result is not None and result["result"] == "DENY"
+    assert "no write-down" in result["reason"]
+
+
+def test_no_write_down_write_inside_compartment_allowed() -> None:
+    """After reading confidential, a write to a confidential file the agent may
+    write (here, in ``write_files``) is allowed — it stays inside the set.
+
+    The no-write-down check does not fire (target is confidential), and the base
+    write rule permits it because the file is explicitly writable.
+    """
+    policy = gdrive_policy(confidential_files=[_CONF_ID], write_files=[_CONF_ID])
+    state = {"gdrive_read_confidential": True}
+    assert policy(tc("mcp__google__drive_file_update", {"file_id": _CONF_ID}, state)) is None
+
+
+def test_no_write_down_confidential_does_not_grant_write() -> None:
+    """Declaring a file confidential does not by itself make it writable.
+
+    Guards against a boundary-widening side effect: a confidential file the
+    agent neither created nor has in ``write_files`` is still denied by the base
+    write rule (``confidential_files`` is a containment declaration, not a
+    write grant).
+    """
+    policy = gdrive_policy(confidential_files=[_CONF_ID])
+    result = policy(tc("mcp__google__drive_file_update", {"file_id": _CONF_ID}))
+    assert result is not None and result["result"] == "DENY"
+
+
+def test_no_write_down_create_after_confidential_denied() -> None:
+    """After reading confidential, creating a new (outside) file is denied.
+
+    A brand-new file is outside the confidential compartment, so writing into it
+    would move confidential data into a less-protected place.
+    """
+    policy = gdrive_policy(confidential_files=[_CONF_ID], allow_create=True)
+    state = {"gdrive_read_confidential": True}
+    result = policy(tc("mcp__google__docs_document_create", {"title": "leak"}, state))
+    assert result is not None and result["result"] == "DENY"
+    assert "no write-down" in result["reason"]
+
+
+def test_no_write_down_ask_action_escalates_instead_of_denying() -> None:
+    """``write_down_action='ASK'`` turns a violation into an approval prompt."""
+    policy = gdrive_policy(confidential_files=[_CONF_ID], write_down_action="ASK")
+    state = {"gdrive_read_confidential": True}
+    result = policy(tc("mcp__google__drive_file_update", {"file_id": _OTHER_ID}, state))
+    assert result is not None and result["result"] == "ASK"
+
+
+def test_no_write_down_no_confidential_read_yet_allows_write() -> None:
+    """Before reading any confidential file, writes are unconstrained by the rule.
+
+    The write still passes the base access rules because the file was created
+    this session.
+    """
+    policy = gdrive_policy(confidential_files=[_CONF_ID])
+    state = {CREATED_FILES_STATE_KEY: [_OTHER_ID]}
+    assert policy(tc("mcp__google__drive_file_update", {"file_id": _OTHER_ID}, state)) is None
+
+
+def test_no_write_down_invalid_action_rejected() -> None:
+    """A bad ``write_down_action`` is rejected at factory-build time."""
+    with pytest.raises(ValueError, match="write_down_action"):
+        gdrive_policy(confidential_files=[_CONF_ID], write_down_action="MAYBE")
+
+
+def test_no_write_down_confidential_files_accepts_urls() -> None:
+    """A Google URL in ``confidential_files`` matches a call targeting the bare ID."""
+    url = f"https://docs.google.com/document/d/{_CONF_ID}/edit"
+    policy = gdrive_policy(confidential_files=[url])
+    # Reading via the bare ID latches, proving the URL normalized to that ID.
+    result = policy(
+        tr(
+            "mcp__google__docs_document_get",
+            "{}",
+            request_arguments={"document_id": _CONF_ID},
+        )
+    )
+    assert result is not None
+    assert result["state_updates"][0]["key"] == "gdrive_read_confidential"
+
+
+@pytest.mark.asyncio
+async def test_no_write_down_roundtrip_read_then_blocked_write(
+    conversation_store: SqlAlchemyConversationStore,
+) -> None:
+    """End-to-end: read a confidential doc, then a later-turn outside write is denied.
+
+    Proves the confidential-read latch survives an engine rebuild via persisted
+    ``session_state`` — the demo's core narrative.
+    """
+    conv = conversation_store.create_conversation()
+    args = {"confidential_files": [_CONF_ID]}
+
+    engine1 = _engine(conversation_store, conv.id, {}, _DRIVE_HANDLER, args)
+    read_result = await engine1.evaluate(
+        EvaluationContext(
+            phase=Phase.TOOL_RESULT,
+            tool_name="mcp__google__docs_document_get",
+            content={"result": "{}"},
+            request_data={
+                "name": "mcp__google__docs_document_get",
+                "arguments": {"document_id": _CONF_ID},
+            },
+        )
+    )
+    assert read_result.action == PolicyAction.ALLOW
+    reloaded = conversation_store.get_conversation(conv.id)
+    assert reloaded is not None
+    assert reloaded.session_state.get("gdrive_read_confidential") is True
+
+    engine2 = _engine(
+        conversation_store, conv.id, dict(reloaded.session_state), _DRIVE_HANDLER, args
+    )
+    write = await engine2.evaluate(
+        EvaluationContext(
+            phase=Phase.TOOL_CALL,
+            tool_name="mcp__google__drive_file_update",
+            content={
+                "name": "mcp__google__drive_file_update",
+                "arguments": {"file_id": _OTHER_ID},
+            },
+        )
+    )
+    assert write.action == PolicyAction.DENY
+
+
 # ══════════════════════════════════════════════════════════════════════════════
 # gmail_policy
 # ══════════════════════════════════════════════════════════════════════════════
@@ -0,0 +1,70 @@
+# Demo for the built-in Google Drive *Bell-LaPadula "no write-down"* policy.
+#
+# Bell-LaPadula is the classic confidentiality model. Its "*-property" forbids a
+# **write down**: once a subject is cleared for confidential material, it must
+# not write that material into a less-protected place (a leak).
+#
+# This policy models a simple two-level compartment. You declare a set of
+# confidential documents up front (``confidential_files``). The rule is:
+#
+#   * The agent may freely read and write anything to start with.
+#   * The moment it *reads* one of the confidential documents, the session is
+#     "in the compartment": from then on, its writes are confined to the
+#     confidential set. Writing to — or creating — any other file is a
+#     write-down and is blocked.
+#
+# The confidential set is declared explicitly rather than inferred from a
+# per-document label, so the policy works on any Google Drive tenant — nothing
+# depends on Databricks-internal classification labels.
+#
+# The very same "create a new doc and write the summary there" action is allowed
+# before the agent has opened the strategy doc, and denied after — the action
+# didn't change, the session state did.
+#
+# Usage:
+#   export DATABRICKS_CONFIG_PROFILE=<your-profile>
+#   python -m omnigent tests/resources/examples/info_flow_agent.yaml \
+#       --prompt "Read the strategy doc, then create a public summary doc and paste the highlights."
+#   # Reading the confidential strategy doc puts the session in the compartment;
+#   # the follow-up docs_document_create is then blocked as a write-down.
+
+name: info_flow_agent
+prompt: |
+  You are an assistant with Google Workspace tools. You can read, edit, and
+  create Google Drive files and Docs. Your actions are governed by an
+  information-flow policy: once you have opened a confidential document, you may
+  only write back into confidential documents — writing that content into any
+  other file (or a new one) is blocked to prevent leaks. If an action is denied,
+  explain that it was gated by the no-write-down policy and suggest keeping the
+  work inside the confidential document.
+
+executor:
+  harness: openai-agents
+  model: databricks-gpt-5-4-mini
+  auth:
+    type: databricks
+    profile: oss
+
+tools:
+  # Point this at your Google MCP server. The policy matches tools by canonical
+  # name regardless of the server prefix (mcp__google__*, google__*, or bare),
+  # so no server-specific configuration is needed here.
+  google:
+    type: mcp
+    command: npx
+    args: ["-y", "@your-org/google-mcp-server"]
+
+policies:
+  confidential_containment:
+    type: function
+    handler: omnigent.policies.builtins.google.gdrive_policy
+    factory_params:
+      # The confidential compartment. Replace these with your own document IDs
+      # or URLs. Reading any of them confines the session's writes to this set.
+      confidential_files:
+        - "1ConfidentialStrategyDocDEMO0000000000000000"
+      # Allow creating new files so the demo can show a *create* being gated as a
+      # write-down (rather than blocked outright by the base access rules).
+      allow_create: true
+      # DENY (hard block) or ASK (require human approval) on a write-down.
+      write_down_action: DENY
@@ -2,21 +2,26 @@
 #
 # The agent accrues a risk score as it acts, and once the score crosses a
 # threshold, sensitive actions (sending mail, sharing files) flip from ALLOW to
-# ASK so a human has to approve them. Risk comes from two sources:
-#   * each web_search adds 10 points;
-#   * reading a doc whose result is tagged "Highly Confidential" / "RESTRICTED"
-#     adds 30 points (the label is read out of the tool result, so it works for
-#     any Google MCP server).
+# ASK so a human has to approve them.
+#
+# Risk comes from two independent sources — this demo drives the threshold using
+# the first one, which works with no special tool support:
+#   * each web_search adds 10 points (per-tool-call scoring via tool_points);
+#   * OPTIONALLY, reading a result annotated with a data-classification label
+#     adds points (sensitive_labels). This only fires if your MCP server tags
+#     its results with a classification field — many do not. It is shown here
+#     commented out; enable it and set the label values your server actually
+#     emits if you have one.
 #
 # The running score persists across turns via session_state, so a session that
-# has touched enough sensitive material stays gated.
+# has taken enough risky actions stays gated.
 #
 # Usage:
 #   export DATABRICKS_CONFIG_PROFILE=<your-profile>
 #   python -m omnigent tests/resources/examples/risk_score_agent.yaml \
 #       --prompt "Search the web for our competitors, then email a summary."
-#   # The web searches raise the score; the gmail_message_send is escalated to
-#   # ASK once the score reaches the threshold.
+#   # Five web searches raise the score to the threshold; the gmail_message_send
+#   # is then escalated to ASK.

 name: risk_score_agent
 prompt: |
@@ -49,13 +54,17 @@ policies:
    handler: omnigent.policies.builtins.risk_score.risk_score_policy
    factory_params:
      threshold: 50
-      # Each of these tool calls adds risk.
+      # Each of these tool calls adds risk. This is the portable path: it needs
+      # no special support from any tool, so the demo drives the threshold here
+      # (5 x 10 = 50).
      tool_points:
        web_search: 10
-      # Reading a result tagged with one of these classifications adds risk.
-      sensitive_labels:
-        "Highly Confidential": 30
-        RESTRICTED: 30
+      # OPTIONAL: score results that carry a data-classification label. Only
+      # uncomment this if your MCP server annotates results with a classification
+      # field, and set the label values it actually emits. (There is no portable,
+      # cross-server classification field, so this is left off by default.)
+      # sensitive_labels:
+      #   confidential: 30
      # Once the score reaches the threshold, these tools require approval.
      guarded_tools:
        - gmail_message_send
@@ -41,7 +41,6 @@ are covered by :mod:`tests.e2e.test_policies_e2e`

 from __future__ import annotations

-import json
 from pathlib import Path

 import pytest
@@ -282,22 +281,6 @@ async def test_secure_research_both_taints_deny_shell(
 # ─── risk_score_agent: built-in session-risk-score policy ───


-def _tool_result_ctx(name: str, result: str) -> EvaluationContext:
-    """
-    Build a TOOL_RESULT :class:`EvaluationContext` carrying a stringified result.
-
-    :param name: Tool that produced the result, e.g. ``"docs_document_get"``.
-    :param result: Stringified tool output under ``content.result``, e.g.
-        ``'{"label_classification": "Highly Confidential"}'``.
-    :returns: A ready-to-enforce TOOL_RESULT context.
-    """
-    return EvaluationContext(
-        phase=Phase.TOOL_RESULT,
-        content={"result": result},
-        tool_name=name,
-    )
-
-
@pytest.mark.asyncio
 async def test_risk_score_below_threshold_allows_guarded_tool(
    conversation_store: SqlAlchemyConversationStore,
@@ -340,28 +323,13 @@ async def test_risk_score_web_searches_accrue_and_gate_send(
    assert gated.deciding_policy == "session_risk"


-@pytest.mark.asyncio
-async def test_risk_score_confidential_reads_accrue_and_gate(
-    conversation_store: SqlAlchemyConversationStore,
-) -> None:
-    """
-    Loaded from YAML: two reads returning a "Highly Confidential" label
-    (2×30 = 60) cross the threshold, gating drive_permission_create.
-
-    Claim: the label-in-result scoring path works through the real load + engine
-    pipeline, not just the unit-level callable.
-    """
-    engine = _load_engine_from_yaml(_RISK_SCORE, conversation_store)
-    confidential = json.dumps({"label_classification": "Highly Confidential"})
-    for _ in range(2):
-        read = await _enforce_policy(engine, _tool_result_ctx("docs_document_get", confidential))
-        assert read.action == PolicyAction.ALLOW  # +30 each
-    gated = await _enforce_policy(
-        engine, _tool_ctx("drive_permission_create", {"file_id": "1AbC"})
-    )
-    # 60 >= 50 → sharing is gated after reading confidential material.
-    assert gated.action == PolicyAction.ASK
-    assert gated.deciding_policy == "session_risk"
+# NOTE: the label-in-result scoring path (``sensitive_labels``) is intentionally
+# NOT exercised from this example YAML. There is no portable, cross-MCP
+# classification field to depend on (the field the demo previously used,
+# ``label_classification``, is specific to the Databricks-internal Google MCP),
+# so the shipped example leaves ``sensitive_labels`` commented out and drives the
+# threshold via ``tool_points`` alone. The label-scoring mechanism itself is
+# fully covered at the unit level in tests/policies/builtins/test_risk_score.py.


 # Scenario 10 (secure_research_agent_os_env.yaml) is NOT tested
@@ -375,3 +343,108 @@ async def test_risk_score_confidential_reads_accrue_and_gate(
 # structurally identical to scenario 9 above, which IS covered,
 # so the policy-engine behavior is not uncovered — only the
 # direct-from-YAML load path is blocked.
+
+
+# ─── info_flow_agent: built-in gdrive Bell-LaPadula "no write-down" ───
+
+
+_INFO_FLOW = _EXAMPLES_DIR / "info_flow_agent.yaml"
+
+# Must match the confidential_files entry declared in info_flow_agent.yaml.
+_CONF_DOC_ID = "1ConfidentialStrategyDocDEMO0000000000000000"
+
+
+def _read_result_ctx(name: str, file_id: str) -> EvaluationContext:
+    """
+    Build a TOOL_RESULT context for a Drive *read*, carrying ``request_data``.
+
+    The policy correlates the read with the file it targeted (via
+    ``request_data``) to decide whether a confidential file was read, so a
+    scenario must supply the target file id.
+
+    :param name: Read tool name, e.g. ``"mcp__google__docs_document_get"``.
+    :param file_id: The file the read targeted, echoed under ``request_data``.
+    :returns: A ready-to-enforce TOOL_RESULT context.
+    """
+    return EvaluationContext(
+        phase=Phase.TOOL_RESULT,
+        content={"result": "{}"},
+        tool_name=name,
+        request_data={"name": name, "arguments": {"document_id": file_id}},
+    )
+
+
+@pytest.mark.asyncio
+async def test_info_flow_write_allowed_before_reading_confidential(
+    conversation_store: SqlAlchemyConversationStore,
+) -> None:
+    """
+    Loaded from YAML: before reading a confidential doc, writing elsewhere is fine.
+
+    Claim: the compartment rule imposes no constraint until the session has read
+    a confidential file.
+    """
+    engine = _load_engine_from_yaml(_INFO_FLOW, conversation_store)
+    # A create is allowed (allow_create: true) since no confidential read yet.
+    created = await _enforce_policy(
+        engine, _tool_ctx("mcp__google__docs_document_create", {"title": "notes"})
+    )
+    assert created.action == PolicyAction.ALLOW
+
+
+@pytest.mark.asyncio
+async def test_info_flow_blocks_write_out_after_reading_confidential(
+    conversation_store: SqlAlchemyConversationStore,
+) -> None:
+    """
+    Loaded from YAML: reading the confidential doc then creating an outside file denies.
+
+    Reading the confidential doc latches the session; the follow-up
+    ``docs_document_create`` targets a brand-new (outside-compartment) file, a
+    write-down, and DENYs — the demo's headline "same action, different outcome,
+    because the state changed".
+
+    Claim: the confidential-read latch persists in session_state across
+    enforcement calls and drives the write-down gate through the real load +
+    engine pipeline.
+    """
+    engine = _load_engine_from_yaml(_INFO_FLOW, conversation_store)
+    read = await _enforce_policy(
+        engine,
+        _read_result_ctx("mcp__google__docs_document_get", _CONF_DOC_ID),
+    )
+    assert read.action == PolicyAction.ALLOW
+    create = await _enforce_policy(
+        engine,
+        _tool_ctx("mcp__google__docs_document_create", {"title": "leak"}),
+    )
+    assert create.action == PolicyAction.DENY
+    assert create.deciding_policy == "confidential_containment"
+
+
+@pytest.mark.asyncio
+async def test_info_flow_confidential_files_does_not_grant_write(
+    conversation_store: SqlAlchemyConversationStore,
+) -> None:
+    """
+    Loaded from YAML: declaring a file confidential does not make it writable.
+
+    The example lists the doc in ``confidential_files`` but not ``write_files``,
+    and the agent never created it, so a write to it is denied by the base write
+    rule — ``confidential_files`` is a containment declaration, not a write
+    grant. (The no-write-down check itself abstains here, since the target is in
+    the confidential set; the denial comes from the base scope rule.)
+
+    Claim: through the real load + engine pipeline, ``confidential_files`` does
+    not widen the write boundary.
+    """
+    engine = _load_engine_from_yaml(_INFO_FLOW, conversation_store)
+    await _enforce_policy(
+        engine,
+        _read_result_ctx("mcp__google__docs_document_get", _CONF_DOC_ID),
+    )
+    write = await _enforce_policy(
+        engine,
+        _tool_ctx("mcp__google__docs_document_batch_update", {"document_id": _CONF_DOC_ID}),
+    )
+    assert write.action == PolicyAction.DENY