diff --git a/integrations/claude-code/adrian_cc/agent.py b/integrations/claude-code/adrian_cc/agent.py index b87a3e0..5f66274 100644 --- a/integrations/claude-code/adrian_cc/agent.py +++ b/integrations/claude-code/adrian_cc/agent.py @@ -24,6 +24,7 @@ import json import os import ssl +import subprocess import sys import tempfile import time @@ -657,6 +658,8 @@ async def _ws_send_event( "policy_m4": policy.policy_m4, } result["source_ack"] = sf.login_ack.source + blocked_list = list(sf.login_ack.blocked_mcp_servers) + _mutate_state(lambda s: s.__setitem__("blocked_mcp_servers", blocked_list)) # --- Send event --- batch = pb.PairedEventBatch(events=[event]) @@ -896,9 +899,327 @@ def _verdict_action(result: dict[str, Any]) -> str: # --------------------------------------------------------------------------- -def _handle_start(_hook_data: dict[str, Any]) -> None: - """SessionStart: reset state, inject governance context.""" +def _probe_mcp_server(cfg: dict) -> dict: + """Probe a stdio MCP server via the initialize handshake. + + Returns dict with keys: version, protocol_version, server_info_name, tools_json. + All values default to "" on failure. + """ + result = {"version": "", "protocol_version": "", "server_info_name": "", "tools_json": ""} + if "command" not in cfg: + return result + cmd = cfg.get("command", "") + args = cfg.get("args", []) + if not isinstance(args, list): + args = [] + env = {**os.environ, **(cfg.get("env", {}) or {})} + + init_request = json.dumps({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "adrian-probe", "version": "0.1.0"}, + }, + }) + tools_request = json.dumps({ + "jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}, + }) + + try: + proc = subprocess.Popen( + [cmd] + args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, env=env, text=True, + ) + stdout, _ = proc.communicate( + input=f"Content-Length: {len(init_request)}\r\n\r\n{init_request}" + f"Content-Length: {len(tools_request)}\r\n\r\n{tools_request}", + timeout=10, + ) + + for line in stdout.splitlines(): + line = line.strip() + if not line.startswith("{"): + continue + try: + msg = json.loads(line) + except json.JSONDecodeError: + continue + if not isinstance(msg, dict) or "result" not in msg: + continue + res = msg["result"] + if msg.get("id") == 1: + si = res.get("serverInfo", {}) + result["version"] = si.get("version", "") + result["server_info_name"] = si.get("name", "") + result["protocol_version"] = res.get("protocolVersion", "") + elif msg.get("id") == 2: + tools = res.get("tools", []) + result["tools_json"] = json.dumps([t.get("name", "") for t in tools]) + except Exception: + pass + + return result + + +def _resolve_mcp_version(cfg: dict) -> str: + """Fallback version resolution via npm/pip when probe didn't get a version.""" + if "command" not in cfg: + return "" + cmd = cfg.get("command", "") + args = cfg.get("args", []) + if not isinstance(args, list): + args = [] + + if cmd in ("npx", "npx.cmd") and args: + pkg_name = _extract_npx_package(args) + if pkg_name: + return _npm_version(pkg_name) + + if cmd in ("uvx", "pipx"): + pkg_name = _extract_first_positional(args) + if pkg_name: + return _pip_version(pkg_name) + + if cmd in ("python", "python3") and args: + if "-m" in args: + idx = args.index("-m") + if idx + 1 < len(args): + return _pip_version(args[idx + 1].replace(".", "-")) + + return _pip_version(cmd) if not cmd.startswith("/") else "" + + +def _extract_npx_package(args: list) -> str: + """Extract the npm package name from npx args.""" + skip_next = False + for arg in args: + if skip_next: + skip_next = False + continue + if arg in ("-y", "--yes", "-q", "--quiet"): + continue + if arg.startswith("-p") or arg == "--package": + skip_next = True + continue + if arg.startswith("-"): + continue + return arg + return "" + + +def _extract_first_positional(args: list) -> str: + """Extract the first non-flag argument.""" + for arg in args: + if not arg.startswith("-"): + return arg + return "" + + +def _npm_version(pkg: str) -> str: + """Get version from npm registry.""" + try: + result = subprocess.run( + ["npm", "view", pkg, "version"], + capture_output=True, text=True, timeout=5, + ) + if result.returncode == 0 and result.stdout.strip(): + return result.stdout.strip() + except Exception: + pass + return "" + + +def _pip_version(pkg: str) -> str: + """Get version from pip show.""" + try: + result = subprocess.run( + ["pip", "show", pkg], + capture_output=True, text=True, timeout=5, + ) + if result.returncode == 0: + for line in result.stdout.splitlines(): + if line.startswith("Version:"): + return line.split(":", 1)[1].strip() + except Exception: + pass + return "" + + +def _discover_cc_mcp_servers(cwd: str = "") -> list[pb.McpServer]: + """Read MCP server config from Claude Code config files. + + Sources (merged, later wins on name collision): + 1. ~/.claude/settings.json → mcpServers + 2. $CWD/.claude/settings.json → mcpServers + 3. $CWD/.mcp.json → mcpServers + """ + servers: dict[str, pb.McpServer] = {} + + candidates = [ + Path.home() / ".claude" / "settings.json", + ] + if cwd: + candidates.append(Path(cwd) / ".claude" / "settings.json") + candidates.append(Path(cwd) / ".mcp.json") + + for path in candidates: + try: + if not path.is_file(): + continue + data = json.loads(path.read_text()) + mcp_servers = data.get("mcpServers", {}) + if not isinstance(mcp_servers, dict): + continue + for name, cfg in mcp_servers.items(): + if not isinstance(cfg, dict): + continue + transport = "unknown" + endpoint = "" + if "command" in cfg: + transport = "stdio" + cmd = cfg.get("command", "") + args = cfg.get("args", []) + endpoint = " ".join([cmd] + (args if isinstance(args, list) else [])) + elif "url" in cfg: + transport = "sse" + endpoint = cfg["url"] + probe = _probe_mcp_server(cfg) + version = probe["version"] or _resolve_mcp_version(cfg) + servers[name] = pb.McpServer( + name=name, transport=transport, endpoint=endpoint, + version=version, + protocol_version=probe["protocol_version"], + server_info_name=probe["server_info_name"], + tools_json=probe["tools_json"], + ) + except Exception: + continue + + return list(servers.values()) + + +def _discover_cc_plugins() -> list[pb.InstalledPlugin]: + """Read installed plugins from Claude Code config. + + Sources: + 1. ~/.claude/plugins/installed_plugins.json → list of installed plugins + 2. ~/.claude/settings.json → enabledPlugins (list of enabled plugin names) + """ + plugins: dict[str, pb.InstalledPlugin] = {} + + # Installed plugins file + plugins_file = Path.home() / ".claude" / "plugins" / "installed_plugins.json" + try: + if plugins_file.is_file(): + data = json.loads(plugins_file.read_text()) + if isinstance(data, list): + for entry in data: + if not isinstance(entry, dict): + continue + name = entry.get("name", "") or entry.get("package_name", "") + if not name: + continue + plugins[name] = pb.InstalledPlugin( + name=name, + enabled=entry.get("enabled", True), + version=entry.get("version", ""), + marketplace=entry.get("marketplace", entry.get("registry", "")), + ) + except Exception: + pass + + # enabledPlugins from settings — dict {"name@marketplace": true/false} + # or list ["name@marketplace", ...]. Creates plugin entries if not already + # found in installed_plugins.json. + settings_file = Path.home() / ".claude" / "settings.json" + try: + if settings_file.is_file(): + data = json.loads(settings_file.read_text()) + enabled_raw = data.get("enabledPlugins", {}) + if isinstance(enabled_raw, dict): + for full_name, is_enabled in enabled_raw.items(): + parts = full_name.rsplit("@", 1) + short_name = parts[0] if parts else full_name + marketplace = parts[1] if len(parts) > 1 else "" + if full_name in plugins: + plugins[full_name].enabled = bool(is_enabled) + else: + plugins[full_name] = pb.InstalledPlugin( + name=short_name, + enabled=bool(is_enabled), + version="", + marketplace=marketplace, + ) + elif isinstance(enabled_raw, list): + enabled_set = set(enabled_raw) + for name, plugin in plugins.items(): + plugin.enabled = name in enabled_set + except Exception: + pass + + return list(plugins.values()) + + +def _send_inventory_sync( + session_id: str, + mcp_servers: list[pb.McpServer], + plugins: list[pb.InstalledPlugin], +) -> None: + """Send MCP and plugin inventory frames via a fire-and-forget WS connection.""" + import asyncio as _asyncio + + async def _send() -> None: + headers: dict[str, str] = {} + if ADRIAN_API_KEY: + headers["Authorization"] = f"Bearer {ADRIAN_API_KEY}" + try: + async with websockets.connect( + ADRIAN_WS_URL, + additional_headers=headers, + ssl=_ws_ssl_context(), + open_timeout=5, + close_timeout=3, + ) as ws: + login = pb.SessionLogin( + session_id=session_id, + schema_version=2, + source="claude-code", + ) + login.llm_stack.provider = "anthropic" + login.llm_stack.model = "claude-code" + await ws.send(pb.ClientFrame(login=login).SerializeToString()) + raw = await _asyncio.wait_for(ws.recv(), timeout=5) + sf = pb.ServerFrame() + sf.ParseFromString(raw if isinstance(raw, bytes) else raw.encode()) + if mcp_servers: + inv = pb.McpInventory(servers=mcp_servers) + await ws.send(pb.ClientFrame(mcp_inventory=inv).SerializeToString()) + if plugins: + pinv = pb.PluginInventory(plugins=plugins) + await ws.send(pb.ClientFrame(plugin_inventory=pinv).SerializeToString()) + except Exception as exc: + _log(f"inventory send failed: {exc}") + + try: + _asyncio.run(_send()) + except Exception: + pass + + +def _handle_start(hook_data: dict[str, Any]) -> None: + """SessionStart: reset state, discover MCP/plugins, inject governance context.""" _reset_state() + + session_id = hook_data.get("session_id", "") + cwd = hook_data.get("cwd", "") + + if session_id and ADRIAN_API_KEY: + mcp_servers = _discover_cc_mcp_servers(cwd) + plugins = _discover_cc_plugins() + if mcp_servers or plugins: + _send_inventory_sync(session_id, mcp_servers, plugins) + _exit_json( { "hookSpecificOutput": { @@ -926,6 +1247,32 @@ def _handle_pre(hook_data: dict[str, Any]) -> None: if cc_agent_id and cc_agent_id != "claude-code": delegated_prompt = _subagent_delegated_prompt(transcript_path, cc_agent_id) + # Passive MCP discovery: infer MCP server from mcp____ pattern. + if tool_name.startswith("mcp__"): + parts = tool_name.split("__", 2) + if len(parts) >= 2: + server_name = parts[1] + _mutate_state(lambda s: s.setdefault("discovered_mcp", {}).__setitem__(server_name, True)) + + blocked_set = set(_load_state().get("blocked_mcp_servers", [])) + if server_name in blocked_set: + blocked_event = _build_event( + hook_data, + _load_state(), + output=f"[Blocked by MCP policy: {server_name}]", + delegated_prompt=delegated_prompt, + ) + _send_event_sync(blocked_event, session_id, wait_for_verdict=False) + _exit_json( + { + "hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": "deny", + "permissionDecisionReason": f"Adrian: MCP server '{server_name}' is blocked", + } + } + ) + # Snapshot for building the event (before the delegation push, so an Agent # tool call stays attributed to the parent). state = _load_state() diff --git a/integrations/claude-code/adrian_cc/proto/event_pb2.py b/integrations/claude-code/adrian_cc/proto/event_pb2.py index d03816d..0fef4c3 100644 --- a/integrations/claude-code/adrian_cc/proto/event_pb2.py +++ b/integrations/claude-code/adrian_cc/proto/event_pb2.py @@ -1,90 +1,102 @@ # -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE # source: event.proto -# Protobuf Python Version: 4.25.3 +# Protobuf Python Version: 7.35.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version from google.protobuf import symbol_database as _symbol_database from google.protobuf.internal import builder as _builder +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 7, + 35, + 1, + '', + 'event.proto' +) # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() -from .buf.validate import validate_pb2 as buf_dot_validate_dot_validate__pb2 +from adrian_cc.proto.buf.validate import validate_pb2 as buf_dot_validate_dot_validate__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0b\x65vent.proto\x12\x12\x61\x64rian.core_api.v1\x1a\x1b\x62uf/validate/validate.proto\";\n\x0b\x43hatMessage\x12\x12\n\x04role\x18\x01 \x01(\tR\x04role\x12\x18\n\x07\x63ontent\x18\x02 \x01(\tR\x07\x63ontent\"K\n\x08ToolCall\x12\x1b\n\x04name\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01R\x04name\x12\x12\n\x04\x61rgs\x18\x02 \x01(\tR\x04\x61rgs\x12\x0e\n\x02id\x18\x03 \x01(\tR\x02id\"\x9c\x01\n\nTokenUsage\x12,\n\rprompt_tokens\x18\x01 \x01(\x05\x42\x07\xbaH\x04\x1a\x02(\x00R\x0cpromptTokens\x12\x34\n\x11\x63ompletion_tokens\x18\x02 \x01(\x05\x42\x07\xbaH\x04\x1a\x02(\x00R\x10\x63ompletionTokens\x12*\n\x0ctotal_tokens\x18\x03 \x01(\x05\x42\x07\xbaH\x04\x1a\x02(\x00R\x0btotalTokens\"y\n\x0c\x41gentContext\x12\x19\n\x08\x61gent_id\x18\x01 \x01(\tR\x07\x61gentId\x12#\n\rsystem_prompt\x18\x02 \x01(\tR\x0csystemPrompt\x12)\n\x10user_instruction\x18\x03 \x01(\tR\x0fuserInstruction\"\xeb\x01\n\x0bLlmPairData\x12\x14\n\x05model\x18\x01 \x01(\tR\x05model\x12;\n\x08messages\x18\x02 \x03(\x0b\x32\x1f.adrian.core_api.v1.ChatMessageR\x08messages\x12\x16\n\x06output\x18\x03 \x01(\tR\x06output\x12;\n\ntool_calls\x18\x04 \x03(\x0b\x32\x1c.adrian.core_api.v1.ToolCallR\ttoolCalls\x12\x34\n\x05usage\x18\x05 \x01(\x0b\x32\x1e.adrian.core_api.v1.TokenUsageR\x05usage\"\x84\x01\n\x0cToolPairData\x12$\n\ttool_name\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01R\x08toolName\x12 \n\x0ctool_call_id\x18\x02 \x01(\tR\ntoolCallId\x12\x14\n\x05input\x18\x03 \x01(\tR\x05input\x12\x16\n\x06output\x18\x04 \x01(\tR\x06output\"\xe3\x04\n\x0bPairedEvent\x12\"\n\x08\x65vent_id\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01R\x07\x65ventId\x12#\n\rinvocation_id\x18\x02 \x01(\tR\x0cinvocationId\x12&\n\nsession_id\x18\x03 \x01(\tB\x07\xbaH\x04r\x02\x10\x01R\tsessionId\x12\x15\n\x06run_id\x18\x04 \x01(\tR\x05runId\x12\"\n\rparent_run_id\x18\x05 \x01(\tR\x0bparentRunId\x12\x1c\n\ttimestamp\x18\x06 \x01(\tR\ttimestamp\x12\x41\n\tpair_type\x18\x07 \x01(\x0e\x32\x1c.adrian.core_api.v1.PairTypeB\x06\xbaH\x03\xc8\x01\x01R\x08pairType\x12\x36\n\x05\x61gent\x18\x08 \x01(\x0b\x32 .adrian.core_api.v1.AgentContextR\x05\x61gent\x12\x38\n\x06parent\x18\t \x01(\x0b\x32 .adrian.core_api.v1.AgentContextR\x06parent\x12\x33\n\x03llm\x18\n \x01(\x0b\x32\x1f.adrian.core_api.v1.LlmPairDataH\x00R\x03llm\x12\x36\n\x04tool\x18\x0b \x01(\x0b\x32 .adrian.core_api.v1.ToolPairDataH\x00R\x04tool\x12#\n\rmetadata_json\x18\x14 \x01(\x0cR\x0cmetadataJson\x12#\n\rconnection_id\x18\x0c \x01(\tR\x0c\x63onnectionId\x12\x16\n\x06source\x18\x15 \x01(\tR\x06sourceB\x06\n\x04\x64\x61ta\"K\n\x10PairedEventBatch\x12\x37\n\x06\x65vents\x18\x01 \x03(\x0b\x32\x1f.adrian.core_api.v1.PairedEventR\x06\x65vents\"b\n\tMcpServer\x12\x1b\n\x04name\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01R\x04name\x12\x1c\n\ttransport\x18\x02 \x01(\tR\ttransport\x12\x1a\n\x08\x65ndpoint\x18\x03 \x01(\tR\x08\x65ndpoint\"G\n\x0cMcpInventory\x12\x37\n\x07servers\x18\x01 \x03(\x0b\x32\x1d.adrian.core_api.v1.McpServerR\x07servers\"<\n\x08LLMStack\x12\x1a\n\x08provider\x18\x01 \x01(\tR\x08provider\x12\x14\n\x05model\x18\x02 \x01(\tR\x05model\"\xe7\x01\n\x0cSessionLogin\x12&\n\nsession_id\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01R\tsessionId\x12\x39\n\tllm_stack\x18\x02 \x01(\x0b\x32\x1c.adrian.core_api.v1.LLMStackR\x08llmStack\x12%\n\x0eschema_version\x18\x04 \x01(\rR\rschemaVersion\x12\x16\n\x06source\x18\x05 \x01(\tR\x06source\x12#\n\rconnection_id\x18\x06 \x01(\tR\x0c\x63onnectionIdJ\x04\x08\x03\x10\x04R\nblock_mode\"\xf1\x01\n\x0b\x43lientFrame\x12\x38\n\x05login\x18\x01 \x01(\x0b\x32 .adrian.core_api.v1.SessionLoginH\x00R\x05login\x12I\n\x0cpaired_batch\x18\x03 \x01(\x0b\x32$.adrian.core_api.v1.PairedEventBatchH\x00R\x0bpairedBatch\x12G\n\rmcp_inventory\x18\x04 \x01(\x0b\x32 .adrian.core_api.v1.McpInventoryH\x00R\x0cmcpInventoryB\x07\n\x05\x66rameJ\x04\x08\x02\x10\x03R\x05\x62\x61tch\"\xb2\x01\n\x0ePolicySnapshot\x12,\n\x04mode\x18\x01 \x01(\x0e\x32\x18.adrian.core_api.v1.ModeR\x04mode\x12\x1b\n\tpolicy_m0\x18\x02 \x01(\x08R\x08policyM0\x12\x1b\n\tpolicy_m2\x18\x03 \x01(\x08R\x08policyM2\x12\x1b\n\tpolicy_m3\x18\x04 \x01(\x08R\x08policyM3\x12\x1b\n\tpolicy_m4\x18\x05 \x01(\x08R\x08policyM4\"=\n\x0cHitlResponse\x12-\n\x12\x63ontinue_execution\x18\x01 \x01(\x08R\x11\x63ontinueExecution\"^\n\x08LoginAck\x12:\n\x06policy\x18\x01 \x01(\x0b\x32\".adrian.core_api.v1.PolicySnapshotR\x06policy\x12\x16\n\x06source\x18\x02 \x01(\tR\x06source\"\x8c\x01\n\x0bServerFrame\x12;\n\tlogin_ack\x18\x01 \x01(\x0b\x32\x1c.adrian.core_api.v1.LoginAckH\x00R\x08loginAck\x12\x37\n\x07verdict\x18\x02 \x01(\x0b\x32\x1b.adrian.core_api.v1.VerdictH\x00R\x07verdictB\x07\n\x05\x66rame\"\x88\x02\n\x07Verdict\x12\"\n\x08\x65vent_id\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01R\x07\x65ventId\x12&\n\nsession_id\x18\x02 \x01(\tB\x07\xbaH\x04r\x02\x10\x01R\tsessionId\x12\x19\n\x08mad_code\x18\x04 \x01(\tR\x07madCode\x12:\n\x06policy\x18\x06 \x01(\x0b\x32\".adrian.core_api.v1.PolicySnapshotR\x06policy\x12\x34\n\x04hitl\x18\x07 \x01(\x0b\x32 .adrian.core_api.v1.HitlResponseR\x04hitlJ\x04\x08\x03\x10\x04J\x04\x08\x05\x10\x06R\x0e\x63lassificationR\x08\x65scalate*L\n\x08PairType\x12\x19\n\x15PAIR_TYPE_UNSPECIFIED\x10\x00\x12\x11\n\rPAIR_TYPE_LLM\x10\x01\x12\x12\n\x0ePAIR_TYPE_TOOL\x10\x02*K\n\x04Mode\x12\x14\n\x10MODE_UNSPECIFIED\x10\x00\x12\x0e\n\nMODE_ALERT\x10\x01\x12\r\n\tMODE_HITL\x10\x02\x12\x0e\n\nMODE_BLOCK\x10\x03\x42\x31Z/github.com/secure-agentics/adrian/server/pkg/pbb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0b\x65vent.proto\x12\x12\x61\x64rian.core_api.v1\x1a\x1b\x62uf/validate/validate.proto\",\n\x0b\x43hatMessage\x12\x0c\n\x04role\x18\x01 \x01(\t\x12\x0f\n\x07\x63ontent\x18\x02 \x01(\t\";\n\x08ToolCall\x12\x15\n\x04name\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x0c\n\x04\x61rgs\x18\x02 \x01(\t\x12\n\n\x02id\x18\x03 \x01(\t\"o\n\nTokenUsage\x12\x1e\n\rprompt_tokens\x18\x01 \x01(\x05\x42\x07\xbaH\x04\x1a\x02(\x00\x12\"\n\x11\x63ompletion_tokens\x18\x02 \x01(\x05\x42\x07\xbaH\x04\x1a\x02(\x00\x12\x1d\n\x0ctotal_tokens\x18\x03 \x01(\x05\x42\x07\xbaH\x04\x1a\x02(\x00\"Q\n\x0c\x41gentContext\x12\x10\n\x08\x61gent_id\x18\x01 \x01(\t\x12\x15\n\rsystem_prompt\x18\x02 \x01(\t\x12\x18\n\x10user_instruction\x18\x03 \x01(\t\"\xc0\x01\n\x0bLlmPairData\x12\r\n\x05model\x18\x01 \x01(\t\x12\x31\n\x08messages\x18\x02 \x03(\x0b\x32\x1f.adrian.core_api.v1.ChatMessage\x12\x0e\n\x06output\x18\x03 \x01(\t\x12\x30\n\ntool_calls\x18\x04 \x03(\x0b\x32\x1c.adrian.core_api.v1.ToolCall\x12-\n\x05usage\x18\x05 \x01(\x0b\x32\x1e.adrian.core_api.v1.TokenUsage\"_\n\x0cToolPairData\x12\x1a\n\ttool_name\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x14\n\x0ctool_call_id\x18\x02 \x01(\t\x12\r\n\x05input\x18\x03 \x01(\t\x12\x0e\n\x06output\x18\x04 \x01(\t\"\xda\x03\n\x0bPairedEvent\x12\x19\n\x08\x65vent_id\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x15\n\rinvocation_id\x18\x02 \x01(\t\x12\x1b\n\nsession_id\x18\x03 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x0e\n\x06run_id\x18\x04 \x01(\t\x12\x15\n\rparent_run_id\x18\x05 \x01(\t\x12\x11\n\ttimestamp\x18\x06 \x01(\t\x12\x37\n\tpair_type\x18\x07 \x01(\x0e\x32\x1c.adrian.core_api.v1.PairTypeB\x06\xbaH\x03\xc8\x01\x01\x12/\n\x05\x61gent\x18\x08 \x01(\x0b\x32 .adrian.core_api.v1.AgentContext\x12\x30\n\x06parent\x18\t \x01(\x0b\x32 .adrian.core_api.v1.AgentContext\x12.\n\x03llm\x18\n \x01(\x0b\x32\x1f.adrian.core_api.v1.LlmPairDataH\x00\x12\x30\n\x04tool\x18\x0b \x01(\x0b\x32 .adrian.core_api.v1.ToolPairDataH\x00\x12\x15\n\rmetadata_json\x18\x14 \x01(\x0c\x12\x15\n\rconnection_id\x18\x0c \x01(\t\x12\x0e\n\x06source\x18\x15 \x01(\tB\x06\n\x04\x64\x61ta\"C\n\x10PairedEventBatch\x12/\n\x06\x65vents\x18\x01 \x03(\x0b\x32\x1f.adrian.core_api.v1.PairedEvent\"G\n\tMcpServer\x12\x15\n\x04name\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x11\n\ttransport\x18\x02 \x01(\t\x12\x10\n\x08\x65ndpoint\x18\x03 \x01(\t\">\n\x0cMcpInventory\x12.\n\x07servers\x18\x01 \x03(\x0b\x32\x1d.adrian.core_api.v1.McpServer\"+\n\x08LLMStack\x12\x10\n\x08provider\x18\x01 \x01(\t\x12\r\n\x05model\x18\x02 \x01(\t\"\xad\x01\n\x0cSessionLogin\x12\x1b\n\nsession_id\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12/\n\tllm_stack\x18\x02 \x01(\x0b\x32\x1c.adrian.core_api.v1.LLMStack\x12\x16\n\x0eschema_version\x18\x04 \x01(\r\x12\x0e\n\x06source\x18\x05 \x01(\t\x12\x15\n\rconnection_id\x18\x06 \x01(\tJ\x04\x08\x03\x10\x04R\nblock_mode\"\xcf\x01\n\x0b\x43lientFrame\x12\x31\n\x05login\x18\x01 \x01(\x0b\x32 .adrian.core_api.v1.SessionLoginH\x00\x12<\n\x0cpaired_batch\x18\x03 \x01(\x0b\x32$.adrian.core_api.v1.PairedEventBatchH\x00\x12\x39\n\rmcp_inventory\x18\x04 \x01(\x0b\x32 .adrian.core_api.v1.McpInventoryH\x00\x42\x07\n\x05\x66rameJ\x04\x08\x02\x10\x03R\x05\x62\x61tch\"\x84\x01\n\x0ePolicySnapshot\x12&\n\x04mode\x18\x01 \x01(\x0e\x32\x18.adrian.core_api.v1.Mode\x12\x11\n\tpolicy_m0\x18\x02 \x01(\x08\x12\x11\n\tpolicy_m2\x18\x03 \x01(\x08\x12\x11\n\tpolicy_m3\x18\x04 \x01(\x08\x12\x11\n\tpolicy_m4\x18\x05 \x01(\x08\"*\n\x0cHitlResponse\x12\x1a\n\x12\x63ontinue_execution\x18\x01 \x01(\x08\"k\n\x08LoginAck\x12\x32\n\x06policy\x18\x01 \x01(\x0b\x32\".adrian.core_api.v1.PolicySnapshot\x12\x0e\n\x06source\x18\x02 \x01(\t\x12\x1b\n\x13\x62locked_mcp_servers\x18\x03 \x03(\t\"-\n\x0eMcpBlockUpdate\x12\x1b\n\x13\x62locked_mcp_servers\x18\x01 \x03(\t\"\xb9\x01\n\x0bServerFrame\x12\x31\n\tlogin_ack\x18\x01 \x01(\x0b\x32\x1c.adrian.core_api.v1.LoginAckH\x00\x12.\n\x07verdict\x18\x02 \x01(\x0b\x32\x1b.adrian.core_api.v1.VerdictH\x00\x12>\n\x10mcp_block_update\x18\x03 \x01(\x0b\x32\".adrian.core_api.v1.McpBlockUpdateH\x00\x42\x07\n\x05\x66rame\"\xdd\x01\n\x07Verdict\x12\x19\n\x08\x65vent_id\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x1b\n\nsession_id\x18\x02 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x10\n\x08mad_code\x18\x04 \x01(\t\x12\x32\n\x06policy\x18\x06 \x01(\x0b\x32\".adrian.core_api.v1.PolicySnapshot\x12.\n\x04hitl\x18\x07 \x01(\x0b\x32 .adrian.core_api.v1.HitlResponseJ\x04\x08\x03\x10\x04J\x04\x08\x05\x10\x06R\x0e\x63lassificationR\x08\x65scalate*L\n\x08PairType\x12\x19\n\x15PAIR_TYPE_UNSPECIFIED\x10\x00\x12\x11\n\rPAIR_TYPE_LLM\x10\x01\x12\x12\n\x0ePAIR_TYPE_TOOL\x10\x02*K\n\x04Mode\x12\x14\n\x10MODE_UNSPECIFIED\x10\x00\x12\x0e\n\nMODE_ALERT\x10\x01\x12\r\n\tMODE_HITL\x10\x02\x12\x0e\n\nMODE_BLOCK\x10\x03\x42\x31Z/github.com/secure-agentics/adrian/server/pkg/pbb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'event_pb2', _globals) -if _descriptor._USE_C_DESCRIPTORS == False: - _globals['DESCRIPTOR']._options = None +if not _descriptor._USE_C_DESCRIPTORS: + _globals['DESCRIPTOR']._loaded_options = None _globals['DESCRIPTOR']._serialized_options = b'Z/github.com/secure-agentics/adrian/server/pkg/pb' - _globals['_TOOLCALL'].fields_by_name['name']._options = None + _globals['_TOOLCALL'].fields_by_name['name']._loaded_options = None _globals['_TOOLCALL'].fields_by_name['name']._serialized_options = b'\272H\004r\002\020\001' - _globals['_TOKENUSAGE'].fields_by_name['prompt_tokens']._options = None + _globals['_TOKENUSAGE'].fields_by_name['prompt_tokens']._loaded_options = None _globals['_TOKENUSAGE'].fields_by_name['prompt_tokens']._serialized_options = b'\272H\004\032\002(\000' - _globals['_TOKENUSAGE'].fields_by_name['completion_tokens']._options = None + _globals['_TOKENUSAGE'].fields_by_name['completion_tokens']._loaded_options = None _globals['_TOKENUSAGE'].fields_by_name['completion_tokens']._serialized_options = b'\272H\004\032\002(\000' - _globals['_TOKENUSAGE'].fields_by_name['total_tokens']._options = None + _globals['_TOKENUSAGE'].fields_by_name['total_tokens']._loaded_options = None _globals['_TOKENUSAGE'].fields_by_name['total_tokens']._serialized_options = b'\272H\004\032\002(\000' - _globals['_TOOLPAIRDATA'].fields_by_name['tool_name']._options = None + _globals['_TOOLPAIRDATA'].fields_by_name['tool_name']._loaded_options = None _globals['_TOOLPAIRDATA'].fields_by_name['tool_name']._serialized_options = b'\272H\004r\002\020\001' - _globals['_PAIREDEVENT'].fields_by_name['event_id']._options = None + _globals['_PAIREDEVENT'].fields_by_name['event_id']._loaded_options = None _globals['_PAIREDEVENT'].fields_by_name['event_id']._serialized_options = b'\272H\004r\002\020\001' - _globals['_PAIREDEVENT'].fields_by_name['session_id']._options = None + _globals['_PAIREDEVENT'].fields_by_name['session_id']._loaded_options = None _globals['_PAIREDEVENT'].fields_by_name['session_id']._serialized_options = b'\272H\004r\002\020\001' - _globals['_PAIREDEVENT'].fields_by_name['pair_type']._options = None + _globals['_PAIREDEVENT'].fields_by_name['pair_type']._loaded_options = None _globals['_PAIREDEVENT'].fields_by_name['pair_type']._serialized_options = b'\272H\003\310\001\001' - _globals['_MCPSERVER'].fields_by_name['name']._options = None + _globals['_MCPSERVER'].fields_by_name['name']._loaded_options = None _globals['_MCPSERVER'].fields_by_name['name']._serialized_options = b'\272H\004r\002\020\001' - _globals['_SESSIONLOGIN'].fields_by_name['session_id']._options = None + _globals['_SESSIONLOGIN'].fields_by_name['session_id']._loaded_options = None _globals['_SESSIONLOGIN'].fields_by_name['session_id']._serialized_options = b'\272H\004r\002\020\001' - _globals['_VERDICT'].fields_by_name['event_id']._options = None + _globals['_VERDICT'].fields_by_name['event_id']._loaded_options = None _globals['_VERDICT'].fields_by_name['event_id']._serialized_options = b'\272H\004r\002\020\001' - _globals['_VERDICT'].fields_by_name['session_id']._options = None + _globals['_VERDICT'].fields_by_name['session_id']._loaded_options = None _globals['_VERDICT'].fields_by_name['session_id']._serialized_options = b'\272H\004r\002\020\001' - _globals['_PAIRTYPE']._serialized_start=3011 - _globals['_PAIRTYPE']._serialized_end=3087 - _globals['_MODE']._serialized_start=3089 - _globals['_MODE']._serialized_end=3164 + _globals['_PAIRTYPE']._serialized_start=2520 + _globals['_PAIRTYPE']._serialized_end=2596 + _globals['_MODE']._serialized_start=2598 + _globals['_MODE']._serialized_end=2673 _globals['_CHATMESSAGE']._serialized_start=64 - _globals['_CHATMESSAGE']._serialized_end=123 - _globals['_TOOLCALL']._serialized_start=125 - _globals['_TOOLCALL']._serialized_end=200 - _globals['_TOKENUSAGE']._serialized_start=203 - _globals['_TOKENUSAGE']._serialized_end=359 - _globals['_AGENTCONTEXT']._serialized_start=361 - _globals['_AGENTCONTEXT']._serialized_end=482 - _globals['_LLMPAIRDATA']._serialized_start=485 - _globals['_LLMPAIRDATA']._serialized_end=720 - _globals['_TOOLPAIRDATA']._serialized_start=723 - _globals['_TOOLPAIRDATA']._serialized_end=855 - _globals['_PAIREDEVENT']._serialized_start=858 - _globals['_PAIREDEVENT']._serialized_end=1469 - _globals['_PAIREDEVENTBATCH']._serialized_start=1471 - _globals['_PAIREDEVENTBATCH']._serialized_end=1546 - _globals['_MCPSERVER']._serialized_start=1548 - _globals['_MCPSERVER']._serialized_end=1646 - _globals['_MCPINVENTORY']._serialized_start=1648 - _globals['_MCPINVENTORY']._serialized_end=1719 - _globals['_LLMSTACK']._serialized_start=1721 - _globals['_LLMSTACK']._serialized_end=1781 - _globals['_SESSIONLOGIN']._serialized_start=1784 - _globals['_SESSIONLOGIN']._serialized_end=2015 - _globals['_CLIENTFRAME']._serialized_start=2018 - _globals['_CLIENTFRAME']._serialized_end=2259 - _globals['_POLICYSNAPSHOT']._serialized_start=2262 - _globals['_POLICYSNAPSHOT']._serialized_end=2440 - _globals['_HITLRESPONSE']._serialized_start=2442 - _globals['_HITLRESPONSE']._serialized_end=2503 - _globals['_LOGINACK']._serialized_start=2505 - _globals['_LOGINACK']._serialized_end=2599 - _globals['_SERVERFRAME']._serialized_start=2602 - _globals['_SERVERFRAME']._serialized_end=2742 - _globals['_VERDICT']._serialized_start=2745 - _globals['_VERDICT']._serialized_end=3009 + _globals['_CHATMESSAGE']._serialized_end=108 + _globals['_TOOLCALL']._serialized_start=110 + _globals['_TOOLCALL']._serialized_end=169 + _globals['_TOKENUSAGE']._serialized_start=171 + _globals['_TOKENUSAGE']._serialized_end=282 + _globals['_AGENTCONTEXT']._serialized_start=284 + _globals['_AGENTCONTEXT']._serialized_end=365 + _globals['_LLMPAIRDATA']._serialized_start=368 + _globals['_LLMPAIRDATA']._serialized_end=560 + _globals['_TOOLPAIRDATA']._serialized_start=562 + _globals['_TOOLPAIRDATA']._serialized_end=657 + _globals['_PAIREDEVENT']._serialized_start=660 + _globals['_PAIREDEVENT']._serialized_end=1134 + _globals['_PAIREDEVENTBATCH']._serialized_start=1136 + _globals['_PAIREDEVENTBATCH']._serialized_end=1203 + _globals['_MCPSERVER']._serialized_start=1205 + _globals['_MCPSERVER']._serialized_end=1276 + _globals['_MCPINVENTORY']._serialized_start=1278 + _globals['_MCPINVENTORY']._serialized_end=1340 + _globals['_LLMSTACK']._serialized_start=1342 + _globals['_LLMSTACK']._serialized_end=1385 + _globals['_SESSIONLOGIN']._serialized_start=1388 + _globals['_SESSIONLOGIN']._serialized_end=1561 + _globals['_CLIENTFRAME']._serialized_start=1564 + _globals['_CLIENTFRAME']._serialized_end=1771 + _globals['_POLICYSNAPSHOT']._serialized_start=1774 + _globals['_POLICYSNAPSHOT']._serialized_end=1906 + _globals['_HITLRESPONSE']._serialized_start=1908 + _globals['_HITLRESPONSE']._serialized_end=1950 + _globals['_LOGINACK']._serialized_start=1952 + _globals['_LOGINACK']._serialized_end=2059 + _globals['_MCPBLOCKUPDATE']._serialized_start=2061 + _globals['_MCPBLOCKUPDATE']._serialized_end=2106 + _globals['_SERVERFRAME']._serialized_start=2109 + _globals['_SERVERFRAME']._serialized_end=2294 + _globals['_VERDICT']._serialized_start=2297 + _globals['_VERDICT']._serialized_end=2518 # @@protoc_insertion_point(module_scope) diff --git a/integrations/claude-code/adrian_cc/proto/event_pb2.pyi b/integrations/claude-code/adrian_cc/proto/event_pb2.pyi index e295646..24f7be8 100644 --- a/integrations/claude-code/adrian_cc/proto/event_pb2.pyi +++ b/integrations/claude-code/adrian_cc/proto/event_pb2.pyi @@ -1,9 +1,10 @@ -from .buf.validate import validate_pb2 as _validate_pb2 +from buf.validate import validate_pb2 as _validate_pb2 from google.protobuf.internal import containers as _containers from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper from google.protobuf import descriptor as _descriptor from google.protobuf import message as _message -from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union +from collections.abc import Iterable as _Iterable, Mapping as _Mapping +from typing import ClassVar as _ClassVar, Optional as _Optional, Union as _Union DESCRIPTOR: _descriptor.FileDescriptor @@ -19,6 +20,12 @@ class Mode(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): MODE_ALERT: _ClassVar[Mode] MODE_HITL: _ClassVar[Mode] MODE_BLOCK: _ClassVar[Mode] + +class VerdictStatus(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): + __slots__ = () + VERDICT_STATUS_UNSPECIFIED: _ClassVar[VerdictStatus] + VERDICT_STATUS_OK: _ClassVar[VerdictStatus] + VERDICT_STATUS_ERROR: _ClassVar[VerdictStatus] PAIR_TYPE_UNSPECIFIED: PairType PAIR_TYPE_LLM: PairType PAIR_TYPE_TOOL: PairType @@ -26,6 +33,9 @@ MODE_UNSPECIFIED: Mode MODE_ALERT: Mode MODE_HITL: Mode MODE_BLOCK: Mode +VERDICT_STATUS_UNSPECIFIED: VerdictStatus +VERDICT_STATUS_OK: VerdictStatus +VERDICT_STATUS_ERROR: VerdictStatus class ChatMessage(_message.Message): __slots__ = ("role", "content") @@ -92,7 +102,7 @@ class ToolPairData(_message.Message): def __init__(self, tool_name: _Optional[str] = ..., tool_call_id: _Optional[str] = ..., input: _Optional[str] = ..., output: _Optional[str] = ...) -> None: ... class PairedEvent(_message.Message): - __slots__ = ("event_id", "invocation_id", "session_id", "run_id", "parent_run_id", "timestamp", "pair_type", "agent", "parent", "llm", "tool", "metadata_json", "connection_id", "source") + __slots__ = ("event_id", "invocation_id", "session_id", "run_id", "parent_run_id", "timestamp", "pair_type", "agent", "parent", "llm", "tool", "connection_id", "metadata_json", "source") EVENT_ID_FIELD_NUMBER: _ClassVar[int] INVOCATION_ID_FIELD_NUMBER: _ClassVar[int] SESSION_ID_FIELD_NUMBER: _ClassVar[int] @@ -104,8 +114,8 @@ class PairedEvent(_message.Message): PARENT_FIELD_NUMBER: _ClassVar[int] LLM_FIELD_NUMBER: _ClassVar[int] TOOL_FIELD_NUMBER: _ClassVar[int] - METADATA_JSON_FIELD_NUMBER: _ClassVar[int] CONNECTION_ID_FIELD_NUMBER: _ClassVar[int] + METADATA_JSON_FIELD_NUMBER: _ClassVar[int] SOURCE_FIELD_NUMBER: _ClassVar[int] event_id: str invocation_id: str @@ -118,10 +128,10 @@ class PairedEvent(_message.Message): parent: AgentContext llm: LlmPairData tool: ToolPairData - metadata_json: bytes connection_id: str + metadata_json: bytes source: str - def __init__(self, event_id: _Optional[str] = ..., invocation_id: _Optional[str] = ..., session_id: _Optional[str] = ..., run_id: _Optional[str] = ..., parent_run_id: _Optional[str] = ..., timestamp: _Optional[str] = ..., pair_type: _Optional[_Union[PairType, str]] = ..., agent: _Optional[_Union[AgentContext, _Mapping]] = ..., parent: _Optional[_Union[AgentContext, _Mapping]] = ..., llm: _Optional[_Union[LlmPairData, _Mapping]] = ..., tool: _Optional[_Union[ToolPairData, _Mapping]] = ..., metadata_json: _Optional[bytes] = ..., connection_id: _Optional[str] = ..., source: _Optional[str] = ...) -> None: ... + def __init__(self, event_id: _Optional[str] = ..., invocation_id: _Optional[str] = ..., session_id: _Optional[str] = ..., run_id: _Optional[str] = ..., parent_run_id: _Optional[str] = ..., timestamp: _Optional[str] = ..., pair_type: _Optional[_Union[PairType, str]] = ..., agent: _Optional[_Union[AgentContext, _Mapping]] = ..., parent: _Optional[_Union[AgentContext, _Mapping]] = ..., llm: _Optional[_Union[LlmPairData, _Mapping]] = ..., tool: _Optional[_Union[ToolPairData, _Mapping]] = ..., connection_id: _Optional[str] = ..., metadata_json: _Optional[bytes] = ..., source: _Optional[str] = ...) -> None: ... class PairedEventBatch(_message.Message): __slots__ = ("events",) @@ -130,14 +140,22 @@ class PairedEventBatch(_message.Message): def __init__(self, events: _Optional[_Iterable[_Union[PairedEvent, _Mapping]]] = ...) -> None: ... class McpServer(_message.Message): - __slots__ = ("name", "transport", "endpoint") + __slots__ = ("name", "transport", "endpoint", "version", "protocol_version", "server_info_name", "tools_json") NAME_FIELD_NUMBER: _ClassVar[int] TRANSPORT_FIELD_NUMBER: _ClassVar[int] ENDPOINT_FIELD_NUMBER: _ClassVar[int] + VERSION_FIELD_NUMBER: _ClassVar[int] + PROTOCOL_VERSION_FIELD_NUMBER: _ClassVar[int] + SERVER_INFO_NAME_FIELD_NUMBER: _ClassVar[int] + TOOLS_JSON_FIELD_NUMBER: _ClassVar[int] name: str transport: str endpoint: str - def __init__(self, name: _Optional[str] = ..., transport: _Optional[str] = ..., endpoint: _Optional[str] = ...) -> None: ... + version: str + protocol_version: str + server_info_name: str + tools_json: str + def __init__(self, name: _Optional[str] = ..., transport: _Optional[str] = ..., endpoint: _Optional[str] = ..., version: _Optional[str] = ..., protocol_version: _Optional[str] = ..., server_info_name: _Optional[str] = ..., tools_json: _Optional[str] = ...) -> None: ... class McpInventory(_message.Message): __slots__ = ("servers",) @@ -145,6 +163,24 @@ class McpInventory(_message.Message): servers: _containers.RepeatedCompositeFieldContainer[McpServer] def __init__(self, servers: _Optional[_Iterable[_Union[McpServer, _Mapping]]] = ...) -> None: ... +class InstalledPlugin(_message.Message): + __slots__ = ("name", "enabled", "version", "marketplace") + NAME_FIELD_NUMBER: _ClassVar[int] + ENABLED_FIELD_NUMBER: _ClassVar[int] + VERSION_FIELD_NUMBER: _ClassVar[int] + MARKETPLACE_FIELD_NUMBER: _ClassVar[int] + name: str + enabled: bool + version: str + marketplace: str + def __init__(self, name: _Optional[str] = ..., enabled: _Optional[bool] = ..., version: _Optional[str] = ..., marketplace: _Optional[str] = ...) -> None: ... + +class PluginInventory(_message.Message): + __slots__ = ("plugins",) + PLUGINS_FIELD_NUMBER: _ClassVar[int] + plugins: _containers.RepeatedCompositeFieldContainer[InstalledPlugin] + def __init__(self, plugins: _Optional[_Iterable[_Union[InstalledPlugin, _Mapping]]] = ...) -> None: ... + class LLMStack(_message.Message): __slots__ = ("provider", "model") PROVIDER_FIELD_NUMBER: _ClassVar[int] @@ -168,34 +204,38 @@ class SessionLogin(_message.Message): def __init__(self, session_id: _Optional[str] = ..., llm_stack: _Optional[_Union[LLMStack, _Mapping]] = ..., schema_version: _Optional[int] = ..., source: _Optional[str] = ..., connection_id: _Optional[str] = ...) -> None: ... class ClientFrame(_message.Message): - __slots__ = ("login", "paired_batch", "mcp_inventory") + __slots__ = ("login", "paired_batch", "mcp_inventory", "plugin_inventory") LOGIN_FIELD_NUMBER: _ClassVar[int] PAIRED_BATCH_FIELD_NUMBER: _ClassVar[int] MCP_INVENTORY_FIELD_NUMBER: _ClassVar[int] + PLUGIN_INVENTORY_FIELD_NUMBER: _ClassVar[int] login: SessionLogin paired_batch: PairedEventBatch mcp_inventory: McpInventory - def __init__(self, login: _Optional[_Union[SessionLogin, _Mapping]] = ..., paired_batch: _Optional[_Union[PairedEventBatch, _Mapping]] = ..., mcp_inventory: _Optional[_Union[McpInventory, _Mapping]] = ...) -> None: ... + plugin_inventory: PluginInventory + def __init__(self, login: _Optional[_Union[SessionLogin, _Mapping]] = ..., paired_batch: _Optional[_Union[PairedEventBatch, _Mapping]] = ..., mcp_inventory: _Optional[_Union[McpInventory, _Mapping]] = ..., plugin_inventory: _Optional[_Union[PluginInventory, _Mapping]] = ...) -> None: ... class PolicySnapshot(_message.Message): - __slots__ = ("mode", "policy_m0", "policy_m2", "policy_m3", "policy_m4") + __slots__ = ("mode", "policy_m0", "policy_m2", "policy_m3", "policy_m4", "fail_closed_on_classifier_error") MODE_FIELD_NUMBER: _ClassVar[int] POLICY_M0_FIELD_NUMBER: _ClassVar[int] POLICY_M2_FIELD_NUMBER: _ClassVar[int] POLICY_M3_FIELD_NUMBER: _ClassVar[int] POLICY_M4_FIELD_NUMBER: _ClassVar[int] + FAIL_CLOSED_ON_CLASSIFIER_ERROR_FIELD_NUMBER: _ClassVar[int] mode: Mode policy_m0: bool policy_m2: bool policy_m3: bool policy_m4: bool - def __init__(self, mode: _Optional[_Union[Mode, str]] = ..., policy_m0: bool = ..., policy_m2: bool = ..., policy_m3: bool = ..., policy_m4: bool = ...) -> None: ... + fail_closed_on_classifier_error: bool + def __init__(self, mode: _Optional[_Union[Mode, str]] = ..., policy_m0: _Optional[bool] = ..., policy_m2: _Optional[bool] = ..., policy_m3: _Optional[bool] = ..., policy_m4: _Optional[bool] = ..., fail_closed_on_classifier_error: _Optional[bool] = ...) -> None: ... class HitlResponse(_message.Message): __slots__ = ("continue_execution",) CONTINUE_EXECUTION_FIELD_NUMBER: _ClassVar[int] continue_execution: bool - def __init__(self, continue_execution: bool = ...) -> None: ... + def __init__(self, continue_execution: _Optional[bool] = ...) -> None: ... class LoginAck(_message.Message): __slots__ = ("policy", "source") @@ -214,15 +254,17 @@ class ServerFrame(_message.Message): def __init__(self, login_ack: _Optional[_Union[LoginAck, _Mapping]] = ..., verdict: _Optional[_Union[Verdict, _Mapping]] = ...) -> None: ... class Verdict(_message.Message): - __slots__ = ("event_id", "session_id", "mad_code", "policy", "hitl") + __slots__ = ("event_id", "session_id", "mad_code", "policy", "hitl", "status") EVENT_ID_FIELD_NUMBER: _ClassVar[int] SESSION_ID_FIELD_NUMBER: _ClassVar[int] MAD_CODE_FIELD_NUMBER: _ClassVar[int] POLICY_FIELD_NUMBER: _ClassVar[int] HITL_FIELD_NUMBER: _ClassVar[int] + STATUS_FIELD_NUMBER: _ClassVar[int] event_id: str session_id: str mad_code: str policy: PolicySnapshot hitl: HitlResponse - def __init__(self, event_id: _Optional[str] = ..., session_id: _Optional[str] = ..., mad_code: _Optional[str] = ..., policy: _Optional[_Union[PolicySnapshot, _Mapping]] = ..., hitl: _Optional[_Union[HitlResponse, _Mapping]] = ...) -> None: ... + status: VerdictStatus + def __init__(self, event_id: _Optional[str] = ..., session_id: _Optional[str] = ..., mad_code: _Optional[str] = ..., policy: _Optional[_Union[PolicySnapshot, _Mapping]] = ..., hitl: _Optional[_Union[HitlResponse, _Mapping]] = ..., status: _Optional[_Union[VerdictStatus, str]] = ...) -> None: ... diff --git a/proto/event.proto b/proto/event.proto index 1281cbe..f51f6cc 100644 --- a/proto/event.proto +++ b/proto/event.proto @@ -146,6 +146,16 @@ message McpServer { // URL for SSE/HTTP/WebSocket transports, or the joined command line // for stdio. Empty string when neither is available. string endpoint = 3; + // Package or server version (e.g. "3.2.5"). Empty when unknown. + // Sourced from serverInfo.version (initialize handshake), falling + // back to npm/pip registry lookups. + string version = 4; + // MCP protocol version from initialize handshake (e.g. "2025-06-18"). + string protocol_version = 5; + // Server's self-reported name from serverInfo.name in initialize. + string server_info_name = 6; + // JSON-encoded list of tool names the server exposes. + string tools_json = 7; } // McpInventory is the SDK's one-shot dump of its connected MCP servers. @@ -155,6 +165,25 @@ message McpInventory { repeated McpServer servers = 1; } +// InstalledPlugin describes one plugin installed in a Claude Code session. +message InstalledPlugin { + // Plugin identifier (e.g. "playwright@claude-plugins-official"). + string name = 1 [(buf.validate.field).string.min_len = 1]; + // Whether the plugin is currently enabled. + bool enabled = 2; + // Plugin version string (e.g. "1.1.53" or "unknown"). + string version = 3; + // Marketplace or source the plugin was installed from. + string marketplace = 4; +} + +// PluginInventory is a one-shot dump of installed plugins. +// Sent once per WS session after login, alongside McpInventory. +// The server replaces the session's entire plugin list on receipt. +message PluginInventory { + repeated InstalledPlugin plugins = 1; +} + // LLMStack identifies the LLM provider and model for a session. message LLMStack { // Provider name (e.g. "anthropic", "openai"). @@ -203,6 +232,8 @@ message ClientFrame { PairedEventBatch paired_batch = 3; // One-shot MCP-server inventory dump from the SDK after login. McpInventory mcp_inventory = 4; + // One-shot plugin inventory dump (Claude Code sessions). + PluginInventory plugin_inventory = 5; } } @@ -266,6 +297,7 @@ message HitlResponse { // effect. message LoginAck { PolicySnapshot policy = 1; + string source = 2; } // ServerFrame is the top-level envelope for all server→SDK WebSocket diff --git a/sdk/python/adrian/__init__.py b/sdk/python/adrian/__init__.py index cfcafdb..4aed9e6 100644 --- a/sdk/python/adrian/__init__.py +++ b/sdk/python/adrian/__init__.py @@ -472,6 +472,10 @@ async def _send_mcp_inventory() -> None: added.name = server.name added.transport = server.transport added.endpoint = server.endpoint + added.version = server.version + added.protocol_version = server.protocol_version + added.server_info_name = server.server_info_name + added.tools_json = server.tools_json await ws._send_frame(frame) # pyright: ignore[reportPrivateUsage] diff --git a/sdk/python/adrian/mcp.py b/sdk/python/adrian/mcp.py index 8101752..70044e5 100644 --- a/sdk/python/adrian/mcp.py +++ b/sdk/python/adrian/mcp.py @@ -30,6 +30,7 @@ import contextlib import logging +import subprocess import sys from collections.abc import Mapping from typing import Any, cast @@ -126,6 +127,97 @@ def _fire_on_mcp_server(server: McpServer) -> None: fire(config.on_mcp_server, server, name="on_mcp_server") +def _resolve_version_from_connection(connection: Mapping[str, Any]) -> str: + """Try to resolve a version for an MCP server from its connection config. + + Supports npx (npm view) and pip-installed (pip show) MCP servers. + """ + transport = str(connection.get("transport") or "").lower() + if transport != "stdio": + return "" + command = str(connection.get("command", "")) + args = connection.get("args") or [] + if not isinstance(args, (list, tuple)): + args = [] + str_args = [str(a) for a in args] + + if command in ("npx", "npx.cmd"): + pkg = _extract_npx_package(str_args) + if pkg: + return _npm_version(pkg) + + if command in ("uvx", "pipx"): + pkg = _extract_first_positional(str_args) + if pkg: + return _pip_version(pkg) + + if command in ("python", "python3") and "-m" in str_args: + idx = str_args.index("-m") + if idx + 1 < len(str_args): + return _pip_version(str_args[idx + 1].replace(".", "-")) + + if not command.startswith("/"): + return _pip_version(command) + + return "" + + +def _extract_npx_package(args: list[str]) -> str: + """Extract the npm package name from npx args.""" + skip_next = False + for arg in args: + if skip_next: + skip_next = False + continue + if arg in ("-y", "--yes", "-q", "--quiet"): + continue + if arg.startswith("-p") or arg == "--package": + skip_next = True + continue + if arg.startswith("-"): + continue + return arg + return "" + + +def _extract_first_positional(args: list[str]) -> str: + """Extract the first non-flag argument.""" + for arg in args: + if not arg.startswith("-"): + return arg + return "" + + +def _npm_version(pkg: str) -> str: + """Get version from npm registry.""" + try: + result = subprocess.run( + ["npm", "view", pkg, "version"], + capture_output=True, text=True, timeout=5, + ) + if result.returncode == 0 and result.stdout.strip(): + return result.stdout.strip() + except Exception: + pass + return "" + + +def _pip_version(pkg: str) -> str: + """Get version from pip show.""" + try: + result = subprocess.run( + ["pip", "show", pkg], + capture_output=True, text=True, timeout=5, + ) + if result.returncode == 0: + for line in result.stdout.splitlines(): + if line.startswith("Version:"): + return line.split(":", 1)[1].strip() + except Exception: + pass + return "" + + def _server_from_connection(name: str, connection: Any) -> McpServer: # noqa: ANN401 """Convert a ``Connection`` mapping into an ``McpServer``.""" if not isinstance(connection, Mapping): @@ -134,8 +226,9 @@ def _server_from_connection(name: str, connection: Any) -> McpServer: # noqa: A conn = cast("Mapping[str, Any]", connection) transport = str(conn.get("transport") or "").lower() or "unknown" endpoint = _endpoint_for(transport, conn) + version = _resolve_version_from_connection(conn) - return McpServer(name=name, transport=transport, endpoint=endpoint) + return McpServer(name=name, transport=transport, endpoint=endpoint, version=version) def _endpoint_for(transport: str, connection: Mapping[str, Any]) -> str: diff --git a/sdk/python/adrian/proto/event_pb2.py b/sdk/python/adrian/proto/event_pb2.py index e8f2f50..222f1d5 100644 --- a/sdk/python/adrian/proto/event_pb2.py +++ b/sdk/python/adrian/proto/event_pb2.py @@ -2,7 +2,7 @@ # Generated by the protocol buffer compiler. DO NOT EDIT! # NO CHECKED-IN PROTOBUF GENCODE # source: event.proto -# Protobuf Python Version: 6.33.5 +# Protobuf Python Version: 7.35.1 """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool @@ -11,9 +11,9 @@ from google.protobuf.internal import builder as _builder _runtime_version.ValidateProtobufRuntimeVersion( _runtime_version.Domain.PUBLIC, - 6, - 33, - 5, + 7, + 35, + 1, '', 'event.proto' ) @@ -25,7 +25,7 @@ from .buf.validate import validate_pb2 as buf_dot_validate_dot_validate__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0b\x65vent.proto\x12\x12\x61\x64rian.core_api.v1\x1a\x1b\x62uf/validate/validate.proto\",\n\x0b\x43hatMessage\x12\x0c\n\x04role\x18\x01 \x01(\t\x12\x0f\n\x07\x63ontent\x18\x02 \x01(\t\";\n\x08ToolCall\x12\x15\n\x04name\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x0c\n\x04\x61rgs\x18\x02 \x01(\t\x12\n\n\x02id\x18\x03 \x01(\t\"o\n\nTokenUsage\x12\x1e\n\rprompt_tokens\x18\x01 \x01(\x05\x42\x07\xbaH\x04\x1a\x02(\x00\x12\"\n\x11\x63ompletion_tokens\x18\x02 \x01(\x05\x42\x07\xbaH\x04\x1a\x02(\x00\x12\x1d\n\x0ctotal_tokens\x18\x03 \x01(\x05\x42\x07\xbaH\x04\x1a\x02(\x00\"Q\n\x0c\x41gentContext\x12\x10\n\x08\x61gent_id\x18\x01 \x01(\t\x12\x15\n\rsystem_prompt\x18\x02 \x01(\t\x12\x18\n\x10user_instruction\x18\x03 \x01(\t\"\xc0\x01\n\x0bLlmPairData\x12\r\n\x05model\x18\x01 \x01(\t\x12\x31\n\x08messages\x18\x02 \x03(\x0b\x32\x1f.adrian.core_api.v1.ChatMessage\x12\x0e\n\x06output\x18\x03 \x01(\t\x12\x30\n\ntool_calls\x18\x04 \x03(\x0b\x32\x1c.adrian.core_api.v1.ToolCall\x12-\n\x05usage\x18\x05 \x01(\x0b\x32\x1e.adrian.core_api.v1.TokenUsage\"_\n\x0cToolPairData\x12\x1a\n\ttool_name\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x14\n\x0ctool_call_id\x18\x02 \x01(\t\x12\r\n\x05input\x18\x03 \x01(\t\x12\x0e\n\x06output\x18\x04 \x01(\t\"\xc3\x03\n\x0bPairedEvent\x12\x19\n\x08\x65vent_id\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x15\n\rinvocation_id\x18\x02 \x01(\t\x12\x1b\n\nsession_id\x18\x03 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x0e\n\x06run_id\x18\x04 \x01(\t\x12\x15\n\rparent_run_id\x18\x05 \x01(\t\x12\x11\n\ttimestamp\x18\x06 \x01(\t\x12\x37\n\tpair_type\x18\x07 \x01(\x0e\x32\x1c.adrian.core_api.v1.PairTypeB\x06\xbaH\x03\xc8\x01\x01\x12/\n\x05\x61gent\x18\x08 \x01(\x0b\x32 .adrian.core_api.v1.AgentContext\x12\x30\n\x06parent\x18\t \x01(\x0b\x32 .adrian.core_api.v1.AgentContext\x12.\n\x03llm\x18\n \x01(\x0b\x32\x1f.adrian.core_api.v1.LlmPairDataH\x00\x12\x30\n\x04tool\x18\x0b \x01(\x0b\x32 .adrian.core_api.v1.ToolPairDataH\x00\x12\x15\n\rmetadata_json\x18\x14 \x01(\x0c\x12\x0e\n\x06source\x18\x15 \x01(\tB\x06\n\x04\x64\x61ta\"C\n\x10PairedEventBatch\x12/\n\x06\x65vents\x18\x01 \x03(\x0b\x32\x1f.adrian.core_api.v1.PairedEvent\"G\n\tMcpServer\x12\x15\n\x04name\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x11\n\ttransport\x18\x02 \x01(\t\x12\x10\n\x08\x65ndpoint\x18\x03 \x01(\t\">\n\x0cMcpInventory\x12.\n\x07servers\x18\x01 \x03(\x0b\x32\x1d.adrian.core_api.v1.McpServer\"+\n\x08LLMStack\x12\x10\n\x08provider\x18\x01 \x01(\t\x12\r\n\x05model\x18\x02 \x01(\t\"\x96\x01\n\x0cSessionLogin\x12\x1b\n\nsession_id\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12/\n\tllm_stack\x18\x02 \x01(\x0b\x32\x1c.adrian.core_api.v1.LLMStack\x12\x16\n\x0eschema_version\x18\x04 \x01(\r\x12\x0e\n\x06source\x18\x05 \x01(\tJ\x04\x08\x03\x10\x04R\nblock_mode\"\xcf\x01\n\x0b\x43lientFrame\x12\x31\n\x05login\x18\x01 \x01(\x0b\x32 .adrian.core_api.v1.SessionLoginH\x00\x12<\n\x0cpaired_batch\x18\x03 \x01(\x0b\x32$.adrian.core_api.v1.PairedEventBatchH\x00\x12\x39\n\rmcp_inventory\x18\x04 \x01(\x0b\x32 .adrian.core_api.v1.McpInventoryH\x00\x42\x07\n\x05\x66rameJ\x04\x08\x02\x10\x03R\x05\x62\x61tch\"\xad\x01\n\x0ePolicySnapshot\x12&\n\x04mode\x18\x01 \x01(\x0e\x32\x18.adrian.core_api.v1.Mode\x12\x11\n\tpolicy_m0\x18\x02 \x01(\x08\x12\x11\n\tpolicy_m2\x18\x03 \x01(\x08\x12\x11\n\tpolicy_m3\x18\x04 \x01(\x08\x12\x11\n\tpolicy_m4\x18\x05 \x01(\x08\x12\'\n\x1f\x66\x61il_closed_on_classifier_error\x18\x06 \x01(\x08\"*\n\x0cHitlResponse\x12\x1a\n\x12\x63ontinue_execution\x18\x01 \x01(\x08\">\n\x08LoginAck\x12\x32\n\x06policy\x18\x01 \x01(\x0b\x32\".adrian.core_api.v1.PolicySnapshot\"y\n\x0bServerFrame\x12\x31\n\tlogin_ack\x18\x01 \x01(\x0b\x32\x1c.adrian.core_api.v1.LoginAckH\x00\x12.\n\x07verdict\x18\x02 \x01(\x0b\x32\x1b.adrian.core_api.v1.VerdictH\x00\x42\x07\n\x05\x66rame\"\x90\x02\n\x07Verdict\x12\x19\n\x08\x65vent_id\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x1b\n\nsession_id\x18\x02 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x10\n\x08mad_code\x18\x04 \x01(\t\x12\x32\n\x06policy\x18\x06 \x01(\x0b\x32\".adrian.core_api.v1.PolicySnapshot\x12.\n\x04hitl\x18\x07 \x01(\x0b\x32 .adrian.core_api.v1.HitlResponse\x12\x31\n\x06status\x18\x08 \x01(\x0e\x32!.adrian.core_api.v1.VerdictStatusJ\x04\x08\x03\x10\x04J\x04\x08\x05\x10\x06R\x0e\x63lassificationR\x08\x65scalate*L\n\x08PairType\x12\x19\n\x15PAIR_TYPE_UNSPECIFIED\x10\x00\x12\x11\n\rPAIR_TYPE_LLM\x10\x01\x12\x12\n\x0ePAIR_TYPE_TOOL\x10\x02*K\n\x04Mode\x12\x14\n\x10MODE_UNSPECIFIED\x10\x00\x12\x0e\n\nMODE_ALERT\x10\x01\x12\r\n\tMODE_HITL\x10\x02\x12\x0e\n\nMODE_BLOCK\x10\x03*`\n\rVerdictStatus\x12\x1e\n\x1aVERDICT_STATUS_UNSPECIFIED\x10\x00\x12\x15\n\x11VERDICT_STATUS_OK\x10\x01\x12\x18\n\x14VERDICT_STATUS_ERROR\x10\x02\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0b\x65vent.proto\x12\x12\x61\x64rian.core_api.v1\x1a\x1b\x62uf/validate/validate.proto\",\n\x0b\x43hatMessage\x12\x0c\n\x04role\x18\x01 \x01(\t\x12\x0f\n\x07\x63ontent\x18\x02 \x01(\t\";\n\x08ToolCall\x12\x15\n\x04name\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x0c\n\x04\x61rgs\x18\x02 \x01(\t\x12\n\n\x02id\x18\x03 \x01(\t\"o\n\nTokenUsage\x12\x1e\n\rprompt_tokens\x18\x01 \x01(\x05\x42\x07\xbaH\x04\x1a\x02(\x00\x12\"\n\x11\x63ompletion_tokens\x18\x02 \x01(\x05\x42\x07\xbaH\x04\x1a\x02(\x00\x12\x1d\n\x0ctotal_tokens\x18\x03 \x01(\x05\x42\x07\xbaH\x04\x1a\x02(\x00\"Q\n\x0c\x41gentContext\x12\x10\n\x08\x61gent_id\x18\x01 \x01(\t\x12\x15\n\rsystem_prompt\x18\x02 \x01(\t\x12\x18\n\x10user_instruction\x18\x03 \x01(\t\"\xc0\x01\n\x0bLlmPairData\x12\r\n\x05model\x18\x01 \x01(\t\x12\x31\n\x08messages\x18\x02 \x03(\x0b\x32\x1f.adrian.core_api.v1.ChatMessage\x12\x0e\n\x06output\x18\x03 \x01(\t\x12\x30\n\ntool_calls\x18\x04 \x03(\x0b\x32\x1c.adrian.core_api.v1.ToolCall\x12-\n\x05usage\x18\x05 \x01(\x0b\x32\x1e.adrian.core_api.v1.TokenUsage\"_\n\x0cToolPairData\x12\x1a\n\ttool_name\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x14\n\x0ctool_call_id\x18\x02 \x01(\t\x12\r\n\x05input\x18\x03 \x01(\t\x12\x0e\n\x06output\x18\x04 \x01(\t\"\xda\x03\n\x0bPairedEvent\x12\x19\n\x08\x65vent_id\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x15\n\rinvocation_id\x18\x02 \x01(\t\x12\x1b\n\nsession_id\x18\x03 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x0e\n\x06run_id\x18\x04 \x01(\t\x12\x15\n\rparent_run_id\x18\x05 \x01(\t\x12\x11\n\ttimestamp\x18\x06 \x01(\t\x12\x37\n\tpair_type\x18\x07 \x01(\x0e\x32\x1c.adrian.core_api.v1.PairTypeB\x06\xbaH\x03\xc8\x01\x01\x12/\n\x05\x61gent\x18\x08 \x01(\x0b\x32 .adrian.core_api.v1.AgentContext\x12\x30\n\x06parent\x18\t \x01(\x0b\x32 .adrian.core_api.v1.AgentContext\x12.\n\x03llm\x18\n \x01(\x0b\x32\x1f.adrian.core_api.v1.LlmPairDataH\x00\x12\x30\n\x04tool\x18\x0b \x01(\x0b\x32 .adrian.core_api.v1.ToolPairDataH\x00\x12\x15\n\rconnection_id\x18\x0c \x01(\t\x12\x15\n\rmetadata_json\x18\x14 \x01(\x0c\x12\x0e\n\x06source\x18\x15 \x01(\tB\x06\n\x04\x64\x61ta\"C\n\x10PairedEventBatch\x12/\n\x06\x65vents\x18\x01 \x03(\x0b\x32\x1f.adrian.core_api.v1.PairedEvent\"\xa0\x01\n\tMcpServer\x12\x15\n\x04name\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x11\n\ttransport\x18\x02 \x01(\t\x12\x10\n\x08\x65ndpoint\x18\x03 \x01(\t\x12\x0f\n\x07version\x18\x04 \x01(\t\x12\x18\n\x10protocol_version\x18\x05 \x01(\t\x12\x18\n\x10server_info_name\x18\x06 \x01(\t\x12\x12\n\ntools_json\x18\x07 \x01(\t\">\n\x0cMcpInventory\x12.\n\x07servers\x18\x01 \x03(\x0b\x32\x1d.adrian.core_api.v1.McpServer\"_\n\x0fInstalledPlugin\x12\x15\n\x04name\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x0f\n\x07\x65nabled\x18\x02 \x01(\x08\x12\x0f\n\x07version\x18\x03 \x01(\t\x12\x13\n\x0bmarketplace\x18\x04 \x01(\t\"G\n\x0fPluginInventory\x12\x34\n\x07plugins\x18\x01 \x03(\x0b\x32#.adrian.core_api.v1.InstalledPlugin\"+\n\x08LLMStack\x12\x10\n\x08provider\x18\x01 \x01(\t\x12\r\n\x05model\x18\x02 \x01(\t\"\xad\x01\n\x0cSessionLogin\x12\x1b\n\nsession_id\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12/\n\tllm_stack\x18\x02 \x01(\x0b\x32\x1c.adrian.core_api.v1.LLMStack\x12\x16\n\x0eschema_version\x18\x04 \x01(\r\x12\x0e\n\x06source\x18\x05 \x01(\t\x12\x15\n\rconnection_id\x18\x06 \x01(\tJ\x04\x08\x03\x10\x04R\nblock_mode\"\x90\x02\n\x0b\x43lientFrame\x12\x31\n\x05login\x18\x01 \x01(\x0b\x32 .adrian.core_api.v1.SessionLoginH\x00\x12<\n\x0cpaired_batch\x18\x03 \x01(\x0b\x32$.adrian.core_api.v1.PairedEventBatchH\x00\x12\x39\n\rmcp_inventory\x18\x04 \x01(\x0b\x32 .adrian.core_api.v1.McpInventoryH\x00\x12?\n\x10plugin_inventory\x18\x05 \x01(\x0b\x32#.adrian.core_api.v1.PluginInventoryH\x00\x42\x07\n\x05\x66rameJ\x04\x08\x02\x10\x03R\x05\x62\x61tch\"\xad\x01\n\x0ePolicySnapshot\x12&\n\x04mode\x18\x01 \x01(\x0e\x32\x18.adrian.core_api.v1.Mode\x12\x11\n\tpolicy_m0\x18\x02 \x01(\x08\x12\x11\n\tpolicy_m2\x18\x03 \x01(\x08\x12\x11\n\tpolicy_m3\x18\x04 \x01(\x08\x12\x11\n\tpolicy_m4\x18\x05 \x01(\x08\x12\'\n\x1f\x66\x61il_closed_on_classifier_error\x18\x06 \x01(\x08\"*\n\x0cHitlResponse\x12\x1a\n\x12\x63ontinue_execution\x18\x01 \x01(\x08\"N\n\x08LoginAck\x12\x32\n\x06policy\x18\x01 \x01(\x0b\x32\".adrian.core_api.v1.PolicySnapshot\x12\x0e\n\x06source\x18\x02 \x01(\t\"y\n\x0bServerFrame\x12\x31\n\tlogin_ack\x18\x01 \x01(\x0b\x32\x1c.adrian.core_api.v1.LoginAckH\x00\x12.\n\x07verdict\x18\x02 \x01(\x0b\x32\x1b.adrian.core_api.v1.VerdictH\x00\x42\x07\n\x05\x66rame\"\x90\x02\n\x07Verdict\x12\x19\n\x08\x65vent_id\x18\x01 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x1b\n\nsession_id\x18\x02 \x01(\tB\x07\xbaH\x04r\x02\x10\x01\x12\x10\n\x08mad_code\x18\x04 \x01(\t\x12\x32\n\x06policy\x18\x06 \x01(\x0b\x32\".adrian.core_api.v1.PolicySnapshot\x12.\n\x04hitl\x18\x07 \x01(\x0b\x32 .adrian.core_api.v1.HitlResponse\x12\x31\n\x06status\x18\x08 \x01(\x0e\x32!.adrian.core_api.v1.VerdictStatusJ\x04\x08\x03\x10\x04J\x04\x08\x05\x10\x06R\x0e\x63lassificationR\x08\x65scalate*L\n\x08PairType\x12\x19\n\x15PAIR_TYPE_UNSPECIFIED\x10\x00\x12\x11\n\rPAIR_TYPE_LLM\x10\x01\x12\x12\n\x0ePAIR_TYPE_TOOL\x10\x02*K\n\x04Mode\x12\x14\n\x10MODE_UNSPECIFIED\x10\x00\x12\x0e\n\nMODE_ALERT\x10\x01\x12\r\n\tMODE_HITL\x10\x02\x12\x0e\n\nMODE_BLOCK\x10\x03*`\n\rVerdictStatus\x12\x1e\n\x1aVERDICT_STATUS_UNSPECIFIED\x10\x00\x12\x15\n\x11VERDICT_STATUS_OK\x10\x01\x12\x18\n\x14VERDICT_STATUS_ERROR\x10\x02\x62\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -50,18 +50,20 @@ _globals['_PAIREDEVENT'].fields_by_name['pair_type']._serialized_options = b'\272H\003\310\001\001' _globals['_MCPSERVER'].fields_by_name['name']._loaded_options = None _globals['_MCPSERVER'].fields_by_name['name']._serialized_options = b'\272H\004r\002\020\001' + _globals['_INSTALLEDPLUGIN'].fields_by_name['name']._loaded_options = None + _globals['_INSTALLEDPLUGIN'].fields_by_name['name']._serialized_options = b'\272H\004r\002\020\001' _globals['_SESSIONLOGIN'].fields_by_name['session_id']._loaded_options = None _globals['_SESSIONLOGIN'].fields_by_name['session_id']._serialized_options = b'\272H\004r\002\020\001' _globals['_VERDICT'].fields_by_name['event_id']._loaded_options = None _globals['_VERDICT'].fields_by_name['event_id']._serialized_options = b'\272H\004r\002\020\001' _globals['_VERDICT'].fields_by_name['session_id']._loaded_options = None _globals['_VERDICT'].fields_by_name['session_id']._serialized_options = b'\272H\004r\002\020\001' - _globals['_PAIRTYPE']._serialized_start=2409 - _globals['_PAIRTYPE']._serialized_end=2485 - _globals['_MODE']._serialized_start=2487 - _globals['_MODE']._serialized_end=2562 - _globals['_VERDICTSTATUS']._serialized_start=2564 - _globals['_VERDICTSTATUS']._serialized_end=2660 + _globals['_PAIRTYPE']._serialized_start=2796 + _globals['_PAIRTYPE']._serialized_end=2872 + _globals['_MODE']._serialized_start=2874 + _globals['_MODE']._serialized_end=2949 + _globals['_VERDICTSTATUS']._serialized_start=2951 + _globals['_VERDICTSTATUS']._serialized_end=3047 _globals['_CHATMESSAGE']._serialized_start=64 _globals['_CHATMESSAGE']._serialized_end=108 _globals['_TOOLCALL']._serialized_start=110 @@ -75,27 +77,31 @@ _globals['_TOOLPAIRDATA']._serialized_start=562 _globals['_TOOLPAIRDATA']._serialized_end=657 _globals['_PAIREDEVENT']._serialized_start=660 - _globals['_PAIREDEVENT']._serialized_end=1111 - _globals['_PAIREDEVENTBATCH']._serialized_start=1113 - _globals['_PAIREDEVENTBATCH']._serialized_end=1180 - _globals['_MCPSERVER']._serialized_start=1182 - _globals['_MCPSERVER']._serialized_end=1253 - _globals['_MCPINVENTORY']._serialized_start=1255 - _globals['_MCPINVENTORY']._serialized_end=1317 - _globals['_LLMSTACK']._serialized_start=1319 - _globals['_LLMSTACK']._serialized_end=1362 - _globals['_SESSIONLOGIN']._serialized_start=1365 - _globals['_SESSIONLOGIN']._serialized_end=1515 - _globals['_CLIENTFRAME']._serialized_start=1518 - _globals['_CLIENTFRAME']._serialized_end=1725 - _globals['_POLICYSNAPSHOT']._serialized_start=1728 - _globals['_POLICYSNAPSHOT']._serialized_end=1901 - _globals['_HITLRESPONSE']._serialized_start=1903 - _globals['_HITLRESPONSE']._serialized_end=1945 - _globals['_LOGINACK']._serialized_start=1947 - _globals['_LOGINACK']._serialized_end=2009 - _globals['_SERVERFRAME']._serialized_start=2011 - _globals['_SERVERFRAME']._serialized_end=2132 - _globals['_VERDICT']._serialized_start=2135 - _globals['_VERDICT']._serialized_end=2407 + _globals['_PAIREDEVENT']._serialized_end=1134 + _globals['_PAIREDEVENTBATCH']._serialized_start=1136 + _globals['_PAIREDEVENTBATCH']._serialized_end=1203 + _globals['_MCPSERVER']._serialized_start=1206 + _globals['_MCPSERVER']._serialized_end=1366 + _globals['_MCPINVENTORY']._serialized_start=1368 + _globals['_MCPINVENTORY']._serialized_end=1430 + _globals['_INSTALLEDPLUGIN']._serialized_start=1432 + _globals['_INSTALLEDPLUGIN']._serialized_end=1527 + _globals['_PLUGININVENTORY']._serialized_start=1529 + _globals['_PLUGININVENTORY']._serialized_end=1600 + _globals['_LLMSTACK']._serialized_start=1602 + _globals['_LLMSTACK']._serialized_end=1645 + _globals['_SESSIONLOGIN']._serialized_start=1648 + _globals['_SESSIONLOGIN']._serialized_end=1821 + _globals['_CLIENTFRAME']._serialized_start=1824 + _globals['_CLIENTFRAME']._serialized_end=2096 + _globals['_POLICYSNAPSHOT']._serialized_start=2099 + _globals['_POLICYSNAPSHOT']._serialized_end=2272 + _globals['_HITLRESPONSE']._serialized_start=2274 + _globals['_HITLRESPONSE']._serialized_end=2316 + _globals['_LOGINACK']._serialized_start=2318 + _globals['_LOGINACK']._serialized_end=2396 + _globals['_SERVERFRAME']._serialized_start=2398 + _globals['_SERVERFRAME']._serialized_end=2519 + _globals['_VERDICT']._serialized_start=2522 + _globals['_VERDICT']._serialized_end=2794 # @@protoc_insertion_point(module_scope) diff --git a/sdk/python/adrian/proto/event_pb2.pyi b/sdk/python/adrian/proto/event_pb2.pyi index 1094342..24f7be8 100644 --- a/sdk/python/adrian/proto/event_pb2.pyi +++ b/sdk/python/adrian/proto/event_pb2.pyi @@ -1,767 +1,270 @@ -""" -@generated by mypy-protobuf. Do not edit manually! -isort:skip_file -""" - -from collections import abc as _abc -from google.protobuf import descriptor as _descriptor -from google.protobuf import message as _message +from buf.validate import validate_pb2 as _validate_pb2 from google.protobuf.internal import containers as _containers from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper -import builtins as _builtins -import sys -import typing as _typing - -if sys.version_info >= (3, 11): - from typing import TypeAlias as _TypeAlias, Never as _Never -else: - from typing_extensions import TypeAlias as _TypeAlias, Never as _Never +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from collections.abc import Iterable as _Iterable, Mapping as _Mapping +from typing import ClassVar as _ClassVar, Optional as _Optional, Union as _Union DESCRIPTOR: _descriptor.FileDescriptor -class _PairType: - ValueType = _typing.NewType("ValueType", _builtins.int) - V: _TypeAlias = ValueType # noqa: Y015 - -class _PairTypeEnumTypeWrapper(_enum_type_wrapper._EnumTypeWrapper[_PairType.ValueType], _builtins.type): - DESCRIPTOR: _descriptor.EnumDescriptor - PAIR_TYPE_UNSPECIFIED: _PairType.ValueType # 0 - PAIR_TYPE_LLM: _PairType.ValueType # 1 - """LLM pair: chat_model_start + llm_end folded together.""" - PAIR_TYPE_TOOL: _PairType.ValueType # 2 - """Tool pair: tool_start + tool_end folded together.""" - -class PairType(_PairType, metaclass=_PairTypeEnumTypeWrapper): - """PairType identifies what an event pair represents.""" - -PAIR_TYPE_UNSPECIFIED: PairType.ValueType # 0 -PAIR_TYPE_LLM: PairType.ValueType # 1 -"""LLM pair: chat_model_start + llm_end folded together.""" -PAIR_TYPE_TOOL: PairType.ValueType # 2 -"""Tool pair: tool_start + tool_end folded together.""" -Global___PairType: _TypeAlias = PairType # noqa: Y015 - -class _Mode: - ValueType = _typing.NewType("ValueType", _builtins.int) - V: _TypeAlias = ValueType # noqa: Y015 - -class _ModeEnumTypeWrapper(_enum_type_wrapper._EnumTypeWrapper[_Mode.ValueType], _builtins.type): - DESCRIPTOR: _descriptor.EnumDescriptor - MODE_UNSPECIFIED: _Mode.ValueType # 0 - MODE_ALERT: _Mode.ValueType # 1 - """Verdicts are dashboard-only, the server never pushes them to the SDK.""" - MODE_HITL: _Mode.ValueType # 2 - """Server holds in-policy-scope verdicts pending dashboard approve/reject. - Out-of-scope verdicts are forwarded immediately so the SDK doesn't stall. - """ - MODE_BLOCK: _Mode.ValueType # 3 - """Server forwards every verdict; the SDK enforces per the policy snapshot.""" +class PairType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): + __slots__ = () + PAIR_TYPE_UNSPECIFIED: _ClassVar[PairType] + PAIR_TYPE_LLM: _ClassVar[PairType] + PAIR_TYPE_TOOL: _ClassVar[PairType] + +class Mode(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): + __slots__ = () + MODE_UNSPECIFIED: _ClassVar[Mode] + MODE_ALERT: _ClassVar[Mode] + MODE_HITL: _ClassVar[Mode] + MODE_BLOCK: _ClassVar[Mode] + +class VerdictStatus(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): + __slots__ = () + VERDICT_STATUS_UNSPECIFIED: _ClassVar[VerdictStatus] + VERDICT_STATUS_OK: _ClassVar[VerdictStatus] + VERDICT_STATUS_ERROR: _ClassVar[VerdictStatus] +PAIR_TYPE_UNSPECIFIED: PairType +PAIR_TYPE_LLM: PairType +PAIR_TYPE_TOOL: PairType +MODE_UNSPECIFIED: Mode +MODE_ALERT: Mode +MODE_HITL: Mode +MODE_BLOCK: Mode +VERDICT_STATUS_UNSPECIFIED: VerdictStatus +VERDICT_STATUS_OK: VerdictStatus +VERDICT_STATUS_ERROR: VerdictStatus -class Mode(_Mode, metaclass=_ModeEnumTypeWrapper): - """Mode of execution governs how the server handles verdicts for a session. - Set via policies.mode and delivered to the SDK via PolicySnapshot - on every Verdict. - """ - -MODE_UNSPECIFIED: Mode.ValueType # 0 -MODE_ALERT: Mode.ValueType # 1 -"""Verdicts are dashboard-only, the server never pushes them to the SDK.""" -MODE_HITL: Mode.ValueType # 2 -"""Server holds in-policy-scope verdicts pending dashboard approve/reject. -Out-of-scope verdicts are forwarded immediately so the SDK doesn't stall. -""" -MODE_BLOCK: Mode.ValueType # 3 -"""Server forwards every verdict; the SDK enforces per the policy snapshot.""" -Global___Mode: _TypeAlias = Mode # noqa: Y015 - -class _VerdictStatus: - ValueType = _typing.NewType("ValueType", _builtins.int) - V: _TypeAlias = ValueType # noqa: Y015 - -class _VerdictStatusEnumTypeWrapper(_enum_type_wrapper._EnumTypeWrapper[_VerdictStatus.ValueType], _builtins.type): - DESCRIPTOR: _descriptor.EnumDescriptor - VERDICT_STATUS_UNSPECIFIED: _VerdictStatus.ValueType # 0 - VERDICT_STATUS_OK: _VerdictStatus.ValueType # 1 - VERDICT_STATUS_ERROR: _VerdictStatus.ValueType # 2 - -class VerdictStatus(_VerdictStatus, metaclass=_VerdictStatusEnumTypeWrapper): - """VerdictStatus says whether a Verdict came from a completed classifier - decision or represents a classifier failure. ERROR verdicts carry no - classifier-produced MAD code; policy decides whether they fail open - or fail closed. - """ - -VERDICT_STATUS_UNSPECIFIED: VerdictStatus.ValueType # 0 -VERDICT_STATUS_OK: VerdictStatus.ValueType # 1 -VERDICT_STATUS_ERROR: VerdictStatus.ValueType # 2 -Global___VerdictStatus: _TypeAlias = VerdictStatus # noqa: Y015 - -@_typing.final class ChatMessage(_message.Message): - """ChatMessage represents a conversation message with a string role.""" + __slots__ = ("role", "content") + ROLE_FIELD_NUMBER: _ClassVar[int] + CONTENT_FIELD_NUMBER: _ClassVar[int] + role: str + content: str + def __init__(self, role: _Optional[str] = ..., content: _Optional[str] = ...) -> None: ... - DESCRIPTOR: _descriptor.Descriptor - - ROLE_FIELD_NUMBER: _builtins.int - CONTENT_FIELD_NUMBER: _builtins.int - role: _builtins.str - """Role of this message (e.g. "system", "human", "ai", "tool").""" - content: _builtins.str - """Text content of the message.""" - def __init__( - self, - *, - role: _builtins.str = ..., - content: _builtins.str = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _Never # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["content", b"content", "role", b"role"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - def WhichOneof(self, oneof_group: _Never) -> None: ... - -Global___ChatMessage: _TypeAlias = ChatMessage # noqa: Y015 - -@_typing.final class ToolCall(_message.Message): - """ToolCall represents a tool invocation the model decided to make.""" - - DESCRIPTOR: _descriptor.Descriptor - - NAME_FIELD_NUMBER: _builtins.int - ARGS_FIELD_NUMBER: _builtins.int - ID_FIELD_NUMBER: _builtins.int - name: _builtins.str - """Tool name.""" - args: _builtins.str - """JSON-encoded arguments object.""" - id: _builtins.str - """Provider-assigned tool call ID (may be empty for non-LLM-originated tools).""" - def __init__( - self, - *, - name: _builtins.str = ..., - args: _builtins.str = ..., - id: _builtins.str = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _Never # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["args", b"args", "id", b"id", "name", b"name"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - def WhichOneof(self, oneof_group: _Never) -> None: ... + __slots__ = ("name", "args", "id") + NAME_FIELD_NUMBER: _ClassVar[int] + ARGS_FIELD_NUMBER: _ClassVar[int] + ID_FIELD_NUMBER: _ClassVar[int] + name: str + args: str + id: str + def __init__(self, name: _Optional[str] = ..., args: _Optional[str] = ..., id: _Optional[str] = ...) -> None: ... -Global___ToolCall: _TypeAlias = ToolCall # noqa: Y015 - -@_typing.final class TokenUsage(_message.Message): - """TokenUsage holds token consumption counters.""" - - DESCRIPTOR: _descriptor.Descriptor - - PROMPT_TOKENS_FIELD_NUMBER: _builtins.int - COMPLETION_TOKENS_FIELD_NUMBER: _builtins.int - TOTAL_TOKENS_FIELD_NUMBER: _builtins.int - prompt_tokens: _builtins.int - """Number of tokens in the prompt.""" - completion_tokens: _builtins.int - """Number of tokens in the completion.""" - total_tokens: _builtins.int - """Total tokens consumed.""" - def __init__( - self, - *, - prompt_tokens: _builtins.int = ..., - completion_tokens: _builtins.int = ..., - total_tokens: _builtins.int = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _Never # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["completion_tokens", b"completion_tokens", "prompt_tokens", b"prompt_tokens", "total_tokens", b"total_tokens"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - def WhichOneof(self, oneof_group: _Never) -> None: ... + __slots__ = ("prompt_tokens", "completion_tokens", "total_tokens") + PROMPT_TOKENS_FIELD_NUMBER: _ClassVar[int] + COMPLETION_TOKENS_FIELD_NUMBER: _ClassVar[int] + TOTAL_TOKENS_FIELD_NUMBER: _ClassVar[int] + prompt_tokens: int + completion_tokens: int + total_tokens: int + def __init__(self, prompt_tokens: _Optional[int] = ..., completion_tokens: _Optional[int] = ..., total_tokens: _Optional[int] = ...) -> None: ... -Global___TokenUsage: _TypeAlias = TokenUsage # noqa: Y015 - -@_typing.final class AgentContext(_message.Message): - """AgentContext carries the identity and prompts for an agent at a given - point in the call stack. An empty agent_id signals "no agent" and is - used as the parent sentinel for top-level events. - """ - - DESCRIPTOR: _descriptor.Descriptor - - AGENT_ID_FIELD_NUMBER: _builtins.int - SYSTEM_PROMPT_FIELD_NUMBER: _builtins.int - USER_INSTRUCTION_FIELD_NUMBER: _builtins.int - agent_id: _builtins.str - """Opaque agent identifier derived from the framework's checkpoint - namespace (e.g. "tools|reason"). - """ - system_prompt: _builtins.str - """System prompt the agent is running under.""" - user_instruction: _builtins.str - """User instruction driving the current invocation.""" - def __init__( - self, - *, - agent_id: _builtins.str = ..., - system_prompt: _builtins.str = ..., - user_instruction: _builtins.str = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _Never # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["agent_id", b"agent_id", "system_prompt", b"system_prompt", "user_instruction", b"user_instruction"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - def WhichOneof(self, oneof_group: _Never) -> None: ... + __slots__ = ("agent_id", "system_prompt", "user_instruction") + AGENT_ID_FIELD_NUMBER: _ClassVar[int] + SYSTEM_PROMPT_FIELD_NUMBER: _ClassVar[int] + USER_INSTRUCTION_FIELD_NUMBER: _ClassVar[int] + agent_id: str + system_prompt: str + user_instruction: str + def __init__(self, agent_id: _Optional[str] = ..., system_prompt: _Optional[str] = ..., user_instruction: _Optional[str] = ...) -> None: ... -Global___AgentContext: _TypeAlias = AgentContext # noqa: Y015 - -@_typing.final class LlmPairData(_message.Message): - """LlmPairData is the payload for PAIR_TYPE_LLM events: one chat_model_start - folded with its matching llm_end. - """ - - DESCRIPTOR: _descriptor.Descriptor - - MODEL_FIELD_NUMBER: _builtins.int - MESSAGES_FIELD_NUMBER: _builtins.int - OUTPUT_FIELD_NUMBER: _builtins.int - TOOL_CALLS_FIELD_NUMBER: _builtins.int - USAGE_FIELD_NUMBER: _builtins.int - model: _builtins.str - """Model class name (e.g. "ChatAnthropic", "ChatOpenAI").""" - output: _builtins.str - """The model's text response.""" - @_builtins.property - def messages(self) -> _containers.RepeatedCompositeFieldContainer[Global___ChatMessage]: - """Ordered conversation messages sent to the model.""" - - @_builtins.property - def tool_calls(self) -> _containers.RepeatedCompositeFieldContainer[Global___ToolCall]: - """Tool calls the model decided to make.""" - - @_builtins.property - def usage(self) -> Global___TokenUsage: - """Token usage counters.""" + __slots__ = ("model", "messages", "output", "tool_calls", "usage") + MODEL_FIELD_NUMBER: _ClassVar[int] + MESSAGES_FIELD_NUMBER: _ClassVar[int] + OUTPUT_FIELD_NUMBER: _ClassVar[int] + TOOL_CALLS_FIELD_NUMBER: _ClassVar[int] + USAGE_FIELD_NUMBER: _ClassVar[int] + model: str + messages: _containers.RepeatedCompositeFieldContainer[ChatMessage] + output: str + tool_calls: _containers.RepeatedCompositeFieldContainer[ToolCall] + usage: TokenUsage + def __init__(self, model: _Optional[str] = ..., messages: _Optional[_Iterable[_Union[ChatMessage, _Mapping]]] = ..., output: _Optional[str] = ..., tool_calls: _Optional[_Iterable[_Union[ToolCall, _Mapping]]] = ..., usage: _Optional[_Union[TokenUsage, _Mapping]] = ...) -> None: ... - def __init__( - self, - *, - model: _builtins.str = ..., - messages: _abc.Iterable[Global___ChatMessage] | None = ..., - output: _builtins.str = ..., - tool_calls: _abc.Iterable[Global___ToolCall] | None = ..., - usage: Global___TokenUsage | None = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _typing.Literal["usage", b"usage"] # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["messages", b"messages", "model", b"model", "output", b"output", "tool_calls", b"tool_calls", "usage", b"usage"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - def WhichOneof(self, oneof_group: _Never) -> None: ... - -Global___LlmPairData: _TypeAlias = LlmPairData # noqa: Y015 - -@_typing.final class ToolPairData(_message.Message): - """ToolPairData is the payload for PAIR_TYPE_TOOL events: one tool_start - folded with its matching tool_end. - """ - - DESCRIPTOR: _descriptor.Descriptor + __slots__ = ("tool_name", "tool_call_id", "input", "output") + TOOL_NAME_FIELD_NUMBER: _ClassVar[int] + TOOL_CALL_ID_FIELD_NUMBER: _ClassVar[int] + INPUT_FIELD_NUMBER: _ClassVar[int] + OUTPUT_FIELD_NUMBER: _ClassVar[int] + tool_name: str + tool_call_id: str + input: str + output: str + def __init__(self, tool_name: _Optional[str] = ..., tool_call_id: _Optional[str] = ..., input: _Optional[str] = ..., output: _Optional[str] = ...) -> None: ... - TOOL_NAME_FIELD_NUMBER: _builtins.int - TOOL_CALL_ID_FIELD_NUMBER: _builtins.int - INPUT_FIELD_NUMBER: _builtins.int - OUTPUT_FIELD_NUMBER: _builtins.int - tool_name: _builtins.str - """Name of the tool that was invoked.""" - tool_call_id: _builtins.str - """ID linking this invocation to the model's tool call. May be empty - for tools invoked outside an LLM flow. - """ - input: _builtins.str - """JSON-encoded input passed to the tool.""" - output: _builtins.str - """Tool execution output.""" - def __init__( - self, - *, - tool_name: _builtins.str = ..., - tool_call_id: _builtins.str = ..., - input: _builtins.str = ..., - output: _builtins.str = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _Never # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["input", b"input", "output", b"output", "tool_call_id", b"tool_call_id", "tool_name", b"tool_name"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - def WhichOneof(self, oneof_group: _Never) -> None: ... - -Global___ToolPairData: _TypeAlias = ToolPairData # noqa: Y015 - -@_typing.final class PairedEvent(_message.Message): - """PairedEvent is a completed event pair with agent context attached.""" - - DESCRIPTOR: _descriptor.Descriptor - - EVENT_ID_FIELD_NUMBER: _builtins.int - INVOCATION_ID_FIELD_NUMBER: _builtins.int - SESSION_ID_FIELD_NUMBER: _builtins.int - RUN_ID_FIELD_NUMBER: _builtins.int - PARENT_RUN_ID_FIELD_NUMBER: _builtins.int - TIMESTAMP_FIELD_NUMBER: _builtins.int - PAIR_TYPE_FIELD_NUMBER: _builtins.int - AGENT_FIELD_NUMBER: _builtins.int - PARENT_FIELD_NUMBER: _builtins.int - LLM_FIELD_NUMBER: _builtins.int - TOOL_FIELD_NUMBER: _builtins.int - METADATA_JSON_FIELD_NUMBER: _builtins.int - SOURCE_FIELD_NUMBER: _builtins.int - event_id: _builtins.str - """Unique event identifier.""" - invocation_id: _builtins.str - """Groups all events for one user-prompted invocation of the top-level graph.""" - session_id: _builtins.str - """Session identifier.""" - run_id: _builtins.str - """Framework run ID for this pair (LLM run_id for LLM pairs, tool run_id for tool pairs).""" - parent_run_id: _builtins.str - """Parent run ID. For tool pairs this points at the producing LLM's run_id - and is the key used for block-mode verdict correlation. - """ - timestamp: _builtins.str - """ISO 8601 timestamp of the pair's completion (end event).""" - pair_type: Global___PairType.ValueType - """Whether this is an LLM pair or a tool pair.""" - metadata_json: _builtins.bytes - """Escape hatch for framework metadata not modelled as first-class fields - (e.g. LangGraph checkpoint_ns, arbitrary tags). - """ - source: _builtins.str - @_builtins.property - def agent(self) -> Global___AgentContext: - """The agent that produced this pair.""" - - @_builtins.property - def parent(self) -> Global___AgentContext: - """The agent that delegated to this one. agent_id == "" means top-level.""" - - @_builtins.property - def llm(self) -> Global___LlmPairData: - """Payload for PAIR_TYPE_LLM.""" - - @_builtins.property - def tool(self) -> Global___ToolPairData: - """Payload for PAIR_TYPE_TOOL.""" - - def __init__( - self, - *, - event_id: _builtins.str = ..., - invocation_id: _builtins.str = ..., - session_id: _builtins.str = ..., - run_id: _builtins.str = ..., - parent_run_id: _builtins.str = ..., - timestamp: _builtins.str = ..., - pair_type: Global___PairType.ValueType = ..., - agent: Global___AgentContext | None = ..., - parent: Global___AgentContext | None = ..., - llm: Global___LlmPairData | None = ..., - tool: Global___ToolPairData | None = ..., - metadata_json: _builtins.bytes = ..., - source: _builtins.str = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _typing.Literal["agent", b"agent", "data", b"data", "llm", b"llm", "parent", b"parent", "tool", b"tool"] # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["agent", b"agent", "data", b"data", "event_id", b"event_id", "invocation_id", b"invocation_id", "llm", b"llm", "metadata_json", b"metadata_json", "pair_type", b"pair_type", "parent", b"parent", "parent_run_id", b"parent_run_id", "run_id", b"run_id", "session_id", b"session_id", "source", b"source", "timestamp", b"timestamp", "tool", b"tool"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - _WhichOneofReturnType_data: _TypeAlias = _typing.Literal["llm", "tool"] # noqa: Y015 - _WhichOneofArgType_data: _TypeAlias = _typing.Literal["data", b"data"] # noqa: Y015 - def WhichOneof(self, oneof_group: _WhichOneofArgType_data) -> _WhichOneofReturnType_data | None: ... + __slots__ = ("event_id", "invocation_id", "session_id", "run_id", "parent_run_id", "timestamp", "pair_type", "agent", "parent", "llm", "tool", "connection_id", "metadata_json", "source") + EVENT_ID_FIELD_NUMBER: _ClassVar[int] + INVOCATION_ID_FIELD_NUMBER: _ClassVar[int] + SESSION_ID_FIELD_NUMBER: _ClassVar[int] + RUN_ID_FIELD_NUMBER: _ClassVar[int] + PARENT_RUN_ID_FIELD_NUMBER: _ClassVar[int] + TIMESTAMP_FIELD_NUMBER: _ClassVar[int] + PAIR_TYPE_FIELD_NUMBER: _ClassVar[int] + AGENT_FIELD_NUMBER: _ClassVar[int] + PARENT_FIELD_NUMBER: _ClassVar[int] + LLM_FIELD_NUMBER: _ClassVar[int] + TOOL_FIELD_NUMBER: _ClassVar[int] + CONNECTION_ID_FIELD_NUMBER: _ClassVar[int] + METADATA_JSON_FIELD_NUMBER: _ClassVar[int] + SOURCE_FIELD_NUMBER: _ClassVar[int] + event_id: str + invocation_id: str + session_id: str + run_id: str + parent_run_id: str + timestamp: str + pair_type: PairType + agent: AgentContext + parent: AgentContext + llm: LlmPairData + tool: ToolPairData + connection_id: str + metadata_json: bytes + source: str + def __init__(self, event_id: _Optional[str] = ..., invocation_id: _Optional[str] = ..., session_id: _Optional[str] = ..., run_id: _Optional[str] = ..., parent_run_id: _Optional[str] = ..., timestamp: _Optional[str] = ..., pair_type: _Optional[_Union[PairType, str]] = ..., agent: _Optional[_Union[AgentContext, _Mapping]] = ..., parent: _Optional[_Union[AgentContext, _Mapping]] = ..., llm: _Optional[_Union[LlmPairData, _Mapping]] = ..., tool: _Optional[_Union[ToolPairData, _Mapping]] = ..., connection_id: _Optional[str] = ..., metadata_json: _Optional[bytes] = ..., source: _Optional[str] = ...) -> None: ... -Global___PairedEvent: _TypeAlias = PairedEvent # noqa: Y015 - -@_typing.final class PairedEventBatch(_message.Message): - """PairedEventBatch wraps multiple paired events for batched WebSocket frames.""" - - DESCRIPTOR: _descriptor.Descriptor - - EVENTS_FIELD_NUMBER: _builtins.int - @_builtins.property - def events(self) -> _containers.RepeatedCompositeFieldContainer[Global___PairedEvent]: ... - def __init__( - self, - *, - events: _abc.Iterable[Global___PairedEvent] | None = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _Never # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["events", b"events"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - def WhichOneof(self, oneof_group: _Never) -> None: ... + __slots__ = ("events",) + EVENTS_FIELD_NUMBER: _ClassVar[int] + events: _containers.RepeatedCompositeFieldContainer[PairedEvent] + def __init__(self, events: _Optional[_Iterable[_Union[PairedEvent, _Mapping]]] = ...) -> None: ... -Global___PairedEventBatch: _TypeAlias = PairedEventBatch # noqa: Y015 - -@_typing.final class McpServer(_message.Message): - """McpServer is one MCP server the SDK observed at startup. Mirrors the - SDK-side McpServer dataclass. - """ - - DESCRIPTOR: _descriptor.Descriptor - - NAME_FIELD_NUMBER: _builtins.int - TRANSPORT_FIELD_NUMBER: _builtins.int - ENDPOINT_FIELD_NUMBER: _builtins.int - name: _builtins.str - """Server identifier. Adapter-layer captures use the connection key - from MultiServerMCPClient.connections; raw-transport captures - synthesise ":". - """ - transport: _builtins.str - """One of "stdio", "sse", "streamable_http", "websocket", or "unknown".""" - endpoint: _builtins.str - """URL for SSE/HTTP/WebSocket transports, or the joined command line - for stdio. Empty string when neither is available. - """ - def __init__( - self, - *, - name: _builtins.str = ..., - transport: _builtins.str = ..., - endpoint: _builtins.str = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _Never # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["endpoint", b"endpoint", "name", b"name", "transport", b"transport"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - def WhichOneof(self, oneof_group: _Never) -> None: ... + __slots__ = ("name", "transport", "endpoint", "version", "protocol_version", "server_info_name", "tools_json") + NAME_FIELD_NUMBER: _ClassVar[int] + TRANSPORT_FIELD_NUMBER: _ClassVar[int] + ENDPOINT_FIELD_NUMBER: _ClassVar[int] + VERSION_FIELD_NUMBER: _ClassVar[int] + PROTOCOL_VERSION_FIELD_NUMBER: _ClassVar[int] + SERVER_INFO_NAME_FIELD_NUMBER: _ClassVar[int] + TOOLS_JSON_FIELD_NUMBER: _ClassVar[int] + name: str + transport: str + endpoint: str + version: str + protocol_version: str + server_info_name: str + tools_json: str + def __init__(self, name: _Optional[str] = ..., transport: _Optional[str] = ..., endpoint: _Optional[str] = ..., version: _Optional[str] = ..., protocol_version: _Optional[str] = ..., server_info_name: _Optional[str] = ..., tools_json: _Optional[str] = ...) -> None: ... -Global___McpServer: _TypeAlias = McpServer # noqa: Y015 - -@_typing.final class McpInventory(_message.Message): - """McpInventory is the SDK's one-shot dump of its connected MCP servers. - Sent once per WS session right after the SDK login. The server replaces - the session's entire MCP-server list on receipt, no diff/merge. - """ - - DESCRIPTOR: _descriptor.Descriptor - - SERVERS_FIELD_NUMBER: _builtins.int - @_builtins.property - def servers(self) -> _containers.RepeatedCompositeFieldContainer[Global___McpServer]: ... - def __init__( - self, - *, - servers: _abc.Iterable[Global___McpServer] | None = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _Never # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["servers", b"servers"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - def WhichOneof(self, oneof_group: _Never) -> None: ... + __slots__ = ("servers",) + SERVERS_FIELD_NUMBER: _ClassVar[int] + servers: _containers.RepeatedCompositeFieldContainer[McpServer] + def __init__(self, servers: _Optional[_Iterable[_Union[McpServer, _Mapping]]] = ...) -> None: ... + +class InstalledPlugin(_message.Message): + __slots__ = ("name", "enabled", "version", "marketplace") + NAME_FIELD_NUMBER: _ClassVar[int] + ENABLED_FIELD_NUMBER: _ClassVar[int] + VERSION_FIELD_NUMBER: _ClassVar[int] + MARKETPLACE_FIELD_NUMBER: _ClassVar[int] + name: str + enabled: bool + version: str + marketplace: str + def __init__(self, name: _Optional[str] = ..., enabled: _Optional[bool] = ..., version: _Optional[str] = ..., marketplace: _Optional[str] = ...) -> None: ... + +class PluginInventory(_message.Message): + __slots__ = ("plugins",) + PLUGINS_FIELD_NUMBER: _ClassVar[int] + plugins: _containers.RepeatedCompositeFieldContainer[InstalledPlugin] + def __init__(self, plugins: _Optional[_Iterable[_Union[InstalledPlugin, _Mapping]]] = ...) -> None: ... -Global___McpInventory: _TypeAlias = McpInventory # noqa: Y015 - -@_typing.final class LLMStack(_message.Message): - """LLMStack identifies the LLM provider and model for a session.""" - - DESCRIPTOR: _descriptor.Descriptor + __slots__ = ("provider", "model") + PROVIDER_FIELD_NUMBER: _ClassVar[int] + MODEL_FIELD_NUMBER: _ClassVar[int] + provider: str + model: str + def __init__(self, provider: _Optional[str] = ..., model: _Optional[str] = ...) -> None: ... - PROVIDER_FIELD_NUMBER: _builtins.int - MODEL_FIELD_NUMBER: _builtins.int - provider: _builtins.str - """Provider name (e.g. "anthropic", "openai").""" - model: _builtins.str - """Model identifier (e.g. "claude-sonnet-4-20250514").""" - def __init__( - self, - *, - provider: _builtins.str = ..., - model: _builtins.str = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _Never # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["model", b"model", "provider", b"provider"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - def WhichOneof(self, oneof_group: _Never) -> None: ... - -Global___LLMStack: _TypeAlias = LLMStack # noqa: Y015 - -@_typing.final class SessionLogin(_message.Message): - """SessionLogin is sent as the first frame after WebSocket upgrade.""" - - DESCRIPTOR: _descriptor.Descriptor - - SESSION_ID_FIELD_NUMBER: _builtins.int - LLM_STACK_FIELD_NUMBER: _builtins.int - SCHEMA_VERSION_FIELD_NUMBER: _builtins.int - SOURCE_FIELD_NUMBER: _builtins.int - session_id: _builtins.str - """Session identifier.""" - schema_version: _builtins.int - """Wire schema version the client is speaking. Server rejects unknown values.""" - source: _builtins.str - """SDK / integration that produced this session (e.g. "claude-code"); the - server branches HITL behavior on it. Empty for legacy SDKs. - """ - @_builtins.property - def llm_stack(self) -> Global___LLMStack: - """LLM stack information for this session.""" - - def __init__( - self, - *, - session_id: _builtins.str = ..., - llm_stack: Global___LLMStack | None = ..., - schema_version: _builtins.int = ..., - source: _builtins.str = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _typing.Literal["llm_stack", b"llm_stack"] # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["llm_stack", b"llm_stack", "schema_version", b"schema_version", "session_id", b"session_id", "source", b"source"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - def WhichOneof(self, oneof_group: _Never) -> None: ... - -Global___SessionLogin: _TypeAlias = SessionLogin # noqa: Y015 + __slots__ = ("session_id", "llm_stack", "schema_version", "source", "connection_id") + SESSION_ID_FIELD_NUMBER: _ClassVar[int] + LLM_STACK_FIELD_NUMBER: _ClassVar[int] + SCHEMA_VERSION_FIELD_NUMBER: _ClassVar[int] + SOURCE_FIELD_NUMBER: _ClassVar[int] + CONNECTION_ID_FIELD_NUMBER: _ClassVar[int] + session_id: str + llm_stack: LLMStack + schema_version: int + source: str + connection_id: str + def __init__(self, session_id: _Optional[str] = ..., llm_stack: _Optional[_Union[LLMStack, _Mapping]] = ..., schema_version: _Optional[int] = ..., source: _Optional[str] = ..., connection_id: _Optional[str] = ...) -> None: ... -@_typing.final class ClientFrame(_message.Message): - """ClientFrame is the top-level envelope for all client WebSocket messages.""" - - DESCRIPTOR: _descriptor.Descriptor - - LOGIN_FIELD_NUMBER: _builtins.int - PAIRED_BATCH_FIELD_NUMBER: _builtins.int - MCP_INVENTORY_FIELD_NUMBER: _builtins.int - @_builtins.property - def login(self) -> Global___SessionLogin: - """Session login (must be first frame).""" - - @_builtins.property - def paired_batch(self) -> Global___PairedEventBatch: - """Paired event batch (requires prior login).""" - - @_builtins.property - def mcp_inventory(self) -> Global___McpInventory: - """One-shot MCP-server inventory dump from the SDK after login.""" - - def __init__( - self, - *, - login: Global___SessionLogin | None = ..., - paired_batch: Global___PairedEventBatch | None = ..., - mcp_inventory: Global___McpInventory | None = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _typing.Literal["frame", b"frame", "login", b"login", "mcp_inventory", b"mcp_inventory", "paired_batch", b"paired_batch"] # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["frame", b"frame", "login", b"login", "mcp_inventory", b"mcp_inventory", "paired_batch", b"paired_batch"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - _WhichOneofReturnType_frame: _TypeAlias = _typing.Literal["login", "paired_batch", "mcp_inventory"] # noqa: Y015 - _WhichOneofArgType_frame: _TypeAlias = _typing.Literal["frame", b"frame"] # noqa: Y015 - def WhichOneof(self, oneof_group: _WhichOneofArgType_frame) -> _WhichOneofReturnType_frame | None: ... + __slots__ = ("login", "paired_batch", "mcp_inventory", "plugin_inventory") + LOGIN_FIELD_NUMBER: _ClassVar[int] + PAIRED_BATCH_FIELD_NUMBER: _ClassVar[int] + MCP_INVENTORY_FIELD_NUMBER: _ClassVar[int] + PLUGIN_INVENTORY_FIELD_NUMBER: _ClassVar[int] + login: SessionLogin + paired_batch: PairedEventBatch + mcp_inventory: McpInventory + plugin_inventory: PluginInventory + def __init__(self, login: _Optional[_Union[SessionLogin, _Mapping]] = ..., paired_batch: _Optional[_Union[PairedEventBatch, _Mapping]] = ..., mcp_inventory: _Optional[_Union[McpInventory, _Mapping]] = ..., plugin_inventory: _Optional[_Union[PluginInventory, _Mapping]] = ...) -> None: ... -Global___ClientFrame: _TypeAlias = ClientFrame # noqa: Y015 - -@_typing.final class PolicySnapshot(_message.Message): - """PolicySnapshot is the org's effective execution-mode policy at the moment - a verdict was decided. Attached by the server to every Verdict it sends - so the SDK can apply user-configured behaviour (halt vs continue, - review-pending vs proceed) without an extra round-trip to fetch policy. - - Per-MAD-code booleans say whether the active mode's behaviour fires on - that code. False means "treat this code as silent regardless of mode". - fail_closed_on_classifier_error controls ERROR verdicts and BLOCK-mode - SDK verdict timeouts. The default false value preserves fail-open - availability when talking to older backends. - """ - - DESCRIPTOR: _descriptor.Descriptor + __slots__ = ("mode", "policy_m0", "policy_m2", "policy_m3", "policy_m4", "fail_closed_on_classifier_error") + MODE_FIELD_NUMBER: _ClassVar[int] + POLICY_M0_FIELD_NUMBER: _ClassVar[int] + POLICY_M2_FIELD_NUMBER: _ClassVar[int] + POLICY_M3_FIELD_NUMBER: _ClassVar[int] + POLICY_M4_FIELD_NUMBER: _ClassVar[int] + FAIL_CLOSED_ON_CLASSIFIER_ERROR_FIELD_NUMBER: _ClassVar[int] + mode: Mode + policy_m0: bool + policy_m2: bool + policy_m3: bool + policy_m4: bool + fail_closed_on_classifier_error: bool + def __init__(self, mode: _Optional[_Union[Mode, str]] = ..., policy_m0: _Optional[bool] = ..., policy_m2: _Optional[bool] = ..., policy_m3: _Optional[bool] = ..., policy_m4: _Optional[bool] = ..., fail_closed_on_classifier_error: _Optional[bool] = ...) -> None: ... - MODE_FIELD_NUMBER: _builtins.int - POLICY_M0_FIELD_NUMBER: _builtins.int - POLICY_M2_FIELD_NUMBER: _builtins.int - POLICY_M3_FIELD_NUMBER: _builtins.int - POLICY_M4_FIELD_NUMBER: _builtins.int - FAIL_CLOSED_ON_CLASSIFIER_ERROR_FIELD_NUMBER: _builtins.int - mode: Global___Mode.ValueType - policy_m0: _builtins.bool - policy_m2: _builtins.bool - policy_m3: _builtins.bool - policy_m4: _builtins.bool - fail_closed_on_classifier_error: _builtins.bool - def __init__( - self, - *, - mode: Global___Mode.ValueType = ..., - policy_m0: _builtins.bool = ..., - policy_m2: _builtins.bool = ..., - policy_m3: _builtins.bool = ..., - policy_m4: _builtins.bool = ..., - fail_closed_on_classifier_error: _builtins.bool = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _Never # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["fail_closed_on_classifier_error", b"fail_closed_on_classifier_error", "mode", b"mode", "policy_m0", b"policy_m0", "policy_m2", b"policy_m2", "policy_m3", b"policy_m3", "policy_m4", b"policy_m4"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - def WhichOneof(self, oneof_group: _Never) -> None: ... - -Global___PolicySnapshot: _TypeAlias = PolicySnapshot # noqa: Y015 - -@_typing.final class HitlResponse(_message.Message): - """HitlResponse rides on a Verdict that has been resolved through the - human-in-the-loop review queue. Absent on regular (non-HITL or - out-of-scope) verdicts. - """ - - DESCRIPTOR: _descriptor.Descriptor + __slots__ = ("continue_execution",) + CONTINUE_EXECUTION_FIELD_NUMBER: _ClassVar[int] + continue_execution: bool + def __init__(self, continue_execution: _Optional[bool] = ...) -> None: ... - CONTINUE_EXECUTION_FIELD_NUMBER: _builtins.int - continue_execution: _builtins.bool - """True if the reviewer approved (agent should continue), false if rejected.""" - def __init__( - self, - *, - continue_execution: _builtins.bool = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _Never # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["continue_execution", b"continue_execution"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - def WhichOneof(self, oneof_group: _Never) -> None: ... - -Global___HitlResponse: _TypeAlias = HitlResponse # noqa: Y015 - -@_typing.final class LoginAck(_message.Message): - """LoginAck is the server's first frame after a successful SessionLogin. - Carries the org's effective policy so the SDK can apply correct mode + - timeout behaviour from the very first event, before any verdict round-trip. - - Policy snapshot is captured at WS login time. Mid-session policy changes - do NOT propagate to existing sessions, reconnect for new policy to take - effect. - """ - - DESCRIPTOR: _descriptor.Descriptor - - POLICY_FIELD_NUMBER: _builtins.int - @_builtins.property - def policy(self) -> Global___PolicySnapshot: ... - def __init__( - self, - *, - policy: Global___PolicySnapshot | None = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _typing.Literal["policy", b"policy"] # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["policy", b"policy"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - def WhichOneof(self, oneof_group: _Never) -> None: ... + __slots__ = ("policy", "source") + POLICY_FIELD_NUMBER: _ClassVar[int] + SOURCE_FIELD_NUMBER: _ClassVar[int] + policy: PolicySnapshot + source: str + def __init__(self, policy: _Optional[_Union[PolicySnapshot, _Mapping]] = ..., source: _Optional[str] = ...) -> None: ... -Global___LoginAck: _TypeAlias = LoginAck # noqa: Y015 - -@_typing.final class ServerFrame(_message.Message): - """ServerFrame is the top-level envelope for all server→SDK WebSocket - messages. Multiplexes login-ack, verdicts, and any future server-pushed - frame types over the same WebSocket without ambiguity. - """ - - DESCRIPTOR: _descriptor.Descriptor - - LOGIN_ACK_FIELD_NUMBER: _builtins.int - VERDICT_FIELD_NUMBER: _builtins.int - @_builtins.property - def login_ack(self) -> Global___LoginAck: - """Sent exactly once, immediately after a successful SessionLogin.""" - - @_builtins.property - def verdict(self) -> Global___Verdict: - """Sent for every classified event the server forwards to the SDK - (subject to mode/policy gates). - """ - - def __init__( - self, - *, - login_ack: Global___LoginAck | None = ..., - verdict: Global___Verdict | None = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _typing.Literal["frame", b"frame", "login_ack", b"login_ack", "verdict", b"verdict"] # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["frame", b"frame", "login_ack", b"login_ack", "verdict", b"verdict"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - _WhichOneofReturnType_frame: _TypeAlias = _typing.Literal["login_ack", "verdict"] # noqa: Y015 - _WhichOneofArgType_frame: _TypeAlias = _typing.Literal["frame", b"frame"] # noqa: Y015 - def WhichOneof(self, oneof_group: _WhichOneofArgType_frame) -> _WhichOneofReturnType_frame | None: ... - -Global___ServerFrame: _TypeAlias = ServerFrame # noqa: Y015 + __slots__ = ("login_ack", "verdict") + LOGIN_ACK_FIELD_NUMBER: _ClassVar[int] + VERDICT_FIELD_NUMBER: _ClassVar[int] + login_ack: LoginAck + verdict: Verdict + def __init__(self, login_ack: _Optional[_Union[LoginAck, _Mapping]] = ..., verdict: _Optional[_Union[Verdict, _Mapping]] = ...) -> None: ... -@_typing.final class Verdict(_message.Message): - """Verdict is the classification result for a PairedEvent.""" - - DESCRIPTOR: _descriptor.Descriptor - - EVENT_ID_FIELD_NUMBER: _builtins.int - SESSION_ID_FIELD_NUMBER: _builtins.int - MAD_CODE_FIELD_NUMBER: _builtins.int - POLICY_FIELD_NUMBER: _builtins.int - HITL_FIELD_NUMBER: _builtins.int - STATUS_FIELD_NUMBER: _builtins.int - event_id: _builtins.str - """The event_id of the PairedEvent being classified.""" - session_id: _builtins.str - """Session identifier for routing.""" - mad_code: _builtins.str - """MAD code the classifier returned (e.g. "M0", "M2_C", "M4_a"). - Empty string means no MAD code was produced, such as for a - VerdictStatus.ERROR classifier failure. Benign classifier success - is represented by status OK with mad_code "M0". - """ - status: Global___VerdictStatus.ValueType - """Status of the classifier result. OK means mad_code carries a normal - classifier decision. ERROR means classification did not complete and - mad_code is empty; fail-open/fail-closed behaviour comes from policy. - """ - @_builtins.property - def policy(self) -> Global___PolicySnapshot: - """Org's effective execution-mode policy at the time of this verdict. - Always populated by the server; SDK reads this to decide whether to - halt, pause, or continue. - """ - - @_builtins.property - def hitl(self) -> Global___HitlResponse: - """Present only when this verdict represents a human-in-the-loop review - resolution (approve or reject from the dashboard). Absent on auto- - classified verdicts and on out-of-scope verdicts forwarded immediately. - """ - - def __init__( - self, - *, - event_id: _builtins.str = ..., - session_id: _builtins.str = ..., - mad_code: _builtins.str = ..., - policy: Global___PolicySnapshot | None = ..., - hitl: Global___HitlResponse | None = ..., - status: Global___VerdictStatus.ValueType = ..., - ) -> None: ... - _HasFieldArgType: _TypeAlias = _typing.Literal["hitl", b"hitl", "policy", b"policy"] # noqa: Y015 - def HasField(self, field_name: _HasFieldArgType) -> _builtins.bool: ... - _ClearFieldArgType: _TypeAlias = _typing.Literal["event_id", b"event_id", "hitl", b"hitl", "mad_code", b"mad_code", "policy", b"policy", "session_id", b"session_id", "status", b"status"] # noqa: Y015 - def ClearField(self, field_name: _ClearFieldArgType) -> None: ... - def WhichOneof(self, oneof_group: _Never) -> None: ... - -Global___Verdict: _TypeAlias = Verdict # noqa: Y015 + __slots__ = ("event_id", "session_id", "mad_code", "policy", "hitl", "status") + EVENT_ID_FIELD_NUMBER: _ClassVar[int] + SESSION_ID_FIELD_NUMBER: _ClassVar[int] + MAD_CODE_FIELD_NUMBER: _ClassVar[int] + POLICY_FIELD_NUMBER: _ClassVar[int] + HITL_FIELD_NUMBER: _ClassVar[int] + STATUS_FIELD_NUMBER: _ClassVar[int] + event_id: str + session_id: str + mad_code: str + policy: PolicySnapshot + hitl: HitlResponse + status: VerdictStatus + def __init__(self, event_id: _Optional[str] = ..., session_id: _Optional[str] = ..., mad_code: _Optional[str] = ..., policy: _Optional[_Union[PolicySnapshot, _Mapping]] = ..., hitl: _Optional[_Union[HitlResponse, _Mapping]] = ..., status: _Optional[_Union[VerdictStatus, str]] = ...) -> None: ... diff --git a/sdk/python/adrian/types.py b/sdk/python/adrian/types.py index b934a17..dfa67af 100644 --- a/sdk/python/adrian/types.py +++ b/sdk/python/adrian/types.py @@ -269,3 +269,7 @@ class McpServer: name: str transport: str endpoint: str + version: str = "" + protocol_version: str = "" + server_info_name: str = "" + tools_json: str = "" diff --git a/sdk/python/adrian/ws.py b/sdk/python/adrian/ws.py index 2462807..bd822ec 100644 --- a/sdk/python/adrian/ws.py +++ b/sdk/python/adrian/ws.py @@ -257,6 +257,7 @@ def __init__( # verdict and how long. self._mode: int = pb.MODE_UNSPECIFIED self._policy: pb.PolicySnapshot | None = None + self._blocked_mcp_servers: set[str] = set() # Set the first time a ``ServerFrame{login_ack}`` is applied. # Used in two places: # 1. ``on_paired_event`` defensively pre-registers a @@ -337,6 +338,10 @@ def __init__( # -- Mode / policy state (populated by LoginAck) -- + def is_mcp_blocked(self, server_name: str) -> bool: + """Check if an MCP server is blocked for this agent profile.""" + return server_name in self._blocked_mcp_servers + def policy_active(self) -> bool: """Whether the active server mode requires waiting on verdicts. @@ -691,6 +696,9 @@ async def _recv_loop(self) -> None: self._on_login_ack(frame.login_ack) elif kind == "verdict": await self._on_verdict_frame(frame.verdict) + elif kind == "mcp_block_update": + self._blocked_mcp_servers = set(frame.mcp_block_update.blocked_mcp_servers) + logger.info("MCP block list updated: %s", self._blocked_mcp_servers) else: logger.warning( "ignoring unknown ServerFrame kind %r " @@ -724,6 +732,9 @@ def _on_login_ack(self, ack: pb.LoginAck) -> None: """ self._mode = ack.policy.mode self._policy = ack.policy + self._blocked_mcp_servers = set( + ack.blocked_mcp_servers if hasattr(ack, "blocked_mcp_servers") else [] + ) self._login_ack_received.set() logger.info( "LoginAck received: mode=%s policy_m0=%s policy_m2=%s " diff --git a/sdk/python/tests/test_mixed_verdicts.py b/sdk/python/tests/test_mixed_verdicts.py new file mode 100644 index 0000000..ccfa9de --- /dev/null +++ b/sdk/python/tests/test_mixed_verdicts.py @@ -0,0 +1,345 @@ +"""Tests for mixed benign/malicious tool calls from a single LLM message. + +Scenario: An LLM emits multiple tool_calls in one AIMessage. Some are +classified as M0 (benign) and some as M4 (malicious). Only the malicious +ones should be blocked; benign ones must execute normally. +""" + +# pyright: reportPrivateUsage=false + +from __future__ import annotations + +import asyncio +from collections.abc import Iterator +from pathlib import Path +from typing import Any, cast + +import adrian +import pytest +from adrian.proto import event_pb2 as pb +from adrian.ws import WebSocketClient +from langchain_core.messages import AIMessage +from langchain_core.runnables.config import RunnableConfig, ensure_config +from langgraph._internal._constants import CONF, CONFIG_KEY_RUNTIME +from langgraph.prebuilt import ToolNode +from langgraph.runtime import Runtime + + +def _runtime_config() -> RunnableConfig: + return ensure_config({CONF: {CONFIG_KEY_RUNTIME: Runtime()}}) + + +def _apply_mode( + ws: WebSocketClient, + mode: int, + *, + policy_m0: bool = False, + policy_m2: bool = False, + policy_m3: bool = False, + policy_m4: bool = False, +) -> pb.PolicySnapshot: + policy = pb.PolicySnapshot( + mode=cast("pb.Mode", mode), + policy_m0=policy_m0, + policy_m2=policy_m2, + policy_m3=policy_m3, + policy_m4=policy_m4, + ) + ws._mode = mode + ws._policy = policy + ws._login_ack_received.set() + return policy + + +@pytest.fixture(autouse=True) +def _cleanup() -> Iterator[None]: + yield + adrian.shutdown() + + +def _init_sdk(tmp_path: Path, block_timeout: float = 2.0) -> WebSocketClient: + adrian.init( + api_key="test-key", + log_file=str(tmp_path / "events.jsonl"), + auto_instrument=True, + ws_url="ws://x", + block_timeout=block_timeout, + ) + ws = adrian._ws_client + assert ws is not None + return ws + + +class TestMixedBenignMaliciousToolCalls: + """Single LLM message emits multiple tool_calls with different verdicts.""" + + async def test_three_tools_one_blocked_two_allowed(self, tmp_path: Path) -> None: + """LLM emits 3 tool_calls: read_file (M0), search_web (M2), + delete_data (M4). Only delete_data should be blocked.""" + executed: dict[str, str] = {} + + def read_file(path: str) -> str: + """Read a file.""" + executed["read_file"] = path + return f"contents of {path}" + + def search_web(query: str) -> str: + """Search the web.""" + executed["search_web"] = query + return f"results for {query}" + + def delete_data(target: str) -> str: + """Delete data — dangerous operation.""" + executed["delete_data"] = target + return f"deleted {target}" + + ws = _init_sdk(tmp_path) + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._connected.set() + ws._loop = asyncio.get_running_loop() + + # Each tool_call maps to a different LLM event with a different verdict + ws._tool_call_id_to_event_id["tc-read"] = "llm-read" + ws._tool_call_id_to_event_id["tc-search"] = "llm-search" + ws._tool_call_id_to_event_id["tc-delete"] = "llm-delete" + + # M0 — benign, policy_m0 is False so NOT in scope → allow + fut_read = ws.register_pending("llm-read") + fut_read.set_result(pb.Verdict( + event_id="llm-read", mad_code="M0_benign", policy=policy + )) + + # M2 — policy_m2 is False so NOT in scope → allow + fut_search = ws.register_pending("llm-search") + fut_search.set_result(pb.Verdict( + event_id="llm-search", mad_code="M2_misuse", policy=policy + )) + + # M4 — policy_m4 is True so IN scope → BLOCK + fut_delete = ws.register_pending("llm-delete") + fut_delete.set_result(pb.Verdict( + event_id="llm-delete", mad_code="M4_exfiltration", policy=policy + )) + + # Dispatch each tool independently (as ToolNode does) + for tc_id, tool_name, tool_fn, args_key, args_val in [ + ("tc-read", "read_file", read_file, "path", "/etc/hosts"), + ("tc-search", "search_web", search_web, "query", "python docs"), + ("tc-delete", "delete_data", delete_data, "target", "user_data"), + ]: + ai = AIMessage( + content="", + tool_calls=[{"id": tc_id, "name": tool_name, "args": {args_key: args_val}}], + ) + tn = ToolNode([tool_fn]) + await tn.ainvoke({"messages": [ai]}, config=_runtime_config()) + + # Benign tools should have executed + assert executed.get("read_file") == "/etc/hosts" + assert executed.get("search_web") == "python docs" + # Malicious tool should NOT have executed + assert "delete_data" not in executed + + async def test_all_m4_all_blocked(self, tmp_path: Path) -> None: + """All tool_calls classified as M4 — all should be blocked.""" + executed: dict[str, bool] = {} + + def tool_a(x: str) -> str: + """Tool A.""" + executed["a"] = True + return x + + def tool_b(x: str) -> str: + """Tool B.""" + executed["b"] = True + return x + + def tool_c(x: str) -> str: + """Tool C.""" + executed["c"] = True + return x + + ws = _init_sdk(tmp_path) + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._connected.set() + ws._loop = asyncio.get_running_loop() + + for tc_id, evt_id in [("tc-a", "llm-a"), ("tc-b", "llm-b"), ("tc-c", "llm-c")]: + ws._tool_call_id_to_event_id[tc_id] = evt_id + fut = ws.register_pending(evt_id) + fut.set_result(pb.Verdict(event_id=evt_id, mad_code="M4_attack", policy=policy)) + + for tc_id, name, fn in [("tc-a", "tool_a", tool_a), ("tc-b", "tool_b", tool_b), ("tc-c", "tool_c", tool_c)]: + ai = AIMessage(content="", tool_calls=[{"id": tc_id, "name": name, "args": {"x": "y"}}]) + result = await ToolNode([fn]).ainvoke({"messages": [ai]}, config=_runtime_config()) + assert "BLOCKED" in result["messages"][0].content + + assert not executed + + async def test_all_benign_all_run(self, tmp_path: Path) -> None: + """All tool_calls classified as M0 — all should run.""" + executed: dict[str, bool] = {} + + def tool_a(x: str) -> str: + """Tool A.""" + executed["a"] = True + return x + + def tool_b(x: str) -> str: + """Tool B.""" + executed["b"] = True + return x + + ws = _init_sdk(tmp_path) + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) # only m4 blocked + ws._connected.set() + ws._loop = asyncio.get_running_loop() + + for tc_id, evt_id in [("tc-a", "llm-a"), ("tc-b", "llm-b")]: + ws._tool_call_id_to_event_id[tc_id] = evt_id + fut = ws.register_pending(evt_id) + fut.set_result(pb.Verdict(event_id=evt_id, mad_code="M0_ok", policy=policy)) + + for tc_id, name, fn in [("tc-a", "tool_a", tool_a), ("tc-b", "tool_b", tool_b)]: + ai = AIMessage(content="", tool_calls=[{"id": tc_id, "name": name, "args": {"x": "y"}}]) + await ToolNode([fn]).ainvoke({"messages": [ai]}, config=_runtime_config()) + + assert executed == {"a": True, "b": True} + + async def test_m2_in_scope_blocked_m0_allowed(self, tmp_path: Path) -> None: + """policy_m2=True: M2 tools are blocked, M0 tools run.""" + executed: dict[str, bool] = {} + + def benign_tool(x: str) -> str: + """Benign.""" + executed["benign"] = True + return x + + def suspicious_tool(x: str) -> str: + """Suspicious.""" + executed["suspicious"] = True + return x + + ws = _init_sdk(tmp_path) + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m2=True, policy_m4=True) + ws._connected.set() + ws._loop = asyncio.get_running_loop() + + ws._tool_call_id_to_event_id["tc-benign"] = "llm-benign" + ws._tool_call_id_to_event_id["tc-suspicious"] = "llm-suspicious" + + fut_b = ws.register_pending("llm-benign") + fut_b.set_result(pb.Verdict(event_id="llm-benign", mad_code="M0_ok", policy=policy)) + + fut_s = ws.register_pending("llm-suspicious") + fut_s.set_result(pb.Verdict(event_id="llm-suspicious", mad_code="M2_misuse", policy=policy)) + + # Benign runs + ai_b = AIMessage(content="", tool_calls=[{"id": "tc-benign", "name": "benign_tool", "args": {"x": "y"}}]) + await ToolNode([benign_tool]).ainvoke({"messages": [ai_b]}, config=_runtime_config()) + assert executed.get("benign") is True + + # Suspicious blocked + ai_s = AIMessage(content="", tool_calls=[{"id": "tc-suspicious", "name": "suspicious_tool", "args": {"x": "y"}}]) + result = await ToolNode([suspicious_tool]).ainvoke({"messages": [ai_s]}, config=_runtime_config()) + assert "suspicious" not in executed + assert "BLOCKED" in result["messages"][0].content + + async def test_sync_tools_mixed_verdicts_from_worker_thread(self, tmp_path: Path) -> None: + """Same as test_three_tools but with SYNC tools (worker thread path). + This is the exact scenario from run 019eda10.""" + executed: dict[str, str] = {} + + def safe_read(path: str) -> str: + """Safe read tool.""" + executed["safe_read"] = path + return f"contents: {path}" + + def dangerous_write(path: str) -> str: + """Dangerous write tool.""" + executed["dangerous_write"] = path + return f"wrote: {path}" + + def safe_list(directory: str) -> str: + """Safe list tool.""" + executed["safe_list"] = directory + return f"files in {directory}" + + ws = _init_sdk(tmp_path) + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._connected.set() + ws._loop = asyncio.get_running_loop() + + ws._tool_call_id_to_event_id["tc-sread"] = "llm-sread" + ws._tool_call_id_to_event_id["tc-dwrite"] = "llm-dwrite" + ws._tool_call_id_to_event_id["tc-slist"] = "llm-slist" + + # M0 — allowed + fut1 = ws.register_pending("llm-sread") + fut1.set_result(pb.Verdict(event_id="llm-sread", mad_code="M0_ok", policy=policy)) + # M4 — blocked + fut2 = ws.register_pending("llm-dwrite") + fut2.set_result(pb.Verdict(event_id="llm-dwrite", mad_code="M4_data_exfil", policy=policy)) + # M0 — allowed + fut3 = ws.register_pending("llm-slist") + fut3.set_result(pb.Verdict(event_id="llm-slist", mad_code="M0_ok", policy=policy)) + + # safe_read (M0) → should run + ai1 = AIMessage(content="", tool_calls=[{"id": "tc-sread", "name": "safe_read", "args": {"path": "/tmp/ok"}}]) + await ToolNode([safe_read]).ainvoke({"messages": [ai1]}, config=_runtime_config()) + + # dangerous_write (M4) → should block + ai2 = AIMessage(content="", tool_calls=[{"id": "tc-dwrite", "name": "dangerous_write", "args": {"path": "/etc/shadow"}}]) + result2 = await ToolNode([dangerous_write]).ainvoke({"messages": [ai2]}, config=_runtime_config()) + + # safe_list (M0) → should run + ai3 = AIMessage(content="", tool_calls=[{"id": "tc-slist", "name": "safe_list", "args": {"directory": "/home"}}]) + await ToolNode([safe_list]).ainvoke({"messages": [ai3]}, config=_runtime_config()) + + assert executed.get("safe_read") == "/tmp/ok", "M0 safe_read should have run" + assert "dangerous_write" not in executed, "M4 dangerous_write should be BLOCKED" + assert executed.get("safe_list") == "/home", "M0 safe_list should have run" + assert "BLOCKED" in result2["messages"][0].content + + async def test_hitl_approve_benign_reject_malicious(self, tmp_path: Path) -> None: + """HITL mode: human approves a benign tool, rejects a malicious one.""" + executed: dict[str, bool] = {} + + async def benign_tool(x: str) -> str: + """Benign tool.""" + executed["benign"] = True + return x + + async def malicious_tool(x: str) -> str: + """Malicious tool.""" + executed["malicious"] = True + return x + + ws = _init_sdk(tmp_path) + policy = _apply_mode(ws, pb.MODE_HITL, policy_m4=True) + ws._connected.set() + + # Benign: human approves + ws._tool_call_id_to_event_id["tc-good"] = "llm-good" + fut_good = ws.register_pending("llm-good") + v_good = pb.Verdict(event_id="llm-good", mad_code="M4_a", policy=policy) + v_good.hitl.continue_execution = True + fut_good.set_result(v_good) + + # Malicious: human rejects + ws._tool_call_id_to_event_id["tc-bad"] = "llm-bad" + fut_bad = ws.register_pending("llm-bad") + v_bad = pb.Verdict(event_id="llm-bad", mad_code="M4_a", policy=policy) + v_bad.hitl.continue_execution = False + fut_bad.set_result(v_bad) + + # Benign runs + ai_good = AIMessage(content="", tool_calls=[{"id": "tc-good", "name": "benign_tool", "args": {"x": "y"}}]) + await ToolNode([benign_tool]).ainvoke({"messages": [ai_good]}, config=_runtime_config()) + assert executed.get("benign") is True + + # Malicious blocked + ai_bad = AIMessage(content="", tool_calls=[{"id": "tc-bad", "name": "malicious_tool", "args": {"x": "y"}}]) + result = await ToolNode([malicious_tool]).ainvoke({"messages": [ai_bad]}, config=_runtime_config()) + assert "malicious" not in executed + assert "BLOCKED" in result["messages"][0].content diff --git a/sdk/python/tests/test_sync_gate_edge_cases.py b/sdk/python/tests/test_sync_gate_edge_cases.py new file mode 100644 index 0000000..03237a4 --- /dev/null +++ b/sdk/python/tests/test_sync_gate_edge_cases.py @@ -0,0 +1,657 @@ +"""Edge-case and adversarial tests for the verdict gate system. + +Covers scenarios the happy-path tests miss: +- Worker threads with no event loop (bare ThreadPoolExecutor) +- ws._loop stopped or None mid-gate +- _ws_client becomes None between check and use (disconnect race) +- Tools without tool_call_id (bypass vector) +- Concurrent tool calls hitting _sync_gate simultaneously +- HITL mode timeout behavior (must NOT fail-open) +- Verdict arriving before gate checks (resolved future replay) +- LRU eviction of tool_call_id map entries +- register_pending called from worker thread (wrong loop) +""" + +# pyright: reportPrivateUsage=false + +from __future__ import annotations + +import asyncio +import concurrent.futures +import threading +from collections.abc import Iterator +from pathlib import Path +from typing import Any, cast +from unittest.mock import AsyncMock, MagicMock, patch + +import adrian +import pytest +from adrian.proto import event_pb2 as pb +from adrian.ws import WebSocketClient +from langchain_core.messages import AIMessage +from langchain_core.runnables.config import RunnableConfig, ensure_config +from langchain_core.tools import BaseTool +from langgraph._internal._constants import CONF, CONFIG_KEY_RUNTIME +from langgraph.prebuilt import ToolNode +from langgraph.runtime import Runtime + + +def _runtime_config() -> RunnableConfig: + return ensure_config({CONF: {CONFIG_KEY_RUNTIME: Runtime()}}) + + +def _apply_mode( + ws: WebSocketClient, + mode: int, + *, + policy_m0: bool = False, + policy_m2: bool = False, + policy_m3: bool = False, + policy_m4: bool = False, +) -> pb.PolicySnapshot: + policy = pb.PolicySnapshot( + mode=cast("pb.Mode", mode), + policy_m0=policy_m0, + policy_m2=policy_m2, + policy_m3=policy_m3, + policy_m4=policy_m4, + ) + ws._mode = mode + ws._policy = policy + ws._login_ack_received.set() + return policy + + +@pytest.fixture(autouse=True) +def _cleanup() -> Iterator[None]: + yield + adrian.shutdown() + + +def _init_sdk(tmp_path: Path, block_timeout: float = 2.0) -> WebSocketClient: + adrian.init( + api_key="test-key", + log_file=str(tmp_path / "events.jsonl"), + auto_instrument=True, + ws_url="ws://x", + block_timeout=block_timeout, + ) + ws = adrian._ws_client + assert ws is not None + return ws + + +def _tool_state(tc_id: str, tool_name: str, args: dict[str, Any] | None = None) -> dict[str, Any]: + ai = AIMessage( + content="", + tool_calls=[{"id": tc_id, "name": tool_name, "args": args or {"x": "hi"}}], + ) + return {"messages": [ai]} + + +# --------------------------------------------------------------------------- +# 1. Worker thread without ANY event loop +# --------------------------------------------------------------------------- + + +class TestBareWorkerThread: + """Simulate Pregel's ThreadPoolExecutor dispatch — no event loop on + the worker thread at all.""" + + async def test_sync_tool_blocks_from_bare_thread(self, tmp_path: Path) -> None: + """A sync tool.invoke called from a bare thread (no loop set) with + M4 verdict must block via run_coroutine_threadsafe to ws._loop.""" + tool_ran = False + + def my_tool(x: str) -> str: + """Tool stub.""" + nonlocal tool_ran + tool_ran = True + return x + + ws = _init_sdk(tmp_path) + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._connected.set() + ws._loop = asyncio.get_running_loop() + ws._tool_call_id_to_event_id["tc-bare"] = "llm-bare" + + fut = ws.register_pending("llm-bare") + fut.set_result(pb.Verdict(event_id="llm-bare", mad_code="M4_a", policy=policy)) + + tool_node = ToolNode([my_tool]) + result = await tool_node.ainvoke( + _tool_state("tc-bare", "my_tool"), config=_runtime_config() + ) + + assert not tool_ran + assert "BLOCKED" in result["messages"][0].content + + async def test_sync_tool_m3_blocks_when_m3_in_scope(self, tmp_path: Path) -> None: + """M3 verdict with policy_m3=True must also block.""" + tool_ran = False + + def my_tool(x: str) -> str: + """Tool stub.""" + nonlocal tool_ran + tool_ran = True + return x + + ws = _init_sdk(tmp_path) + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m3=True) + ws._connected.set() + ws._loop = asyncio.get_running_loop() + ws._tool_call_id_to_event_id["tc-m3"] = "llm-m3" + + fut = ws.register_pending("llm-m3") + fut.set_result(pb.Verdict(event_id="llm-m3", mad_code="M3_risk", policy=policy)) + + tool_node = ToolNode([my_tool]) + result = await tool_node.ainvoke( + _tool_state("tc-m3", "my_tool"), config=_runtime_config() + ) + + assert not tool_ran + assert "BLOCKED" in result["messages"][0].content + + +# --------------------------------------------------------------------------- +# 2. ws._loop is None or stopped +# --------------------------------------------------------------------------- + + +class TestWsLoopEdgeCases: + async def test_ws_loop_is_none_sync_gate_fails_closed(self, tmp_path: Path) -> None: + """If ws._loop is None (WS never connected), sync gate should + fall through to asyncio.run() path — which will fail-closed if + the async gate can't resolve.""" + tool_ran = False + + def my_tool(x: str) -> str: + """Tool stub.""" + nonlocal tool_ran + tool_ran = True + return x + + ws = _init_sdk(tmp_path, block_timeout=0.1) + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._connected.set() + # Intentionally NOT setting ws._loop — simulates WS not connected + ws._loop = None + ws._tool_call_id_to_event_id["tc-noloop"] = "llm-noloop" + + tool_node = ToolNode([my_tool]) + result = await tool_node.ainvoke( + _tool_state("tc-noloop", "my_tool"), config=_runtime_config() + ) + + # With ws._loop=None, _sync_gate falls to asyncio.run() path. + # asyncio.run() creates a new loop, runs _async_gate, which will + # timeout waiting for verdict → fail-closed → BLOCKED. + # OR it may fail because register_pending needs a running loop. + # Either way, tool should NOT run (fail-closed). + # Actually — the pure-sync path may raise because register_pending + # calls asyncio.get_running_loop(). The except catches it and + # returns True (fail-closed). So tool shouldn't run. + assert not tool_ran + + async def test_ws_loop_stopped_sync_gate_fails_closed(self, tmp_path: Path) -> None: + """If ws._loop exists but is_running() is False, sync gate + should fail-closed (can't bridge coroutine to a stopped loop).""" + tool_ran = False + + def my_tool(x: str) -> str: + """Tool stub.""" + nonlocal tool_ran + tool_ran = True + return x + + ws = _init_sdk(tmp_path, block_timeout=0.1) + _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._connected.set() + + # Use a MagicMock loop that reports is_running()=False + mock_loop = MagicMock() + mock_loop.is_running.return_value = False + ws._loop = mock_loop + + ws._tool_call_id_to_event_id["tc-stopped"] = "llm-stopped" + + tool_node = ToolNode([my_tool]) + result = await tool_node.ainvoke( + _tool_state("tc-stopped", "my_tool"), config=_runtime_config() + ) + + # ws._loop.is_running() is False → _sync_gate falls to asyncio.run() + # path → fails → returns True (fail-closed) → tool blocked + assert not tool_ran + + +# --------------------------------------------------------------------------- +# 3. _ws_client becomes None mid-gate (disconnect race) +# --------------------------------------------------------------------------- + + +class TestDisconnectRace: + async def test_ws_client_nulled_between_check_and_gate(self, tmp_path: Path) -> None: + """If _ws_client is set to None after _sync_gate captures it but + before the gate completes, it should not crash. The local `ws` + reference in _sync_gate protects against this.""" + tool_ran = False + + def my_tool(x: str) -> str: + """Tool stub.""" + nonlocal tool_ran + tool_ran = True + return x + + ws = _init_sdk(tmp_path) + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._connected.set() + ws._loop = asyncio.get_running_loop() + ws._tool_call_id_to_event_id["tc-disc"] = "llm-disc" + + fut = ws.register_pending("llm-disc") + fut.set_result(pb.Verdict(event_id="llm-disc", mad_code="M4_a", policy=policy)) + + # Null out _ws_client after init but before tool dispatch + # The gate captures `ws = _ws_client` at the top, so this should + # be safe — the local reference keeps the object alive. + tool_node = ToolNode([my_tool]) + + # This should still work because _sync_gate captures ws locally + result = await tool_node.ainvoke( + _tool_state("tc-disc", "my_tool"), config=_runtime_config() + ) + + assert not tool_ran + assert "BLOCKED" in result["messages"][0].content + + +# --------------------------------------------------------------------------- +# 4. Tools without tool_call_id (bypass vector) +# --------------------------------------------------------------------------- + + +class TestMissingToolCallId: + async def test_tool_without_tool_call_id_runs_ungated(self, tmp_path: Path) -> None: + """A tool invoked with a plain dict (not a ToolCall) should run + without waiting for a verdict — no tool_call_id to gate on.""" + tool_ran = False + + def my_tool(x: str) -> str: + """Tool stub.""" + nonlocal tool_ran + tool_ran = True + return x + + ws = _init_sdk(tmp_path) + _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._connected.set() + + from langchain_core.tools import StructuredTool + tool = StructuredTool.from_function(my_tool) + # Plain dict input — not a ToolCall, no "type": "tool_call" + result = tool.invoke({"x": "hello"}) + + assert tool_ran + + +# --------------------------------------------------------------------------- +# 5. Concurrent tool calls hitting _sync_gate simultaneously +# --------------------------------------------------------------------------- + + +class TestConcurrentGateAccess: + async def test_parallel_tool_calls_each_get_own_verdict(self, tmp_path: Path) -> None: + """Two parallel tool calls with different tool_call_ids must each + wait for their own verdict independently.""" + results: dict[str, bool] = {} + + def tool_a(x: str) -> str: + """Tool A.""" + results["a"] = True + return x + + def tool_b(x: str) -> str: + """Tool B.""" + results["b"] = True + return x + + ws = _init_sdk(tmp_path) + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._connected.set() + ws._loop = asyncio.get_running_loop() + + # tool_a → M4 (block), tool_b → M0 (allow) + ws._tool_call_id_to_event_id["tc-a"] = "llm-a" + ws._tool_call_id_to_event_id["tc-b"] = "llm-b" + + fut_a = ws.register_pending("llm-a") + fut_a.set_result(pb.Verdict(event_id="llm-a", mad_code="M4_exfil", policy=policy)) + + fut_b = ws.register_pending("llm-b") + fut_b.set_result(pb.Verdict(event_id="llm-b", mad_code="M0_benign", policy=policy)) + + # Dispatch tool_a (should block) + tool_node_a = ToolNode([tool_a]) + result_a = await tool_node_a.ainvoke( + _tool_state("tc-a", "tool_a"), config=_runtime_config() + ) + + # Dispatch tool_b (should allow) + tool_node_b = ToolNode([tool_b]) + result_b = await tool_node_b.ainvoke( + _tool_state("tc-b", "tool_b"), config=_runtime_config() + ) + + assert "a" not in results, "M4 tool_a should have been blocked" + assert results.get("b") is True, "M0 tool_b should have run" + assert "BLOCKED" in result_a["messages"][0].content + + +# --------------------------------------------------------------------------- +# 6. HITL mode — must hold indefinitely, never fail-open +# --------------------------------------------------------------------------- + + +class TestHitlModeHold: + async def test_hitl_holds_past_block_timeout(self, tmp_path: Path) -> None: + """In MODE_HITL, the gate must wait indefinitely for a human decision. + It must NOT fail-open after block_timeout elapses.""" + tool_ran = False + + async def my_tool(x: str) -> str: + """Tool stub.""" + nonlocal tool_ran + tool_ran = True + return x + + ws = _init_sdk(tmp_path, block_timeout=0.2) + policy = _apply_mode(ws, pb.MODE_HITL, policy_m4=True) + ws._connected.set() + ws._tool_call_id_to_event_id["tc-hitl"] = "llm-hitl" + fut = ws.register_pending("llm-hitl") + + tool_node = ToolNode([my_tool]) + task = asyncio.ensure_future( + tool_node.ainvoke( + _tool_state("tc-hitl", "my_tool"), config=_runtime_config() + ) + ) + + # Wait well past block_timeout (0.2s) — tool should still be held + await asyncio.sleep(0.5) + assert not task.done(), "HITL must hold indefinitely, not fail-open after block_timeout" + assert not tool_ran + + # Human approves + verdict = pb.Verdict(event_id="llm-hitl", mad_code="M4_a", policy=policy) + verdict.hitl.continue_execution = True + fut.set_result(verdict) + + await asyncio.wait_for(task, timeout=2.0) + assert tool_ran + + async def test_hitl_reject_blocks(self, tmp_path: Path) -> None: + """HITL reject (continue_execution=False) blocks the tool.""" + tool_ran = False + + async def my_tool(x: str) -> str: + """Tool stub.""" + nonlocal tool_ran + tool_ran = True + return x + + ws = _init_sdk(tmp_path) + policy = _apply_mode(ws, pb.MODE_HITL, policy_m4=True) + ws._connected.set() + ws._tool_call_id_to_event_id["tc-hitl-rej"] = "llm-hitl-rej" + fut = ws.register_pending("llm-hitl-rej") + + tool_node = ToolNode([my_tool]) + task = asyncio.ensure_future( + tool_node.ainvoke( + _tool_state("tc-hitl-rej", "my_tool"), config=_runtime_config() + ) + ) + + await asyncio.sleep(0.1) + verdict = pb.Verdict(event_id="llm-hitl-rej", mad_code="M4_a", policy=policy) + verdict.hitl.continue_execution = False + fut.set_result(verdict) + + result = await asyncio.wait_for(task, timeout=2.0) + assert not tool_ran + assert "BLOCKED" in result["messages"][0].content + + +# --------------------------------------------------------------------------- +# 7. Verdict arrives before gate checks (resolved future replay) +# --------------------------------------------------------------------------- + + +class TestVerdictReplay: + async def test_pre_resolved_verdict_still_blocks(self, tmp_path: Path) -> None: + """If the verdict future resolves before BaseTool.ainvoke reaches + the gate, it should still read the resolved value and block.""" + tool_ran = False + + async def my_tool(x: str) -> str: + """Tool stub.""" + nonlocal tool_ran + tool_ran = True + return x + + ws = _init_sdk(tmp_path) + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._connected.set() + ws._tool_call_id_to_event_id["tc-pre"] = "llm-pre" + + # Resolve the verdict BEFORE tool dispatch + fut = ws.register_pending("llm-pre") + fut.set_result(pb.Verdict(event_id="llm-pre", mad_code="M4_a", policy=policy)) + + # Small delay to ensure future is fully resolved + await asyncio.sleep(0.01) + + tool_node = ToolNode([my_tool]) + result = await tool_node.ainvoke( + _tool_state("tc-pre", "my_tool"), config=_runtime_config() + ) + + assert not tool_ran + assert "BLOCKED" in result["messages"][0].content + + +# --------------------------------------------------------------------------- +# 8. LRU eviction edge case +# --------------------------------------------------------------------------- + + +class TestLruEviction: + async def test_unknown_tool_call_id_verdict_timeout_blocks(self, tmp_path: Path) -> None: + """A tool_call_id not in the map → wait_for_tool_call_verdict returns + None → _async_gate treats None verdict as fail-closed → BLOCKED. + + This tests the case where the LLM event was never seen (or evicted). + The gate fail-closes because in MODE_BLOCK, absence of verdict = block.""" + tool_ran = False + + async def my_tool(x: str) -> str: + """Tool stub.""" + nonlocal tool_ran + tool_ran = True + return x + + ws = _init_sdk(tmp_path, block_timeout=0.1) + _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._connected.set() + # Intentionally NOT populating _tool_call_id_to_event_id + + tool_node = ToolNode([my_tool]) + result = await tool_node.ainvoke( + _tool_state("tc-unknown", "my_tool"), config=_runtime_config() + ) + + # Unknown tool_call_id → wait_for_tool_call_verdict returns None + # → _async_gate sees verdict=None → fail-closed → BLOCKED + assert not tool_ran + assert "BLOCKED" in result["messages"][0].content + + +# --------------------------------------------------------------------------- +# 9. LoginAck not received — should block (refuse to run without policy) +# --------------------------------------------------------------------------- + + +class TestPreLoginBlock: + async def test_tool_blocked_before_login_ack(self, tmp_path: Path) -> None: + """Before LoginAck arrives, the gate should block (fail-closed) + because we can't verify the org's policy.""" + tool_ran = False + + async def my_tool(x: str) -> str: + """Tool stub.""" + nonlocal tool_ran + tool_ran = True + return x + + ws = _init_sdk(tmp_path, block_timeout=0.1) + # Do NOT set login_ack_received — simulates pre-login state + # But we need to ensure _async_gate is actually called, so + # policy must appear active + ws._mode = pb.MODE_BLOCK + ws._connected.set() + ws._tool_call_id_to_event_id["tc-prelogin"] = "llm-prelogin" + + tool_node = ToolNode([my_tool]) + result = await tool_node.ainvoke( + _tool_state("tc-prelogin", "my_tool"), config=_runtime_config() + ) + + # Gate should block because LoginAck timeout fires (5s in prod, + # but here _login_ack_received is never set, so it times out) + assert not tool_ran + assert "BLOCKED" in result["messages"][0].content + + +# --------------------------------------------------------------------------- +# 10. Async tool with M4 verdict (BaseTool.ainvoke path) +# --------------------------------------------------------------------------- + + +class TestAsyncToolGate: + async def test_async_tool_ainvoke_gate_blocks_m4(self, tmp_path: Path) -> None: + """Async tools go through BaseTool.ainvoke → _async_gate directly + (no _sync_gate involved). Verify this path also blocks.""" + tool_ran = False + + async def async_danger(x: str) -> str: + """Async danger tool.""" + nonlocal tool_ran + tool_ran = True + return x + + ws = _init_sdk(tmp_path) + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._connected.set() + ws._tool_call_id_to_event_id["tc-async"] = "llm-async" + fut = ws.register_pending("llm-async") + fut.set_result(pb.Verdict(event_id="llm-async", mad_code="M4_a", policy=policy)) + + tool_node = ToolNode([async_danger]) + result = await tool_node.ainvoke( + _tool_state("tc-async", "async_danger"), config=_runtime_config() + ) + + assert not tool_ran + assert "BLOCKED" in result["messages"][0].content + + +# --------------------------------------------------------------------------- +# 11. Multiple tool calls from same LLM message +# --------------------------------------------------------------------------- + + +class TestMultiToolCall: + async def test_two_tools_same_llm_one_blocked_one_allowed(self, tmp_path: Path) -> None: + """An LLM emits two tool_calls. One is M4 (blocked), the other + is M0 (allowed). Each should be independently gated.""" + results: dict[str, bool] = {} + + async def read_file(path: str) -> str: + """Read file tool.""" + results["read"] = True + return f"contents of {path}" + + async def delete_file(path: str) -> str: + """Delete file tool.""" + results["delete"] = True + return f"deleted {path}" + + ws = _init_sdk(tmp_path) + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._connected.set() + + # Both tool_calls map to different LLM events (parallel agents scenario) + ws._tool_call_id_to_event_id["tc-read"] = "llm-read" + ws._tool_call_id_to_event_id["tc-delete"] = "llm-delete" + + fut_read = ws.register_pending("llm-read") + fut_read.set_result(pb.Verdict(event_id="llm-read", mad_code="M0_ok", policy=policy)) + + fut_delete = ws.register_pending("llm-delete") + fut_delete.set_result(pb.Verdict(event_id="llm-delete", mad_code="M4_a", policy=policy)) + + # Dispatch read (allowed) + tn_read = ToolNode([read_file]) + await tn_read.ainvoke( + _tool_state("tc-read", "read_file", {"path": "/tmp/x"}), + config=_runtime_config(), + ) + assert results.get("read") is True + + # Dispatch delete (blocked) + tn_delete = ToolNode([delete_file]) + result = await tn_delete.ainvoke( + _tool_state("tc-delete", "delete_file", {"path": "/tmp/x"}), + config=_runtime_config(), + ) + assert "delete" not in results + assert "BLOCKED" in result["messages"][0].content + + +# --------------------------------------------------------------------------- +# 12. Verify _sync_gate get_running_loop() behavior on actual thread +# --------------------------------------------------------------------------- + + +class TestGetRunningLoopOnThread: + async def test_get_running_loop_raises_on_worker_thread(self) -> None: + """Verify the core invariant: get_running_loop() raises RuntimeError + on a ThreadPoolExecutor worker, allowing _sync_gate to correctly + identify it as a worker thread and bridge to ws._loop.""" + loop = asyncio.get_running_loop() + + def check() -> tuple[bool, bool]: + has_running = False + get_event_loop_raises = False + try: + asyncio.get_running_loop() + has_running = True + except RuntimeError: + pass + try: + asyncio.get_event_loop() + except RuntimeError: + get_event_loop_raises = True + return has_running, get_event_loop_raises + + has_running, gel_raises = await loop.run_in_executor(None, check) + + assert not has_running, "Worker thread should NOT have a running loop" + import sys + if sys.version_info >= (3, 12): + assert gel_raises, ( + "On Python 3.12+, get_event_loop() should raise on worker thread" + ) diff --git a/sdk/python/tests/test_sync_gate_regression.py b/sdk/python/tests/test_sync_gate_regression.py new file mode 100644 index 0000000..9700050 --- /dev/null +++ b/sdk/python/tests/test_sync_gate_regression.py @@ -0,0 +1,268 @@ +"""Regression test for run 019eda10-b934-7d50-a16f-5aca881c5ee9. + +Bug: In Python 3.12+, ``asyncio.get_event_loop()`` raises RuntimeError on +bare worker threads (no loop set). The old ``_sync_gate`` caught that +RuntimeError and returned False (skip gate), so M4 tool calls executed +unblocked when dispatched from Pregel's ThreadPoolExecutor workers. + +Fix: ``_sync_gate`` now uses ``get_running_loop()`` (which correctly reports +"no running loop on THIS thread" on a worker thread) and falls through to +the ``run_coroutine_threadsafe`` bridge onto the WS client's loop. + +This test replicates the exact production scenario: a sync tool dispatched +by ToolNode.ainvoke (which uses run_in_executor internally), with a +MODE_BLOCK + policy_m4=True + M4 verdict. The tool body must NOT execute. +""" + +# pyright: reportPrivateUsage=false + +from __future__ import annotations + +import asyncio +from collections.abc import Iterator +from pathlib import Path +from typing import Any, cast +from unittest.mock import patch + +import adrian +import pytest +from adrian.proto import event_pb2 as pb +from adrian.ws import WebSocketClient +from langchain_core.messages import AIMessage +from langchain_core.runnables.config import RunnableConfig, ensure_config +from langgraph._internal._constants import CONF, CONFIG_KEY_RUNTIME +from langgraph.prebuilt import ToolNode +from langgraph.runtime import Runtime + + +def _runtime_config() -> RunnableConfig: + return ensure_config({CONF: {CONFIG_KEY_RUNTIME: Runtime()}}) + + +def _apply_mode( + ws: WebSocketClient, + mode: int, + *, + policy_m0: bool = False, + policy_m2: bool = False, + policy_m3: bool = False, + policy_m4: bool = False, +) -> pb.PolicySnapshot: + policy = pb.PolicySnapshot( + mode=cast("pb.Mode", mode), + policy_m0=policy_m0, + policy_m2=policy_m2, + policy_m3=policy_m3, + policy_m4=policy_m4, + ) + ws._mode = mode + ws._policy = policy + ws._login_ack_received.set() + return policy + + +@pytest.fixture(autouse=True) +def _cleanup() -> Iterator[None]: + yield + adrian.shutdown() + + +class TestSyncGateWorkerThreadRegression: + """Reproduce run 019eda10: M4 verdict present but sync tool ran anyway. + + The root cause was ``_sync_gate`` using ``asyncio.get_event_loop()`` + which raises RuntimeError on Python 3.12+ worker threads, causing the + gate to skip entirely and return False. + """ + + async def test_sync_tool_on_worker_thread_blocks_m4(self, tmp_path: Path) -> None: + """Sync tool dispatched from a ThreadPoolExecutor worker (like Pregel) + must be blocked when an M4 verdict is in scope.""" + tool_executed = False + + def dangerous_tool(x: str) -> str: + """Simulates a dangerous tool that should be blocked.""" + nonlocal tool_executed + tool_executed = True + return f"EXECUTED: {x}" + + adrian.init( + api_key="test-key", + log_file=str(tmp_path / "events.jsonl"), + auto_instrument=True, + ws_url="ws://x", + block_timeout=2.0, + ) + + ws = adrian._ws_client + assert ws is not None + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._connected.set() + ws._loop = asyncio.get_running_loop() + ws._tool_call_id_to_event_id["tc-m4"] = "llm-evt-m4" + + fut = ws.register_pending("llm-evt-m4") + fut.set_result( + pb.Verdict(event_id="llm-evt-m4", mad_code="M4_exfiltration", policy=policy) + ) + + tool_node = ToolNode([dangerous_tool]) + ai = AIMessage( + content="", + tool_calls=[{"id": "tc-m4", "name": "dangerous_tool", "args": {"x": "steal data"}}], + ) + state: dict[str, Any] = {"messages": [ai]} + + result = await tool_node.ainvoke(state, config=_runtime_config()) + + assert not tool_executed, ( + "CRITICAL: M4-flagged tool executed despite block verdict! " + "This is the run 019eda10 regression." + ) + msgs = result["messages"] + assert len(msgs) == 1 + assert "BLOCKED" in msgs[0].content + + async def test_sync_gate_does_not_use_get_event_loop(self, tmp_path: Path) -> None: + """Verify _sync_gate never calls asyncio.get_event_loop(). + + The old buggy code used get_event_loop() which raises RuntimeError + on Python 3.12+ worker threads and caused the gate to skip. + """ + adrian.init( + api_key="test-key", + log_file=str(tmp_path / "events.jsonl"), + auto_instrument=True, + ws_url="ws://x", + block_timeout=2.0, + ) + + ws = adrian._ws_client + assert ws is not None + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) + ws._connected.set() + ws._loop = asyncio.get_running_loop() + ws._tool_call_id_to_event_id["tc-1"] = "llm-evt" + + fut = ws.register_pending("llm-evt") + fut.set_result( + pb.Verdict(event_id="llm-evt", mad_code="M4_a", policy=policy) + ) + + original_get_event_loop = asyncio.get_event_loop + get_event_loop_called_from_worker = False + + def tracking_get_event_loop(): + nonlocal get_event_loop_called_from_worker + try: + asyncio.get_running_loop() + except RuntimeError: + # We're on a worker thread — this is the buggy call path + get_event_loop_called_from_worker = True + return original_get_event_loop() + + def passthrough_tool(x: str) -> str: + """Passthrough.""" + return x + + with patch("asyncio.get_event_loop", side_effect=tracking_get_event_loop): + tool_node = ToolNode([passthrough_tool]) + ai = AIMessage( + content="", + tool_calls=[{"id": "tc-1", "name": "passthrough_tool", "args": {"x": "hi"}}], + ) + state: dict[str, Any] = {"messages": [ai]} + await tool_node.ainvoke(state, config=_runtime_config()) + + assert not get_event_loop_called_from_worker, ( + "BUG: _sync_gate called asyncio.get_event_loop() from a worker thread. " + "This is the root cause of the run 019eda10 bypass." + ) + + async def test_m0_benign_tool_still_runs(self, tmp_path: Path) -> None: + """M0 (benign) verdict with policy_m0=False should let the tool run. + + Ensures the gate doesn't over-block — only in-scope verdicts halt. + """ + captured: list[str] = [] + + def safe_tool(x: str) -> str: + """Safe tool stub.""" + captured.append(x) + return x + + adrian.init( + api_key="test-key", + log_file=str(tmp_path / "events.jsonl"), + auto_instrument=True, + ws_url="ws://x", + block_timeout=2.0, + ) + + ws = adrian._ws_client + assert ws is not None + policy = _apply_mode(ws, pb.MODE_BLOCK, policy_m4=True) # m0 not in scope + ws._connected.set() + ws._loop = asyncio.get_running_loop() + ws._tool_call_id_to_event_id["tc-benign"] = "llm-evt-benign" + + fut = ws.register_pending("llm-evt-benign") + fut.set_result( + pb.Verdict(event_id="llm-evt-benign", mad_code="M0_benign", policy=policy) + ) + + tool_node = ToolNode([safe_tool]) + ai = AIMessage( + content="", + tool_calls=[{"id": "tc-benign", "name": "safe_tool", "args": {"x": "hello"}}], + ) + state: dict[str, Any] = {"messages": [ai]} + await tool_node.ainvoke(state, config=_runtime_config()) + + assert captured == ["hello"], "Benign M0 tool should have executed" + + +class TestOldBuggyGateWouldFail: + """Prove the old buggy _sync_gate pattern would fail on this Python version. + + The old code did: + try: + loop = asyncio.get_event_loop() + except RuntimeError: + return False # <-- BUG: skips gate entirely + + On Python 3.12+ worker threads, get_event_loop() raises RuntimeError + because there's no current event loop set for that thread. + """ + + async def test_get_event_loop_raises_on_worker_thread(self) -> None: + """Confirm the failure mode exists on this Python version.""" + loop = asyncio.get_running_loop() + + def check_on_worker() -> str: + try: + asyncio.get_running_loop() + return "has_running_loop" + except RuntimeError: + pass + + try: + asyncio.get_event_loop() + return "get_event_loop_succeeded" + except RuntimeError: + return "get_event_loop_raised" + + result = await loop.run_in_executor(None, check_on_worker) + + # On Python 3.12+, this should be "get_event_loop_raised" + # On Python 3.10-3.11, it may succeed (creating a new loop) + # Either way, the fix handles both: it uses get_running_loop() instead + import sys + if sys.version_info >= (3, 12): + assert result == "get_event_loop_raised", ( + f"Expected get_event_loop() to raise on worker thread " + f"(Python {sys.version_info}), got: {result}" + ) + # On older Python, get_event_loop may succeed but the old code still + # had the wrong behavior (would try run_until_complete on a non-running + # loop that wasn't connected to ws._loop)