fix(mcp): scope late OAuth attempts by profile
_LATE_ATTEMPTS was keyed by session key alone while live.py keys the open-operation table by (profile_key, session_key) for exactly this collision: two multiplexed profiles can carry the same session key (the api_server binds X-Hermes-Session-Key verbatim, and it is client-chosen). A parked OAuth attempt could therefore be adopted by a different profile's next turn: adopt_late_connections would poll it, register the same-named server from that profile's config, and enable it - a grant authorized under profile A materializing under profile B, or being discarded so the owning profile never adopts it. Key the table by (profile_key, session_key), matching live.py's documented identity pairing. The detached no-card path never opens an operation, so its profile_key is empty; fall back to the calling thread's home, which close() runs under the turn's profile scope.
This commit is contained in:
beardthelion
Teknium
@@ -10,6 +10,7 @@ Contracts:
|
||||
- deadline ownership: fixed operation deadline + sequential-deadline exemption
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
@@ -355,3 +356,148 @@ def test_a_desktop_session_with_no_callback_gets_the_link_at_once_and_opens_no_o
|
||||
assert out["status"] == "initiated"
|
||||
assert out["targets"][0]["connect_url"] == "https://auth.example/paper/1"
|
||||
assert live.current("s1") is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# late-attempt parking is scoped by (profile, session), never by session alone
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _as_home(home):
|
||||
"""Bind a profile home the way the multiplex gateway binds one per activity."""
|
||||
from hermes_constants import reset_hermes_home_override, set_hermes_home_override
|
||||
|
||||
token = set_hermes_home_override(str(home))
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
reset_hermes_home_override(token)
|
||||
|
||||
|
||||
def _profile_home(tmp_path, name):
|
||||
"""A real profile dir with the same-named MCP server configured, as both multiplexed
|
||||
homes carry in production."""
|
||||
home = tmp_path / name
|
||||
home.mkdir(parents=True)
|
||||
(home / "config.yaml").write_text(
|
||||
"mcp_servers:\n linear:\n url: https://mcp.linear.app/mcp\n")
|
||||
return home
|
||||
|
||||
|
||||
def _park_attempt(home, name="linear", session_key="s1", profile_key="stamped"):
|
||||
"""Close a runner whose operation parked one approved OAuth attempt, the way
|
||||
``_Runner.close`` does at the end of a tool call bound to ``home``."""
|
||||
import tools.connectors.mcp as mcp
|
||||
from hermes_constants import hermes_home_key
|
||||
from tools.connectors.operation import ConnectionOperation
|
||||
|
||||
attempt = FakeAttempt("https://auth.example/linear")
|
||||
attempt.approve(["read"])
|
||||
operation = ConnectionOperation(targets=[], session_key=session_key)
|
||||
runner = mcp._Runner("authorize", backend=None)
|
||||
runner.operation = operation
|
||||
runner.op_id = None
|
||||
runner.work = {name: mcp._Work(attempt=attempt)}
|
||||
with _as_home(home):
|
||||
if profile_key == "stamped":
|
||||
operation.profile_key = hermes_home_key() # what live.open stamps
|
||||
else:
|
||||
operation.profile_key = profile_key # "" on the detached no-card path
|
||||
runner.close()
|
||||
return attempt
|
||||
|
||||
|
||||
def _adopt_as(home, agent=None, session_id="s1"):
|
||||
"""adopt_late_connections as it runs inside a turn bound to ``home``."""
|
||||
import tools.connectors.mcp as mcp
|
||||
|
||||
registered = []
|
||||
agent = agent or SimpleNamespace(session_id=session_id, enabled_toolsets=[])
|
||||
with _as_home(home), \
|
||||
patch("tools.mcp_tool_discovery.register_mcp_servers",
|
||||
side_effect=lambda servers: registered.extend(servers) or list(servers)):
|
||||
adopted = mcp.adopt_late_connections(agent)
|
||||
return adopted, registered, agent
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_late_attempts():
|
||||
import tools.connectors.mcp as mcp
|
||||
|
||||
mcp._LATE_ATTEMPTS.clear()
|
||||
yield
|
||||
mcp._LATE_ATTEMPTS.clear()
|
||||
|
||||
|
||||
def test_late_attempt_is_never_adopted_by_another_profile(tmp_path):
|
||||
"""Two multiplexed profiles can carry the same session key (the api_server's
|
||||
X-Hermes-Session-Key header is client-chosen, and live.py keys _open by
|
||||
(profile, session) for exactly this reason). A parked grant must not leak."""
|
||||
import tools.connectors.mcp as mcp
|
||||
|
||||
home_a = _profile_home(tmp_path, "home-a")
|
||||
home_b = _profile_home(tmp_path, "home-b")
|
||||
_park_attempt(home_a)
|
||||
|
||||
adopted, registered, agent = _adopt_as(home_b)
|
||||
assert adopted == [] and registered == []
|
||||
assert mcp._LATE_ATTEMPTS # still parked for its owner
|
||||
|
||||
adopted, registered, agent = _adopt_as(home_a)
|
||||
assert adopted == ["linear"] and registered == ["linear"]
|
||||
assert agent.enabled_toolsets == ["linear"]
|
||||
assert mcp._LATE_ATTEMPTS == {}
|
||||
|
||||
|
||||
def test_late_attempt_keyed_by_detached_path_uses_calling_profile(tmp_path):
|
||||
"""The no-card DetachedOperation never passes through live.open, so profile_key is empty;
|
||||
the park must still record the home the tool thread was scoped to."""
|
||||
import tools.connectors.mcp as mcp
|
||||
|
||||
home_a = _profile_home(tmp_path, "home-a")
|
||||
home_b = _profile_home(tmp_path, "home-b")
|
||||
_park_attempt(home_a, profile_key="")
|
||||
|
||||
from hermes_constants import hermes_home_key
|
||||
with _as_home(home_a):
|
||||
assert list(mcp._LATE_ATTEMPTS) == [(hermes_home_key(), "s1")]
|
||||
adopted, registered, _ = _adopt_as(home_b)
|
||||
assert adopted == []
|
||||
adopted, registered, _ = _adopt_as(home_a)
|
||||
assert adopted == ["linear"]
|
||||
|
||||
|
||||
def test_e2e_carded_oauth_park_and_adopt_stay_inside_their_profile(tmp_path):
|
||||
"""The production path end to end: an authorize card bound to home A deadline-settles
|
||||
while its OAuth attempt is still pending, so _Runner.close parks the attempt under the
|
||||
profile key live.open stamped. The browser-side approval that lands afterwards may only
|
||||
be adopted by a turn bound to the same home — never by the other multiplexed profile,
|
||||
even one that configures the same-named server."""
|
||||
import tools.connectors.mcp as mcp
|
||||
|
||||
home_a = _profile_home(tmp_path, "home-a")
|
||||
home_b = _profile_home(tmp_path, "home-b")
|
||||
|
||||
backend = FakeBackend()
|
||||
with _as_home(home_a):
|
||||
with patch("tools.connectors.operation.OPERATION_DEADLINE_SECONDS", 0.05):
|
||||
out = _mcp({"action": "authorize", "connectors": [_linear()]},
|
||||
_answering(None), mcp_backend=backend)
|
||||
assert out["settled_by"] == SettleReason.deadline.value
|
||||
backend.attempts["linear"].approve(["read"]) # the browser flow lands after the card closed
|
||||
|
||||
registered = []
|
||||
with patch("tools.mcp_tool_discovery.register_mcp_servers",
|
||||
side_effect=lambda servers: registered.extend(servers) or list(servers)):
|
||||
agent_b = SimpleNamespace(session_id="s1", enabled_toolsets=[])
|
||||
with _as_home(home_b):
|
||||
assert mcp.adopt_late_connections(agent_b) == []
|
||||
assert registered == [] and agent_b.enabled_toolsets == []
|
||||
assert mcp._LATE_ATTEMPTS # still parked for its owner
|
||||
|
||||
agent_a = SimpleNamespace(session_id="s1", enabled_toolsets=[])
|
||||
with _as_home(home_a):
|
||||
assert mcp.adopt_late_connections(agent_a) == ["linear"]
|
||||
assert registered == ["linear"] and agent_a.enabled_toolsets == ["linear"]
|
||||
assert mcp._LATE_ATTEMPTS == {}
|
||||
|
||||
@@ -7,8 +7,9 @@ import logging
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
|
||||
from hermes_constants import hermes_home_key
|
||||
from tools.connectors.contract import Actor, SettleReason, TargetState
|
||||
from tools.connectors.gateway.config import operation_session_key
|
||||
from tools.connectors.operation import ConnectionOperation, DetachedOperation, IllegalTransition, Target
|
||||
@@ -319,12 +320,22 @@ class _Runner:
|
||||
|
||||
cancel_attempt(work.attempt.flow)
|
||||
else:
|
||||
_LATE_ATTEMPTS.setdefault(operation.session_key, {})[name] = work.attempt
|
||||
_LATE_ATTEMPTS.setdefault(_late_key(operation), {})[name] = work.attempt
|
||||
self.work.clear()
|
||||
|
||||
|
||||
# session key -> {server: attempt} for OAuth attempts that outlived their card.
|
||||
_LATE_ATTEMPTS: Dict[str, Dict[str, Any]] = {}
|
||||
def _late_key(operation: ConnectionOperation) -> Tuple[str, str]:
|
||||
"""The ``(profile, session)`` pairing ``live.open`` keys an operation by. ``profile_key``
|
||||
is stamped there; the detached no-card path never opens, so fall back to the calling
|
||||
thread's home — ``close`` runs on the tool thread under the turn's profile scope."""
|
||||
return (operation.profile_key or hermes_home_key(), operation.session_key)
|
||||
|
||||
|
||||
# (profile key, session key) -> {server: attempt} for OAuth attempts that outlived their card.
|
||||
# The profile is part of the key for the same reason live.py keys _open by it: two multiplexed
|
||||
# profiles can carry the same session key, and an attempt must only ever be adopted by the
|
||||
# profile whose card authorized it.
|
||||
_LATE_ATTEMPTS: Dict[Tuple[str, str], Dict[str, Any]] = {}
|
||||
|
||||
|
||||
def adopt_late_connections(agent: Any) -> List[str]:
|
||||
@@ -332,7 +343,8 @@ def adopt_late_connections(agent: Any) -> List[str]:
|
||||
them to the agent's toolset selection. Runs between turns, so the result that said "not
|
||||
connected" is followed by a turn in which the tools are there."""
|
||||
session_key = operation_session_key(getattr(agent, "session_id", None))
|
||||
attempts = _LATE_ATTEMPTS.get(session_key)
|
||||
key = (hermes_home_key(), session_key)
|
||||
attempts = _LATE_ATTEMPTS.get(key)
|
||||
if not attempts:
|
||||
return []
|
||||
adopted: List[str] = []
|
||||
@@ -354,7 +366,7 @@ def adopt_late_connections(agent: Any) -> List[str]:
|
||||
except Exception:
|
||||
logger.debug("late MCP connection %s was not adopted", name, exc_info=True)
|
||||
if not attempts:
|
||||
_LATE_ATTEMPTS.pop(session_key, None)
|
||||
_LATE_ATTEMPTS.pop(key, None)
|
||||
enabled = getattr(agent, "enabled_toolsets", None)
|
||||
if adopted and enabled is not None and "no_mcp" not in enabled:
|
||||
agent.enabled_toolsets = [*enabled, *(n for n in adopted if n not in enabled)]
|
||||
|
||||
Reference in New Issue
Block a user