From 89d236a774df6bd6649baf0ec8903a45c58b17fe Mon Sep 17 00:00:00 2001 From: Kaiki Date: Sun, 12 Jul 2026 20:32:16 -0400 Subject: [PATCH 01/12] fix(skills): validate serialized list fields --- routes/skills_routes.py | 2 +- services/memory/skill_format.py | 31 ++- services/memory/skills.py | 110 ++++++++-- src/tools/system.py | 72 ++++--- tests/test_manage_skills_list_fields.py | 267 ++++++++++++++++++++++++ 5 files changed, 429 insertions(+), 53 deletions(-) create mode 100644 tests/test_manage_skills_list_fields.py diff --git a/routes/skills_routes.py b/routes/skills_routes.py index 711baa2e57..4cce3659d9 100644 --- a/routes/skills_routes.py +++ b/routes/skills_routes.py @@ -1625,7 +1625,7 @@ async def update_skill(request: Request, skill_id: str, body: SkillUpdateRequest raise HTTPException(404, "Skill not found") _verify_owner(match, user) - updates = body.dict(exclude_none=True) + updates = body.model_dump(exclude_none=True) if not updates: return {"ok": True} ok = skills_manager.update_skill(match.get("name"), updates, owner=user) diff --git a/services/memory/skill_format.py b/services/memory/skill_format.py index 2b2dfb1b30..c6051525d9 100644 --- a/services/memory/skill_format.py +++ b/services/memory/skill_format.py @@ -186,12 +186,14 @@ def _emit_scalar(v: Any) -> str: return s -def _as_list(v: Any) -> List[str]: +def _as_string_list(v: Any, field_name: str) -> List[str]: if v is None: return [] if isinstance(v, list): - return [str(x) for x in v if x not in (None, "")] - return [str(v)] + if not all(isinstance(x, str) for x in v): + raise ValueError(f"{field_name} must be a list of strings") + return [x for x in v if x != ""] + raise ValueError(f"{field_name} must be a list of strings") def _as_float(v: Any, default: float = 0.8) -> float: @@ -284,12 +286,21 @@ def _parse_list_lines(text: str) -> List[str]: items.append(m.group(1).strip()) elif items: # continuation of previous bullet - items[-1] = items[-1] + " " + s + items[-1] = items[-1] + "\n" + s else: items.append(s) return items +def _emit_list_item(prefix: str, item: Any) -> str: + lines = str(item).splitlines() or [""] + first, rest = lines[0], lines[1:] + if not rest: + return f"{prefix} {first}" + continuation = "\n".join(f" {line}" if line else "" for line in rest) + return f"{prefix} {first}\n{continuation}" + + def emit_body(sections: Dict[str, Any]) -> str: parts: List[str] = [] when = (sections.get("when_to_use") or "").strip() @@ -301,9 +312,9 @@ def emit_body(sections: Dict[str, Any]) -> str: continue heading = _KEY_TO_HEADING[key] if key == "procedure": - body = "\n".join(f"{i + 1}. {x}" for i, x in enumerate(items)) + body = "\n".join(_emit_list_item(f"{i + 1}.", x) for i, x in enumerate(items)) else: - body = "\n".join(f"- {x}" for x in items) + body = "\n".join(_emit_list_item("-", x) for x in items) parts.append(f"## {heading}\n\n{body}") extra = (sections.get("body_extra") or "").strip() if extra: @@ -410,10 +421,10 @@ def from_markdown(cls, text: str, *, path: Optional[str] = None) -> "Skill": description=str(fm.get("description", "") or ""), version=str(fm.get("version", "1.0.0") or "1.0.0"), category=str(fm.get("category", "general") or "general"), - tags=_as_list(fm.get("tags")), - platforms=_as_list(fm.get("platforms")), - requires_toolsets=_as_list(fm.get("requires_toolsets")), - fallback_for_toolsets=_as_list(fm.get("fallback_for_toolsets")), + tags=_as_string_list(fm.get("tags"), "tags"), + platforms=_as_string_list(fm.get("platforms"), "platforms"), + requires_toolsets=_as_string_list(fm.get("requires_toolsets"), "requires_toolsets"), + fallback_for_toolsets=_as_string_list(fm.get("fallback_for_toolsets"), "fallback_for_toolsets"), status=str(fm.get("status", "draft") or "draft"), confidence=_as_float(fm.get("confidence", 0.8), 0.8), source=str(fm.get("source", "learned") or "learned"), diff --git a/services/memory/skills.py b/services/memory/skills.py index 5baaa88c5b..07f51af8d3 100644 --- a/services/memory/skills.py +++ b/services/memory/skills.py @@ -22,8 +22,9 @@ import json import logging import os +import re import time -from typing import Dict, Iterable, List, Optional +from typing import Any, Dict, Iterable, List, Optional from .skill_format import Skill, slugify @@ -54,6 +55,67 @@ def _to_float(x, default: float = 0.0) -> float: return default +def _normalize_procedure_text(value: str) -> List[str]: + """Normalize a plain-text procedure into list entries. + + A top-level Markdown bullet or numbered item starts a new entry; indented + continuation lines stay with that entry. If the text is not a Markdown list, + blank-line-separated blocks become entries. This keeps multiline Markdown + ordered and stable without ever treating the string as an iterable of chars. + """ + text = value.strip("\n") + if not text.strip(): + return [] + + marker_re = re.compile(r"^(?:[-*]|\d+[.)])\s+(.*)$") + has_top_level_markers = any( + line == line.lstrip() and marker_re.match(line.strip()) + for line in text.splitlines() + ) + + if not has_top_level_markers: + blocks = re.split(r"\n\s*\n", text) + return [block.strip() for block in blocks if block.strip()] + + items: List[List[str]] = [] + current: List[str] = [] + for line in text.splitlines(): + stripped = line.strip() + if not stripped: + continue + marker = marker_re.match(stripped) if line == line.lstrip() else None + if marker: + if current: + items.append(current) + current = [marker.group(1).rstrip()] + elif current: + current.append(stripped) + else: + current = [stripped] + if current: + items.append(current) + return ["\n".join(lines).strip() for lines in items if "\n".join(lines).strip()] + + +def _normalize_string_list_field( + field: str, + value: Any, + *, + allow_procedure_string: bool = False, +) -> List[str]: + if value is None: + return [] + if allow_procedure_string and isinstance(value, str): + return _normalize_procedure_text(value) + if not isinstance(value, list): + if allow_procedure_string: + raise ValueError(f"{field} must be a list of strings or a string") + raise ValueError(f"{field} must be a list of strings") + if not all(isinstance(item, str) for item in value): + raise ValueError(f"{field} must be a list of strings") + return list(value) + + # --------------------------------------------------------------------------- # SkillsManager # --------------------------------------------------------------------------- @@ -318,6 +380,26 @@ def add_skill( ) -> Dict: # Normalize name nm = slugify(name or title or description or "skill") + steps_norm = _normalize_string_list_field("steps", steps) if steps is not None else [] + if procedure is not None: + proc = _normalize_string_list_field("procedure", procedure, allow_procedure_string=True) + else: + proc = steps_norm + tag_list = _normalize_string_list_field("tags", tags if tags is not None else []) + platform_list = _normalize_string_list_field("platforms", platforms if platforms is not None else []) + requires_toolset_list = _normalize_string_list_field( + "requires_toolsets", + requires_toolsets if requires_toolsets is not None else [], + ) + fallback_toolset_list = _normalize_string_list_field( + "fallback_for_toolsets", + fallback_for_toolsets if fallback_for_toolsets is not None else [], + ) + pitfall_list = _normalize_string_list_field("pitfalls", pitfalls if pitfalls is not None else []) + verification_list = _normalize_string_list_field( + "verification", + verification if verification is not None else [], + ) # Free dedup-at-creation (always, no API): for LLM-authored skills, # skip if a near-identical skill already exists (Jaccard over @@ -330,7 +412,7 @@ def add_skill( cand = _tokenize(" ".join([ nm, (description or title or ""), (when_to_use if when_to_use is not None else (problem or "")), - " ".join(procedure if procedure is not None else (steps or [])), + " ".join(proc), ])) if cand: for s in _dedup_pool: @@ -362,20 +444,20 @@ def add_skill( description=(description or title or "").strip(), version=version, category=category or "general", - tags=list(tags or []), - platforms=list(platforms or []), - requires_toolsets=list(requires_toolsets or []), - fallback_for_toolsets=list(fallback_for_toolsets or []), + tags=tag_list, + platforms=platform_list, + requires_toolsets=requires_toolset_list, + fallback_for_toolsets=fallback_toolset_list, status=status or "draft", confidence=float(confidence), source=source, teacher_model=teacher_model, owner=owner, when_to_use=(when_to_use if when_to_use is not None else (problem or "")), - procedure=list(procedure if procedure is not None else (steps or [])), - pitfalls=list(pitfalls or []), - verification=list(verification or []), - body_extra=(solution if solution and not procedure else ""), + procedure=proc, + pitfalls=pitfall_list, + verification=verification_list, + body_extra=(solution if solution and not proc else ""), ) self._write_skill(sk) @@ -465,7 +547,11 @@ def update_skill(self, skill_id: str, updates: Dict, owner: Optional[str] = None "platforms", "requires_toolsets", "fallback_for_toolsets") for k in list_keys: if k in updates: - setattr(sk, k, list(updates[k] or [])) + setattr(sk, k, _normalize_string_list_field( + k, + updates[k], + allow_procedure_string=(k == "procedure"), + )) # Old-schema field aliases if "title" in updates and "description" not in updates: @@ -475,7 +561,7 @@ def update_skill(self, skill_id: str, updates: Dict, owner: Optional[str] = None if "solution" in updates and "body_extra" not in updates and not sk.procedure: sk.body_extra = updates["solution"] if "steps" in updates and "procedure" not in updates: - sk.procedure = list(updates["steps"] or []) + sk.procedure = _normalize_string_list_field("steps", updates["steps"]) # Rename new_name = slugify(updates.get("name") or sk.name) diff --git a/src/tools/system.py b/src/tools/system.py index a901992c26..78accbd151 100644 --- a/src/tools/system.py +++ b/src/tools/system.py @@ -97,10 +97,13 @@ async def do_manage_skills(content: str, owner: Optional[str] = None) -> Dict: "error": "name is required for add. Provide the exact slug the user should see, then report the returned name.", "exit_code": 1, } - proc = args.get("procedure") - if proc is None: - proc = args.get("steps") or [] - if not proc and not args.get("body_extra") and not args.get("solution"): + if "procedure" in args: + proc = args.get("procedure") + elif "steps" in args: + proc = args.get("steps") + else: + proc = [] + if proc in (None, [], "") and not args.get("body_extra") and not args.get("solution"): return {"error": "procedure (or solution body) is required", "exit_code": 1} # Same auto-publish gate as the extractor path — when the user # has auto_approve_skills on and the caller didn't pin an explicit @@ -113,30 +116,33 @@ async def do_manage_skills(content: str, owner: Optional[str] = None) -> Dict: _status_arg = "published" if _prefs.get("auto_approve_skills", True) else "draft" except Exception: _status_arg = "draft" - entry = sm.add_skill( - name=args.get("name"), - description=(args.get("description") or args.get("title") or "").strip(), - category=args.get("category") or "general", - tags=args.get("tags") or [], - platforms=args.get("platforms") or [], - requires_toolsets=args.get("requires_toolsets") or [], - fallback_for_toolsets=args.get("fallback_for_toolsets") or [], - when_to_use=(args.get("when_to_use") if args.get("when_to_use") is not None - else args.get("problem", "")), - procedure=proc, - pitfalls=args.get("pitfalls") or [], - verification=args.get("verification") or [], - status=_status_arg, - version=args.get("version") or "1.0.0", - confidence=args.get("confidence", 0.8), - source=args.get("source", "learned"), - teacher_model=args.get("teacher_model"), - owner=owner, - title=args.get("title", ""), - problem=args.get("problem", ""), - solution=args.get("solution", ""), - steps=args.get("steps") or [], - ) + try: + entry = sm.add_skill( + name=args.get("name"), + description=(args.get("description") or args.get("title") or "").strip(), + category=args.get("category") or "general", + tags=args["tags"] if "tags" in args else [], + platforms=args["platforms"] if "platforms" in args else [], + requires_toolsets=args["requires_toolsets"] if "requires_toolsets" in args else [], + fallback_for_toolsets=args["fallback_for_toolsets"] if "fallback_for_toolsets" in args else [], + when_to_use=(args.get("when_to_use") if args.get("when_to_use") is not None + else args.get("problem", "")), + procedure=proc, + pitfalls=args["pitfalls"] if "pitfalls" in args else [], + verification=args["verification"] if "verification" in args else [], + status=_status_arg, + version=args.get("version") or "1.0.0", + confidence=args.get("confidence", 0.8), + source=args.get("source", "learned"), + teacher_model=args.get("teacher_model"), + owner=owner, + title=args.get("title", ""), + problem=args.get("problem", ""), + solution=args.get("solution", ""), + steps=args["steps"] if "steps" in args else None, + ) + except ValueError as e: + return {"error": str(e), "exit_code": 1} if entry.get("_deduped"): return {"results": ( f"A near-identical skill already exists: `{entry['name']}` — not creating " @@ -172,7 +178,10 @@ async def do_manage_skills(content: str, owner: Optional[str] = None) -> Dict: return {"error": f"Skill {name!r} not found", "exit_code": 1} if not sk_new.owner: sk_new.owner = match.get("owner") or owner - ok = sm.update_skill(name, _skill_dump(sk_new), owner=owner) + try: + ok = sm.update_skill(name, _skill_dump(sk_new), owner=owner) + except ValueError as e: + return {"error": str(e), "exit_code": 1} return {"results": f"Edited skill `{sk_new.name}`."} if ok else {"error": "Update failed", "exit_code": 1} if action == "patch": @@ -196,7 +205,10 @@ async def do_manage_skills(content: str, owner: Optional[str] = None) -> Dict: except Exception as e: return {"error": f"Patched content is not valid SKILL.md: {e}", "exit_code": 1} sk_new.name = slugify(sk_new.name or name) - ok = sm.update_skill(name, _skill_dump(sk_new), owner=owner) + try: + ok = sm.update_skill(name, _skill_dump(sk_new), owner=owner) + except ValueError as e: + return {"error": str(e), "exit_code": 1} return {"results": f"Patched skill `{sk_new.name}`."} if ok else {"error": "Patch update failed", "exit_code": 1} if action == "publish": diff --git a/tests/test_manage_skills_list_fields.py b/tests/test_manage_skills_list_fields.py new file mode 100644 index 0000000000..7b9d8f3aa5 --- /dev/null +++ b/tests/test_manage_skills_list_fields.py @@ -0,0 +1,267 @@ +import json + +import pytest + +from services.memory.skill_format import Skill +from services.memory.skills import SkillsManager +from src.tools.system import do_manage_skills + + +MULTILINE_PROCEDURE = """1. Open the existing file. + Keep related context visible. + +2. Apply the minimal patch. + ```python + print("still one step") + ``` + +3. Run the focused test.""" + + +def test_add_skill_normalizes_plain_string_procedure_without_character_entries(tmp_path): + sm = SkillsManager(str(tmp_path)) + + entry = sm.add_skill( + name="plain-string-procedure", + description="plain string procedure", + procedure=MULTILINE_PROCEDURE, + tags=["dev"], + owner="alice", + source="user", + ) + + assert entry["procedure"] == [ + 'Open the existing file.\nKeep related context visible.', + 'Apply the minimal patch.\n```python\nprint("still one step")\n```', + "Run the focused test.", + ] + assert entry["procedure"] != list(MULTILINE_PROCEDURE) + + markdown = sm.read_skill_md("plain-string-procedure", owner="alice") + assert markdown is not None + assert "\n1. O\n2. p\n3. e\n" not in markdown + assert 'print("still one step")' in markdown + + reparsed = Skill.from_markdown(markdown) + assert reparsed.procedure == entry["procedure"] + + +def test_update_skill_normalizes_plain_string_procedure_without_character_entries(tmp_path): + sm = SkillsManager(str(tmp_path)) + sm.add_skill( + name="update-string-procedure", + description="before", + procedure=["before"], + owner="alice", + source="user", + ) + + assert sm.update_skill( + "update-string-procedure", + {"procedure": MULTILINE_PROCEDURE}, + owner="alice", + ) + + updated = sm.load(owner="alice")[0] + assert updated["procedure"][0] == "Open the existing file.\nKeep related context visible." + assert updated["procedure"] != list(MULTILINE_PROCEDURE) + markdown = sm.read_skill_md("update-string-procedure", owner="alice") + assert markdown is not None + assert "\n1. O\n2. p\n3. e\n" not in markdown + + +def test_add_and_update_accept_canonical_list_fields(tmp_path): + sm = SkillsManager(str(tmp_path)) + entry = sm.add_skill( + name="canonical-lists", + description="canonical", + procedure=["one", "two"], + pitfalls=["avoid this"], + verification=["passes"], + tags=["dev", "skill"], + platforms=["linux"], + requires_toolsets=["shell"], + fallback_for_toolsets=["web"], + owner="alice", + source="user", + ) + + assert entry["procedure"] == ["one", "two"] + assert entry["pitfalls"] == ["avoid this"] + assert entry["verification"] == ["passes"] + assert entry["tags"] == ["dev", "skill"] + assert entry["platforms"] == ["linux"] + assert entry["requires_toolsets"] == ["shell"] + assert entry["fallback_for_toolsets"] == ["web"] + + assert sm.update_skill( + "canonical-lists", + { + "steps": ["updated one", "updated two"], + "tags": ["updated"], + "platforms": [], + "requires_toolsets": ["read_file"], + "fallback_for_toolsets": [], + }, + owner="alice", + ) + updated = sm.load(owner="alice")[0] + assert updated["procedure"] == ["updated one", "updated two"] + assert updated["tags"] == ["updated"] + assert updated["platforms"] == [] + assert updated["requires_toolsets"] == ["read_file"] + assert updated["fallback_for_toolsets"] == [] + + +@pytest.mark.parametrize( + "field,value", + [ + ("procedure", {"bad": "dict"}), + ("procedure", 123), + ("procedure", [["nested"]]), + ("procedure", ["ok", 3]), + ("steps", "plain steps string"), + ("pitfalls", {"bad": "dict"}), + ("pitfalls", 123), + ("pitfalls", [["nested"]]), + ("pitfalls", ["ok", 3]), + ("verification", {"bad": "dict"}), + ("tags", {"bad": "dict"}), + ("tags", 123), + ("tags", [["nested"]]), + ("tags", ["ok", 3]), + ("platforms", ["linux", 3]), + ("requires_toolsets", ["shell", ["nested"]]), + ("fallback_for_toolsets", {"bad": "dict"}), + ], +) +def test_add_skill_rejects_invalid_list_field_values(tmp_path, field, value): + sm = SkillsManager(str(tmp_path)) + kwargs = { + "name": f"invalid-{field.replace('_', '-')}", + "description": "invalid", + "procedure": ["valid procedure"], + "owner": "alice", + "source": "user", + } + kwargs[field] = value + + with pytest.raises(ValueError, match=field): + sm.add_skill(**kwargs) + + +@pytest.mark.parametrize( + "field,value", + [ + ("procedure", {"bad": "dict"}), + ("procedure", 123), + ("procedure", [["nested"]]), + ("procedure", ["ok", 3]), + ("steps", "plain steps string"), + ("pitfalls", ["ok", 3]), + ("verification", [["nested"]]), + ("tags", {"bad": "dict"}), + ("platforms", 123), + ("requires_toolsets", ["shell", ["nested"]]), + ("fallback_for_toolsets", ["web", 3]), + ], +) +def test_update_skill_rejects_invalid_list_field_values(tmp_path, field, value): + sm = SkillsManager(str(tmp_path)) + sm.add_skill( + name="invalid-update", + description="invalid update", + procedure=["valid procedure"], + owner="alice", + source="user", + ) + + with pytest.raises(ValueError, match=field): + sm.update_skill("invalid-update", {field: value}, owner="alice") + + +@pytest.mark.asyncio +async def test_manage_skills_add_returns_sanitized_validation_error(monkeypatch, tmp_path): + import src.constants as constants + + monkeypatch.setattr(constants, "DATA_DIR", str(tmp_path)) + result = await do_manage_skills( + json.dumps({ + "action": "add", + "name": "invalid-tool-add", + "description": "invalid tool add", + "procedure": ["valid"], + "tags": ["ok", 3], + "status": "draft", + }), + owner="alice", + ) + + assert result["exit_code"] == 1 + assert result["error"] == "tags must be a list of strings" + + +@pytest.mark.asyncio +async def test_manage_skills_edit_rejects_invalid_frontmatter_list(monkeypatch, tmp_path): + import src.constants as constants + + monkeypatch.setattr(constants, "DATA_DIR", str(tmp_path)) + sm = SkillsManager(str(tmp_path)) + sm.add_skill( + name="invalid-tool-edit", + description="invalid tool edit", + procedure=["valid"], + owner="alice", + source="user", + ) + + result = await do_manage_skills( + json.dumps({ + "action": "edit", + "name": "invalid-tool-edit", + "content": """--- +name: invalid-tool-edit +description: invalid tool edit +tags: [ok, 3] +status: draft +confidence: 0.8 +source: user +owner: alice +created: 2026-07-12T00:00:00Z +--- + +## Procedure + +1. valid +""", + }), + owner="alice", + ) + + assert result["exit_code"] == 1 + assert result["error"] == "Could not parse content as SKILL.md: tags must be a list of strings" + + +def test_skill_from_markdown_preserves_valid_list_round_trip(): + original = Skill( + name="valid-round-trip", + description="valid round trip", + tags=["dev", "skill"], + platforms=["linux"], + requires_toolsets=["shell"], + fallback_for_toolsets=["web"], + procedure=["first", "second"], + pitfalls=["avoid this"], + verification=["passes"], + owner="alice", + ) + + reparsed = Skill.from_markdown(original.to_markdown()) + + assert reparsed.tags == ["dev", "skill"] + assert reparsed.platforms == ["linux"] + assert reparsed.requires_toolsets == ["shell"] + assert reparsed.fallback_for_toolsets == ["web"] + assert reparsed.procedure == ["first", "second"] + assert reparsed.pitfalls == ["avoid this"] + assert reparsed.verification == ["passes"] From f751409c395c7520e68e0170ddea11404e294b4a Mon Sep 17 00:00:00 2001 From: Kaiki Date: Sun, 12 Jul 2026 20:32:16 -0400 Subject: [PATCH 02/12] fix(mcp): bound connections and redact errors --- src/mcp_manager.py | 149 +++++++-- .../test_mcp_connection_timeout_redaction.py | 304 ++++++++++++++++++ 2 files changed, 434 insertions(+), 19 deletions(-) create mode 100644 tests/test_mcp_connection_timeout_redaction.py diff --git a/src/mcp_manager.py b/src/mcp_manager.py index 90430c78fa..b1d8f8b825 100644 --- a/src/mcp_manager.py +++ b/src/mcp_manager.py @@ -9,16 +9,89 @@ import logging import os import re +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit from typing import Any, Dict, List, Optional, Set, Tuple from src.runtime_paths import get_app_root logger = logging.getLogger(__name__) + +def _mcp_connect_timeout_seconds() -> float: + try: + return max(0.1, float(os.getenv("ODYSSEUS_MCP_CONNECT_TIMEOUT_SECONDS", "15"))) + except (TypeError, ValueError): + return 15.0 + + +MCP_CONNECT_TIMEOUT_SECONDS = _mcp_connect_timeout_seconds() +_MCP_REDACTED = "[REDACTED]" +_SENSITIVE_QUERY_KEYS = { + "access_token", "api_key", "apikey", "auth", "authorization", "client_secret", + "code", "key", "password", "refresh_token", "secret", "signature", "sig", + "token", +} +_SENSITIVE_QUERY_KEYS_NORMALIZED = { + re.sub(r"[^a-z0-9]", "", key.lower()) for key in _SENSITIVE_QUERY_KEYS +} + + +def sanitize_mcp_error(error: Any) -> str: + """Redact credentials from MCP-facing errors before status/API/log use.""" + text = str(error) if error is not None else "Unknown error" + text = _redact_urls(text) + text = re.sub( + r"""(?ix)(["']?authorization["']?\s*[:=]\s*["']?)(?:(?:bearer|basic)\s+)?[^\s"',;}\]]+""", + rf"\1{_MCP_REDACTED}", + text, + ) + text = re.sub( + r"(?i)\b(bearer|basic)\s+[A-Za-z0-9._~+/=-]+", + rf"\1 {_MCP_REDACTED}", + text, + ) + text = re.sub( + r"""(?ix) + \b([A-Za-z0-9_.-]*(?:api[_-]?key|token|password|passwd|pwd|secret)[A-Za-z0-9_.-]*) + (\s*[:=]\s*) + (?: + "([^"]*)"|'([^']*)'|([^\s,;&}\]]+) + ) + """, + lambda m: f"{m.group(1)}{m.group(2)}{_MCP_REDACTED}", + text, + ) + return text + + +def _redact_urls(text: str) -> str: + def repl(match): + raw = match.group(0) + try: + parts = urlsplit(raw) + netloc = parts.netloc + if "@" in netloc: + netloc = f"{_MCP_REDACTED}@{netloc.rsplit('@', 1)[1]}" + query = urlencode([ + (key, _MCP_REDACTED if _is_sensitive_query_key(key) else value) + for key, value in parse_qsl(parts.query, keep_blank_values=True) + ], doseq=True) + return urlunsplit((parts.scheme, netloc, parts.path, query, parts.fragment)) + except Exception: + return raw + + return re.sub(r"https?://[^\s\"'<>]+", repl, text) + + +def _is_sensitive_query_key(key: str) -> bool: + normalized = re.sub(r"[^a-z0-9]", "", key.lower()) + return key.lower() in _SENSITIVE_QUERY_KEYS or normalized in _SENSITIVE_QUERY_KEYS_NORMALIZED + + def _format_mcp_connection_error(name: str, command: str = "", args: Optional[List[str]] = None, error: Exception = None) -> str: """Return a user-actionable MCP connection error message.""" args = args or [] - raw_error = str(error) if error else "Unknown error" + raw_error = sanitize_mcp_error(error) command_line = " ".join([command or "", *args]).strip() lower_command = command_line.lower() @@ -172,13 +245,28 @@ async def connect_server( self._generation += 1 return res except Exception as e: - logger.error(f"Failed to connect MCP server {name} ({server_id}): {e}") + safe_error = sanitize_mcp_error(e) + logger.error(f"Failed to connect MCP server {name} ({server_id}): {safe_error}") error_message = _format_mcp_connection_error(name, command or "", args or [], e) self._connections[server_id] = {"status": "error", "error": error_message, "name": name} self._generation += 1 return False async def _connect_stdio(self, server_id: str, name: str, command: str, args: List[str], env: Dict[str, str]) -> bool: + """Connect to an MCP server via stdio transport.""" + import asyncio + try: + return await asyncio.wait_for( + self._connect_stdio_unbounded(server_id, name, command, args, env), + timeout=MCP_CONNECT_TIMEOUT_SECONDS, + ) + except asyncio.TimeoutError: + msg = f"MCP connection timed out after {MCP_CONNECT_TIMEOUT_SECONDS:g}s" + logger.warning(f"Failed to connect MCP server {name} ({server_id}): {msg}") + self._connections[server_id] = {"status": "error", "error": msg, "name": name} + return False + + async def _connect_stdio_unbounded(self, server_id: str, name: str, command: str, args: List[str], env: Dict[str, str]) -> bool: """Connect to an MCP server via stdio transport.""" try: from mcp import ClientSession, StdioServerParameters @@ -201,7 +289,7 @@ async def _connect_stdio(self, server_id: str, name: str, command: str, args: Li # Discover tools tools_result = await session.list_tools() - except Exception: + except BaseException: await stack.aclose() raise tools = [] @@ -246,6 +334,20 @@ async def _connect_stdio(self, server_id: str, name: str, command: str, args: Li return False async def _connect_sse(self, server_id: str, name: str, url: str) -> bool: + """Connect to an MCP server via SSE transport.""" + import asyncio + try: + return await asyncio.wait_for( + self._connect_sse_unbounded(server_id, name, url), + timeout=MCP_CONNECT_TIMEOUT_SECONDS, + ) + except asyncio.TimeoutError: + msg = f"MCP connection timed out after {MCP_CONNECT_TIMEOUT_SECONDS:g}s" + logger.warning(f"Failed to connect MCP server {name} ({server_id}): {msg}") + self._connections[server_id] = {"status": "error", "error": msg, "name": name} + return False + + async def _connect_sse_unbounded(self, server_id: str, name: str, url: str) -> bool: """Connect to an MCP server via SSE transport.""" try: from mcp import ClientSession @@ -262,7 +364,7 @@ async def _connect_sse(self, server_id: str, name: str, url: str) -> bool: # Discover tools tools_result = await session.list_tools() - except Exception: + except BaseException: await stack.aclose() raise tools = [] @@ -308,7 +410,9 @@ async def _start_http_connect(self, server_id: str, name: str, url: str, wait: f try: return task.result() except Exception as e: - self._connections[server_id] = {"status": "error", "error": str(e), "name": name} + safe_error = sanitize_mcp_error(e) + logger.error(f"Failed to connect HTTP MCP server {name} ({server_id}): {safe_error}") + self._connections[server_id] = {"status": "error", "error": safe_error, "name": name} return False # Still running → either awaiting authorization, or discovery/DCR is # still in flight. If _on_redirect already published needs_auth+auth_url, @@ -340,12 +444,16 @@ def _on_redirect(auth_url): provider = build_provider(server_id, url, on_redirect=_on_redirect) stack = AsyncExitStack() - transport = await stack.enter_async_context(streamablehttp_client(url, auth=provider)) - read_stream, write_stream, _get_session_id = transport - session = await stack.enter_async_context(ClientSession(read_stream, write_stream)) - await session.initialize() + try: + transport = await stack.enter_async_context(streamablehttp_client(url, auth=provider)) + read_stream, write_stream, _get_session_id = transport + session = await stack.enter_async_context(ClientSession(read_stream, write_stream)) + await session.initialize() - tools_result = await session.list_tools() + tools_result = await session.list_tools() + except BaseException: + await stack.aclose() + raise tools = [] for tool in tools_result.tools: tools.append({ @@ -373,8 +481,9 @@ def _on_redirect(auth_url): self._connections[server_id] = {"status": "error", "error": "mcp package not installed", "name": name} return False except Exception as e: - logger.error(f"Failed to connect HTTP MCP server {name} ({server_id}): {e}") - self._connections[server_id] = {"status": "error", "error": str(e), "name": name} + safe_error = sanitize_mcp_error(e) + logger.error(f"Failed to connect HTTP MCP server {name} ({server_id}): {safe_error}") + self._connections[server_id] = {"status": "error", "error": safe_error, "name": name} return False async def disconnect_server(self, server_id: str): @@ -395,7 +504,7 @@ async def disconnect_server(self, server_id: str): try: await stack.aclose() except Exception as e: - logger.warning(f"Error closing MCP server {server_id}: {e}") + logger.warning(f"Error closing MCP server {server_id}: {sanitize_mcp_error(e)}") self._sessions.pop(server_id, None) self._tools.pop(server_id, None) @@ -450,9 +559,10 @@ async def call_tool(self, qualified_name: str, arguments: Dict) -> Dict: try: result = await self._do_call(session, tool_name, arguments) except Exception as e: + safe_error = sanitize_mcp_error(e) # Auto-reconnect for builtin servers whose subprocess may have died if self.is_builtin(server_id): - logger.warning(f"MCP call failed for {qualified_name}, attempting reconnect: {e}") + logger.warning(f"MCP call failed for {qualified_name}, attempting reconnect: {safe_error}") reconnected = await self._reconnect_builtin(server_id) if reconnected: session = self._sessions.get(server_id) @@ -460,16 +570,17 @@ async def call_tool(self, qualified_name: str, arguments: Dict) -> Dict: try: result = await self._do_call(session, tool_name, arguments) except Exception as e2: - logger.error(f"MCP tool call failed after reconnect: {qualified_name}: {e2}") - return {"error": str(e2), "exit_code": 1} + safe_error2 = sanitize_mcp_error(e2) + logger.error(f"MCP tool call failed after reconnect: {qualified_name}: {safe_error2}") + return {"error": safe_error2, "exit_code": 1} else: return {"error": f"Reconnected but no session for {server_id}", "exit_code": 1} else: logger.error(f"MCP reconnect failed for {server_id}") return {"error": f"MCP server crashed and reconnect failed: {server_id}", "exit_code": 1} else: - logger.error(f"MCP tool call failed: {qualified_name}: {e}") - return {"error": str(e), "exit_code": 1} + logger.error(f"MCP tool call failed: {qualified_name}: {safe_error}") + return {"error": safe_error, "exit_code": 1} return result @@ -529,7 +640,7 @@ async def _reconnect_builtin(self, server_id: str) -> bool: logger.info(f"Reconnected builtin MCP server: {name}") return ok except Exception as e: - logger.error(f"Failed to reconnect builtin MCP server {name}: {e}") + logger.error(f"Failed to reconnect builtin MCP server {name}: {sanitize_mcp_error(e)}") return False def get_all_openai_schemas(self, disabled_map: Optional[Dict[str, set]] = None) -> List[Dict]: diff --git a/tests/test_mcp_connection_timeout_redaction.py b/tests/test_mcp_connection_timeout_redaction.py new file mode 100644 index 0000000000..45f9ad91e9 --- /dev/null +++ b/tests/test_mcp_connection_timeout_redaction.py @@ -0,0 +1,304 @@ +import asyncio +import sys +import types +from types import SimpleNamespace + +import pytest + +import src.mcp_manager as mcp_manager +from src.mcp_manager import McpManager, sanitize_mcp_error + + +SECRET = "super-secret-token" + + +class _TransportContext: + def __init__(self, state, *, hang_enter=False, fail_enter=None): + self.state = state + self.hang_enter = hang_enter + self.fail_enter = fail_enter + + async def __aenter__(self): + self.state["transport_entered"] = True + if self.fail_enter: + raise self.fail_enter + if self.hang_enter: + await asyncio.Event().wait() + return object(), object() + + async def __aexit__(self, exc_type, exc, tb): + self.state["transport_closed"] = True + + +class _SessionContext: + def __init__(self, state, read_stream, write_stream, *, hang_initialize=False): + self.state = state + self.hang_initialize = hang_initialize + + async def __aenter__(self): + self.state["session_entered"] = True + return self + + async def __aexit__(self, exc_type, exc, tb): + self.state["session_closed"] = True + + async def initialize(self): + self.state["initialized"] = True + if self.hang_initialize: + await asyncio.Event().wait() + + async def list_tools(self): + tool = SimpleNamespace( + name="echo", + description="Echo tool", + inputSchema={"type": "object", "properties": {}}, + annotations={"readOnlyHint": True}, + ) + return SimpleNamespace(tools=[tool]) + + async def call_tool(self, tool_name, arguments): + raise RuntimeError( + "Authorization: Bearer " + + SECRET + + " url=https://user:pass@example.com/mcp?api_key=" + + SECRET + + " password=" + + SECRET + ) + + +@pytest.fixture +def short_mcp_timeout(monkeypatch): + monkeypatch.setattr(mcp_manager, "MCP_CONNECT_TIMEOUT_SECONDS", 0.01) + + +def _install_fake_mcp(monkeypatch, *, hang_enter=False, hang_initialize=False, fail_enter=None): + state = {} + mcp = types.ModuleType("mcp") + mcp.ClientSession = lambda read, write: _SessionContext( + state, read, write, hang_initialize=hang_initialize + ) + mcp.StdioServerParameters = lambda command, args, env: SimpleNamespace( + command=command, args=args, env=env + ) + + client = types.ModuleType("mcp.client") + stdio = types.ModuleType("mcp.client.stdio") + sse = types.ModuleType("mcp.client.sse") + stdio.stdio_client = lambda params: _TransportContext( + state, hang_enter=hang_enter, fail_enter=fail_enter + ) + sse.sse_client = lambda url: _TransportContext( + state, hang_enter=hang_enter, fail_enter=fail_enter + ) + + monkeypatch.setitem(sys.modules, "mcp", mcp) + monkeypatch.setitem(sys.modules, "mcp.client", client) + monkeypatch.setitem(sys.modules, "mcp.client.stdio", stdio) + monkeypatch.setitem(sys.modules, "mcp.client.sse", sse) + return state + + +async def test_stdio_connection_times_out_and_cleans_partial_resources(monkeypatch, short_mcp_timeout): + state = _install_fake_mcp(monkeypatch, hang_initialize=True) + mgr = McpManager() + + ok = await mgr.connect_server( + "srv", "Hanging stdio", "stdio", command="fake", args=[], env={} + ) + + assert ok is False + assert mgr.get_server_status("srv")["status"] == "error" + assert "timed out" in mgr.get_server_status("srv")["error"] + assert state["session_closed"] is True + assert state["transport_closed"] is True + assert "srv" not in mgr._sessions + assert "srv" not in mgr._stacks + + +async def test_sse_connection_times_out(monkeypatch, short_mcp_timeout): + state = _install_fake_mcp(monkeypatch, hang_initialize=True) + mgr = McpManager() + + ok = await mgr.connect_server("srv", "Hanging SSE", "sse", url="https://example.com/sse") + + assert ok is False + assert mgr.get_server_status("srv")["status"] == "error" + assert "timed out" in mgr.get_server_status("srv")["error"] + assert state["transport_closed"] is True + + +async def test_builtin_reconnect_returns_after_connection_timeout(monkeypatch, short_mcp_timeout): + _install_fake_mcp(monkeypatch, hang_initialize=True) + import src.builtin_mcp as builtin_mcp + + monkeypatch.setitem(builtin_mcp._BUILTIN_SERVERS, "builtin_hang", ("server.py", "Hang")) + monkeypatch.setattr(builtin_mcp, "builtin_python_env", lambda base_dir: {}) + + mgr = McpManager() + ok = await mgr._reconnect_builtin("builtin_hang") + + assert ok is False + status = mgr.get_server_status("builtin_hang") + assert status["status"] == "error" + assert "timed out" in status["error"] + + +async def test_cancelled_stdio_connection_cleans_partial_resources(monkeypatch): + state = _install_fake_mcp(monkeypatch, hang_initialize=True) + monkeypatch.setattr(mcp_manager, "MCP_CONNECT_TIMEOUT_SECONDS", 3600) + mgr = McpManager() + + task = asyncio.create_task( + mgr.connect_server("srv", "Cancelled stdio", "stdio", command="fake", args=[], env={}) + ) + while not state.get("initialized"): + await asyncio.sleep(0) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert state["session_closed"] is True + assert state["transport_closed"] is True + assert "srv" not in mgr._sessions + assert "srv" not in mgr._stacks + + +async def test_connection_status_api_and_logs_use_sanitized_errors(monkeypatch, caplog): + raw = RuntimeError( + "Authorization: Bearer " + + SECRET + + " token=" + + SECRET + + " https://user:pass@example.com/mcp?api_key=" + + SECRET + + "&ok=1" + ) + _install_fake_mcp(monkeypatch, fail_enter=raw) + mgr = McpManager() + + with caplog.at_level("ERROR", logger="src.mcp_manager"): + ok = await mgr.connect_server("srv", "Leaky stdio", "stdio", command="fake", args=[], env={}) + + status_error = mgr.get_server_status("srv")["error"] + assert ok is False + assert SECRET not in status_error + assert "user:pass" not in status_error + assert _has_redaction(status_error) + assert SECRET not in caplog.text + assert "user:pass" not in caplog.text + assert _has_redaction(caplog.text) + + mgr._sessions["remote"] = _SessionContext({}, object(), object()) + api_result = await mgr.call_tool("mcp__remote__echo", {}) + + assert api_result["exit_code"] == 1 + assert SECRET not in api_result["error"] + assert "user:pass" not in api_result["error"] + assert _has_redaction(api_result["error"]) + + +@pytest.mark.parametrize( + ("category", "raw", "forbidden", "expected"), + [ + pytest.param( + "authorization_bearer", + "Authorization: Bearer " + SECRET, + [SECRET], + ["Authorization: [REDACTED]"], + id="authorization_bearer", + ), + pytest.param( + "authorization_basic", + "Authorization: Basic " + SECRET, + [SECRET], + ["Authorization: [REDACTED]"], + id="authorization_basic", + ), + pytest.param( + "standalone_bearer", + "Bearer " + SECRET, + [SECRET], + ["Bearer [REDACTED]"], + id="standalone_bearer", + ), + pytest.param( + "standalone_basic", + "Basic " + SECRET, + [SECRET], + ["Basic [REDACTED]"], + id="standalone_basic", + ), + pytest.param( + "api_key_assignment", + "API_KEY='" + SECRET + "'", + [SECRET], + ["API_KEY=[REDACTED]"], + id="api_key_assignment", + ), + pytest.param( + "token_assignment", + "token=" + SECRET, + [SECRET], + ["token=[REDACTED]"], + id="token_assignment", + ), + pytest.param( + "password_assignment", + "password=" + SECRET, + [SECRET], + ["password=[REDACTED]"], + id="password_assignment", + ), + pytest.param( + "secret_assignment", + "client_secret=" + SECRET, + [SECRET], + ["client_secret=[REDACTED]"], + id="secret_assignment", + ), + pytest.param( + "url_userinfo", + "https://user:pass@example.test:8443/mcp/path?safe=value&mode=ok#frag", + ["user:pass"], + ["https://[REDACTED]@example.test:8443/mcp/path?safe=value&mode=ok#frag"], + id="url_userinfo", + ), + pytest.param( + "url_sensitive_query", + "https://example.test/mcp?token=" + SECRET + "&safe=value&mode=ok#frag", + [SECRET], + ["token=[REDACTED]", "safe=value", "mode=ok", "#frag"], + id="url_sensitive_query", + ), + ], +) +def test_sanitizer_redacts_supported_secret_shapes(category, raw, forbidden, expected): + text = sanitize_mcp_error(raw) + + for value in forbidden: + if value in text: + pytest.fail(f"{category} leaked a synthetic secret") + for value in expected: + if value not in text: + pytest.fail(f"{category} missing expected sanitized output") + if not _has_redaction(text): + pytest.fail(f"{category} missing redaction marker") + + +async def test_successful_stdio_connection_still_registers_tools(monkeypatch): + _install_fake_mcp(monkeypatch) + mgr = McpManager() + + ok = await mgr.connect_server("srv", "Working stdio", "stdio", command="fake", args=[], env={}) + + assert ok is True + assert mgr.get_server_status("srv")["status"] == "connected" + assert mgr.get_server_status("srv")["tool_count"] == 1 + assert mgr.get_all_tools()[0]["qualified_name"] == "mcp__srv__echo" + await mgr.disconnect_server("srv") + + +def _has_redaction(text): + return "[REDACTED]" in text From 4587b1174c4f5d471def52d438df0d053dc8beac Mon Sep 17 00:00:00 2001 From: Kaiki Date: Sun, 12 Jul 2026 20:32:16 -0400 Subject: [PATCH 03/12] fix(io): harden atomic file replacement --- core/atomic_io.py | 87 ++++++++++++++++++++++++------- tests/test_atomic_io_hardening.py | 74 ++++++++++++++++++++++++++ 2 files changed, 142 insertions(+), 19 deletions(-) create mode 100644 tests/test_atomic_io_hardening.py diff --git a/core/atomic_io.py b/core/atomic_io.py index 81c640d8aa..71406a4e4c 100644 --- a/core/atomic_io.py +++ b/core/atomic_io.py @@ -15,31 +15,80 @@ import json import os +import stat from typing import Any, Optional +def _fsync_parent_directory(path: str) -> None: + """Best-effort durability for a completed same-directory rename.""" + if os.name != "posix": + return + + directory = os.path.dirname(os.path.abspath(path)) or "." + flags = os.O_RDONLY + if hasattr(os, "O_DIRECTORY"): + flags |= os.O_DIRECTORY + + descriptor = os.open(directory, flags) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def _atomic_write(path: str, writer, *, mode: Optional[int] = None) -> None: + import tempfile + + directory = os.path.dirname(os.path.abspath(path)) or "." + os.makedirs(directory, exist_ok=True) + + existing_mode = None + try: + existing_mode = stat.S_IMODE(os.stat(path, follow_symlinks=False).st_mode) + except FileNotFoundError: + pass + + effective_mode = mode if mode is not None else existing_mode + descriptor, temporary = tempfile.mkstemp( + dir=directory, + prefix=f".{os.path.basename(path)}.", + suffix=".tmp", + ) + + try: + if effective_mode is not None: + os.fchmod(descriptor, effective_mode) + + with os.fdopen(descriptor, "w", encoding="utf-8") as stream: + descriptor = -1 + writer(stream) + stream.flush() + os.fsync(stream.fileno()) + + os.replace(temporary, path) + _fsync_parent_directory(path) + except BaseException: + if descriptor >= 0: + os.close(descriptor) + try: + os.unlink(temporary) + except FileNotFoundError: + pass + raise + + def atomic_write_json(path: str, data: Any, *, indent: Optional[int] = None) -> None: - """Atomically persist `data` as JSON at `path`. + def write(stream) -> None: + json.dump(data, stream, ensure_ascii=False, indent=indent) - The temp file uses the live PID as a suffix so two processes saving the - same file (e.g. unit tests) don't collide on the rename target. - """ - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - tmp = f"{path}.tmp.{os.getpid()}" - with open(tmp, "w", encoding="utf-8") as f: - json.dump(data, f, indent=indent) - f.flush() - os.fsync(f.fileno()) - os.replace(tmp, path) + _atomic_write(path, write) -def atomic_write_text(path: str, text: str) -> None: +def atomic_write_text(path: str, text: str, *, mode: Optional[int] = None) -> None: if not isinstance(text, str): raise TypeError("atomic_write_text expects a string") - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - tmp = f"{path}.tmp.{os.getpid()}" - with open(tmp, "w", encoding="utf-8") as f: - f.write(text) - f.flush() - os.fsync(f.fileno()) - os.replace(tmp, path) + + def write(stream) -> None: + stream.write(text) + + _atomic_write(path, write, mode=mode) diff --git a/tests/test_atomic_io_hardening.py b/tests/test_atomic_io_hardening.py new file mode 100644 index 0000000000..ea5b3dca85 --- /dev/null +++ b/tests/test_atomic_io_hardening.py @@ -0,0 +1,74 @@ +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import pytest + +import core.atomic_io as atomic_io + + +def test_concurrent_writes_use_unique_temporary_files(tmp_path): + target = tmp_path / "state.txt" + payloads = [f"payload-{index}-" * 1000 for index in range(32)] + + with ThreadPoolExecutor(max_workers=8) as executor: + futures = [ + executor.submit(atomic_io.atomic_write_text, str(target), payload) + for payload in payloads + ] + for future in futures: + future.result() + + assert target.read_text(encoding="utf-8") in payloads + assert list(tmp_path.glob(".*.tmp")) == [] + + +def test_failed_replace_preserves_target_and_cleans_temporary_file( + tmp_path, monkeypatch +): + target = tmp_path / "state.txt" + target.write_text("original", encoding="utf-8") + + def fail_replace(_source, _target): + raise OSError("synthetic replace failure") + + monkeypatch.setattr(atomic_io.os, "replace", fail_replace) + + with pytest.raises(OSError, match="synthetic replace failure"): + atomic_io.atomic_write_text(str(target), "replacement") + + assert target.read_text(encoding="utf-8") == "original" + assert list(tmp_path.glob(".*.tmp")) == [] + + +def test_existing_mode_is_preserved(tmp_path): + target = tmp_path / "state.txt" + target.write_text("old", encoding="utf-8") + target.chmod(0o640) + + atomic_io.atomic_write_text(str(target), "new") + + assert target.read_text(encoding="utf-8") == "new" + assert target.stat().st_mode & 0o777 == 0o640 + + +def test_explicit_mode_is_applied_to_new_file(tmp_path): + target = tmp_path / "state.txt" + + atomic_io.atomic_write_text(str(target), "new", mode=0o600) + + assert target.stat().st_mode & 0o777 == 0o600 + + +def test_parent_directory_is_fsynced_after_replace(tmp_path, monkeypatch): + target = tmp_path / "state.txt" + calls = [] + + monkeypatch.setattr( + atomic_io, + "_fsync_parent_directory", + lambda path: calls.append(path), + ) + + atomic_io.atomic_write_text(str(target), "new") + + assert calls == [str(target)] From 3525f633dc2336d4e0686ee7a3c93db23bca6e97 Mon Sep 17 00:00:00 2001 From: Kaiki Date: Sun, 12 Jul 2026 20:32:16 -0400 Subject: [PATCH 04/12] fix(tools): make agent file writes atomic --- src/agent_tools/filesystem_tools.py | 23 ++- tests/test_agent_file_write_atomicity.py | 193 +++++++++++++++++++++++ 2 files changed, 210 insertions(+), 6 deletions(-) create mode 100644 tests/test_agent_file_write_atomicity.py diff --git a/src/agent_tools/filesystem_tools.py b/src/agent_tools/filesystem_tools.py index 52dcb60f86..16aa492af2 100644 --- a/src/agent_tools/filesystem_tools.py +++ b/src/agent_tools/filesystem_tools.py @@ -7,6 +7,7 @@ import shutil from typing import Optional, Dict, Any, Tuple +from core.atomic_io import atomic_write_text from src.constants import MAX_READ_CHARS, MAX_DIFF_LINES, MAX_OUTPUT_CHARS _CODENAV_SKIP_DIRS = frozenset({ @@ -70,6 +71,15 @@ def _unified_diff(old: str, new: str, path: str) -> Optional[Dict[str, Any]]: "file": os.path.basename(path) or (path or "file"), } +def _regular_file_mode(path: str) -> Optional[int]: + try: + st = os.stat(path) + except FileNotFoundError: + return None + if not os.path.isfile(path): + raise IsADirectoryError(path) + return st.st_mode & 0o7777 + class EditFileTool: async def execute(self, content: str, ctx: dict) -> dict: from src.tool_execution import _resolve_tool_path, _resolve_search_root, _truncate @@ -94,6 +104,7 @@ async def execute(self, content: str, ctx: dict) -> dict: def _apply(): """Helper function that performs the actual string replacement and file writing logic.""" + mode = _regular_file_mode(path) with open(path, "r", encoding="utf-8") as f: original = f.read() count = original.count(old) @@ -102,8 +113,7 @@ def _apply(): if count > 1 and not replace_all: return original, None, f"not_unique:{count}" updated = original.replace(old, new) if replace_all else original.replace(old, new, 1) - with open(path, "w", encoding="utf-8") as f: - f.write(updated) + atomic_write_text(path, updated, mode=mode) return original, updated, "ok" try: @@ -208,16 +218,17 @@ async def execute(self, content: str, ctx: dict) -> dict: try: def _write(): old = "" + mode = _regular_file_mode(path) try: - with open(path, "r", encoding="utf-8") as f: - old = f.read() + if mode is not None: + with open(path, "r", encoding="utf-8") as f: + old = f.read() except (FileNotFoundError, IsADirectoryError, UnicodeDecodeError, OSError): old = "" d = os.path.dirname(path) if d: os.makedirs(d, exist_ok=True) - with open(path, "w", encoding="utf-8") as f: - f.write(body) + atomic_write_text(path, body, mode=mode) return old, len(body) old_content, size = await asyncio.to_thread(_write) except PermissionError: diff --git a/tests/test_agent_file_write_atomicity.py b/tests/test_agent_file_write_atomicity.py new file mode 100644 index 0000000000..7e27dee457 --- /dev/null +++ b/tests/test_agent_file_write_atomicity.py @@ -0,0 +1,193 @@ +import asyncio +import json +import os +import stat + +import pytest + +import core.atomic_io as atomic_io +import src.agent_tools.filesystem_tools as filesystem_tools +from src.agent_tools.filesystem_tools import EditFileTool, WriteFileTool + + +def _write_payload(path, content): + return json.dumps({"path": str(path), "content": content}) + + +def _edit_payload(path, old, new, **extra): + payload = {"path": str(path), "old_string": old, "new_string": new} + payload.update(extra) + return json.dumps(payload) + + +@pytest.mark.asyncio +async def test_write_file_successful_replacement_preserves_result_and_diff(tmp_path): + target = tmp_path / "note.txt" + target.write_text("old line\nsame\n", encoding="utf-8") + + result = await WriteFileTool().execute(_write_payload(target, "new line\nsame\n"), {}) + + assert result["exit_code"] == 0 + assert result["output"] == f"Wrote 14 bytes to {target}" + assert target.read_text(encoding="utf-8") == "new line\nsame\n" + assert result["diff"]["added"] == 1 + assert result["diff"]["removed"] == 1 + assert result["diff"]["file"] == "note.txt" + assert result["diff"]["new_file"] is False + + +@pytest.mark.asyncio +async def test_edit_file_successful_replacement_preserves_result_and_diff(tmp_path): + target = tmp_path / "code.py" + target.write_text("def f():\n return 1\n", encoding="utf-8") + + result = await EditFileTool().execute( + _edit_payload(target, "return 1", "return 2"), + {}, + ) + + assert result["exit_code"] == 0 + assert result["output"] == f"Edited {target} (1 replacement)" + assert target.read_text(encoding="utf-8") == "def f():\n return 2\n" + assert result["diff"]["added"] == 1 + assert result["diff"]["removed"] == 1 + assert result["diff"]["file"] == "code.py" + assert result["diff"]["new_file"] is False + + +@pytest.mark.asyncio +async def test_write_file_partial_final_write_failure_preserves_existing_target( + tmp_path, monkeypatch +): + target = tmp_path / "partial.txt" + target.write_text("original-content", encoding="utf-8") + target.chmod(0o640) + + real_fdopen = atomic_io.os.fdopen + + class FailingStream: + def __init__(self, stream): + self._stream = stream + + def __enter__(self): + self._stream.__enter__() + return self + + def __exit__(self, exc_type, exc_value, traceback): + return self._stream.__exit__(exc_type, exc_value, traceback) + + def write(self, content): + self._stream.write(content[: max(1, len(content) // 2)]) + self._stream.flush() + raise OSError("simulated partial write") + + def __getattr__(self, name): + return getattr(self._stream, name) + + def flaky_fdopen(*args, **kwargs): + return FailingStream(real_fdopen(*args, **kwargs)) + + monkeypatch.setattr(atomic_io.os, "fdopen", flaky_fdopen) + + result = await WriteFileTool().execute( + _write_payload(target, "replacement-content"), + {}, + ) + + assert result["exit_code"] == 1 + assert "simulated partial write" in result["error"] + assert target.read_text(encoding="utf-8") == "original-content" + assert target.stat().st_mode & 0o777 == 0o640 + assert list(tmp_path.glob(".*.tmp")) == [] + + +async def test_write_file_cancellation_before_replace_preserves_existing_target(tmp_path, monkeypatch): + target = tmp_path / "cancelled.txt" + original = "stable content\n" + target.write_text(original, encoding="utf-8") + + def cancelled(path, text, **kwargs): + raise asyncio.CancelledError() + + monkeypatch.setattr(filesystem_tools, "atomic_write_text", cancelled) + + with pytest.raises(asyncio.CancelledError): + await WriteFileTool().execute(_write_payload(target, "new content\n"), {}) + + assert target.read_text(encoding="utf-8") == original + + +@pytest.mark.asyncio +async def test_write_file_failure_before_replace_preserves_existing_target(tmp_path, monkeypatch): + target = tmp_path / "failed.txt" + original = b"original bytes\n" + target.write_bytes(original) + + def failed(path, text, **kwargs): + tmp = f"{path}.tmp.test" + with open(tmp, "w", encoding="utf-8") as f: + f.write(text[:4]) + raise OSError("staged write failed") + + monkeypatch.setattr(filesystem_tools, "atomic_write_text", failed) + + result = await WriteFileTool().execute(_write_payload(target, "replacement\n"), {}) + + assert result["exit_code"] == 1 + assert "staged write failed" in result["error"] + assert target.read_bytes() == original + + +@pytest.mark.asyncio +async def test_write_file_new_file_creation_still_supported(tmp_path): + target = tmp_path / "new" / "created.txt" + + result = await WriteFileTool().execute(_write_payload(target, "created\n"), {}) + + assert result["exit_code"] == 0 + assert target.read_text(encoding="utf-8") == "created\n" + assert result["diff"]["new_file"] is True + + +@pytest.mark.asyncio +async def test_write_file_preserves_existing_regular_file_permissions(tmp_path): + target = tmp_path / "mode.txt" + target.write_text("old\n", encoding="utf-8") + target.chmod(0o640) + + result = await WriteFileTool().execute(_write_payload(target, "new\n"), {}) + + assert result["exit_code"] == 0 + assert stat.S_IMODE(target.stat().st_mode) == 0o640 + assert target.read_text(encoding="utf-8") == "new\n" + + +@pytest.mark.asyncio +async def test_edit_file_preserves_existing_regular_file_permissions(tmp_path): + target = tmp_path / "edit-mode.txt" + target.write_text("old value\n", encoding="utf-8") + target.chmod(0o640) + + result = await EditFileTool().execute(_edit_payload(target, "old", "new"), {}) + + assert result["exit_code"] == 0 + assert stat.S_IMODE(target.stat().st_mode) == 0o640 + assert target.read_text(encoding="utf-8") == "new value\n" + + +@pytest.mark.asyncio +async def test_edit_file_failure_before_replace_preserves_existing_target(tmp_path, monkeypatch): + target = tmp_path / "edit-failed.txt" + original = "alpha beta gamma\n" + target.write_text(original, encoding="utf-8") + + def failed(path, text, **kwargs): + raise OSError("replace never reached") + + monkeypatch.setattr(filesystem_tools, "atomic_write_text", failed) + + result = await EditFileTool().execute(_edit_payload(target, "beta", "BETA"), {}) + + assert result["exit_code"] == 1 + assert "replace never reached" in result["error"] + assert target.read_text(encoding="utf-8") == original From 6baba35ff112b3dcb8cb3efabd558c66a5617b56 Mon Sep 17 00:00:00 2001 From: Kaiki Date: Sun, 12 Jul 2026 20:32:16 -0400 Subject: [PATCH 05/12] fix(runtime): update compatibility and resource lifecycles --- core/database.py | 2 +- routes/gallery/gallery_routes.py | 31 +++++++++++++++++++++++++------ 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/core/database.py b/core/database.py index a9ad90b8b7..354d670b96 100644 --- a/core/database.py +++ b/core/database.py @@ -8,7 +8,7 @@ from sqlalchemy import event, create_engine, Column, String, Text, Boolean, DateTime, Integer, ForeignKey, JSON, Index, func, text from sqlalchemy.engine import Engine, make_url from sqlalchemy.types import TypeDecorator -from sqlalchemy.ext.declarative import declarative_base, declared_attr +from sqlalchemy.orm import declarative_base, declared_attr from sqlalchemy.orm import relationship, sessionmaker, backref from src.runtime_paths import get_app_root diff --git a/routes/gallery/gallery_routes.py b/routes/gallery/gallery_routes.py index 5b20086dd9..c296059fea 100644 --- a/routes/gallery/gallery_routes.py +++ b/routes/gallery/gallery_routes.py @@ -257,10 +257,15 @@ async def gallery_upload(request: Request): @router.post("/api/gallery/{image_id}/replace") async def gallery_replace(request: Request, image_id: str): """Replace an existing gallery image file with a new one.""" - user = get_current_user(request) + form = None db = SessionLocal() try: - img = db.query(GalleryImage).filter(GalleryImage.id == image_id).first() + user = get_current_user(request) + img = ( + db.query(GalleryImage) + .filter(GalleryImage.id == image_id) + .first() + ) if not img: raise HTTPException(404, "Image not found") if not user or img.owner != user: @@ -268,10 +273,14 @@ async def gallery_replace(request: Request, image_id: str): form = await request.form() file = form.get("image") - if not file or not hasattr(file, 'read'): + if not file or not hasattr(file, "read"): raise HTTPException(400, "No image provided") - content = await read_upload_limited(file, GALLERY_UPLOAD_MAX_BYTES, "Gallery replacement") + content = await read_upload_limited( + file, + GALLERY_UPLOAD_MAX_BYTES, + "Gallery replacement", + ) GALLERY_IMAGE_DIR.mkdir(parents=True, exist_ok=True) img_path = _gallery_image_path(img.filename) img_path.write_bytes(content) @@ -279,21 +288,31 @@ async def gallery_replace(request: Request, image_id: str): # Refresh dimensions in case the editor resized the canvas. # updated_at auto-bumps via TimestampMixin's onupdate hook. try: - from PIL import Image from io import BytesIO + + from PIL import Image + with Image.open(BytesIO(content)) as new_im: img.width = new_im.width img.height = new_im.height except Exception: pass + try: db.commit() except Exception: db.rollback() logger.exception("gallery_replace: DB commit failed") raise HTTPException(500, "Image update failed") - return {"ok": True, "width": img.width, "height": img.height} + + return { + "ok": True, + "width": img.width, + "height": img.height, + } finally: + if form is not None: + await form.close() db.close() # ---- POST /api/gallery/{image_id}/rename ---- From 9016b3285b6d1d334a10ee3497ee3db58af20f27 Mon Sep 17 00:00:00 2001 From: Kaiki Date: Sun, 12 Jul 2026 20:32:16 -0400 Subject: [PATCH 06/12] test: close resources and enforce clean async tests --- pyproject.toml | 4 +++ tests/test_auth_regressions.py | 9 ++++-- tests/test_caldav_redirect_hardening.py | 29 +++++++++++++++---- tests/test_document_tool_owner_scope.py | 3 +- tests/test_edit_file.py | 13 +++++---- tests/test_gallery_filename_confinement.py | 10 +++---- .../test_memory_extractor_vector_degraded.py | 8 ++++- tests/test_memory_owner_isolation.py | 1 + tests/test_ollama_runner_hint.py | 3 +- tests/test_research_service.py | 8 ++++- tests/test_research_utils.py | 2 ++ tests/test_scheduler_restart_doublefire.py | 7 ++++- tests/test_tool_index_schema_parity.py | 9 ++++-- tests/test_upload_handler_atomicity.py | 16 +++++++--- tests/test_vault_password_not_in_argv.py | 3 +- tests/test_webhook_dns_rebinding_pin.py | 11 +++++-- 16 files changed, 101 insertions(+), 35 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index da00ee2599..4f09934a40 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,4 +1,5 @@ [tool.pytest.ini_options] +asyncio_default_fixture_loop_scope = "function" testpaths = ["tests"] asyncio_mode = "auto" # Test-taxonomy markers added at collection time by tests/conftest.py. The @@ -19,4 +20,7 @@ markers = [ # taxonomy. The fast lane runs `not slow`; mark a test slow only with # duration evidence (see tests/run_focus.py --durations and tests/README.md). "slow: opt-in marker for known-slow tests; excluded by the fast lane (not slow)", + "asyncio: repository test taxonomy marker", + "parametrize: repository test taxonomy marker", + "skipif: repository test taxonomy marker", ] diff --git a/tests/test_auth_regressions.py b/tests/test_auth_regressions.py index 62b4797482..4c6c8b1ee5 100644 --- a/tests/test_auth_regressions.py +++ b/tests/test_auth_regressions.py @@ -356,8 +356,10 @@ def test_ship_paused_housekeeping_stays_paused_by_default(): from routes import task_routes from src import task_scheduler - route_src = open(task_routes.__file__, encoding="utf-8").read() - scheduler_src = open(task_scheduler.__file__, encoding="utf-8").read() + with open(task_routes.__file__, encoding="utf-8") as fh: + route_src = fh.read() + with open(task_scheduler.__file__, encoding="utf-8") as fh: + scheduler_src = fh.read() assert '"ship_paused": True' in scheduler_src assert 'defs.get("ship_paused")' in scheduler_src assert 'defs.get("ship_paused")' in route_src @@ -366,5 +368,6 @@ def test_ship_paused_housekeeping_stays_paused_by_default(): def test_task_payload_exposes_crew_member_id_for_ui_category(): from routes import task_routes - src = open(task_routes.__file__, encoding="utf-8").read() + with open(task_routes.__file__, encoding="utf-8") as fh: + src = fh.read() assert '"crew_member_id"' in src diff --git a/tests/test_caldav_redirect_hardening.py b/tests/test_caldav_redirect_hardening.py index 0d3ce91b76..9efe972986 100644 --- a/tests/test_caldav_redirect_hardening.py +++ b/tests/test_caldav_redirect_hardening.py @@ -60,16 +60,23 @@ def log_message(self, *a): internal = socketserver.TCPServer(("127.0.0.1", 0), _Internal) internal_port = internal.server_address[1] - public = socketserver.TCPServer(("127.0.0.1", 0), _Public) - public_port = public.server_address[1] - threading.Thread(target=internal.serve_forever, daemon=True).start() - threading.Thread(target=public.serve_forever, daemon=True).start() + public = None + internal_thread = threading.Thread(target=internal.serve_forever, daemon=True) + public_thread = None + client = None + response = None try: + public = socketserver.TCPServer(("127.0.0.1", 0), _Public) + public_port = public.server_address[1] + public_thread = threading.Thread(target=public.serve_forever, daemon=True) + internal_thread.start() + public_thread.start() + public_url = f"http://127.0.0.1:{public_port}/dav" client = caldav_sync._build_dav_client(public_url, "u", "p") client.timeout = 5 try: - client.request(public_url, "PROPFIND", "") + response = client.request(public_url, "PROPFIND", "") except Exception: # Refusing the redirect surfaces as an exception (TooManyRedirects); # that is the intended fail-closed behavior. The security assertion @@ -80,8 +87,18 @@ def log_message(self, *a): assert public_methods == ["PROPFIND"], "the PROPFIND must reach the public server first" assert sink_hits == [], "redirect toward an internal host must not be followed" finally: + if response is not None and hasattr(response, "close"): + response.close() + if client is not None and hasattr(client, "session") and hasattr(client.session, "close"): + client.session.close() + if public is not None: + public.shutdown() + public.server_close() internal.shutdown() - public.shutdown() + internal.server_close() + if public_thread is not None: + public_thread.join() + internal_thread.join() def test_sync_and_writeback_construct_clients_through_the_helper(): diff --git a/tests/test_document_tool_owner_scope.py b/tests/test_document_tool_owner_scope.py index 21d5ad9ce9..056b93a516 100644 --- a/tests/test_document_tool_owner_scope.py +++ b/tests/test_document_tool_owner_scope.py @@ -154,7 +154,8 @@ def test_suggest_document_active_id_filters_to_calling_owner(monkeypatch): def test_document_tool_dispatch_forwards_owner(): - source = open("src/tool_execution.py", encoding="utf-8").read() + with open("src/tool_execution.py", encoding="utf-8") as fh: + source = fh.read() assert "_document_tool_dispatch(tool, content, session_id, owner)" in source diff --git a/tests/test_edit_file.py b/tests/test_edit_file.py index 6af22fb5dc..31a3d78cf6 100644 --- a/tests/test_edit_file.py +++ b/tests/test_edit_file.py @@ -2,6 +2,7 @@ import json import os import tempfile +from pathlib import Path import pytest @@ -46,7 +47,7 @@ async def test_edit_file_blocked_at_execution_for_non_admin(monkeypatch): monkeypatch.setattr(te, "_owner_is_admin", lambda owner: False) ws = tempfile.mkdtemp() p = os.path.join("/tmp", "ef_block.txt") - open(p, "w").write("a\n") + Path(p).write_text("a\n", encoding="utf-8") _desc, result = await te.execute_tool_block( ToolBlock("edit_file", json.dumps({"path": p, "old_string": "a", "new_string": "b"})), owner="bob", @@ -59,10 +60,10 @@ async def test_edit_file_blocked_at_execution_for_non_admin(monkeypatch): @pytest.mark.asyncio async def test_edit_file_success(): p = os.path.join("/tmp", "ef_ok.py") - open(p, "w").write("def f():\n return 1\n") + Path(p).write_text("def f():\n return 1\n", encoding="utf-8") res = await EditFileTool().execute(json.dumps({"path": p, "old_string": "return 1", "new_string": "return 2"}), {}) assert res["exit_code"] == 0 - assert open(p).read() == "def f():\n return 2\n" + assert Path(p).read_text(encoding="utf-8") == "def f():\n return 2\n" assert res["diff"]["added"] == 1 and res["diff"]["removed"] == 1 and res["diff"]["file"] == "ef_ok.py" os.unlink(p) @@ -70,7 +71,7 @@ async def test_edit_file_success(): @pytest.mark.asyncio async def test_edit_file_not_found(): p = os.path.join("/tmp", "ef_nf.txt") - open(p, "w").write("hello\n") + Path(p).write_text("hello\n", encoding="utf-8") res = await EditFileTool().execute(json.dumps({"path": p, "old_string": "nope", "new_string": "x"}), {}) assert res["exit_code"] == 1 and "not found" in res["error"] os.unlink(p) @@ -79,12 +80,12 @@ async def test_edit_file_not_found(): @pytest.mark.asyncio async def test_edit_file_non_unique(): p = os.path.join("/tmp", "ef_dup.txt") - open(p, "w").write("x\nx\n") + Path(p).write_text("x\nx\n", encoding="utf-8") res = await EditFileTool().execute(json.dumps({"path": p, "old_string": "x", "new_string": "y"}), {}) assert res["exit_code"] == 1 and "not unique" in res["error"] # replace_all resolves it res = await EditFileTool().execute(json.dumps({"path": p, "old_string": "x", "new_string": "y", "replace_all": True}), {}) - assert res["exit_code"] == 0 and open(p).read() == "y\ny\n" + assert res["exit_code"] == 0 and Path(p).read_text(encoding="utf-8") == "y\ny\n" os.unlink(p) diff --git a/tests/test_gallery_filename_confinement.py b/tests/test_gallery_filename_confinement.py index 25e41ad5b6..82127f7c3a 100644 --- a/tests/test_gallery_filename_confinement.py +++ b/tests/test_gallery_filename_confinement.py @@ -116,12 +116,12 @@ def test_gallery_replace_rejects_symlink_escape(tmp_path, monkeypatch): app = FastAPI() app.include_router(gallery_routes.setup_gallery_routes()) - client = TestClient(app) - response = client.post( - "/api/gallery/img-1/replace", - files={"image": ("replacement.png", b"replacement bytes", "image/png")}, - ) + with TestClient(app) as client: + response = client.post( + "/api/gallery/img-1/replace", + files={"image": ("replacement.png", b"replacement bytes", "image/png")}, + ) assert response.status_code == 400 assert outside.read_bytes() == b"outside image root" diff --git a/tests/test_memory_extractor_vector_degraded.py b/tests/test_memory_extractor_vector_degraded.py index 1b3bd24756..61c8291944 100644 --- a/tests/test_memory_extractor_vector_degraded.py +++ b/tests/test_memory_extractor_vector_degraded.py @@ -49,7 +49,13 @@ def add(self, memory_id, text): def _run(coro): - return asyncio.new_event_loop().run_until_complete(coro) + loop = asyncio.new_event_loop() + try: + return loop.run_until_complete( + coro + ) + finally: + loop.close() def test_extraction_persists_facts_when_vector_store_fails_at_runtime(monkeypatch): diff --git a/tests/test_memory_owner_isolation.py b/tests/test_memory_owner_isolation.py index ff32b9cd11..475b612612 100644 --- a/tests/test_memory_owner_isolation.py +++ b/tests/test_memory_owner_isolation.py @@ -1,5 +1,6 @@ from unittest.mock import MagicMock + import routes.memory_routes as memory_routes from src.memory import MemoryManager diff --git a/tests/test_ollama_runner_hint.py b/tests/test_ollama_runner_hint.py index 97a5d69e86..5816c3d349 100644 --- a/tests/test_ollama_runner_hint.py +++ b/tests/test_ollama_runner_hint.py @@ -34,7 +34,8 @@ def test_hint_still_tells_the_user_how_to_install(): def test_no_runner_echo_line_uses_backticks_in_double_quotes(): # Source-level guard: generated-script echo lines must never carry # backticks inside a double-quoted bash string again. - src = open(os.path.join(ROOT, "routes", "cookbook_routes.py"), encoding="utf-8").read() + with open(os.path.join(ROOT, "routes", "cookbook_routes.py"), encoding="utf-8") as fh: + src = fh.read() offenders = [ line.strip() for line in src.splitlines() diff --git a/tests/test_research_service.py b/tests/test_research_service.py index cc6e57a7df..14cd10bfd4 100644 --- a/tests/test_research_service.py +++ b/tests/test_research_service.py @@ -43,7 +43,13 @@ def _run(coro): - return asyncio.new_event_loop().run_until_complete(coro) + loop = asyncio.new_event_loop() + try: + return loop.run_until_complete( + coro + ) + finally: + loop.close() class _StubHandler: diff --git a/tests/test_research_utils.py b/tests/test_research_utils.py index 52001d06f3..8c8a733fc0 100644 --- a/tests/test_research_utils.py +++ b/tests/test_research_utils.py @@ -1,5 +1,7 @@ """Tests for research_utils.py — thinking block stripping and quality filtering.""" + + from src.research_utils import strip_thinking, is_low_quality diff --git a/tests/test_scheduler_restart_doublefire.py b/tests/test_scheduler_restart_doublefire.py index 9f0c873725..507378e72c 100644 --- a/tests/test_scheduler_restart_doublefire.py +++ b/tests/test_scheduler_restart_doublefire.py @@ -100,6 +100,8 @@ async def _never(): dispatched = [] def _fake_create_task(coro): dispatched.append(coro) + if hasattr(coro, "close"): + coro.close() class _T: def cancel(self): pass return _T() @@ -115,7 +117,10 @@ async def _drive(): # start() also fires the long-lived _loop and _note_pings_loop as tasks # (stubbed to _never here); filter those out so the test only counts # real per-poll task dispatches. - real_dispatches = [c for c in all_dispatched if c.__name__ != "_never"] + real_dispatches = [ + c for c in all_dispatched + if getattr(getattr(c, "cr_code", None), "co_name", None) != "_never" + ] return cd, ScheduledTask, TaskRun, real_dispatches diff --git a/tests/test_tool_index_schema_parity.py b/tests/test_tool_index_schema_parity.py index 77aae50232..1fcb9adf4e 100644 --- a/tests/test_tool_index_schema_parity.py +++ b/tests/test_tool_index_schema_parity.py @@ -27,13 +27,15 @@ def _assigned_value(tree, name): def _schema_tool_names(): - src = open(os.path.join(ROOT, "src", "tool_schemas.py"), encoding="utf-8").read() + with open(os.path.join(ROOT, "src", "tool_schemas.py"), encoding="utf-8") as fh: + src = fh.read() value = _assigned_value(ast.parse(src), "FUNCTION_TOOL_SCHEMAS") return {item["function"]["name"] for item in ast.literal_eval(value)} def _indexed_tool_names(): - src = open(os.path.join(ROOT, "src", "tool_index.py"), encoding="utf-8").read() + with open(os.path.join(ROOT, "src", "tool_index.py"), encoding="utf-8") as fh: + src = fh.read() value = _assigned_value(ast.parse(src), "BUILTIN_TOOL_DESCRIPTIONS") return {ast.literal_eval(key) for key in value.keys} @@ -47,7 +49,8 @@ def test_every_schema_tool_has_an_index_description(): def test_api_call_is_indexed_with_a_real_description(): - src = open(os.path.join(ROOT, "src", "tool_index.py"), encoding="utf-8").read() + with open(os.path.join(ROOT, "src", "tool_index.py"), encoding="utf-8") as fh: + src = fh.read() value = _assigned_value(ast.parse(src), "BUILTIN_TOOL_DESCRIPTIONS") descriptions = { ast.literal_eval(k): ast.literal_eval(v) for k, v in zip(value.keys, value.values) diff --git a/tests/test_upload_handler_atomicity.py b/tests/test_upload_handler_atomicity.py index 73cf27917c..e2dadbb4d7 100644 --- a/tests/test_upload_handler_atomicity.py +++ b/tests/test_upload_handler_atomicity.py @@ -92,7 +92,11 @@ def test_concurrent_inserts_lose_entries(tmp_path): def insert(idx: int) -> None: with handler._index_lock: - current = json.load(open(db_path)) if os.path.exists(db_path) else {} + if os.path.exists(db_path): + with open(db_path, encoding="utf-8") as f: + current = json.load(f) + else: + current = {} current[f"owner:hash_{idx}"] = {"id": f"file_{idx}", "owner": "owner"} handler._atomic_write_json(db_path, current) @@ -246,7 +250,8 @@ def test_partial_write_recovery_via_bak(tmp_path): "Production _atomic_write_json must create a .bak sibling on subsequent writes." ) - full = open(db_path, "rb").read() + with open(db_path, "rb") as f: + full = f.read() truncated_len = max(1, len(full) // 2) with open(db_path, "wb") as f: f.write(full[:truncated_len]) @@ -383,15 +388,18 @@ def test_smoke_info_lookup_after_bak_recovery(tmp_path): "owner_a", ) # Force a .bak by writing a second time. + with open(db_path, encoding="utf-8") as f: + current = json.load(f) handler._atomic_write_json( db_path, - json.load(open(db_path)), + current, ) handler._atomic_write_json(db_path, {"sentinel": True}) assert os.path.exists(db_path + ".bak") # Truncate the live file. - full = open(db_path, "rb").read() + with open(db_path, "rb") as f: + full = f.read() with open(db_path, "wb") as f: f.write(full[: max(1, len(full) // 2)]) diff --git a/tests/test_vault_password_not_in_argv.py b/tests/test_vault_password_not_in_argv.py index f23cddcd8b..8aeb95fe4e 100644 --- a/tests/test_vault_password_not_in_argv.py +++ b/tests/test_vault_password_not_in_argv.py @@ -102,7 +102,8 @@ def test_unlock_handler_feeds_password_on_stdin_not_argv(): def test_tool_vault_unlock_feeds_password_on_stdin_not_argv(): - text = open("src/tools/vault.py", encoding="utf-8").read() + with open("src/tools/vault.py", encoding="utf-8") as fh: + text = fh.read() assert '["unlock", master_password, "--raw"]' not in text assert '_run_bw(["unlock", master_password' not in text diff --git a/tests/test_webhook_dns_rebinding_pin.py b/tests/test_webhook_dns_rebinding_pin.py index 144a19e961..52e37e68b8 100644 --- a/tests/test_webhook_dns_rebinding_pin.py +++ b/tests/test_webhook_dns_rebinding_pin.py @@ -68,6 +68,13 @@ def _serve(handler): return srv, port +def _stop_server(srv): + try: + srv.shutdown() + finally: + srv.server_close() + + def test_pinned_transport_connects_to_pinned_ip(): """A request whose URL host is a throwaway hostname is still delivered to the pinned loopback IP — proving the socket destination comes from the pin, @@ -104,7 +111,7 @@ async def go(): assert resp.status_code == 204 assert hits == ["/hook"] finally: - srv.shutdown() + _stop_server(srv) def test_deliver_pins_to_validated_ip_end_to_end(monkeypatch): @@ -153,4 +160,4 @@ def close(self): pass assert received.get("event") == "webhook.test" assert b'"ok": true' in received["body"] finally: - srv.shutdown() + _stop_server(srv) From ee17d8977202ad093666f7e5a9f909e668828df6 Mon Sep 17 00:00:00 2001 From: Kaiki Date: Sun, 12 Jul 2026 22:08:49 -0400 Subject: [PATCH 07/12] fix(tools): stop batches at user prompts --- src/agent_loop.py | 5 ++++ src/agent_tools/__init__.py | 2 +- tests/test_tool_lifecycle_batch_barriers.py | 27 +++++++++++++++++++++ 3 files changed, 33 insertions(+), 1 deletion(-) create mode 100644 tests/test_tool_lifecycle_batch_barriers.py diff --git a/src/agent_loop.py b/src/agent_loop.py index 6f7dde6054..84e8eb0866 100644 --- a/src/agent_loop.py +++ b/src/agent_loop.py @@ -4403,6 +4403,11 @@ async def _run_tool(): ): _ody_doc_tool_completed = True + # ASK_USER_BATCH_BARRIER: a prompt to the user terminates the + # current parsed tool batch after the prompt has been emitted. + if _pending_ask_user_event: + break + # If budget was hit, stop the loop if budget_hit: break diff --git a/src/agent_tools/__init__.py b/src/agent_tools/__init__.py index 848acf6958..0aad7ccd29 100644 --- a/src/agent_tools/__init__.py +++ b/src/agent_tools/__init__.py @@ -97,7 +97,7 @@ # the dispatcher — silent failure for the whole cookbook # surface. "download_model", "serve_model", - "list_served_models", "stop_served_model", + "list_served_models", "stop_served_model", "tail_serve_output", "list_downloads", "cancel_download", "search_hf_models", "list_cached_models", "list_serve_presets", "serve_preset", "adopt_served_model", diff --git a/tests/test_tool_lifecycle_batch_barriers.py b/tests/test_tool_lifecycle_batch_barriers.py new file mode 100644 index 0000000000..e841d70581 --- /dev/null +++ b/tests/test_tool_lifecycle_batch_barriers.py @@ -0,0 +1,27 @@ +from pathlib import Path + +from src.agent_tools import TOOL_TAGS, parse_tool_blocks + + +def test_tail_serve_output_is_parseable_fenced_tool(): + assert "tail_serve_output" in TOOL_TAGS + + blocks = parse_tool_blocks( + '```' + 'tail_serve_output\n{"session_id": "abc", "tail": 400}\n' + '```' + ) + + assert len(blocks) == 1 + assert blocks[0].tool_type == "tail_serve_output" + assert '"session_id": "abc"' in blocks[0].content + assert '"tail": 400' in blocks[0].content + + +def test_ask_user_branch_sets_same_batch_barrier(): + source = Path("src/agent_loop.py").read_text(encoding="utf-8") + + ask_start = source.index('if "ask_user" in result:') + round_stop = source.index("if _awaiting_user:", ask_start) + branch = source[ask_start:round_stop] + + assert "ASK_USER_BATCH_BARRIER" in branch + assert "break" in branch From 1d28e8b67e6d86af71e6920f25f128c93c0eb8ca Mon Sep 17 00:00:00 2001 From: Kaiki Date: Sun, 12 Jul 2026 22:25:55 -0400 Subject: [PATCH 08/12] fix(research): persist generated report links --- src/agent_loop.py | 1 + tests/test_research_anchor_persistence.py | 32 +++++++++++++++++++++++ 2 files changed, 33 insertions(+) create mode 100644 tests/test_research_anchor_persistence.py diff --git a/src/agent_loop.py b/src/agent_loop.py index 84e8eb0866..716c33879e 100644 --- a/src/agent_loop.py +++ b/src/agent_loop.py @@ -4343,6 +4343,7 @@ async def _run_tool(): _rsid = result.get("research_session_id") if _rsid: _anchor = f"\n\n[Open in Deep Research](#research-{_rsid})\n" + full_response = (full_response.rstrip() + _anchor).strip() yield 'data: ' + json.dumps({"delta": _anchor}) + '\n\n' # Same pattern for notes: when manage_notes creates a note diff --git a/tests/test_research_anchor_persistence.py b/tests/test_research_anchor_persistence.py new file mode 100644 index 0000000000..1b4a96b14e --- /dev/null +++ b/tests/test_research_anchor_persistence.py @@ -0,0 +1,32 @@ +from pathlib import Path + + +def _research_anchor_branch() -> str: + source = Path("src/agent_loop.py").read_text(encoding="utf-8") + start = source.index('result.get("research_session_id")') + end = source.index("# Same pattern for notes:", start) + return source[start:end] + + +def test_research_anchor_is_added_to_persisted_assistant_response(): + """The link shown in SSE must survive session reload via full_response.""" + branch = _research_anchor_branch() + + assert 'json.dumps({"delta": _anchor})' in branch + assert "full_response" in branch + assert "_anchor" in branch + + persist_position = branch.index("full_response") + stream_position = branch.index('json.dumps({"delta": _anchor})') + assert persist_position < stream_position + + +def test_research_anchor_is_not_persisted_twice(): + branch = _research_anchor_branch() + + persistence_statements = [ + line + for line in branch.splitlines() + if "full_response" in line and "_anchor" in line + ] + assert len(persistence_statements) == 1 From 9ab71e620bd47d48ad17e6313a89c7abdaf027b8 Mon Sep 17 00:00:00 2001 From: Kaiki Date: Sun, 12 Jul 2026 22:29:52 -0400 Subject: [PATCH 09/12] fix(mcp): close replaced server connections --- src/mcp_manager.py | 12 +++ .../test_mcp_duplicate_connection_cleanup.py | 80 +++++++++++++++++++ 2 files changed, 92 insertions(+) create mode 100644 tests/test_mcp_duplicate_connection_cleanup.py diff --git a/src/mcp_manager.py b/src/mcp_manager.py index b1d8f8b825..c22d15f6a0 100644 --- a/src/mcp_manager.py +++ b/src/mcp_manager.py @@ -231,6 +231,18 @@ async def connect_server( url: Optional[str] = None, ) -> bool: """Connect to an MCP server via stdio, SSE, or Streamable HTTP transport.""" + # Reusing an ID must release its prior transport stack, session, tools, + # and in-flight connection task before any dictionaries are replaced. + # Otherwise the previous process/session loses its cleanup handle. + if ( + server_id in self._stacks + or server_id in self._sessions + or server_id in self._tools + or server_id in self._connect_tasks + or server_id in self._connections + ): + await self.disconnect_server(server_id) + try: if transport == "stdio": res = await self._connect_stdio(server_id, name, command, args or [], env or {}) diff --git a/tests/test_mcp_duplicate_connection_cleanup.py b/tests/test_mcp_duplicate_connection_cleanup.py new file mode 100644 index 0000000000..c1f9b68be6 --- /dev/null +++ b/tests/test_mcp_duplicate_connection_cleanup.py @@ -0,0 +1,80 @@ +from unittest.mock import AsyncMock + +import pytest + +from src.mcp_manager import McpManager + + +@pytest.mark.asyncio +async def test_connecting_existing_server_id_closes_old_resources(monkeypatch): + manager = McpManager() + old_stack = AsyncMock() + old_session = object() + + manager._stacks["same-id"] = old_stack + manager._sessions["same-id"] = old_session + manager._tools["same-id"] = [{"name": "old_tool"}] + manager._connections["same-id"] = { + "status": "connected", + "name": "Old server", + "transport": "stdio", + } + + async def fake_connect_stdio( + server_id, + name, + command, + args, + env, + ): + assert server_id == "same-id" + assert "same-id" not in manager._stacks + assert "same-id" not in manager._sessions + assert "same-id" not in manager._tools + return True + + monkeypatch.setattr(manager, "_connect_stdio", fake_connect_stdio) + + connected = await manager.connect_server( + server_id="same-id", + name="Replacement server", + transport="stdio", + command="synthetic-command", + args=[], + env={}, + url="", + ) + + assert connected is True + old_stack.aclose.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_first_connection_does_not_call_disconnect(monkeypatch): + manager = McpManager() + disconnect = AsyncMock(wraps=manager.disconnect_server) + monkeypatch.setattr(manager, "disconnect_server", disconnect) + + async def fake_connect_stdio( + server_id, + name, + command, + args, + env, + ): + return True + + monkeypatch.setattr(manager, "_connect_stdio", fake_connect_stdio) + + connected = await manager.connect_server( + server_id="new-id", + name="New server", + transport="stdio", + command="synthetic-command", + args=[], + env={}, + url="", + ) + + assert connected is True + disconnect.assert_not_awaited() From 286957fd5487ffbd391b2960bb61e1027b9b4086 Mon Sep 17 00:00:00 2001 From: Kaiki Date: Sun, 12 Jul 2026 22:31:23 -0400 Subject: [PATCH 10/12] fix(mcp): cancel pending connections on shutdown --- src/mcp_manager.py | 23 +++++++- .../test_mcp_disconnect_all_pending_tasks.py | 59 +++++++++++++++++++ 2 files changed, 79 insertions(+), 3 deletions(-) create mode 100644 tests/test_mcp_disconnect_all_pending_tasks.py diff --git a/src/mcp_manager.py b/src/mcp_manager.py index c22d15f6a0..288da4bc62 100644 --- a/src/mcp_manager.py +++ b/src/mcp_manager.py @@ -5,6 +5,7 @@ Each server exposes tools that are made available to the agent loop. """ +import asyncio import json import logging import os @@ -505,6 +506,16 @@ async def disconnect_server(self, server_id: str): task = self._connect_tasks.pop(server_id, None) if task is not None and not task.done(): task.cancel() + try: + await task + except asyncio.CancelledError: + pass + except Exception as exc: + logger.debug( + "MCP connect task ended during disconnect for %s: %s", + server_id, + sanitize_mcp_error(exc), + ) try: from src.mcp_oauth import clear_auth_url clear_auth_url(server_id) @@ -525,9 +536,15 @@ async def disconnect_server(self, server_id: str): logger.info(f"MCP server disconnected: {server_id}") async def disconnect_all(self): - """Disconnect from all MCP servers.""" - ids = list(self._sessions.keys()) - for sid in ids: + """Disconnect every tracked server, including pending connections.""" + ids = ( + set(self._sessions) + | set(self._stacks) + | set(self._tools) + | set(self._connections) + | set(self._connect_tasks) + ) + for sid in list(ids): await self.disconnect_server(sid) async def connect_all_enabled(self): diff --git a/tests/test_mcp_disconnect_all_pending_tasks.py b/tests/test_mcp_disconnect_all_pending_tasks.py new file mode 100644 index 0000000000..d72b293d63 --- /dev/null +++ b/tests/test_mcp_disconnect_all_pending_tasks.py @@ -0,0 +1,59 @@ +import asyncio + +import pytest + +from src.mcp_manager import McpManager + + +@pytest.mark.asyncio +async def test_disconnect_all_cancels_pending_task_without_session(): + manager = McpManager() + started = asyncio.Event() + cancelled = asyncio.Event() + + async def pending_connect(): + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + task = asyncio.create_task(pending_connect()) + await started.wait() + + manager._connect_tasks["pending-http"] = task + manager._connections["pending-http"] = { + "status": "connecting", + "name": "Pending HTTP", + "transport": "http", + } + + await manager.disconnect_all() + await asyncio.wait_for(cancelled.wait(), timeout=1) + + assert task.cancelled() or task.done() + assert "pending-http" not in manager._connect_tasks + assert "pending-http" not in manager._connections + + +@pytest.mark.asyncio +async def test_disconnect_server_waits_for_cancelled_connect_task(): + manager = McpManager() + cleanup_finished = asyncio.Event() + + async def pending_connect(): + try: + await asyncio.Event().wait() + finally: + await asyncio.sleep(0) + cleanup_finished.set() + + task = asyncio.create_task(pending_connect()) + await asyncio.sleep(0) + + manager._connect_tasks["pending-http"] = task + + await manager.disconnect_server("pending-http") + + assert cleanup_finished.is_set() + assert task.done() From 7238b430a93bebc2c18eaa8e47003f12fe4af070 Mon Sep 17 00:00:00 2001 From: Kaiki Date: Sun, 12 Jul 2026 23:48:13 -0400 Subject: [PATCH 11/12] fix(mcp): clear completed HTTP connection tasks --- src/mcp_manager.py | 12 +++ tests/test_mcp_http_timeout_cleanup.py | 108 +++++++++++++++++++++++++ 2 files changed, 120 insertions(+) create mode 100644 tests/test_mcp_http_timeout_cleanup.py diff --git a/src/mcp_manager.py b/src/mcp_manager.py index 288da4bc62..8b1ecd5eaf 100644 --- a/src/mcp_manager.py +++ b/src/mcp_manager.py @@ -418,8 +418,20 @@ async def _start_http_connect(self, server_id: str, name: str, url: str, wait: f self._connections[server_id] = {"status": "connecting", "name": name, "transport": "http"} task = asyncio.create_task(self._connect_http(server_id, name, url)) self._connect_tasks[server_id] = task + + def _remove_finished_task(finished_task): + # A reconnect may already have installed a newer task under this + # ID. Only remove the task that actually completed. + if self._connect_tasks.get(server_id) is finished_task: + self._connect_tasks.pop(server_id, None) + + task.add_done_callback(_remove_finished_task) done, _ = await asyncio.wait({task}, timeout=wait) if task in done: + # Done callbacks are scheduled by the event loop and may not have + # run before this coroutine resumes, so clear immediate completion + # synchronously as well. + _remove_finished_task(task) try: return task.result() except Exception as e: diff --git a/tests/test_mcp_http_timeout_cleanup.py b/tests/test_mcp_http_timeout_cleanup.py new file mode 100644 index 0000000000..d1ed6b7781 --- /dev/null +++ b/tests/test_mcp_http_timeout_cleanup.py @@ -0,0 +1,108 @@ +import asyncio + +import pytest + +from src.mcp_manager import McpManager + + +@pytest.mark.asyncio +async def test_http_disconnect_cancels_and_awaits_pending_connection(): + manager = McpManager() + started = asyncio.Event() + cleanup_finished = asyncio.Event() + + async def pending_http_connect(): + started.set() + try: + await asyncio.Event().wait() + finally: + await asyncio.sleep(0) + cleanup_finished.set() + + task = asyncio.create_task(pending_http_connect()) + await started.wait() + + manager._connect_tasks["http-id"] = task + manager._connections["http-id"] = { + "status": "connecting", + "name": "HTTP", + "transport": "http", + } + + await manager.disconnect_server("http-id") + + assert cleanup_finished.is_set() + assert task.done() + assert "http-id" not in manager._connect_tasks + assert "http-id" not in manager._connections + + +@pytest.mark.asyncio +async def test_completed_http_task_is_removed_from_tracking(monkeypatch): + manager = McpManager() + + async def completed_connect(server_id, name, url): + return True + + monkeypatch.setattr(manager, "_connect_http", completed_connect) + + result = await manager._start_http_connect( + "http-id", + "HTTP server", + "https://example.invalid/mcp", + ) + + assert result is True + assert "http-id" not in manager._connect_tasks + + +@pytest.mark.asyncio +async def test_failed_http_task_is_removed_from_tracking(monkeypatch): + manager = McpManager() + + async def failed_connect(server_id, name, url): + raise RuntimeError("synthetic HTTP connection failure") + + monkeypatch.setattr(manager, "_connect_http", failed_connect) + + result = await manager._start_http_connect( + "http-id", + "HTTP server", + "https://example.invalid/mcp", + ) + + # Immediate failures are converted into status=False by the public helper. + assert result is False + assert manager._connections["http-id"]["status"] == "error" + assert "http-id" not in manager._connect_tasks + + +@pytest.mark.asyncio +async def test_background_http_task_remains_tracked_until_completion( + monkeypatch, +): + manager = McpManager() + release = asyncio.Event() + + async def delayed_connect(server_id, name, url): + await release.wait() + return True + + monkeypatch.setattr(manager, "_connect_http", delayed_connect) + + result = await manager._start_http_connect( + "http-id", + "HTTP server", + "https://example.invalid/mcp", + wait=0, + ) + + assert result is False + task = manager._connect_tasks["http-id"] + assert not task.done() + + release.set() + await task + await asyncio.sleep(0) + + assert "http-id" not in manager._connect_tasks From 504ce42a346d4551a209566e174d128366e46147 Mon Sep 17 00:00:00 2001 From: Kaiki Date: Mon, 13 Jul 2026 00:09:28 -0400 Subject: [PATCH 12/12] fix(mcp): validate route input payloads --- routes/mcp_routes.py | 75 +++++++++++++++++++--- tests/test_mcp_oauth_error_redaction.py | 27 ++++++++ tests/test_mcp_route_input_validation.py | 80 ++++++++++++++++++++++++ 3 files changed, 172 insertions(+), 10 deletions(-) create mode 100644 tests/test_mcp_oauth_error_redaction.py create mode 100644 tests/test_mcp_route_input_validation.py diff --git a/routes/mcp_routes.py b/routes/mcp_routes.py index a0ade88b6d..9085e75f1f 100644 --- a/routes/mcp_routes.py +++ b/routes/mcp_routes.py @@ -21,6 +21,67 @@ router = APIRouter(prefix="/api/mcp", tags=["mcp"]) + +_SUPPORTED_MCP_TRANSPORTS = frozenset({"stdio", "sse", "http"}) + + +def _validate_mcp_transport(raw_transport) -> str: + transport = str(raw_transport or "").lower() + if transport not in _SUPPORTED_MCP_TRANSPORTS: + raise HTTPException( + status_code=400, + detail=( + "Unsupported MCP transport; expected one of: " + "stdio, sse, http" + ), + ) + return transport + + +def _parse_mcp_args(raw_args) -> list[str]: + try: + parsed = json.loads(raw_args) + except (TypeError, ValueError, json.JSONDecodeError) as exc: + raise HTTPException( + status_code=400, + detail="Invalid MCP args: expected a JSON string list", + ) from exc + + if ( + not isinstance(parsed, list) + or not all(isinstance(item, str) for item in parsed) + ): + raise HTTPException( + status_code=400, + detail="Invalid MCP args: expected a JSON string list", + ) + + return parsed + + +def _parse_mcp_env(raw_env) -> dict[str, str]: + try: + parsed = json.loads(raw_env) + except (TypeError, ValueError, json.JSONDecodeError) as exc: + raise HTTPException( + status_code=400, + detail="Invalid MCP env: expected a JSON string mapping", + ) from exc + + if ( + not isinstance(parsed, dict) + or not all( + isinstance(key, str) and isinstance(value, str) + for key, value in parsed.items() + ) + ): + raise HTTPException( + status_code=400, + detail="Invalid MCP env: expected a JSON string mapping", + ) + + return parsed + def _mcp_oauth_base_dir() -> Path: """Directory that may contain OAuth files managed by Odysseus.""" return Path(MCP_OAUTH_DIR).resolve(strict=False) @@ -174,6 +235,8 @@ async def add_server( server_id = str(uuid.uuid4())[:8] # Validate + transport = _validate_mcp_transport(transport) + if transport == "stdio" and not command: raise HTTPException(400, "command is required for stdio transport") if transport == "sse" and not url: @@ -182,16 +245,8 @@ async def add_server( raise HTTPException(400, "url is required for HTTP transport") # Parse JSON fields - try: - parsed_args = json.loads(args) if args else [] - except json.JSONDecodeError: - parsed_args = [] - try: - parsed_env = json.loads(env) if env else {} - except json.JSONDecodeError: - parsed_env = {} - if not isinstance(parsed_env, dict): - parsed_env = {} + parsed_args = _parse_mcp_args(args if args else "[]") + parsed_env = _parse_mcp_env(env if env else "{}") # Parse OAuth config parsed_oauth_config = None diff --git a/tests/test_mcp_oauth_error_redaction.py b/tests/test_mcp_oauth_error_redaction.py new file mode 100644 index 0000000000..c14eb4863e --- /dev/null +++ b/tests/test_mcp_oauth_error_redaction.py @@ -0,0 +1,27 @@ +from src.mcp_manager import sanitize_mcp_error + + +def test_oauth_error_redacts_query_and_body_credentials(): + error = RuntimeError( + "OAuth failed: " + "https://example.invalid/callback?" + "client_secret=query-secret&code=auth-code " + "client_secret=body-secret " + "access_token=access-secret " + "refresh_token=refresh-secret" + ) + + sanitized = sanitize_mcp_error(error) + + for secret in ( + "query-secret", + "auth-code", + "body-secret", + "access-secret", + "refresh-secret", + ): + assert secret not in sanitized + + assert "client_secret" in sanitized + assert "access_token" in sanitized + assert "refresh_token" in sanitized diff --git a/tests/test_mcp_route_input_validation.py b/tests/test_mcp_route_input_validation.py new file mode 100644 index 0000000000..e956ed5723 --- /dev/null +++ b/tests/test_mcp_route_input_validation.py @@ -0,0 +1,80 @@ +import pytest +from fastapi import HTTPException + +from routes import mcp_routes + + +@pytest.mark.parametrize( + "transport", + ["", "bogus", "websocket", "streamable_http", "stdio "], +) +def test_validate_transport_rejects_unsupported_values(transport): + with pytest.raises(HTTPException) as exc: + mcp_routes._validate_mcp_transport(transport) + + assert exc.value.status_code == 400 + assert "transport" in str(exc.value.detail).lower() + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("stdio", "stdio"), + ("STDIO", "stdio"), + ("sse", "sse"), + ("http", "http"), + ], +) +def test_validate_transport_accepts_supported_values(raw, expected): + assert mcp_routes._validate_mcp_transport(raw) == expected + + +@pytest.mark.parametrize("raw", ["{", "{}", '"value"', "1", "null"]) +def test_parse_args_rejects_invalid_or_wrong_type(raw): + with pytest.raises(HTTPException) as exc: + mcp_routes._parse_mcp_args(raw) + + assert exc.value.status_code == 400 + assert "args" in str(exc.value.detail).lower() + + +@pytest.mark.parametrize( + "raw", + ['["ok", 3]', '["ok", null]', '["ok", true]'], +) +def test_parse_args_requires_string_items(raw): + with pytest.raises(HTTPException): + mcp_routes._parse_mcp_args(raw) + + +def test_parse_args_accepts_string_list(): + assert mcp_routes._parse_mcp_args( + '["--flag", "value"]' + ) == ["--flag", "value"] + + +@pytest.mark.parametrize("raw", ["{", "[]", '"value"', "1", "null"]) +def test_parse_env_rejects_invalid_or_wrong_type(raw): + with pytest.raises(HTTPException) as exc: + mcp_routes._parse_mcp_env(raw) + + assert exc.value.status_code == 400 + assert "env" in str(exc.value.detail).lower() + + +@pytest.mark.parametrize( + "raw", + ['{"TOKEN": 3}', '{"TOKEN": null}', '{"TOKEN": true}'], +) +def test_parse_env_requires_string_values(raw): + with pytest.raises(HTTPException): + mcp_routes._parse_mcp_env(raw) + + +def test_parse_env_accepts_string_mapping(): + assert mcp_routes._parse_mcp_env( + '{"TOKEN": "value", "EMPTY": ""}' + ) == { + "TOKEN": "value", + "EMPTY": "", + }