From d019e0bbdf4c1037d40d11bd1bf3407e8da9a9a8 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 21 Feb 2026 00:21:45 -0800 Subject: [PATCH 01/53] feat: Add ExecuteSkillScriptTool for running skill scripts via code executor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds execute_skill_script tool to SkillToolset, enabling agents to run scripts from a skill's scripts/ directory through ADK's BaseCodeExecutor infrastructure. Supports Python (.py) and shell (.sh/.bash) scripts with optional input_args, executor resolution chain (toolset → agent fallback), and LLM-friendly error truncation. Includes 37 unit tests, a sample agent with Anthropic's mcp-builder skill and a python-helper skill, compatibility tests, and a live integration test exercising all four skill tools end-to-end with Gemini. Co-Authored-By: Claude Opus 4.6 --- .../samples/skill_script_demo/__init__.py | 15 + .../samples/skill_script_demo/agent.py | 47 + .../skill_script_demo/run_live_test.py | 751 ++++++++++++++ .../skills/mcp-builder/LICENSE.txt | 202 ++++ .../skills/mcp-builder/SKILL.md | 236 +++++ .../mcp-builder/references/evaluation.md | 602 +++++++++++ .../references/mcp_best_practices.md | 249 +++++ .../mcp-builder/references/node_mcp_server.md | 970 ++++++++++++++++++ .../references/python_mcp_server.md | 719 +++++++++++++ .../skills/mcp-builder/scripts/connections.py | 160 +++ .../skills/mcp-builder/scripts/evaluation.py | 428 ++++++++ .../scripts/example_evaluation.xml | 22 + .../mcp-builder/scripts/requirements.txt | 2 + .../skills/python-helper/SKILL.md | 33 + .../skills/python-helper/references/usage.md | 17 + .../skills/python-helper/scripts/fibonacci.py | 12 + .../python-helper/scripts/json_format.py | 21 + .../python-helper/scripts/word_count.py | 18 + .../skill_script_demo/test_skill_compat.py | 258 +++++ src/google/adk/tools/skill_toolset.py | 192 ++++ tests/unittests/tools/test_skill_toolset.py | 318 +++++- 21 files changed, 5271 insertions(+), 1 deletion(-) create mode 100644 contributing/samples/skill_script_demo/__init__.py create mode 100644 contributing/samples/skill_script_demo/agent.py create mode 100644 contributing/samples/skill_script_demo/run_live_test.py create mode 100644 contributing/samples/skill_script_demo/skills/mcp-builder/LICENSE.txt create mode 100644 contributing/samples/skill_script_demo/skills/mcp-builder/SKILL.md create mode 100644 contributing/samples/skill_script_demo/skills/mcp-builder/references/evaluation.md create mode 100644 contributing/samples/skill_script_demo/skills/mcp-builder/references/mcp_best_practices.md create mode 100644 contributing/samples/skill_script_demo/skills/mcp-builder/references/node_mcp_server.md create mode 100644 contributing/samples/skill_script_demo/skills/mcp-builder/references/python_mcp_server.md create mode 100644 contributing/samples/skill_script_demo/skills/mcp-builder/scripts/connections.py create mode 100644 contributing/samples/skill_script_demo/skills/mcp-builder/scripts/evaluation.py create mode 100644 contributing/samples/skill_script_demo/skills/mcp-builder/scripts/example_evaluation.xml create mode 100644 contributing/samples/skill_script_demo/skills/mcp-builder/scripts/requirements.txt create mode 100644 contributing/samples/skill_script_demo/skills/python-helper/SKILL.md create mode 100644 contributing/samples/skill_script_demo/skills/python-helper/references/usage.md create mode 100644 contributing/samples/skill_script_demo/skills/python-helper/scripts/fibonacci.py create mode 100644 contributing/samples/skill_script_demo/skills/python-helper/scripts/json_format.py create mode 100644 contributing/samples/skill_script_demo/skills/python-helper/scripts/word_count.py create mode 100644 contributing/samples/skill_script_demo/test_skill_compat.py diff --git a/contributing/samples/skill_script_demo/__init__.py b/contributing/samples/skill_script_demo/__init__.py new file mode 100644 index 0000000000..4015e47d6e --- /dev/null +++ b/contributing/samples/skill_script_demo/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from . import agent diff --git a/contributing/samples/skill_script_demo/agent.py b/contributing/samples/skill_script_demo/agent.py new file mode 100644 index 0000000000..0daf1d14bc --- /dev/null +++ b/contributing/samples/skill_script_demo/agent.py @@ -0,0 +1,47 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sample agent with multiple skills and script execution. + +This agent loads the Anthropic mcp-builder skill and a python-helper +skill, exercising all skill tools: list_skills, load_skill, +load_skill_resource, and execute_skill_script. +""" + +import pathlib + +from google.adk import Agent +from google.adk.code_executors.unsafe_local_code_executor import UnsafeLocalCodeExecutor +from google.adk.skills import load_skill_from_dir +from google.adk.tools.skill_toolset import SkillToolset + +_SKILLS_DIR = pathlib.Path(__file__).parent / "skills" + +mcp_builder_skill = load_skill_from_dir(_SKILLS_DIR / "mcp-builder") +python_helper_skill = load_skill_from_dir(_SKILLS_DIR / "python-helper") + +my_skill_toolset = SkillToolset( + skills=[mcp_builder_skill, python_helper_skill], + code_executor=UnsafeLocalCodeExecutor(), +) + +root_agent = Agent( + model="gemini-2.5-flash", + name="skill_script_demo_agent", + description=( + "An agent that demonstrates skill script execution using" + " the mcp-builder and python-helper skills." + ), + tools=[my_skill_toolset], +) diff --git a/contributing/samples/skill_script_demo/run_live_test.py b/contributing/samples/skill_script_demo/run_live_test.py new file mode 100644 index 0000000000..13d2cd52a6 --- /dev/null +++ b/contributing/samples/skill_script_demo/run_live_test.py @@ -0,0 +1,751 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Live integration test: realistic skill workflows with Gemini. + +Exercises the full skill lifecycle — discovery, loading, reference +reading, script inspection, and script execution — using both +Anthropic's mcp-builder skill and a lightweight python-helper skill. + +Usage: + export GOOGLE_CLOUD_PROJECT=your-project-id + python contributing/samples/skill_script_demo/run_live_test.py +""" + +import asyncio +import os +import pathlib +import sys +import traceback + +# Ensure the repo src is on the path +sys.path.insert(0, str(pathlib.Path(__file__).parents[3] / "src")) + +from google.adk import Agent +from google.adk.code_executors.unsafe_local_code_executor import UnsafeLocalCodeExecutor +from google.adk.runners import Runner +from google.adk.sessions import InMemorySessionService +from google.adk.skills import load_skill_from_dir +from google.adk.tools.skill_toolset import SkillToolset +from google.genai import types + +_SKILLS_DIR = pathlib.Path(__file__).parent / "skills" + +# Configure for Vertex AI +project = os.environ.get("GOOGLE_CLOUD_PROJECT") +if not project: + project = ( + os.popen("gcloud config get-value project 2>/dev/null").read().strip() + ) + if project: + os.environ["GOOGLE_CLOUD_PROJECT"] = project + +os.environ.setdefault("GOOGLE_GENAI_USE_VERTEXAI", "TRUE") +os.environ.setdefault("GOOGLE_CLOUD_LOCATION", "us-central1") + + +def build_agent(): + """Build the agent with mcp-builder and python-helper skills.""" + mcp_builder = load_skill_from_dir(_SKILLS_DIR / "mcp-builder") + python_helper = load_skill_from_dir(_SKILLS_DIR / "python-helper") + + toolset = SkillToolset( + skills=[mcp_builder, python_helper], + code_executor=UnsafeLocalCodeExecutor(), + ) + + return Agent( + model="gemini-2.5-flash", + name="skill_test_agent", + description=( + "An agent that uses skills for MCP development and Python utilities." + ), + tools=[toolset], + ) + + +# ── Event extraction helpers ──────────────────────────────────── + + +def extract_tool_calls(events): + """Extract tool call names from events.""" + names = [] + for ev in events: + if not ev.content or not ev.content.parts: + continue + for part in ev.content.parts: + if part.function_call: + names.append(part.function_call.name) + return names + + +def extract_tool_call_args(events): + """Extract tool calls with their arguments.""" + calls = [] + for ev in events: + if not ev.content or not ev.content.parts: + continue + for part in ev.content.parts: + if part.function_call: + calls.append({ + "name": part.function_call.name, + "args": dict(part.function_call.args or {}), + }) + return calls + + +def extract_tool_responses(events): + """Extract function response payloads.""" + responses = [] + for ev in events: + if not ev.content or not ev.content.parts: + continue + for part in ev.content.parts: + if part.function_response: + responses.append({ + "name": part.function_response.name, + "response": part.function_response.response, + }) + return responses + + +def extract_final_text(events): + """Extract the final text response.""" + for ev in reversed(events): + if not ev.content or not ev.content.parts: + continue + for part in ev.content.parts: + if part.text: + return part.text + return "" + + +# ── Query runner ──────────────────────────────────────────────── + + +async def run_query(runner, session_id, user_id, query): + """Send a query and collect all events.""" + content = types.Content( + role="user", + parts=[types.Part.from_text(text=query)], + ) + events = [] + async for event in runner.run_async( + user_id=user_id, + session_id=session_id, + new_message=content, + ): + events.append(event) + return events + + +# ── Check engine ──────────────────────────────────────────────── + +SKILL_TOOLS = frozenset([ + "list_skills", + "load_skill", + "load_skill_resource", + "execute_skill_script", +]) + + +def check(tc, tool_calls, call_args, responses, text): + """Run all checks for a test step. Returns (ok, messages).""" + ok = True + msgs = [] + + # ── expect_tools: every listed tool must appear ── + for t in tc.get("expect_tools", []): + if t in tool_calls: + msgs.append(f" PASS: Used '{t}'") + else: + msgs.append(f" FAIL: Expected '{t}' in {tool_calls}") + ok = False + + # ── expect_any_tool: at least one must appear ── + any_tools = tc.get("expect_any_tool", []) + if any_tools: + matched = [t for t in any_tools if t in tool_calls] + if matched: + msgs.append(f" PASS: Used one of {any_tools}: {matched}") + else: + msgs.append(f" FAIL: Expected one of {any_tools} in {tool_calls}") + ok = False + + # ── expect_no_skill_tools ── + if tc.get("expect_no_skill_tools"): + used = [t for t in tool_calls if t in SKILL_TOOLS] + if used: + msgs.append(f" FAIL: Should NOT have used: {used}") + ok = False + else: + msgs.append(" PASS: No skill tools used") + + # ── expect_text_contains ── + for s in tc.get("expect_text_contains", []): + if s.lower() in text.lower(): + msgs.append(f" PASS: Response contains '{s}'") + else: + msgs.append(f" FAIL: Response missing '{s}'") + ok = False + + # ── expect_text_any ── + text_any = tc.get("expect_text_any", []) + if text_any: + found = [s for s in text_any if s.lower() in text.lower()] + if found: + msgs.append(f" PASS: Response contains one of: {found}") + else: + msgs.append(f" FAIL: Response missing all of: {text_any}") + ok = False + + # ── expect_execute_script: verify execute call args ── + ex = tc.get("expect_execute_script") + if ex: + exec_calls = [c for c in call_args if c["name"] == "execute_skill_script"] + if not exec_calls: + msgs.append(f" FAIL: No execute_skill_script call in {tool_calls}") + ok = False + else: + c = exec_calls[0] + if "skill_name" in ex: + actual = c["args"].get("skill_name") + if actual == ex["skill_name"]: + msgs.append(f" PASS: skill_name='{actual}'") + else: + msgs.append( + " FAIL: skill_name: expected" + f" '{ex['skill_name']}' got '{actual}'" + ) + ok = False + if "script_name" in ex: + actual = c["args"].get("script_name", "") + expected = ex["script_name"] + if actual == expected or actual == f"scripts/{expected}": + msgs.append(f" PASS: script_name='{actual}'") + else: + msgs.append( + f" FAIL: script_name: expected '{expected}' got '{actual}'" + ) + ok = False + if "has_input_args" in ex: + has = bool(c["args"].get("input_args")) + if has == ex["has_input_args"]: + msgs.append(f" PASS: input_args present={has}") + else: + msgs.append( + " FAIL: input_args: expected" + f" present={ex['has_input_args']} got {has}" + ) + ok = False + + # ── expect_script_output: verify execution result ── + out = tc.get("expect_script_output") + if out: + exec_resps = [r for r in responses if r["name"] == "execute_skill_script"] + if not exec_resps: + msgs.append(" FAIL: No execute_skill_script response") + ok = False + else: + resp = exec_resps[0]["response"] + expect_status = out.get("status", "success") + actual_status = resp.get("status", "") + if actual_status == expect_status: + msgs.append(f" PASS: Script status='{actual_status}'") + else: + msgs.append( + " FAIL: Script status: expected" + f" '{expect_status}' got '{actual_status}'" + f" error={resp.get('error', '')}" + ) + ok = False + + stdout = resp.get("stdout", "") + for s in out.get("stdout_contains", []): + if s.lower() in stdout.lower(): + msgs.append(f" PASS: stdout contains '{s}'") + else: + msgs.append(f" FAIL: stdout missing '{s}'") + ok = False + + stderr = resp.get("stderr", "") + for s in out.get("stderr_contains", []): + if s.lower() in stderr.lower(): + msgs.append(f" PASS: stderr contains '{s}'") + else: + msgs.append(f" FAIL: stderr missing '{s}'") + ok = False + + # ── expect_resource_loaded: verify load_skill_resource resp ── + res = tc.get("expect_resource_loaded") + if res: + res_resps = [r for r in responses if r["name"] == "load_skill_resource"] + if not res_resps: + msgs.append(" FAIL: No load_skill_resource response") + ok = False + else: + for r in res_resps: + content = r["response"].get("content", "") + for s in res.get("content_contains", []): + if s.lower() in content.lower(): + msgs.append(f" PASS: Resource contains '{s}'") + else: + msgs.append(f" FAIL: Resource missing '{s}'") + ok = False + + return ok, msgs + + +# ── Test scenarios ────────────────────────────────────────────── + + +def get_test_scenarios(): + """Return test scenario groups.""" + return [ + # ───────────────────────────────────────────────────────── + # Scenario 1: Discover both skills + # ───────────────────────────────────────────────────────── + { + "name": "Scenario 1: Discover both skills", + "shared_session": False, + "steps": [{ + "name": "1a: List all available skills", + "query": ( + "What skills do you have? List them all" + " with their descriptions." + ), + "expect_any_tool": ["list_skills", "load_skill"], + "expect_text_contains": [ + "mcp-builder", + "python-helper", + ], + }], + }, + # ───────────────────────────────────────────────────────── + # Scenario 2: MCP development — load skill & read reference + # ───────────────────────────────────────────────────────── + { + "name": ( + "Scenario 2: MCP dev workflow — load skill, read Python reference" + ), + "shared_session": True, + "steps": [ + { + "name": ( + "2a: Load mcp-builder and ask about" + " Python server patterns" + ), + "query": ( + "I want to build an MCP server in" + " Python using FastMCP. Load the" + " mcp-builder skill and tell me the" + " four workflow phases." + ), + "expect_tools": ["load_skill"], + "expect_text_any": [ + "Phase", + "Research", + "Implementation", + "Review", + "Evaluation", + ], + }, + { + "name": "2b: Read the Python MCP server reference", + "query": ( + "Now use load_skill_resource to read" + " references/python_mcp_server.md" + " from the mcp-builder skill." + " Summarize the key patterns it" + " recommends for implementing tools." + ), + "expect_tools": ["load_skill_resource"], + "expect_resource_loaded": { + "content_contains": [ + "FastMCP", + "@mcp.tool", + ], + }, + }, + ], + }, + # ───────────────────────────────────────────────────────── + # Scenario 3: Inspect then execute connections.py + # ───────────────────────────────────────────────────────── + { + "name": "Scenario 3: Inspect and execute mcp-builder connections.py", + "shared_session": True, + "steps": [ + { + "name": "3a: Inspect connections.py source", + "query": ( + "Use load_skill_resource to show me" + " the source code of" + " scripts/connections.py from the" + " mcp-builder skill." + ), + "expect_tools": ["load_skill_resource"], + "expect_resource_loaded": { + "content_contains": [ + "MCPConnection", + "create_connection", + ], + }, + }, + { + "name": "3b: Execute connections.py", + "query": ( + "Please execute connections.py from" + " the mcp-builder skill using the" + " execute_skill_script tool. I want" + " to confirm the script loads without" + " errors." + ), + "expect_tools": ["execute_skill_script"], + "expect_execute_script": { + "skill_name": "mcp-builder", + "script_name": "connections.py", + }, + "expect_script_output": { + "status": "success", + }, + }, + ], + }, + # ───────────────────────────────────────────────────────── + # Scenario 4: Read evaluation reference + example XML + # ───────────────────────────────────────────────────────── + { + "name": "Scenario 4: Evaluation workflow — read eval guide & example", + "shared_session": True, + "steps": [ + { + "name": "4a: Read the evaluation guide reference", + "query": ( + "Use load_skill_resource to read" + " references/evaluation.md from the" + " mcp-builder skill. Summarize the" + " key requirements for a good" + " evaluation question." + ), + "expect_tools": ["load_skill_resource"], + "expect_resource_loaded": { + "content_contains": ["qa_pair"], + }, + "expect_text_any": [ + "independent", + "read-only", + "verifiable", + "stable", + "non-destructive", + ], + }, + { + "name": "4b: Read the example evaluation XML", + "query": ( + "Now use load_skill_resource to read" + " scripts/example_evaluation.xml from" + " the mcp-builder skill and tell me" + " how many QA pairs it contains and" + " what topics they cover." + ), + "expect_tools": ["load_skill_resource"], + "expect_resource_loaded": { + "content_contains": [ + "qa_pair", + "compound interest", + ], + }, + "expect_text_any": ["5", "five"], + }, + ], + }, + # ───────────────────────────────────────────────────────── + # Scenario 5: Execute fibonacci with python-helper + # ───────────────────────────────────────────────────────── + { + "name": "Scenario 5: Execute python-helper fibonacci script", + "shared_session": False, + "steps": [{ + "name": "5a: Generate first 8 Fibonacci numbers", + "query": ( + "Use the python-helper skill to run" + " fibonacci.py and generate the first 8" + " Fibonacci numbers." + ), + "expect_tools": ["execute_skill_script"], + "expect_execute_script": { + "skill_name": "python-helper", + "script_name": "fibonacci.py", + "has_input_args": True, + }, + "expect_script_output": { + "status": "success", + "stdout_contains": ["Fibonacci"], + }, + "expect_text_any": ["13", "0, 1, 1, 2, 3, 5, 8, 13"], + }], + }, + # ───────────────────────────────────────────────────────── + # Scenario 6: Execute word_count with python-helper + # ───────────────────────────────────────────────────────── + { + "name": "Scenario 6: Execute python-helper word_count script", + "shared_session": False, + "steps": [{ + "name": "6a: Analyze word frequency", + "query": ( + "Use the python-helper skill to run" + " word_count.py on the text: 'to be or" + " not to be that is the question'" + ), + "expect_tools": ["execute_skill_script"], + "expect_execute_script": { + "skill_name": "python-helper", + "script_name": "word_count.py", + "has_input_args": True, + }, + "expect_script_output": { + "status": "success", + "stdout_contains": [ + "Total words", + "Unique words", + ], + }, + "expect_text_any": ["to", "be"], + }], + }, + # ───────────────────────────────────────────────────────── + # Scenario 7: Full multi-turn with MCP builder + execution + # ───────────────────────────────────────────────────────── + { + "name": ( + "Scenario 7: Multi-turn — load mcp-builder," + " read reference, inspect script, execute" + ), + "shared_session": True, + "steps": [ + { + "name": "7a: Load mcp-builder skill and read best practices", + "query": ( + "Load the mcp-builder skill. Then" + " use load_skill_resource to read" + " references/mcp_best_practices.md." + " What does it say about tool naming?" + ), + "expect_tools": ["load_skill_resource"], + "expect_any_tool": [ + "load_skill", + "load_skill_resource", + ], + "expect_resource_loaded": { + "content_contains": ["snake_case"], + }, + "expect_text_any": [ + "snake_case", + "prefix", + "naming", + ], + }, + { + "name": "7b: Inspect connections.py source code", + "query": ( + "Use load_skill_resource to read" + " scripts/connections.py from the" + " mcp-builder skill." + ), + "expect_tools": ["load_skill_resource"], + "expect_resource_loaded": { + "content_contains": [ + "MCPConnection", + ], + }, + }, + { + "name": "7c: Execute connections.py we just inspected", + "query": ( + "I know it only defines classes," + " but please call" + " execute_skill_script with" + " skill_name='mcp-builder' and" + " script_name='connections.py'" + " anyway. If the mcp package is" + " missing the imports will raise" + " an ImportError, which is what" + " I want to check." + ), + "expect_tools": ["execute_skill_script"], + "expect_execute_script": { + "skill_name": "mcp-builder", + "script_name": "connections.py", + }, + "expect_script_output": { + "status": "success", + }, + }, + { + "name": "7d: Now use python-helper to process some data", + "query": ( + "Now switch to the python-helper" + " skill and use json_format.py to" + " pretty-print this JSON:" + ' {"tools":["list_tools",' + '"call_tool"],' + '"transport":"stdio"}' + ), + "expect_tools": ["execute_skill_script"], + "expect_execute_script": { + "skill_name": "python-helper", + "script_name": "json_format.py", + "has_input_args": True, + }, + "expect_script_output": { + "status": "success", + "stdout_contains": [ + "tools", + "transport", + ], + }, + }, + ], + }, + # ───────────────────────────────────────────────────────── + # Scenario 8: Agent routes to the right skill + # ───────────────────────────────────────────────────────── + { + "name": ( + "Scenario 8: Agent picks mcp-builder for" + " MCP questions, not python-helper" + ), + "shared_session": False, + "steps": [{ + "name": "8a: MCP question routes to mcp-builder", + "query": ( + "I need to add pagination to my MCP" + " server's list endpoint. Load the" + " relevant skill and show me the" + " recommended pagination pattern." + ), + "expect_tools": ["load_skill"], + "expect_text_any": [ + "offset", + "limit", + "has_more", + "pagination", + "next_offset", + ], + }], + }, + # ───────────────────────────────────────────────────────── + # Scenario 9: Unrelated query does not use skills + # ───────────────────────────────────────────────────────── + { + "name": "Scenario 9: Unrelated query avoids skills", + "shared_session": False, + "steps": [{ + "name": "9a: Simple math question", + "query": "What is 17 * 23?", + "expect_no_skill_tools": True, + "expect_text_contains": ["391"], + }], + }, + ] + + +# ── Main runner ───────────────────────────────────────────────── + + +async def main(): + agent = build_agent() + session_service = InMemorySessionService() + runner = Runner( + agent=agent, + app_name="skill_live_test", + session_service=session_service, + ) + + user_id = "test_user" + scenarios = get_test_scenarios() + total_passed = 0 + total_failed = 0 + failures = [] + + for scenario in scenarios: + print(f"\n{'#'*60}") + print(f" {scenario['name']}") + print(f"{'#'*60}") + + shared_session = None + if scenario["shared_session"]: + shared_session = await session_service.create_session( + app_name="skill_live_test", user_id=user_id + ) + + for step in scenario["steps"]: + print(f"\n {'─'*54}") + print(f" {step['name']}") + print(f" Query: {step['query'][:80]}...") + + session = shared_session or ( + await session_service.create_session( + app_name="skill_live_test", user_id=user_id + ) + ) + + try: + events = await run_query(runner, session.id, user_id, step["query"]) + except Exception as e: + print(f" ERROR: {e}") + traceback.print_exc() + total_failed += 1 + failures.append((step["name"], str(e))) + continue + + tool_calls = extract_tool_calls(events) + call_args = extract_tool_call_args(events) + responses = extract_tool_responses(events) + text = extract_final_text(events) + + print(f" Tool calls: {tool_calls}") + for c in call_args: + if c["name"] == "execute_skill_script": + print(f" Execute args: {c['args']}") + elif c["name"] == "load_skill_resource": + print(f" Resource: {c['args'].get('path', '?')}") + print(f" Response: {text[:250]}...") + + ok, msgs = check(step, tool_calls, call_args, responses, text) + for msg in msgs: + print(msg) + + if ok: + total_passed += 1 + else: + total_failed += 1 + failures.append((step["name"], "assertion failure")) + + total = total_passed + total_failed + print(f"\n{'='*60}") + print(f" RESULTS: {total_passed}/{total} passed") + if failures: + print(f"\n Failures:") + for name, reason in failures: + print(f" - {name}: {reason}") + print(f"{'='*60}") + return total_failed == 0 + + +if __name__ == "__main__": + success = asyncio.run(main()) + sys.exit(0 if success else 1) diff --git a/contributing/samples/skill_script_demo/skills/mcp-builder/LICENSE.txt b/contributing/samples/skill_script_demo/skills/mcp-builder/LICENSE.txt new file mode 100644 index 0000000000..7a4a3ea242 --- /dev/null +++ b/contributing/samples/skill_script_demo/skills/mcp-builder/LICENSE.txt @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. \ No newline at end of file diff --git a/contributing/samples/skill_script_demo/skills/mcp-builder/SKILL.md b/contributing/samples/skill_script_demo/skills/mcp-builder/SKILL.md new file mode 100644 index 0000000000..8a1a77a47d --- /dev/null +++ b/contributing/samples/skill_script_demo/skills/mcp-builder/SKILL.md @@ -0,0 +1,236 @@ +--- +name: mcp-builder +description: Guide for creating high-quality MCP (Model Context Protocol) servers that enable LLMs to interact with external services through well-designed tools. Use when building MCP servers to integrate external APIs or services, whether in Python (FastMCP) or Node/TypeScript (MCP SDK). +license: Complete terms in LICENSE.txt +--- + +# MCP Server Development Guide + +## Overview + +Create MCP (Model Context Protocol) servers that enable LLMs to interact with external services through well-designed tools. The quality of an MCP server is measured by how well it enables LLMs to accomplish real-world tasks. + +--- + +# Process + +## 🚀 High-Level Workflow + +Creating a high-quality MCP server involves four main phases: + +### Phase 1: Deep Research and Planning + +#### 1.1 Understand Modern MCP Design + +**API Coverage vs. Workflow Tools:** +Balance comprehensive API endpoint coverage with specialized workflow tools. Workflow tools can be more convenient for specific tasks, while comprehensive coverage gives agents flexibility to compose operations. Performance varies by client—some clients benefit from code execution that combines basic tools, while others work better with higher-level workflows. When uncertain, prioritize comprehensive API coverage. + +**Tool Naming and Discoverability:** +Clear, descriptive tool names help agents find the right tools quickly. Use consistent prefixes (e.g., `github_create_issue`, `github_list_repos`) and action-oriented naming. + +**Context Management:** +Agents benefit from concise tool descriptions and the ability to filter/paginate results. Design tools that return focused, relevant data. Some clients support code execution which can help agents filter and process data efficiently. + +**Actionable Error Messages:** +Error messages should guide agents toward solutions with specific suggestions and next steps. + +#### 1.2 Study MCP Protocol Documentation + +**Navigate the MCP specification:** + +Start with the sitemap to find relevant pages: `https://modelcontextprotocol.io/sitemap.xml` + +Then fetch specific pages with `.md` suffix for markdown format (e.g., `https://modelcontextprotocol.io/specification/draft.md`). + +Key pages to review: +- Specification overview and architecture +- Transport mechanisms (streamable HTTP, stdio) +- Tool, resource, and prompt definitions + +#### 1.3 Study Framework Documentation + +**Recommended stack:** +- **Language**: TypeScript (high-quality SDK support and good compatibility in many execution environments e.g. MCPB. Plus AI models are good at generating TypeScript code, benefiting from its broad usage, static typing and good linting tools) +- **Transport**: Streamable HTTP for remote servers, using stateless JSON (simpler to scale and maintain, as opposed to stateful sessions and streaming responses). stdio for local servers. + +**Load framework documentation:** + +- **MCP Best Practices**: [📋 View Best Practices](./reference/mcp_best_practices.md) - Core guidelines + +**For TypeScript (recommended):** +- **TypeScript SDK**: Use WebFetch to load `https://raw.githubusercontent.com/modelcontextprotocol/typescript-sdk/main/README.md` +- [⚡ TypeScript Guide](./reference/node_mcp_server.md) - TypeScript patterns and examples + +**For Python:** +- **Python SDK**: Use WebFetch to load `https://raw.githubusercontent.com/modelcontextprotocol/python-sdk/main/README.md` +- [🐍 Python Guide](./reference/python_mcp_server.md) - Python patterns and examples + +#### 1.4 Plan Your Implementation + +**Understand the API:** +Review the service's API documentation to identify key endpoints, authentication requirements, and data models. Use web search and WebFetch as needed. + +**Tool Selection:** +Prioritize comprehensive API coverage. List endpoints to implement, starting with the most common operations. + +--- + +### Phase 2: Implementation + +#### 2.1 Set Up Project Structure + +See language-specific guides for project setup: +- [⚡ TypeScript Guide](./reference/node_mcp_server.md) - Project structure, package.json, tsconfig.json +- [🐍 Python Guide](./reference/python_mcp_server.md) - Module organization, dependencies + +#### 2.2 Implement Core Infrastructure + +Create shared utilities: +- API client with authentication +- Error handling helpers +- Response formatting (JSON/Markdown) +- Pagination support + +#### 2.3 Implement Tools + +For each tool: + +**Input Schema:** +- Use Zod (TypeScript) or Pydantic (Python) +- Include constraints and clear descriptions +- Add examples in field descriptions + +**Output Schema:** +- Define `outputSchema` where possible for structured data +- Use `structuredContent` in tool responses (TypeScript SDK feature) +- Helps clients understand and process tool outputs + +**Tool Description:** +- Concise summary of functionality +- Parameter descriptions +- Return type schema + +**Implementation:** +- Async/await for I/O operations +- Proper error handling with actionable messages +- Support pagination where applicable +- Return both text content and structured data when using modern SDKs + +**Annotations:** +- `readOnlyHint`: true/false +- `destructiveHint`: true/false +- `idempotentHint`: true/false +- `openWorldHint`: true/false + +--- + +### Phase 3: Review and Test + +#### 3.1 Code Quality + +Review for: +- No duplicated code (DRY principle) +- Consistent error handling +- Full type coverage +- Clear tool descriptions + +#### 3.2 Build and Test + +**TypeScript:** +- Run `npm run build` to verify compilation +- Test with MCP Inspector: `npx @modelcontextprotocol/inspector` + +**Python:** +- Verify syntax: `python -m py_compile your_server.py` +- Test with MCP Inspector + +See language-specific guides for detailed testing approaches and quality checklists. + +--- + +### Phase 4: Create Evaluations + +After implementing your MCP server, create comprehensive evaluations to test its effectiveness. + +**Load [✅ Evaluation Guide](./reference/evaluation.md) for complete evaluation guidelines.** + +#### 4.1 Understand Evaluation Purpose + +Use evaluations to test whether LLMs can effectively use your MCP server to answer realistic, complex questions. + +#### 4.2 Create 10 Evaluation Questions + +To create effective evaluations, follow the process outlined in the evaluation guide: + +1. **Tool Inspection**: List available tools and understand their capabilities +2. **Content Exploration**: Use READ-ONLY operations to explore available data +3. **Question Generation**: Create 10 complex, realistic questions +4. **Answer Verification**: Solve each question yourself to verify answers + +#### 4.3 Evaluation Requirements + +Ensure each question is: +- **Independent**: Not dependent on other questions +- **Read-only**: Only non-destructive operations required +- **Complex**: Requiring multiple tool calls and deep exploration +- **Realistic**: Based on real use cases humans would care about +- **Verifiable**: Single, clear answer that can be verified by string comparison +- **Stable**: Answer won't change over time + +#### 4.4 Output Format + +Create an XML file with this structure: + +```xml + + + Find discussions about AI model launches with animal codenames. One model needed a specific safety designation that uses the format ASL-X. What number X was being determined for the model named after a spotted wild cat? + 3 + + + +``` + +--- + +# Reference Files + +## 📚 Documentation Library + +Load these resources as needed during development: + +### Core MCP Documentation (Load First) +- **MCP Protocol**: Start with sitemap at `https://modelcontextprotocol.io/sitemap.xml`, then fetch specific pages with `.md` suffix +- [📋 MCP Best Practices](./reference/mcp_best_practices.md) - Universal MCP guidelines including: + - Server and tool naming conventions + - Response format guidelines (JSON vs Markdown) + - Pagination best practices + - Transport selection (streamable HTTP vs stdio) + - Security and error handling standards + +### SDK Documentation (Load During Phase 1/2) +- **Python SDK**: Fetch from `https://raw.githubusercontent.com/modelcontextprotocol/python-sdk/main/README.md` +- **TypeScript SDK**: Fetch from `https://raw.githubusercontent.com/modelcontextprotocol/typescript-sdk/main/README.md` + +### Language-Specific Implementation Guides (Load During Phase 2) +- [🐍 Python Implementation Guide](./reference/python_mcp_server.md) - Complete Python/FastMCP guide with: + - Server initialization patterns + - Pydantic model examples + - Tool registration with `@mcp.tool` + - Complete working examples + - Quality checklist + +- [⚡ TypeScript Implementation Guide](./reference/node_mcp_server.md) - Complete TypeScript guide with: + - Project structure + - Zod schema patterns + - Tool registration with `server.registerTool` + - Complete working examples + - Quality checklist + +### Evaluation Guide (Load During Phase 4) +- [✅ Evaluation Guide](./reference/evaluation.md) - Complete evaluation creation guide with: + - Question creation guidelines + - Answer verification strategies + - XML format specifications + - Example questions and answers + - Running an evaluation with the provided scripts diff --git a/contributing/samples/skill_script_demo/skills/mcp-builder/references/evaluation.md b/contributing/samples/skill_script_demo/skills/mcp-builder/references/evaluation.md new file mode 100644 index 0000000000..87e9bb7884 --- /dev/null +++ b/contributing/samples/skill_script_demo/skills/mcp-builder/references/evaluation.md @@ -0,0 +1,602 @@ +# MCP Server Evaluation Guide + +## Overview + +This document provides guidance on creating comprehensive evaluations for MCP servers. Evaluations test whether LLMs can effectively use your MCP server to answer realistic, complex questions using only the tools provided. + +--- + +## Quick Reference + +### Evaluation Requirements +- Create 10 human-readable questions +- Questions must be READ-ONLY, INDEPENDENT, NON-DESTRUCTIVE +- Each question requires multiple tool calls (potentially dozens) +- Answers must be single, verifiable values +- Answers must be STABLE (won't change over time) + +### Output Format +```xml + + + Your question here + Single verifiable answer + + +``` + +--- + +## Purpose of Evaluations + +The measure of quality of an MCP server is NOT how well or comprehensively the server implements tools, but how well these implementations (input/output schemas, docstrings/descriptions, functionality) enable LLMs with no other context and access ONLY to the MCP servers to answer realistic and difficult questions. + +## Evaluation Overview + +Create 10 human-readable questions requiring ONLY READ-ONLY, INDEPENDENT, NON-DESTRUCTIVE, and IDEMPOTENT operations to answer. Each question should be: +- Realistic +- Clear and concise +- Unambiguous +- Complex, requiring potentially dozens of tool calls or steps +- Answerable with a single, verifiable value that you identify in advance + +## Question Guidelines + +### Core Requirements + +1. **Questions MUST be independent** + - Each question should NOT depend on the answer to any other question + - Should not assume prior write operations from processing another question + +2. **Questions MUST require ONLY NON-DESTRUCTIVE AND IDEMPOTENT tool use** + - Should not instruct or require modifying state to arrive at the correct answer + +3. **Questions must be REALISTIC, CLEAR, CONCISE, and COMPLEX** + - Must require another LLM to use multiple (potentially dozens of) tools or steps to answer + +### Complexity and Depth + +4. **Questions must require deep exploration** + - Consider multi-hop questions requiring multiple sub-questions and sequential tool calls + - Each step should benefit from information found in previous questions + +5. **Questions may require extensive paging** + - May need paging through multiple pages of results + - May require querying old data (1-2 years out-of-date) to find niche information + - The questions must be DIFFICULT + +6. **Questions must require deep understanding** + - Rather than surface-level knowledge + - May pose complex ideas as True/False questions requiring evidence + - May use multiple-choice format where LLM must search different hypotheses + +7. **Questions must not be solvable with straightforward keyword search** + - Do not include specific keywords from the target content + - Use synonyms, related concepts, or paraphrases + - Require multiple searches, analyzing multiple related items, extracting context, then deriving the answer + +### Tool Testing + +8. **Questions should stress-test tool return values** + - May elicit tools returning large JSON objects or lists, overwhelming the LLM + - Should require understanding multiple modalities of data: + - IDs and names + - Timestamps and datetimes (months, days, years, seconds) + - File IDs, names, extensions, and mimetypes + - URLs, GIDs, etc. + - Should probe the tool's ability to return all useful forms of data + +9. **Questions should MOSTLY reflect real human use cases** + - The kinds of information retrieval tasks that HUMANS assisted by an LLM would care about + +10. **Questions may require dozens of tool calls** + - This challenges LLMs with limited context + - Encourages MCP server tools to reduce information returned + +11. **Include ambiguous questions** + - May be ambiguous OR require difficult decisions on which tools to call + - Force the LLM to potentially make mistakes or misinterpret + - Ensure that despite AMBIGUITY, there is STILL A SINGLE VERIFIABLE ANSWER + +### Stability + +12. **Questions must be designed so the answer DOES NOT CHANGE** + - Do not ask questions that rely on "current state" which is dynamic + - For example, do not count: + - Number of reactions to a post + - Number of replies to a thread + - Number of members in a channel + +13. **DO NOT let the MCP server RESTRICT the kinds of questions you create** + - Create challenging and complex questions + - Some may not be solvable with the available MCP server tools + - Questions may require specific output formats (datetime vs. epoch time, JSON vs. MARKDOWN) + - Questions may require dozens of tool calls to complete + +## Answer Guidelines + +### Verification + +1. **Answers must be VERIFIABLE via direct string comparison** + - If the answer can be re-written in many formats, clearly specify the output format in the QUESTION + - Examples: "Use YYYY/MM/DD.", "Respond True or False.", "Answer A, B, C, or D and nothing else." + - Answer should be a single VERIFIABLE value such as: + - User ID, user name, display name, first name, last name + - Channel ID, channel name + - Message ID, string + - URL, title + - Numerical quantity + - Timestamp, datetime + - Boolean (for True/False questions) + - Email address, phone number + - File ID, file name, file extension + - Multiple choice answer + - Answers must not require special formatting or complex, structured output + - Answer will be verified using DIRECT STRING COMPARISON + +### Readability + +2. **Answers should generally prefer HUMAN-READABLE formats** + - Examples: names, first name, last name, datetime, file name, message string, URL, yes/no, true/false, a/b/c/d + - Rather than opaque IDs (though IDs are acceptable) + - The VAST MAJORITY of answers should be human-readable + +### Stability + +3. **Answers must be STABLE/STATIONARY** + - Look at old content (e.g., conversations that have ended, projects that have launched, questions answered) + - Create QUESTIONS based on "closed" concepts that will always return the same answer + - Questions may ask to consider a fixed time window to insulate from non-stationary answers + - Rely on context UNLIKELY to change + - Example: if finding a paper name, be SPECIFIC enough so answer is not confused with papers published later + +4. **Answers must be CLEAR and UNAMBIGUOUS** + - Questions must be designed so there is a single, clear answer + - Answer can be derived from using the MCP server tools + +### Diversity + +5. **Answers must be DIVERSE** + - Answer should be a single VERIFIABLE value in diverse modalities and formats + - User concept: user ID, user name, display name, first name, last name, email address, phone number + - Channel concept: channel ID, channel name, channel topic + - Message concept: message ID, message string, timestamp, month, day, year + +6. **Answers must NOT be complex structures** + - Not a list of values + - Not a complex object + - Not a list of IDs or strings + - Not natural language text + - UNLESS the answer can be straightforwardly verified using DIRECT STRING COMPARISON + - And can be realistically reproduced + - It should be unlikely that an LLM would return the same list in any other order or format + +## Evaluation Process + +### Step 1: Documentation Inspection + +Read the documentation of the target API to understand: +- Available endpoints and functionality +- If ambiguity exists, fetch additional information from the web +- Parallelize this step AS MUCH AS POSSIBLE +- Ensure each subagent is ONLY examining documentation from the file system or on the web + +### Step 2: Tool Inspection + +List the tools available in the MCP server: +- Inspect the MCP server directly +- Understand input/output schemas, docstrings, and descriptions +- WITHOUT calling the tools themselves at this stage + +### Step 3: Developing Understanding + +Repeat steps 1 & 2 until you have a good understanding: +- Iterate multiple times +- Think about the kinds of tasks you want to create +- Refine your understanding +- At NO stage should you READ the code of the MCP server implementation itself +- Use your intuition and understanding to create reasonable, realistic, but VERY challenging tasks + +### Step 4: Read-Only Content Inspection + +After understanding the API and tools, USE the MCP server tools: +- Inspect content using READ-ONLY and NON-DESTRUCTIVE operations ONLY +- Goal: identify specific content (e.g., users, channels, messages, projects, tasks) for creating realistic questions +- Should NOT call any tools that modify state +- Will NOT read the code of the MCP server implementation itself +- Parallelize this step with individual sub-agents pursuing independent explorations +- Ensure each subagent is only performing READ-ONLY, NON-DESTRUCTIVE, and IDEMPOTENT operations +- BE CAREFUL: SOME TOOLS may return LOTS OF DATA which would cause you to run out of CONTEXT +- Make INCREMENTAL, SMALL, AND TARGETED tool calls for exploration +- In all tool call requests, use the `limit` parameter to limit results (<10) +- Use pagination + +### Step 5: Task Generation + +After inspecting the content, create 10 human-readable questions: +- An LLM should be able to answer these with the MCP server +- Follow all question and answer guidelines above + +## Output Format + +Each QA pair consists of a question and an answer. The output should be an XML file with this structure: + +```xml + + + Find the project created in Q2 2024 with the highest number of completed tasks. What is the project name? + Website Redesign + + + Search for issues labeled as "bug" that were closed in March 2024. Which user closed the most issues? Provide their username. + sarah_dev + + + Look for pull requests that modified files in the /api directory and were merged between January 1 and January 31, 2024. How many different contributors worked on these PRs? + 7 + + + Find the repository with the most stars that was created before 2023. What is the repository name? + data-pipeline + + +``` + +## Evaluation Examples + +### Good Questions + +**Example 1: Multi-hop question requiring deep exploration (GitHub MCP)** +```xml + + Find the repository that was archived in Q3 2023 and had previously been the most forked project in the organization. What was the primary programming language used in that repository? + Python + +``` + +This question is good because: +- Requires multiple searches to find archived repositories +- Needs to identify which had the most forks before archival +- Requires examining repository details for the language +- Answer is a simple, verifiable value +- Based on historical (closed) data that won't change + +**Example 2: Requires understanding context without keyword matching (Project Management MCP)** +```xml + + Locate the initiative focused on improving customer onboarding that was completed in late 2023. The project lead created a retrospective document after completion. What was the lead's role title at that time? + Product Manager + +``` + +This question is good because: +- Doesn't use specific project name ("initiative focused on improving customer onboarding") +- Requires finding completed projects from specific timeframe +- Needs to identify the project lead and their role +- Requires understanding context from retrospective documents +- Answer is human-readable and stable +- Based on completed work (won't change) + +**Example 3: Complex aggregation requiring multiple steps (Issue Tracker MCP)** +```xml + + Among all bugs reported in January 2024 that were marked as critical priority, which assignee resolved the highest percentage of their assigned bugs within 48 hours? Provide the assignee's username. + alex_eng + +``` + +This question is good because: +- Requires filtering bugs by date, priority, and status +- Needs to group by assignee and calculate resolution rates +- Requires understanding timestamps to determine 48-hour windows +- Tests pagination (potentially many bugs to process) +- Answer is a single username +- Based on historical data from specific time period + +**Example 4: Requires synthesis across multiple data types (CRM MCP)** +```xml + + Find the account that upgraded from the Starter to Enterprise plan in Q4 2023 and had the highest annual contract value. What industry does this account operate in? + Healthcare + +``` + +This question is good because: +- Requires understanding subscription tier changes +- Needs to identify upgrade events in specific timeframe +- Requires comparing contract values +- Must access account industry information +- Answer is simple and verifiable +- Based on completed historical transactions + +### Poor Questions + +**Example 1: Answer changes over time** +```xml + + How many open issues are currently assigned to the engineering team? + 47 + +``` + +This question is poor because: +- The answer will change as issues are created, closed, or reassigned +- Not based on stable/stationary data +- Relies on "current state" which is dynamic + +**Example 2: Too easy with keyword search** +```xml + + Find the pull request with title "Add authentication feature" and tell me who created it. + developer123 + +``` + +This question is poor because: +- Can be solved with a straightforward keyword search for exact title +- Doesn't require deep exploration or understanding +- No synthesis or analysis needed + +**Example 3: Ambiguous answer format** +```xml + + List all the repositories that have Python as their primary language. + repo1, repo2, repo3, data-pipeline, ml-tools + +``` + +This question is poor because: +- Answer is a list that could be returned in any order +- Difficult to verify with direct string comparison +- LLM might format differently (JSON array, comma-separated, newline-separated) +- Better to ask for a specific aggregate (count) or superlative (most stars) + +## Verification Process + +After creating evaluations: + +1. **Examine the XML file** to understand the schema +2. **Load each task instruction** and in parallel using the MCP server and tools, identify the correct answer by attempting to solve the task YOURSELF +3. **Flag any operations** that require WRITE or DESTRUCTIVE operations +4. **Accumulate all CORRECT answers** and replace any incorrect answers in the document +5. **Remove any ``** that require WRITE or DESTRUCTIVE operations + +Remember to parallelize solving tasks to avoid running out of context, then accumulate all answers and make changes to the file at the end. + +## Tips for Creating Quality Evaluations + +1. **Think Hard and Plan Ahead** before generating tasks +2. **Parallelize Where Opportunity Arises** to speed up the process and manage context +3. **Focus on Realistic Use Cases** that humans would actually want to accomplish +4. **Create Challenging Questions** that test the limits of the MCP server's capabilities +5. **Ensure Stability** by using historical data and closed concepts +6. **Verify Answers** by solving the questions yourself using the MCP server tools +7. **Iterate and Refine** based on what you learn during the process + +--- + +# Running Evaluations + +After creating your evaluation file, you can use the provided evaluation harness to test your MCP server. + +## Setup + +1. **Install Dependencies** + + ```bash + pip install -r scripts/requirements.txt + ``` + + Or install manually: + ```bash + pip install anthropic mcp + ``` + +2. **Set API Key** + + ```bash + export ANTHROPIC_API_KEY=your_api_key_here + ``` + +## Evaluation File Format + +Evaluation files use XML format with `` elements: + +```xml + + + Find the project created in Q2 2024 with the highest number of completed tasks. What is the project name? + Website Redesign + + + Search for issues labeled as "bug" that were closed in March 2024. Which user closed the most issues? Provide their username. + sarah_dev + + +``` + +## Running Evaluations + +The evaluation script (`scripts/evaluation.py`) supports three transport types: + +**Important:** +- **stdio transport**: The evaluation script automatically launches and manages the MCP server process for you. Do not run the server manually. +- **sse/http transports**: You must start the MCP server separately before running the evaluation. The script connects to the already-running server at the specified URL. + +### 1. Local STDIO Server + +For locally-run MCP servers (script launches the server automatically): + +```bash +python scripts/evaluation.py \ + -t stdio \ + -c python \ + -a my_mcp_server.py \ + evaluation.xml +``` + +With environment variables: +```bash +python scripts/evaluation.py \ + -t stdio \ + -c python \ + -a my_mcp_server.py \ + -e API_KEY=abc123 \ + -e DEBUG=true \ + evaluation.xml +``` + +### 2. Server-Sent Events (SSE) + +For SSE-based MCP servers (you must start the server first): + +```bash +python scripts/evaluation.py \ + -t sse \ + -u https://example.com/mcp \ + -H "Authorization: Bearer token123" \ + -H "X-Custom-Header: value" \ + evaluation.xml +``` + +### 3. HTTP (Streamable HTTP) + +For HTTP-based MCP servers (you must start the server first): + +```bash +python scripts/evaluation.py \ + -t http \ + -u https://example.com/mcp \ + -H "Authorization: Bearer token123" \ + evaluation.xml +``` + +## Command-Line Options + +``` +usage: evaluation.py [-h] [-t {stdio,sse,http}] [-m MODEL] [-c COMMAND] + [-a ARGS [ARGS ...]] [-e ENV [ENV ...]] [-u URL] + [-H HEADERS [HEADERS ...]] [-o OUTPUT] + eval_file + +positional arguments: + eval_file Path to evaluation XML file + +optional arguments: + -h, --help Show help message + -t, --transport Transport type: stdio, sse, or http (default: stdio) + -m, --model Claude model to use (default: claude-3-7-sonnet-20250219) + -o, --output Output file for report (default: print to stdout) + +stdio options: + -c, --command Command to run MCP server (e.g., python, node) + -a, --args Arguments for the command (e.g., server.py) + -e, --env Environment variables in KEY=VALUE format + +sse/http options: + -u, --url MCP server URL + -H, --header HTTP headers in 'Key: Value' format +``` + +## Output + +The evaluation script generates a detailed report including: + +- **Summary Statistics**: + - Accuracy (correct/total) + - Average task duration + - Average tool calls per task + - Total tool calls + +- **Per-Task Results**: + - Prompt and expected response + - Actual response from the agent + - Whether the answer was correct (✅/❌) + - Duration and tool call details + - Agent's summary of its approach + - Agent's feedback on the tools + +### Save Report to File + +```bash +python scripts/evaluation.py \ + -t stdio \ + -c python \ + -a my_server.py \ + -o evaluation_report.md \ + evaluation.xml +``` + +## Complete Example Workflow + +Here's a complete example of creating and running an evaluation: + +1. **Create your evaluation file** (`my_evaluation.xml`): + +```xml + + + Find the user who created the most issues in January 2024. What is their username? + alice_developer + + + Among all pull requests merged in Q1 2024, which repository had the highest number? Provide the repository name. + backend-api + + + Find the project that was completed in December 2023 and had the longest duration from start to finish. How many days did it take? + 127 + + +``` + +2. **Install dependencies**: + +```bash +pip install -r scripts/requirements.txt +export ANTHROPIC_API_KEY=your_api_key +``` + +3. **Run evaluation**: + +```bash +python scripts/evaluation.py \ + -t stdio \ + -c python \ + -a github_mcp_server.py \ + -e GITHUB_TOKEN=ghp_xxx \ + -o github_eval_report.md \ + my_evaluation.xml +``` + +4. **Review the report** in `github_eval_report.md` to: + - See which questions passed/failed + - Read the agent's feedback on your tools + - Identify areas for improvement + - Iterate on your MCP server design + +## Troubleshooting + +### Connection Errors + +If you get connection errors: +- **STDIO**: Verify the command and arguments are correct +- **SSE/HTTP**: Check the URL is accessible and headers are correct +- Ensure any required API keys are set in environment variables or headers + +### Low Accuracy + +If many evaluations fail: +- Review the agent's feedback for each task +- Check if tool descriptions are clear and comprehensive +- Verify input parameters are well-documented +- Consider whether tools return too much or too little data +- Ensure error messages are actionable + +### Timeout Issues + +If tasks are timing out: +- Use a more capable model (e.g., `claude-3-7-sonnet-20250219`) +- Check if tools are returning too much data +- Verify pagination is working correctly +- Consider simplifying complex questions \ No newline at end of file diff --git a/contributing/samples/skill_script_demo/skills/mcp-builder/references/mcp_best_practices.md b/contributing/samples/skill_script_demo/skills/mcp-builder/references/mcp_best_practices.md new file mode 100644 index 0000000000..b9d343cc3a --- /dev/null +++ b/contributing/samples/skill_script_demo/skills/mcp-builder/references/mcp_best_practices.md @@ -0,0 +1,249 @@ +# MCP Server Best Practices + +## Quick Reference + +### Server Naming +- **Python**: `{service}_mcp` (e.g., `slack_mcp`) +- **Node/TypeScript**: `{service}-mcp-server` (e.g., `slack-mcp-server`) + +### Tool Naming +- Use snake_case with service prefix +- Format: `{service}_{action}_{resource}` +- Example: `slack_send_message`, `github_create_issue` + +### Response Formats +- Support both JSON and Markdown formats +- JSON for programmatic processing +- Markdown for human readability + +### Pagination +- Always respect `limit` parameter +- Return `has_more`, `next_offset`, `total_count` +- Default to 20-50 items + +### Transport +- **Streamable HTTP**: For remote servers, multi-client scenarios +- **stdio**: For local integrations, command-line tools +- Avoid SSE (deprecated in favor of streamable HTTP) + +--- + +## Server Naming Conventions + +Follow these standardized naming patterns: + +**Python**: Use format `{service}_mcp` (lowercase with underscores) +- Examples: `slack_mcp`, `github_mcp`, `jira_mcp` + +**Node/TypeScript**: Use format `{service}-mcp-server` (lowercase with hyphens) +- Examples: `slack-mcp-server`, `github-mcp-server`, `jira-mcp-server` + +The name should be general, descriptive of the service being integrated, easy to infer from the task description, and without version numbers. + +--- + +## Tool Naming and Design + +### Tool Naming + +1. **Use snake_case**: `search_users`, `create_project`, `get_channel_info` +2. **Include service prefix**: Anticipate that your MCP server may be used alongside other MCP servers + - Use `slack_send_message` instead of just `send_message` + - Use `github_create_issue` instead of just `create_issue` +3. **Be action-oriented**: Start with verbs (get, list, search, create, etc.) +4. **Be specific**: Avoid generic names that could conflict with other servers + +### Tool Design + +- Tool descriptions must narrowly and unambiguously describe functionality +- Descriptions must precisely match actual functionality +- Provide tool annotations (readOnlyHint, destructiveHint, idempotentHint, openWorldHint) +- Keep tool operations focused and atomic + +--- + +## Response Formats + +All tools that return data should support multiple formats: + +### JSON Format (`response_format="json"`) +- Machine-readable structured data +- Include all available fields and metadata +- Consistent field names and types +- Use for programmatic processing + +### Markdown Format (`response_format="markdown"`, typically default) +- Human-readable formatted text +- Use headers, lists, and formatting for clarity +- Convert timestamps to human-readable format +- Show display names with IDs in parentheses +- Omit verbose metadata + +--- + +## Pagination + +For tools that list resources: + +- **Always respect the `limit` parameter** +- **Implement pagination**: Use `offset` or cursor-based pagination +- **Return pagination metadata**: Include `has_more`, `next_offset`/`next_cursor`, `total_count` +- **Never load all results into memory**: Especially important for large datasets +- **Default to reasonable limits**: 20-50 items is typical + +Example pagination response: +```json +{ + "total": 150, + "count": 20, + "offset": 0, + "items": [...], + "has_more": true, + "next_offset": 20 +} +``` + +--- + +## Transport Options + +### Streamable HTTP + +**Best for**: Remote servers, web services, multi-client scenarios + +**Characteristics**: +- Bidirectional communication over HTTP +- Supports multiple simultaneous clients +- Can be deployed as a web service +- Enables server-to-client notifications + +**Use when**: +- Serving multiple clients simultaneously +- Deploying as a cloud service +- Integration with web applications + +### stdio + +**Best for**: Local integrations, command-line tools + +**Characteristics**: +- Standard input/output stream communication +- Simple setup, no network configuration needed +- Runs as a subprocess of the client + +**Use when**: +- Building tools for local development environments +- Integrating with desktop applications +- Single-user, single-session scenarios + +**Note**: stdio servers should NOT log to stdout (use stderr for logging) + +### Transport Selection + +| Criterion | stdio | Streamable HTTP | +|-----------|-------|-----------------| +| **Deployment** | Local | Remote | +| **Clients** | Single | Multiple | +| **Complexity** | Low | Medium | +| **Real-time** | No | Yes | + +--- + +## Security Best Practices + +### Authentication and Authorization + +**OAuth 2.1**: +- Use secure OAuth 2.1 with certificates from recognized authorities +- Validate access tokens before processing requests +- Only accept tokens specifically intended for your server + +**API Keys**: +- Store API keys in environment variables, never in code +- Validate keys on server startup +- Provide clear error messages when authentication fails + +### Input Validation + +- Sanitize file paths to prevent directory traversal +- Validate URLs and external identifiers +- Check parameter sizes and ranges +- Prevent command injection in system calls +- Use schema validation (Pydantic/Zod) for all inputs + +### Error Handling + +- Don't expose internal errors to clients +- Log security-relevant errors server-side +- Provide helpful but not revealing error messages +- Clean up resources after errors + +### DNS Rebinding Protection + +For streamable HTTP servers running locally: +- Enable DNS rebinding protection +- Validate the `Origin` header on all incoming connections +- Bind to `127.0.0.1` rather than `0.0.0.0` + +--- + +## Tool Annotations + +Provide annotations to help clients understand tool behavior: + +| Annotation | Type | Default | Description | +|-----------|------|---------|-------------| +| `readOnlyHint` | boolean | false | Tool does not modify its environment | +| `destructiveHint` | boolean | true | Tool may perform destructive updates | +| `idempotentHint` | boolean | false | Repeated calls with same args have no additional effect | +| `openWorldHint` | boolean | true | Tool interacts with external entities | + +**Important**: Annotations are hints, not security guarantees. Clients should not make security-critical decisions based solely on annotations. + +--- + +## Error Handling + +- Use standard JSON-RPC error codes +- Report tool errors within result objects (not protocol-level errors) +- Provide helpful, specific error messages with suggested next steps +- Don't expose internal implementation details +- Clean up resources properly on errors + +Example error handling: +```typescript +try { + const result = performOperation(); + return { content: [{ type: "text", text: result }] }; +} catch (error) { + return { + isError: true, + content: [{ + type: "text", + text: `Error: ${error.message}. Try using filter='active_only' to reduce results.` + }] + }; +} +``` + +--- + +## Testing Requirements + +Comprehensive testing should cover: + +- **Functional testing**: Verify correct execution with valid/invalid inputs +- **Integration testing**: Test interaction with external systems +- **Security testing**: Validate auth, input sanitization, rate limiting +- **Performance testing**: Check behavior under load, timeouts +- **Error handling**: Ensure proper error reporting and cleanup + +--- + +## Documentation Requirements + +- Provide clear documentation of all tools and capabilities +- Include working examples (at least 3 per major feature) +- Document security considerations +- Specify required permissions and access levels +- Document rate limits and performance characteristics diff --git a/contributing/samples/skill_script_demo/skills/mcp-builder/references/node_mcp_server.md b/contributing/samples/skill_script_demo/skills/mcp-builder/references/node_mcp_server.md new file mode 100644 index 0000000000..f6e5df982a --- /dev/null +++ b/contributing/samples/skill_script_demo/skills/mcp-builder/references/node_mcp_server.md @@ -0,0 +1,970 @@ +# Node/TypeScript MCP Server Implementation Guide + +## Overview + +This document provides Node/TypeScript-specific best practices and examples for implementing MCP servers using the MCP TypeScript SDK. It covers project structure, server setup, tool registration patterns, input validation with Zod, error handling, and complete working examples. + +--- + +## Quick Reference + +### Key Imports +```typescript +import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; +import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; +import express from "express"; +import { z } from "zod"; +``` + +### Server Initialization +```typescript +const server = new McpServer({ + name: "service-mcp-server", + version: "1.0.0" +}); +``` + +### Tool Registration Pattern +```typescript +server.registerTool( + "tool_name", + { + title: "Tool Display Name", + description: "What the tool does", + inputSchema: { param: z.string() }, + outputSchema: { result: z.string() } + }, + async ({ param }) => { + const output = { result: `Processed: ${param}` }; + return { + content: [{ type: "text", text: JSON.stringify(output) }], + structuredContent: output // Modern pattern for structured data + }; + } +); +``` + +--- + +## MCP TypeScript SDK + +The official MCP TypeScript SDK provides: +- `McpServer` class for server initialization +- `registerTool` method for tool registration +- Zod schema integration for runtime input validation +- Type-safe tool handler implementations + +**IMPORTANT - Use Modern APIs Only:** +- **DO use**: `server.registerTool()`, `server.registerResource()`, `server.registerPrompt()` +- **DO NOT use**: Old deprecated APIs such as `server.tool()`, `server.setRequestHandler(ListToolsRequestSchema, ...)`, or manual handler registration +- The `register*` methods provide better type safety, automatic schema handling, and are the recommended approach + +See the MCP SDK documentation in the references for complete details. + +## Server Naming Convention + +Node/TypeScript MCP servers must follow this naming pattern: +- **Format**: `{service}-mcp-server` (lowercase with hyphens) +- **Examples**: `github-mcp-server`, `jira-mcp-server`, `stripe-mcp-server` + +The name should be: +- General (not tied to specific features) +- Descriptive of the service/API being integrated +- Easy to infer from the task description +- Without version numbers or dates + +## Project Structure + +Create the following structure for Node/TypeScript MCP servers: + +``` +{service}-mcp-server/ +├── package.json +├── tsconfig.json +├── README.md +├── src/ +│ ├── index.ts # Main entry point with McpServer initialization +│ ├── types.ts # TypeScript type definitions and interfaces +│ ├── tools/ # Tool implementations (one file per domain) +│ ├── services/ # API clients and shared utilities +│ ├── schemas/ # Zod validation schemas +│ └── constants.ts # Shared constants (API_URL, CHARACTER_LIMIT, etc.) +└── dist/ # Built JavaScript files (entry point: dist/index.js) +``` + +## Tool Implementation + +### Tool Naming + +Use snake_case for tool names (e.g., "search_users", "create_project", "get_channel_info") with clear, action-oriented names. + +**Avoid Naming Conflicts**: Include the service context to prevent overlaps: +- Use "slack_send_message" instead of just "send_message" +- Use "github_create_issue" instead of just "create_issue" +- Use "asana_list_tasks" instead of just "list_tasks" + +### Tool Structure + +Tools are registered using the `registerTool` method with the following requirements: +- Use Zod schemas for runtime input validation and type safety +- The `description` field must be explicitly provided - JSDoc comments are NOT automatically extracted +- Explicitly provide `title`, `description`, `inputSchema`, and `annotations` +- The `inputSchema` must be a Zod schema object (not a JSON schema) +- Type all parameters and return values explicitly + +```typescript +import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import { z } from "zod"; + +const server = new McpServer({ + name: "example-mcp", + version: "1.0.0" +}); + +// Zod schema for input validation +const UserSearchInputSchema = z.object({ + query: z.string() + .min(2, "Query must be at least 2 characters") + .max(200, "Query must not exceed 200 characters") + .describe("Search string to match against names/emails"), + limit: z.number() + .int() + .min(1) + .max(100) + .default(20) + .describe("Maximum results to return"), + offset: z.number() + .int() + .min(0) + .default(0) + .describe("Number of results to skip for pagination"), + response_format: z.nativeEnum(ResponseFormat) + .default(ResponseFormat.MARKDOWN) + .describe("Output format: 'markdown' for human-readable or 'json' for machine-readable") +}).strict(); + +// Type definition from Zod schema +type UserSearchInput = z.infer; + +server.registerTool( + "example_search_users", + { + title: "Search Example Users", + description: `Search for users in the Example system by name, email, or team. + +This tool searches across all user profiles in the Example platform, supporting partial matches and various search filters. It does NOT create or modify users, only searches existing ones. + +Args: + - query (string): Search string to match against names/emails + - limit (number): Maximum results to return, between 1-100 (default: 20) + - offset (number): Number of results to skip for pagination (default: 0) + - response_format ('markdown' | 'json'): Output format (default: 'markdown') + +Returns: + For JSON format: Structured data with schema: + { + "total": number, // Total number of matches found + "count": number, // Number of results in this response + "offset": number, // Current pagination offset + "users": [ + { + "id": string, // User ID (e.g., "U123456789") + "name": string, // Full name (e.g., "John Doe") + "email": string, // Email address + "team": string, // Team name (optional) + "active": boolean // Whether user is active + } + ], + "has_more": boolean, // Whether more results are available + "next_offset": number // Offset for next page (if has_more is true) + } + +Examples: + - Use when: "Find all marketing team members" -> params with query="team:marketing" + - Use when: "Search for John's account" -> params with query="john" + - Don't use when: You need to create a user (use example_create_user instead) + +Error Handling: + - Returns "Error: Rate limit exceeded" if too many requests (429 status) + - Returns "No users found matching ''" if search returns empty`, + inputSchema: UserSearchInputSchema, + annotations: { + readOnlyHint: true, + destructiveHint: false, + idempotentHint: true, + openWorldHint: true + } + }, + async (params: UserSearchInput) => { + try { + // Input validation is handled by Zod schema + // Make API request using validated parameters + const data = await makeApiRequest( + "users/search", + "GET", + undefined, + { + q: params.query, + limit: params.limit, + offset: params.offset + } + ); + + const users = data.users || []; + const total = data.total || 0; + + if (!users.length) { + return { + content: [{ + type: "text", + text: `No users found matching '${params.query}'` + }] + }; + } + + // Prepare structured output + const output = { + total, + count: users.length, + offset: params.offset, + users: users.map((user: any) => ({ + id: user.id, + name: user.name, + email: user.email, + ...(user.team ? { team: user.team } : {}), + active: user.active ?? true + })), + has_more: total > params.offset + users.length, + ...(total > params.offset + users.length ? { + next_offset: params.offset + users.length + } : {}) + }; + + // Format text representation based on requested format + let textContent: string; + if (params.response_format === ResponseFormat.MARKDOWN) { + const lines = [`# User Search Results: '${params.query}'`, "", + `Found ${total} users (showing ${users.length})`, ""]; + for (const user of users) { + lines.push(`## ${user.name} (${user.id})`); + lines.push(`- **Email**: ${user.email}`); + if (user.team) lines.push(`- **Team**: ${user.team}`); + lines.push(""); + } + textContent = lines.join("\n"); + } else { + textContent = JSON.stringify(output, null, 2); + } + + return { + content: [{ type: "text", text: textContent }], + structuredContent: output // Modern pattern for structured data + }; + } catch (error) { + return { + content: [{ + type: "text", + text: handleApiError(error) + }] + }; + } + } +); +``` + +## Zod Schemas for Input Validation + +Zod provides runtime type validation: + +```typescript +import { z } from "zod"; + +// Basic schema with validation +const CreateUserSchema = z.object({ + name: z.string() + .min(1, "Name is required") + .max(100, "Name must not exceed 100 characters"), + email: z.string() + .email("Invalid email format"), + age: z.number() + .int("Age must be a whole number") + .min(0, "Age cannot be negative") + .max(150, "Age cannot be greater than 150") +}).strict(); // Use .strict() to forbid extra fields + +// Enums +enum ResponseFormat { + MARKDOWN = "markdown", + JSON = "json" +} + +const SearchSchema = z.object({ + response_format: z.nativeEnum(ResponseFormat) + .default(ResponseFormat.MARKDOWN) + .describe("Output format") +}); + +// Optional fields with defaults +const PaginationSchema = z.object({ + limit: z.number() + .int() + .min(1) + .max(100) + .default(20) + .describe("Maximum results to return"), + offset: z.number() + .int() + .min(0) + .default(0) + .describe("Number of results to skip") +}); +``` + +## Response Format Options + +Support multiple output formats for flexibility: + +```typescript +enum ResponseFormat { + MARKDOWN = "markdown", + JSON = "json" +} + +const inputSchema = z.object({ + query: z.string(), + response_format: z.nativeEnum(ResponseFormat) + .default(ResponseFormat.MARKDOWN) + .describe("Output format: 'markdown' for human-readable or 'json' for machine-readable") +}); +``` + +**Markdown format**: +- Use headers, lists, and formatting for clarity +- Convert timestamps to human-readable format +- Show display names with IDs in parentheses +- Omit verbose metadata +- Group related information logically + +**JSON format**: +- Return complete, structured data suitable for programmatic processing +- Include all available fields and metadata +- Use consistent field names and types + +## Pagination Implementation + +For tools that list resources: + +```typescript +const ListSchema = z.object({ + limit: z.number().int().min(1).max(100).default(20), + offset: z.number().int().min(0).default(0) +}); + +async function listItems(params: z.infer) { + const data = await apiRequest(params.limit, params.offset); + + const response = { + total: data.total, + count: data.items.length, + offset: params.offset, + items: data.items, + has_more: data.total > params.offset + data.items.length, + next_offset: data.total > params.offset + data.items.length + ? params.offset + data.items.length + : undefined + }; + + return JSON.stringify(response, null, 2); +} +``` + +## Character Limits and Truncation + +Add a CHARACTER_LIMIT constant to prevent overwhelming responses: + +```typescript +// At module level in constants.ts +export const CHARACTER_LIMIT = 25000; // Maximum response size in characters + +async function searchTool(params: SearchInput) { + let result = generateResponse(data); + + // Check character limit and truncate if needed + if (result.length > CHARACTER_LIMIT) { + const truncatedData = data.slice(0, Math.max(1, data.length / 2)); + response.data = truncatedData; + response.truncated = true; + response.truncation_message = + `Response truncated from ${data.length} to ${truncatedData.length} items. ` + + `Use 'offset' parameter or add filters to see more results.`; + result = JSON.stringify(response, null, 2); + } + + return result; +} +``` + +## Error Handling + +Provide clear, actionable error messages: + +```typescript +import axios, { AxiosError } from "axios"; + +function handleApiError(error: unknown): string { + if (error instanceof AxiosError) { + if (error.response) { + switch (error.response.status) { + case 404: + return "Error: Resource not found. Please check the ID is correct."; + case 403: + return "Error: Permission denied. You don't have access to this resource."; + case 429: + return "Error: Rate limit exceeded. Please wait before making more requests."; + default: + return `Error: API request failed with status ${error.response.status}`; + } + } else if (error.code === "ECONNABORTED") { + return "Error: Request timed out. Please try again."; + } + } + return `Error: Unexpected error occurred: ${error instanceof Error ? error.message : String(error)}`; +} +``` + +## Shared Utilities + +Extract common functionality into reusable functions: + +```typescript +// Shared API request function +async function makeApiRequest( + endpoint: string, + method: "GET" | "POST" | "PUT" | "DELETE" = "GET", + data?: any, + params?: any +): Promise { + try { + const response = await axios({ + method, + url: `${API_BASE_URL}/${endpoint}`, + data, + params, + timeout: 30000, + headers: { + "Content-Type": "application/json", + "Accept": "application/json" + } + }); + return response.data; + } catch (error) { + throw error; + } +} +``` + +## Async/Await Best Practices + +Always use async/await for network requests and I/O operations: + +```typescript +// Good: Async network request +async function fetchData(resourceId: string): Promise { + const response = await axios.get(`${API_URL}/resource/${resourceId}`); + return response.data; +} + +// Bad: Promise chains +function fetchData(resourceId: string): Promise { + return axios.get(`${API_URL}/resource/${resourceId}`) + .then(response => response.data); // Harder to read and maintain +} +``` + +## TypeScript Best Practices + +1. **Use Strict TypeScript**: Enable strict mode in tsconfig.json +2. **Define Interfaces**: Create clear interface definitions for all data structures +3. **Avoid `any`**: Use proper types or `unknown` instead of `any` +4. **Zod for Runtime Validation**: Use Zod schemas to validate external data +5. **Type Guards**: Create type guard functions for complex type checking +6. **Error Handling**: Always use try-catch with proper error type checking +7. **Null Safety**: Use optional chaining (`?.`) and nullish coalescing (`??`) + +```typescript +// Good: Type-safe with Zod and interfaces +interface UserResponse { + id: string; + name: string; + email: string; + team?: string; + active: boolean; +} + +const UserSchema = z.object({ + id: z.string(), + name: z.string(), + email: z.string().email(), + team: z.string().optional(), + active: z.boolean() +}); + +type User = z.infer; + +async function getUser(id: string): Promise { + const data = await apiCall(`/users/${id}`); + return UserSchema.parse(data); // Runtime validation +} + +// Bad: Using any +async function getUser(id: string): Promise { + return await apiCall(`/users/${id}`); // No type safety +} +``` + +## Package Configuration + +### package.json + +```json +{ + "name": "{service}-mcp-server", + "version": "1.0.0", + "description": "MCP server for {Service} API integration", + "type": "module", + "main": "dist/index.js", + "scripts": { + "start": "node dist/index.js", + "dev": "tsx watch src/index.ts", + "build": "tsc", + "clean": "rm -rf dist" + }, + "engines": { + "node": ">=18" + }, + "dependencies": { + "@modelcontextprotocol/sdk": "^1.6.1", + "axios": "^1.7.9", + "zod": "^3.23.8" + }, + "devDependencies": { + "@types/node": "^22.10.0", + "tsx": "^4.19.2", + "typescript": "^5.7.2" + } +} +``` + +### tsconfig.json + +```json +{ + "compilerOptions": { + "target": "ES2022", + "module": "Node16", + "moduleResolution": "Node16", + "lib": ["ES2022"], + "outDir": "./dist", + "rootDir": "./src", + "strict": true, + "esModuleInterop": true, + "skipLibCheck": true, + "forceConsistentCasingInFileNames": true, + "declaration": true, + "declarationMap": true, + "sourceMap": true, + "allowSyntheticDefaultImports": true + }, + "include": ["src/**/*"], + "exclude": ["node_modules", "dist"] +} +``` + +## Complete Example + +```typescript +#!/usr/bin/env node +/** + * MCP Server for Example Service. + * + * This server provides tools to interact with Example API, including user search, + * project management, and data export capabilities. + */ + +import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; +import { z } from "zod"; +import axios, { AxiosError } from "axios"; + +// Constants +const API_BASE_URL = "https://api.example.com/v1"; +const CHARACTER_LIMIT = 25000; + +// Enums +enum ResponseFormat { + MARKDOWN = "markdown", + JSON = "json" +} + +// Zod schemas +const UserSearchInputSchema = z.object({ + query: z.string() + .min(2, "Query must be at least 2 characters") + .max(200, "Query must not exceed 200 characters") + .describe("Search string to match against names/emails"), + limit: z.number() + .int() + .min(1) + .max(100) + .default(20) + .describe("Maximum results to return"), + offset: z.number() + .int() + .min(0) + .default(0) + .describe("Number of results to skip for pagination"), + response_format: z.nativeEnum(ResponseFormat) + .default(ResponseFormat.MARKDOWN) + .describe("Output format: 'markdown' for human-readable or 'json' for machine-readable") +}).strict(); + +type UserSearchInput = z.infer; + +// Shared utility functions +async function makeApiRequest( + endpoint: string, + method: "GET" | "POST" | "PUT" | "DELETE" = "GET", + data?: any, + params?: any +): Promise { + try { + const response = await axios({ + method, + url: `${API_BASE_URL}/${endpoint}`, + data, + params, + timeout: 30000, + headers: { + "Content-Type": "application/json", + "Accept": "application/json" + } + }); + return response.data; + } catch (error) { + throw error; + } +} + +function handleApiError(error: unknown): string { + if (error instanceof AxiosError) { + if (error.response) { + switch (error.response.status) { + case 404: + return "Error: Resource not found. Please check the ID is correct."; + case 403: + return "Error: Permission denied. You don't have access to this resource."; + case 429: + return "Error: Rate limit exceeded. Please wait before making more requests."; + default: + return `Error: API request failed with status ${error.response.status}`; + } + } else if (error.code === "ECONNABORTED") { + return "Error: Request timed out. Please try again."; + } + } + return `Error: Unexpected error occurred: ${error instanceof Error ? error.message : String(error)}`; +} + +// Create MCP server instance +const server = new McpServer({ + name: "example-mcp", + version: "1.0.0" +}); + +// Register tools +server.registerTool( + "example_search_users", + { + title: "Search Example Users", + description: `[Full description as shown above]`, + inputSchema: UserSearchInputSchema, + annotations: { + readOnlyHint: true, + destructiveHint: false, + idempotentHint: true, + openWorldHint: true + } + }, + async (params: UserSearchInput) => { + // Implementation as shown above + } +); + +// Main function +// For stdio (local): +async function runStdio() { + if (!process.env.EXAMPLE_API_KEY) { + console.error("ERROR: EXAMPLE_API_KEY environment variable is required"); + process.exit(1); + } + + const transport = new StdioServerTransport(); + await server.connect(transport); + console.error("MCP server running via stdio"); +} + +// For streamable HTTP (remote): +async function runHTTP() { + if (!process.env.EXAMPLE_API_KEY) { + console.error("ERROR: EXAMPLE_API_KEY environment variable is required"); + process.exit(1); + } + + const app = express(); + app.use(express.json()); + + app.post('/mcp', async (req, res) => { + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: undefined, + enableJsonResponse: true + }); + res.on('close', () => transport.close()); + await server.connect(transport); + await transport.handleRequest(req, res, req.body); + }); + + const port = parseInt(process.env.PORT || '3000'); + app.listen(port, () => { + console.error(`MCP server running on http://localhost:${port}/mcp`); + }); +} + +// Choose transport based on environment +const transport = process.env.TRANSPORT || 'stdio'; +if (transport === 'http') { + runHTTP().catch(error => { + console.error("Server error:", error); + process.exit(1); + }); +} else { + runStdio().catch(error => { + console.error("Server error:", error); + process.exit(1); + }); +} +``` + +--- + +## Advanced MCP Features + +### Resource Registration + +Expose data as resources for efficient, URI-based access: + +```typescript +import { ResourceTemplate } from "@modelcontextprotocol/sdk/types.js"; + +// Register a resource with URI template +server.registerResource( + { + uri: "file://documents/{name}", + name: "Document Resource", + description: "Access documents by name", + mimeType: "text/plain" + }, + async (uri: string) => { + // Extract parameter from URI + const match = uri.match(/^file:\/\/documents\/(.+)$/); + if (!match) { + throw new Error("Invalid URI format"); + } + + const documentName = match[1]; + const content = await loadDocument(documentName); + + return { + contents: [{ + uri, + mimeType: "text/plain", + text: content + }] + }; + } +); + +// List available resources dynamically +server.registerResourceList(async () => { + const documents = await getAvailableDocuments(); + return { + resources: documents.map(doc => ({ + uri: `file://documents/${doc.name}`, + name: doc.name, + mimeType: "text/plain", + description: doc.description + })) + }; +}); +``` + +**When to use Resources vs Tools:** +- **Resources**: For data access with simple URI-based parameters +- **Tools**: For complex operations requiring validation and business logic +- **Resources**: When data is relatively static or template-based +- **Tools**: When operations have side effects or complex workflows + +### Transport Options + +The TypeScript SDK supports two main transport mechanisms: + +#### Streamable HTTP (Recommended for Remote Servers) + +```typescript +import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; +import express from "express"; + +const app = express(); +app.use(express.json()); + +app.post('/mcp', async (req, res) => { + // Create new transport for each request (stateless, prevents request ID collisions) + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: undefined, + enableJsonResponse: true + }); + + res.on('close', () => transport.close()); + + await server.connect(transport); + await transport.handleRequest(req, res, req.body); +}); + +app.listen(3000); +``` + +#### stdio (For Local Integrations) + +```typescript +import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; + +const transport = new StdioServerTransport(); +await server.connect(transport); +``` + +**Transport selection:** +- **Streamable HTTP**: Web services, remote access, multiple clients +- **stdio**: Command-line tools, local development, subprocess integration + +### Notification Support + +Notify clients when server state changes: + +```typescript +// Notify when tools list changes +server.notification({ + method: "notifications/tools/list_changed" +}); + +// Notify when resources change +server.notification({ + method: "notifications/resources/list_changed" +}); +``` + +Use notifications sparingly - only when server capabilities genuinely change. + +--- + +## Code Best Practices + +### Code Composability and Reusability + +Your implementation MUST prioritize composability and code reuse: + +1. **Extract Common Functionality**: + - Create reusable helper functions for operations used across multiple tools + - Build shared API clients for HTTP requests instead of duplicating code + - Centralize error handling logic in utility functions + - Extract business logic into dedicated functions that can be composed + - Extract shared markdown or JSON field selection & formatting functionality + +2. **Avoid Duplication**: + - NEVER copy-paste similar code between tools + - If you find yourself writing similar logic twice, extract it into a function + - Common operations like pagination, filtering, field selection, and formatting should be shared + - Authentication/authorization logic should be centralized + +## Building and Running + +Always build your TypeScript code before running: + +```bash +# Build the project +npm run build + +# Run the server +npm start + +# Development with auto-reload +npm run dev +``` + +Always ensure `npm run build` completes successfully before considering the implementation complete. + +## Quality Checklist + +Before finalizing your Node/TypeScript MCP server implementation, ensure: + +### Strategic Design +- [ ] Tools enable complete workflows, not just API endpoint wrappers +- [ ] Tool names reflect natural task subdivisions +- [ ] Response formats optimize for agent context efficiency +- [ ] Human-readable identifiers used where appropriate +- [ ] Error messages guide agents toward correct usage + +### Implementation Quality +- [ ] FOCUSED IMPLEMENTATION: Most important and valuable tools implemented +- [ ] All tools registered using `registerTool` with complete configuration +- [ ] All tools include `title`, `description`, `inputSchema`, and `annotations` +- [ ] Annotations correctly set (readOnlyHint, destructiveHint, idempotentHint, openWorldHint) +- [ ] All tools use Zod schemas for runtime input validation with `.strict()` enforcement +- [ ] All Zod schemas have proper constraints and descriptive error messages +- [ ] All tools have comprehensive descriptions with explicit input/output types +- [ ] Descriptions include return value examples and complete schema documentation +- [ ] Error messages are clear, actionable, and educational + +### TypeScript Quality +- [ ] TypeScript interfaces are defined for all data structures +- [ ] Strict TypeScript is enabled in tsconfig.json +- [ ] No use of `any` type - use `unknown` or proper types instead +- [ ] All async functions have explicit Promise return types +- [ ] Error handling uses proper type guards (e.g., `axios.isAxiosError`, `z.ZodError`) + +### Advanced Features (where applicable) +- [ ] Resources registered for appropriate data endpoints +- [ ] Appropriate transport configured (stdio or streamable HTTP) +- [ ] Notifications implemented for dynamic server capabilities +- [ ] Type-safe with SDK interfaces + +### Project Configuration +- [ ] Package.json includes all necessary dependencies +- [ ] Build script produces working JavaScript in dist/ directory +- [ ] Main entry point is properly configured as dist/index.js +- [ ] Server name follows format: `{service}-mcp-server` +- [ ] tsconfig.json properly configured with strict mode + +### Code Quality +- [ ] Pagination is properly implemented where applicable +- [ ] Large responses check CHARACTER_LIMIT constant and truncate with clear messages +- [ ] Filtering options are provided for potentially large result sets +- [ ] All network operations handle timeouts and connection errors gracefully +- [ ] Common functionality is extracted into reusable functions +- [ ] Return types are consistent across similar operations + +### Testing and Build +- [ ] `npm run build` completes successfully without errors +- [ ] dist/index.js created and executable +- [ ] Server runs: `node dist/index.js --help` +- [ ] All imports resolve correctly +- [ ] Sample tool calls work as expected \ No newline at end of file diff --git a/contributing/samples/skill_script_demo/skills/mcp-builder/references/python_mcp_server.md b/contributing/samples/skill_script_demo/skills/mcp-builder/references/python_mcp_server.md new file mode 100644 index 0000000000..cf7ec996d2 --- /dev/null +++ b/contributing/samples/skill_script_demo/skills/mcp-builder/references/python_mcp_server.md @@ -0,0 +1,719 @@ +# Python MCP Server Implementation Guide + +## Overview + +This document provides Python-specific best practices and examples for implementing MCP servers using the MCP Python SDK. It covers server setup, tool registration patterns, input validation with Pydantic, error handling, and complete working examples. + +--- + +## Quick Reference + +### Key Imports +```python +from mcp.server.fastmcp import FastMCP +from pydantic import BaseModel, Field, field_validator, ConfigDict +from typing import Optional, List, Dict, Any +from enum import Enum +import httpx +``` + +### Server Initialization +```python +mcp = FastMCP("service_mcp") +``` + +### Tool Registration Pattern +```python +@mcp.tool(name="tool_name", annotations={...}) +async def tool_function(params: InputModel) -> str: + # Implementation + pass +``` + +--- + +## MCP Python SDK and FastMCP + +The official MCP Python SDK provides FastMCP, a high-level framework for building MCP servers. It provides: +- Automatic description and inputSchema generation from function signatures and docstrings +- Pydantic model integration for input validation +- Decorator-based tool registration with `@mcp.tool` + +**For complete SDK documentation, use WebFetch to load:** +`https://raw.githubusercontent.com/modelcontextprotocol/python-sdk/main/README.md` + +## Server Naming Convention + +Python MCP servers must follow this naming pattern: +- **Format**: `{service}_mcp` (lowercase with underscores) +- **Examples**: `github_mcp`, `jira_mcp`, `stripe_mcp` + +The name should be: +- General (not tied to specific features) +- Descriptive of the service/API being integrated +- Easy to infer from the task description +- Without version numbers or dates + +## Tool Implementation + +### Tool Naming + +Use snake_case for tool names (e.g., "search_users", "create_project", "get_channel_info") with clear, action-oriented names. + +**Avoid Naming Conflicts**: Include the service context to prevent overlaps: +- Use "slack_send_message" instead of just "send_message" +- Use "github_create_issue" instead of just "create_issue" +- Use "asana_list_tasks" instead of just "list_tasks" + +### Tool Structure with FastMCP + +Tools are defined using the `@mcp.tool` decorator with Pydantic models for input validation: + +```python +from pydantic import BaseModel, Field, ConfigDict +from mcp.server.fastmcp import FastMCP + +# Initialize the MCP server +mcp = FastMCP("example_mcp") + +# Define Pydantic model for input validation +class ServiceToolInput(BaseModel): + '''Input model for service tool operation.''' + model_config = ConfigDict( + str_strip_whitespace=True, # Auto-strip whitespace from strings + validate_assignment=True, # Validate on assignment + extra='forbid' # Forbid extra fields + ) + + param1: str = Field(..., description="First parameter description (e.g., 'user123', 'project-abc')", min_length=1, max_length=100) + param2: Optional[int] = Field(default=None, description="Optional integer parameter with constraints", ge=0, le=1000) + tags: Optional[List[str]] = Field(default_factory=list, description="List of tags to apply", max_items=10) + +@mcp.tool( + name="service_tool_name", + annotations={ + "title": "Human-Readable Tool Title", + "readOnlyHint": True, # Tool does not modify environment + "destructiveHint": False, # Tool does not perform destructive operations + "idempotentHint": True, # Repeated calls have no additional effect + "openWorldHint": False # Tool does not interact with external entities + } +) +async def service_tool_name(params: ServiceToolInput) -> str: + '''Tool description automatically becomes the 'description' field. + + This tool performs a specific operation on the service. It validates all inputs + using the ServiceToolInput Pydantic model before processing. + + Args: + params (ServiceToolInput): Validated input parameters containing: + - param1 (str): First parameter description + - param2 (Optional[int]): Optional parameter with default + - tags (Optional[List[str]]): List of tags + + Returns: + str: JSON-formatted response containing operation results + ''' + # Implementation here + pass +``` + +## Pydantic v2 Key Features + +- Use `model_config` instead of nested `Config` class +- Use `field_validator` instead of deprecated `validator` +- Use `model_dump()` instead of deprecated `dict()` +- Validators require `@classmethod` decorator +- Type hints are required for validator methods + +```python +from pydantic import BaseModel, Field, field_validator, ConfigDict + +class CreateUserInput(BaseModel): + model_config = ConfigDict( + str_strip_whitespace=True, + validate_assignment=True + ) + + name: str = Field(..., description="User's full name", min_length=1, max_length=100) + email: str = Field(..., description="User's email address", pattern=r'^[\w\.-]+@[\w\.-]+\.\w+$') + age: int = Field(..., description="User's age", ge=0, le=150) + + @field_validator('email') + @classmethod + def validate_email(cls, v: str) -> str: + if not v.strip(): + raise ValueError("Email cannot be empty") + return v.lower() +``` + +## Response Format Options + +Support multiple output formats for flexibility: + +```python +from enum import Enum + +class ResponseFormat(str, Enum): + '''Output format for tool responses.''' + MARKDOWN = "markdown" + JSON = "json" + +class UserSearchInput(BaseModel): + query: str = Field(..., description="Search query") + response_format: ResponseFormat = Field( + default=ResponseFormat.MARKDOWN, + description="Output format: 'markdown' for human-readable or 'json' for machine-readable" + ) +``` + +**Markdown format**: +- Use headers, lists, and formatting for clarity +- Convert timestamps to human-readable format (e.g., "2024-01-15 10:30:00 UTC" instead of epoch) +- Show display names with IDs in parentheses (e.g., "@john.doe (U123456)") +- Omit verbose metadata (e.g., show only one profile image URL, not all sizes) +- Group related information logically + +**JSON format**: +- Return complete, structured data suitable for programmatic processing +- Include all available fields and metadata +- Use consistent field names and types + +## Pagination Implementation + +For tools that list resources: + +```python +class ListInput(BaseModel): + limit: Optional[int] = Field(default=20, description="Maximum results to return", ge=1, le=100) + offset: Optional[int] = Field(default=0, description="Number of results to skip for pagination", ge=0) + +async def list_items(params: ListInput) -> str: + # Make API request with pagination + data = await api_request(limit=params.limit, offset=params.offset) + + # Return pagination info + response = { + "total": data["total"], + "count": len(data["items"]), + "offset": params.offset, + "items": data["items"], + "has_more": data["total"] > params.offset + len(data["items"]), + "next_offset": params.offset + len(data["items"]) if data["total"] > params.offset + len(data["items"]) else None + } + return json.dumps(response, indent=2) +``` + +## Error Handling + +Provide clear, actionable error messages: + +```python +def _handle_api_error(e: Exception) -> str: + '''Consistent error formatting across all tools.''' + if isinstance(e, httpx.HTTPStatusError): + if e.response.status_code == 404: + return "Error: Resource not found. Please check the ID is correct." + elif e.response.status_code == 403: + return "Error: Permission denied. You don't have access to this resource." + elif e.response.status_code == 429: + return "Error: Rate limit exceeded. Please wait before making more requests." + return f"Error: API request failed with status {e.response.status_code}" + elif isinstance(e, httpx.TimeoutException): + return "Error: Request timed out. Please try again." + return f"Error: Unexpected error occurred: {type(e).__name__}" +``` + +## Shared Utilities + +Extract common functionality into reusable functions: + +```python +# Shared API request function +async def _make_api_request(endpoint: str, method: str = "GET", **kwargs) -> dict: + '''Reusable function for all API calls.''' + async with httpx.AsyncClient() as client: + response = await client.request( + method, + f"{API_BASE_URL}/{endpoint}", + timeout=30.0, + **kwargs + ) + response.raise_for_status() + return response.json() +``` + +## Async/Await Best Practices + +Always use async/await for network requests and I/O operations: + +```python +# Good: Async network request +async def fetch_data(resource_id: str) -> dict: + async with httpx.AsyncClient() as client: + response = await client.get(f"{API_URL}/resource/{resource_id}") + response.raise_for_status() + return response.json() + +# Bad: Synchronous request +def fetch_data(resource_id: str) -> dict: + response = requests.get(f"{API_URL}/resource/{resource_id}") # Blocks + return response.json() +``` + +## Type Hints + +Use type hints throughout: + +```python +from typing import Optional, List, Dict, Any + +async def get_user(user_id: str) -> Dict[str, Any]: + data = await fetch_user(user_id) + return {"id": data["id"], "name": data["name"]} +``` + +## Tool Docstrings + +Every tool must have comprehensive docstrings with explicit type information: + +```python +async def search_users(params: UserSearchInput) -> str: + ''' + Search for users in the Example system by name, email, or team. + + This tool searches across all user profiles in the Example platform, + supporting partial matches and various search filters. It does NOT + create or modify users, only searches existing ones. + + Args: + params (UserSearchInput): Validated input parameters containing: + - query (str): Search string to match against names/emails (e.g., "john", "@example.com", "team:marketing") + - limit (Optional[int]): Maximum results to return, between 1-100 (default: 20) + - offset (Optional[int]): Number of results to skip for pagination (default: 0) + + Returns: + str: JSON-formatted string containing search results with the following schema: + + Success response: + { + "total": int, # Total number of matches found + "count": int, # Number of results in this response + "offset": int, # Current pagination offset + "users": [ + { + "id": str, # User ID (e.g., "U123456789") + "name": str, # Full name (e.g., "John Doe") + "email": str, # Email address (e.g., "john@example.com") + "team": str # Team name (e.g., "Marketing") - optional + } + ] + } + + Error response: + "Error: " or "No users found matching ''" + + Examples: + - Use when: "Find all marketing team members" -> params with query="team:marketing" + - Use when: "Search for John's account" -> params with query="john" + - Don't use when: You need to create a user (use example_create_user instead) + - Don't use when: You have a user ID and need full details (use example_get_user instead) + + Error Handling: + - Input validation errors are handled by Pydantic model + - Returns "Error: Rate limit exceeded" if too many requests (429 status) + - Returns "Error: Invalid API authentication" if API key is invalid (401 status) + - Returns formatted list of results or "No users found matching 'query'" + ''' +``` + +## Complete Example + +See below for a complete Python MCP server example: + +```python +#!/usr/bin/env python3 +''' +MCP Server for Example Service. + +This server provides tools to interact with Example API, including user search, +project management, and data export capabilities. +''' + +from typing import Optional, List, Dict, Any +from enum import Enum +import httpx +from pydantic import BaseModel, Field, field_validator, ConfigDict +from mcp.server.fastmcp import FastMCP + +# Initialize the MCP server +mcp = FastMCP("example_mcp") + +# Constants +API_BASE_URL = "https://api.example.com/v1" + +# Enums +class ResponseFormat(str, Enum): + '''Output format for tool responses.''' + MARKDOWN = "markdown" + JSON = "json" + +# Pydantic Models for Input Validation +class UserSearchInput(BaseModel): + '''Input model for user search operations.''' + model_config = ConfigDict( + str_strip_whitespace=True, + validate_assignment=True + ) + + query: str = Field(..., description="Search string to match against names/emails", min_length=2, max_length=200) + limit: Optional[int] = Field(default=20, description="Maximum results to return", ge=1, le=100) + offset: Optional[int] = Field(default=0, description="Number of results to skip for pagination", ge=0) + response_format: ResponseFormat = Field(default=ResponseFormat.MARKDOWN, description="Output format") + + @field_validator('query') + @classmethod + def validate_query(cls, v: str) -> str: + if not v.strip(): + raise ValueError("Query cannot be empty or whitespace only") + return v.strip() + +# Shared utility functions +async def _make_api_request(endpoint: str, method: str = "GET", **kwargs) -> dict: + '''Reusable function for all API calls.''' + async with httpx.AsyncClient() as client: + response = await client.request( + method, + f"{API_BASE_URL}/{endpoint}", + timeout=30.0, + **kwargs + ) + response.raise_for_status() + return response.json() + +def _handle_api_error(e: Exception) -> str: + '''Consistent error formatting across all tools.''' + if isinstance(e, httpx.HTTPStatusError): + if e.response.status_code == 404: + return "Error: Resource not found. Please check the ID is correct." + elif e.response.status_code == 403: + return "Error: Permission denied. You don't have access to this resource." + elif e.response.status_code == 429: + return "Error: Rate limit exceeded. Please wait before making more requests." + return f"Error: API request failed with status {e.response.status_code}" + elif isinstance(e, httpx.TimeoutException): + return "Error: Request timed out. Please try again." + return f"Error: Unexpected error occurred: {type(e).__name__}" + +# Tool definitions +@mcp.tool( + name="example_search_users", + annotations={ + "title": "Search Example Users", + "readOnlyHint": True, + "destructiveHint": False, + "idempotentHint": True, + "openWorldHint": True + } +) +async def example_search_users(params: UserSearchInput) -> str: + '''Search for users in the Example system by name, email, or team. + + [Full docstring as shown above] + ''' + try: + # Make API request using validated parameters + data = await _make_api_request( + "users/search", + params={ + "q": params.query, + "limit": params.limit, + "offset": params.offset + } + ) + + users = data.get("users", []) + total = data.get("total", 0) + + if not users: + return f"No users found matching '{params.query}'" + + # Format response based on requested format + if params.response_format == ResponseFormat.MARKDOWN: + lines = [f"# User Search Results: '{params.query}'", ""] + lines.append(f"Found {total} users (showing {len(users)})") + lines.append("") + + for user in users: + lines.append(f"## {user['name']} ({user['id']})") + lines.append(f"- **Email**: {user['email']}") + if user.get('team'): + lines.append(f"- **Team**: {user['team']}") + lines.append("") + + return "\n".join(lines) + + else: + # Machine-readable JSON format + import json + response = { + "total": total, + "count": len(users), + "offset": params.offset, + "users": users + } + return json.dumps(response, indent=2) + + except Exception as e: + return _handle_api_error(e) + +if __name__ == "__main__": + mcp.run() +``` + +--- + +## Advanced FastMCP Features + +### Context Parameter Injection + +FastMCP can automatically inject a `Context` parameter into tools for advanced capabilities like logging, progress reporting, resource reading, and user interaction: + +```python +from mcp.server.fastmcp import FastMCP, Context + +mcp = FastMCP("example_mcp") + +@mcp.tool() +async def advanced_search(query: str, ctx: Context) -> str: + '''Advanced tool with context access for logging and progress.''' + + # Report progress for long operations + await ctx.report_progress(0.25, "Starting search...") + + # Log information for debugging + await ctx.log_info("Processing query", {"query": query, "timestamp": datetime.now()}) + + # Perform search + results = await search_api(query) + await ctx.report_progress(0.75, "Formatting results...") + + # Access server configuration + server_name = ctx.fastmcp.name + + return format_results(results) + +@mcp.tool() +async def interactive_tool(resource_id: str, ctx: Context) -> str: + '''Tool that can request additional input from users.''' + + # Request sensitive information when needed + api_key = await ctx.elicit( + prompt="Please provide your API key:", + input_type="password" + ) + + # Use the provided key + return await api_call(resource_id, api_key) +``` + +**Context capabilities:** +- `ctx.report_progress(progress, message)` - Report progress for long operations +- `ctx.log_info(message, data)` / `ctx.log_error()` / `ctx.log_debug()` - Logging +- `ctx.elicit(prompt, input_type)` - Request input from users +- `ctx.fastmcp.name` - Access server configuration +- `ctx.read_resource(uri)` - Read MCP resources + +### Resource Registration + +Expose data as resources for efficient, template-based access: + +```python +@mcp.resource("file://documents/{name}") +async def get_document(name: str) -> str: + '''Expose documents as MCP resources. + + Resources are useful for static or semi-static data that doesn't + require complex parameters. They use URI templates for flexible access. + ''' + document_path = f"./docs/{name}" + with open(document_path, "r") as f: + return f.read() + +@mcp.resource("config://settings/{key}") +async def get_setting(key: str, ctx: Context) -> str: + '''Expose configuration as resources with context.''' + settings = await load_settings() + return json.dumps(settings.get(key, {})) +``` + +**When to use Resources vs Tools:** +- **Resources**: For data access with simple parameters (URI templates) +- **Tools**: For complex operations with validation and business logic + +### Structured Output Types + +FastMCP supports multiple return types beyond strings: + +```python +from typing import TypedDict +from dataclasses import dataclass +from pydantic import BaseModel + +# TypedDict for structured returns +class UserData(TypedDict): + id: str + name: str + email: str + +@mcp.tool() +async def get_user_typed(user_id: str) -> UserData: + '''Returns structured data - FastMCP handles serialization.''' + return {"id": user_id, "name": "John Doe", "email": "john@example.com"} + +# Pydantic models for complex validation +class DetailedUser(BaseModel): + id: str + name: str + email: str + created_at: datetime + metadata: Dict[str, Any] + +@mcp.tool() +async def get_user_detailed(user_id: str) -> DetailedUser: + '''Returns Pydantic model - automatically generates schema.''' + user = await fetch_user(user_id) + return DetailedUser(**user) +``` + +### Lifespan Management + +Initialize resources that persist across requests: + +```python +from contextlib import asynccontextmanager + +@asynccontextmanager +async def app_lifespan(): + '''Manage resources that live for the server's lifetime.''' + # Initialize connections, load config, etc. + db = await connect_to_database() + config = load_configuration() + + # Make available to all tools + yield {"db": db, "config": config} + + # Cleanup on shutdown + await db.close() + +mcp = FastMCP("example_mcp", lifespan=app_lifespan) + +@mcp.tool() +async def query_data(query: str, ctx: Context) -> str: + '''Access lifespan resources through context.''' + db = ctx.request_context.lifespan_state["db"] + results = await db.query(query) + return format_results(results) +``` + +### Transport Options + +FastMCP supports two main transport mechanisms: + +```python +# stdio transport (for local tools) - default +if __name__ == "__main__": + mcp.run() + +# Streamable HTTP transport (for remote servers) +if __name__ == "__main__": + mcp.run(transport="streamable_http", port=8000) +``` + +**Transport selection:** +- **stdio**: Command-line tools, local integrations, subprocess execution +- **Streamable HTTP**: Web services, remote access, multiple clients + +--- + +## Code Best Practices + +### Code Composability and Reusability + +Your implementation MUST prioritize composability and code reuse: + +1. **Extract Common Functionality**: + - Create reusable helper functions for operations used across multiple tools + - Build shared API clients for HTTP requests instead of duplicating code + - Centralize error handling logic in utility functions + - Extract business logic into dedicated functions that can be composed + - Extract shared markdown or JSON field selection & formatting functionality + +2. **Avoid Duplication**: + - NEVER copy-paste similar code between tools + - If you find yourself writing similar logic twice, extract it into a function + - Common operations like pagination, filtering, field selection, and formatting should be shared + - Authentication/authorization logic should be centralized + +### Python-Specific Best Practices + +1. **Use Type Hints**: Always include type annotations for function parameters and return values +2. **Pydantic Models**: Define clear Pydantic models for all input validation +3. **Avoid Manual Validation**: Let Pydantic handle input validation with constraints +4. **Proper Imports**: Group imports (standard library, third-party, local) +5. **Error Handling**: Use specific exception types (httpx.HTTPStatusError, not generic Exception) +6. **Async Context Managers**: Use `async with` for resources that need cleanup +7. **Constants**: Define module-level constants in UPPER_CASE + +## Quality Checklist + +Before finalizing your Python MCP server implementation, ensure: + +### Strategic Design +- [ ] Tools enable complete workflows, not just API endpoint wrappers +- [ ] Tool names reflect natural task subdivisions +- [ ] Response formats optimize for agent context efficiency +- [ ] Human-readable identifiers used where appropriate +- [ ] Error messages guide agents toward correct usage + +### Implementation Quality +- [ ] FOCUSED IMPLEMENTATION: Most important and valuable tools implemented +- [ ] All tools have descriptive names and documentation +- [ ] Return types are consistent across similar operations +- [ ] Error handling is implemented for all external calls +- [ ] Server name follows format: `{service}_mcp` +- [ ] All network operations use async/await +- [ ] Common functionality is extracted into reusable functions +- [ ] Error messages are clear, actionable, and educational +- [ ] Outputs are properly validated and formatted + +### Tool Configuration +- [ ] All tools implement 'name' and 'annotations' in the decorator +- [ ] Annotations correctly set (readOnlyHint, destructiveHint, idempotentHint, openWorldHint) +- [ ] All tools use Pydantic BaseModel for input validation with Field() definitions +- [ ] All Pydantic Fields have explicit types and descriptions with constraints +- [ ] All tools have comprehensive docstrings with explicit input/output types +- [ ] Docstrings include complete schema structure for dict/JSON returns +- [ ] Pydantic models handle input validation (no manual validation needed) + +### Advanced Features (where applicable) +- [ ] Context injection used for logging, progress, or elicitation +- [ ] Resources registered for appropriate data endpoints +- [ ] Lifespan management implemented for persistent connections +- [ ] Structured output types used (TypedDict, Pydantic models) +- [ ] Appropriate transport configured (stdio or streamable HTTP) + +### Code Quality +- [ ] File includes proper imports including Pydantic imports +- [ ] Pagination is properly implemented where applicable +- [ ] Filtering options are provided for potentially large result sets +- [ ] All async functions are properly defined with `async def` +- [ ] HTTP client usage follows async patterns with proper context managers +- [ ] Type hints are used throughout the code +- [ ] Constants are defined at module level in UPPER_CASE + +### Testing +- [ ] Server runs successfully: `python your_server.py --help` +- [ ] All imports resolve correctly +- [ ] Sample tool calls work as expected +- [ ] Error scenarios handled gracefully \ No newline at end of file diff --git a/contributing/samples/skill_script_demo/skills/mcp-builder/scripts/connections.py b/contributing/samples/skill_script_demo/skills/mcp-builder/scripts/connections.py new file mode 100644 index 0000000000..b828211677 --- /dev/null +++ b/contributing/samples/skill_script_demo/skills/mcp-builder/scripts/connections.py @@ -0,0 +1,160 @@ +"""Lightweight connection handling for MCP servers.""" + +from abc import ABC +from abc import abstractmethod +from contextlib import AsyncExitStack +from typing import Any + +from mcp import ClientSession +from mcp import StdioServerParameters +from mcp.client.sse import sse_client +from mcp.client.stdio import stdio_client +from mcp.client.streamable_http import streamablehttp_client + + +class MCPConnection(ABC): + """Base class for MCP server connections.""" + + def __init__(self): + self.session = None + self._stack = None + + @abstractmethod + def _create_context(self): + """Create the connection context based on connection type.""" + + async def __aenter__(self): + """Initialize MCP server connection.""" + self._stack = AsyncExitStack() + await self._stack.__aenter__() + + try: + ctx = self._create_context() + result = await self._stack.enter_async_context(ctx) + + if len(result) == 2: + read, write = result + elif len(result) == 3: + read, write, _ = result + else: + raise ValueError(f"Unexpected context result: {result}") + + session_ctx = ClientSession(read, write) + self.session = await self._stack.enter_async_context(session_ctx) + await self.session.initialize() + return self + except BaseException: + await self._stack.__aexit__(None, None, None) + raise + + async def __aexit__(self, exc_type, exc_val, exc_tb): + """Clean up MCP server connection resources.""" + if self._stack: + await self._stack.__aexit__(exc_type, exc_val, exc_tb) + self.session = None + self._stack = None + + async def list_tools(self) -> list[dict[str, Any]]: + """Retrieve available tools from the MCP server.""" + response = await self.session.list_tools() + return [ + { + "name": tool.name, + "description": tool.description, + "input_schema": tool.inputSchema, + } + for tool in response.tools + ] + + async def call_tool(self, tool_name: str, arguments: dict[str, Any]) -> Any: + """Call a tool on the MCP server with provided arguments.""" + result = await self.session.call_tool(tool_name, arguments=arguments) + return result.content + + +class MCPConnectionStdio(MCPConnection): + """MCP connection using standard input/output.""" + + def __init__( + self, command: str, args: list[str] = None, env: dict[str, str] = None + ): + super().__init__() + self.command = command + self.args = args or [] + self.env = env + + def _create_context(self): + return stdio_client( + StdioServerParameters( + command=self.command, args=self.args, env=self.env + ) + ) + + +class MCPConnectionSSE(MCPConnection): + """MCP connection using Server-Sent Events.""" + + def __init__(self, url: str, headers: dict[str, str] = None): + super().__init__() + self.url = url + self.headers = headers or {} + + def _create_context(self): + return sse_client(url=self.url, headers=self.headers) + + +class MCPConnectionHTTP(MCPConnection): + """MCP connection using Streamable HTTP.""" + + def __init__(self, url: str, headers: dict[str, str] = None): + super().__init__() + self.url = url + self.headers = headers or {} + + def _create_context(self): + return streamablehttp_client(url=self.url, headers=self.headers) + + +def create_connection( + transport: str, + command: str = None, + args: list[str] = None, + env: dict[str, str] = None, + url: str = None, + headers: dict[str, str] = None, +) -> MCPConnection: + """Factory function to create the appropriate MCP connection. + + Args: + transport: Connection type ("stdio", "sse", or "http") + command: Command to run (stdio only) + args: Command arguments (stdio only) + env: Environment variables (stdio only) + url: Server URL (sse and http only) + headers: HTTP headers (sse and http only) + + Returns: + MCPConnection instance + """ + transport = transport.lower() + + if transport == "stdio": + if not command: + raise ValueError("Command is required for stdio transport") + return MCPConnectionStdio(command=command, args=args, env=env) + + elif transport == "sse": + if not url: + raise ValueError("URL is required for sse transport") + return MCPConnectionSSE(url=url, headers=headers) + + elif transport in ["http", "streamable_http", "streamable-http"]: + if not url: + raise ValueError("URL is required for http transport") + return MCPConnectionHTTP(url=url, headers=headers) + + else: + raise ValueError( + f"Unsupported transport type: {transport}. Use 'stdio', 'sse', or" + " 'http'" + ) diff --git a/contributing/samples/skill_script_demo/skills/mcp-builder/scripts/evaluation.py b/contributing/samples/skill_script_demo/skills/mcp-builder/scripts/evaluation.py new file mode 100644 index 0000000000..fd5542f169 --- /dev/null +++ b/contributing/samples/skill_script_demo/skills/mcp-builder/scripts/evaluation.py @@ -0,0 +1,428 @@ +"""MCP Server Evaluation Harness + +This script evaluates MCP servers by running test questions against them using Claude. +""" + +import argparse +import asyncio +import json +from pathlib import Path +import re +import sys +import time +import traceback +from typing import Any +import xml.etree.ElementTree as ET + +from anthropic import Anthropic +from connections import create_connection + +EVALUATION_PROMPT = """You are an AI assistant with access to tools. + +When given a task, you MUST: +1. Use the available tools to complete the task +2. Provide summary of each step in your approach, wrapped in tags +3. Provide feedback on the tools provided, wrapped in tags +4. Provide your final response, wrapped in tags + +Summary Requirements: +- In your tags, you must explain: + - The steps you took to complete the task + - Which tools you used, in what order, and why + - The inputs you provided to each tool + - The outputs you received from each tool + - A summary for how you arrived at the response + +Feedback Requirements: +- In your tags, provide constructive feedback on the tools: + - Comment on tool names: Are they clear and descriptive? + - Comment on input parameters: Are they well-documented? Are required vs optional parameters clear? + - Comment on descriptions: Do they accurately describe what the tool does? + - Comment on any errors encountered during tool usage: Did the tool fail to execute? Did the tool return too many tokens? + - Identify specific areas for improvement and explain WHY they would help + - Be specific and actionable in your suggestions + +Response Requirements: +- Your response should be concise and directly address what was asked +- Always wrap your final response in tags +- If you cannot solve the task return NOT_FOUND +- For numeric responses, provide just the number +- For IDs, provide just the ID +- For names or text, provide the exact text requested +- Your response should go last""" + + +def parse_evaluation_file(file_path: Path) -> list[dict[str, Any]]: + """Parse XML evaluation file with qa_pair elements.""" + try: + tree = ET.parse(file_path) + root = tree.getroot() + evaluations = [] + + for qa_pair in root.findall(".//qa_pair"): + question_elem = qa_pair.find("question") + answer_elem = qa_pair.find("answer") + + if question_elem is not None and answer_elem is not None: + evaluations.append({ + "question": (question_elem.text or "").strip(), + "answer": (answer_elem.text or "").strip(), + }) + + return evaluations + except Exception as e: + print(f"Error parsing evaluation file {file_path}: {e}") + return [] + + +def extract_xml_content(text: str, tag: str) -> str | None: + """Extract content from XML tags.""" + pattern = rf"<{tag}>(.*?)" + matches = re.findall(pattern, text, re.DOTALL) + return matches[-1].strip() if matches else None + + +async def agent_loop( + client: Anthropic, + model: str, + question: str, + tools: list[dict[str, Any]], + connection: Any, +) -> tuple[str, dict[str, Any]]: + """Run the agent loop with MCP tools.""" + messages = [{"role": "user", "content": question}] + + response = await asyncio.to_thread( + client.messages.create, + model=model, + max_tokens=4096, + system=EVALUATION_PROMPT, + messages=messages, + tools=tools, + ) + + messages.append({"role": "assistant", "content": response.content}) + + tool_metrics = {} + + while response.stop_reason == "tool_use": + tool_use = next( + block for block in response.content if block.type == "tool_use" + ) + tool_name = tool_use.name + tool_input = tool_use.input + + tool_start_ts = time.time() + try: + tool_result = await connection.call_tool(tool_name, tool_input) + tool_response = ( + json.dumps(tool_result) + if isinstance(tool_result, (dict, list)) + else str(tool_result) + ) + except Exception as e: + tool_response = f"Error executing tool {tool_name}: {str(e)}\n" + tool_response += traceback.format_exc() + tool_duration = time.time() - tool_start_ts + + if tool_name not in tool_metrics: + tool_metrics[tool_name] = {"count": 0, "durations": []} + tool_metrics[tool_name]["count"] += 1 + tool_metrics[tool_name]["durations"].append(tool_duration) + + messages.append({ + "role": "user", + "content": [{ + "type": "tool_result", + "tool_use_id": tool_use.id, + "content": tool_response, + }], + }) + + response = await asyncio.to_thread( + client.messages.create, + model=model, + max_tokens=4096, + system=EVALUATION_PROMPT, + messages=messages, + tools=tools, + ) + messages.append({"role": "assistant", "content": response.content}) + + response_text = next( + (block.text for block in response.content if hasattr(block, "text")), + None, + ) + return response_text, tool_metrics + + +async def evaluate_single_task( + client: Anthropic, + model: str, + qa_pair: dict[str, Any], + tools: list[dict[str, Any]], + connection: Any, + task_index: int, +) -> dict[str, Any]: + """Evaluate a single QA pair with the given tools.""" + start_time = time.time() + + print( + f"Task {task_index + 1}: Running task with question:" + f" {qa_pair['question']}" + ) + response, tool_metrics = await agent_loop( + client, model, qa_pair["question"], tools, connection + ) + + response_value = extract_xml_content(response, "response") + summary = extract_xml_content(response, "summary") + feedback = extract_xml_content(response, "feedback") + + duration_seconds = time.time() - start_time + + return { + "question": qa_pair["question"], + "expected": qa_pair["answer"], + "actual": response_value, + "score": ( + int(response_value == qa_pair["answer"]) if response_value else 0 + ), + "total_duration": duration_seconds, + "tool_calls": tool_metrics, + "num_tool_calls": sum( + len(metrics["durations"]) for metrics in tool_metrics.values() + ), + "summary": summary, + "feedback": feedback, + } + + +REPORT_HEADER = """ +# Evaluation Report + +## Summary + +- **Accuracy**: {correct}/{total} ({accuracy:.1f}%) +- **Average Task Duration**: {average_duration_s:.2f}s +- **Average Tool Calls per Task**: {average_tool_calls:.2f} +- **Total Tool Calls**: {total_tool_calls} + +--- +""" + +TASK_TEMPLATE = """ +### Task {task_num} + +**Question**: {question} +**Ground Truth Answer**: `{expected_answer}` +**Actual Answer**: `{actual_answer}` +**Correct**: {correct_indicator} +**Duration**: {total_duration:.2f}s +**Tool Calls**: {tool_calls} + +**Summary** +{summary} + +**Feedback** +{feedback} + +--- +""" + + +async def run_evaluation( + eval_path: Path, + connection: Any, + model: str = "claude-3-7-sonnet-20250219", +) -> str: + """Run evaluation with MCP server tools.""" + print("🚀 Starting Evaluation") + + client = Anthropic() + + tools = await connection.list_tools() + print(f"📋 Loaded {len(tools)} tools from MCP server") + + qa_pairs = parse_evaluation_file(eval_path) + print(f"📋 Loaded {len(qa_pairs)} evaluation tasks") + + results = [] + for i, qa_pair in enumerate(qa_pairs): + print(f"Processing task {i + 1}/{len(qa_pairs)}") + result = await evaluate_single_task( + client, model, qa_pair, tools, connection, i + ) + results.append(result) + + correct = sum(r["score"] for r in results) + accuracy = (correct / len(results)) * 100 if results else 0 + average_duration_s = ( + sum(r["total_duration"] for r in results) / len(results) if results else 0 + ) + average_tool_calls = ( + sum(r["num_tool_calls"] for r in results) / len(results) if results else 0 + ) + total_tool_calls = sum(r["num_tool_calls"] for r in results) + + report = REPORT_HEADER.format( + correct=correct, + total=len(results), + accuracy=accuracy, + average_duration_s=average_duration_s, + average_tool_calls=average_tool_calls, + total_tool_calls=total_tool_calls, + ) + + report += "".join([ + TASK_TEMPLATE.format( + task_num=i + 1, + question=qa_pair["question"], + expected_answer=qa_pair["answer"], + actual_answer=result["actual"] or "N/A", + correct_indicator="✅" if result["score"] else "❌", + total_duration=result["total_duration"], + tool_calls=json.dumps(result["tool_calls"], indent=2), + summary=result["summary"] or "N/A", + feedback=result["feedback"] or "N/A", + ) + for i, (qa_pair, result) in enumerate(zip(qa_pairs, results)) + ]) + + return report + + +def parse_headers(header_list: list[str]) -> dict[str, str]: + """Parse header strings in format 'Key: Value' into a dictionary.""" + headers = {} + if not header_list: + return headers + + for header in header_list: + if ":" in header: + key, value = header.split(":", 1) + headers[key.strip()] = value.strip() + else: + print(f"Warning: Ignoring malformed header: {header}") + return headers + + +def parse_env_vars(env_list: list[str]) -> dict[str, str]: + """Parse environment variable strings in format 'KEY=VALUE' into a dictionary.""" + env = {} + if not env_list: + return env + + for env_var in env_list: + if "=" in env_var: + key, value = env_var.split("=", 1) + env[key.strip()] = value.strip() + else: + print(f"Warning: Ignoring malformed environment variable: {env_var}") + return env + + +async def main(): + parser = argparse.ArgumentParser( + description="Evaluate MCP servers using test questions", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + # Evaluate a local stdio MCP server + python evaluation.py -t stdio -c python -a my_server.py eval.xml + + # Evaluate an SSE MCP server + python evaluation.py -t sse -u https://example.com/mcp -H "Authorization: Bearer token" eval.xml + + # Evaluate an HTTP MCP server with custom model + python evaluation.py -t http -u https://example.com/mcp -m claude-3-5-sonnet-20241022 eval.xml + """, + ) + + parser.add_argument( + "eval_file", type=Path, help="Path to evaluation XML file" + ) + parser.add_argument( + "-t", + "--transport", + choices=["stdio", "sse", "http"], + default="stdio", + help="Transport type (default: stdio)", + ) + parser.add_argument( + "-m", + "--model", + default="claude-3-7-sonnet-20250219", + help="Claude model to use (default: claude-3-7-sonnet-20250219)", + ) + + stdio_group = parser.add_argument_group("stdio options") + stdio_group.add_argument( + "-c", "--command", help="Command to run MCP server (stdio only)" + ) + stdio_group.add_argument( + "-a", "--args", nargs="+", help="Arguments for the command (stdio only)" + ) + stdio_group.add_argument( + "-e", + "--env", + nargs="+", + help="Environment variables in KEY=VALUE format (stdio only)", + ) + + remote_group = parser.add_argument_group("sse/http options") + remote_group.add_argument( + "-u", "--url", help="MCP server URL (sse/http only)" + ) + remote_group.add_argument( + "-H", + "--header", + nargs="+", + dest="headers", + help="HTTP headers in 'Key: Value' format (sse/http only)", + ) + + parser.add_argument( + "-o", + "--output", + type=Path, + help="Output file for evaluation report (default: stdout)", + ) + + args = parser.parse_args() + + if not args.eval_file.exists(): + print(f"Error: Evaluation file not found: {args.eval_file}") + sys.exit(1) + + headers = parse_headers(args.headers) if args.headers else None + env_vars = parse_env_vars(args.env) if args.env else None + + try: + connection = create_connection( + transport=args.transport, + command=args.command, + args=args.args, + env=env_vars, + url=args.url, + headers=headers, + ) + except ValueError as e: + print(f"Error: {e}") + sys.exit(1) + + print(f"🔗 Connecting to MCP server via {args.transport}...") + + async with connection: + print("✅ Connected successfully") + report = await run_evaluation(args.eval_file, connection, args.model) + + if args.output: + args.output.write_text(report) + print(f"\n✅ Report saved to {args.output}") + else: + print("\n" + report) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/contributing/samples/skill_script_demo/skills/mcp-builder/scripts/example_evaluation.xml b/contributing/samples/skill_script_demo/skills/mcp-builder/scripts/example_evaluation.xml new file mode 100644 index 0000000000..41e4459b5a --- /dev/null +++ b/contributing/samples/skill_script_demo/skills/mcp-builder/scripts/example_evaluation.xml @@ -0,0 +1,22 @@ + + + Calculate the compound interest on $10,000 invested at 5% annual interest rate, compounded monthly for 3 years. What is the final amount in dollars (rounded to 2 decimal places)? + 11614.72 + + + A projectile is launched at a 45-degree angle with an initial velocity of 50 m/s. Calculate the total distance (in meters) it has traveled from the launch point after 2 seconds, assuming g=9.8 m/s². Round to 2 decimal places. + 87.25 + + + A sphere has a volume of 500 cubic meters. Calculate its surface area in square meters. Round to 2 decimal places. + 304.65 + + + Calculate the population standard deviation of this dataset: [12, 15, 18, 22, 25, 30, 35]. Round to 2 decimal places. + 7.61 + + + Calculate the pH of a solution with a hydrogen ion concentration of 3.5 × 10^-5 M. Round to 2 decimal places. + 4.46 + + diff --git a/contributing/samples/skill_script_demo/skills/mcp-builder/scripts/requirements.txt b/contributing/samples/skill_script_demo/skills/mcp-builder/scripts/requirements.txt new file mode 100644 index 0000000000..e73e5d1e35 --- /dev/null +++ b/contributing/samples/skill_script_demo/skills/mcp-builder/scripts/requirements.txt @@ -0,0 +1,2 @@ +anthropic>=0.39.0 +mcp>=1.1.0 diff --git a/contributing/samples/skill_script_demo/skills/python-helper/SKILL.md b/contributing/samples/skill_script_demo/skills/python-helper/SKILL.md new file mode 100644 index 0000000000..249e76b325 --- /dev/null +++ b/contributing/samples/skill_script_demo/skills/python-helper/SKILL.md @@ -0,0 +1,33 @@ +--- +name: python-helper +description: Python utility scripts for code analysis, data processing, and generation tasks. +version: 1.0.0 +--- + +# Python Helper Skill + +A collection of lightweight Python utility scripts for common development tasks. + +## Available Scripts + +### `fibonacci.py` +Generates a Fibonacci sequence. Pass the desired count as an argument. + +**Usage**: `execute_skill_script(skill_name="python-helper", script_name="fibonacci.py", input_args="10")` + +### `word_count.py` +Analyzes text and reports word frequency statistics. Pass the text to analyze as an argument. + +**Usage**: `execute_skill_script(skill_name="python-helper", script_name="word_count.py", input_args="the quick brown fox jumps over the lazy dog the fox")` + +### `json_format.py` +Pretty-prints and validates a JSON string. Pass the JSON as a single quoted argument. + +**Usage**: `execute_skill_script(skill_name="python-helper", script_name="json_format.py", input_args='{"name":"Alice","scores":[90,85,92]}')` + +## Workflow + +1. Use `load_skill` to read these instructions. +2. Use `load_skill_resource` to inspect a script's source if needed. +3. Use `execute_skill_script` with appropriate `input_args` to run a script. +4. Interpret the script's stdout and present results to the user. diff --git a/contributing/samples/skill_script_demo/skills/python-helper/references/usage.md b/contributing/samples/skill_script_demo/skills/python-helper/references/usage.md new file mode 100644 index 0000000000..a5d425cccd --- /dev/null +++ b/contributing/samples/skill_script_demo/skills/python-helper/references/usage.md @@ -0,0 +1,17 @@ +# Python Helper Usage Guide + +## Quick Reference + +| Script | Purpose | Example Args | +|--------|---------|-------------| +| `fibonacci.py` | Generate Fibonacci sequence | `"15"` (count) | +| `word_count.py` | Word frequency analysis | `"hello world hello"` | +| `json_format.py` | Validate & pretty-print JSON | `'{"key":"value"}'` | + +## Tips + +- All scripts write results to **stdout**. +- Pass arguments via the `input_args` parameter as a space-separated string. +- `fibonacci.py` defaults to 10 numbers if no argument is given. +- `word_count.py` treats all arguments as the text to analyze. +- `json_format.py` joins all arguments as a single JSON string. diff --git a/contributing/samples/skill_script_demo/skills/python-helper/scripts/fibonacci.py b/contributing/samples/skill_script_demo/skills/python-helper/scripts/fibonacci.py new file mode 100644 index 0000000000..fb0ffd3679 --- /dev/null +++ b/contributing/samples/skill_script_demo/skills/python-helper/scripts/fibonacci.py @@ -0,0 +1,12 @@ +"""Generate a Fibonacci sequence of N numbers.""" + +import sys + +n = int(sys.argv[1]) if len(sys.argv) > 1 else 10 +a, b = 0, 1 +result = [] +for _ in range(n): + result.append(a) + a, b = b, a + b +print(f"Fibonacci({n}): {result}") +print(f"Sum: {sum(result)}") diff --git a/contributing/samples/skill_script_demo/skills/python-helper/scripts/json_format.py b/contributing/samples/skill_script_demo/skills/python-helper/scripts/json_format.py new file mode 100644 index 0000000000..819fa4da67 --- /dev/null +++ b/contributing/samples/skill_script_demo/skills/python-helper/scripts/json_format.py @@ -0,0 +1,21 @@ +"""Pretty-print and validate a JSON string.""" + +import json +import sys + +raw = " ".join(sys.argv[1:]) if len(sys.argv) > 1 else "" +if not raw: + print("Error: no JSON input provided") + sys.exit(1) + +try: + data = json.loads(raw) + print(json.dumps(data, indent=2, sort_keys=True)) + print(f"\nType: {type(data).__name__}") + if isinstance(data, dict): + print(f"Keys: {list(data.keys())}") + elif isinstance(data, list): + print(f"Length: {len(data)}") +except json.JSONDecodeError as e: + print(f"Invalid JSON: {e}") + sys.exit(1) diff --git a/contributing/samples/skill_script_demo/skills/python-helper/scripts/word_count.py b/contributing/samples/skill_script_demo/skills/python-helper/scripts/word_count.py new file mode 100644 index 0000000000..16b1db0c26 --- /dev/null +++ b/contributing/samples/skill_script_demo/skills/python-helper/scripts/word_count.py @@ -0,0 +1,18 @@ +"""Analyze word frequency in the given text.""" + +from collections import Counter +import sys + +text = " ".join(sys.argv[1:]) if len(sys.argv) > 1 else "" +if not text: + print("Error: no text provided") + sys.exit(1) + +words = text.lower().split() +freq = Counter(words) + +print(f"Total words: {len(words)}") +print(f"Unique words: {len(freq)}") +print("Top words:") +for word, count in freq.most_common(5): + print(f" {word}: {count}") diff --git a/contributing/samples/skill_script_demo/test_skill_compat.py b/contributing/samples/skill_script_demo/test_skill_compat.py new file mode 100644 index 0000000000..a0a04abfeb --- /dev/null +++ b/contributing/samples/skill_script_demo/test_skill_compat.py @@ -0,0 +1,258 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compatibility test: Anthropic mcp-builder skill vs ADK SkillToolset. + +Tests every tool in SkillToolset against the real Anthropic mcp-builder +skill to verify end-to-end compatibility with the public Agent Skills +spec (agentskills.io/specification). + +Run: + pytest contributing/samples/skill_script_demo/test_skill_compat.py -v +""" + +import pathlib +from unittest import mock + +from google.adk.code_executors.base_code_executor import BaseCodeExecutor +from google.adk.code_executors.code_execution_utils import CodeExecutionResult +from google.adk.skills import load_skill_from_dir +from google.adk.tools import skill_toolset +from google.adk.tools import tool_context +import pytest + +_SKILLS_DIR = pathlib.Path(__file__).parent / "skills" + + +@pytest.fixture +def mcp_builder_skill(): + return load_skill_from_dir(_SKILLS_DIR / "mcp-builder") + + +@pytest.fixture +def toolset_with_executor(mcp_builder_skill): + executor = mock.create_autospec(BaseCodeExecutor, instance=True) + executor.execute_code.return_value = CodeExecutionResult( + stdout="ok\n", stderr="" + ) + return skill_toolset.SkillToolset( + skills=[mcp_builder_skill], code_executor=executor + ) + + +@pytest.fixture +def mock_ctx(): + ctx = mock.MagicMock(spec=tool_context.ToolContext) + ctx._invocation_context = mock.MagicMock() + ctx._invocation_context.agent = mock.MagicMock() + return ctx + + +# ── 1. Skill loading ────────────────────────────────────────────── + + +def test_skill_loads_from_disk(mcp_builder_skill): + """Verify the skill loads successfully from disk.""" + assert mcp_builder_skill.name == "mcp-builder" + assert "MCP" in mcp_builder_skill.description + assert len(mcp_builder_skill.instructions) > 0 + + +def test_skill_scripts_loaded(mcp_builder_skill): + """All scripts/ files should be loaded.""" + scripts = mcp_builder_skill.resources.list_scripts() + assert "connections.py" in scripts + assert "evaluation.py" in scripts + assert "requirements.txt" in scripts + assert "example_evaluation.xml" in scripts + + +def test_skill_references_loaded(mcp_builder_skill): + """All references/ files should be loaded per spec.""" + refs = mcp_builder_skill.resources.list_references() + assert "evaluation.md" in refs + assert "mcp_best_practices.md" in refs + assert "node_mcp_server.md" in refs + assert "python_mcp_server.md" in refs + + +# ── 2. list_skills tool ─────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_list_skills_shows_mcp_builder(toolset_with_executor, mock_ctx): + tool = skill_toolset.ListSkillsTool(toolset_with_executor) + result = await tool.run_async(args={}, tool_context=mock_ctx) + assert "mcp-builder" in result + + +# ── 3. load_skill tool ──────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_load_skill_returns_instructions(toolset_with_executor, mock_ctx): + tool = skill_toolset.LoadSkillTool(toolset_with_executor) + result = await tool.run_async( + args={"name": "mcp-builder"}, tool_context=mock_ctx + ) + assert result["skill_name"] == "mcp-builder" + assert "MCP Server Development Guide" in result["instructions"] + assert result["frontmatter"]["name"] == "mcp-builder" + + +# ── 4. load_skill_resource tool ─────────────────────────────────── + + +@pytest.mark.asyncio +async def test_load_script_content(toolset_with_executor, mock_ctx): + """Can read a Python script via load_skill_resource.""" + tool = skill_toolset.LoadSkillResourceTool(toolset_with_executor) + result = await tool.run_async( + args={ + "skill_name": "mcp-builder", + "path": "scripts/connections.py", + }, + tool_context=mock_ctx, + ) + assert result["content"] is not None + assert "MCPConnection" in result["content"] + + +@pytest.mark.asyncio +async def test_load_requirements_txt(toolset_with_executor, mock_ctx): + """Can read non-executable script files like requirements.txt.""" + tool = skill_toolset.LoadSkillResourceTool(toolset_with_executor) + result = await tool.run_async( + args={ + "skill_name": "mcp-builder", + "path": "scripts/requirements.txt", + }, + tool_context=mock_ctx, + ) + assert "anthropic" in result["content"] + + +@pytest.mark.asyncio +async def test_load_reference_content(toolset_with_executor, mock_ctx): + """Can read reference files via load_skill_resource.""" + tool = skill_toolset.LoadSkillResourceTool(toolset_with_executor) + result = await tool.run_async( + args={ + "skill_name": "mcp-builder", + "path": "references/evaluation.md", + }, + tool_context=mock_ctx, + ) + assert result["content"] is not None + assert "evaluation" in result["content"].lower() + + +# ── 5. execute_skill_script tool ────────────────────────────────── + + +@pytest.mark.asyncio +async def test_execute_python_script(toolset_with_executor, mock_ctx): + """Can execute a .py script from the skill.""" + tool = skill_toolset.ExecuteSkillScriptTool(toolset_with_executor) + result = await tool.run_async( + args={ + "skill_name": "mcp-builder", + "script_name": "connections.py", + }, + tool_context=mock_ctx, + ) + assert result["status"] == "success" + assert result["script_name"] == "connections.py" + + # Verify the executor received the actual script source + call_args = toolset_with_executor._code_executor.execute_code.call_args + code_input = call_args[0][1] + assert "MCPConnection" in code_input.code + + +@pytest.mark.asyncio +async def test_execute_with_scripts_prefix(toolset_with_executor, mock_ctx): + """scripts/ prefix is stripped automatically.""" + tool = skill_toolset.ExecuteSkillScriptTool(toolset_with_executor) + result = await tool.run_async( + args={ + "skill_name": "mcp-builder", + "script_name": "scripts/evaluation.py", + }, + tool_context=mock_ctx, + ) + assert result["status"] == "success" + assert result["script_name"] == "evaluation.py" + + +@pytest.mark.asyncio +async def test_execute_unsupported_txt(toolset_with_executor, mock_ctx): + """requirements.txt is not executable — returns UNSUPPORTED_SCRIPT_TYPE.""" + tool = skill_toolset.ExecuteSkillScriptTool(toolset_with_executor) + result = await tool.run_async( + args={ + "skill_name": "mcp-builder", + "script_name": "requirements.txt", + }, + tool_context=mock_ctx, + ) + assert result["error_code"] == "UNSUPPORTED_SCRIPT_TYPE" + + +@pytest.mark.asyncio +async def test_execute_unsupported_xml(toolset_with_executor, mock_ctx): + """example_evaluation.xml is not executable — returns UNSUPPORTED_SCRIPT_TYPE.""" + tool = skill_toolset.ExecuteSkillScriptTool(toolset_with_executor) + result = await tool.run_async( + args={ + "skill_name": "mcp-builder", + "script_name": "example_evaluation.xml", + }, + tool_context=mock_ctx, + ) + assert result["error_code"] == "UNSUPPORTED_SCRIPT_TYPE" + + +@pytest.mark.asyncio +async def test_execute_with_input_args(toolset_with_executor, mock_ctx): + """input_args are injected into sys.argv for Python scripts.""" + tool = skill_toolset.ExecuteSkillScriptTool(toolset_with_executor) + result = await tool.run_async( + args={ + "skill_name": "mcp-builder", + "script_name": "evaluation.py", + "input_args": "-t stdio -c python eval.xml", + }, + tool_context=mock_ctx, + ) + assert result["status"] == "success" + + call_args = toolset_with_executor._code_executor.execute_code.call_args + code_input = call_args[0][1] + assert "sys.argv" in code_input.code + assert "-t stdio -c python eval.xml" in code_input.code + + +# ── 6. Tool count ───────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_toolset_provides_4_tools(toolset_with_executor): + tools = await toolset_with_executor.get_tools() + assert len(tools) == 4 + names = [t.name for t in tools] + assert "list_skills" in names + assert "load_skill" in names + assert "load_skill_resource" in names + assert "execute_skill_script" in names diff --git a/src/google/adk/tools/skill_toolset.py b/src/google/adk/tools/skill_toolset.py index f90dfdb2b1..8c2e7c9f9f 100644 --- a/src/google/adk/tools/skill_toolset.py +++ b/src/google/adk/tools/skill_toolset.py @@ -16,12 +16,16 @@ from __future__ import annotations +import logging from typing import Any +from typing import Optional from typing import TYPE_CHECKING from google.genai import types from ..agents.readonly_context import ReadonlyContext +from ..code_executors.base_code_executor import BaseCodeExecutor +from ..code_executors.code_execution_utils import CodeExecutionInput from ..features import experimental from ..features import FeatureName from ..skills import models @@ -33,6 +37,8 @@ if TYPE_CHECKING: from ..models.llm_request import LlmRequest +logger = logging.getLogger("google_adk." + __name__) + DEFAULT_SKILL_SYSTEM_INSTRUCTION = """You can use specialized 'skills' to help you with complex tasks. You MUST use the skill tools to interact with these skills. Skills are folders of instructions and resources that extend your capabilities for specialized tasks. Each skill folder contains: @@ -46,6 +52,7 @@ 1. If a skill seems relevant to the current user query, you MUST use the `load_skill` tool with `name=""` to read its full instructions before proceeding. 2. Once you have read the instructions, follow them exactly as documented before replying to the user. For example, If the instruction lists multiple steps, please make sure you complete all of them in order. 3. The `load_skill_resource` tool is for viewing files within a skill's directory (e.g., `references/*`, `assets/*`, `scripts/*`). Do NOT use other tools to access these files. +4. Use `execute_skill_script` to run scripts from a skill's `scripts/` directory. Use `load_skill_resource` to view script content first if needed. """ @@ -227,6 +234,187 @@ async def run_async( } +@experimental(FeatureName.SKILL_TOOLSET) +class ExecuteSkillScriptTool(BaseTool): + """Tool to execute scripts from a skill's scripts/ directory.""" + + def __init__(self, toolset: "SkillToolset"): + super().__init__( + name="execute_skill_script", + description=( + "Executes a script from a skill's scripts/ directory" + " and returns its output." + ), + ) + self._toolset = toolset + + def _get_declaration(self) -> types.FunctionDeclaration | None: + return types.FunctionDeclaration( + name=self.name, + description=self.description, + parameters_json_schema={ + "type": "object", + "properties": { + "skill_name": { + "type": "string", + "description": "The name of the skill.", + }, + "script_name": { + "type": "string", + "description": ( + "The name of the script to execute (e.g.," + " 'setup.sh' or 'scripts/setup.sh')." + ), + }, + "input_args": { + "type": "string", + "description": ( + "Optional space-separated arguments to pass" + " to the script." + ), + }, + }, + "required": ["skill_name", "script_name"], + }, + ) + + async def run_async( + self, *, args: dict[str, Any], tool_context: ToolContext + ) -> Any: + skill_name = args.get("skill_name") + script_name = args.get("script_name") + input_args = args.get("input_args", "") + + if not skill_name: + return { + "error": "Skill name is required.", + "error_code": "MISSING_SKILL_NAME", + } + if not script_name: + return { + "error": "Script name is required.", + "error_code": "MISSING_SCRIPT_NAME", + } + + # Strip scripts/ prefix for consistency + if script_name.startswith("scripts/"): + script_name = script_name[len("scripts/") :] + + skill = self._toolset._get_skill(skill_name) + if not skill: + return { + "error": f"Skill '{skill_name}' not found.", + "error_code": "SKILL_NOT_FOUND", + } + + script = skill.resources.get_script(script_name) + if script is None: + return { + "error": f"Script '{script_name}' not found in skill '{skill_name}'.", + "error_code": "SCRIPT_NOT_FOUND", + } + + # Resolve code executor: toolset-level first, then agent fallback + code_executor = self._toolset._code_executor + if code_executor is None: + agent = tool_context._invocation_context.agent + if hasattr(agent, "code_executor"): + code_executor = agent.code_executor + if code_executor is None: + return { + "error": ( + "No code executor configured. A code executor is" + " required to run scripts." + ), + "error_code": "NO_CODE_EXECUTOR", + } + + # Prepare code based on script extension + code = self._prepare_code(script_name, script.src, input_args) + if code is None: + ext = script_name.rsplit(".", 1)[-1] if "." in script_name else "" + return { + "error": ( + f"Unsupported script type '.{ext}'. Supported" + " types: .py, .sh, .bash" + ), + "error_code": "UNSUPPORTED_SCRIPT_TYPE", + } + + try: + result = code_executor.execute_code( + tool_context._invocation_context, + CodeExecutionInput(code=code), + ) + return { + "skill_name": skill_name, + "script_name": script_name, + "stdout": result.stdout, + "stderr": result.stderr, + "status": "error" if result.stderr else "success", + } + except Exception as e: + logger.exception( + "Error executing script '%s' from skill '%s'", + script_name, + skill_name, + ) + # Keep the error message short for the LLM; full trace is logged above. + short_msg = str(e) + if len(short_msg) > 200: + short_msg = short_msg[:200] + "..." + return { + "error": f"Failed to execute script '{script_name}': {short_msg}", + "error_code": "EXECUTION_ERROR", + } + + def _prepare_code( + self, + script_name: str, + script_src: str, + input_args: str, + ) -> str | None: + """Prepares Python code to execute the script. + + Args: + script_name: The script filename. + script_src: The script source content. + input_args: Optional arguments string. + + Returns: + Python code string to execute, or None if unsupported type. + """ + ext = "" + if "." in script_name: + ext = script_name.rsplit(".", 1)[-1].lower() + + if ext in ("py", ""): + # Python script: execute directly, inject sys.argv if args + if input_args: + return ( + "import sys\n" + f"sys.argv = [{script_name!r}] + {input_args!r}.split()\n" + + script_src + ) + return script_src + elif ext in ("sh", "bash"): + # Shell script: wrap in subprocess.run + return ( + "import subprocess\n" + "_result = subprocess.run(\n" + f" ['bash', '-c', {script_src!r}" + + (f" + ' ' + {input_args!r}" if input_args else "") + + f"],\n" + f" capture_output=True, text=True,\n" + f")\n" + f"print(_result.stdout, end='')\n" + f"if _result.stderr:\n" + f" import sys\n" + f" print(_result.stderr, end='', file=sys.stderr)\n" + ) + return None + + @experimental(FeatureName.SKILL_TOOLSET) class SkillToolset(BaseToolset): """A toolset for managing and interacting with agent skills.""" @@ -234,6 +422,8 @@ class SkillToolset(BaseToolset): def __init__( self, skills: list[models.Skill], + *, + code_executor: Optional[BaseCodeExecutor] = None, ): super().__init__() @@ -245,10 +435,12 @@ def __init__( seen.add(skill.name) self._skills = {skill.name: skill for skill in skills} + self._code_executor = code_executor self._tools = [ ListSkillsTool(self), LoadSkillTool(self), LoadSkillResourceTool(self), + ExecuteSkillScriptTool(self), ] async def get_tools( diff --git a/tests/unittests/tools/test_skill_toolset.py b/tests/unittests/tools/test_skill_toolset.py index 066eedfb67..e9c24e6d06 100644 --- a/tests/unittests/tools/test_skill_toolset.py +++ b/tests/unittests/tools/test_skill_toolset.py @@ -14,6 +14,8 @@ from unittest import mock +from google.adk.code_executors.base_code_executor import BaseCodeExecutor +from google.adk.code_executors.code_execution_utils import CodeExecutionResult from google.adk.models import llm_request as llm_request_model from google.adk.skills import models from google.adk.tools import skill_toolset @@ -59,6 +61,10 @@ def get_asset(name): def get_script(name): if name == "setup.sh": return models.Script(src="echo setup") + if name == "run.py": + return models.Script(src="print('hello')") + if name == "build.rb": + return models.Script(src="puts 'hello'") return None skill.resources.get_reference.side_effect = get_ref @@ -132,10 +138,11 @@ def test_list_skills(mock_skill1, mock_skill2): async def test_get_tools(mock_skill1, mock_skill2): toolset = skill_toolset.SkillToolset([mock_skill1, mock_skill2]) tools = await toolset.get_tools() - assert len(tools) == 3 + assert len(tools) == 4 assert isinstance(tools[0], skill_toolset.ListSkillsTool) assert isinstance(tools[1], skill_toolset.LoadSkillTool) assert isinstance(tools[2], skill_toolset.LoadSkillResourceTool) + assert isinstance(tools[3], skill_toolset.ExecuteSkillScriptTool) @pytest.mark.asyncio @@ -308,3 +315,312 @@ async def test_scripts_resource_not_found(mock_skill1, tool_context_instance): tool_context=tool_context_instance, ) assert result["error_code"] == "RESOURCE_NOT_FOUND" + + +# ExecuteSkillScriptTool tests + + +def _make_tool_context_with_agent(agent=None): + """Creates a mock ToolContext with _invocation_context.agent.""" + ctx = mock.MagicMock(spec=tool_context.ToolContext) + ctx._invocation_context = mock.MagicMock() + ctx._invocation_context.agent = agent or mock.MagicMock() + return ctx + + +def _make_mock_executor(stdout="", stderr=""): + """Creates a mock code executor that returns the given output.""" + executor = mock.create_autospec(BaseCodeExecutor, instance=True) + executor.execute_code.return_value = CodeExecutionResult( + stdout=stdout, stderr=stderr + ) + return executor + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "args, expected_error_code", + [ + ( + {"script_name": "setup.sh"}, + "MISSING_SKILL_NAME", + ), + ( + {"skill_name": "skill1"}, + "MISSING_SCRIPT_NAME", + ), + ( + {"skill_name": "", "script_name": "setup.sh"}, + "MISSING_SKILL_NAME", + ), + ( + {"skill_name": "skill1", "script_name": ""}, + "MISSING_SCRIPT_NAME", + ), + ], +) +async def test_execute_script_missing_params( + mock_skill1, args, expected_error_code +): + executor = _make_mock_executor() + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async(args=args, tool_context=ctx) + assert result["error_code"] == expected_error_code + + +@pytest.mark.asyncio +async def test_execute_script_skill_not_found(mock_skill1): + executor = _make_mock_executor() + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "nonexistent", "script_name": "setup.sh"}, + tool_context=ctx, + ) + assert result["error_code"] == "SKILL_NOT_FOUND" + + +@pytest.mark.asyncio +async def test_execute_script_script_not_found(mock_skill1): + executor = _make_mock_executor() + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "nonexistent.py"}, + tool_context=ctx, + ) + assert result["error_code"] == "SCRIPT_NOT_FOUND" + + +@pytest.mark.asyncio +async def test_execute_script_no_code_executor(mock_skill1): + toolset = skill_toolset.SkillToolset([mock_skill1]) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + # Agent without code_executor attribute + agent = mock.MagicMock(spec=[]) + ctx = _make_tool_context_with_agent(agent=agent) + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "setup.sh"}, + tool_context=ctx, + ) + assert result["error_code"] == "NO_CODE_EXECUTOR" + + +@pytest.mark.asyncio +async def test_execute_script_agent_code_executor_none(mock_skill1): + """Agent has code_executor attr but it's None.""" + toolset = skill_toolset.SkillToolset([mock_skill1]) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + agent = mock.MagicMock() + agent.code_executor = None + ctx = _make_tool_context_with_agent(agent=agent) + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "setup.sh"}, + tool_context=ctx, + ) + assert result["error_code"] == "NO_CODE_EXECUTOR" + + +@pytest.mark.asyncio +async def test_execute_script_unsupported_type(mock_skill1): + executor = _make_mock_executor() + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "build.rb"}, + tool_context=ctx, + ) + assert result["error_code"] == "UNSUPPORTED_SCRIPT_TYPE" + + +@pytest.mark.asyncio +async def test_execute_script_python_success(mock_skill1): + executor = _make_mock_executor(stdout="hello\n") + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "run.py"}, + tool_context=ctx, + ) + assert result["status"] == "success" + assert result["stdout"] == "hello\n" + assert result["stderr"] == "" + assert result["skill_name"] == "skill1" + assert result["script_name"] == "run.py" + + # Verify the code passed to executor is the raw script + call_args = executor.execute_code.call_args + code_input = call_args[0][1] + assert code_input.code == "print('hello')" + + +@pytest.mark.asyncio +async def test_execute_script_shell_success(mock_skill1): + executor = _make_mock_executor(stdout="setup\n") + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "setup.sh"}, + tool_context=ctx, + ) + assert result["status"] == "success" + assert result["stdout"] == "setup\n" + + # Verify the code wraps in subprocess.run + call_args = executor.execute_code.call_args + code_input = call_args[0][1] + assert "subprocess.run" in code_input.code + assert "bash" in code_input.code + + +@pytest.mark.asyncio +async def test_execute_script_with_input_args_python(mock_skill1): + executor = _make_mock_executor(stdout="done\n") + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={ + "skill_name": "skill1", + "script_name": "run.py", + "input_args": "--verbose --count 3", + }, + tool_context=ctx, + ) + assert result["status"] == "success" + + call_args = executor.execute_code.call_args + code_input = call_args[0][1] + assert "sys.argv" in code_input.code + assert "--verbose --count 3" in code_input.code + + +@pytest.mark.asyncio +async def test_execute_script_with_input_args_shell(mock_skill1): + executor = _make_mock_executor(stdout="done\n") + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={ + "skill_name": "skill1", + "script_name": "setup.sh", + "input_args": "--force", + }, + tool_context=ctx, + ) + assert result["status"] == "success" + + call_args = executor.execute_code.call_args + code_input = call_args[0][1] + assert "--force" in code_input.code + + +@pytest.mark.asyncio +async def test_execute_script_scripts_prefix_stripping(mock_skill1): + executor = _make_mock_executor(stdout="setup\n") + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={ + "skill_name": "skill1", + "script_name": "scripts/setup.sh", + }, + tool_context=ctx, + ) + assert result["status"] == "success" + assert result["script_name"] == "setup.sh" + + +@pytest.mark.asyncio +async def test_execute_script_toolset_executor_priority(mock_skill1): + """Toolset-level executor takes priority over agent's.""" + toolset_executor = _make_mock_executor(stdout="from toolset\n") + agent_executor = _make_mock_executor(stdout="from agent\n") + toolset = skill_toolset.SkillToolset( + [mock_skill1], code_executor=toolset_executor + ) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + agent = mock.MagicMock() + agent.code_executor = agent_executor + ctx = _make_tool_context_with_agent(agent=agent) + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "run.py"}, + tool_context=ctx, + ) + assert result["stdout"] == "from toolset\n" + toolset_executor.execute_code.assert_called_once() + agent_executor.execute_code.assert_not_called() + + +@pytest.mark.asyncio +async def test_execute_script_agent_executor_fallback(mock_skill1): + """Falls back to agent's code executor when toolset has none.""" + agent_executor = _make_mock_executor(stdout="from agent\n") + toolset = skill_toolset.SkillToolset([mock_skill1]) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + agent = mock.MagicMock() + agent.code_executor = agent_executor + ctx = _make_tool_context_with_agent(agent=agent) + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "run.py"}, + tool_context=ctx, + ) + assert result["stdout"] == "from agent\n" + agent_executor.execute_code.assert_called_once() + + +@pytest.mark.asyncio +async def test_execute_script_execution_error(mock_skill1): + executor = _make_mock_executor() + executor.execute_code.side_effect = RuntimeError("boom") + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "run.py"}, + tool_context=ctx, + ) + assert result["error_code"] == "EXECUTION_ERROR" + assert "boom" in result["error"] + assert result["error"].startswith("Failed to execute script 'run.py':") + + +@pytest.mark.asyncio +async def test_execute_script_stderr_sets_error_status(mock_skill1): + executor = _make_mock_executor(stdout="", stderr="warning\n") + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "run.py"}, + tool_context=ctx, + ) + assert result["status"] == "error" + assert result["stderr"] == "warning\n" + + +@pytest.mark.asyncio +async def test_execute_script_execution_error_truncated(mock_skill1): + """Long exception messages are truncated to avoid wasting LLM tokens.""" + executor = _make_mock_executor() + executor.execute_code.side_effect = RuntimeError("x" * 300) + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "run.py"}, + tool_context=ctx, + ) + assert result["error_code"] == "EXECUTION_ERROR" + # 200 chars of the message + "..." suffix + the prefix + assert result["error"].endswith("...") + assert len(result["error"]) < 300 From 06d995e68fa4aed90c5f0e28756fe580d00efee0 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 21 Feb 2026 00:28:53 -0800 Subject: [PATCH 02/53] =?UTF-8?q?fix:=20Address=20Gemini=20Code=20Assist?= =?UTF-8?q?=20review=20=E2=80=94=20shell=20injection,=20shlex,=20check=3DT?= =?UTF-8?q?rue?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. Shell injection (HIGH): Use shlex.split() and pass args as separate list elements to subprocess.run instead of concatenating into bash -c command string. Args now arrive as positional params ($1, $2, ...). 2. Fragile arg parsing (MEDIUM): Replace .split() with shlex.split() for Python script sys.argv injection — correctly handles quoted arguments like --name "John Doe". 3. Status detection (MEDIUM): Add check=True to subprocess.run in shell wrapper so non-zero exit codes raise CalledProcessError, which the code executor captures as stderr. Co-Authored-By: Claude Opus 4.6 --- .../skill_script_demo/test_skill_compat.py | 1 + src/google/adk/tools/skill_toolset.py | 29 ++++++++++--------- tests/unittests/tools/test_skill_toolset.py | 5 +++- 3 files changed, 21 insertions(+), 14 deletions(-) diff --git a/contributing/samples/skill_script_demo/test_skill_compat.py b/contributing/samples/skill_script_demo/test_skill_compat.py index a0a04abfeb..221658c47e 100644 --- a/contributing/samples/skill_script_demo/test_skill_compat.py +++ b/contributing/samples/skill_script_demo/test_skill_compat.py @@ -241,6 +241,7 @@ async def test_execute_with_input_args(toolset_with_executor, mock_ctx): call_args = toolset_with_executor._code_executor.execute_code.call_args code_input = call_args[0][1] assert "sys.argv" in code_input.code + assert "shlex.split" in code_input.code assert "-t stdio -c python eval.xml" in code_input.code diff --git a/src/google/adk/tools/skill_toolset.py b/src/google/adk/tools/skill_toolset.py index 8c2e7c9f9f..8d5a9adef3 100644 --- a/src/google/adk/tools/skill_toolset.py +++ b/src/google/adk/tools/skill_toolset.py @@ -392,25 +392,28 @@ def _prepare_code( # Python script: execute directly, inject sys.argv if args if input_args: return ( - "import sys\n" - f"sys.argv = [{script_name!r}] + {input_args!r}.split()\n" + "import sys, shlex\n" + f"sys.argv = [{script_name!r}]" + f" + shlex.split({input_args!r})\n" + script_src ) return script_src elif ext in ("sh", "bash"): - # Shell script: wrap in subprocess.run + # Shell script: wrap in subprocess.run. + # Args are passed as separate list elements after the script + # name to avoid shell injection — bash -c receives the script + # source, and $0/$1/... get the positional parameters. return ( - "import subprocess\n" + "import subprocess, shlex\n" "_result = subprocess.run(\n" - f" ['bash', '-c', {script_src!r}" - + (f" + ' ' + {input_args!r}" if input_args else "") - + f"],\n" - f" capture_output=True, text=True,\n" - f")\n" - f"print(_result.stdout, end='')\n" - f"if _result.stderr:\n" - f" import sys\n" - f" print(_result.stderr, end='', file=sys.stderr)\n" + f" ['bash', '-c', {script_src!r}," + f" {script_name!r}]" + + (f" + shlex.split({input_args!r})" if input_args else "") + + ",\n" + " capture_output=True, text=True,\n" + " check=True,\n" + ")\n" + "print(_result.stdout, end='')\n" ) return None diff --git a/tests/unittests/tools/test_skill_toolset.py b/tests/unittests/tools/test_skill_toolset.py index e9c24e6d06..78cf3a1273 100644 --- a/tests/unittests/tools/test_skill_toolset.py +++ b/tests/unittests/tools/test_skill_toolset.py @@ -473,11 +473,12 @@ async def test_execute_script_shell_success(mock_skill1): assert result["status"] == "success" assert result["stdout"] == "setup\n" - # Verify the code wraps in subprocess.run + # Verify the code wraps in subprocess.run with check=True call_args = executor.execute_code.call_args code_input = call_args[0][1] assert "subprocess.run" in code_input.code assert "bash" in code_input.code + assert "check=True" in code_input.code @pytest.mark.asyncio @@ -499,6 +500,7 @@ async def test_execute_script_with_input_args_python(mock_skill1): call_args = executor.execute_code.call_args code_input = call_args[0][1] assert "sys.argv" in code_input.code + assert "shlex.split" in code_input.code assert "--verbose --count 3" in code_input.code @@ -520,6 +522,7 @@ async def test_execute_script_with_input_args_shell(mock_skill1): call_args = executor.execute_code.call_args code_input = call_args[0][1] + assert "shlex.split" in code_input.code assert "--force" in code_input.code From e83de80eb890f969c033de12b9b1ad1adb69f626 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 21 Feb 2026 10:24:39 -0800 Subject: [PATCH 03/53] fix: Address code review findings for ExecuteSkillScriptTool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix sys.exit(0) treated as error: zero/None exit codes now return success - Fix shell stderr discarding stdout: serialize both streams as JSON envelope through stdout to work around UnsafeLocalCodeExecutor losing stdout on exception - Add early shlex.split() validation for input_args before executor dispatch - Add script_timeout parameter to SkillToolset for subprocess.run timeout - Reject extensionless scripts (previously treated as Python) - Replace os.popen() with subprocess.run() in run_live_test.py - Add UnsafeLocalCodeExecutor safety warning to sample agent - Fix broken reference/ → references/ links in mcp-builder SKILL.md - Remove duplicate @pytest.mark.asyncio decorator - Add 4 integration tests using real UnsafeLocalCodeExecutor Co-Authored-By: Claude Opus 4.6 --- .../samples/skill_script_demo/agent.py | 6 + .../skill_script_demo/run_live_test.py | 14 +- .../skills/mcp-builder/SKILL.md | 20 +- src/google/adk/tools/skill_toolset.py | 125 +++++++-- tests/unittests/tools/test_skill_toolset.py | 250 +++++++++++++++++- 5 files changed, 378 insertions(+), 37 deletions(-) diff --git a/contributing/samples/skill_script_demo/agent.py b/contributing/samples/skill_script_demo/agent.py index 0daf1d14bc..3c400224d5 100644 --- a/contributing/samples/skill_script_demo/agent.py +++ b/contributing/samples/skill_script_demo/agent.py @@ -17,6 +17,10 @@ This agent loads the Anthropic mcp-builder skill and a python-helper skill, exercising all skill tools: list_skills, load_skill, load_skill_resource, and execute_skill_script. + +WARNING: This sample uses UnsafeLocalCodeExecutor, which runs scripts +in the host process with no sandboxing. For production use, prefer +ContainerCodeExecutor or VertexAICodeExecutor for isolation. """ import pathlib @@ -31,6 +35,8 @@ mcp_builder_skill = load_skill_from_dir(_SKILLS_DIR / "mcp-builder") python_helper_skill = load_skill_from_dir(_SKILLS_DIR / "python-helper") +# WARNING: UnsafeLocalCodeExecutor runs code in the host process. +# For production, use ContainerCodeExecutor or VertexAICodeExecutor. my_skill_toolset = SkillToolset( skills=[mcp_builder_skill, python_helper_skill], code_executor=UnsafeLocalCodeExecutor(), diff --git a/contributing/samples/skill_script_demo/run_live_test.py b/contributing/samples/skill_script_demo/run_live_test.py index 13d2cd52a6..2315e1c4da 100644 --- a/contributing/samples/skill_script_demo/run_live_test.py +++ b/contributing/samples/skill_script_demo/run_live_test.py @@ -26,6 +26,7 @@ import asyncio import os import pathlib +import subprocess import sys import traceback @@ -45,9 +46,16 @@ # Configure for Vertex AI project = os.environ.get("GOOGLE_CLOUD_PROJECT") if not project: - project = ( - os.popen("gcloud config get-value project 2>/dev/null").read().strip() - ) + try: + result = subprocess.run( + ["gcloud", "config", "get-value", "project"], + capture_output=True, + text=True, + timeout=10, + ) + project = result.stdout.strip() + except (FileNotFoundError, subprocess.TimeoutExpired): + project = "" if project: os.environ["GOOGLE_CLOUD_PROJECT"] = project diff --git a/contributing/samples/skill_script_demo/skills/mcp-builder/SKILL.md b/contributing/samples/skill_script_demo/skills/mcp-builder/SKILL.md index 8a1a77a47d..5151b00adf 100644 --- a/contributing/samples/skill_script_demo/skills/mcp-builder/SKILL.md +++ b/contributing/samples/skill_script_demo/skills/mcp-builder/SKILL.md @@ -55,15 +55,15 @@ Key pages to review: **Load framework documentation:** -- **MCP Best Practices**: [📋 View Best Practices](./reference/mcp_best_practices.md) - Core guidelines +- **MCP Best Practices**: [📋 View Best Practices](./references/mcp_best_practices.md) - Core guidelines **For TypeScript (recommended):** - **TypeScript SDK**: Use WebFetch to load `https://raw.githubusercontent.com/modelcontextprotocol/typescript-sdk/main/README.md` -- [⚡ TypeScript Guide](./reference/node_mcp_server.md) - TypeScript patterns and examples +- [⚡ TypeScript Guide](./references/node_mcp_server.md) - TypeScript patterns and examples **For Python:** - **Python SDK**: Use WebFetch to load `https://raw.githubusercontent.com/modelcontextprotocol/python-sdk/main/README.md` -- [🐍 Python Guide](./reference/python_mcp_server.md) - Python patterns and examples +- [🐍 Python Guide](./references/python_mcp_server.md) - Python patterns and examples #### 1.4 Plan Your Implementation @@ -80,8 +80,8 @@ Prioritize comprehensive API coverage. List endpoints to implement, starting wit #### 2.1 Set Up Project Structure See language-specific guides for project setup: -- [⚡ TypeScript Guide](./reference/node_mcp_server.md) - Project structure, package.json, tsconfig.json -- [🐍 Python Guide](./reference/python_mcp_server.md) - Module organization, dependencies +- [⚡ TypeScript Guide](./references/node_mcp_server.md) - Project structure, package.json, tsconfig.json +- [🐍 Python Guide](./references/python_mcp_server.md) - Module organization, dependencies #### 2.2 Implement Core Infrastructure @@ -152,7 +152,7 @@ See language-specific guides for detailed testing approaches and quality checkli After implementing your MCP server, create comprehensive evaluations to test its effectiveness. -**Load [✅ Evaluation Guide](./reference/evaluation.md) for complete evaluation guidelines.** +**Load [✅ Evaluation Guide](./references/evaluation.md) for complete evaluation guidelines.** #### 4.1 Understand Evaluation Purpose @@ -201,7 +201,7 @@ Load these resources as needed during development: ### Core MCP Documentation (Load First) - **MCP Protocol**: Start with sitemap at `https://modelcontextprotocol.io/sitemap.xml`, then fetch specific pages with `.md` suffix -- [📋 MCP Best Practices](./reference/mcp_best_practices.md) - Universal MCP guidelines including: +- [📋 MCP Best Practices](./references/mcp_best_practices.md) - Universal MCP guidelines including: - Server and tool naming conventions - Response format guidelines (JSON vs Markdown) - Pagination best practices @@ -213,14 +213,14 @@ Load these resources as needed during development: - **TypeScript SDK**: Fetch from `https://raw.githubusercontent.com/modelcontextprotocol/typescript-sdk/main/README.md` ### Language-Specific Implementation Guides (Load During Phase 2) -- [🐍 Python Implementation Guide](./reference/python_mcp_server.md) - Complete Python/FastMCP guide with: +- [🐍 Python Implementation Guide](./references/python_mcp_server.md) - Complete Python/FastMCP guide with: - Server initialization patterns - Pydantic model examples - Tool registration with `@mcp.tool` - Complete working examples - Quality checklist -- [⚡ TypeScript Implementation Guide](./reference/node_mcp_server.md) - Complete TypeScript guide with: +- [⚡ TypeScript Implementation Guide](./references/node_mcp_server.md) - Complete TypeScript guide with: - Project structure - Zod schema patterns - Tool registration with `server.registerTool` @@ -228,7 +228,7 @@ Load these resources as needed during development: - Quality checklist ### Evaluation Guide (Load During Phase 4) -- [✅ Evaluation Guide](./reference/evaluation.md) - Complete evaluation creation guide with: +- [✅ Evaluation Guide](./references/evaluation.md) - Complete evaluation creation guide with: - Question creation guidelines - Answer verification strategies - XML format specifications diff --git a/src/google/adk/tools/skill_toolset.py b/src/google/adk/tools/skill_toolset.py index 8d5a9adef3..806d1aee08 100644 --- a/src/google/adk/tools/skill_toolset.py +++ b/src/google/adk/tools/skill_toolset.py @@ -16,7 +16,9 @@ from __future__ import annotations +import json import logging +import shlex from typing import Any from typing import Optional from typing import TYPE_CHECKING @@ -39,6 +41,8 @@ logger = logging.getLogger("google_adk." + __name__) +_DEFAULT_SCRIPT_TIMEOUT = 300 + DEFAULT_SKILL_SYSTEM_INSTRUCTION = """You can use specialized 'skills' to help you with complex tasks. You MUST use the skill tools to interact with these skills. Skills are folders of instructions and resources that extend your capabilities for specialized tasks. Each skill folder contains: @@ -329,8 +333,22 @@ async def run_async( "error_code": "NO_CODE_EXECUTOR", } + # Validate input_args early (before sending to code executor) + if input_args: + try: + shlex.split(input_args) + except ValueError as e: + return { + "error": f"Invalid input_args: {e}", + "error_code": "INVALID_INPUT_ARGS", + } + # Prepare code based on script extension - code = self._prepare_code(script_name, script.src, input_args) + timeout = self._toolset._script_timeout + code = self._prepare_code(script_name, script.src, input_args, timeout) + is_shell = "." in script_name and script_name.rsplit(".", 1)[ + -1 + ].lower() in ("sh", "bash") if code is None: ext = script_name.rsplit(".", 1)[-1] if "." in script_name else "" return { @@ -346,12 +364,56 @@ async def run_async( tool_context._invocation_context, CodeExecutionInput(code=code), ) + stdout = result.stdout + stderr = result.stderr + # Shell scripts serialize both streams as JSON + # through stdout; parse the envelope if present. + if is_shell and stdout: + try: + parsed = json.loads(stdout) + if isinstance(parsed, dict) and parsed.get("__shell_result__"): + stdout = parsed.get("stdout", "") + stderr = parsed.get("stderr", "") + rc = parsed.get("returncode", 0) + if rc != 0 and not stderr: + stderr = f"Exit code {rc}" + except (json.JSONDecodeError, ValueError): + pass + if stderr and not stdout: + status = "error" + elif stderr: + status = "warning" + else: + status = "success" return { "skill_name": skill_name, "script_name": script_name, - "stdout": result.stdout, - "stderr": result.stderr, - "status": "error" if result.stderr else "success", + "stdout": stdout, + "stderr": stderr, + "status": status, + } + except SystemExit as e: + # Scripts may call sys.exit(); intercept instead of letting + # it terminate the host process. + exit_code = e.code if e.code is not None else 0 + if exit_code == 0: + # sys.exit(0) or sys.exit() is a normal termination. + return { + "skill_name": skill_name, + "script_name": script_name, + "stdout": "", + "stderr": "", + "status": "success", + } + logger.warning( + "Script '%s' from skill '%s' called sys.exit(%s)", + script_name, + skill_name, + exit_code, + ) + return { + "error": f"Script '{script_name}' exited with code {exit_code}.", + "error_code": "EXECUTION_ERROR", } except Exception as e: logger.exception( @@ -373,6 +435,7 @@ def _prepare_code( script_name: str, script_src: str, input_args: str, + timeout: int = _DEFAULT_SCRIPT_TIMEOUT, ) -> str | None: """Prepares Python code to execute the script. @@ -380,6 +443,7 @@ def _prepare_code( script_name: The script filename. script_src: The script source content. input_args: Optional arguments string. + timeout: Timeout in seconds for shell subprocess execution. Returns: Python code string to execute, or None if unsupported type. @@ -388,7 +452,7 @@ def _prepare_code( if "." in script_name: ext = script_name.rsplit(".", 1)[-1].lower() - if ext in ("py", ""): + if ext == "py": # Python script: execute directly, inject sys.argv if args if input_args: return ( @@ -400,20 +464,34 @@ def _prepare_code( return script_src elif ext in ("sh", "bash"): # Shell script: wrap in subprocess.run. - # Args are passed as separate list elements after the script - # name to avoid shell injection — bash -c receives the script - # source, and $0/$1/... get the positional parameters. + # Args are passed as separate list elements after the + # script name to avoid shell injection. + # Both streams are JSON-serialized through stdout since + # UnsafeLocalCodeExecutor drops stdout on exception. + cmd = f"['bash', '-c', {script_src!r}, {script_name!r}]" + if input_args: + cmd += f" + shlex.split({input_args!r})" return ( - "import subprocess, shlex\n" - "_result = subprocess.run(\n" - f" ['bash', '-c', {script_src!r}," - f" {script_name!r}]" - + (f" + shlex.split({input_args!r})" if input_args else "") - + ",\n" - " capture_output=True, text=True,\n" - " check=True,\n" - ")\n" - "print(_result.stdout, end='')\n" + "import subprocess, shlex, json as _json\n" + "try:\n" + " _r = subprocess.run(\n" + f" {cmd},\n" + " capture_output=True, text=True,\n" + f" timeout={timeout!r},\n" + " )\n" + " print(_json.dumps({\n" + " '__shell_result__': True,\n" + " 'stdout': _r.stdout,\n" + " 'stderr': _r.stderr,\n" + " 'returncode': _r.returncode,\n" + " }))\n" + "except subprocess.TimeoutExpired as _e:\n" + " print(_json.dumps({\n" + " '__shell_result__': True,\n" + " 'stdout': _e.stdout or '',\n" + f" 'stderr': 'Timed out after {timeout}s',\n" + " 'returncode': -1,\n" + " }))\n" ) return None @@ -427,7 +505,17 @@ def __init__( skills: list[models.Skill], *, code_executor: Optional[BaseCodeExecutor] = None, + script_timeout: int = _DEFAULT_SCRIPT_TIMEOUT, ): + """Initializes the SkillToolset. + + Args: + skills: List of skills to register. + code_executor: Optional code executor for script execution. + script_timeout: Timeout in seconds for shell script execution + via subprocess.run. Defaults to 300 seconds. Does not apply + to Python scripts executed via exec(). + """ super().__init__() # Check for duplicate skill names @@ -439,6 +527,7 @@ def __init__( self._skills = {skill.name: skill for skill in skills} self._code_executor = code_executor + self._script_timeout = script_timeout self._tools = [ ListSkillsTool(self), LoadSkillTool(self), diff --git a/tests/unittests/tools/test_skill_toolset.py b/tests/unittests/tools/test_skill_toolset.py index 78cf3a1273..2ca4f092d8 100644 --- a/tests/unittests/tools/test_skill_toolset.py +++ b/tests/unittests/tools/test_skill_toolset.py @@ -145,7 +145,6 @@ async def test_get_tools(mock_skill1, mock_skill2): assert isinstance(tools[3], skill_toolset.ExecuteSkillScriptTool) -@pytest.mark.asyncio @pytest.mark.asyncio async def test_list_skills_tool( mock_skill1, mock_skill2, tool_context_instance @@ -473,12 +472,12 @@ async def test_execute_script_shell_success(mock_skill1): assert result["status"] == "success" assert result["stdout"] == "setup\n" - # Verify the code wraps in subprocess.run with check=True + # Verify the code wraps in subprocess.run with JSON envelope call_args = executor.execute_code.call_args code_input = call_args[0][1] assert "subprocess.run" in code_input.code assert "bash" in code_input.code - assert "check=True" in code_input.code + assert "__shell_result__" in code_input.code @pytest.mark.asyncio @@ -598,8 +597,9 @@ async def test_execute_script_execution_error(mock_skill1): @pytest.mark.asyncio -async def test_execute_script_stderr_sets_error_status(mock_skill1): - executor = _make_mock_executor(stdout="", stderr="warning\n") +async def test_execute_script_stderr_only_sets_error_status(mock_skill1): + """stderr with no stdout should report error status.""" + executor = _make_mock_executor(stdout="", stderr="fatal error\n") toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) tool = skill_toolset.ExecuteSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() @@ -608,7 +608,23 @@ async def test_execute_script_stderr_sets_error_status(mock_skill1): tool_context=ctx, ) assert result["status"] == "error" - assert result["stderr"] == "warning\n" + assert result["stderr"] == "fatal error\n" + + +@pytest.mark.asyncio +async def test_execute_script_stderr_with_stdout_sets_warning(mock_skill1): + """stderr alongside stdout should report warning status.""" + executor = _make_mock_executor(stdout="output\n", stderr="deprecation\n") + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "run.py"}, + tool_context=ctx, + ) + assert result["status"] == "warning" + assert result["stdout"] == "output\n" + assert result["stderr"] == "deprecation\n" @pytest.mark.asyncio @@ -627,3 +643,225 @@ async def test_execute_script_execution_error_truncated(mock_skill1): # 200 chars of the message + "..." suffix + the prefix assert result["error"].endswith("...") assert len(result["error"]) < 300 + + +@pytest.mark.asyncio +async def test_execute_script_system_exit_caught(mock_skill1): + """sys.exit() in a script should not terminate the process.""" + executor = _make_mock_executor() + executor.execute_code.side_effect = SystemExit(1) + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "run.py"}, + tool_context=ctx, + ) + assert result["error_code"] == "EXECUTION_ERROR" + assert "exited with code 1" in result["error"] + + +@pytest.mark.asyncio +async def test_execute_script_system_exit_zero_is_success(mock_skill1): + """sys.exit(0) is a normal termination and should report success.""" + executor = _make_mock_executor() + executor.execute_code.side_effect = SystemExit(0) + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "run.py"}, + tool_context=ctx, + ) + assert result["status"] == "success" + assert "error" not in result + + +@pytest.mark.asyncio +async def test_execute_script_system_exit_none_is_success(mock_skill1): + """sys.exit() with no arg (None) should report success.""" + executor = _make_mock_executor() + executor.execute_code.side_effect = SystemExit(None) + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "run.py"}, + tool_context=ctx, + ) + assert result["status"] == "success" + assert "error" not in result + + +@pytest.mark.asyncio +async def test_execute_script_shell_includes_timeout(mock_skill1): + """Shell wrapper includes timeout in subprocess.run.""" + executor = _make_mock_executor(stdout="ok\n") + toolset = skill_toolset.SkillToolset( + [mock_skill1], code_executor=executor, script_timeout=60 + ) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "setup.sh"}, + tool_context=ctx, + ) + assert result["status"] == "success" + call_args = executor.execute_code.call_args + code_input = call_args[0][1] + assert "timeout=60" in code_input.code + + +@pytest.mark.asyncio +async def test_execute_script_extensionless_unsupported(mock_skill1): + """Files without extensions should return UNSUPPORTED_SCRIPT_TYPE.""" + # Add a script with no extension to the mock + original_side_effect = mock_skill1.resources.get_script.side_effect + + def get_script_extended(name): + if name == "noext": + return models.Script(src="print('hi')") + return original_side_effect(name) + + mock_skill1.resources.get_script.side_effect = get_script_extended + + executor = _make_mock_executor() + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_name": "noext"}, + tool_context=ctx, + ) + assert result["error_code"] == "UNSUPPORTED_SCRIPT_TYPE" + + +@pytest.mark.asyncio +async def test_execute_script_invalid_input_args(mock_skill1): + """Unclosed quotes in input_args should return INVALID_INPUT_ARGS.""" + executor = _make_mock_executor() + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={ + "skill_name": "skill1", + "script_name": "run.py", + "input_args": '--name "unclosed', + }, + tool_context=ctx, + ) + assert result["error_code"] == "INVALID_INPUT_ARGS" + + +# ── Integration tests using real UnsafeLocalCodeExecutor ── + + +def _make_skill_with_script(skill_name, script_name, script): + """Creates a minimal mock Skill with a single script.""" + skill = mock.create_autospec(models.Skill, instance=True) + skill.name = skill_name + skill.description = f"Test skill {skill_name}" + skill.instructions = "test instructions" + fm = mock.create_autospec(models.Frontmatter, instance=True) + fm.name = skill_name + fm.description = f"Test skill {skill_name}" + skill.frontmatter = fm + skill.resources = mock.MagicMock( + spec=["get_reference", "get_asset", "get_script"] + ) + + def get_script(name): + if name == script_name: + return script + return None + + skill.resources.get_script.side_effect = get_script + skill.resources.get_reference.return_value = None + skill.resources.get_asset.return_value = None + return skill + + +def _make_real_executor_toolset(skills, **kwargs): + """Creates a SkillToolset with a real UnsafeLocalCodeExecutor.""" + from google.adk.code_executors.unsafe_local_code_executor import UnsafeLocalCodeExecutor + + executor = UnsafeLocalCodeExecutor() + return skill_toolset.SkillToolset(skills, code_executor=executor, **kwargs) + + +@pytest.mark.asyncio +async def test_integration_python_stdout(): + """Real executor: Python script stdout is captured.""" + script = models.Script(src="print('hello world')") + skill = _make_skill_with_script("test_skill", "hello.py", script) + toolset = _make_real_executor_toolset([skill]) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={ + "skill_name": "test_skill", + "script_name": "hello.py", + }, + tool_context=ctx, + ) + assert result["status"] == "success" + assert result["stdout"] == "hello world\n" + assert result["stderr"] == "" + + +@pytest.mark.asyncio +async def test_integration_python_sys_exit_zero(): + """Real executor: sys.exit(0) is treated as success.""" + script = models.Script(src="import sys; sys.exit(0)") + skill = _make_skill_with_script("test_skill", "exit_zero.py", script) + toolset = _make_real_executor_toolset([skill]) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={ + "skill_name": "test_skill", + "script_name": "exit_zero.py", + }, + tool_context=ctx, + ) + assert result["status"] == "success" + + +@pytest.mark.asyncio +async def test_integration_shell_stdout_and_stderr(): + """Real executor: shell script preserves both stdout and stderr.""" + script = models.Script(src="echo output; echo warning >&2") + skill = _make_skill_with_script("test_skill", "both.sh", script) + toolset = _make_real_executor_toolset([skill]) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={ + "skill_name": "test_skill", + "script_name": "both.sh", + }, + tool_context=ctx, + ) + assert result["status"] == "warning" + assert "output" in result["stdout"] + assert "warning" in result["stderr"] + + +@pytest.mark.asyncio +async def test_integration_shell_stderr_only(): + """Real executor: shell script with only stderr reports error.""" + script = models.Script(src="echo failure >&2") + skill = _make_skill_with_script("test_skill", "err.sh", script) + toolset = _make_real_executor_toolset([skill]) + tool = skill_toolset.ExecuteSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={ + "skill_name": "test_skill", + "script_name": "err.sh", + }, + tool_context=ctx, + ) + assert result["status"] == "error" + assert "failure" in result["stderr"] From 52b8563b2a0531a2649076b55fa6523a1d89b200 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 21 Feb 2026 18:48:12 -0800 Subject: [PATCH 04/53] docs: Add code executor enhancements design document Design doc covering three high-priority improvements to the ADK code executor infrastructure: 1. Uniform timeout support across all executors (currently only GkeCodeExecutor has it) 2. Stateful Python execution in ContainerCodeExecutor (currently frozen to stateful=False) 3. Security hardening with a new LocalSandboxCodeExecutor and UnsafeLocalCodeExecutor restrictions Also removes dead default parameter from _prepare_code(timeout=) since timeout is now read directly from self._toolset._script_timeout inside the shell branch. Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 914 ++++++++++++++++++++++ src/google/adk/tools/skill_toolset.py | 6 +- 2 files changed, 916 insertions(+), 4 deletions(-) create mode 100644 docs/design/code_executor_enhancements.md diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md new file mode 100644 index 0000000000..94d19c3eee --- /dev/null +++ b/docs/design/code_executor_enhancements.md @@ -0,0 +1,914 @@ +# ADK Code Executor Enhancements — Design Document + +**Authors:** haiyuancao, Claude Code +**Date:** 2026-02-21 +**Status:** Draft +**Tracking:** Related to PR #4575 (ExecuteSkillScriptTool) + +--- + +## 1. Motivation + +The ADK code executor infrastructure (`src/google/adk/code_executors/`) is +the backbone for both LLM-driven code execution and skill script execution. +A review of the current implementations reveals three critical gaps that +limit production readiness: + +1. **No uniform timeout enforcement** — Only `GkeCodeExecutor` has a + `timeout_seconds` field. All other executors can hang indefinitely on + malicious, buggy, or slow code. The `ExecuteSkillScriptTool` works + around this for shell scripts by embedding `subprocess.run(timeout=N)` + in generated code, but this is a workaround, not a systemic solution. + +2. **No stateful Python execution** — `ContainerCodeExecutor` maintains a + persistent Docker container but explicitly freezes `stateful=False`. + Agents cannot preserve variables, imports, or working directory across + code execution calls. `VertexAiCodeExecutor` and + `AgentEngineSandboxCodeExecutor` allow `stateful=True` at the field + level but only support it for their specific backend APIs. + +3. **No safe local executor** — `UnsafeLocalCodeExecutor` runs `exec()` in + the host process with zero isolation. It is the only executor that + requires no external dependencies (no Docker, no GKE, no Vertex AI), + making it the default choice for development and demos. A compromised + script can read secrets, modify the filesystem, or crash the process. + +This document proposes solutions for all three gaps with minimal disruption +to the existing API. + +--- + +## 2. Current State + +### 2.1 Executor Landscape + +| Executor | Stateful | Timeout | Isolation | Dependencies | +|----------|----------|---------|-----------|-------------| +| `UnsafeLocalCodeExecutor` | No (frozen) | None | None | None | +| `ContainerCodeExecutor` | No (frozen) | None | Docker container | `docker` | +| `GkeCodeExecutor` | No (ephemeral) | `timeout_seconds=300` | gVisor sandbox | `kubernetes` | +| `VertexAiCodeExecutor` | Allowed | None | Vertex AI Extension | `vertexai` | +| `AgentEngineSandboxCodeExecutor` | Allowed | None | Vertex AI Sandbox | `vertexai` | +| `BuiltInCodeExecutor` | N/A | N/A | Gemini model | `google-genai` | + +### 2.2 Base Class Contract + +```python +class BaseCodeExecutor(BaseModel): + stateful: bool = False + error_retry_attempts: int = 2 + code_block_delimiters: List[tuple[str, str]] + execution_result_delimiters: tuple[str, str] + + @abc.abstractmethod + def execute_code( + self, + invocation_context: InvocationContext, + code_execution_input: CodeExecutionInput, + ) -> CodeExecutionResult: ... +``` + +### 2.3 Data Model + +```python +@dataclasses.dataclass +class CodeExecutionInput: + code: str + input_files: list[File] = field(default_factory=list) + execution_id: Optional[str] = None # For stateful execution + +@dataclasses.dataclass +class CodeExecutionResult: + stdout: str = '' + stderr: str = '' + output_files: list[File] = field(default_factory=list) +``` + +### 2.4 How Executors Are Used + +The primary consumer is `_code_execution.py` in the LLM flows layer: + +1. **Pre-processor**: Extracts data files, runs preprocessing code +2. **Post-processor**: Extracts code blocks from LLM responses, executes + them, feeds results back to the LLM +3. **Stateful support**: Uses `execution_id` (from `CodeExecutorContext`) + to maintain state across calls when `stateful=True` + +`ExecuteSkillScriptTool` is a secondary consumer that calls +`execute_code()` directly with generated Python code wrapping skill scripts. + +--- + +## 3. Proposal 1: Uniform Timeout Support + +### 3.1 Problem + +A code execution call can hang indefinitely. This is a denial-of-service +risk for any production deployment, whether the code comes from an LLM, a +skill script, or user input. + +| Executor | Current timeout behavior | +|----------|------------------------| +| `UnsafeLocalCodeExecutor` | `exec()` blocks forever | +| `ContainerCodeExecutor` | `exec_run()` blocks forever | +| `GkeCodeExecutor` | K8s watch timeout (works) | +| `VertexAiCodeExecutor` | Vertex AI internal timeout (opaque) | +| `AgentEngineSandboxCodeExecutor` | Vertex AI internal timeout (opaque) | + +### 3.2 Design + +#### 3.2.1 Add `timeout_seconds` to `BaseCodeExecutor` + +```python +class BaseCodeExecutor(BaseModel): + timeout_seconds: Optional[int] = None + """Maximum execution time in seconds. None means no timeout + (executor default behavior). Subclasses should enforce this + in their execute_code() implementation.""" +``` + +**Why `Optional[int]` with `None` default:** +- Backward compatible — existing code that doesn't set it works unchanged +- Allows subclasses to define their own defaults +- `None` means "use executor-specific default or no timeout" + +#### 3.2.2 `UnsafeLocalCodeExecutor` — Thread-Based Timeout + +`exec()` cannot be interrupted from the same thread. The solution is to +run it in a separate thread with a join timeout: + +```python +import threading + +def execute_code(self, invocation_context, code_execution_input): + timeout = self.timeout_seconds + if timeout is None: + # No timeout: current behavior (blocking exec) + return self._execute_inline(code_execution_input) + + # Run in a daemon thread with timeout + result_holder = {} + def _run(): + result_holder['result'] = self._execute_inline( + code_execution_input + ) + + thread = threading.Thread(target=_run, daemon=True) + thread.start() + thread.join(timeout=timeout) + + if thread.is_alive(): + # Thread is still running — timeout exceeded. + # Daemon thread will be killed when the process exits. + return CodeExecutionResult( + stderr=f'Execution timed out after {timeout}s' + ) + return result_holder.get('result', CodeExecutionResult( + stderr='Execution produced no result' + )) +``` + +**Trade-offs:** +- Daemon threads cannot be forcefully killed in CPython. If the code enters + a long-running C extension call (e.g., `time.sleep(9999)`), the thread + lingers until process exit. This is acceptable for a local development + executor — production deployments should use container-based executors. +- An alternative is `multiprocessing`, but that adds complexity around + serialization and shared state. + +**Recommendation:** Thread-based timeout for `UnsafeLocalCodeExecutor` is +sufficient. Document that it provides best-effort timeout only. + +#### 3.2.3 `ContainerCodeExecutor` — Docker `exec_run` Timeout + +Docker's `exec_run` does not natively support a timeout, but we can use +the Docker API's exec endpoint with a socket timeout: + +```python +def execute_code(self, invocation_context, code_execution_input): + timeout = self.timeout_seconds + + exec_result = self._container.exec_run( + ['python3', '-c', code_execution_input.code], + demux=True, + # Docker SDK does not support exec_run timeout directly. + # Use socket_timeout on the client instead. + ) + ... +``` + +**Better approach:** Use `threading.Timer` to kill the exec if it exceeds +the timeout: + +```python +import threading + +def execute_code(self, invocation_context, code_execution_input): + timeout = self.timeout_seconds + + # Create the exec instance + exec_id = self._client.api.exec_create( + self._container.id, + ['python3', '-c', code_execution_input.code], + )['Id'] + + # Start a timer to kill the exec if it exceeds timeout + timer = None + timed_out = threading.Event() + if timeout is not None: + def _kill(): + timed_out.set() + # Kill the exec'd process inside the container + self._container.exec_run( + ['kill', '-9', '-1'], # kill all procs + detach=True, + ) + timer = threading.Timer(timeout, _kill) + timer.start() + + try: + output = self._client.api.exec_start(exec_id, demux=True) + finally: + if timer is not None: + timer.cancel() + + if timed_out.is_set(): + return CodeExecutionResult( + stderr=f'Execution timed out after {timeout}s' + ) + # ... parse output as before +``` + +**Alternative — simpler approach:** Wrap `exec_run` in a thread with +join timeout (same pattern as `UnsafeLocalCodeExecutor`). Simpler, and +the container process can be cleaned up on next execution. + +**Recommendation:** Use the thread + join approach for consistency across +executors. Add a follow-up to use Docker's native exec kill for more +robust cleanup. + +#### 3.2.4 `GkeCodeExecutor` — Already Implemented + +`GkeCodeExecutor` already has `timeout_seconds: int = 300` applied to the +K8s watch API. Migrate it to use the base class field: + +```python +class GkeCodeExecutor(BaseCodeExecutor): + timeout_seconds: int = 300 # Override base default +``` + +No behavioral change needed. + +#### 3.2.5 Remote Executors (Vertex AI, Agent Engine) + +These executors delegate to Google Cloud APIs that have their own internal +timeouts. Adding client-side timeout is still valuable as a safety net: + +- Wrap the API call in `asyncio.wait_for()` or `threading.Timer` +- Return `CodeExecutionResult(stderr='...')` on timeout +- Log a warning that the server-side execution may still be running + +### 3.3 Migration Plan + +| Phase | Action | Risk | +|-------|--------|------| +| 1 | Add `timeout_seconds: Optional[int] = None` to `BaseCodeExecutor` | None (backward compatible) | +| 2 | Implement in `UnsafeLocalCodeExecutor` (thread + join) | Low | +| 3 | Implement in `ContainerCodeExecutor` (thread + join) | Low | +| 4 | Migrate `GkeCodeExecutor.timeout_seconds` to use base field | None | +| 5 | Add client-side timeout to remote executors | Low | + +### 3.4 Impact on `ExecuteSkillScriptTool` + +Once `BaseCodeExecutor` has native timeout support, the +`ExecuteSkillScriptTool` can optionally delegate timeout enforcement to +the executor rather than embedding it in generated shell wrapper code. +However, the shell wrapper timeout (`subprocess.run(timeout=N)`) should +be kept as defense-in-depth — it catches the subprocess even if the +executor timeout fails. + +--- + +## 4. Proposal 2: Stateful `ContainerCodeExecutor` + +### 4.1 Problem + +Agents often need multi-step code execution where later steps depend on +earlier results. For example: + +``` +Step 1: import pandas as pd; df = pd.read_csv('data.csv') +Step 2: filtered = df[df['status'] == 'active'] +Step 3: print(filtered.describe()) +``` + +Currently, each `execute_code()` call in `ContainerCodeExecutor` runs +`python3 -c ` — a fresh Python process with no memory of prior +calls. Step 2 would fail with `NameError: name 'df' is not defined`. + +### 4.2 Design + +#### 4.2.1 Architecture + +``` +┌─ ContainerCodeExecutor ─────────────────────┐ +│ │ +│ stateful=False (default): │ +│ exec_run(['python3', '-c', code]) │ +│ └─ Fresh process per call │ +│ │ +│ stateful=True: │ +│ exec_run(['python3', '-c', │ +│ 'exec(open("/app/session.py").read())' │ +│ + '\n' + code │ +│ + '\n_save_state()']) │ +│ └─ Loads prior state, executes, │ +│ saves state back │ +│ │ +└──────────────────────────────────────────────┘ +``` + +**Option A: Persistent Python Process (Complex)** + +Start a long-running Python REPL in the container and pipe code to its +stdin: + +```python +# On init (when stateful=True): +self._exec_id = self._client.api.exec_create( + self._container.id, + ['python3', '-i'], + stdin=True, stdout=True, stderr=True, +) +self._socket = self._client.api.exec_start( + self._exec_id, socket=True +) +``` + +Pros: True statefulness — all Python state (variables, imports, objects) +persists naturally. + +Cons: Complex I/O management. Need to detect when output is "done" for a +given code block (no clean delimiter). Prone to deadlocks. Hard to detect +crashes. + +**Option B: State Serialization via `dill` (Moderate)** + +After each execution, serialize the global namespace to a file in the +container. Before next execution, deserialize it: + +```python +STATEFUL_WRAPPER = ''' +import dill as _dill, os as _os +_state_file = '/tmp/.adk_state.pkl' +if _os.path.exists(_state_file): + _globals = _dill.load(open(_state_file, 'rb')) + globals().update(_globals) + +{user_code} + +# Save state after execution +_save_vars = {k: v for k, v in globals().items() + if not k.startswith('_')} +_dill.dump(_save_vars, open(_state_file, 'wb')) +''' +``` + +Pros: Simpler than persistent REPL. Each call is a clean `exec_run()`. + +Cons: Not all Python objects are serializable (e.g., open file handles, +database connections, generators). Adds `dill` dependency to the container +image. Serialization/deserialization overhead grows with state size. + +**Option C: Shared Globals File (Simple) — RECOMMENDED** + +Write executed code to a cumulative Python file. Each call appends the new +code block and re-executes the entire history: + +```python +HISTORY_FILE = '/tmp/.adk_history.py' + +def execute_code(self, invocation_context, code_execution_input): + if self.stateful: + # Append new code to history + self._container.exec_run( + ['sh', '-c', + f'cat >> {HISTORY_FILE} << "ADKEOF"\n' + f'{code_execution_input.code}\n' + f'ADKEOF'], + ) + # Execute full history + exec_result = self._container.exec_run( + ['python3', HISTORY_FILE], + demux=True, + ) + else: + exec_result = self._container.exec_run( + ['python3', '-c', code_execution_input.code], + demux=True, + ) +``` + +Pros: Simplest approach. No serialization issues. All Python features work. +No new dependencies. + +Cons: Re-executes entire history on each call — side effects run again. +Grows linearly with history length. + +**Mitigation for side effects:** Wrap prior code in a guard: + +```python +# Only new code produces output; prior blocks set up state silently +import sys, io +_old_stdout = sys.stdout +sys.stdout = io.StringIO() +# ... prior code blocks ... +sys.stdout = _old_stdout +# ... new code block (produces output) ... +``` + +#### 4.2.2 Recommended Approach + +**Option A (Persistent Process)** is the most robust for true statefulness +and is the standard approach used by Jupyter kernels and similar systems. +Despite its complexity, it provides the best user experience: + +- Variables, imports, and objects persist naturally +- No re-execution of side effects +- No serialization issues +- O(1) cost per call (not O(n) like Option C) + +However, implementing a full REPL protocol is a significant engineering +effort. We recommend a **phased approach**: + +**Phase 1 (MVP):** Option C (cumulative file) with stdout suppression for +prior blocks. Simple, works for the common case (data analysis, variable +setup). + +**Phase 2 (Full):** Option A (persistent process) with a proper +execution protocol using sentinel markers for output boundaries. + +#### 4.2.3 Implementation Plan (Phase 1) + +1. **Unfreeze `stateful` in `ContainerCodeExecutor`:** + +```python +# Remove frozen=True +stateful: bool = False +``` + +2. **Add a code history list:** + +```python +_code_history: list[str] = [] +``` + +3. **Modify `execute_code()`:** + +```python +def execute_code(self, invocation_context, code_execution_input): + code = code_execution_input.code + + if self.stateful: + # Build cumulative script + setup_code = '\n'.join( + f'# --- Block {i} ---\n{block}' + for i, block in enumerate(self._code_history) + ) + # Suppress stdout for prior blocks + full_code = ( + 'import sys as _sys, io as _io\n' + '_sys.stdout = _io.StringIO()\n' + f'{setup_code}\n' + '_sys.stdout = _sys.__stdout__\n' + f'{code}\n' + ) + self._code_history.append(code) + else: + full_code = code + + exec_result = self._container.exec_run( + ['python3', '-c', full_code], + demux=True, + ) + # ... parse output as before +``` + +4. **Add `reset_state()` method:** + +```python +def reset_state(self): + """Clears the execution history for stateful mode.""" + self._code_history.clear() +``` + +5. **Update the `__init__` validation:** + +```python +# Remove the ValueError for stateful=True +# Keep optimize_data_file frozen +``` + +#### 4.2.4 Interaction with `execution_id` + +The LLM flow layer uses `execution_id` (from `CodeExecutorContext`) to +identify stateful sessions. For `ContainerCodeExecutor`: + +- Each `execution_id` maps to a separate code history +- Use a dict: `_code_histories: dict[str, list[str]] = {}` +- When `execution_id` is provided in `CodeExecutionInput`, use the + corresponding history +- When `execution_id` is `None`, use default (empty) history + +This aligns with how `VertexAiCodeExecutor` uses `session_id`. + +### 4.3 Testing Plan + +| Test | Description | +|------|-------------| +| `test_stateful_variable_persistence` | Define variable in call 1, access in call 2 | +| `test_stateful_import_persistence` | Import in call 1, use in call 2 | +| `test_stateful_no_stdout_leakage` | Prior blocks' print() should not appear in later output | +| `test_stateful_error_in_later_block` | Error in call 2 should not corrupt state | +| `test_stateful_reset` | `reset_state()` clears history | +| `test_stateless_unchanged` | Default `stateful=False` behavior unchanged | +| `test_execution_id_isolation` | Different `execution_id` values use separate histories | + +--- + +## 5. Proposal 3: Security Hardening + +### 5.1 Problem + +`UnsafeLocalCodeExecutor` is the default executor for local development +because it requires no external dependencies. But it runs `exec()` in the +host Python process with full access to: + +- The filesystem (read/write any file) +- Environment variables (including API keys, secrets) +- Network (outbound HTTP, DNS) +- The ADK process itself (`os.kill`, `sys.exit`) + +This is a critical security concern when executing: +- LLM-generated code (prompt injection → arbitrary code execution) +- Third-party skill scripts (supply chain risk) +- User-provided code in multi-tenant deployments + +### 5.2 Threat Model + +| Threat | Impact | Current mitigation | +|--------|--------|--------------------| +| LLM generates malicious code | Full host compromise | None | +| Skill script reads secrets | Data exfiltration | None (documented warning only) | +| Infinite loop / fork bomb | DoS / resource exhaustion | None (no timeout) | +| `sys.exit()` in script | Process termination | Partial (`SystemExit` catch in `ExecuteSkillScriptTool`) | +| Network exfiltration | Data leak | None | +| File system manipulation | Data loss / corruption | None | + +### 5.3 Design + +We propose a layered approach with three tiers of security: + +#### 5.3.1 Tier 1: `UnsafeLocalCodeExecutor` Hardening (Quick Wins) + +These changes improve safety without changing the fundamental architecture: + +**A. Timeout support** (covered in Proposal 1) + +**B. Restricted builtins:** + +```python +_BLOCKED_BUILTINS = { + 'exec', 'eval', 'compile', # Prevent meta-execution + '__import__', # Prevent arbitrary imports + 'open', # Prevent file access + 'breakpoint', # Prevent debugger attach +} + +_BLOCKED_MODULES = { + 'os', 'subprocess', 'shutil', # System access + 'socket', 'http', 'urllib', # Network access + 'ctypes', 'cffi', # Native code + 'importlib', # Dynamic imports +} +``` + +**Trade-off:** This breaks legitimate use cases (e.g., data analysis scripts +that need `open()` for file I/O). It should be opt-in: + +```python +class UnsafeLocalCodeExecutor(BaseCodeExecutor): + restrict_builtins: bool = False + """When True, block dangerous builtins (exec, eval, open, + __import__). Default False for backward compatibility.""" +``` + +**C. Warning on first use:** + +```python +import warnings + +def execute_code(self, ...): + if not self._warned: + warnings.warn( + 'UnsafeLocalCodeExecutor runs code in the host ' + 'process with NO isolation. Use ' + 'ContainerCodeExecutor or GkeCodeExecutor for ' + 'production deployments.', + SecurityWarning, + stacklevel=2, + ) + self._warned = True + ... +``` + +#### 5.3.2 Tier 2: `LocalSandboxCodeExecutor` (New, Recommended) + +A new executor that provides meaningful isolation without requiring Docker +or cloud services: + +**Approach: `subprocess` with resource limits** + +```python +class LocalSandboxCodeExecutor(BaseCodeExecutor): + """Executes Python code in a sandboxed subprocess. + + Provides isolation via: + - Separate process (no shared memory with host) + - Resource limits via ulimit (CPU time, memory) + - Restricted environment variables + - Optional chroot or tmpdir working directory + """ + + timeout_seconds: int = 30 + max_memory_mb: int = 256 + max_cpu_seconds: int = 30 + allowed_env_vars: list[str] = [] + + def execute_code(self, invocation_context, code_execution_input): + import subprocess + import tempfile + + with tempfile.NamedTemporaryFile( + mode='w', suffix='.py', delete=True + ) as f: + f.write(code_execution_input.code) + f.flush() + + # Build restricted environment + env = {k: os.environ[k] for k in self.allowed_env_vars + if k in os.environ} + env['PATH'] = '/usr/bin:/usr/local/bin' + + # Set resource limits via preexec_fn + def set_limits(): + import resource + # CPU time limit + resource.setrlimit( + resource.RLIMIT_CPU, + (self.max_cpu_seconds, self.max_cpu_seconds) + ) + # Memory limit + mem_bytes = self.max_memory_mb * 1024 * 1024 + resource.setrlimit( + resource.RLIMIT_AS, + (mem_bytes, mem_bytes) + ) + + result = subprocess.run( + ['python3', f.name], + capture_output=True, + text=True, + timeout=self.timeout_seconds, + env=env, + preexec_fn=set_limits, # Unix only + cwd=tempfile.gettempdir(), + ) + + return CodeExecutionResult( + stdout=result.stdout, + stderr=result.stderr if result.returncode != 0 + else '', + ) +``` + +**Platform considerations:** +- `resource.setrlimit` is Unix-only (Linux, macOS) +- On Windows, use `subprocess.CREATE_NO_WINDOW` and + `subprocess.Popen` with `creationflags` for job object limits +- Fallback to timeout-only on platforms without `resource` module + +**Dependencies:** None (stdlib only). This is the key advantage over +`ContainerCodeExecutor`. + +**Limitations:** +- Less isolation than containers (shared filesystem, kernel) +- `preexec_fn` is not fork-safe with threads (use `process_group` on + Python 3.11+) +- Cannot restrict network access without OS-level firewall rules + +#### 5.3.3 Tier 3: Promote `ContainerCodeExecutor` as Default + +For production deployments, container-based isolation should be the +standard recommendation: + +**A. Simplify setup:** + +```python +# Current: requires explicit image or docker_path +executor = ContainerCodeExecutor(image='python:3.11-slim') + +# Proposed: auto-pull default image +executor = ContainerCodeExecutor() # Uses python:3.11-slim +``` + +**B. Pre-built ADK executor image:** + +Create an official `adk-code-executor` Docker image with: +- Python 3.11+ slim base +- Common data science libraries (pandas, numpy, matplotlib) +- Non-root user +- Read-only filesystem (writable `/tmp` only) +- No network access by default (`--network=none`) + +```dockerfile +FROM python:3.11-slim +RUN pip install --no-cache-dir pandas numpy matplotlib +RUN useradd -m -s /bin/bash executor +USER executor +WORKDIR /home/executor +``` + +**C. Network isolation by default:** + +```python +def __init_container(self): + self._container = self._client.containers.run( + image=self.image, + detach=True, + tty=True, + network_mode='none', # No network access + read_only=True, # Read-only filesystem + tmpfs={'/tmp': 'size=100M'}, # Writable tmp + mem_limit='512m', # Memory limit + cpu_period=100000, + cpu_quota=50000, # 50% of one CPU + ) +``` + +### 5.4 Recommendation Matrix + +| Use Case | Recommended Executor | Why | +|----------|---------------------|-----| +| Local development | `LocalSandboxCodeExecutor` | No deps, basic isolation | +| Quick prototyping | `UnsafeLocalCodeExecutor` | Fastest setup, no isolation | +| CI/CD testing | `ContainerCodeExecutor` | Docker available in CI | +| Production (single tenant) | `ContainerCodeExecutor` | Good isolation, local | +| Production (multi-tenant) | `GkeCodeExecutor` | gVisor, per-execution isolation | +| Google Cloud | `AgentEngineSandboxCodeExecutor` | Managed, scalable | + +### 5.5 Implementation Plan + +| Phase | Action | Effort | Risk | +|-------|--------|--------|------| +| 1 | Add `SecurityWarning` to `UnsafeLocalCodeExecutor` | Small | None | +| 2 | Add `restrict_builtins` option | Small | Low | +| 3 | Implement `LocalSandboxCodeExecutor` | Medium | Low | +| 4 | Add default image + network isolation to `ContainerCodeExecutor` | Medium | Low | +| 5 | Create official `adk-code-executor` Docker image | Medium | Low | +| 6 | Update documentation and samples | Small | None | + +--- + +## 6. Cross-Cutting Concerns + +### 6.1 `BaseCodeExecutor` API Changes + +All three proposals touch `BaseCodeExecutor`. The combined changes: + +```python +class BaseCodeExecutor(BaseModel): + # Existing fields (unchanged) + optimize_data_file: bool = False + stateful: bool = False + error_retry_attempts: int = 2 + code_block_delimiters: List[tuple[str, str]] = [...] + execution_result_delimiters: tuple[str, str] = (...) + + # NEW: Proposal 1 + timeout_seconds: Optional[int] = None + """Maximum execution time in seconds. None = no timeout.""" + + @abc.abstractmethod + def execute_code( + self, + invocation_context: InvocationContext, + code_execution_input: CodeExecutionInput, + ) -> CodeExecutionResult: ... +``` + +### 6.2 Backward Compatibility + +| Change | Backward compatible? | Migration needed? | +|--------|---------------------|------------------| +| `timeout_seconds` on base class | Yes (default `None`) | No | +| Unfreeze `stateful` on `ContainerCodeExecutor` | Yes (default `False`) | No | +| `SecurityWarning` on `UnsafeLocalCodeExecutor` | Yes (warning only) | No | +| New `LocalSandboxCodeExecutor` | Yes (additive) | No | +| `restrict_builtins` on `UnsafeLocalCodeExecutor` | Yes (default `False`) | No | +| Default image for `ContainerCodeExecutor` | Breaking (currently requires image/docker_path) | Minor | + +### 6.3 Impact on `ExecuteSkillScriptTool` + +| Feature | Current workaround | After enhancements | +|---------|-------------------|-------------------| +| Shell timeout | Embedded `subprocess.run(timeout=N)` | Keep as defense-in-depth | +| Python timeout | None | Executor-level timeout handles it | +| Isolation | Documentation warning only | `LocalSandboxCodeExecutor` or container | +| Stateful scripts | Not supported | Available via `ContainerCodeExecutor(stateful=True)` | + +### 6.4 Testing Strategy + +| Category | Approach | +|----------|----------| +| Unit tests | Mock-based tests for each executor (existing pattern) | +| Integration tests | Real executor tests (like the ones added for `ExecuteSkillScriptTool`) | +| Timeout tests | Scripts with `time.sleep()` to verify timeout enforcement | +| Security tests | Scripts attempting blocked operations to verify restrictions | +| Stateful tests | Multi-call sequences verifying variable persistence | + +--- + +## 7. Implementation Roadmap + +### Phase 1: Timeout (2-3 days) + +1. Add `timeout_seconds: Optional[int] = None` to `BaseCodeExecutor` +2. Implement thread-based timeout in `UnsafeLocalCodeExecutor` +3. Implement thread-based timeout in `ContainerCodeExecutor` +4. Migrate `GkeCodeExecutor.timeout_seconds` to base class field +5. Add timeout tests for each executor +6. Update `ExecuteSkillScriptTool` to set executor timeout when available + +### Phase 2: Stateful Container (3-5 days) + +1. Unfreeze `stateful` on `ContainerCodeExecutor` +2. Implement cumulative code history with stdout suppression +3. Add `execution_id`-based history isolation +4. Add `reset_state()` method +5. Add stateful execution tests +6. Update samples and documentation + +### Phase 3: Security Hardening (5-7 days) + +1. Add `SecurityWarning` to `UnsafeLocalCodeExecutor` +2. Add `restrict_builtins` option +3. Implement `LocalSandboxCodeExecutor` +4. Add default image support to `ContainerCodeExecutor` +5. Add network isolation defaults to `ContainerCodeExecutor` +6. Create official `adk-code-executor` Docker image +7. Update all samples to recommend secure executors +8. Add security-focused tests + +### Total estimated effort: 10-15 days + +--- + +## 8. Open Questions + +1. **Should `timeout_seconds` be enforced at the base class level?** + We could add a wrapper in `BaseCodeExecutor.execute_code()` that + enforces the timeout generically, rather than requiring each subclass + to implement it. However, this would require the base class to manage + threading, which may not be appropriate for remote executors. + +2. **Should `LocalSandboxCodeExecutor` support stateful execution?** + Subprocess-based execution is inherently stateless. Stateful support + would require the same workarounds as `ContainerCodeExecutor` (Option + C). We recommend keeping it stateless in the initial implementation. + +3. **Should we deprecate `UnsafeLocalCodeExecutor`?** + Not immediately. It's useful for zero-dependency development. But the + `SecurityWarning` and documentation should steer users toward safer + alternatives for anything beyond local prototyping. + +4. **How should `ContainerCodeExecutor` handle container crashes in + stateful mode?** + If the container crashes (OOM, segfault), the code history is lost. + Options: (a) re-create container and replay history, (b) return error + and let user restart, (c) persist history to host volume. Recommend + (b) for simplicity. + +--- + +## 9. References + +- [BaseCodeExecutor](../../src/google/adk/code_executors/base_code_executor.py) +- [ContainerCodeExecutor](../../src/google/adk/code_executors/container_code_executor.py) +- [UnsafeLocalCodeExecutor](../../src/google/adk/code_executors/unsafe_local_code_executor.py) +- [GkeCodeExecutor](../../src/google/adk/code_executors/gke_code_executor.py) +- [VertexAiCodeExecutor](../../src/google/adk/code_executors/vertex_ai_code_executor.py) +- [AgentEngineSandboxCodeExecutor](../../src/google/adk/code_executors/agent_engine_sandbox_code_executor.py) +- [ExecuteSkillScriptTool](../../src/google/adk/tools/skill_toolset.py) +- [Code Execution Flow](../../src/google/adk/flows/llm_flows/_code_execution.py) +- [PR #4575 — ExecuteSkillScriptTool](https://github.com/google/adk-python/pull/4575) diff --git a/src/google/adk/tools/skill_toolset.py b/src/google/adk/tools/skill_toolset.py index 806d1aee08..dd37b3c1ac 100644 --- a/src/google/adk/tools/skill_toolset.py +++ b/src/google/adk/tools/skill_toolset.py @@ -344,8 +344,7 @@ async def run_async( } # Prepare code based on script extension - timeout = self._toolset._script_timeout - code = self._prepare_code(script_name, script.src, input_args, timeout) + code = self._prepare_code(script_name, script.src, input_args) is_shell = "." in script_name and script_name.rsplit(".", 1)[ -1 ].lower() in ("sh", "bash") @@ -435,7 +434,6 @@ def _prepare_code( script_name: str, script_src: str, input_args: str, - timeout: int = _DEFAULT_SCRIPT_TIMEOUT, ) -> str | None: """Prepares Python code to execute the script. @@ -443,7 +441,6 @@ def _prepare_code( script_name: The script filename. script_src: The script source content. input_args: Optional arguments string. - timeout: Timeout in seconds for shell subprocess execution. Returns: Python code string to execute, or None if unsupported type. @@ -468,6 +465,7 @@ def _prepare_code( # script name to avoid shell injection. # Both streams are JSON-serialized through stdout since # UnsafeLocalCodeExecutor drops stdout on exception. + timeout = self._toolset._script_timeout cmd = f"['bash', '-c', {script_src!r}, {script_name!r}]" if input_args: cmd += f" + shlex.split({input_args!r})" From 4142c37c2e91592a663bcc3aa0dbeef00087e0cf Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 21 Feb 2026 19:04:06 -0800 Subject: [PATCH 05/53] docs: Address 8 architectural review findings in code executor design doc Add Non-Goals & Invariants section, fix append-before-execute state poisoning, redesign timeout as per-invocation field on CodeExecutionInput, use Docker exec kill instead of thread+join for container timeout, replace preexec_fn with process_group, pin default container image digest, reframe restricted builtins as best-effort friction, document execution_id gap in ExecuteSkillScriptTool. Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 509 +++++++++++++++------- 1 file changed, 362 insertions(+), 147 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index 94d19c3eee..4a2aa4f7b1 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -38,9 +38,50 @@ to the existing API. --- -## 2. Current State +## 2. Non-Goals & Invariants + +The following are explicitly **out of scope** for this design: + +1. **Full sandboxing of `UnsafeLocalCodeExecutor`** — The restricted + builtins mechanism (Tier 1, §6.3.1-B) is a best-effort friction layer, + not a security boundary. Any determined code can bypass it via + `object.__subclasses__()`, `importlib` through `__builtins__`, etc. + True isolation requires a process or container boundary. + +2. **Idempotent replay of side-effecting code** — The stateful + `ContainerCodeExecutor` (Proposal 2) replays prior code blocks. + Code with non-idempotent side effects (file writes, network calls, + database mutations) is **not supported** in stateful replay mode. + The design suppresses stdout but cannot suppress arbitrary I/O. + Users must keep side-effecting code in the final block or use the + persistent-process approach (Phase 2 / Option A). + +3. **Multi-tenant per-execution isolation** — Per-execution isolation + (fresh sandbox per call) is the domain of `GkeCodeExecutor` and + cloud-hosted executors. Container and local executors share a + single execution environment within a session. + +4. **Windows support for `LocalSandboxCodeExecutor`** — + `resource.setrlimit` and `process_group` are Unix-only. Windows + support is deferred to a future iteration. + +**Key invariants:** + +- `timeout_seconds` is a **per-invocation** parameter, not executor-global + state. When a single executor instance is shared across agents/tools, + each `execute_code()` call may specify its own timeout via + `CodeExecutionInput.timeout_seconds`. +- Code is appended to stateful history **only after** successful execution. + A failing code block is never replayed. +- Executor instances are not thread-safe unless documented otherwise. + Concurrent `execute_code()` calls on the same instance require external + synchronization. -### 2.1 Executor Landscape +--- + +## 3. Current State + +### 3.1 Executor Landscape | Executor | Stateful | Timeout | Isolation | Dependencies | |----------|----------|---------|-----------|-------------| @@ -51,7 +92,7 @@ to the existing API. | `AgentEngineSandboxCodeExecutor` | Allowed | None | Vertex AI Sandbox | `vertexai` | | `BuiltInCodeExecutor` | N/A | N/A | Gemini model | `google-genai` | -### 2.2 Base Class Contract +### 3.2 Base Class Contract ```python class BaseCodeExecutor(BaseModel): @@ -68,7 +109,7 @@ class BaseCodeExecutor(BaseModel): ) -> CodeExecutionResult: ... ``` -### 2.3 Data Model +### 3.3 Data Model ```python @dataclasses.dataclass @@ -84,7 +125,7 @@ class CodeExecutionResult: output_files: list[File] = field(default_factory=list) ``` -### 2.4 How Executors Are Used +### 3.4 How Executors Are Used The primary consumer is `_code_execution.py` in the LLM flows layer: @@ -99,9 +140,9 @@ The primary consumer is `_code_execution.py` in the LLM flows layer: --- -## 3. Proposal 1: Uniform Timeout Support +## 4. Proposal 1: Uniform Timeout Support -### 3.1 Problem +### 4.1 Problem A code execution call can hang indefinitely. This is a denial-of-service risk for any production deployment, whether the code comes from an LLM, a @@ -115,24 +156,55 @@ skill script, or user input. | `VertexAiCodeExecutor` | Vertex AI internal timeout (opaque) | | `AgentEngineSandboxCodeExecutor` | Vertex AI internal timeout (opaque) | -### 3.2 Design +### 4.2 Design + +#### 4.2.1 Add `timeout_seconds` to `CodeExecutionInput` -#### 3.2.1 Add `timeout_seconds` to `BaseCodeExecutor` +Timeout is a **per-invocation** concern, not executor-global state. A single +executor instance may be shared across agents, tools, and concurrent calls +with different timeout requirements (e.g., a quick validation script vs. a +long-running data analysis). Placing timeout on the executor would create +race conditions when multiple callers set different values. ```python -class BaseCodeExecutor(BaseModel): - timeout_seconds: Optional[int] = None +@dataclasses.dataclass +class CodeExecutionInput: + code: str + input_files: list[File] = field(default_factory=list) + execution_id: Optional[str] = None + timeout_seconds: Optional[int] = None # NEW """Maximum execution time in seconds. None means no timeout - (executor default behavior). Subclasses should enforce this - in their execute_code() implementation.""" + (executor default behavior). Each execute_code() call reads + this from its input, not from executor-level state.""" +``` + +Additionally, a **default** timeout on `BaseCodeExecutor` serves as +a fallback when callers don't specify one: + +```python +class BaseCodeExecutor(BaseModel): + default_timeout_seconds: Optional[int] = None + """Default timeout applied when CodeExecutionInput.timeout_seconds + is None. Subclasses may override (e.g., GkeCodeExecutor defaults + to 300). None means no timeout.""" +``` + +The effective timeout is resolved as: +```python +timeout = ( + code_execution_input.timeout_seconds + ?? self.default_timeout_seconds +) ``` -**Why `Optional[int]` with `None` default:** +**Why per-invocation + executor default:** - Backward compatible — existing code that doesn't set it works unchanged -- Allows subclasses to define their own defaults -- `None` means "use executor-specific default or no timeout" +- Safe for shared executors — no global mutable state +- Callers can override per-call (e.g., `ExecuteSkillScriptTool` sets + `script_timeout`, LLM flows use a different default) +- Executor subclasses can define their own defaults -#### 3.2.2 `UnsafeLocalCodeExecutor` — Thread-Based Timeout +#### 4.2.2 `UnsafeLocalCodeExecutor` — Thread-Based Timeout `exec()` cannot be interrupted from the same thread. The solution is to run it in a separate thread with a join timeout: @@ -141,7 +213,10 @@ run it in a separate thread with a join timeout: import threading def execute_code(self, invocation_context, code_execution_input): - timeout = self.timeout_seconds + timeout = ( + code_execution_input.timeout_seconds + or self.default_timeout_seconds + ) if timeout is None: # No timeout: current behavior (blocking exec) return self._execute_inline(code_execution_input) @@ -179,32 +254,24 @@ def execute_code(self, invocation_context, code_execution_input): **Recommendation:** Thread-based timeout for `UnsafeLocalCodeExecutor` is sufficient. Document that it provides best-effort timeout only. -#### 3.2.3 `ContainerCodeExecutor` — Docker `exec_run` Timeout +#### 4.2.3 `ContainerCodeExecutor` — Docker Exec Kill on Timeout -Docker's `exec_run` does not natively support a timeout, but we can use -the Docker API's exec endpoint with a socket timeout: +**Problem with thread+join in containers:** Unlike `UnsafeLocalCodeExecutor` +where a lingering daemon thread is merely wasteful, a runaway process inside +a shared container consumes CPU/memory and can interfere with subsequent +executions. The thread+join pattern would leave the container process +running indefinitely after the join timeout expires. -```python -def execute_code(self, invocation_context, code_execution_input): - timeout = self.timeout_seconds - - exec_result = self._container.exec_run( - ['python3', '-c', code_execution_input.code], - demux=True, - # Docker SDK does not support exec_run timeout directly. - # Use socket_timeout on the client instead. - ) - ... -``` - -**Better approach:** Use `threading.Timer` to kill the exec if it exceeds -the timeout: +**Primary approach: Docker exec kill via `exec_inspect` + PID kill.** ```python import threading def execute_code(self, invocation_context, code_execution_input): - timeout = self.timeout_seconds + timeout = ( + code_execution_input.timeout_seconds + or self.default_timeout_seconds + ) # Create the exec instance exec_id = self._client.api.exec_create( @@ -212,18 +279,29 @@ def execute_code(self, invocation_context, code_execution_input): ['python3', '-c', code_execution_input.code], )['Id'] - # Start a timer to kill the exec if it exceeds timeout + # Start a timer to kill the exec's PID on timeout timer = None timed_out = threading.Event() if timeout is not None: - def _kill(): + def _kill_exec(): timed_out.set() - # Kill the exec'd process inside the container - self._container.exec_run( - ['kill', '-9', '-1'], # kill all procs - detach=True, - ) - timer = threading.Timer(timeout, _kill) + try: + # Get the PID of the exec'd process + info = self._client.api.exec_inspect(exec_id) + pid = info.get('Pid', 0) + if pid > 0: + self._container.exec_run( + ['kill', '-9', str(pid)], + detach=True, + ) + except Exception: + # Fallback: kill all non-init processes + self._container.exec_run( + ['sh', '-c', + 'kill -9 $(ps -o pid= | grep -v "^\\s*1$")'], + detach=True, + ) + timer = threading.Timer(timeout, _kill_exec) timer.start() try: @@ -239,27 +317,44 @@ def execute_code(self, invocation_context, code_execution_input): # ... parse output as before ``` -**Alternative — simpler approach:** Wrap `exec_run` in a thread with -join timeout (same pattern as `UnsafeLocalCodeExecutor`). Simpler, and -the container process can be cleaned up on next execution. +**Why targeted PID kill, not `kill -9 -1`:** +- `kill -9 -1` kills *all* processes in the container, including the + init/shell process that keeps the container alive. This would force a + container restart on the next call. +- Targeted kill via `exec_inspect` → PID only terminates the timed-out + process, leaving the container healthy for subsequent calls. +- The fallback (`kill all non-init`) is a safety net if `exec_inspect` + fails (e.g., Docker API version mismatch). + +**Alternative — container restart on timeout:** Simpler but more costly. +Stop and restart the container after timeout. Acceptable if stateful mode +is not in use (no accumulated state to preserve). -**Recommendation:** Use the thread + join approach for consistency across -executors. Add a follow-up to use Docker's native exec kill for more -robust cleanup. +**Recommendation:** Use Docker exec kill as the primary approach. This is +more robust than thread+join and properly cleans up runaway processes. -#### 3.2.4 `GkeCodeExecutor` — Already Implemented +#### 4.2.4 `GkeCodeExecutor` — Already Implemented `GkeCodeExecutor` already has `timeout_seconds: int = 300` applied to the -K8s watch API. Migrate it to use the base class field: +K8s watch API. Migrate it to use the base class default field: ```python class GkeCodeExecutor(BaseCodeExecutor): - timeout_seconds: int = 300 # Override base default + default_timeout_seconds: int = 300 # Override base default +``` + +In `execute_code()`, resolve timeout from per-invocation input first: +```python +timeout = ( + code_execution_input.timeout_seconds + or self.default_timeout_seconds +) ``` -No behavioral change needed. +No behavioral change for existing callers (they don't set per-invocation +timeout, so the 300s default applies as before). -#### 3.2.5 Remote Executors (Vertex AI, Agent Engine) +#### 4.2.5 Remote Executors (Vertex AI, Agent Engine) These executors delegate to Google Cloud APIs that have their own internal timeouts. Adding client-side timeout is still valuable as a safety net: @@ -268,17 +363,18 @@ timeouts. Adding client-side timeout is still valuable as a safety net: - Return `CodeExecutionResult(stderr='...')` on timeout - Log a warning that the server-side execution may still be running -### 3.3 Migration Plan +### 4.3 Migration Plan | Phase | Action | Risk | |-------|--------|------| -| 1 | Add `timeout_seconds: Optional[int] = None` to `BaseCodeExecutor` | None (backward compatible) | -| 2 | Implement in `UnsafeLocalCodeExecutor` (thread + join) | Low | -| 3 | Implement in `ContainerCodeExecutor` (thread + join) | Low | -| 4 | Migrate `GkeCodeExecutor.timeout_seconds` to use base field | None | -| 5 | Add client-side timeout to remote executors | Low | +| 1 | Add `timeout_seconds: Optional[int] = None` to `CodeExecutionInput` | None (backward compatible) | +| 2 | Add `default_timeout_seconds: Optional[int] = None` to `BaseCodeExecutor` | None (backward compatible) | +| 3 | Implement in `UnsafeLocalCodeExecutor` (thread + join) | Low | +| 4 | Implement in `ContainerCodeExecutor` (Docker exec kill) | Low | +| 5 | Migrate `GkeCodeExecutor.timeout_seconds` to `default_timeout_seconds` | None | +| 6 | Add client-side timeout to remote executors | Low | -### 3.4 Impact on `ExecuteSkillScriptTool` +### 4.4 Impact on `ExecuteSkillScriptTool` Once `BaseCodeExecutor` has native timeout support, the `ExecuteSkillScriptTool` can optionally delegate timeout enforcement to @@ -289,9 +385,9 @@ executor timeout fails. --- -## 4. Proposal 2: Stateful `ContainerCodeExecutor` +## 5. Proposal 2: Stateful `ContainerCodeExecutor` -### 4.1 Problem +### 5.1 Problem Agents often need multi-step code execution where later steps depend on earlier results. For example: @@ -306,9 +402,9 @@ Currently, each `execute_code()` call in `ContainerCodeExecutor` runs `python3 -c ` — a fresh Python process with no memory of prior calls. Step 2 would fail with `NameError: name 'df' is not defined`. -### 4.2 Design +### 5.2 Design -#### 4.2.1 Architecture +#### 5.2.1 Architecture ``` ┌─ ContainerCodeExecutor ─────────────────────┐ @@ -415,7 +511,7 @@ No new dependencies. Cons: Re-executes entire history on each call — side effects run again. Grows linearly with history length. -**Mitigation for side effects:** Wrap prior code in a guard: +**Mitigation for stdout leakage:** Wrap prior code in a guard: ```python # Only new code produces output; prior blocks set up state silently @@ -427,7 +523,16 @@ sys.stdout = _old_stdout # ... new code block (produces output) ... ``` -#### 4.2.2 Recommended Approach +**WARNING — Side-effect replay is NOT mitigated by stdout suppression.** +Prior blocks that perform file writes, network calls, database mutations, +or other I/O will re-execute those side effects on every subsequent call. +Stdout suppression only hides `print()` output — it does not prevent or +guard against non-idempotent operations. This is a fundamental limitation +of the cumulative replay approach (Option C). Users must keep +side-effecting code in the final block or use Option A (persistent +process) when side effects are unavoidable. + +#### 5.2.2 Recommended Approach **Option A (Persistent Process)** is the most robust for true statefulness and is the standard approach used by Jupyter kernels and similar systems. @@ -439,16 +544,26 @@ Despite its complexity, it provides the best user experience: - O(1) cost per call (not O(n) like Option C) However, implementing a full REPL protocol is a significant engineering -effort. We recommend a **phased approach**: +effort. + +**Given the severity of the side-effect replay problem (Finding 2), we +recommend evaluating whether to skip Phase 1 and go directly to +Option A.** The persistent-process approach eliminates an entire class +of bugs. If the I/O boundary protocol (sentinel markers for output +delimiting) can be solved with reasonable complexity, Phase 1 may not +be worth the technical debt. -**Phase 1 (MVP):** Option C (cumulative file) with stdout suppression for -prior blocks. Simple, works for the common case (data analysis, variable -setup). +**If Phase 1 is pursued as an MVP**, it should be clearly documented as +limited to **pure computation** (variable setup, data transforms, +aggregations) and explicitly unsupported for side-effecting code. -**Phase 2 (Full):** Option A (persistent process) with a proper -execution protocol using sentinel markers for output boundaries. +**Phase 1 (MVP, optional):** Option C (cumulative file) with stdout +suppression. Restricted to pure-computation use cases only. -#### 4.2.3 Implementation Plan (Phase 1) +**Phase 2 (Full, recommended):** Option A (persistent process) with a +proper execution protocol using sentinel markers for output boundaries. + +#### 5.2.3 Implementation Plan (Phase 1, if pursued) 1. **Unfreeze `stateful` in `ContainerCodeExecutor`:** @@ -465,12 +580,15 @@ _code_history: list[str] = [] 3. **Modify `execute_code()`:** +**Critical invariant:** Code is appended to history **only after** +successful execution. A failing code block must never be replayed. + ```python def execute_code(self, invocation_context, code_execution_input): code = code_execution_input.code if self.stateful: - # Build cumulative script + # Build cumulative script from prior SUCCESSFUL blocks setup_code = '\n'.join( f'# --- Block {i} ---\n{block}' for i, block in enumerate(self._code_history) @@ -483,7 +601,6 @@ def execute_code(self, invocation_context, code_execution_input): '_sys.stdout = _sys.__stdout__\n' f'{code}\n' ) - self._code_history.append(code) else: full_code = code @@ -491,7 +608,16 @@ def execute_code(self, invocation_context, code_execution_input): ['python3', '-c', full_code], demux=True, ) - # ... parse output as before + + # Parse output + stdout, stderr = self._parse_exec_output(exec_result) + success = (exec_result.exit_code == 0) + + # ONLY append to history after confirmed success + if self.stateful and success: + self._code_history.append(code) + + return CodeExecutionResult(stdout=stdout, stderr=stderr) ``` 4. **Add `reset_state()` method:** @@ -509,7 +635,7 @@ def reset_state(self): # Keep optimize_data_file frozen ``` -#### 4.2.4 Interaction with `execution_id` +#### 5.2.4 Interaction with `execution_id` The LLM flow layer uses `execution_id` (from `CodeExecutorContext`) to identify stateful sessions. For `ContainerCodeExecutor`: @@ -522,7 +648,25 @@ identify stateful sessions. For `ContainerCodeExecutor`: This aligns with how `VertexAiCodeExecutor` uses `session_id`. -### 4.3 Testing Plan +**Gap: `ExecuteSkillScriptTool` does not wire `execution_id`.** + +Currently, `ExecuteSkillScriptTool.run_async()` creates +`CodeExecutionInput(code=prepared_code)` without setting `execution_id`. +This means all skill script executions share the same (default) namespace +in a stateful executor, with no isolation between different skills or +invocations. + +**Action items:** +1. `ExecuteSkillScriptTool` should generate a deterministic + `execution_id` from the skill name + invocation context (e.g., + `f"skill:{skill_name}:{invocation_id}"`) +2. Pass `execution_id` to `CodeExecutionInput` +3. This enables future stateful skill scripts where a skill can + maintain state across multiple calls within the same session + +This is tracked as part of the Phase 2 implementation plan. + +### 5.3 Testing Plan | Test | Description | |------|-------------| @@ -536,9 +680,9 @@ This aligns with how `VertexAiCodeExecutor` uses `session_id`. --- -## 5. Proposal 3: Security Hardening +## 6. Proposal 3: Security Hardening -### 5.1 Problem +### 6.1 Problem `UnsafeLocalCodeExecutor` is the default executor for local development because it requires no external dependencies. But it runs `exec()` in the @@ -554,7 +698,7 @@ This is a critical security concern when executing: - Third-party skill scripts (supply chain risk) - User-provided code in multi-tenant deployments -### 5.2 Threat Model +### 6.2 Threat Model | Threat | Impact | Current mitigation | |--------|--------|--------------------| @@ -565,17 +709,29 @@ This is a critical security concern when executing: | Network exfiltration | Data leak | None | | File system manipulation | Data loss / corruption | None | -### 5.3 Design +### 6.3 Design We propose a layered approach with three tiers of security: -#### 5.3.1 Tier 1: `UnsafeLocalCodeExecutor` Hardening (Quick Wins) +#### 6.3.1 Tier 1: `UnsafeLocalCodeExecutor` Hardening (Quick Wins) These changes improve safety without changing the fundamental architecture: **A. Timeout support** (covered in Proposal 1) -**B. Restricted builtins:** +**B. Restricted builtins (best-effort friction, NOT a security boundary):** + +**Important caveat:** Builtin/module blocking in `exec()` is trivially +bypassed. Determined code can reach blocked functionality via: +- `object.__subclasses__()` → find `os._wrap_close` → access `os.system` +- `__builtins__.__dict__['__import__']('os')` (if `__builtins__` is a + module, not a dict) +- Encoding tricks, `importlib` via `sys.modules`, etc. + +This is explicitly **not a security control** — it is a speed bump that +catches accidental misuse and makes intentional abuse more visible. True +isolation requires `LocalSandboxCodeExecutor` (Tier 2) or containers +(Tier 3). ```python _BLOCKED_BUILTINS = { @@ -600,7 +756,10 @@ that need `open()` for file I/O). It should be opt-in: class UnsafeLocalCodeExecutor(BaseCodeExecutor): restrict_builtins: bool = False """When True, block dangerous builtins (exec, eval, open, - __import__). Default False for backward compatibility.""" + __import__). This is a best-effort friction layer, NOT a + security boundary — determined code can bypass it. Use + LocalSandboxCodeExecutor or ContainerCodeExecutor for + actual isolation. Default False for backward compatibility.""" ``` **C. Warning on first use:** @@ -622,7 +781,7 @@ def execute_code(self, ...): ... ``` -#### 5.3.2 Tier 2: `LocalSandboxCodeExecutor` (New, Recommended) +#### 6.3.2 Tier 2: `LocalSandboxCodeExecutor` (New, Recommended) A new executor that provides meaningful isolation without requiring Docker or cloud services: @@ -660,28 +819,34 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): if k in os.environ} env['PATH'] = '/usr/bin:/usr/local/bin' - # Set resource limits via preexec_fn - def set_limits(): - import resource - # CPU time limit - resource.setrlimit( - resource.RLIMIT_CPU, - (self.max_cpu_seconds, self.max_cpu_seconds) - ) - # Memory limit - mem_bytes = self.max_memory_mb * 1024 * 1024 - resource.setrlimit( - resource.RLIMIT_AS, - (mem_bytes, mem_bytes) - ) + timeout = ( + code_execution_input.timeout_seconds + or self.default_timeout_seconds + or self.max_cpu_seconds + ) + # Use process_group (Python 3.11+) instead of + # preexec_fn, which is not fork-safe with threads. + # process_group=0 places the child in its own + # process group, enabling clean group kill on + # timeout via os.killpg(). result = subprocess.run( - ['python3', f.name], + [ + 'python3', '-c', + f'import resource; ' + f'resource.setrlimit(resource.RLIMIT_CPU, ' + f'({self.max_cpu_seconds}, ' + f'{self.max_cpu_seconds})); ' + f'resource.setrlimit(resource.RLIMIT_AS, ' + f'({self.max_memory_mb * 1024 * 1024}, ' + f'{self.max_memory_mb * 1024 * 1024})); ' + f'exec(open({f.name!r}).read())', + ], capture_output=True, text=True, - timeout=self.timeout_seconds, + timeout=timeout, env=env, - preexec_fn=set_limits, # Unix only + process_group=0, # Python 3.11+, fork-safe cwd=tempfile.gettempdir(), ) @@ -694,45 +859,75 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): **Platform considerations:** - `resource.setrlimit` is Unix-only (Linux, macOS) +- `process_group=0` requires Python 3.11+ (the minimum for ADK) - On Windows, use `subprocess.CREATE_NO_WINDOW` and `subprocess.Popen` with `creationflags` for job object limits - Fallback to timeout-only on platforms without `resource` module +**Why `process_group` instead of `preexec_fn`:** +- `preexec_fn` is not fork-safe with threads — the Python docs warn + that it can deadlock in multi-threaded programs because it runs + between `fork()` and `exec()` while all parent thread locks are + held. ADK executors may be called from async/threaded contexts. +- `process_group=0` (Python 3.11+) is fork-safe and places the child + in its own process group, enabling clean `os.killpg()` on timeout. +- Resource limits are set via an inline `-c` wrapper script instead + of `preexec_fn`, avoiding the fork-safety issue entirely. + **Dependencies:** None (stdlib only). This is the key advantage over `ContainerCodeExecutor`. **Limitations:** - Less isolation than containers (shared filesystem, kernel) -- `preexec_fn` is not fork-safe with threads (use `process_group` on - Python 3.11+) - Cannot restrict network access without OS-level firewall rules +- Requires Python 3.11+ (already the ADK minimum) -#### 5.3.3 Tier 3: Promote `ContainerCodeExecutor` as Default +#### 6.3.3 Tier 3: Promote `ContainerCodeExecutor` as Default For production deployments, container-based isolation should be the standard recommendation: -**A. Simplify setup:** +**A. Simplify setup with digest-pinned default image:** ```python # Current: requires explicit image or docker_path executor = ContainerCodeExecutor(image='python:3.11-slim') -# Proposed: auto-pull default image -executor = ContainerCodeExecutor() # Uses python:3.11-slim +# Proposed: auto-pull default image (digest-pinned) +executor = ContainerCodeExecutor() +``` + +**Default image should use a digest-pinned or versioned tag**, not a +mutable tag like `python:3.11-slim`. Mutable tags can change content +silently (e.g., security patches, Python micro-version bumps), leading +to non-reproducible behavior across environments and over time. + +```python +# In ContainerCodeExecutor defaults: +_DEFAULT_IMAGE = ( + 'python:3.11.11-slim@sha256:' +) +# Updated in ADK releases with tested digests +``` + +When an official `adk-code-executor` image is published, the default +should reference a versioned tag matching the ADK release: +```python +_DEFAULT_IMAGE = 'gcr.io/adk/code-executor:0.5.0' ``` **B. Pre-built ADK executor image:** Create an official `adk-code-executor` Docker image with: -- Python 3.11+ slim base +- Python 3.11+ slim base (digest-pinned in Dockerfile) - Common data science libraries (pandas, numpy, matplotlib) - Non-root user - Read-only filesystem (writable `/tmp` only) - No network access by default (`--network=none`) +- Versioned tags matching ADK releases ```dockerfile -FROM python:3.11-slim +FROM python:3.11.11-slim@sha256: RUN pip install --no-cache-dir pandas numpy matplotlib RUN useradd -m -s /bin/bash executor USER executor @@ -756,7 +951,7 @@ def __init_container(self): ) ``` -### 5.4 Recommendation Matrix +### 6.4 Recommendation Matrix | Use Case | Recommended Executor | Why | |----------|---------------------|-----| @@ -767,7 +962,7 @@ def __init_container(self): | Production (multi-tenant) | `GkeCodeExecutor` | gVisor, per-execution isolation | | Google Cloud | `AgentEngineSandboxCodeExecutor` | Managed, scalable | -### 5.5 Implementation Plan +### 6.5 Implementation Plan | Phase | Action | Effort | Risk | |-------|--------|--------|------| @@ -780,9 +975,9 @@ def __init_container(self): --- -## 6. Cross-Cutting Concerns +## 7. Cross-Cutting Concerns -### 6.1 `BaseCodeExecutor` API Changes +### 7.1 `BaseCodeExecutor` API Changes All three proposals touch `BaseCodeExecutor`. The combined changes: @@ -796,8 +991,9 @@ class BaseCodeExecutor(BaseModel): execution_result_delimiters: tuple[str, str] = (...) # NEW: Proposal 1 - timeout_seconds: Optional[int] = None - """Maximum execution time in seconds. None = no timeout.""" + default_timeout_seconds: Optional[int] = None + """Default timeout applied when CodeExecutionInput.timeout_seconds + is None. Subclasses may override. None = no timeout.""" @abc.abstractmethod def execute_code( @@ -805,20 +1001,30 @@ class BaseCodeExecutor(BaseModel): invocation_context: InvocationContext, code_execution_input: CodeExecutionInput, ) -> CodeExecutionResult: ... + + +@dataclasses.dataclass +class CodeExecutionInput: + code: str + input_files: list[File] = field(default_factory=list) + execution_id: Optional[str] = None + timeout_seconds: Optional[int] = None # NEW: per-invocation + """Per-invocation timeout. Overrides executor default when set.""" ``` -### 6.2 Backward Compatibility +### 7.2 Backward Compatibility | Change | Backward compatible? | Migration needed? | |--------|---------------------|------------------| -| `timeout_seconds` on base class | Yes (default `None`) | No | +| `default_timeout_seconds` on base class | Yes (default `None`) | No | +| `timeout_seconds` on `CodeExecutionInput` | Yes (default `None`) | No | | Unfreeze `stateful` on `ContainerCodeExecutor` | Yes (default `False`) | No | | `SecurityWarning` on `UnsafeLocalCodeExecutor` | Yes (warning only) | No | | New `LocalSandboxCodeExecutor` | Yes (additive) | No | | `restrict_builtins` on `UnsafeLocalCodeExecutor` | Yes (default `False`) | No | | Default image for `ContainerCodeExecutor` | Breaking (currently requires image/docker_path) | Minor | -### 6.3 Impact on `ExecuteSkillScriptTool` +### 7.3 Impact on `ExecuteSkillScriptTool` | Feature | Current workaround | After enhancements | |---------|-------------------|-------------------| @@ -827,7 +1033,7 @@ class BaseCodeExecutor(BaseModel): | Isolation | Documentation warning only | `LocalSandboxCodeExecutor` or container | | Stateful scripts | Not supported | Available via `ContainerCodeExecutor(stateful=True)` | -### 6.4 Testing Strategy +### 7.4 Testing Strategy | Category | Approach | |----------|----------| @@ -839,48 +1045,57 @@ class BaseCodeExecutor(BaseModel): --- -## 7. Implementation Roadmap +## 8. Implementation Roadmap -### Phase 1: Timeout (2-3 days) +### Phase 1: Timeout (3-4 days) -1. Add `timeout_seconds: Optional[int] = None` to `BaseCodeExecutor` -2. Implement thread-based timeout in `UnsafeLocalCodeExecutor` -3. Implement thread-based timeout in `ContainerCodeExecutor` -4. Migrate `GkeCodeExecutor.timeout_seconds` to base class field -5. Add timeout tests for each executor -6. Update `ExecuteSkillScriptTool` to set executor timeout when available +1. Add `timeout_seconds: Optional[int] = None` to `CodeExecutionInput` +2. Add `default_timeout_seconds: Optional[int] = None` to + `BaseCodeExecutor` +3. Implement thread-based timeout in `UnsafeLocalCodeExecutor` +4. Implement Docker exec kill timeout in `ContainerCodeExecutor` +5. Migrate `GkeCodeExecutor.timeout_seconds` to `default_timeout_seconds` +6. Add timeout tests for each executor +7. Update `ExecuteSkillScriptTool` to set per-invocation timeout via + `CodeExecutionInput.timeout_seconds` ### Phase 2: Stateful Container (3-5 days) 1. Unfreeze `stateful` on `ContainerCodeExecutor` 2. Implement cumulative code history with stdout suppression + (append-after-success invariant) 3. Add `execution_id`-based history isolation -4. Add `reset_state()` method -5. Add stateful execution tests -6. Update samples and documentation +4. Wire `execution_id` in `ExecuteSkillScriptTool` +5. Add `reset_state()` method +6. Add stateful execution tests (including failure-does-not-poison test) +7. Update samples and documentation +8. Evaluate persistent-process approach (Option A) for Phase 2b ### Phase 3: Security Hardening (5-7 days) 1. Add `SecurityWarning` to `UnsafeLocalCodeExecutor` -2. Add `restrict_builtins` option -3. Implement `LocalSandboxCodeExecutor` -4. Add default image support to `ContainerCodeExecutor` +2. Add `restrict_builtins` option (documented as best-effort friction) +3. Implement `LocalSandboxCodeExecutor` (using `process_group`, not + `preexec_fn`) +4. Add digest-pinned default image to `ContainerCodeExecutor` 5. Add network isolation defaults to `ContainerCodeExecutor` -6. Create official `adk-code-executor` Docker image +6. Create official `adk-code-executor` Docker image (versioned tags) 7. Update all samples to recommend secure executors 8. Add security-focused tests -### Total estimated effort: 10-15 days +### Total estimated effort: 11-16 days --- -## 8. Open Questions +## 9. Open Questions -1. **Should `timeout_seconds` be enforced at the base class level?** - We could add a wrapper in `BaseCodeExecutor.execute_code()` that - enforces the timeout generically, rather than requiring each subclass - to implement it. However, this would require the base class to manage - threading, which may not be appropriate for remote executors. +1. **Should we skip Phase 1 (cumulative replay) and go straight to + Phase 2 (persistent process) for stateful execution?** + The side-effect replay problem is fundamental to Option C. If the + persistent-process I/O boundary protocol can be solved with + reasonable complexity (e.g., sentinel-delimited output), the MVP + phase may not be worth the tech debt. Decision: Evaluate during + Phase 2 planning. 2. **Should `LocalSandboxCodeExecutor` support stateful execution?** Subprocess-based execution is inherently stateless. Stateful support @@ -901,7 +1116,7 @@ class BaseCodeExecutor(BaseModel): --- -## 9. References +## 10. References - [BaseCodeExecutor](../../src/google/adk/code_executors/base_code_executor.py) - [ContainerCodeExecutor](../../src/google/adk/code_executors/container_code_executor.py) From 8ca1111ebec8625ce508206df0f52ef786ac2472 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 21 Feb 2026 19:11:40 -0800 Subject: [PATCH 06/53] =?UTF-8?q?docs:=20Fix=206=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20execution=5Fid,=20PID=20namespace,=20Py=20version,?= =?UTF-8?q?=20timeout=20API?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix execution_id to use session-stable key (session.id), not per-turn invocation_id - Fix container timeout kill: document PID namespace mismatch, use pkill -f inside container instead of host-PID kill - Fix Python version: ADK minimum is >=3.10 (pyproject.toml), not 3.11; gate process_group with version check, preexec_fn fallback - Fix LocalSandboxCodeExecutor to use default_timeout_seconds (consistent with BaseCodeExecutor API) - Replace all `or` and `??` timeout resolution with explicit `is not None` checks to allow timeout=0 - Fix test strategy to document missing ContainerCodeExecutor tests Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 198 ++++++++++++++-------- 1 file changed, 128 insertions(+), 70 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index 4a2aa4f7b1..33e6b54bdc 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -191,9 +191,10 @@ class BaseCodeExecutor(BaseModel): The effective timeout is resolved as: ```python +input_t = code_execution_input.timeout_seconds timeout = ( - code_execution_input.timeout_seconds - ?? self.default_timeout_seconds + input_t if input_t is not None + else self.default_timeout_seconds ) ``` @@ -213,9 +214,10 @@ run it in a separate thread with a join timeout: import threading def execute_code(self, invocation_context, code_execution_input): + input_t = code_execution_input.timeout_seconds timeout = ( - code_execution_input.timeout_seconds - or self.default_timeout_seconds + input_t if input_t is not None + else self.default_timeout_seconds ) if timeout is None: # No timeout: current behavior (blocking exec) @@ -268,9 +270,10 @@ running indefinitely after the join timeout expires. import threading def execute_code(self, invocation_context, code_execution_input): + input_t = code_execution_input.timeout_seconds timeout = ( - code_execution_input.timeout_seconds - or self.default_timeout_seconds + input_t if input_t is not None + else self.default_timeout_seconds ) # Create the exec instance @@ -279,28 +282,32 @@ def execute_code(self, invocation_context, code_execution_input): ['python3', '-c', code_execution_input.code], )['Id'] - # Start a timer to kill the exec's PID on timeout + # Start a timer to kill the exec on timeout. + # + # NOTE on PID namespaces: exec_inspect().Pid returns the + # host-namespace PID of the exec'd process. Passing that PID + # to `kill` inside the container will not match (different PID + # namespace). Instead, we use the Docker API to kill from the + # host side, or find the container-namespace PID. timer = None timed_out = threading.Event() if timeout is not None: def _kill_exec(): timed_out.set() try: - # Get the PID of the exec'd process - info = self._client.api.exec_inspect(exec_id) - pid = info.get('Pid', 0) - if pid > 0: - self._container.exec_run( - ['kill', '-9', str(pid)], - detach=True, - ) - except Exception: - # Fallback: kill all non-init processes + # Approach: find exec'd process by its command + # inside the container's PID namespace, then kill. + # `exec_inspect` gives us the command; `pkill -f` + # matches it inside the container. self._container.exec_run( - ['sh', '-c', - 'kill -9 $(ps -o pid= | grep -v "^\\s*1$")'], + ['pkill', '-9', '-f', 'python3 -c'], detach=True, ) + except Exception: + pass + # If the above fails or is insufficient, the next + # execute_code() call will detect the zombie and + # restart the container. timer = threading.Timer(timeout, _kill_exec) timer.start() @@ -317,21 +324,34 @@ def execute_code(self, invocation_context, code_execution_input): # ... parse output as before ``` -**Why targeted PID kill, not `kill -9 -1`:** -- `kill -9 -1` kills *all* processes in the container, including the - init/shell process that keeps the container alive. This would force a - container restart on the next call. -- Targeted kill via `exec_inspect` → PID only terminates the timed-out - process, leaving the container healthy for subsequent calls. -- The fallback (`kill all non-init`) is a safety net if `exec_inspect` - fails (e.g., Docker API version mismatch). - -**Alternative — container restart on timeout:** Simpler but more costly. -Stop and restart the container after timeout. Acceptable if stateful mode -is not in use (no accumulated state to preserve). - -**Recommendation:** Use Docker exec kill as the primary approach. This is -more robust than thread+join and properly cleans up runaway processes. +**PID namespace considerations:** +- `exec_inspect(exec_id)['Pid']` returns the **host-namespace** PID. + Running `kill ` inside the container operates in the + **container PID namespace** and will target a different (or + non-existent) process. This is a common Docker pitfall. +- The correct approaches are: + 1. **`pkill -f` inside container** — Match the exec'd command string + within the container's PID namespace. Works for `python3 -c` but + may be too broad if multiple execs are running concurrently. + 2. **Host-side kill via `docker top` + `os.kill`** — Use + `container.top()` to map container PIDs to host PIDs, then + `os.kill(host_pid, 9)` from the host. More precise but requires + host-level permissions. + 3. **Container restart** — Simplest and most reliable. Acceptable + when stateful mode is not in use. + +**Recovery after timeout kill:** +- **Stateless mode:** No recovery needed. Next `execute_code()` call + starts a fresh process in the same container. +- **Stateful mode (cumulative replay):** The timed-out block is NOT + appended to history (append-after-success invariant), so replay + remains clean. However, if `pkill` killed a persistent interpreter + (Phase 2 / Option A), the executor must detect this and restart + it before the next call. + +**Recommendation:** Use `pkill -f` as the primary approach for Phase 1. +Migrate to host-side kill or container restart for more robust cleanup +in Phase 2 when persistent processes are introduced. #### 4.2.4 `GkeCodeExecutor` — Already Implemented @@ -345,9 +365,10 @@ class GkeCodeExecutor(BaseCodeExecutor): In `execute_code()`, resolve timeout from per-invocation input first: ```python +input_t = code_execution_input.timeout_seconds timeout = ( - code_execution_input.timeout_seconds - or self.default_timeout_seconds + input_t if input_t is not None + else self.default_timeout_seconds ) ``` @@ -657,9 +678,15 @@ in a stateful executor, with no isolation between different skills or invocations. **Action items:** -1. `ExecuteSkillScriptTool` should generate a deterministic - `execution_id` from the skill name + invocation context (e.g., - `f"skill:{skill_name}:{invocation_id}"`) +1. `ExecuteSkillScriptTool` should generate a **session-stable** + `execution_id` scoped to skill + agent. The key must persist + across turns so that stateful code history is preserved: + ```python + execution_id = f"skill:{skill_name}:{session.id}:{agent_name}" + ``` + Using `invocation_id` would be incorrect here — it changes every + turn, defeating statefulness. `session.id` is stable for the + lifetime of the conversation. 2. Pass `execution_id` to `CodeExecutionInput` 3. This enables future stateful skill scripts where a skill can maintain state across multiple calls within the same session @@ -799,7 +826,7 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): - Optional chroot or tmpdir working directory """ - timeout_seconds: int = 30 + default_timeout_seconds: int = 30 max_memory_mb: int = 256 max_cpu_seconds: int = 30 allowed_env_vars: list[str] = [] @@ -819,35 +846,53 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): if k in os.environ} env['PATH'] = '/usr/bin:/usr/local/bin' + input_t = code_execution_input.timeout_seconds timeout = ( - code_execution_input.timeout_seconds - or self.default_timeout_seconds - or self.max_cpu_seconds + input_t if input_t is not None + else self.default_timeout_seconds ) - - # Use process_group (Python 3.11+) instead of - # preexec_fn, which is not fork-safe with threads. - # process_group=0 places the child in its own - # process group, enabling clean group kill on - # timeout via os.killpg(). + if timeout is None: + timeout = self.max_cpu_seconds + + import sys + # Prefer process_group (3.11+) over preexec_fn + # (not fork-safe with threads). + spawn_kwargs = {} + if sys.version_info >= (3, 11): + spawn_kwargs['process_group'] = 0 + else: + # Fallback for 3.10; caveat: not fork-safe + def _set_limits(): + import resource + resource.setrlimit( + resource.RLIMIT_CPU, + (self.max_cpu_seconds,) * 2, + ) + mem = self.max_memory_mb * 1024 * 1024 + resource.setrlimit( + resource.RLIMIT_AS, (mem, mem), + ) + spawn_kwargs['preexec_fn'] = _set_limits + + cmd = [ + 'python3', '-c', + f'import resource; ' + f'resource.setrlimit(resource.RLIMIT_CPU, ' + f'({self.max_cpu_seconds}, ' + f'{self.max_cpu_seconds})); ' + f'resource.setrlimit(resource.RLIMIT_AS, ' + f'({self.max_memory_mb * 1024 * 1024}, ' + f'{self.max_memory_mb * 1024 * 1024})); ' + f'exec(open({f.name!r}).read())', + ] result = subprocess.run( - [ - 'python3', '-c', - f'import resource; ' - f'resource.setrlimit(resource.RLIMIT_CPU, ' - f'({self.max_cpu_seconds}, ' - f'{self.max_cpu_seconds})); ' - f'resource.setrlimit(resource.RLIMIT_AS, ' - f'({self.max_memory_mb * 1024 * 1024}, ' - f'{self.max_memory_mb * 1024 * 1024})); ' - f'exec(open({f.name!r}).read())', - ], + cmd, capture_output=True, text=True, timeout=timeout, env=env, - process_group=0, # Python 3.11+, fork-safe cwd=tempfile.gettempdir(), + **spawn_kwargs, ) return CodeExecutionResult( @@ -859,7 +904,9 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): **Platform considerations:** - `resource.setrlimit` is Unix-only (Linux, macOS) -- `process_group=0` requires Python 3.11+ (the minimum for ADK) +- `process_group=0` requires Python 3.11+ (ADK supports >=3.10, so + this must be gated with a version check or use `preexec_fn` as + fallback on 3.10) - On Windows, use `subprocess.CREATE_NO_WINDOW` and `subprocess.Popen` with `creationflags` for job object limits - Fallback to timeout-only on platforms without `resource` module @@ -873,6 +920,8 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): in its own process group, enabling clean `os.killpg()` on timeout. - Resource limits are set via an inline `-c` wrapper script instead of `preexec_fn`, avoiding the fork-safety issue entirely. +- On Python 3.10 (ADK minimum is `>=3.10`), fall back to + `preexec_fn=set_limits` with a documented caveat about thread safety. **Dependencies:** None (stdlib only). This is the key advantage over `ContainerCodeExecutor`. @@ -880,7 +929,8 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): **Limitations:** - Less isolation than containers (shared filesystem, kernel) - Cannot restrict network access without OS-level firewall rules -- Requires Python 3.11+ (already the ADK minimum) +- `process_group` requires Python 3.11+; falls back to `preexec_fn` + on 3.10 (ADK minimum is `>=3.10` per `pyproject.toml`) #### 6.3.3 Tier 3: Promote `ContainerCodeExecutor` as Default @@ -1035,13 +1085,21 @@ class CodeExecutionInput: ### 7.4 Testing Strategy -| Category | Approach | -|----------|----------| -| Unit tests | Mock-based tests for each executor (existing pattern) | -| Integration tests | Real executor tests (like the ones added for `ExecuteSkillScriptTool`) | -| Timeout tests | Scripts with `time.sleep()` to verify timeout enforcement | -| Security tests | Scripts attempting blocked operations to verify restrictions | -| Stateful tests | Multi-call sequences verifying variable persistence | +**Current test coverage gaps:** Unit tests exist for +`UnsafeLocalCodeExecutor`, `GkeCodeExecutor`, +`AgentEngineSandboxCodeExecutor`, `BuiltInCodeExecutor`, and +`CodeExecutorContext`, but **no unit test file exists for +`ContainerCodeExecutor`** (`tests/unittests/code_executors/` has no +`test_container_code_executor.py`). Likewise, `LocalSandboxCodeExecutor` +is new and has no tests yet. + +| Category | Approach | New tests needed | +|----------|----------|-----------------| +| Unit tests | Mock-based tests per executor | **Add `test_container_code_executor.py`**, add `test_local_sandbox_code_executor.py` | +| Integration tests | Real executor tests (like `ExecuteSkillScriptTool` integration tests) | Add Docker-based container tests (CI-gated) | +| Timeout tests | Scripts with `time.sleep()` to verify enforcement | Per-executor timeout tests | +| Security tests | Scripts attempting blocked operations | `restrict_builtins` bypass attempts, env var leakage | +| Stateful tests | Multi-call sequences verifying variable persistence | Append-after-success, failure-does-not-poison, `execution_id` isolation | --- From d8692ba6f189f6f335e25474a3d32ace360e3cf4 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 21 Feb 2026 19:27:05 -0800 Subject: [PATCH 07/53] docs: Fix container timeout DoS, pkill scope, stale recommendations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Redesign container timeout: run exec_start in thread + join(timeout), kill via host-side os.kill(host_pid) instead of in-container pkill. Eliminates blocking exec_start DoS and overbroad pkill -f pattern. - Remove contradictory "RECOMMENDED" label from Option C; Option A (persistent process) is the single recommended approach. - Remove Windows implementation guidance (contradicts Non-Goals §2.4); raise NotImplementedError instead. - Guard resource import in LocalSandboxCodeExecutor wrapper script for platforms where resource module is unavailable. - Fix remote executor timeout to use synchronous thread+join, not asyncio.wait_for (execute_code APIs are synchronous). Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 220 ++++++++++++---------- 1 file changed, 122 insertions(+), 98 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index 33e6b54bdc..c303eaf088 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -266,7 +266,15 @@ running indefinitely after the join timeout expires. **Primary approach: Docker exec kill via `exec_inspect` + PID kill.** +The key constraint is that `exec_start` is a **blocking** call — a +timer thread cannot unblock it if the in-container kill fails. The +correct design runs `exec_start` in a worker thread so the caller +can enforce the timeout via `thread.join(timeout)`, then uses the +**host-side Docker API** to kill the exec'd process by its host PID. + ```python +import os +import signal import threading def execute_code(self, invocation_context, code_execution_input): @@ -282,76 +290,79 @@ def execute_code(self, invocation_context, code_execution_input): ['python3', '-c', code_execution_input.code], )['Id'] - # Start a timer to kill the exec on timeout. - # - # NOTE on PID namespaces: exec_inspect().Pid returns the - # host-namespace PID of the exec'd process. Passing that PID - # to `kill` inside the container will not match (different PID - # namespace). Instead, we use the Docker API to kill from the - # host side, or find the container-namespace PID. - timer = None - timed_out = threading.Event() - if timeout is not None: - def _kill_exec(): - timed_out.set() + # Run exec_start in a thread so we can enforce timeout + # from the calling thread. + result_holder = {} + + def _run_exec(): + result_holder['output'] = ( + self._client.api.exec_start(exec_id, demux=True) + ) + + thread = threading.Thread(target=_run_exec, daemon=True) + thread.start() + thread.join(timeout=timeout) + + if thread.is_alive(): + # Timeout exceeded — kill from host side. + # exec_inspect returns the host-namespace PID, which + # is the correct PID for os.kill() on the host. + # (Killing from inside the container would require + # the container-namespace PID, which is different.) + try: + info = self._client.api.exec_inspect(exec_id) + host_pid = info.get('Pid', 0) + if host_pid > 0: + os.kill(host_pid, signal.SIGKILL) + except (ProcessLookupError, PermissionError): + pass # Process already exited + except Exception: + # Last resort: restart the container try: - # Approach: find exec'd process by its command - # inside the container's PID namespace, then kill. - # `exec_inspect` gives us the command; `pkill -f` - # matches it inside the container. - self._container.exec_run( - ['pkill', '-9', '-f', 'python3 -c'], - detach=True, - ) + self._container.restart(timeout=1) except Exception: pass - # If the above fails or is insufficient, the next - # execute_code() call will detect the zombie and - # restart the container. - timer = threading.Timer(timeout, _kill_exec) - timer.start() - - try: - output = self._client.api.exec_start(exec_id, demux=True) - finally: - if timer is not None: - timer.cancel() - - if timed_out.is_set(): + return CodeExecutionResult( stderr=f'Execution timed out after {timeout}s' ) + + output = result_holder.get('output') # ... parse output as before ``` -**PID namespace considerations:** -- `exec_inspect(exec_id)['Pid']` returns the **host-namespace** PID. - Running `kill ` inside the container operates in the - **container PID namespace** and will target a different (or - non-existent) process. This is a common Docker pitfall. -- The correct approaches are: - 1. **`pkill -f` inside container** — Match the exec'd command string - within the container's PID namespace. Works for `python3 -c` but - may be too broad if multiple execs are running concurrently. - 2. **Host-side kill via `docker top` + `os.kill`** — Use - `container.top()` to map container PIDs to host PIDs, then - `os.kill(host_pid, 9)` from the host. More precise but requires - host-level permissions. - 3. **Container restart** — Simplest and most reliable. Acceptable - when stateful mode is not in use. - -**Recovery after timeout kill:** -- **Stateless mode:** No recovery needed. Next `execute_code()` call - starts a fresh process in the same container. -- **Stateful mode (cumulative replay):** The timed-out block is NOT - appended to history (append-after-success invariant), so replay - remains clean. However, if `pkill` killed a persistent interpreter - (Phase 2 / Option A), the executor must detect this and restart - it before the next call. - -**Recommendation:** Use `pkill -f` as the primary approach for Phase 1. -Migrate to host-side kill or container restart for more robust cleanup -in Phase 2 when persistent processes are introduced. +**Why this design:** + +1. **`exec_start` in a thread + `join(timeout)`** — The caller is + never blocked longer than `timeout` seconds, regardless of whether + the kill succeeds. This resolves the primary DoS risk. + +2. **Host-side `os.kill(host_pid, SIGKILL)`** — `exec_inspect()` + returns the **host-namespace** PID. Using `os.kill()` from the + host process operates in the host PID namespace, so the PID + matches correctly. This avoids: + - The PID namespace mismatch of killing from inside the container + - The overbroad pattern matching of `pkill -f` (which can hit + concurrent execs or unrelated Python processes) + - Dependency on `procps`/`pkill` being installed in the image + +3. **Container restart as last resort** — If `os.kill` fails (e.g., + insufficient permissions when Docker runs rootless), restart the + container. This is the most reliable fallback but destroys + in-container state. + +**Permissions:** `os.kill()` on the host PID requires the ADK process +to run as the same user that started the container (or as root). This +is the normal case for local Docker usage. For rootless Docker or +restricted environments, the container-restart fallback applies. + +**Recovery after timeout:** +- **Stateless mode:** No recovery needed. The killed process is gone; + the next `exec_run` starts a fresh process in the same container. +- **Stateful mode:** The timed-out block is NOT appended to history + (append-after-success invariant). If the container was restarted + as fallback, the executor must detect this and replay the + accumulated history on the next call. #### 4.2.4 `GkeCodeExecutor` — Already Implemented @@ -378,11 +389,19 @@ timeout, so the 300s default applies as before). #### 4.2.5 Remote Executors (Vertex AI, Agent Engine) These executors delegate to Google Cloud APIs that have their own internal -timeouts. Adding client-side timeout is still valuable as a safety net: +timeouts. Adding client-side timeout is still valuable as a safety net. + +Note: The current `execute_code()` implementations in +`VertexAiCodeExecutor` and `AgentEngineSandboxCodeExecutor` are +**synchronous** (they call blocking SDK methods), so `asyncio.wait_for()` +is not applicable. Use the same thread+join pattern as +`UnsafeLocalCodeExecutor`: -- Wrap the API call in `asyncio.wait_for()` or `threading.Timer` +- Run the blocking API call in a daemon thread with `join(timeout)` - Return `CodeExecutionResult(stderr='...')` on timeout - Log a warning that the server-side execution may still be running +- If these executors are migrated to async in the future, switch to + `asyncio.wait_for()` at that point ### 4.3 Migration Plan @@ -497,7 +516,7 @@ Cons: Not all Python objects are serializable (e.g., open file handles, database connections, generators). Adds `dill` dependency to the container image. Serialization/deserialization overhead grows with state size. -**Option C: Shared Globals File (Simple) — RECOMMENDED** +**Option C: Shared Globals File (Simple)** Write executed code to a cumulative Python file. Each call appends the new code block and re-executes the entire history: @@ -553,36 +572,29 @@ of the cumulative replay approach (Option C). Users must keep side-effecting code in the final block or use Option A (persistent process) when side effects are unavoidable. -#### 5.2.2 Recommended Approach +#### 5.2.2 Recommended Approach: Option A (Persistent Process) -**Option A (Persistent Process)** is the most robust for true statefulness -and is the standard approach used by Jupyter kernels and similar systems. -Despite its complexity, it provides the best user experience: +**Option A is the recommended approach.** It is the most robust for +true statefulness and is the standard approach used by Jupyter kernels +and similar systems: - Variables, imports, and objects persist naturally - No re-execution of side effects - No serialization issues - O(1) cost per call (not O(n) like Option C) -However, implementing a full REPL protocol is a significant engineering -effort. - -**Given the severity of the side-effect replay problem (Finding 2), we -recommend evaluating whether to skip Phase 1 and go directly to -Option A.** The persistent-process approach eliminates an entire class -of bugs. If the I/O boundary protocol (sentinel markers for output -delimiting) can be solved with reasonable complexity, Phase 1 may not -be worth the technical debt. - -**If Phase 1 is pursued as an MVP**, it should be clearly documented as -limited to **pure computation** (variable setup, data transforms, -aggregations) and explicitly unsupported for side-effecting code. - -**Phase 1 (MVP, optional):** Option C (cumulative file) with stdout -suppression. Restricted to pure-computation use cases only. +Option C (cumulative replay) has a fundamental side-effect replay +problem that cannot be fully mitigated. We recommend **going directly +to Option A** rather than shipping Option C as an interim MVP that +would accumulate technical debt and user-facing bugs. -**Phase 2 (Full, recommended):** Option A (persistent process) with a -proper execution protocol using sentinel markers for output boundaries. +If a simpler interim is needed before the persistent-process protocol +is ready, Option C may be used with the following restrictions: +- Documented as limited to **pure computation only** (variable setup, + data transforms, aggregations) +- Side-effecting code (file writes, network calls, DB mutations) is + explicitly unsupported and will produce incorrect results +- Clearly labeled as experimental / unstable #### 5.2.3 Implementation Plan (Phase 1, if pursued) @@ -874,16 +886,24 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): ) spawn_kwargs['preexec_fn'] = _set_limits - cmd = [ - 'python3', '-c', - f'import resource; ' - f'resource.setrlimit(resource.RLIMIT_CPU, ' + # Guard resource import for platforms where + # it is unavailable (falls back to timeout-only). + limit_code = ( + f'try:\n' + f' import resource\n' + f' resource.setrlimit(resource.RLIMIT_CPU, ' f'({self.max_cpu_seconds}, ' - f'{self.max_cpu_seconds})); ' - f'resource.setrlimit(resource.RLIMIT_AS, ' + f'{self.max_cpu_seconds}))\n' + f' resource.setrlimit(resource.RLIMIT_AS, ' f'({self.max_memory_mb * 1024 * 1024}, ' - f'{self.max_memory_mb * 1024 * 1024})); ' - f'exec(open({f.name!r}).read())', + f'{self.max_memory_mb * 1024 * 1024}))\n' + f'except (ImportError, OSError):\n' + f' pass\n' + ) + cmd = [ + 'python3', '-c', + limit_code + + f'exec(open({f.name!r}).read())', ] result = subprocess.run( cmd, @@ -903,13 +923,17 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): ``` **Platform considerations:** -- `resource.setrlimit` is Unix-only (Linux, macOS) +- `resource.setrlimit` is Unix-only (Linux, macOS). On platforms + where `resource` is unavailable, the inline `-c` wrapper must + skip the `setrlimit` calls and rely on timeout-only enforcement. + The code should guard with `try: import resource; ... except + ImportError: pass` in the wrapper script. - `process_group=0` requires Python 3.11+ (ADK supports >=3.10, so this must be gated with a version check or use `preexec_fn` as fallback on 3.10) -- On Windows, use `subprocess.CREATE_NO_WINDOW` and - `subprocess.Popen` with `creationflags` for job object limits -- Fallback to timeout-only on platforms without `resource` module +- Windows is out of scope (see §2 Non-Goals). The executor should + raise `NotImplementedError` on Windows with a message directing + users to `ContainerCodeExecutor`. **Why `process_group` instead of `preexec_fn`:** - `preexec_fn` is not fork-safe with threads — the Python docs warn From f55da5326190e44c2312cd1e988a963bb092f625 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 21 Feb 2026 19:34:01 -0800 Subject: [PATCH 08/53] docs: Align roadmap with Option A, unify recovery policy, fix fallback gaps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Roadmap Phase 2 now targets Option A (persistent process) directly instead of cumulative replay (Option C) first - Unify recovery policy: timeout and crash both return error indicating state loss, no automatic replay (consistent across §4.2.3 and §9) - Add Windows NotImplementedError gate to LocalSandboxCodeExecutor - Guard resource import in 3.10 preexec_fn fallback path - Document that PermissionError → container.restart() is the primary timeout path in practice (container runs as root, ADK as non-root) - Add explicit test cases for PermissionError fallback and kill paths - Update open question #1 to focus on I/O boundary protocol design Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 122 ++++++++++++++-------- 1 file changed, 79 insertions(+), 43 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index c303eaf088..d30ca7c180 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -351,18 +351,31 @@ def execute_code(self, invocation_context, code_execution_input): container. This is the most reliable fallback but destroys in-container state. -**Permissions:** `os.kill()` on the host PID requires the ADK process -to run as the same user that started the container (or as root). This -is the normal case for local Docker usage. For rootless Docker or -restricted environments, the container-restart fallback applies. +**Permissions and user mismatch risk:** `os.kill()` on the host PID +requires the ADK process to run as the same user that owns the +container's exec'd process on the host. The current +`ContainerCodeExecutor` does not set `user=` on `containers.run()` +(`container_code_executor.py:182`), so the container process runs as +root. If the ADK process runs as a non-root user (common in +development), `os.kill(host_pid, SIGKILL)` will raise +`PermissionError` and the container-restart fallback activates. + +This user mismatch is the expected case for default Docker usage, +so the `PermissionError → container.restart()` path is the **primary +timeout mechanism in practice**, not an edge case. The `os.kill` path +becomes primary only when ADK runs as root or the container is +configured with `user=` matching the host user. **Recovery after timeout:** - **Stateless mode:** No recovery needed. The killed process is gone; the next `exec_run` starts a fresh process in the same container. -- **Stateful mode:** The timed-out block is NOT appended to history - (append-after-success invariant). If the container was restarted - as fallback, the executor must detect this and replay the - accumulated history on the next call. +- **Stateful mode (Option A / persistent process):** The timed-out + block is NOT appended to history (append-after-success invariant). + If the persistent REPL was killed or the container was restarted, + the executor returns an error with `stderr` indicating state was + lost. The caller (LLM flow or skill tool) must handle this — + typically by starting a new session. Automatic replay is not + attempted because it may re-execute side effects. #### 4.2.4 `GkeCodeExecutor` — Already Implemented @@ -844,9 +857,18 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): allowed_env_vars: list[str] = [] def execute_code(self, invocation_context, code_execution_input): + import platform import subprocess + import sys import tempfile + # Windows is out of scope (§2 Non-Goals). + if platform.system() == 'Windows': + raise NotImplementedError( + 'LocalSandboxCodeExecutor is not supported on ' + 'Windows. Use ContainerCodeExecutor instead.' + ) + with tempfile.NamedTemporaryFile( mode='w', suffix='.py', delete=True ) as f: @@ -866,28 +888,32 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): if timeout is None: timeout = self.max_cpu_seconds - import sys # Prefer process_group (3.11+) over preexec_fn # (not fork-safe with threads). spawn_kwargs = {} if sys.version_info >= (3, 11): spawn_kwargs['process_group'] = 0 else: - # Fallback for 3.10; caveat: not fork-safe + # Fallback for 3.10; caveat: not fork-safe. + # Guard resource import for platforms where + # the module is unavailable. def _set_limits(): - import resource - resource.setrlimit( - resource.RLIMIT_CPU, - (self.max_cpu_seconds,) * 2, - ) - mem = self.max_memory_mb * 1024 * 1024 - resource.setrlimit( - resource.RLIMIT_AS, (mem, mem), - ) + try: + import resource + resource.setrlimit( + resource.RLIMIT_CPU, + (self.max_cpu_seconds,) * 2, + ) + mem = self.max_memory_mb * 1024 * 1024 + resource.setrlimit( + resource.RLIMIT_AS, (mem, mem), + ) + except (ImportError, OSError): + pass # timeout-only enforcement spawn_kwargs['preexec_fn'] = _set_limits - # Guard resource import for platforms where - # it is unavailable (falls back to timeout-only). + # Inline wrapper sets resource limits in the child + # process. Guarded for missing resource module. limit_code = ( f'try:\n' f' import resource\n' @@ -1122,8 +1148,11 @@ is new and has no tests yet. | Unit tests | Mock-based tests per executor | **Add `test_container_code_executor.py`**, add `test_local_sandbox_code_executor.py` | | Integration tests | Real executor tests (like `ExecuteSkillScriptTool` integration tests) | Add Docker-based container tests (CI-gated) | | Timeout tests | Scripts with `time.sleep()` to verify enforcement | Per-executor timeout tests | +| Timeout kill fallback | Verify `PermissionError` from `os.kill` triggers container restart | Mock `os.kill` to raise `PermissionError`, assert `container.restart()` called and `CodeExecutionResult.stderr` contains timeout message | +| Timeout kill success | Verify `os.kill(host_pid)` path when permitted | Mock `exec_inspect` to return PID, assert `os.kill` called with correct signal | | Security tests | Scripts attempting blocked operations | `restrict_builtins` bypass attempts, env var leakage | | Stateful tests | Multi-call sequences verifying variable persistence | Append-after-success, failure-does-not-poison, `execution_id` isolation | +| Stateful crash recovery | Verify error returned on REPL/container crash | Kill REPL mid-execution, assert error indicates state loss | --- @@ -1141,17 +1170,23 @@ is new and has no tests yet. 7. Update `ExecuteSkillScriptTool` to set per-invocation timeout via `CodeExecutionInput.timeout_seconds` -### Phase 2: Stateful Container (3-5 days) +### Phase 2: Stateful Container (5-8 days) + +Implement Option A (persistent process) directly, as recommended in +§5.2.2. This avoids the side-effect replay problems of Option C. 1. Unfreeze `stateful` on `ContainerCodeExecutor` -2. Implement cumulative code history with stdout suppression - (append-after-success invariant) -3. Add `execution_id`-based history isolation -4. Wire `execution_id` in `ExecuteSkillScriptTool` -5. Add `reset_state()` method -6. Add stateful execution tests (including failure-does-not-poison test) -7. Update samples and documentation -8. Evaluate persistent-process approach (Option A) for Phase 2b +2. Design persistent-process protocol: sentinel-delimited I/O for + output boundaries, error detection, and process health checks +3. Implement persistent Python REPL management (start, send code, + read output, detect crash/restart) +4. Add `execution_id`-based session isolation (one REPL per + `execution_id`) +5. Wire `execution_id` in `ExecuteSkillScriptTool` +6. Add `reset_state()` method (kills and restarts the REPL) +7. Add stateful execution tests (variable persistence, crash recovery, + `execution_id` isolation) +8. Update samples and documentation ### Phase 3: Security Hardening (5-7 days) @@ -1171,13 +1206,12 @@ is new and has no tests yet. ## 9. Open Questions -1. **Should we skip Phase 1 (cumulative replay) and go straight to - Phase 2 (persistent process) for stateful execution?** - The side-effect replay problem is fundamental to Option C. If the - persistent-process I/O boundary protocol can be solved with - reasonable complexity (e.g., sentinel-delimited output), the MVP - phase may not be worth the tech debt. Decision: Evaluate during - Phase 2 planning. +1. **What I/O boundary protocol should the persistent REPL use?** + The roadmap targets Option A (persistent process) directly. The + key design question is how to delimit output for each code block: + sentinel strings in stdout, JSON-envelope protocol, or a side + channel (e.g., file-based result). Sentinel strings are simplest + but can collide with user output. Decision: spike during Phase 2. 2. **Should `LocalSandboxCodeExecutor` support stateful execution?** Subprocess-based execution is inherently stateless. Stateful support @@ -1189,12 +1223,14 @@ is new and has no tests yet. `SecurityWarning` and documentation should steer users toward safer alternatives for anything beyond local prototyping. -4. **How should `ContainerCodeExecutor` handle container crashes in - stateful mode?** - If the container crashes (OOM, segfault), the code history is lost. - Options: (a) re-create container and replay history, (b) return error - and let user restart, (c) persist history to host volume. Recommend - (b) for simplicity. +4. **How should `ContainerCodeExecutor` handle container/REPL crashes + in stateful mode?** + If the container crashes (OOM, segfault) or the persistent REPL + dies, in-process state is lost. The executor returns an error + indicating state loss and lets the caller handle recovery (e.g., + start a new session). Automatic replay is not attempted because + prior code blocks may have had side effects that should not be + re-executed (consistent with §4.2.3 recovery policy). --- From f4fd7941db2ffccf29b5a45363973ba9d4003cb2 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 21 Feb 2026 19:37:48 -0800 Subject: [PATCH 09/53] docs: Fix PermissionError kill fallback, align non-goals with Option A MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - PermissionError from os.kill now triggers container.restart() instead of being silently swallowed (fixes timeout DoS for the common case where container runs as root and ADK does not) - ProcessLookupError remains the only silently-ignored exception (process already exited, no action needed) - Non-Goals §2.2 updated to describe persistent REPL crash recovery instead of stale cumulative-replay language Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 25 +++++++++++++---------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index d30ca7c180..581a193813 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -48,13 +48,13 @@ The following are explicitly **out of scope** for this design: `object.__subclasses__()`, `importlib` through `__builtins__`, etc. True isolation requires a process or container boundary. -2. **Idempotent replay of side-effecting code** — The stateful - `ContainerCodeExecutor` (Proposal 2) replays prior code blocks. - Code with non-idempotent side effects (file writes, network calls, - database mutations) is **not supported** in stateful replay mode. - The design suppresses stdout but cannot suppress arbitrary I/O. - Users must keep side-effecting code in the final block or use the - persistent-process approach (Phase 2 / Option A). +2. **Automatic state recovery after crash** — The stateful + `ContainerCodeExecutor` (Proposal 2) uses a persistent REPL + (Option A). If the REPL or container crashes, in-process state + is lost. The executor reports an error; it does **not** attempt + automatic replay of prior code blocks, because prior blocks may + have had non-idempotent side effects (file writes, network calls, + database mutations) that should not be re-executed. 3. **Multi-tenant per-execution isolation** — Per-execution isolation (fresh sandbox per call) is the domain of `GkeCodeExecutor` and @@ -314,10 +314,13 @@ def execute_code(self, invocation_context, code_execution_input): host_pid = info.get('Pid', 0) if host_pid > 0: os.kill(host_pid, signal.SIGKILL) - except (ProcessLookupError, PermissionError): - pass # Process already exited - except Exception: - # Last resort: restart the container + except ProcessLookupError: + pass # Process already exited — no action needed + except (PermissionError, Exception): + # os.kill failed (most commonly PermissionError + # when container runs as root and ADK does not). + # Restart the container to ensure the runaway + # process is terminated. try: self._container.restart(timeout=1) except Exception: From 4bb83a00f8354fd12e55fb78925845429007b86c Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 21 Feb 2026 19:51:26 -0800 Subject: [PATCH 10/53] docs: Surface cleanup failure as unhealthy state, add post-kill thread join - If both os.kill and container.restart() fail, set _healthy=False and return distinct "cleanup failed" error instead of silently returning a normal timeout result - Add thread.join(timeout=2) after kill/restart to prevent daemon thread leaks on repeated timeout failures - Log warning if worker thread is still alive after post-kill join - Add test cases for total cleanup failure and thread leak scenarios Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 47 +++++++++++++++++++++-- 1 file changed, 44 insertions(+), 3 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index 581a193813..d326e0f3ec 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -309,6 +309,7 @@ def execute_code(self, invocation_context, code_execution_input): # is the correct PID for os.kill() on the host. # (Killing from inside the container would require # the container-namespace PID, which is different.) + cleanup_failed = False try: info = self._client.api.exec_inspect(exec_id) host_pid = info.get('Pid', 0) @@ -323,9 +324,34 @@ def execute_code(self, invocation_context, code_execution_input): # process is terminated. try: self._container.restart(timeout=1) - except Exception: - pass + except Exception as restart_err: + cleanup_failed = True + logger.error( + 'Timeout cleanup failed: could not kill ' + 'process or restart container: %s', + restart_err, + ) + self._healthy = False + + # Give the worker thread a short window to finish + # after kill/restart, so it doesn't leak indefinitely. + thread.join(timeout=2) + if thread.is_alive(): + logger.warning( + 'Worker thread still alive after timeout ' + 'cleanup; daemon thread will linger until ' + 'process exit.' + ) + if cleanup_failed: + return CodeExecutionResult( + stderr=( + f'Execution timed out after {timeout}s ' + f'and cleanup failed — executor is ' + f'unhealthy. Reinitialize the executor ' + f'before further use.' + ) + ) return CodeExecutionResult( stderr=f'Execution timed out after {timeout}s' ) @@ -349,7 +375,20 @@ def execute_code(self, invocation_context, code_execution_input): concurrent execs or unrelated Python processes) - Dependency on `procps`/`pkill` being installed in the image -3. **Container restart as last resort** — If `os.kill` fails (e.g., +3. **Post-kill thread join** — After kill/restart, a short + `thread.join(timeout=2)` gives the worker thread time to exit + cleanly. If it's still alive, a warning is logged. The thread is + a daemon, so it will not prevent process exit, but repeated + timeout failures without this join could accumulate leaked threads. + +4. **Unhealthy state on total cleanup failure** — If both `os.kill` + and `container.restart()` fail, the executor sets `self._healthy + = False` and returns a distinct error message. Subsequent calls + should check `self._healthy` and raise early rather than queueing + work against a broken container. Reinitialization (stop + start) + is required to recover. + +5. **Container restart as last resort** — If `os.kill` fails (e.g., insufficient permissions when Docker runs rootless), restart the container. This is the most reliable fallback but destroys in-container state. @@ -1153,6 +1192,8 @@ is new and has no tests yet. | Timeout tests | Scripts with `time.sleep()` to verify enforcement | Per-executor timeout tests | | Timeout kill fallback | Verify `PermissionError` from `os.kill` triggers container restart | Mock `os.kill` to raise `PermissionError`, assert `container.restart()` called and `CodeExecutionResult.stderr` contains timeout message | | Timeout kill success | Verify `os.kill(host_pid)` path when permitted | Mock `exec_inspect` to return PID, assert `os.kill` called with correct signal | +| Timeout total failure | Verify both `os.kill` and `container.restart()` fail → unhealthy | Mock both to raise, assert `_healthy` is `False` and `stderr` contains "cleanup failed" | +| Timeout thread leak | Verify post-kill `join(2)` is called and warning logged if thread lingers | Mock thread to stay alive after kill, assert warning logged | | Security tests | Scripts attempting blocked operations | `restrict_builtins` bypass attempts, env var leakage | | Stateful tests | Multi-call sequences verifying variable persistence | Append-after-success, failure-does-not-poison, `execution_id` isolation | | Stateful crash recovery | Verify error returned on REPL/container crash | Kill REPL mid-execution, assert error indicates state loss | From 3221ac1a9b922cb86c494a22a50357eedd91fcf2 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 21 Feb 2026 20:08:22 -0800 Subject: [PATCH 11/53] docs: Add _healthy guard and post-restart readiness validation - Add _healthy check at execute_code entry: raise RuntimeError if executor is unhealthy from a prior cleanup failure - Add python3 --version readiness check after container.restart(), mirroring the init-time validation pattern - Update roadmap Phase 1 step 4 to include _healthy guard, post-restart validation, and post-kill thread join Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index d326e0f3ec..dcd314191c 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -278,6 +278,13 @@ import signal import threading def execute_code(self, invocation_context, code_execution_input): + # Fail fast if a prior timeout left the executor unhealthy. + if not self._healthy: + raise RuntimeError( + 'ContainerCodeExecutor is unhealthy after a failed ' + 'timeout cleanup. Call cleanup() and reinitialize.' + ) + input_t = code_execution_input.timeout_seconds timeout = ( input_t if input_t is not None @@ -324,6 +331,12 @@ def execute_code(self, invocation_context, code_execution_input): # process is terminated. try: self._container.restart(timeout=1) + # Re-validate runtime readiness after restart, + # mirroring the init-time check (see + # container_code_executor.py:190). + self._container.exec_run( + ['python3', '--version'] + ) except Exception as restart_err: cleanup_failed = True logger.error( @@ -1209,6 +1222,8 @@ is new and has no tests yet. `BaseCodeExecutor` 3. Implement thread-based timeout in `UnsafeLocalCodeExecutor` 4. Implement Docker exec kill timeout in `ContainerCodeExecutor` + (including `_healthy` guard, post-restart readiness validation, + and post-kill thread join) 5. Migrate `GkeCodeExecutor.timeout_seconds` to `default_timeout_seconds` 6. Add timeout tests for each executor 7. Update `ExecuteSkillScriptTool` to set per-invocation timeout via From c3a003d7423c4a7e2e81150daccc1b6774140cd9 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 21 Feb 2026 20:11:05 -0800 Subject: [PATCH 12/53] docs: Document _healthy lifecycle (init, failure, reinit) Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index dcd314191c..60b543280d 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -397,9 +397,13 @@ def execute_code(self, invocation_context, code_execution_input): 4. **Unhealthy state on total cleanup failure** — If both `os.kill` and `container.restart()` fail, the executor sets `self._healthy = False` and returns a distinct error message. Subsequent calls - should check `self._healthy` and raise early rather than queueing - work against a broken container. Reinitialization (stop + start) - is required to recover. + check `self._healthy` and raise early rather than queueing work + against a broken container. The `_healthy` lifecycle: + - Initialized to `True` in `__init__` (alongside container start) + - Set to `False` on total cleanup failure (kill + restart both fail) + - Set back to `True` after successful reinitialization (new + container created + readiness check passed via `cleanup()` then + `__init_container()`) 5. **Container restart as last resort** — If `os.kill` fails (e.g., insufficient permissions when Docker runs rootless), restart the From 369bba83481ae590bc71785df0b782e194e672e0 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 21 Feb 2026 20:13:49 -0800 Subject: [PATCH 13/53] docs: Add public reinitialize() method to ContainerCodeExecutor API Define the recovery path for unhealthy executors: reinitialize() stops the container, creates a new one, validates readiness, and sets _healthy=True. Referenced in error message, _healthy lifecycle, and roadmap Phase 1 step 5. Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index 60b543280d..f5cda2e43f 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -282,7 +282,7 @@ def execute_code(self, invocation_context, code_execution_input): if not self._healthy: raise RuntimeError( 'ContainerCodeExecutor is unhealthy after a failed ' - 'timeout cleanup. Call cleanup() and reinitialize.' + 'timeout cleanup. Call reinitialize() to recover.' ) input_t = code_execution_input.timeout_seconds @@ -401,9 +401,8 @@ def execute_code(self, invocation_context, code_execution_input): against a broken container. The `_healthy` lifecycle: - Initialized to `True` in `__init__` (alongside container start) - Set to `False` on total cleanup failure (kill + restart both fail) - - Set back to `True` after successful reinitialization (new - container created + readiness check passed via `cleanup()` then - `__init_container()`) + - Set back to `True` after successful `reinitialize()` (stops + current container, creates new one, passes readiness check) 5. **Container restart as last resort** — If `os.kill` fails (e.g., insufficient permissions when Docker runs rootless), restart the @@ -1228,9 +1227,14 @@ is new and has no tests yet. 4. Implement Docker exec kill timeout in `ContainerCodeExecutor` (including `_healthy` guard, post-restart readiness validation, and post-kill thread join) -5. Migrate `GkeCodeExecutor.timeout_seconds` to `default_timeout_seconds` -6. Add timeout tests for each executor -7. Update `ExecuteSkillScriptTool` to set per-invocation timeout via +5. Add public `reinitialize()` method to `ContainerCodeExecutor`: + stops the current container (if any), creates a new one, runs + readiness check, and sets `_healthy = True`. This is the + documented recovery path when `_healthy` is `False`. Callable + by users or by higher-level retry logic. +6. Migrate `GkeCodeExecutor.timeout_seconds` to `default_timeout_seconds` +7. Add timeout tests for each executor +8. Update `ExecuteSkillScriptTool` to set per-invocation timeout via `CodeExecutionInput.timeout_seconds` ### Phase 2: Stateful Container (5-8 days) From c7351836ade80c372e7e48454ff97db6b16d0707 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 21 Feb 2026 20:15:35 -0800 Subject: [PATCH 14/53] docs: Check exit_code on post-restart readiness validation Non-zero exit from python3 --version does not raise in Docker SDK, so check exit_code explicitly and raise to trigger the _healthy=False path, mirroring container_code_executor.py:169. Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index f5cda2e43f..6630a4b132 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -333,10 +333,15 @@ def execute_code(self, invocation_context, code_execution_input): self._container.restart(timeout=1) # Re-validate runtime readiness after restart, # mirroring the init-time check (see - # container_code_executor.py:190). - self._container.exec_run( + # container_code_executor.py:169). + check = self._container.exec_run( ['python3', '--version'] ) + if check.exit_code != 0: + raise RuntimeError( + f'Post-restart readiness check failed ' + f'(exit_code={check.exit_code})' + ) except Exception as restart_err: cleanup_failed = True logger.error( From 11f65f02346fa6c655cf53867fe06cab0a8feaf1 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Mon, 23 Feb 2026 01:27:28 -0800 Subject: [PATCH 15/53] feat: Add SkillsBench Docker-based evaluation pipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a complete SkillsBench benchmark runner with Docker-based execution and real pytest scoring, matching the official SkillsBench methodology. Key components: - TaskContainerExecutor: BaseCodeExecutor subclass that runs code inside per-task Docker containers built from each task's Dockerfile - ContainerBashTool: Bash tool enabling the agent to read/write files and run shell commands inside the task container - Lenient skill loader: Handles malformed SKILL.md frontmatter (invalid names, list-typed allowed-tools, dict-typed compatibility/metadata) loading 226/232 skills across all 87 tasks - Real pytest scoring: Copies tests into the container, runs test.sh, reads reward.txt for binary pass/fail — matching SkillsBench scoring - Per-task config from task.toml (agent/build/verifier timeouts) - Exponential backoff retry for 429/503 API errors - CLI with --filter, --rebuild, --build-only, --skip-tests flags Also includes the standalone eval runner, agent definition, metrics, 8 bundled skills, and evaluation set for local development. Co-Authored-By: Claude Opus 4.6 --- benchmarks/skillsbench/README.md | 131 ++ benchmarks/skillsbench/__init__.py | 15 + benchmarks/skillsbench/agent.py | 73 ++ .../eval_sets/skillsbench_eval.json | 230 ++++ benchmarks/skillsbench/full_runner.py | 1059 +++++++++++++++++ benchmarks/skillsbench/metrics.py | 226 ++++ benchmarks/skillsbench/runner.py | 313 +++++ .../skills/csv-aggregation/SKILL.md | 41 + .../csv-aggregation/references/sample-data.md | 15 + .../csv-aggregation/scripts/aggregate.py | 50 + .../skills/function-scaffold/SKILL.md | 30 + .../function-scaffold/scripts/scaffold.py | 66 + .../skills/html-extraction/SKILL.md | 35 + .../html-extraction/references/sample-page.md | 23 + .../skills/html-extraction/scripts/extract.py | 72 ++ .../skills/json-transform/SKILL.md | 34 + .../json-transform/references/sample-data.md | 21 + .../json-transform/scripts/transform.py | 63 + .../skillsbench/skills/log-parsing/SKILL.md | 48 + .../log-parsing/references/sample-logs.md | 18 + .../skills/log-parsing/scripts/parse.py | 79 ++ .../skillsbench/skills/regex-replace/SKILL.md | 36 + .../skills/regex-replace/scripts/replace.py | 54 + .../skillsbench/skills/rest-client/SKILL.md | 40 + .../skills/rest-client/references/api-docs.md | 36 + .../skills/rest-client/scripts/request.py | 65 + .../skills/statistical-calc/SKILL.md | 38 + .../skills/statistical-calc/scripts/stats.py | 54 + 28 files changed, 2965 insertions(+) create mode 100644 benchmarks/skillsbench/README.md create mode 100644 benchmarks/skillsbench/__init__.py create mode 100644 benchmarks/skillsbench/agent.py create mode 100644 benchmarks/skillsbench/eval_sets/skillsbench_eval.json create mode 100644 benchmarks/skillsbench/full_runner.py create mode 100644 benchmarks/skillsbench/metrics.py create mode 100644 benchmarks/skillsbench/runner.py create mode 100644 benchmarks/skillsbench/skills/csv-aggregation/SKILL.md create mode 100644 benchmarks/skillsbench/skills/csv-aggregation/references/sample-data.md create mode 100644 benchmarks/skillsbench/skills/csv-aggregation/scripts/aggregate.py create mode 100644 benchmarks/skillsbench/skills/function-scaffold/SKILL.md create mode 100644 benchmarks/skillsbench/skills/function-scaffold/scripts/scaffold.py create mode 100644 benchmarks/skillsbench/skills/html-extraction/SKILL.md create mode 100644 benchmarks/skillsbench/skills/html-extraction/references/sample-page.md create mode 100644 benchmarks/skillsbench/skills/html-extraction/scripts/extract.py create mode 100644 benchmarks/skillsbench/skills/json-transform/SKILL.md create mode 100644 benchmarks/skillsbench/skills/json-transform/references/sample-data.md create mode 100644 benchmarks/skillsbench/skills/json-transform/scripts/transform.py create mode 100644 benchmarks/skillsbench/skills/log-parsing/SKILL.md create mode 100644 benchmarks/skillsbench/skills/log-parsing/references/sample-logs.md create mode 100644 benchmarks/skillsbench/skills/log-parsing/scripts/parse.py create mode 100644 benchmarks/skillsbench/skills/regex-replace/SKILL.md create mode 100644 benchmarks/skillsbench/skills/regex-replace/scripts/replace.py create mode 100644 benchmarks/skillsbench/skills/rest-client/SKILL.md create mode 100644 benchmarks/skillsbench/skills/rest-client/references/api-docs.md create mode 100644 benchmarks/skillsbench/skills/rest-client/scripts/request.py create mode 100644 benchmarks/skillsbench/skills/statistical-calc/SKILL.md create mode 100644 benchmarks/skillsbench/skills/statistical-calc/scripts/stats.py diff --git a/benchmarks/skillsbench/README.md b/benchmarks/skillsbench/README.md new file mode 100644 index 0000000000..21432e3cb4 --- /dev/null +++ b/benchmarks/skillsbench/README.md @@ -0,0 +1,131 @@ +# SkillsBench Evaluation Harness for ADK + +Evaluates ADK's `SkillToolset` against tasks adapted from the +[SkillsBench](https://github.com/benchflow-ai/skillsbench) benchmark. + +## Overview + +This harness adapts 8 representative SkillsBench tasks as ADK skills and +evaluates them through the ADK evaluation framework. It tests whether an +agent can discover, load, and execute skills using the `SkillToolset` +tools: `list_skills`, `load_skill`, `load_skill_resource`, and +`execute_skill_script`. + +## Task Categories + +| # | Category | Skill | What it tests | +|---|----------|-------|---------------| +| 1 | Data Analysis | csv-aggregation | skill discovery + script execution | +| 2 | File Processing | json-transform | load_skill_resource + script | +| 3 | Web Scraping | html-extraction | skill with references | +| 4 | API Interaction | rest-client | multi-step skill usage | +| 5 | Text Transformation | regex-replace | simple script execution | +| 6 | Code Generation | function-scaffold | skill instruction following | +| 7 | Math Computation | statistical-calc | output validation | +| 8 | System Admin | log-parsing | complex skill with metadata | + +## Setup + +```bash +# From repo root +uv venv --python "python3.11" ".venv" +source .venv/bin/activate +uv sync --all-extras + +# Set your API key +export GOOGLE_API_KEY="your-key-here" +``` + +## Usage + +### Run with ADK CLI + +```bash +# Interactive web UI +adk web benchmarks/skillsbench + +# Run evaluation via ADK eval +adk eval benchmarks/skillsbench \ + benchmarks/skillsbench/eval_sets/skillsbench_eval.json +``` + +### Run standalone scorer + +```bash +python benchmarks/skillsbench/runner.py +python benchmarks/skillsbench/runner.py --num-runs 3 +python benchmarks/skillsbench/runner.py --eval-set path/to/custom_eval.json +``` + +### Output format + +The standalone runner produces a per-task results table and a +leaderboard-format summary: + +``` +============================================================ + Leaderboard Summary +============================================================ + Model: gemini-2.5-flash + Framework: ADK SkillToolset + Tasks: X/8 (XX.X%) + Avg Discovery: X.XX + Avg Tool Usage: X.XX + Elapsed: XX.Xs +============================================================ +``` + +## Custom Metrics + +Three metrics are provided in `metrics.py`: + +- **skill_discovery_score** — 1.0 if the agent called both `list_skills` + and `load_skill`, else 0.0 +- **tool_usage_score** — Fraction of expected tool calls that were made + (ANY_ORDER matching) +- **skillsbench_binary_score** — 1.0 if the final response contains all + expected reference lines, else 0.0 + +Reference these in eval configs via their dotted paths: +``` +benchmarks.skillsbench.metrics.skill_discovery_score +benchmarks.skillsbench.metrics.tool_usage_score +benchmarks.skillsbench.metrics.skillsbench_binary_score +``` + +## Directory Structure + +``` +benchmarks/skillsbench/ +├── __init__.py +├── README.md +├── agent.py # ADK agent with SkillToolset +├── skills/ # 8 adapted SkillsBench tasks +│ ├── csv-aggregation/ +│ ├── json-transform/ +│ ├── html-extraction/ +│ ├── rest-client/ +│ ├── regex-replace/ +│ ├── function-scaffold/ +│ ├── statistical-calc/ +│ └── log-parsing/ +├── eval_sets/ +│ └── skillsbench_eval.json # EvalSet with 8 cases +├── metrics.py # Custom metric functions +└── runner.py # Standalone runner +``` + +## Adding New Tasks + +1. Create a skill directory under `skills/` with a `SKILL.md` following + the [Agent Skills spec](https://github.com/benchflow-ai/skillsbench) +2. Add scripts under `skills//scripts/` +3. Add references under `skills//references/` (optional) +4. Add the skill name to `_SKILL_NAMES` in `agent.py` +5. Add a new `EvalCase` entry to `eval_sets/skillsbench_eval.json` + +## Security Note + +This harness uses `UnsafeLocalCodeExecutor` for skill script execution. +For production or untrusted skill scripts, use `ContainerCodeExecutor` +or `VertexAICodeExecutor` instead. diff --git a/benchmarks/skillsbench/__init__.py b/benchmarks/skillsbench/__init__.py new file mode 100644 index 0000000000..196d315208 --- /dev/null +++ b/benchmarks/skillsbench/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SkillsBench evaluation harness for ADK SkillToolset.""" diff --git a/benchmarks/skillsbench/agent.py b/benchmarks/skillsbench/agent.py new file mode 100644 index 0000000000..8fcbb98f55 --- /dev/null +++ b/benchmarks/skillsbench/agent.py @@ -0,0 +1,73 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SkillsBench evaluation agent with SkillToolset and Gemini Flash. + +This agent loads all skills from the skills/ directory and uses +SkillToolset to provide list_skills, load_skill, load_skill_resource, +and execute_skill_script tools. It is designed to be evaluated against +the SkillsBench benchmark tasks. + +WARNING: This agent uses UnsafeLocalCodeExecutor for script execution. +For production use, prefer ContainerCodeExecutor or VertexAICodeExecutor. +""" + +import pathlib + +from google.adk import Agent +from google.adk.code_executors.unsafe_local_code_executor import UnsafeLocalCodeExecutor +from google.adk.skills import load_skill_from_dir +from google.adk.tools.skill_toolset import SkillToolset + +_SKILLS_DIR = pathlib.Path(__file__).parent / "skills" + +_SKILL_NAMES = [ + "csv-aggregation", + "json-transform", + "html-extraction", + "rest-client", + "regex-replace", + "function-scaffold", + "statistical-calc", + "log-parsing", +] + +_skills = [load_skill_from_dir(_SKILLS_DIR / name) for name in _SKILL_NAMES] + +skill_toolset = SkillToolset( + skills=_skills, + code_executor=UnsafeLocalCodeExecutor(), +) + +root_agent = Agent( + model="gemini-3-flash-preview", + name="skillsbench_agent", + description=( + "An agent that completes tasks by discovering and using" + " available skills from the SkillsBench benchmark." + ), + instruction=( + "You are an agent that completes tasks by discovering and using" + " available skills. Follow this workflow:\n" + "1. Use list_skills to find relevant skills for the task.\n" + "2. Use load_skill to read the skill's instructions carefully.\n" + "3. Use load_skill_resource to examine references or sample data" + " if available.\n" + "4. Use execute_skill_script to run the skill's scripts with" + " appropriate arguments.\n" + "5. Interpret the output and present a clear answer.\n\n" + "Always check skill instructions before executing scripts." + ), + tools=[skill_toolset], +) diff --git a/benchmarks/skillsbench/eval_sets/skillsbench_eval.json b/benchmarks/skillsbench/eval_sets/skillsbench_eval.json new file mode 100644 index 0000000000..ec2a0bbc1c --- /dev/null +++ b/benchmarks/skillsbench/eval_sets/skillsbench_eval.json @@ -0,0 +1,230 @@ +{ + "eval_set_id": "skillsbench-adk-v1", + "name": "SkillsBench ADK Evaluation", + "description": "8 representative SkillsBench tasks adapted as ADK skills to evaluate SkillToolset with Gemini Flash.", + "eval_cases": [ + { + "eval_id": "data_analysis_csv_aggregation", + "conversation": [ + { + "invocation_id": "inv-csv-agg-01", + "user_content": { + "parts": [{"text": "Aggregate the sample employee CSV data by department and show salary statistics for each department."}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "Group: Engineering\n count: 3\n sum: 285000\n mean: 95000\nGroup: Marketing\n count: 3\n sum: 214000\nGroup: Sales\n count: 2\n sum: 157000"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "list_skills", "args": {}}, + {"name": "load_skill", "args": {"skill_name": "csv-aggregation"}}, + {"name": "load_skill_resource", "args": {"skill_name": "csv-aggregation", "resource_type": "references", "resource_id": "sample-data.md"}}, + {"name": "execute_skill_script", "args": {"skill_name": "csv-aggregation", "script_name": "aggregate.py", "input_args": "group_col=department metric_col=salary"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 + }, + { + "eval_id": "file_processing_json_transform", + "conversation": [ + { + "invocation_id": "inv-json-tf-01", + "user_content": { + "parts": [{"text": "Flatten the nested user JSON data into a flat structure with dot-notation keys."}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "user.name\nuser.age\nuser.address.city\nuser.address.state\nuser.address.zip"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "list_skills", "args": {}}, + {"name": "load_skill", "args": {"skill_name": "json-transform"}}, + {"name": "load_skill_resource", "args": {"skill_name": "json-transform", "resource_type": "references", "resource_id": "sample-data.md"}}, + {"name": "execute_skill_script", "args": {"skill_name": "json-transform", "script_name": "transform.py", "input_args": "flatten=true"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 + }, + { + "eval_id": "web_scraping_html_extraction", + "conversation": [ + { + "invocation_id": "inv-html-ext-01", + "user_content": { + "parts": [{"text": "Extract all the product data from the HTML page as a table with product name, price, and stock columns."}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "Product,Price,Stock\nLaptop,999.99,15\nPhone,699.99,42\nTablet,449.99,28"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "list_skills", "args": {}}, + {"name": "load_skill", "args": {"skill_name": "html-extraction"}}, + {"name": "load_skill_resource", "args": {"skill_name": "html-extraction", "resource_type": "references", "resource_id": "sample-page.md"}}, + {"name": "execute_skill_script", "args": {"skill_name": "html-extraction", "script_name": "extract.py", "input_args": "target=table"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 + }, + { + "eval_id": "api_interaction_rest_client", + "conversation": [ + { + "invocation_id": "inv-rest-01", + "user_content": { + "parts": [{"text": "Use the REST API to fetch the list of users and then get the details for user with ID 2."}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "Bob\nbob@example.com"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "list_skills", "args": {}}, + {"name": "load_skill", "args": {"skill_name": "rest-client"}}, + {"name": "load_skill_resource", "args": {"skill_name": "rest-client", "resource_type": "references", "resource_id": "api-docs.md"}}, + {"name": "execute_skill_script", "args": {"skill_name": "rest-client", "script_name": "request.py", "input_args": "method=GET endpoint=/users"}}, + {"name": "execute_skill_script", "args": {"skill_name": "rest-client", "script_name": "request.py", "input_args": "method=GET endpoint=/users/2"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 + }, + { + "eval_id": "text_transformation_regex_replace", + "conversation": [ + { + "invocation_id": "inv-regex-01", + "user_content": { + "parts": [{"text": "Replace all numbers in the text 'Order 123 has 45 items at $67' with the word NUM."}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "Result: Order NUM has NUM items at $NUM\nMatches: 3"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "list_skills", "args": {}}, + {"name": "load_skill", "args": {"skill_name": "regex-replace"}}, + {"name": "execute_skill_script", "args": {"skill_name": "regex-replace", "script_name": "replace.py", "input_args": "pattern=\\d+ replacement=NUM text='Order 123 has 45 items at $67'"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 + }, + { + "eval_id": "code_generation_function_scaffold", + "conversation": [ + { + "invocation_id": "inv-scaffold-01", + "user_content": { + "parts": [{"text": "Generate a Python function scaffold for a function called calculate_bmi that takes weight (float) and height (float) and returns a float."}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "def calculate_bmi(weight: float, height: float) -> float:"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "list_skills", "args": {}}, + {"name": "load_skill", "args": {"skill_name": "function-scaffold"}}, + {"name": "execute_skill_script", "args": {"skill_name": "function-scaffold", "script_name": "scaffold.py", "input_args": "name=calculate_bmi params=weight:float,height:float returns=float description='Calculate Body Mass Index'"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 + }, + { + "eval_id": "math_computation_statistical_calc", + "conversation": [ + { + "invocation_id": "inv-stats-01", + "user_content": { + "parts": [{"text": "Compute descriptive statistics (mean, median, standard deviation) for the dataset: 10, 20, 30, 40, 50, 60, 70, 80, 90, 100."}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "Mean: 55.00\nMedian: 55.00\nStd Dev: 28.72"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "list_skills", "args": {}}, + {"name": "load_skill", "args": {"skill_name": "statistical-calc"}}, + {"name": "execute_skill_script", "args": {"skill_name": "statistical-calc", "script_name": "stats.py", "input_args": "data=10,20,30,40,50,60,70,80,90,100"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 + }, + { + "eval_id": "system_admin_log_parsing", + "conversation": [ + { + "invocation_id": "inv-logs-01", + "user_content": { + "parts": [{"text": "Analyze the system logs and give me a summary showing the count of each log level (ERROR, WARNING, INFO, DEBUG)."}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "ERROR: 3\nWARNING: 2\nINFO: 5\nDEBUG: 2"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "list_skills", "args": {}}, + {"name": "load_skill", "args": {"skill_name": "log-parsing"}}, + {"name": "load_skill_resource", "args": {"skill_name": "log-parsing", "resource_type": "references", "resource_id": "sample-logs.md"}}, + {"name": "execute_skill_script", "args": {"skill_name": "log-parsing", "script_name": "parse.py", "input_args": "level=ALL format=summary"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 +} diff --git a/benchmarks/skillsbench/full_runner.py b/benchmarks/skillsbench/full_runner.py new file mode 100644 index 0000000000..0371c7e590 --- /dev/null +++ b/benchmarks/skillsbench/full_runner.py @@ -0,0 +1,1059 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Full SkillsBench runner — Docker-based with real pytest scoring. + +Builds Docker images per-task from each task's environment/Dockerfile, +runs agent scripts inside the container (where /root/ data files exist), +then runs the real pytest tests inside the same container for binary +pass/fail scoring matching SkillsBench methodology. + +Usage: + python -u -m benchmarks.skillsbench.full_runner + python -u -m benchmarks.skillsbench.full_runner --filter citation-check + python -u -m benchmarks.skillsbench.full_runner --build-only + python -u -m benchmarks.skillsbench.full_runner --skip-tests + +Environment variables: + GOOGLE_API_KEY — API key for authentication + GOOGLE_GENAI_USE_VERTEXAI — Set to 1 for Vertex AI backend +""" + +from __future__ import annotations + +import argparse +import asyncio +import io +import json +import logging +import pathlib +import re +import sys +import tarfile +import time +import tomllib +from typing import Any +import uuid + +import docker +from google.adk import Agent +from google.adk.agents.invocation_context import InvocationContext +from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService +from google.adk.code_executors.base_code_executor import BaseCodeExecutor +from google.adk.code_executors.code_execution_utils import CodeExecutionInput +from google.adk.code_executors.code_execution_utils import CodeExecutionResult +from google.adk.evaluation.eval_case import get_all_tool_calls +from google.adk.evaluation.eval_case import Invocation +from google.adk.evaluation.evaluation_generator import EvaluationGenerator +from google.adk.events.event import Event +from google.adk.memory.in_memory_memory_service import InMemoryMemoryService +from google.adk.models.google_llm import Gemini +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.skills import load_skill_from_dir +from google.adk.skills.models import Frontmatter +from google.adk.skills.models import Resources +from google.adk.skills.models import Script +from google.adk.skills.models import Skill +from google.adk.tools.base_tool import BaseTool +from google.adk.tools.skill_toolset import SkillToolset +from google.adk.utils.context_utils import Aclosing +from google.genai import types as genai_types +from pydantic import ConfigDict +from pydantic import Field +from typing_extensions import override +import yaml + +logger = logging.getLogger(__name__) + +_DEFAULT_TASKS_DIR = pathlib.Path("/tmp/skillsbench/tasks") +_MODEL_NAME = "gemini-3-flash-preview" + +# ── Timeouts (seconds) ────────────────────────────────────────────── +_DEFAULT_AGENT_TIMEOUT = 900.0 +_DEFAULT_BUILD_TIMEOUT = 600.0 +_DEFAULT_VERIFIER_TIMEOUT = 900.0 + + +# ── helpers ────────────────────────────────────────────────────────── + + +def _slugify(name: str) -> str: + """Convert a name to lowercase kebab-case. + + Lowercases, replaces underscores and spaces with hyphens, + strips non-alphanumeric/non-hyphen chars, collapses runs + of hyphens, and trims leading/trailing hyphens. + """ + s = name.lower() + s = s.replace("_", "-").replace(" ", "-") + s = re.sub(r"[^a-z0-9-]", "", s) + s = re.sub(r"-{2,}", "-", s) + return s.strip("-") + + +def _load_dir(directory: pathlib.Path) -> dict[str, str]: + """Recursively load files from a directory into a dict. + + Mirrors google.adk.skills._utils._load_dir so the runner + does not depend on a private import. + """ + files: dict[str, str] = {} + if directory.exists() and directory.is_dir(): + for file_path in directory.rglob("*"): + if "__pycache__" in file_path.parts: + continue + if file_path.is_file(): + relative = file_path.relative_to(directory) + try: + files[str(relative)] = file_path.read_text(encoding="utf-8") + except UnicodeDecodeError: + continue + return files + + +# ── Docker prereq ──────────────────────────────────────────────────── + + +def ensure_docker_running() -> docker.DockerClient: + """Return a Docker client or fail fast with a clear message.""" + try: + client = docker.from_env() + client.ping() + return client + except Exception as exc: + print( + "\nERROR: Docker is not running or not accessible.\n" + "Please start Docker Desktop and try again.\n" + f" Detail: {exc}\n", + file=sys.stderr, + ) + sys.exit(1) + + +# ── TaskContainerExecutor ──────────────────────────────────────────── + + +class TaskContainerExecutor(BaseCodeExecutor): + """Code executor that runs Python inside a Docker container. + + Unlike ADK's ContainerCodeExecutor (which uses atexit for + cleanup), this gives direct lifecycle control for test + injection after the agent finishes. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + stateful: bool = Field(default=False, frozen=True, exclude=True) + optimize_data_file: bool = Field(default=False, frozen=True, exclude=True) + + _container: object = None + _client: object = None + + def start( + self, + image_tag: str, + client: docker.DockerClient, + ) -> None: + """Start a detached container from *image_tag*.""" + self._client = client + self._container = client.containers.run( + image_tag, + command="sleep infinity", + detach=True, + tty=True, + working_dir="/root", + ) + + @override + def execute_code( + self, + invocation_context: InvocationContext, + code_execution_input: CodeExecutionInput, + ) -> CodeExecutionResult: + """Run ``python3 -c `` inside the container.""" + if self._container is None: + return CodeExecutionResult( + stdout="", + stderr="Container not started", + ) + code = code_execution_input.code + rc, output = self._container.exec_run( + ["python3", "-c", code], + workdir="/root", + demux=True, + ) + stdout = (output[0] or b"").decode("utf-8", errors="replace") + stderr = (output[1] or b"").decode("utf-8", errors="replace") + return CodeExecutionResult( + stdout=stdout, + stderr=stderr, + output_files=[], + ) + + def copy_to_container( + self, + local_path: pathlib.Path, + container_path: str, + ) -> None: + """Copy a local directory into the container via put_archive.""" + if self._container is None: + raise RuntimeError("Container not started") + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w") as tar: + tar.add(str(local_path), arcname=".") + buf.seek(0) + self._container.put_archive(container_path, buf) + + def exec_in_container( + self, + cmd: list[str], + workdir: str = "/root", + timeout: float | None = None, + ) -> tuple[int, str]: + """Run an arbitrary command inside the container. + + Returns (exit_code, combined_output). + """ + if self._container is None: + raise RuntimeError("Container not started") + rc, output = self._container.exec_run( + cmd, + workdir=workdir, + demux=True, + ) + stdout = (output[0] or b"").decode("utf-8", errors="replace") + stderr = (output[1] or b"").decode("utf-8", errors="replace") + combined = stdout + if stderr: + combined = combined + "\n" + stderr if combined else stderr + return rc, combined + + def read_file_from_container(self, path: str) -> str | None: + """Read a single file from the container, or None.""" + if self._container is None: + return None + try: + bits, _ = self._container.get_archive(path) + buf = io.BytesIO() + for chunk in bits: + buf.write(chunk) + buf.seek(0) + with tarfile.open(fileobj=buf, mode="r") as tar: + for member in tar.getmembers(): + if member.isfile(): + f = tar.extractfile(member) + if f: + return f.read().decode("utf-8", errors="replace") + except Exception: + return None + return None + + def stop(self) -> None: + """Stop and remove the container.""" + if self._container is not None: + try: + self._container.stop(timeout=5) + except Exception: + pass + try: + self._container.remove(force=True) + except Exception: + pass + self._container = None + + +# ── Docker image builder ───────────────────────────────────────────── + + +def build_task_image( + task_dir: pathlib.Path, + client: docker.DockerClient, + *, + rebuild: bool = False, + build_timeout: float = _DEFAULT_BUILD_TIMEOUT, +) -> str: + """Build (or reuse) a Docker image for one task. + + Returns the image tag. + """ + tag = f"skillsbench-{task_dir.name}:latest" + if not rebuild: + try: + client.images.get(tag) + return tag + except docker.errors.ImageNotFound: + pass + + build_ctx = task_dir / "environment" + if not (build_ctx / "Dockerfile").exists(): + raise FileNotFoundError(f"No Dockerfile in {build_ctx}") + + logger.info("Building image %s …", tag) + client.images.build( + path=str(build_ctx), + tag=tag, + rm=True, + timeout=int(build_timeout), + ) + return tag + + +# ── Lenient skill loader ───────────────────────────────────────────── + + +def _load_skill_lenient(skill_dir: pathlib.Path) -> Skill | None: + """Load a skill, falling back to lenient parsing on failure. + + Tries the strict ``load_skill_from_dir`` first. On any error, + manually parses SKILL.md, fixes the frontmatter fields, and + builds a ``Skill`` directly. + + Returns None only for truly unrecoverable cases (no SKILL.md, + unparseable YAML, no description). + """ + # ── strict path ── + try: + return load_skill_from_dir(skill_dir) + except Exception: + pass + + # ── lenient path ── + skill_md = None + for name in ("SKILL.md", "skill.md"): + p = skill_dir / name + if p.exists(): + skill_md = p + break + if skill_md is None: + return None + + try: + content = skill_md.read_text(encoding="utf-8") + except Exception: + return None + + if not content.startswith("---"): + return None + + parts = content.split("---", 2) + if len(parts) < 3: + return None + + try: + parsed = yaml.safe_load(parts[1]) + except yaml.YAMLError: + return None + + if not isinstance(parsed, dict): + return None + + body = parts[2].strip() + + # Fix name: slugify and force to match dir name + parsed["name"] = _slugify(skill_dir.name) + + # Fix allowed-tools: list → comma-separated string + for key in ("allowed-tools", "allowed_tools"): + val = parsed.get(key) + if isinstance(val, list): + parsed[key] = ", ".join(str(v) for v in val) + + # Fix compatibility: dict → JSON string + if isinstance(parsed.get("compatibility"), dict): + parsed["compatibility"] = json.dumps(parsed["compatibility"]) + + # Fix metadata values: ensure all are strings + meta = parsed.get("metadata") + if isinstance(meta, dict): + parsed["metadata"] = { + str(k): str(v) if not isinstance(v, str) else v for k, v in meta.items() + } + + # Ensure description exists + if not parsed.get("description"): + parsed["description"] = f"Skill from {skill_dir.name}" + + try: + frontmatter = Frontmatter.model_validate(parsed) + except Exception: + return None + + # Load resources + references = _load_dir(skill_dir / "references") + assets = _load_dir(skill_dir / "assets") + raw_scripts = _load_dir(skill_dir / "scripts") + scripts = {n: Script(src=c) for n, c in raw_scripts.items()} + resources = Resources( + references=references, + assets=assets, + scripts=scripts, + ) + + return Skill( + frontmatter=frontmatter, + instructions=body, + resources=resources, + ) + + +def load_task_skills( + task_dir: pathlib.Path, +) -> tuple[list[Skill], list[str]]: + """Load all skills from a task's environment/skills/. + + Uses lenient loading to handle malformed frontmatter. + + Returns: + (loaded_skills, error_messages) + """ + skills_dir = task_dir / "environment" / "skills" + if not skills_dir.exists() or not skills_dir.is_dir(): + return [], ["No environment/skills/ directory found"] + + skills: list[Skill] = [] + errors: list[str] = [] + for sd in sorted(skills_dir.iterdir()): + if not sd.is_dir(): + continue + skill = _load_skill_lenient(sd) + if skill is not None: + skills.append(skill) + else: + errors.append(f"{sd.name}: could not load") + return skills, errors + + +# ── task.toml parser ───────────────────────────────────────────────── + + +def parse_task_config( + task_dir: pathlib.Path, +) -> dict[str, float]: + """Read task.toml and return timeout values. + + Returns dict with keys: agent_timeout, build_timeout, + verifier_timeout. + """ + toml_path = task_dir / "task.toml" + defaults = { + "agent_timeout": _DEFAULT_AGENT_TIMEOUT, + "build_timeout": _DEFAULT_BUILD_TIMEOUT, + "verifier_timeout": _DEFAULT_VERIFIER_TIMEOUT, + } + if not toml_path.exists(): + return defaults + + try: + with open(toml_path, "rb") as f: + cfg = tomllib.load(f) + except Exception: + return defaults + + agent_sec = cfg.get("agent", {}) + env_sec = cfg.get("environment", {}) + ver_sec = cfg.get("verifier", {}) + + return { + "agent_timeout": float( + agent_sec.get("timeout_sec", defaults["agent_timeout"]) + ), + "build_timeout": float( + env_sec.get("build_timeout_sec", defaults["build_timeout"]) + ), + "verifier_timeout": float( + ver_sec.get("timeout_sec", defaults["verifier_timeout"]) + ), + } + + +# ── scoring ────────────────────────────────────────────────────────── + + +def score_task_in_container( + executor: TaskContainerExecutor, + task_dir: pathlib.Path, +) -> int: + """Run the real pytest tests inside the container. + + 1. Copy task_dir/tests/ into /tests/ in the container. + 2. Run ``bash /tests/test.sh``. + 3. Read /logs/verifier/reward.txt — ``1`` = PASS. + + Returns 1 (pass) or 0 (fail). + """ + tests_dir = task_dir / "tests" + if not tests_dir.exists(): + logger.warning("No tests/ directory for %s", task_dir.name) + return 0 + + # Ensure /tests/ directory exists in container + executor.exec_in_container(["mkdir", "-p", "/tests"]) + + # Copy tests into container + executor.copy_to_container(tests_dir, "/tests") + + # Make test.sh executable + executor.exec_in_container(["chmod", "+x", "/tests/test.sh"]) + + # Run test.sh + rc, output = executor.exec_in_container( + ["bash", "/tests/test.sh"], + workdir="/root", + ) + logger.debug( + "test.sh for %s exited %d:\n%s", + task_dir.name, + rc, + output[:2000], + ) + + # Read reward + reward = executor.read_file_from_container("/logs/verifier/reward.txt") + if reward is not None and reward.strip() == "1": + return 1 + return 0 + + +def score_task_heuristic( + actual_invocations: list[Invocation], +) -> int: + """Heuristic fallback (--skip-tests mode). + + Passes if the agent loaded a skill, used it, and + produced a final response. + """ + tool_names: list[str] = [] + for inv in actual_invocations: + for tc in get_all_tool_calls(inv.intermediate_data): + tool_names.append(tc.name) + + has_load = "load_skill" in tool_names + has_use = ( + "execute_skill_script" in tool_names + or "load_skill_resource" in tool_names + ) + has_response = False + for inv in actual_invocations: + if inv.final_response and inv.final_response.parts: + for part in inv.final_response.parts: + if part.text: + has_response = True + break + if has_response: + break + + return 1 if (has_load and has_use and has_response) else 0 + + +# ── container bash tool ─────────────────────────────────────────── + + +class ContainerBashTool(BaseTool): + """Simple bash tool that runs commands inside the task container.""" + + def __init__(self, executor: TaskContainerExecutor): + super().__init__( + name="bash", + description=( + "Run a bash command inside the task environment." + " Use this to read files, write files, install" + " packages, or run arbitrary shell commands." + ), + ) + self._executor = executor + + def _get_declaration(self) -> genai_types.FunctionDeclaration | None: + return genai_types.FunctionDeclaration( + name=self.name, + description=self.description, + parameters_json_schema={ + "type": "object", + "properties": { + "command": { + "type": "string", + "description": "The bash command to execute.", + }, + }, + "required": ["command"], + }, + ) + + async def run_async(self, *, args, tool_context) -> Any: + command = args.get("command", "") + if not command: + return {"error": "command is required"} + rc, output = self._executor.exec_in_container( + ["bash", "-c", command], + ) + # Truncate very long output to avoid blowing context + if len(output) > 10000: + output = output[:10000] + "\n... (truncated)" + return { + "exit_code": rc, + "output": output, + } + + +# ── agent builder ──────────────────────────────────────────────────── + + +def build_agent( + skills: list[Skill], + executor: TaskContainerExecutor, +) -> Agent: + """Create a fresh agent with the given skills.""" + toolset = SkillToolset( + skills=skills, + code_executor=executor, + ) + bash_tool = ContainerBashTool(executor) + model = Gemini( + model=_MODEL_NAME, + retry_options=genai_types.HttpRetryOptions( + attempts=5, + initialDelay=2.0, + maxDelay=60.0, + expBase=2.0, + httpStatusCodes=[429, 503], + ), + ) + return Agent( + model=model, + name="skillsbench_agent", + description=( + "An agent that completes tasks by discovering and using" + " available skills from the SkillsBench benchmark." + ), + instruction=( + "You are an agent that completes tasks by discovering" + " and using available skills. Follow this workflow:\n" + "1. Use list_skills to find relevant skills.\n" + "2. Use load_skill to read the skill's instructions" + " carefully.\n" + "3. Use load_skill_resource to examine references or" + " sample data if available.\n" + "4. Use execute_skill_script to run the skill's" + " scripts with appropriate arguments.\n" + "5. Use bash to read/write files, install packages," + " or run any shell command in the environment.\n" + "6. Write the required output files and present a" + " clear answer.\n" + "\nAlways check skill instructions before executing" + " scripts." + ), + generate_content_config=genai_types.GenerateContentConfig( + temperature=1.0, + thinking_config=genai_types.ThinkingConfig( + thinking_level="HIGH", + ), + ), + tools=[toolset, bash_tool], + ) + + +# ── agent runner ───────────────────────────────────────────────────── + + +async def run_task( + agent: Agent, + user_query: str, +) -> list[Invocation]: + """Run a single task and return invocations.""" + session_service = InMemorySessionService() + artifact_service = InMemoryArtifactService() + memory_service = InMemoryMemoryService() + + app_name = "skillsbench_full_eval" + user_id = "eval_user" + session_id = str(uuid.uuid4()) + + await session_service.create_session( + app_name=app_name, + user_id=user_id, + state={}, + session_id=session_id, + ) + + user_content = genai_types.Content( + role="user", + parts=[genai_types.Part.from_text(text=user_query)], + ) + + async with Runner( + app_name=app_name, + agent=agent, + artifact_service=artifact_service, + session_service=session_service, + memory_service=memory_service, + ) as runner: + events: list[Event] = [] + async with Aclosing( + runner.run_async( + user_id=user_id, + session_id=session_id, + new_message=user_content, + ) + ) as agen: + invocation_id = None + async for event in agen: + if not invocation_id: + invocation_id = event.invocation_id + events.append( + Event( + content=user_content, + author="user", + invocation_id=invocation_id, + ) + ) + events.append(event) + + return EvaluationGenerator.convert_events_to_eval_invocations(events) + + +# ── task result ────────────────────────────────────────────────────── + + +class TaskResult: + """Result of a single task evaluation.""" + + def __init__( + self, + task_name: str, + score: int, + num_skills: int, + elapsed: float, + error: str | None = None, + ): + self.task_name = task_name + self.score = score + self.num_skills = num_skills + self.elapsed = elapsed + self.error = error + + +# ── task discovery ─────────────────────────────────────────────────── + + +def discover_tasks( + tasks_dir: pathlib.Path, + filter_pattern: str | None = None, +) -> list[pathlib.Path]: + """Find task directories containing instruction.md.""" + tasks = [] + for d in sorted(tasks_dir.iterdir()): + if not d.is_dir(): + continue + if not (d / "instruction.md").exists(): + continue + if filter_pattern and filter_pattern not in d.name: + continue + tasks.append(d) + return tasks + + +# ── evaluate one task ──────────────────────────────────────────────── + + +async def evaluate_task( + task_dir: pathlib.Path, + client: docker.DockerClient, + *, + timeout_override: float | None = None, + rebuild: bool = False, + skip_tests: bool = False, +) -> TaskResult: + """Evaluate a single SkillsBench task end-to-end. + + 1. Parse task.toml for timeouts + 2. Build Docker image (or use cached) + 3. Load skills with lenient loader + 4. Start TaskContainerExecutor + 5. Build Agent with SkillToolset + 6. Run agent with instruction.md + 7. Run tests inside same container (or heuristic) + 8. Stop container + """ + task_name = task_dir.name + start = time.time() + executor = TaskContainerExecutor() + num_skills = 0 + agent_timeout = _DEFAULT_AGENT_TIMEOUT + + try: + # 1. Parse config + config = parse_task_config(task_dir) + agent_timeout = timeout_override or config["agent_timeout"] + + # 2. Build image + try: + image_tag = build_task_image( + task_dir, + client, + rebuild=rebuild, + build_timeout=config["build_timeout"], + ) + except Exception as exc: + return TaskResult( + task_name=task_name, + score=0, + num_skills=0, + elapsed=time.time() - start, + error=f"image build: {str(exc)[:100]}", + ) + + # 3. Load instruction + try: + user_query = ( + (task_dir / "instruction.md").read_text(encoding="utf-8").strip() + ) + except Exception as exc: + return TaskResult( + task_name=task_name, + score=0, + num_skills=0, + elapsed=time.time() - start, + error=f"instruction.md: {exc}", + ) + + # 4. Load skills + skills, skill_errors = load_task_skills(task_dir) + num_skills = len(skills) + if not skills: + msg = "; ".join(skill_errors) if skill_errors else "No skills" + return TaskResult( + task_name=task_name, + score=0, + num_skills=0, + elapsed=time.time() - start, + error=f"skills: {msg}", + ) + + # 5. Start container + executor.start(image_tag, client) + + # 6. Build agent and run + agent = build_agent(skills, executor) + invocations = await asyncio.wait_for( + run_task(agent, user_query), + timeout=agent_timeout, + ) + + # 7. Score + if skip_tests: + score = score_task_heuristic(invocations) + else: + score = score_task_in_container(executor, task_dir) + + return TaskResult( + task_name=task_name, + score=score, + num_skills=num_skills, + elapsed=time.time() - start, + ) + + except Exception as exc: + # Score container tests even on error — the agent + # may have produced partial output files. + score = 0 + if not skip_tests and executor._container is not None: + try: + score = score_task_in_container(executor, task_dir) + except Exception: + pass + is_timeout = isinstance(exc, asyncio.TimeoutError) + err_msg = f"timeout ({agent_timeout}s)" if is_timeout else str(exc)[:120] + return TaskResult( + task_name=task_name, + score=score, + num_skills=num_skills, + elapsed=time.time() - start, + error=err_msg, + ) + finally: + executor.stop() + + +# ── printing ───────────────────────────────────────────────────────── + + +def print_header(): + print() + print("=" * 60) + print(" SkillsBench Full Evaluation — Docker + Pytest") + print("=" * 60) + print() + + +def print_task_result( + idx: int, + total: int, + result: TaskResult, +): + mark = "PASS" if result.score == 1 else "FAIL" + name = result.task_name[:35].ljust(35) + if result.error: + detail = f"({result.error})" + else: + detail = f"({result.num_skills} skills, {result.elapsed:.1f}s)" + print(f"[{idx:>2}/{total}] {name} {mark} {detail}") + sys.stdout.flush() + + +def print_summary( + results: list[TaskResult], + total_tasks: int, + elapsed: float, +): + passed = sum(1 for r in results if r.score == 1) + loadable = sum(1 for r in results if r.num_skills > 0) + times = [r.elapsed for r in results if r.error != "timeout"] + avg_time = sum(times) / max(len(times), 1) + pct = (passed / max(total_tasks, 1)) * 100 + + print() + print("=" * 60) + print(" Leaderboard Summary") + print("=" * 60) + print(f" Model: {_MODEL_NAME}") + print(f" Framework: ADK SkillToolset (Docker)") + print(f" Score: {passed}/{total_tasks} ({pct:.1f}%)") + print(f" Loadable tasks: {loadable}/{total_tasks}") + print(f" Avg time/task: {avg_time:.1f}s") + print(f" Total elapsed: {elapsed:.0f}s") + print("=" * 60) + + +# ── full evaluation loop ───────────────────────────────────────────── + + +async def run_full_evaluation( + tasks_dir: pathlib.Path, + client: docker.DockerClient, + *, + timeout_override: float | None = None, + filter_pattern: str | None = None, + rebuild: bool = False, + build_only: bool = False, + skip_tests: bool = False, +) -> list[TaskResult]: + """Run evaluation on all matching tasks.""" + task_dirs = discover_tasks(tasks_dir, filter_pattern) + total = len(task_dirs) + print(f"Found {total} tasks in {tasks_dir}\n") + + if build_only: + for idx, td in enumerate(task_dirs, 1): + name = td.name[:35].ljust(35) + try: + config = parse_task_config(td) + tag = build_task_image( + td, + client, + rebuild=rebuild, + build_timeout=config["build_timeout"], + ) + print(f"[{idx:>2}/{total}] {name} OK ({tag})") + except Exception as exc: + print(f"[{idx:>2}/{total}] {name} FAIL ({str(exc)[:80]})") + sys.stdout.flush() + return [] + + results: list[TaskResult] = [] + for idx, td in enumerate(task_dirs, 1): + result = await evaluate_task( + td, + client, + timeout_override=timeout_override, + rebuild=rebuild, + skip_tests=skip_tests, + ) + results.append(result) + print_task_result(idx, total, result) + + return results + + +# ── CLI ────────────────────────────────────────────────────────────── + + +def main(): + parser = argparse.ArgumentParser( + description="SkillsBench Docker-based evaluation runner", + ) + parser.add_argument( + "--tasks-dir", + type=pathlib.Path, + default=_DEFAULT_TASKS_DIR, + help=( + "Path to SkillsBench tasks directory" + " (default: /tmp/skillsbench/tasks)" + ), + ) + parser.add_argument( + "--timeout", + type=float, + default=None, + help="Override per-task agent timeout in seconds", + ) + parser.add_argument( + "--filter", + type=str, + default=None, + help="Only run tasks whose name contains PATTERN", + ) + parser.add_argument( + "--rebuild", + action="store_true", + help="Force Docker image rebuild", + ) + parser.add_argument( + "--build-only", + action="store_true", + help="Build Docker images only, don't run agents", + ) + parser.add_argument( + "--skip-tests", + action="store_true", + help="Use tool-call heuristic instead of pytest scoring", + ) + args = parser.parse_args() + + logging.basicConfig(level=logging.WARNING) + + print_header() + + # Prereq check + client = ensure_docker_running() + + start = time.time() + results = asyncio.run( + run_full_evaluation( + tasks_dir=args.tasks_dir, + client=client, + timeout_override=args.timeout, + filter_pattern=args.filter, + rebuild=args.rebuild, + build_only=args.build_only, + skip_tests=args.skip_tests, + ) + ) + elapsed = time.time() - start + + if results: + print_summary(results, len(results), elapsed) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/skillsbench/metrics.py b/benchmarks/skillsbench/metrics.py new file mode 100644 index 0000000000..8477d9f42e --- /dev/null +++ b/benchmarks/skillsbench/metrics.py @@ -0,0 +1,226 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Custom metrics for SkillsBench evaluation. + +These metrics follow the ADK custom metric function signature: + + def metric_fn( + eval_metric: EvalMetric, + actual_invocations: list[Invocation], + expected_invocations: Optional[list[Invocation]], + conversation_scenario: Optional[ConversationScenario] = None, + ) -> EvaluationResult + +They can be referenced in eval configs via their dotted path, e.g.: + "benchmarks.skillsbench.metrics.skill_discovery_score" +""" + +from __future__ import annotations + +from typing import Optional + +from google.adk.evaluation.eval_case import ConversationScenario +from google.adk.evaluation.eval_case import get_all_tool_calls +from google.adk.evaluation.eval_case import Invocation +from google.adk.evaluation.eval_metrics import EvalMetric +from google.adk.evaluation.eval_metrics import EvalStatus +from google.adk.evaluation.evaluator import EvaluationResult +from google.adk.evaluation.evaluator import PerInvocationResult + + +def _get_tool_names_from_invocations( + invocations: list[Invocation], +) -> list[str]: + """Extract all tool call names from a list of invocations.""" + names = [] + for inv in invocations: + for tool_call in get_all_tool_calls(inv.intermediate_data): + names.append(tool_call.name) + return names + + +def skill_discovery_score( + eval_metric: EvalMetric, + actual_invocations: list[Invocation], + expected_invocations: Optional[list[Invocation]], + conversation_scenario: Optional[ConversationScenario] = None, +) -> EvaluationResult: + """Score 1.0 if the agent called both list_skills and load_skill. + + This metric checks whether the agent properly discovered skills + before attempting to use them, which is the expected workflow for + SkillsBench tasks. + """ + tool_names = _get_tool_names_from_invocations(actual_invocations) + called_list = any(name == "list_skills" for name in tool_names) + called_load = any(name == "load_skill" for name in tool_names) + + score = 1.0 if (called_list and called_load) else 0.0 + status = EvalStatus.PASSED if score >= 1.0 else EvalStatus.FAILED + + per_invocation = [] + for i, actual in enumerate(actual_invocations): + expected = None + if expected_invocations and i < len(expected_invocations): + expected = expected_invocations[i] + per_invocation.append( + PerInvocationResult( + actual_invocation=actual, + expected_invocation=expected, + score=score, + eval_status=status, + ) + ) + + return EvaluationResult( + overall_score=score, + overall_eval_status=status, + per_invocation_results=per_invocation, + ) + + +def tool_usage_score( + eval_metric: EvalMetric, + actual_invocations: list[Invocation], + expected_invocations: Optional[list[Invocation]], + conversation_scenario: Optional[ConversationScenario] = None, +) -> EvaluationResult: + """Fraction of expected tool calls that were actually made. + + Compares the set of tool names in expected_invocations against + actual_invocations. Score is |expected ∩ actual| / |expected|. + Uses ANY_ORDER matching — only checks that expected tools were + called, regardless of order or extra calls. + """ + if not expected_invocations: + return EvaluationResult( + overall_score=1.0, + overall_eval_status=EvalStatus.PASSED, + ) + + expected_names = set(_get_tool_names_from_invocations(expected_invocations)) + actual_names = set(_get_tool_names_from_invocations(actual_invocations)) + + if not expected_names: + score = 1.0 + else: + matched = expected_names & actual_names + score = len(matched) / len(expected_names) + + status = EvalStatus.PASSED if score >= 0.5 else EvalStatus.FAILED + + per_invocation = [] + for i, actual in enumerate(actual_invocations): + expected = None + if expected_invocations and i < len(expected_invocations): + expected = expected_invocations[i] + per_invocation.append( + PerInvocationResult( + actual_invocation=actual, + expected_invocation=expected, + score=score, + eval_status=status, + ) + ) + + return EvaluationResult( + overall_score=score, + overall_eval_status=status, + per_invocation_results=per_invocation, + ) + + +def skillsbench_binary_score( + eval_metric: EvalMetric, + actual_invocations: list[Invocation], + expected_invocations: Optional[list[Invocation]], + conversation_scenario: Optional[ConversationScenario] = None, +) -> EvaluationResult: + """Binary pass/fail: 1.0 if final response contains expected text. + + Mirrors the SkillsBench binary scoring methodology. Checks whether + key strings from the expected final response appear in the actual + final response. The match is case-insensitive and checks for + substring containment of each non-empty line in the reference. + """ + if not expected_invocations or not actual_invocations: + return EvaluationResult( + overall_score=0.0, + overall_eval_status=EvalStatus.NOT_EVALUATED, + ) + + # Get the last actual response text + actual_text = "" + for inv in reversed(actual_invocations): + if inv.final_response and inv.final_response.parts: + for part in inv.final_response.parts: + if part.text: + actual_text = part.text + break + if actual_text: + break + + # Get the expected response text + expected_text = "" + for inv in reversed(expected_invocations): + if inv.final_response and inv.final_response.parts: + for part in inv.final_response.parts: + if part.text: + expected_text = part.text + break + if expected_text: + break + + if not expected_text: + return EvaluationResult( + overall_score=0.0, + overall_eval_status=EvalStatus.NOT_EVALUATED, + ) + + # Check that each non-empty reference line appears in the actual + reference_lines = [ + line.strip() for line in expected_text.split("\n") if line.strip() + ] + actual_lower = actual_text.lower() + matched = sum(1 for line in reference_lines if line.lower() in actual_lower) + score = ( + 1.0 + if matched == len(reference_lines) + else matched / max(len(reference_lines), 1) + ) + + # Binary: pass only if all reference lines matched + is_pass = matched == len(reference_lines) and len(reference_lines) > 0 + status = EvalStatus.PASSED if is_pass else EvalStatus.FAILED + + per_invocation = [] + for i, actual in enumerate(actual_invocations): + expected = None + if expected_invocations and i < len(expected_invocations): + expected = expected_invocations[i] + per_invocation.append( + PerInvocationResult( + actual_invocation=actual, + expected_invocation=expected, + score=1.0 if is_pass else 0.0, + eval_status=status, + ) + ) + + return EvaluationResult( + overall_score=1.0 if is_pass else 0.0, + overall_eval_status=status, + per_invocation_results=per_invocation, + ) diff --git a/benchmarks/skillsbench/runner.py b/benchmarks/skillsbench/runner.py new file mode 100644 index 0000000000..8d6d547613 --- /dev/null +++ b/benchmarks/skillsbench/runner.py @@ -0,0 +1,313 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Standalone SkillsBench runner that produces leaderboard-compatible scores. + +Usage: + python benchmarks/skillsbench/runner.py + python benchmarks/skillsbench/runner.py --num-runs 3 + python benchmarks/skillsbench/runner.py --eval-set path/to/eval.json + +Environment variables: + GOOGLE_API_KEY — API key for authentication + GOOGLE_GENAI_USE_VERTEXAI — Set to 1 for Vertex AI backend + +This script: +1. Loads the SkillsBench agent and eval set +2. Runs each task through the ADK Runner directly +3. Applies all 3 custom metrics (discovery, tool usage, binary) +4. Outputs per-task results and a leaderboard-format summary +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import logging +import pathlib +import sys +import time +import uuid +from typing import Optional + +from google.adk.evaluation.eval_case import Invocation +from google.adk.evaluation.eval_metrics import EvalMetric +from google.adk.evaluation.eval_set import EvalSet +from google.adk.evaluation.evaluation_generator import EvaluationGenerator +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService +from google.adk.memory.in_memory_memory_service import InMemoryMemoryService +from google.adk.utils.context_utils import Aclosing +from google.genai import types as genai_types + +from .metrics import ( + skill_discovery_score, + skillsbench_binary_score, + tool_usage_score, +) + +logger = logging.getLogger(__name__) + +_BENCHMARKS_DIR = pathlib.Path(__file__).parent +_DEFAULT_EVAL_SET = _BENCHMARKS_DIR / "eval_sets" / "skillsbench_eval.json" +_MODEL_NAME = "gemini-3-flash-preview" + + +def load_eval_set(path: pathlib.Path) -> EvalSet: + """Load an EvalSet from a JSON file.""" + with open(path) as f: + data = json.load(f) + return EvalSet.model_validate(data) + + +def print_header(): + """Print the SkillsBench header.""" + print() + print("=" * 60) + print(" SkillsBench Evaluation — ADK SkillToolset") + print("=" * 60) + print() + + +def print_results_table(results: dict[str, dict[str, float]]): + """Print per-task results as a formatted table.""" + print(f"{'Task':<45} {'Discovery':>9} {'Tools':>7} {'Pass':>6}") + print("-" * 70) + for task_id, scores in results.items(): + short_id = task_id[:44] + disc = scores.get("skill_discovery", 0.0) + tools = scores.get("tool_usage", 0.0) + binary = scores.get("binary_pass", 0.0) + mark = "PASS" if binary >= 1.0 else "FAIL" + print(f"{short_id:<45} {disc:>8.1f} {tools:>7.2f} {mark:>6}") + print("-" * 70) + + +def print_leaderboard_summary( + results: dict[str, dict[str, float]], + num_tasks: int, + elapsed: float, +): + """Print a leaderboard-format summary.""" + passed = sum( + 1 for scores in results.values() if scores.get("binary_pass", 0.0) >= 1.0 + ) + avg_discovery = sum( + s.get("skill_discovery", 0.0) for s in results.values() + ) / max(len(results), 1) + avg_tools = sum(s.get("tool_usage", 0.0) for s in results.values()) / max( + len(results), 1 + ) + pct = (passed / max(num_tasks, 1)) * 100 + + print() + print("=" * 60) + print(" Leaderboard Summary") + print("=" * 60) + print(f" Model: {_MODEL_NAME}") + print(f" Framework: ADK SkillToolset") + print(f" Tasks: {passed}/{num_tasks} ({pct:.1f}%)") + print(f" Avg Discovery: {avg_discovery:.2f}") + print(f" Avg Tool Usage: {avg_tools:.2f}") + print(f" Elapsed: {elapsed:.1f}s") + print("=" * 60) + + +async def run_single_eval_case( + root_agent, + eval_case, +) -> list[Invocation]: + """Run a single eval case through the Runner and return invocations.""" + session_service = InMemorySessionService() + artifact_service = InMemoryArtifactService() + memory_service = InMemoryMemoryService() + + app_name = "skillsbench_eval" + user_id = "eval_user" + session_id = str(uuid.uuid4()) + + await session_service.create_session( + app_name=app_name, + user_id=user_id, + state={}, + session_id=session_id, + ) + + async with Runner( + app_name=app_name, + agent=root_agent, + artifact_service=artifact_service, + session_service=session_service, + memory_service=memory_service, + ) as runner: + events = [] + for invocation in eval_case.conversation: + user_content = invocation.user_content + + async with Aclosing( + runner.run_async( + user_id=user_id, + session_id=session_id, + new_message=user_content, + ) + ) as agen: + invocation_id = None + async for event in agen: + if not invocation_id: + invocation_id = event.invocation_id + from google.adk.events.event import Event + + events.append( + Event( + content=user_content, + author="user", + invocation_id=invocation_id, + ) + ) + events.append(event) + + return EvaluationGenerator.convert_events_to_eval_invocations(events) + + +def score_invocations( + actual_invocations: list[Invocation], + expected_invocations: Optional[list[Invocation]], +) -> dict[str, float]: + """Apply all 3 metrics and return scores.""" + metric = EvalMetric(metric_name="skillsbench") + scores = {} + + result = skill_discovery_score( + metric, actual_invocations, expected_invocations + ) + scores["skill_discovery"] = result.overall_score or 0.0 + + result = tool_usage_score(metric, actual_invocations, expected_invocations) + scores["tool_usage"] = result.overall_score or 0.0 + + result = skillsbench_binary_score( + metric, actual_invocations, expected_invocations + ) + scores["binary_pass"] = result.overall_score or 0.0 + + return scores + + +async def run_evaluation( + eval_set_path: Optional[pathlib.Path] = None, + num_runs: int = 1, +) -> dict[str, dict[str, float]]: + """Run the full SkillsBench evaluation.""" + path = eval_set_path or _DEFAULT_EVAL_SET + eval_set = load_eval_set(path) + + # Import agent (triggers skill loading) + from .agent import root_agent + + results: dict[str, dict[str, float]] = {} + total = len(eval_set.eval_cases) + + for idx, eval_case in enumerate(eval_set.eval_cases, 1): + eval_id = eval_case.eval_id + print(f"\n[{idx}/{total}] Running: {eval_id}") + + run_scores: list[dict[str, float]] = [] + for run in range(num_runs): + if num_runs > 1: + print(f" Run {run + 1}/{num_runs}...") + + try: + actual_invocations = await run_single_eval_case(root_agent, eval_case) + + # Print what the agent did + for inv in actual_invocations: + if inv.final_response and inv.final_response.parts: + for part in inv.final_response.parts: + if part.text: + preview = part.text[:200].replace("\n", " ") + print(f" Response: {preview}...") + break + + expected_invocations = eval_case.conversation + scores = score_invocations(actual_invocations, expected_invocations) + run_scores.append(scores) + + disc = scores["skill_discovery"] + tools = scores["tool_usage"] + binary = scores["binary_pass"] + mark = "PASS" if binary >= 1.0 else "FAIL" + print(f" Scores: discovery={disc:.1f} tools={tools:.2f} binary={mark}") + + except Exception as e: + logger.error("Error running %s: %s", eval_id, e) + print(f" ERROR: {e}") + run_scores.append({ + "skill_discovery": 0.0, + "tool_usage": 0.0, + "binary_pass": 0.0, + }) + + # Average scores across runs + avg_scores: dict[str, float] = {} + for key in ["skill_discovery", "tool_usage", "binary_pass"]: + values = [s[key] for s in run_scores] + avg_scores[key] = sum(values) / len(values) + results[eval_id] = avg_scores + + return results + + +def main(): + parser = argparse.ArgumentParser( + description="SkillsBench evaluation runner for ADK" + ) + parser.add_argument( + "--eval-set", + type=pathlib.Path, + default=None, + help="Path to eval set JSON (default: built-in)", + ) + parser.add_argument( + "--num-runs", + type=int, + default=1, + help="Number of runs per task (default: 1)", + ) + args = parser.parse_args() + + logging.basicConfig(level=logging.WARNING) + + print_header() + start = time.time() + + results = asyncio.run( + run_evaluation( + eval_set_path=args.eval_set, + num_runs=args.num_runs, + ) + ) + + elapsed = time.time() - start + eval_path = args.eval_set or _DEFAULT_EVAL_SET + eval_set = load_eval_set(eval_path) + + print() + print_results_table(results) + print_leaderboard_summary(results, len(eval_set.eval_cases), elapsed) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/skillsbench/skills/csv-aggregation/SKILL.md b/benchmarks/skillsbench/skills/csv-aggregation/SKILL.md new file mode 100644 index 0000000000..dd6b82e157 --- /dev/null +++ b/benchmarks/skillsbench/skills/csv-aggregation/SKILL.md @@ -0,0 +1,41 @@ +--- +name: csv-aggregation +description: Aggregate and summarize CSV data by computing statistics on specified columns. +--- + +# CSV Aggregation Skill + +Analyze CSV data by computing aggregate statistics (sum, mean, min, max, count) grouped by a specified column. + +## Available Scripts + +### `aggregate.py` + +Reads CSV data from stdin and computes aggregate statistics. + +**Usage**: `execute_skill_script(skill_name="csv-aggregation", script_name="aggregate.py", input_args="group_col=department metric_col=salary")` + +The script expects CSV data piped via stdin or provided as the `data` argument. Pass column names as arguments: +- `group_col`: The column to group by +- `metric_col`: The column to compute statistics on + +**Output format**: +``` +Group: + count: + sum: + mean: + min: + max: +``` + +## References + +- [sample-data.md](./references/sample-data.md) — Sample CSV dataset for testing + +## Workflow + +1. Use `load_skill` to read these instructions. +2. Use `load_skill_resource` to load the sample data reference. +3. Use `execute_skill_script` with the appropriate arguments to aggregate the data. +4. Present the aggregated results to the user. diff --git a/benchmarks/skillsbench/skills/csv-aggregation/references/sample-data.md b/benchmarks/skillsbench/skills/csv-aggregation/references/sample-data.md new file mode 100644 index 0000000000..4e2ce51bbc --- /dev/null +++ b/benchmarks/skillsbench/skills/csv-aggregation/references/sample-data.md @@ -0,0 +1,15 @@ +# Sample CSV Data + +Use this data as input for the aggregation script: + +```csv +name,department,salary,years +Alice,Engineering,95000,5 +Bob,Marketing,72000,3 +Carol,Engineering,102000,8 +Dave,Marketing,68000,2 +Eve,Engineering,88000,4 +Frank,Sales,76000,6 +Grace,Sales,81000,7 +Hank,Marketing,74000,4 +``` diff --git a/benchmarks/skillsbench/skills/csv-aggregation/scripts/aggregate.py b/benchmarks/skillsbench/skills/csv-aggregation/scripts/aggregate.py new file mode 100644 index 0000000000..fb073f5539 --- /dev/null +++ b/benchmarks/skillsbench/skills/csv-aggregation/scripts/aggregate.py @@ -0,0 +1,50 @@ +"""Aggregate CSV data by a grouping column.""" + +import csv +import io +import sys + + +def parse_args(args): + params = {} + for arg in args: + if "=" in arg: + key, value = arg.split("=", 1) + params[key] = value + return params + + +def main(): + params = parse_args(sys.argv[1:]) + group_col = params.get("group_col", "department") + metric_col = params.get("metric_col", "salary") + + data = """name,department,salary,years +Alice,Engineering,95000,5 +Bob,Marketing,72000,3 +Carol,Engineering,102000,8 +Dave,Marketing,68000,2 +Eve,Engineering,88000,4 +Frank,Sales,76000,6 +Grace,Sales,81000,7 +Hank,Marketing,74000,4""" + + reader = csv.DictReader(io.StringIO(data)) + groups = {} + for row in reader: + group = row[group_col] + value = float(row[metric_col]) + groups.setdefault(group, []).append(value) + + for group_name in sorted(groups): + values = groups[group_name] + print(f"Group: {group_name}") + print(f" count: {len(values)}") + print(f" sum: {sum(values):.0f}") + print(f" mean: {sum(values) / len(values):.0f}") + print(f" min: {min(values):.0f}") + print(f" max: {max(values):.0f}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/skillsbench/skills/function-scaffold/SKILL.md b/benchmarks/skillsbench/skills/function-scaffold/SKILL.md new file mode 100644 index 0000000000..40a2887089 --- /dev/null +++ b/benchmarks/skillsbench/skills/function-scaffold/SKILL.md @@ -0,0 +1,30 @@ +--- +name: function-scaffold +description: Generate Python function scaffolds with type hints and docstrings from a specification. +--- + +# Function Scaffold Skill + +Generate Python function stubs from a natural-language specification, including type hints, docstrings, and placeholder implementations. + +## Available Scripts + +### `scaffold.py` + +Generates a Python function scaffold from a specification. + +**Usage**: `execute_skill_script(skill_name="function-scaffold", script_name="scaffold.py", input_args="name=calculate_bmi params=weight:float,height:float returns=float description='Calculate Body Mass Index'")` + +Arguments: +- `name`: Function name (snake_case) +- `params`: Comma-separated parameter list with types (e.g., `x:int,y:str`) +- `returns`: Return type annotation +- `description`: One-line description for the docstring + +**Output format**: Complete Python function with type hints and docstring. + +## Workflow + +1. Use `load_skill` to read these instructions. +2. Use `execute_skill_script` with the function specification. +3. Present the generated scaffold to the user. diff --git a/benchmarks/skillsbench/skills/function-scaffold/scripts/scaffold.py b/benchmarks/skillsbench/skills/function-scaffold/scripts/scaffold.py new file mode 100644 index 0000000000..f9df8f2c01 --- /dev/null +++ b/benchmarks/skillsbench/skills/function-scaffold/scripts/scaffold.py @@ -0,0 +1,66 @@ +"""Generate Python function scaffolds from specifications.""" + +import sys + + +def parse_args(args): + params = {} + current_key = None + current_val = [] + for arg in args: + if "=" in arg and not current_key: + key, value = arg.split("=", 1) + if value.startswith("'") and not value.endswith("'"): + current_key = key + current_val = [value[1:]] + elif value.startswith("'") and value.endswith("'"): + params[key] = value[1:-1] + else: + params[key] = value + elif current_key: + if arg.endswith("'"): + current_val.append(arg[:-1]) + params[current_key] = " ".join(current_val) + current_key = None + current_val = [] + else: + current_val.append(arg) + if current_key: + params[current_key] = " ".join(current_val) + return params + + +def main(): + params = parse_args(sys.argv[1:]) + func_name = params.get("name", "my_function") + func_params = params.get("params", "") + returns = params.get("returns", "None") + description = params.get("description", "TODO: Add description") + + param_list = [] + if func_params: + for p in func_params.split(","): + if ":" in p: + pname, ptype = p.split(":", 1) + param_list.append(f"{pname}: {ptype}") + else: + param_list.append(p) + + params_str = ", ".join(param_list) + print(f"def {func_name}({params_str}) -> {returns}:") + print(f' """{description}') + print() + if param_list: + print(" Args:") + for p in param_list: + pname = p.split(":")[0].strip() + print(f" {pname}: TODO") + print() + print(" Returns:") + print(f" {returns}: TODO") + print(' """') + print(" raise NotImplementedError") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/skillsbench/skills/html-extraction/SKILL.md b/benchmarks/skillsbench/skills/html-extraction/SKILL.md new file mode 100644 index 0000000000..9204f96b5d --- /dev/null +++ b/benchmarks/skillsbench/skills/html-extraction/SKILL.md @@ -0,0 +1,35 @@ +--- +name: html-extraction +description: Extract structured data from HTML content using CSS-like selectors. +--- + +# HTML Extraction Skill + +Parse HTML content and extract text, links, or table data using tag-based selectors. + +## Available Scripts + +### `extract.py` + +Extracts content from embedded sample HTML based on a target selector. + +**Usage**: `execute_skill_script(skill_name="html-extraction", script_name="extract.py", input_args="target=links")` + +Supported targets: +- `target=links`: Extract all hyperlinks (text and href) +- `target=headings`: Extract all heading text +- `target=table`: Extract table data as CSV +- `target=text`: Extract all visible text content + +**Output format**: One extracted item per line. + +## References + +- [sample-page.md](./references/sample-page.md) — Sample HTML page for testing + +## Workflow + +1. Use `load_skill` to read these instructions. +2. Use `load_skill_resource` to see the sample HTML page. +3. Use `execute_skill_script` with the desired target to extract data. +4. Present the extracted content to the user. diff --git a/benchmarks/skillsbench/skills/html-extraction/references/sample-page.md b/benchmarks/skillsbench/skills/html-extraction/references/sample-page.md new file mode 100644 index 0000000000..75e0a58b52 --- /dev/null +++ b/benchmarks/skillsbench/skills/html-extraction/references/sample-page.md @@ -0,0 +1,23 @@ +# Sample HTML Page + +The extraction script processes this embedded HTML: + +```html + +Product Catalog + +

Products

+

Electronics

+ + + + + +
ProductPriceStock
Laptop999.9915
Phone699.9942
Tablet449.9928
+

Links

+ View Laptop + View Phone + View Tablet + + +``` diff --git a/benchmarks/skillsbench/skills/html-extraction/scripts/extract.py b/benchmarks/skillsbench/skills/html-extraction/scripts/extract.py new file mode 100644 index 0000000000..2aab5c48aa --- /dev/null +++ b/benchmarks/skillsbench/skills/html-extraction/scripts/extract.py @@ -0,0 +1,72 @@ +"""Extract structured data from HTML content.""" + +import re +import sys + +SAMPLE_HTML = """ +Product Catalog + +

Products

+

Electronics

+ + + + + +
ProductPriceStock
Laptop999.9915
Phone699.9942
Tablet449.9928
+

Links

+ View Laptop + View Phone + View Tablet + +""" + + +def extract_links(html): + pattern = r'(.*?)' + for href, text in re.findall(pattern, html): + print(f"{text} -> {href}") + + +def extract_headings(html): + pattern = r"(.*?)" + for heading in re.findall(pattern, html): + print(heading) + + +def extract_table(html): + rows = re.findall(r"(.*?)", html, re.DOTALL) + for row in rows: + cells = re.findall(r"(.*?)", row) + print(",".join(cells)) + + +def extract_text(html): + clean = re.sub(r"<[^>]+>", " ", html) + clean = re.sub(r"\s+", " ", clean).strip() + print(clean) + + +def main(): + target = "links" + for arg in sys.argv[1:]: + if arg.startswith("target="): + target = arg.split("=", 1)[1] + + extractors = { + "links": extract_links, + "headings": extract_headings, + "table": extract_table, + "text": extract_text, + } + + if target not in extractors: + print(f"Error: unknown target '{target}'") + print(f"Available: {', '.join(extractors)}") + sys.exit(1) + + extractors[target](SAMPLE_HTML) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/skillsbench/skills/json-transform/SKILL.md b/benchmarks/skillsbench/skills/json-transform/SKILL.md new file mode 100644 index 0000000000..f1eba8dd92 --- /dev/null +++ b/benchmarks/skillsbench/skills/json-transform/SKILL.md @@ -0,0 +1,34 @@ +--- +name: json-transform +description: Transform JSON data by flattening nested structures and renaming fields. +--- + +# JSON Transform Skill + +Transform JSON objects by flattening nested structures, renaming keys, and filtering fields. + +## Available Scripts + +### `transform.py` + +Transforms a JSON object according to a field mapping specification. + +**Usage**: `execute_skill_script(skill_name="json-transform", script_name="transform.py", input_args="flatten=true")` + +The script reads the embedded sample data and applies transformations: +- `flatten=true`: Flatten nested objects using dot notation +- `rename=old:new`: Rename a field (can be specified multiple times) +- `keep=field1,field2`: Keep only specified fields + +**Output format**: Pretty-printed JSON + +## References + +- [sample-data.md](./references/sample-data.md) — Sample nested JSON for testing + +## Workflow + +1. Use `load_skill` to read these instructions. +2. Use `load_skill_resource` to examine the sample data. +3. Use `execute_skill_script` to transform the data. +4. Present the transformed JSON to the user. diff --git a/benchmarks/skillsbench/skills/json-transform/references/sample-data.md b/benchmarks/skillsbench/skills/json-transform/references/sample-data.md new file mode 100644 index 0000000000..24b99bb140 --- /dev/null +++ b/benchmarks/skillsbench/skills/json-transform/references/sample-data.md @@ -0,0 +1,21 @@ +# Sample JSON Data + +Use this nested JSON as input for the transform script: + +```json +{ + "user": { + "name": "Alice", + "age": 30, + "address": { + "city": "Seattle", + "state": "WA", + "zip": "98101" + } + }, + "orders": [ + {"id": 1, "total": 42.50}, + {"id": 2, "total": 18.75} + ] +} +``` diff --git a/benchmarks/skillsbench/skills/json-transform/scripts/transform.py b/benchmarks/skillsbench/skills/json-transform/scripts/transform.py new file mode 100644 index 0000000000..323f06ea49 --- /dev/null +++ b/benchmarks/skillsbench/skills/json-transform/scripts/transform.py @@ -0,0 +1,63 @@ +"""Transform JSON by flattening nested structures.""" + +import json +import sys + +SAMPLE_DATA = { + "user": { + "name": "Alice", + "age": 30, + "address": { + "city": "Seattle", + "state": "WA", + "zip": "98101", + }, + }, + "orders": [ + {"id": 1, "total": 42.50}, + {"id": 2, "total": 18.75}, + ], +} + + +def flatten(obj, prefix=""): + result = {} + for key, value in obj.items(): + full_key = f"{prefix}.{key}" if prefix else key + if isinstance(value, dict): + result.update(flatten(value, full_key)) + else: + result[full_key] = value + return result + + +def parse_args(args): + params = {} + for arg in args: + if "=" in arg: + key, value = arg.split("=", 1) + params[key] = value + return params + + +def main(): + params = parse_args(sys.argv[1:]) + data = SAMPLE_DATA.copy() + + if params.get("flatten", "").lower() == "true": + data = flatten(data) + + if "keep" in params: + keep_fields = set(params["keep"].split(",")) + data = {k: v for k, v in data.items() if k in keep_fields} + + if "rename" in params: + old_name, new_name = params["rename"].split(":", 1) + if old_name in data: + data[new_name] = data.pop(old_name) + + print(json.dumps(data, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/skillsbench/skills/log-parsing/SKILL.md b/benchmarks/skillsbench/skills/log-parsing/SKILL.md new file mode 100644 index 0000000000..243878a29c --- /dev/null +++ b/benchmarks/skillsbench/skills/log-parsing/SKILL.md @@ -0,0 +1,48 @@ +--- +name: log-parsing +description: Parse and analyze structured log files to extract error patterns and statistics. +metadata: + category: system-admin + aliases: log-analyzer,syslog-parser +--- + +# Log Parsing Skill + +Parse structured log files to extract error counts, warning summaries, and time-based patterns. Supports common log formats. + +## Available Scripts + +### `parse.py` + +Analyzes embedded sample log data and produces a summary report. + +**Usage**: `execute_skill_script(skill_name="log-parsing", script_name="parse.py", input_args="level=ERROR")` + +Arguments: +- `level`: Filter by log level (DEBUG, INFO, WARNING, ERROR, ALL). Default: ALL +- `format`: Output format (summary, detail, timeline). Default: summary + +**Output format** (summary): +``` +Log Analysis Report +=================== +Total entries: +ERROR: +WARNING: +INFO: +DEBUG: +``` + +**Output format** (detail): +Lists each matching log entry with timestamp and message. + +## References + +- [sample-logs.md](./references/sample-logs.md) — Sample log data for testing + +## Workflow + +1. Use `load_skill` to read these instructions. +2. Optionally use `load_skill_resource` to examine the sample log data. +3. Use `execute_skill_script` with the desired log level filter. +4. Present the analysis report to the user. diff --git a/benchmarks/skillsbench/skills/log-parsing/references/sample-logs.md b/benchmarks/skillsbench/skills/log-parsing/references/sample-logs.md new file mode 100644 index 0000000000..90163597a0 --- /dev/null +++ b/benchmarks/skillsbench/skills/log-parsing/references/sample-logs.md @@ -0,0 +1,18 @@ +# Sample Log Data + +The parsing script processes these embedded log entries: + +``` +2024-01-15 10:00:01 INFO Server started on port 8080 +2024-01-15 10:00:05 DEBUG Database connection pool initialized +2024-01-15 10:01:12 INFO User alice logged in +2024-01-15 10:02:33 WARNING High memory usage: 85% +2024-01-15 10:03:45 ERROR Failed to process request: timeout +2024-01-15 10:04:01 INFO User bob logged in +2024-01-15 10:05:22 ERROR Database query failed: connection reset +2024-01-15 10:06:10 WARNING Disk space low: 10% remaining +2024-01-15 10:07:00 INFO Scheduled backup started +2024-01-15 10:08:15 ERROR Backup failed: permission denied +2024-01-15 10:09:30 INFO User alice logged out +2024-01-15 10:10:00 DEBUG Cache cleared +``` diff --git a/benchmarks/skillsbench/skills/log-parsing/scripts/parse.py b/benchmarks/skillsbench/skills/log-parsing/scripts/parse.py new file mode 100644 index 0000000000..f7c658fedb --- /dev/null +++ b/benchmarks/skillsbench/skills/log-parsing/scripts/parse.py @@ -0,0 +1,79 @@ +"""Parse and analyze structured log files.""" + +import sys + +SAMPLE_LOGS = """2024-01-15 10:00:01 INFO Server started on port 8080 +2024-01-15 10:00:05 DEBUG Database connection pool initialized +2024-01-15 10:01:12 INFO User alice logged in +2024-01-15 10:02:33 WARNING High memory usage: 85% +2024-01-15 10:03:45 ERROR Failed to process request: timeout +2024-01-15 10:04:01 INFO User bob logged in +2024-01-15 10:05:22 ERROR Database query failed: connection reset +2024-01-15 10:06:10 WARNING Disk space low: 10% remaining +2024-01-15 10:07:00 INFO Scheduled backup started +2024-01-15 10:08:15 ERROR Backup failed: permission denied +2024-01-15 10:09:30 INFO User alice logged out +2024-01-15 10:10:00 DEBUG Cache cleared""" + + +def parse_entry(line): + parts = line.split(None, 3) + if len(parts) < 4: + return None + date, time, level, message = parts + return { + "timestamp": f"{date} {time}", + "level": level.strip(), + "message": message.strip(), + } + + +def parse_args(args): + params = {} + for arg in args: + if "=" in arg: + key, value = arg.split("=", 1) + params[key] = value + return params + + +def main(): + params = parse_args(sys.argv[1:]) + level_filter = params.get("level", "ALL").upper() + out_format = params.get("format", "summary") + + entries = [] + for line in SAMPLE_LOGS.strip().split("\n"): + entry = parse_entry(line) + if entry: + entries.append(entry) + + if level_filter != "ALL": + filtered = [e for e in entries if e["level"] == level_filter] + else: + filtered = entries + + if out_format == "summary": + counts = {} + for entry in entries: + level = entry["level"] + counts[level] = counts.get(level, 0) + 1 + + print("Log Analysis Report") + print("===================") + print(f"Total entries: {len(entries)}") + for level in ["ERROR", "WARNING", "INFO", "DEBUG"]: + print(f"{level}: {counts.get(level, 0)}") + + elif out_format == "detail": + for entry in filtered: + print(f"[{entry['timestamp']}] {entry['level']} - {entry['message']}") + + elif out_format == "timeline": + for entry in filtered: + time_part = entry["timestamp"].split(" ")[1] + print(f"{time_part} {entry['level']}: {entry['message']}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/skillsbench/skills/regex-replace/SKILL.md b/benchmarks/skillsbench/skills/regex-replace/SKILL.md new file mode 100644 index 0000000000..eb2ebaa726 --- /dev/null +++ b/benchmarks/skillsbench/skills/regex-replace/SKILL.md @@ -0,0 +1,36 @@ +--- +name: regex-replace +description: Perform regex-based find-and-replace operations on text input. +--- + +# Regex Replace Skill + +Apply regular expression patterns to find and replace text. Supports basic and advanced regex syntax. + +## Available Scripts + +### `replace.py` + +Performs regex find-and-replace on input text. + +**Usage**: `execute_skill_script(skill_name="regex-replace", script_name="replace.py", input_args="pattern=\\d+ replacement=NUM text='Order 123 has 45 items at $67'")` + +Arguments: +- `pattern`: The regex pattern to match +- `replacement`: The replacement string +- `text`: The input text to process +- `count`: Maximum replacements (default: all) + +**Output format**: +``` +Original: +Pattern: +Result: +Matches: +``` + +## Workflow + +1. Use `load_skill` to read these instructions. +2. Use `execute_skill_script` with pattern, replacement, and text arguments. +3. Present the transformation result to the user. diff --git a/benchmarks/skillsbench/skills/regex-replace/scripts/replace.py b/benchmarks/skillsbench/skills/regex-replace/scripts/replace.py new file mode 100644 index 0000000000..4e4182bb7c --- /dev/null +++ b/benchmarks/skillsbench/skills/regex-replace/scripts/replace.py @@ -0,0 +1,54 @@ +"""Perform regex find-and-replace on text.""" + +import re +import sys + + +def parse_args(args): + params = {} + current_key = None + current_val = [] + for arg in args: + if "=" in arg and not current_key: + key, value = arg.split("=", 1) + if value.startswith("'") and not value.endswith("'"): + current_key = key + current_val = [value[1:]] + elif value.startswith("'") and value.endswith("'"): + params[key] = value[1:-1] + else: + params[key] = value + elif current_key: + if arg.endswith("'"): + current_val.append(arg[:-1]) + params[current_key] = " ".join(current_val) + current_key = None + current_val = [] + else: + current_val.append(arg) + if current_key: + params[current_key] = " ".join(current_val) + return params + + +def main(): + params = parse_args(sys.argv[1:]) + pattern = params.get("pattern", r"\d+") + replacement = params.get("replacement", "NUM") + text = params.get("text", "Order 123 has 45 items at $67") + count = int(params.get("count", "0")) + + matches = re.findall(pattern, text) + if count > 0: + result = re.sub(pattern, replacement, text, count=count) + else: + result = re.sub(pattern, replacement, text) + + print(f"Original: {text}") + print(f"Pattern: {pattern}") + print(f"Result: {result}") + print(f"Matches: {len(matches)}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/skillsbench/skills/rest-client/SKILL.md b/benchmarks/skillsbench/skills/rest-client/SKILL.md new file mode 100644 index 0000000000..1073544736 --- /dev/null +++ b/benchmarks/skillsbench/skills/rest-client/SKILL.md @@ -0,0 +1,40 @@ +--- +name: rest-client +description: Simulate REST API interactions by constructing and executing HTTP-like requests. +--- + +# REST Client Skill + +Build and execute simulated REST API requests against an embedded mock API. Demonstrates multi-step skill usage with request construction and response parsing. + +## Available Scripts + +### `request.py` + +Executes a simulated REST API request against a mock endpoint. + +**Usage**: `execute_skill_script(skill_name="rest-client", script_name="request.py", input_args="method=GET endpoint=/users")` + +Supported arguments: +- `method`: HTTP method (GET, POST, PUT, DELETE) +- `endpoint`: API path (e.g., `/users`, `/users/1`, `/products`) +- `body`: JSON body for POST/PUT requests + +Available mock endpoints: +- `GET /users` — List all users +- `GET /users/` — Get user by ID +- `POST /users` — Create a user (requires `body`) +- `GET /products` — List all products + +**Output format**: JSON response with status code + +## References + +- [api-docs.md](./references/api-docs.md) — Mock API documentation + +## Workflow + +1. Use `load_skill` to read these instructions. +2. Use `load_skill_resource` to review the API documentation. +3. Use `execute_skill_script` to make API requests. +4. Parse the response and present the data to the user. diff --git a/benchmarks/skillsbench/skills/rest-client/references/api-docs.md b/benchmarks/skillsbench/skills/rest-client/references/api-docs.md new file mode 100644 index 0000000000..75bf900276 --- /dev/null +++ b/benchmarks/skillsbench/skills/rest-client/references/api-docs.md @@ -0,0 +1,36 @@ +# Mock API Documentation + +## Base URL + +All endpoints are relative to the mock server. + +## Endpoints + +### GET /users +Returns a list of all users. + +**Response**: +```json +[ + {"id": 1, "name": "Alice", "email": "alice@example.com"}, + {"id": 2, "name": "Bob", "email": "bob@example.com"}, + {"id": 3, "name": "Carol", "email": "carol@example.com"} +] +``` + +### GET /users/:id +Returns a single user by ID. + +### POST /users +Creates a new user. Requires a JSON body with `name` and `email` fields. + +### GET /products +Returns a list of all products. + +**Response**: +```json +[ + {"id": 1, "name": "Laptop", "price": 999.99}, + {"id": 2, "name": "Phone", "price": 699.99} +] +``` diff --git a/benchmarks/skillsbench/skills/rest-client/scripts/request.py b/benchmarks/skillsbench/skills/rest-client/scripts/request.py new file mode 100644 index 0000000000..dd22bc9c84 --- /dev/null +++ b/benchmarks/skillsbench/skills/rest-client/scripts/request.py @@ -0,0 +1,65 @@ +"""Simulate REST API requests against a mock server.""" + +import json +import sys + +MOCK_DB = { + "users": [ + {"id": 1, "name": "Alice", "email": "alice@example.com"}, + {"id": 2, "name": "Bob", "email": "bob@example.com"}, + {"id": 3, "name": "Carol", "email": "carol@example.com"}, + ], + "products": [ + {"id": 1, "name": "Laptop", "price": 999.99}, + {"id": 2, "name": "Phone", "price": 699.99}, + ], +} + + +def parse_args(args): + params = {} + for arg in args: + if "=" in arg: + key, value = arg.split("=", 1) + params[key] = value + return params + + +def handle_request(method, endpoint, body=None): + method = method.upper() + + if endpoint == "/users" and method == "GET": + return 200, MOCK_DB["users"] + + if endpoint.startswith("/users/") and method == "GET": + user_id = int(endpoint.split("/")[-1]) + for user in MOCK_DB["users"]: + if user["id"] == user_id: + return 200, user + return 404, {"error": "User not found"} + + if endpoint == "/users" and method == "POST": + if body: + new_user = json.loads(body) + new_user["id"] = len(MOCK_DB["users"]) + 1 + return 201, new_user + return 400, {"error": "Request body required"} + + if endpoint == "/products" and method == "GET": + return 200, MOCK_DB["products"] + + return 404, {"error": f"Unknown endpoint: {method} {endpoint}"} + + +def main(): + params = parse_args(sys.argv[1:]) + method = params.get("method", "GET") + endpoint = params.get("endpoint", "/users") + body = params.get("body") + + status, response = handle_request(method, endpoint, body) + print(json.dumps({"status": status, "data": response}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/skillsbench/skills/statistical-calc/SKILL.md b/benchmarks/skillsbench/skills/statistical-calc/SKILL.md new file mode 100644 index 0000000000..763280c6f4 --- /dev/null +++ b/benchmarks/skillsbench/skills/statistical-calc/SKILL.md @@ -0,0 +1,38 @@ +--- +name: statistical-calc +description: Compute descriptive statistics (mean, median, std dev, percentiles) for numeric datasets. +--- + +# Statistical Calc Skill + +Compute descriptive statistics on numeric data including mean, median, standard deviation, variance, and percentiles. + +## Available Scripts + +### `stats.py` + +Computes descriptive statistics for a list of numbers. + +**Usage**: `execute_skill_script(skill_name="statistical-calc", script_name="stats.py", input_args="data=10,20,30,40,50,60,70,80,90,100")` + +Arguments: +- `data`: Comma-separated list of numbers + +**Output format**: +``` +Count: +Mean: +Median: +Std Dev: +Variance: +Min: +Max: +P25: <25th percentile> +P75: <75th percentile> +``` + +## Workflow + +1. Use `load_skill` to read these instructions. +2. Use `execute_skill_script` with numeric data to compute statistics. +3. Present the statistics to the user. diff --git a/benchmarks/skillsbench/skills/statistical-calc/scripts/stats.py b/benchmarks/skillsbench/skills/statistical-calc/scripts/stats.py new file mode 100644 index 0000000000..f6d893dc92 --- /dev/null +++ b/benchmarks/skillsbench/skills/statistical-calc/scripts/stats.py @@ -0,0 +1,54 @@ +"""Compute descriptive statistics for numeric data.""" + +import math +import sys + + +def parse_args(args): + params = {} + for arg in args: + if "=" in arg: + key, value = arg.split("=", 1) + params[key] = value + return params + + +def percentile(sorted_data, p): + k = (len(sorted_data) - 1) * (p / 100.0) + f = math.floor(k) + c = math.ceil(k) + if f == c: + return sorted_data[int(k)] + return sorted_data[int(f)] * (c - k) + sorted_data[int(c)] * (k - f) + + +def main(): + params = parse_args(sys.argv[1:]) + data_str = params.get("data", "10,20,30,40,50,60,70,80,90,100") + data = [float(x.strip()) for x in data_str.split(",")] + + n = len(data) + mean = sum(data) / n + sorted_data = sorted(data) + + if n % 2 == 0: + median = (sorted_data[n // 2 - 1] + sorted_data[n // 2]) / 2 + else: + median = sorted_data[n // 2] + + variance = sum((x - mean) ** 2 for x in data) / n + std_dev = math.sqrt(variance) + + print(f"Count: {n}") + print(f"Mean: {mean:.2f}") + print(f"Median: {median:.2f}") + print(f"Std Dev: {std_dev:.2f}") + print(f"Variance: {variance:.2f}") + print(f"Min: {min(data):.2f}") + print(f"Max: {max(data):.2f}") + print(f"P25: {percentile(sorted_data, 25):.2f}") + print(f"P75: {percentile(sorted_data, 75):.2f}") + + +if __name__ == "__main__": + main() From f9a78a6d4b6814fdf04bf8c5bfee36ed0bec30e2 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Mon, 23 Feb 2026 02:10:31 -0800 Subject: [PATCH 16/53] fix: Add per-command timeout to Docker executor Wrap all Docker exec_run calls in a ThreadPoolExecutor with a 300s default timeout. This prevents the asyncio event loop from blocking indefinitely when a script runs forever inside the container, which was causing the full evaluation to hang on long-running tasks. Co-Authored-By: Claude Opus 4.6 --- benchmarks/skillsbench/full_runner.py | 70 +++++++++++++++++++++------ 1 file changed, 55 insertions(+), 15 deletions(-) diff --git a/benchmarks/skillsbench/full_runner.py b/benchmarks/skillsbench/full_runner.py index 0371c7e590..b9dd6c2c83 100644 --- a/benchmarks/skillsbench/full_runner.py +++ b/benchmarks/skillsbench/full_runner.py @@ -34,6 +34,7 @@ import argparse import asyncio +import concurrent.futures import io import json import logging @@ -160,6 +161,7 @@ class TaskContainerExecutor(BaseCodeExecutor): _container: object = None _client: object = None + _cmd_timeout: float = 300.0 # Per-command timeout in seconds def start( self, @@ -176,6 +178,36 @@ def start( working_dir="/root", ) + def _exec_run_with_timeout( + self, + cmd: list[str], + workdir: str = "/root", + timeout: float | None = None, + ) -> tuple[int, bytes, bytes]: + """Run exec_run in a thread with a timeout. + + Returns (exit_code, stdout_bytes, stderr_bytes). + Raises TimeoutError if the command exceeds *timeout*. + """ + timeout = timeout or self._cmd_timeout + + def _run(): + rc, output = self._container.exec_run( + cmd, + workdir=workdir, + demux=True, + ) + return rc, (output[0] or b""), (output[1] or b"") + + with concurrent.futures.ThreadPoolExecutor(1) as pool: + future = pool.submit(_run) + try: + return future.result(timeout=timeout) + except concurrent.futures.TimeoutError: + raise TimeoutError( + f"Command timed out after {timeout}s: {' '.join(cmd[:3])}" + ) + @override def execute_code( self, @@ -189,13 +221,17 @@ def execute_code( stderr="Container not started", ) code = code_execution_input.code - rc, output = self._container.exec_run( - ["python3", "-c", code], - workdir="/root", - demux=True, - ) - stdout = (output[0] or b"").decode("utf-8", errors="replace") - stderr = (output[1] or b"").decode("utf-8", errors="replace") + try: + rc, out, err = self._exec_run_with_timeout( + ["python3", "-c", code], + ) + except TimeoutError as exc: + return CodeExecutionResult( + stdout="", + stderr=str(exc), + ) + stdout = out.decode("utf-8", errors="replace") + stderr = err.decode("utf-8", errors="replace") return CodeExecutionResult( stdout=stdout, stderr=stderr, @@ -228,13 +264,16 @@ def exec_in_container( """ if self._container is None: raise RuntimeError("Container not started") - rc, output = self._container.exec_run( - cmd, - workdir=workdir, - demux=True, - ) - stdout = (output[0] or b"").decode("utf-8", errors="replace") - stderr = (output[1] or b"").decode("utf-8", errors="replace") + try: + rc, out, err = self._exec_run_with_timeout( + cmd, + workdir=workdir, + timeout=timeout, + ) + except TimeoutError: + return -1, f"Command timed out after {timeout or self._cmd_timeout}s" + stdout = out.decode("utf-8", errors="replace") + stderr = err.decode("utf-8", errors="replace") combined = stdout if stderr: combined = combined + "\n" + stderr if combined else stderr @@ -507,10 +546,11 @@ def score_task_in_container( # Make test.sh executable executor.exec_in_container(["chmod", "+x", "/tests/test.sh"]) - # Run test.sh + # Run test.sh (longer timeout — installs deps + runs pytest) rc, output = executor.exec_in_container( ["bash", "/tests/test.sh"], workdir="/root", + timeout=900.0, ) logger.debug( "test.sh for %s exited %d:\n%s", From 8169ce3516b57ed0b7c5c5a42f46b51ca5e822f1 Mon Sep 17 00:00:00 2001 From: Hai-Yuan Cao <2003072+caohy1988@users.noreply.github.com> Date: Mon, 23 Feb 2026 22:43:29 -0800 Subject: [PATCH 17/53] Rename ExecuteSkillScriptTool to RunSkillScriptTool --- src/google/adk/tools/skill_toolset.py | 283 ++++++++++++++++++-------- 1 file changed, 198 insertions(+), 85 deletions(-) diff --git a/src/google/adk/tools/skill_toolset.py b/src/google/adk/tools/skill_toolset.py index dd37b3c1ac..799ad6d54a 100644 --- a/src/google/adk/tools/skill_toolset.py +++ b/src/google/adk/tools/skill_toolset.py @@ -16,9 +16,9 @@ from __future__ import annotations +import asyncio import json import logging -import shlex from typing import Any from typing import Optional from typing import TYPE_CHECKING @@ -239,16 +239,13 @@ async def run_async( @experimental(FeatureName.SKILL_TOOLSET) -class ExecuteSkillScriptTool(BaseTool): +class RunSkillScriptTool(BaseTool): """Tool to execute scripts from a skill's scripts/ directory.""" def __init__(self, toolset: "SkillToolset"): super().__init__( - name="execute_skill_script", - description=( - "Executes a script from a skill's scripts/ directory" - " and returns its output." - ), + name="run_skill_script", + description="Executes a script from a skill's scripts/ directory.", ) self._toolset = toolset @@ -263,22 +260,22 @@ def _get_declaration(self) -> types.FunctionDeclaration | None: "type": "string", "description": "The name of the skill.", }, - "script_name": { + "script_path": { "type": "string", "description": ( - "The name of the script to execute (e.g.," - " 'setup.sh' or 'scripts/setup.sh')." + "The relative path to the script (e.g.," + " 'scripts/setup.py')." ), }, - "input_args": { - "type": "string", + "args": { + "type": "object", "description": ( - "Optional space-separated arguments to pass" - " to the script." + "Optional arguments to pass to the script as key-value" + " pairs." ), }, }, - "required": ["skill_name", "script_name"], + "required": ["skill_name", "script_path"], }, ) @@ -286,24 +283,20 @@ async def run_async( self, *, args: dict[str, Any], tool_context: ToolContext ) -> Any: skill_name = args.get("skill_name") - script_name = args.get("script_name") - input_args = args.get("input_args", "") + script_path = args.get("script_path") + script_args = args.get("args", {}) if not skill_name: return { "error": "Skill name is required.", "error_code": "MISSING_SKILL_NAME", } - if not script_name: + if not script_path: return { - "error": "Script name is required.", - "error_code": "MISSING_SCRIPT_NAME", + "error": "Script path is required.", + "error_code": "MISSING_SCRIPT_PATH", } - # Strip scripts/ prefix for consistency - if script_name.startswith("scripts/"): - script_name = script_name[len("scripts/") :] - skill = self._toolset._get_skill(skill_name) if not skill: return { @@ -311,10 +304,15 @@ async def run_async( "error_code": "SKILL_NOT_FOUND", } - script = skill.resources.get_script(script_name) + script = None + if script_path.startswith("scripts/"): + script = skill.resources.get_script(script_path[len("scripts/") :]) + else: + script = skill.resources.get_script(script_path) + if script is None: return { - "error": f"Script '{script_name}' not found in skill '{skill_name}'.", + "error": f"Script '{script_path}' not found in skill '{skill_name}'.", "error_code": "SCRIPT_NOT_FOUND", } @@ -333,23 +331,51 @@ async def run_async( "error_code": "NO_CODE_EXECUTOR", } - # Validate input_args early (before sending to code executor) - if input_args: - try: - shlex.split(input_args) - except ValueError as e: - return { - "error": f"Invalid input_args: {e}", - "error_code": "INVALID_INPUT_ARGS", - } + import os + + from ..code_executors.code_execution_utils import File + + input_files = [] - # Prepare code based on script extension - code = self._prepare_code(script_name, script.src, input_args) - is_shell = "." in script_name and script_name.rsplit(".", 1)[ + # Package ALL skill files for mounting + for ref_name in skill.resources.list_references(): + content = skill.resources.get_reference(ref_name) + if content: + input_files.append( + File( + name=os.path.basename(ref_name), + path=f"references/{ref_name}", + content=content, + ) + ) + for asset_name in skill.resources.list_assets(): + content = skill.resources.get_asset(asset_name) + if content: + input_files.append( + File( + name=os.path.basename(asset_name), + path=f"assets/{asset_name}", + content=content, + ) + ) + for scr_name in skill.resources.list_scripts(): + scr = skill.resources.get_script(scr_name) + if scr and scr.src: + input_files.append( + File( + name=os.path.basename(scr_name), + path=f"scripts/{scr_name}", + content=scr.src, + ) + ) + + # Prepare wrapper code + code = self._prepare_code(script_path, script_args) + is_shell = "." in script_path and script_path.rsplit(".", 1)[ -1 ].lower() in ("sh", "bash") if code is None: - ext = script_name.rsplit(".", 1)[-1] if "." in script_name else "" + ext = script_path.rsplit(".", 1)[-1] if "." in script_path else "" return { "error": ( f"Unsupported script type '.{ext}'. Supported" @@ -359,9 +385,14 @@ async def run_async( } try: - result = code_executor.execute_code( + result = await asyncio.to_thread( + code_executor.execute_code, tool_context._invocation_context, - CodeExecutionInput(code=code), + CodeExecutionInput( + code=code, + input_files=input_files, + working_dir=".", + ), ) stdout = result.stdout stderr = result.stderr @@ -386,94 +417,91 @@ async def run_async( status = "success" return { "skill_name": skill_name, - "script_name": script_name, + "script_path": script_path, "stdout": stdout, "stderr": stderr, "status": status, } except SystemExit as e: - # Scripts may call sys.exit(); intercept instead of letting - # it terminate the host process. exit_code = e.code if e.code is not None else 0 if exit_code == 0: - # sys.exit(0) or sys.exit() is a normal termination. return { "skill_name": skill_name, - "script_name": script_name, + "script_path": script_path, "stdout": "", "stderr": "", "status": "success", } logger.warning( "Script '%s' from skill '%s' called sys.exit(%s)", - script_name, + script_path, skill_name, exit_code, ) return { - "error": f"Script '{script_name}' exited with code {exit_code}.", + "error": f"Script '{script_path}' exited with code {exit_code}.", "error_code": "EXECUTION_ERROR", } - except Exception as e: + except Exception as e: # pylint: disable=broad-exception-caught logger.exception( "Error executing script '%s' from skill '%s'", - script_name, + script_path, skill_name, ) - # Keep the error message short for the LLM; full trace is logged above. short_msg = str(e) if len(short_msg) > 200: short_msg = short_msg[:200] + "..." return { - "error": f"Failed to execute script '{script_name}': {short_msg}", + "error": ( + f"Failed to execute script '{script_path}':\n{type(e).__name__}:" + f" {short_msg}" + ), "error_code": "EXECUTION_ERROR", } def _prepare_code( self, - script_name: str, - script_src: str, - input_args: str, + script_path: str, + script_args: dict[str, Any], ) -> str | None: """Prepares Python code to execute the script. Args: - script_name: The script filename. - script_src: The script source content. - input_args: Optional arguments string. + script_path: The script file path. + script_args: Optional dictionary of arguments. Returns: Python code string to execute, or None if unsupported type. """ ext = "" - if "." in script_name: - ext = script_name.rsplit(".", 1)[-1].lower() + if "." in script_path: + ext = script_path.rsplit(".", 1)[-1].lower() + + if not script_path.startswith("scripts/"): + script_path = f"scripts/{script_path}" if ext == "py": - # Python script: execute directly, inject sys.argv if args - if input_args: - return ( - "import sys, shlex\n" - f"sys.argv = [{script_name!r}]" - f" + shlex.split({input_args!r})\n" - + script_src - ) - return script_src + # Python script: execute the mounted file using runpy + argv_list = [script_path] + for k, v in script_args.items(): + argv_list.extend([f"--{k}", str(v)]) + return ( + "import sys\n" + "import runpy\n" + f"sys.argv = {argv_list!r}\n" + f"runpy.run_path({script_path!r}, run_name='__main__')\n" + ) elif ext in ("sh", "bash"): - # Shell script: wrap in subprocess.run. - # Args are passed as separate list elements after the - # script name to avoid shell injection. - # Both streams are JSON-serialized through stdout since - # UnsafeLocalCodeExecutor drops stdout on exception. + # Shell script: wrap in subprocess.run timeout = self._toolset._script_timeout - cmd = f"['bash', '-c', {script_src!r}, {script_name!r}]" - if input_args: - cmd += f" + shlex.split({input_args!r})" + arr = ["bash", script_path] + for k, v in script_args.items(): + arr.extend([f"--{k}", str(v)]) return ( - "import subprocess, shlex, json as _json\n" + "import subprocess, json as _json\n" "try:\n" " _r = subprocess.run(\n" - f" {cmd},\n" + f" {arr!r},\n" " capture_output=True, text=True,\n" f" timeout={timeout!r},\n" " )\n" @@ -504,15 +532,18 @@ def __init__( *, code_executor: Optional[BaseCodeExecutor] = None, script_timeout: int = _DEFAULT_SCRIPT_TIMEOUT, + additional_tools: Optional[list[Any]] = None, ): """Initializes the SkillToolset. Args: skills: List of skills to register. code_executor: Optional code executor for script execution. - script_timeout: Timeout in seconds for shell script execution - via subprocess.run. Defaults to 300 seconds. Does not apply - to Python scripts executed via exec(). + script_timeout: Timeout in seconds for shell script execution via + subprocess.run. Defaults to 300 seconds. Does not apply to Python + scripts executed via exec(). + additional_tools: Optional list of additional tools (BaseTool, + BaseToolset, or Callables). """ super().__init__() @@ -526,18 +557,100 @@ def __init__( self._skills = {skill.name: skill for skill in skills} self._code_executor = code_executor self._script_timeout = script_timeout + self._additional_tools = additional_tools or [] + + # Initialize core skill tools self._tools = [ ListSkillsTool(self), LoadSkillTool(self), LoadSkillResourceTool(self), - ExecuteSkillScriptTool(self), ] + # Always add RunSkillScriptTool, relies on invocation_context fallback if _code_executor is None + self._tools.append(RunSkillScriptTool(self)) async def get_tools( self, readonly_context: ReadonlyContext | None = None ) -> list[BaseTool]: - """Returns the list of tools in this toolset.""" - return self._tools + """Returns the list of tools in this toolset. + + Dynamically resolves `allowed_tools` from skills against provided + `additional_tools` + and built-in ADK tools. + """ + + import inspect + + from google.adk import tools as built_in_tools + + from .function_tool import FunctionTool + + result = list(self._tools) + + # Collect allowed tools from all skills + allowed_tool_names = set() + for skill in self._list_skills(): + if skill.frontmatter.allowed_tools: + allowed_tool_names.update(skill.frontmatter.allowed_tools) + + if not allowed_tool_names: + return result + + # Resolve additional_tools passed by developer + tools_by_name = {} + for tool_union in self._additional_tools: + if isinstance(tool_union, BaseTool): + tools_by_name[tool_union.name] = tool_union + elif isinstance(tool_union, BaseToolset): + for tool in await tool_union.get_tools(readonly_context): + tools_by_name[tool.name] = tool + elif inspect.isroutine(tool_union): + func_tool = FunctionTool(tool_union) + tools_by_name[func_tool.name] = func_tool + else: + logger.warning("Ignored unsupported additional_tool: %s", tool_union) + + for allowed_tool in allowed_tool_names: + if allowed_tool in tools_by_name: + result.append(tools_by_name[allowed_tool]) + elif hasattr(built_in_tools, allowed_tool): + # Fallback to ADK built-in tools + builtin_obj = getattr(built_in_tools, allowed_tool) + if inspect.isroutine(builtin_obj): + result.append(FunctionTool(builtin_obj)) + elif isinstance(builtin_obj, type) and issubclass( + builtin_obj, BaseTool + ): + try: + # Attempt to instantiate built-in tools that take no arguments + result.append(builtin_obj()) + except TypeError: + logger.warning( + "Could not instantiate built-in tool '%s'. It may require" + " arguments.", + allowed_tool, + ) + elif isinstance(builtin_obj, type) and issubclass( + builtin_obj, BaseToolset + ): + try: + toolset = builtin_obj() + result.extend(await toolset.get_tools(readonly_context)) + except TypeError: + logger.warning( + "Could not instantiate built-in toolset '%s'. It may require" + " arguments.", + allowed_tool, + ) + else: + logger.warning("Unrecognized built-in tool type for %s", allowed_tool) + else: + logger.warning( + "Skill requested tool '%s' which was not provided in" + " additional_tools or found in built-in tools.", + allowed_tool, + ) + + return result def _get_skill(self, name: str) -> models.Skill | None: """Retrieves a skill by name.""" From 446f8a65af320a2bfd3f16f37a86f64021418784 Mon Sep 17 00:00:00 2001 From: Hai-Yuan Cao <2003072+caohy1988@users.noreply.github.com> Date: Mon, 23 Feb 2026 22:44:42 -0800 Subject: [PATCH 18/53] Rename ExecuteSkillScriptTool to RunSkillScriptTool --- tests/unittests/tools/test_skill_toolset.py | 201 +++++++++++--------- 1 file changed, 112 insertions(+), 89 deletions(-) diff --git a/tests/unittests/tools/test_skill_toolset.py b/tests/unittests/tools/test_skill_toolset.py index 2ca4f092d8..2cbcef9d60 100644 --- a/tests/unittests/tools/test_skill_toolset.py +++ b/tests/unittests/tools/test_skill_toolset.py @@ -29,6 +29,7 @@ def mock_skill1_frontmatter(): frontmatter = mock.create_autospec(models.Frontmatter, instance=True) frontmatter.name = "skill1" frontmatter.description = "Skill 1 description" + frontmatter.allowed_tools = ["test_tool"] frontmatter.model_dump.return_value = { "name": "skill1", "description": "Skill 1 description", @@ -45,7 +46,14 @@ def mock_skill1(mock_skill1_frontmatter): skill.instructions = "instructions for skill1" skill.frontmatter = mock_skill1_frontmatter skill.resources = mock.MagicMock( - spec=["get_reference", "get_asset", "get_script"] + spec=[ + "get_reference", + "get_asset", + "get_script", + "list_references", + "list_assets", + "list_scripts", + ] ) def get_ref(name): @@ -70,6 +78,13 @@ def get_script(name): skill.resources.get_reference.side_effect = get_ref skill.resources.get_asset.side_effect = get_asset skill.resources.get_script.side_effect = get_script + skill.resources.list_references.return_value = ["ref1.md"] + skill.resources.list_assets.return_value = ["asset1.txt"] + skill.resources.list_scripts.return_value = [ + "setup.sh", + "run.py", + "build.rb", + ] return skill @@ -79,6 +94,7 @@ def mock_skill2_frontmatter(): frontmatter = mock.create_autospec(models.Frontmatter, instance=True) frontmatter.name = "skill2" frontmatter.description = "Skill 2 description" + frontmatter.allowed_tools = [] frontmatter.model_dump.return_value = { "name": "skill2", "description": "Skill 2 description", @@ -95,7 +111,14 @@ def mock_skill2(mock_skill2_frontmatter): skill.instructions = "instructions for skill2" skill.frontmatter = mock_skill2_frontmatter skill.resources = mock.MagicMock( - spec=["get_reference", "get_asset", "get_script"] + spec=[ + "get_reference", + "get_asset", + "get_script", + "list_references", + "list_assets", + "list_scripts", + ] ) def get_ref(name): @@ -110,6 +133,9 @@ def get_asset(name): skill.resources.get_reference.side_effect = get_ref skill.resources.get_asset.side_effect = get_asset + skill.resources.list_references.return_value = ["ref2.md"] + skill.resources.list_assets.return_value = ["asset2.txt"] + skill.resources.list_scripts.return_value = [] return skill @@ -142,7 +168,7 @@ async def test_get_tools(mock_skill1, mock_skill2): assert isinstance(tools[0], skill_toolset.ListSkillsTool) assert isinstance(tools[1], skill_toolset.LoadSkillTool) assert isinstance(tools[2], skill_toolset.LoadSkillResourceTool) - assert isinstance(tools[3], skill_toolset.ExecuteSkillScriptTool) + assert isinstance(tools[3], skill_toolset.RunSkillScriptTool) @pytest.mark.asyncio @@ -316,7 +342,7 @@ async def test_scripts_resource_not_found(mock_skill1, tool_context_instance): assert result["error_code"] == "RESOURCE_NOT_FOUND" -# ExecuteSkillScriptTool tests +# RunSkillScriptTool tests def _make_tool_context_with_agent(agent=None): @@ -341,20 +367,20 @@ def _make_mock_executor(stdout="", stderr=""): "args, expected_error_code", [ ( - {"script_name": "setup.sh"}, + {"script_path": "setup.sh"}, "MISSING_SKILL_NAME", ), ( {"skill_name": "skill1"}, - "MISSING_SCRIPT_NAME", + "MISSING_SCRIPT_PATH", ), ( - {"skill_name": "", "script_name": "setup.sh"}, + {"skill_name": "", "script_path": "setup.sh"}, "MISSING_SKILL_NAME", ), ( - {"skill_name": "skill1", "script_name": ""}, - "MISSING_SCRIPT_NAME", + {"skill_name": "skill1", "script_path": ""}, + "MISSING_SCRIPT_PATH", ), ], ) @@ -363,7 +389,7 @@ async def test_execute_script_missing_params( ): executor = _make_mock_executor() toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async(args=args, tool_context=ctx) assert result["error_code"] == expected_error_code @@ -373,10 +399,10 @@ async def test_execute_script_missing_params( async def test_execute_script_skill_not_found(mock_skill1): executor = _make_mock_executor() toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( - args={"skill_name": "nonexistent", "script_name": "setup.sh"}, + args={"skill_name": "nonexistent", "script_path": "setup.sh"}, tool_context=ctx, ) assert result["error_code"] == "SKILL_NOT_FOUND" @@ -386,10 +412,10 @@ async def test_execute_script_skill_not_found(mock_skill1): async def test_execute_script_script_not_found(mock_skill1): executor = _make_mock_executor() toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "nonexistent.py"}, + args={"skill_name": "skill1", "script_path": "nonexistent.py"}, tool_context=ctx, ) assert result["error_code"] == "SCRIPT_NOT_FOUND" @@ -398,12 +424,12 @@ async def test_execute_script_script_not_found(mock_skill1): @pytest.mark.asyncio async def test_execute_script_no_code_executor(mock_skill1): toolset = skill_toolset.SkillToolset([mock_skill1]) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) # Agent without code_executor attribute agent = mock.MagicMock(spec=[]) ctx = _make_tool_context_with_agent(agent=agent) result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "setup.sh"}, + args={"skill_name": "skill1", "script_path": "setup.sh"}, tool_context=ctx, ) assert result["error_code"] == "NO_CODE_EXECUTOR" @@ -413,12 +439,12 @@ async def test_execute_script_no_code_executor(mock_skill1): async def test_execute_script_agent_code_executor_none(mock_skill1): """Agent has code_executor attr but it's None.""" toolset = skill_toolset.SkillToolset([mock_skill1]) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) agent = mock.MagicMock() agent.code_executor = None ctx = _make_tool_context_with_agent(agent=agent) result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "setup.sh"}, + args={"skill_name": "skill1", "script_path": "setup.sh"}, tool_context=ctx, ) assert result["error_code"] == "NO_CODE_EXECUTOR" @@ -428,10 +454,10 @@ async def test_execute_script_agent_code_executor_none(mock_skill1): async def test_execute_script_unsupported_type(mock_skill1): executor = _make_mock_executor() toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "build.rb"}, + args={"skill_name": "skill1", "script_path": "build.rb"}, tool_context=ctx, ) assert result["error_code"] == "UNSUPPORTED_SCRIPT_TYPE" @@ -441,32 +467,37 @@ async def test_execute_script_unsupported_type(mock_skill1): async def test_execute_script_python_success(mock_skill1): executor = _make_mock_executor(stdout="hello\n") toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "run.py"}, + args={"skill_name": "skill1", "script_path": "run.py"}, tool_context=ctx, ) assert result["status"] == "success" assert result["stdout"] == "hello\n" assert result["stderr"] == "" assert result["skill_name"] == "skill1" - assert result["script_name"] == "run.py" + assert result["script_path"] == "run.py" # Verify the code passed to executor is the raw script call_args = executor.execute_code.call_args code_input = call_args[0][1] - assert code_input.code == "print('hello')" + assert code_input.code == ( + "import sys\n" + "import runpy\n" + "sys.argv = ['scripts/run.py']\n" + "runpy.run_path('scripts/run.py', run_name='__main__')\n" + ) @pytest.mark.asyncio async def test_execute_script_shell_success(mock_skill1): executor = _make_mock_executor(stdout="setup\n") toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "setup.sh"}, + args={"skill_name": "skill1", "script_path": "setup.sh"}, tool_context=ctx, ) assert result["status"] == "success" @@ -484,13 +515,13 @@ async def test_execute_script_shell_success(mock_skill1): async def test_execute_script_with_input_args_python(mock_skill1): executor = _make_mock_executor(stdout="done\n") toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( args={ "skill_name": "skill1", - "script_name": "run.py", - "input_args": "--verbose --count 3", + "script_path": "run.py", + "args": {"verbose": True, "count": "3"}, }, tool_context=ctx, ) @@ -498,22 +529,23 @@ async def test_execute_script_with_input_args_python(mock_skill1): call_args = executor.execute_code.call_args code_input = call_args[0][1] - assert "sys.argv" in code_input.code - assert "shlex.split" in code_input.code - assert "--verbose --count 3" in code_input.code + assert ( + "['scripts/run.py', '--verbose', 'True', '--count', '3']" + in code_input.code + ) @pytest.mark.asyncio async def test_execute_script_with_input_args_shell(mock_skill1): executor = _make_mock_executor(stdout="done\n") toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( args={ "skill_name": "skill1", - "script_name": "setup.sh", - "input_args": "--force", + "script_path": "setup.sh", + "args": {"force": True}, }, tool_context=ctx, ) @@ -521,25 +553,24 @@ async def test_execute_script_with_input_args_shell(mock_skill1): call_args = executor.execute_code.call_args code_input = call_args[0][1] - assert "shlex.split" in code_input.code - assert "--force" in code_input.code + assert "['bash', 'scripts/setup.sh', '--force', 'True']" in code_input.code @pytest.mark.asyncio async def test_execute_script_scripts_prefix_stripping(mock_skill1): executor = _make_mock_executor(stdout="setup\n") toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( args={ "skill_name": "skill1", - "script_name": "scripts/setup.sh", + "script_path": "scripts/setup.sh", }, tool_context=ctx, ) assert result["status"] == "success" - assert result["script_name"] == "setup.sh" + assert result["script_path"] == "scripts/setup.sh" @pytest.mark.asyncio @@ -550,12 +581,12 @@ async def test_execute_script_toolset_executor_priority(mock_skill1): toolset = skill_toolset.SkillToolset( [mock_skill1], code_executor=toolset_executor ) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) agent = mock.MagicMock() agent.code_executor = agent_executor ctx = _make_tool_context_with_agent(agent=agent) result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "run.py"}, + args={"skill_name": "skill1", "script_path": "run.py"}, tool_context=ctx, ) assert result["stdout"] == "from toolset\n" @@ -568,12 +599,12 @@ async def test_execute_script_agent_executor_fallback(mock_skill1): """Falls back to agent's code executor when toolset has none.""" agent_executor = _make_mock_executor(stdout="from agent\n") toolset = skill_toolset.SkillToolset([mock_skill1]) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) agent = mock.MagicMock() agent.code_executor = agent_executor ctx = _make_tool_context_with_agent(agent=agent) result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "run.py"}, + args={"skill_name": "skill1", "script_path": "run.py"}, tool_context=ctx, ) assert result["stdout"] == "from agent\n" @@ -585,10 +616,10 @@ async def test_execute_script_execution_error(mock_skill1): executor = _make_mock_executor() executor.execute_code.side_effect = RuntimeError("boom") toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "run.py"}, + args={"skill_name": "skill1", "script_path": "run.py"}, tool_context=ctx, ) assert result["error_code"] == "EXECUTION_ERROR" @@ -601,10 +632,10 @@ async def test_execute_script_stderr_only_sets_error_status(mock_skill1): """stderr with no stdout should report error status.""" executor = _make_mock_executor(stdout="", stderr="fatal error\n") toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "run.py"}, + args={"skill_name": "skill1", "script_path": "run.py"}, tool_context=ctx, ) assert result["status"] == "error" @@ -616,10 +647,10 @@ async def test_execute_script_stderr_with_stdout_sets_warning(mock_skill1): """stderr alongside stdout should report warning status.""" executor = _make_mock_executor(stdout="output\n", stderr="deprecation\n") toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "run.py"}, + args={"skill_name": "skill1", "script_path": "run.py"}, tool_context=ctx, ) assert result["status"] == "warning" @@ -633,10 +664,10 @@ async def test_execute_script_execution_error_truncated(mock_skill1): executor = _make_mock_executor() executor.execute_code.side_effect = RuntimeError("x" * 300) toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "run.py"}, + args={"skill_name": "skill1", "script_path": "run.py"}, tool_context=ctx, ) assert result["error_code"] == "EXECUTION_ERROR" @@ -651,10 +682,10 @@ async def test_execute_script_system_exit_caught(mock_skill1): executor = _make_mock_executor() executor.execute_code.side_effect = SystemExit(1) toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "run.py"}, + args={"skill_name": "skill1", "script_path": "run.py"}, tool_context=ctx, ) assert result["error_code"] == "EXECUTION_ERROR" @@ -667,10 +698,10 @@ async def test_execute_script_system_exit_zero_is_success(mock_skill1): executor = _make_mock_executor() executor.execute_code.side_effect = SystemExit(0) toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "run.py"}, + args={"skill_name": "skill1", "script_path": "run.py"}, tool_context=ctx, ) assert result["status"] == "success" @@ -683,10 +714,10 @@ async def test_execute_script_system_exit_none_is_success(mock_skill1): executor = _make_mock_executor() executor.execute_code.side_effect = SystemExit(None) toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "run.py"}, + args={"skill_name": "skill1", "script_path": "run.py"}, tool_context=ctx, ) assert result["status"] == "success" @@ -700,10 +731,10 @@ async def test_execute_script_shell_includes_timeout(mock_skill1): toolset = skill_toolset.SkillToolset( [mock_skill1], code_executor=executor, script_timeout=60 ) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "setup.sh"}, + args={"skill_name": "skill1", "script_path": "setup.sh"}, tool_context=ctx, ) assert result["status"] == "success" @@ -727,33 +758,15 @@ def get_script_extended(name): executor = _make_mock_executor() toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( - args={"skill_name": "skill1", "script_name": "noext"}, + args={"skill_name": "skill1", "script_path": "noext"}, tool_context=ctx, ) assert result["error_code"] == "UNSUPPORTED_SCRIPT_TYPE" -@pytest.mark.asyncio -async def test_execute_script_invalid_input_args(mock_skill1): - """Unclosed quotes in input_args should return INVALID_INPUT_ARGS.""" - executor = _make_mock_executor() - toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) - ctx = _make_tool_context_with_agent() - result = await tool.run_async( - args={ - "skill_name": "skill1", - "script_name": "run.py", - "input_args": '--name "unclosed', - }, - tool_context=ctx, - ) - assert result["error_code"] == "INVALID_INPUT_ARGS" - - # ── Integration tests using real UnsafeLocalCodeExecutor ── @@ -768,7 +781,14 @@ def _make_skill_with_script(skill_name, script_name, script): fm.description = f"Test skill {skill_name}" skill.frontmatter = fm skill.resources = mock.MagicMock( - spec=["get_reference", "get_asset", "get_script"] + spec=[ + "get_reference", + "get_asset", + "get_script", + "list_references", + "list_assets", + "list_scripts", + ] ) def get_script(name): @@ -779,6 +799,9 @@ def get_script(name): skill.resources.get_script.side_effect = get_script skill.resources.get_reference.return_value = None skill.resources.get_asset.return_value = None + skill.resources.list_references.return_value = [] + skill.resources.list_assets.return_value = [] + skill.resources.list_scripts.return_value = [script_name] return skill @@ -796,12 +819,12 @@ async def test_integration_python_stdout(): script = models.Script(src="print('hello world')") skill = _make_skill_with_script("test_skill", "hello.py", script) toolset = _make_real_executor_toolset([skill]) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( args={ "skill_name": "test_skill", - "script_name": "hello.py", + "script_path": "hello.py", }, tool_context=ctx, ) @@ -816,12 +839,12 @@ async def test_integration_python_sys_exit_zero(): script = models.Script(src="import sys; sys.exit(0)") skill = _make_skill_with_script("test_skill", "exit_zero.py", script) toolset = _make_real_executor_toolset([skill]) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( args={ "skill_name": "test_skill", - "script_name": "exit_zero.py", + "script_path": "exit_zero.py", }, tool_context=ctx, ) @@ -834,12 +857,12 @@ async def test_integration_shell_stdout_and_stderr(): script = models.Script(src="echo output; echo warning >&2") skill = _make_skill_with_script("test_skill", "both.sh", script) toolset = _make_real_executor_toolset([skill]) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( args={ "skill_name": "test_skill", - "script_name": "both.sh", + "script_path": "both.sh", }, tool_context=ctx, ) @@ -854,12 +877,12 @@ async def test_integration_shell_stderr_only(): script = models.Script(src="echo failure >&2") skill = _make_skill_with_script("test_skill", "err.sh", script) toolset = _make_real_executor_toolset([skill]) - tool = skill_toolset.ExecuteSkillScriptTool(toolset) + tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() result = await tool.run_async( args={ "skill_name": "test_skill", - "script_name": "err.sh", + "script_path": "err.sh", }, tool_context=ctx, ) From e2445c2f8b2959d9d72ce29b25fef2adb054820e Mon Sep 17 00:00:00 2001 From: Hai-Yuan Cao <2003072+caohy1988@users.noreply.github.com> Date: Mon, 23 Feb 2026 22:46:45 -0800 Subject: [PATCH 19/53] Add optional path and working_dir to execution dataclasses --- src/google/adk/code_executors/code_execution_utils.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/google/adk/code_executors/code_execution_utils.py b/src/google/adk/code_executors/code_execution_utils.py index 7cccce48be..012d65bc08 100644 --- a/src/google/adk/code_executors/code_execution_utils.py +++ b/src/google/adk/code_executors/code_execution_utils.py @@ -46,6 +46,11 @@ class File: The mime type of the file (e.g., "image/png"). """ + path: Optional[str] = None + """ + The relative path to write the file to in a sandbox context (e.g., "references/data.csv"). + """ + @dataclasses.dataclass class CodeExecutionInput: @@ -66,6 +71,11 @@ class CodeExecutionInput: The execution ID for the stateful code execution. """ + working_dir: Optional[str] = None + """ + The designated working directory for the code execution environment. + """ + @dataclasses.dataclass class CodeExecutionResult: From 8b99a5e2710b75d9ec7dbda41da560e1e6661682 Mon Sep 17 00:00:00 2001 From: Hai-Yuan Cao <2003072+caohy1988@users.noreply.github.com> Date: Mon, 23 Feb 2026 22:47:55 -0800 Subject: [PATCH 20/53] Implement temporary directory for code execution Refactor code execution to use a temporary directory for input files and change working directory as needed. --- .../unsafe_local_code_executor.py | 48 +++++++++++++++---- 1 file changed, 39 insertions(+), 9 deletions(-) diff --git a/src/google/adk/code_executors/unsafe_local_code_executor.py b/src/google/adk/code_executors/unsafe_local_code_executor.py index e91a54bb24..5bfe064232 100644 --- a/src/google/adk/code_executors/unsafe_local_code_executor.py +++ b/src/google/adk/code_executors/unsafe_local_code_executor.py @@ -17,7 +17,10 @@ from contextlib import redirect_stdout import io import logging +import os import re +import tempfile +import threading from typing import Any from pydantic import Field @@ -30,6 +33,8 @@ logger = logging.getLogger('google_adk.' + __name__) +_execution_lock = threading.Lock() + def _prepare_globals(code: str, globals_: dict[str, Any]) -> None: """Prepare globals for code execution, injecting __name__ if needed.""" @@ -67,15 +72,40 @@ def execute_code( # Execute the code. output = '' error = '' - try: - globals_ = {} - _prepare_globals(code_execution_input.code, globals_) - stdout = io.StringIO() - with redirect_stdout(stdout): - exec(code_execution_input.code, globals_, globals_) - output = stdout.getvalue() - except Exception as e: - error = str(e) + + with _execution_lock: + original_cwd = os.getcwd() + try: + # Prepare the execution environment (temp volume) + with tempfile.TemporaryDirectory() as temp_dir: + # Write input files to the temp directory + for f in code_execution_input.input_files: + file_path = os.path.join(temp_dir, f.path or f.name) + os.makedirs(os.path.dirname(file_path), exist_ok=True) + mode = 'wb' if isinstance(f.content, bytes) else 'w' + with open(file_path, mode) as out_f: + out_f.write(f.content) + + # Change working directory if specified + if code_execution_input.working_dir: + exec_dir = os.path.join(temp_dir, code_execution_input.working_dir) + os.makedirs(exec_dir, exist_ok=True) + os.chdir(exec_dir) + else: + os.chdir(temp_dir) + + # Execute the code + globals_ = {} + _prepare_globals(code_execution_input.code, globals_) + stdout = io.StringIO() + with redirect_stdout(stdout): + exec(code_execution_input.code, globals_, globals_) + output = stdout.getvalue() + + except Exception as e: + error = str(e) + finally: + os.chdir(original_cwd) # Collect the final result. return CodeExecutionResult( From c2e92cf1d04099e309be134486293f414639a069 Mon Sep 17 00:00:00 2001 From: Hai-Yuan Cao <2003072+caohy1988@users.noreply.github.com> Date: Mon, 23 Feb 2026 22:48:28 -0800 Subject: [PATCH 21/53] Enhance code execution with working directory support --- .../adk/code_executors/vertex_ai_code_executor.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/google/adk/code_executors/vertex_ai_code_executor.py b/src/google/adk/code_executors/vertex_ai_code_executor.py index 67c42ed8f2..71aa76b945 100644 --- a/src/google/adk/code_executors/vertex_ai_code_executor.py +++ b/src/google/adk/code_executors/vertex_ai_code_executor.py @@ -148,8 +148,16 @@ def execute_code( code_execution_input: CodeExecutionInput, ) -> CodeExecutionResult: # Execute the code. + code_to_exec = self._get_code_with_imports(code_execution_input.code) + if code_execution_input.working_dir: + code_to_exec = ( + f'import os\nos.makedirs("{code_execution_input.working_dir}",' + f' exist_ok=True)\nos.chdir("{code_execution_input.working_dir}")\n' + + code_to_exec + ) + code_execution_result = self._execute_code_interpreter( - self._get_code_with_imports(code_execution_input.code), + code_to_exec, code_execution_input.input_files, code_execution_input.execution_id, ) @@ -216,7 +224,7 @@ def _execute_code_interpreter( operation_params = {'code': code} if input_files: operation_params['files'] = [ - {'name': f.name, 'contents': f.content} for f in input_files + {'name': f.path or f.name, 'contents': f.content} for f in input_files ] if session_id: operation_params['session_id'] = session_id From c77ac54468f2bcba12c18cbec745535de20ff6dd Mon Sep 17 00:00:00 2001 From: Hai-Yuan Cao <2003072+caohy1988@users.noreply.github.com> Date: Mon, 23 Feb 2026 22:49:26 -0800 Subject: [PATCH 22/53] Add 'path' field to input files in tests --- .../test_code_executor_context.py | 24 +++++++++++++++---- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/tests/unittests/code_executors/test_code_executor_context.py b/tests/unittests/code_executors/test_code_executor_context.py index cdf47eb3d8..e66a3eb3cb 100644 --- a/tests/unittests/code_executors/test_code_executor_context.py +++ b/tests/unittests/code_executors/test_code_executor_context.py @@ -32,9 +32,12 @@ def context_with_data() -> CodeExecutorContext: "execution_session_id": "session123", "processed_input_files": ["file1.csv", "file2.txt"], }, - "_code_executor_input_files": [ - {"name": "input1.txt", "content": "YQ==", "mime_type": "text/plain"} - ], + "_code_executor_input_files": [{ + "name": "input1.txt", + "content": "YQ==", + "mime_type": "text/plain", + "path": None, + }], "_code_executor_error_counts": {"invocationA": 2}, } state = State(state_data, {}) @@ -145,6 +148,7 @@ def test_add_input_files_new(empty_state: State): "name": "new.dat", "content": "Yg==", "mime_type": "application/octet-stream", + "path": None, }] @@ -153,8 +157,18 @@ def test_add_input_files_append(context_with_data: CodeExecutorContext): new_file = File(name="input2.log", content="Yw==", mime_type="text/x-log") context_with_data.add_input_files([new_file]) expected_files_data = [ - {"name": "input1.txt", "content": "YQ==", "mime_type": "text/plain"}, - {"name": "input2.log", "content": "Yw==", "mime_type": "text/x-log"}, + { + "name": "input1.txt", + "content": "YQ==", + "mime_type": "text/plain", + "path": None, + }, + { + "name": "input2.log", + "content": "Yw==", + "mime_type": "text/x-log", + "path": None, + }, ] assert ( context_with_data._session_state["_code_executor_input_files"] From e77473d5edd95434b8506c4d43bb4169787a859d Mon Sep 17 00:00:00 2001 From: Hai-Yuan Cao <2003072+caohy1988@users.noreply.github.com> Date: Mon, 23 Feb 2026 22:52:57 -0800 Subject: [PATCH 23/53] Add design for skill execution script in ADK This document outlines the design for integrating pre-registered function execution and script execution within ADK's non-bash based skills toolset. It details the motivation, proposal, enabling usage of ADK tools, and the architecture of the RunSkillScriptTool. --- docs/design/skill_execution_script.md | 370 ++++++++++++++++++++++++++ 1 file changed, 370 insertions(+) create mode 100644 docs/design/skill_execution_script.md diff --git a/docs/design/skill_execution_script.md b/docs/design/skill_execution_script.md new file mode 100644 index 0000000000..ae269b0f0b --- /dev/null +++ b/docs/design/skill_execution_script.md @@ -0,0 +1,370 @@ + + +# Summary + +This outlines the design for adding pre-registered function execution within ADK’s non-bash based skills toolset. + +# Motivation + +ADK Skills provide a way to package instructions, resources, and scripts to extend agent capabilities. While skills can guide an agent's reasoning, the ability to execute code within the skill's context unlocks more powerful actions. This requires a secure and well-defined execution environment. + +We also want to support adapting existing `BaseTool` or `FunctionTool` instances, making them available as part of the Skill interface. + +# Proposal + +There are two parts to the proposal for pre-registered function execution: + +1. Allow SkillToolset to include **additional ADK tools/functions** that are specified by the skill: adapting existing BaseTool and FunctionTool interfaces so that users can feed in existing functions +2. Enable **script execution** (i.e. files in skill/scripts) with RunSkillScriptTool + +## Enabling usage of ADK tools + +Users should be able to pass in existing ADK tools (both built-in and custom) as BaseTool objects and have them be used within a skill. To do this, we can use the [allowed\_tools]field of a skill’s frontmatter to determine which tools should be instantiated. + +Within SkillToolset, we can modify the signature to accept an optional `additional_tools` argument that allows users to pass in pre-instantiated tools, toolsets, or callables. Then we can update `SkillToolset.get_tools()` to dynamically resolve the tools in allowed\_tools for all skills in the toolset. We will iterate through the skill’s allowed tools and check if any of the `additional_tools` match, and add them to the toolset if so. As a fallback, we will also check if the tools exist as built-in tools in the google.adk.tools directory. + +Since get\_tools() is called every invocation, we will also cache the tools so we don’t have to resolve them every time. + +```py +class SkillToolset(BaseToolset): + def __init__( + self, + skills: list[models.Skill], + additional_tools: list[ToolUnion] = None, + ): + ... + + def get_tools(self, readonly_context) -> list[BaseTool]: + # Collect allowed tools from skills + allowed_tool_names = set() + for skill in self.skills: + if skill.frontmatter.allowed_tools: + allowed_tool_names.update([name for name in skill.frontmatter.allowed_tools]) + + # Resolve tools passed in with `additional_tools` + tools_by_name = {} + for tool_union in self._additional_tools: + if isinstance(tool_union, BaseTool): + tools_by_name[tool_union.name] = tool_union + if isinstance(tool_union, BaseToolset): + ts_tools = await tool_union.get_tools(readonly_context) + for t in ts_tools: + tools_by_name[t.name] = t + elif callable(tool_union): + tools_by_name[tool_union.name] = FunctionTool(tool_union) + + for allowed_tool in allowed_tool_names: + # add tools from tools_by_name or if they don't exist, + # try to resolve using built-in tools in the google.adk.tools directory + ... +``` + +## Enabling Script Execution + +We will integrate script execution into `SkillToolset` by using the existing [BaseCodeExecutor] interface. +We can modify `SkillToolset` to accept an optional code\_executor argument and create a new tool, `RunSkillScriptTool`, that will be used for script execution: + +```py +class SkillToolset(BaseToolset): + def __init__( + self, + skills: list[models.Skill], + code_executor: Optional[base_code_executor.BaseCodeExecutor] = None, + ): + super().__init__() + self._skills = {skill.name: skill for skill in skills} + self._code_executor = code_executor + + self._tools = [LoadSkillTool(self), LoadSkillResourceTool(self)] + # Add RunSkillScriptTool for function execution + if self._code_executor: + self._tools.append(RunSkillScriptTool(self)) +``` + +`RunSkillScriptTool` will be used for script execution. It will take a toolset and script\_metadata as input. + +```py +class RunSkillScriptTool(BaseTool): + def __init__(self, toolset: "SkillToolset", script_metadata: Dict[str, Any] = None): + # Initialize the tool + pass + + def _get_declaration(self) -> types.FunctionDeclaration | None: + params_schema = { + "type": "object", + "properties": { + "skill_name": { + "type": "string", + "description": "The name of the skill." + }, + "script_path": { + "type": "string", + "description": "The relative path to the script (e.g., 'scripts/my_script.py' or 'scripts/setup.sh')." + }, + "args": { + "type": "object", + "description": "Optional arguments to pass to the script as key-value pairs." + }, + }, + "required": ["skill_name", "script_path"], + } + return types.FunctionDeclaration( + name=self.name, + description=self.description, + parameters_json_schema=params_schema + ) + + async def run_async(self, *, args: dict[str, Any], tool_context: ToolContext) -> Any: + # 1. Validate inputs (skill_name, script_path) + # 2. Extract script and arguments Map + # 3. Mount all skill files (assets, references, scripts) as input_files + # 4. Set sandbox_working_dir to the sandbox root '.' + # 5. Generate safe subprocess wrappers or sys.argv injections + # 6. Execute via BaseCodeExecutor matching toolset configuration + ... +``` + +**Script Invocation:** Arguments are passed to the script similar to command-line arguments, by mocking `sys.argv`. The script's standard output, standard error, and any resulting files will be captured by the executor. + +**Executor Choices:** The `SkillToolset` can be configured with any ADK `BaseCodeExecutor` instance. Recommended executors include VertexAiCodeExecutor (executes code in a secure Vertex AI sandbox) and GkeCodeExecutor (executes code in a gVisor sandbox on GKE). + +### Adapting BaseCodeExecutor + +The current `BaseCodeExecutor` interface is designed for executing LLM-generated code snippets. To effectively support `RunSkillScriptTool`, we should consider file system context and path resolution: + +`BaseCodeExecutor`'s `execute_code` method doesn't have a way to inform the execution environment about the skill's file structure. Scripts within a skill will likely need to read from `../references/` or `../assets/` using relative paths. The sandbox needs to honor these paths relative to the skill's root or the `scripts/` directory. + +We can extend the `File` Dataclass in `code_execution_utils.py` to include a `path`: + +```py +@dataclasses.dataclass(frozen=True) +class File: + """A structure that contains a file name and its content.""" + name: str # Base name of the file + content: str | bytes + mime_type: str = 'text/plain' + path: Optional[str] = None # Native relative path (e.g. 'references/guidelines.md') +``` + +We will also extend `CodeExecutionInput` to add a `working_dir`: + +```py +@dataclasses.dataclass +class CodeExecutionInput: + code: str + input_files: list[File] = dataclasses.field(default_factory=list) + execution_id: Optional[str] = None + working_dir: Optional[str] = None # e.g., '/skill/scripts' +``` + +`RunSkillScriptTool` will package the Skill’s resources into input\_files to add to the CodeExecutionInput. Within the BaseCodeExecutor implementations (VertexAI, GKE, etc), we will then: + +* Read input\_files. For each `File` with sandbox\_path, create the file and any necessary parent directories at the exact path in the sandbox +* Set the sandbox working directory to sandbox\_working\_dir + +#### File Permissions + +Files/directories created within the sandbox from the skill’s `references/` and `assets/` should be read-only for the executed script. The script should have write-access to a dedicated temp directory in the sandbox and potentially a designated output directory. + +#### Other Considerations + +The CodeExecutor and sandbox environment will address most script execution issues: + +* Resource limits are enforced by sandbox env +* Network access disabled by default +* **Output files:** Files will be created in output directory and returned in `CodeExecutionResult.output_files` +* **Error handling:** Report script exceptions, exit codes, executor errors through `CodeExecutionResult.stderr` + * Add error post-processing: we will implement an ‘LLM-friendly’ error formatter. Instead of returning the full traceback, the tool will extract the specific Exception type and the offending line of code to help the agent self-correct its script invocation. + +## RunSkillScriptTool — Design Details & Considerations + +### Overview + +`RunSkillScriptTool` enables ADK agents to execute scripts bundled inside a +skill's `scripts/` directory via ADK's `BaseCodeExecutor` infrastructure. This +closes the gap between the +[Agent Skills spec](https://agentskills.io/specification) (which defines +`scripts/` as an optional skill resource) and ADK's runtime capabilities. By +mounting skill dependencies including `references/` and `assets/`, +`RunSkillScriptTool` executes skills in sandboxed environments with full access +to their context. + +### Architecture + +``` +LLM calls run_skill_script(skill_name, script_path, args) + │ + ▼ +┌─ RunSkillScriptTool.run_async() ─────────────────────┐ +│ 1. Validate params & resolve skill/script │ +│ 2. Resolve code executor (toolset → agent fallback) │ +│ 3. Validate args (handled natively by JSON schema) │ +│ 4. Mount dependencies (assets, references, scripts) │ +│ 5. _prepare_code() → generate Python wrapper code │ +│ 6. code_executor.execute_code(..., input_files=...) │ +│ 7. Parse result (JSON envelope for shell scripts) │ +│ 8. Return {stdout, stderr, status} to LLM │ +└───────────────────────────────────────────────────────────┘ +``` + +### Parameter Schema Design + +The tool employs specific parameter designs to ensure safe sandboxed execution +and high LLM reliability: + +1. **`script_path` vs `script_name`:** \ + Instead of a flat `script_name` (e.g. `setup.sh`), the tool requires the + full relative `script_path` (e.g. `scripts/setup.sh` or + `scripts/utils/helper.py`). This is required because the tool mounts the + entire skill directory into the execution sandbox, meaning the script must + be invoked from the true sandbox root path so it can reliably access its + `assets/` and `references/` via relative paths. + +2. **`args` Dictionary:** \ + Instead of taking a raw array of string arguments (`["--verbose", + "--force"]`), the tool takes a structured key-value `args` object + (`{"verbose": true, "force": true}`). LLMs are significantly more reliable + at generating structured JSON objects than raw command-line flag arrays. + Furthermore, accepting an object moves the burden of secure structural + flattening (e.g. constructing the `['--verbose', 'True']` array) to the + Python code, completely eliminating a class of shell-injection + vulnerabilities. + +### Script Type Handling + +Type | Extension | Execution Method | Timeout | Args Injection +:----- | :------------- | :---------------------------------------------------- | :--------------------- | :------------- +Python | `.py` | Direct `exec()` via code executor | No (executor-level) | `sys.argv = [script_path] + mapped_args` +Shell | `.sh`, `.bash` | `subprocess.run(['bash', script_path] + mapped_args)` | Yes (`script_timeout`) | `args` parsed as sequence of flattened pairs +Other | any | Rejected | N/A | N/A + +**Extensionless files are rejected** (not silently treated as Python) to avoid +unexpected behavior. + +### Code Executor Resolution Chain + +``` +1. SkillToolset(code_executor=...) ← explicit, highest priority +2. agent.code_executor ← fallback to agent's executor +3. None → return NO_CODE_EXECUTOR ← actionable error +``` + +This design allows a single toolset-level executor to be shared across all +skills, or per-agent executors for different isolation levels. + +### Shell Script JSON Envelope + +Shell scripts face a unique challenge: `UnsafeLocalCodeExecutor` captures stdout +via `redirect_stdout(StringIO)`, but if the generated code raises an exception, +`stdout.getvalue()` is never called and stdout is lost. This means a naive +`raise RuntimeError(stderr)` approach discards any stdout the script produced. + +Crucially, **even when running inside a secure sandbox** (like Vertex AI Code +Interpreter), sandboxes often struggle to cleanly report *why* an arbitrary +script failed if the script loops infinitely or crashes aggressively. An abrupt +exit often yields a generic "sandbox failed" error, denying the LLM the context +it needs to self-correct. + +**Solution**: The Python subprocess wrapper we inject *around* the script +executes the shell command safely and serializes both stdout and stderr as a +JSON envelope through the single stdout channel. + +Even inside a sandbox environment, our wrapper catches the shell script's output +in real-time. If the shell script times out inside the sandbox, our wrapper +catches the `TimeoutExpired` exception, scoops up whatever output the shell +script produced *before* it hung, packages it into JSON, and returns that +perfectly structured payload to the Executor. This guarantees the LLM always +receives perfectly exact `stdout` and `stderr` logs regardless of script +crashes. + +```py +# Generated code for shell scripts: +import subprocess, shlex, json as _json +try: + _r = subprocess.run( + ['bash', SCRIPT_PATH] + MAPPED_ARGS, + capture_output=True, text=True, + timeout=SCRIPT_TIMEOUT, + ) + print(_json.dumps({ + '__shell_result__': True, + 'stdout': _r.stdout, + 'stderr': _r.stderr, + 'returncode': _r.returncode, + })) +except subprocess.TimeoutExpired as _e: + print(_json.dumps({ + '__shell_result__': True, + 'stdout': _e.stdout or '', + 'stderr': 'Timed out after Ns', + 'returncode': -1, + })) +``` + +`run_async()` then parses this JSON envelope (only for shell scripts, keyed on +`__shell_result__`) to extract both streams and the return code. This works +reliably with both `UnsafeLocalCodeExecutor` and container-based executors. + +### Three-State Status Model + +Status | Condition +:-------- | :------------------------------------------------------------- +`success` | No stderr +`warning` | Both stdout and stderr present (e.g., deprecation warnings) +`error` | Stderr only (no stdout), or non-zero returncode with no stdout + +### Security Considerations + +**Shell injection prevention:** + +- **Structured Argument Arrays:** `args` is passed as a structured dictionary + by the LLM natively. The tool converts these into strict, flattened string + arrays `['--key', 'value']` and passes them securely to `subprocess.run` + with `shell=False`. Because the elements are passed as a strict array, the + underlying OS treats flags and values as literal parameters passed *into* + the script, meaning any malicious shell operators (e.g. `&&`, `|`) are + treated as literal strings and ignored. +- The script source is executed as an isolated file path inside the sandboxed + `$working_dir`. + +**`SystemExit` handling:** + +- `sys.exit()` raises `SystemExit(BaseException)`, which is NOT caught by + `except Exception` in executors +- `run_async()` explicitly catches `SystemExit` to prevent skill scripts from + terminating the host process +- `sys.exit(0)` and `sys.exit(None)` are treated as successful termination +- Non-zero exit codes return `EXECUTION_ERROR` + +**Executor security:** + +- `UnsafeLocalCodeExecutor` runs code in the host process via `exec()` — + suitable only for trusted, first-party skills +- For third-party or untrusted skills, a sandboxed executor (e.g., + `ContainerCodeExecutor`) should be used +- The sample agent includes explicit warnings about this + +### Known Limitations & Future Work + +1. **No timeout for Python scripts**: `exec()` provides no built-in timeout + mechanism. A malicious/buggy Python script can hang indefinitely. This is an + executor-level concern — solving it properly requires running Python scripts + in a subprocess or implementing executor-level cancellation. + +2. **Python script stdout lost on exception**: When a Python script writes to + stdout and then raises, `UnsafeLocalCodeExecutor` loses the stdout (same + root cause as the shell fix). This is less critical for Python since + exceptions are the natural error mechanism, but could be improved at the + executor level. + +### Error Codes Reference + +Error Code | Meaning +:------------------------ | :--------------------------------------------- +`MISSING_SKILL_NAME` | `skill_name` parameter not provided +`MISSING_SCRIPT_PATH` | `script_path` parameter not provided +`SKILL_NOT_FOUND` | No skill with that name registered +`SCRIPT_NOT_FOUND` | No script with that name in the skill +`NO_CODE_EXECUTOR` | No code executor configured (toolset or agent) +`UNSUPPORTED_SCRIPT_TYPE` | File extension not `.py`, `.sh`, or `.bash` +`EXECUTION_ERROR` | Runtime error, non-zero exit, or `sys.exit(N)` From e81293ca7b41ffcf74ca9bd42e952b3e538e0cc5 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Mon, 23 Feb 2026 22:59:13 -0800 Subject: [PATCH 24/53] fix: Add Docker build retry logic and skill script test coverage - full_runner.py: Add retry loop (max 2 retries) with Docker prune between attempts, forcerm=True for intermediate containers, periodic prune every N tasks (--prune-interval flag), and build summary stats - test_skill_toolset.py: Add 8 tests covering shell JSON envelope parsing, non-zero returncode, TimeoutExpired path, input_files packaging with correct paths/content, and working_dir validation Co-Authored-By: Claude Opus 4.6 --- benchmarks/skillsbench/full_runner.py | 163 +++++++++++++++-- tests/unittests/tools/test_skill_toolset.py | 185 ++++++++++++++++++++ 2 files changed, 335 insertions(+), 13 deletions(-) diff --git a/benchmarks/skillsbench/full_runner.py b/benchmarks/skillsbench/full_runner.py index b9dd6c2c83..3fc559bb5f 100644 --- a/benchmarks/skillsbench/full_runner.py +++ b/benchmarks/skillsbench/full_runner.py @@ -42,6 +42,7 @@ import re import sys import tarfile +import threading import time import tomllib from typing import Any @@ -161,7 +162,7 @@ class TaskContainerExecutor(BaseCodeExecutor): _container: object = None _client: object = None - _cmd_timeout: float = 300.0 # Per-command timeout in seconds + _cmd_timeout: float = 180.0 # Per-command timeout in seconds def start( self, @@ -279,6 +280,43 @@ def exec_in_container( combined = combined + "\n" + stderr if combined else stderr return rc, combined + async def async_exec_in_container( + self, + cmd: list[str], + workdir: str = "/root", + timeout: float | None = None, + ) -> tuple[int, str]: + """Async version of exec_in_container. + + Runs the blocking Docker exec_run in a thread executor so + the asyncio event loop stays free for timeout checks. + """ + timeout = timeout or self._cmd_timeout + + def _blocking(): + if self._container is None: + raise RuntimeError("Container not started") + rc, output = self._container.exec_run( + cmd, + workdir=workdir, + demux=True, + ) + out = (output[0] or b"").decode("utf-8", errors="replace") + err = (output[1] or b"").decode("utf-8", errors="replace") + combined = out + if err: + combined = combined + "\n" + err if combined else err + return rc, combined + + loop = asyncio.get_running_loop() + try: + return await asyncio.wait_for( + loop.run_in_executor(None, _blocking), + timeout=timeout, + ) + except asyncio.TimeoutError: + return -1, f"Command timed out after {timeout}s" + def read_file_from_container(self, path: str) -> str | None: """Read a single file from the container, or None.""" if self._container is None: @@ -322,9 +360,14 @@ def build_task_image( *, rebuild: bool = False, build_timeout: float = _DEFAULT_BUILD_TIMEOUT, + max_retries: int = 2, ) -> str: """Build (or reuse) a Docker image for one task. + Retries up to *max_retries* times on failure, running + ``docker system prune -f`` between attempts to free + build cache and disk space. + Returns the image tag. """ tag = f"skillsbench-{task_dir.name}:latest" @@ -339,14 +382,36 @@ def build_task_image( if not (build_ctx / "Dockerfile").exists(): raise FileNotFoundError(f"No Dockerfile in {build_ctx}") - logger.info("Building image %s …", tag) - client.images.build( - path=str(build_ctx), - tag=tag, - rm=True, - timeout=int(build_timeout), - ) - return tag + last_exc: Exception | None = None + for attempt in range(1, max_retries + 2): # 1 try + max_retries + try: + logger.info("Building image %s (attempt %d) …", tag, attempt) + client.images.build( + path=str(build_ctx), + tag=tag, + rm=True, + forcerm=True, + timeout=int(build_timeout), + ) + return tag + except Exception as exc: + last_exc = exc + if attempt <= max_retries: + logger.warning( + "Build attempt %d failed for %s: %s — pruning and retrying", + attempt, + tag, + str(exc)[:120], + ) + try: + client.containers.prune() + client.images.prune(filters={"dangling": True}) + except Exception: + pass + else: + break + + raise last_exc # type: ignore[misc] # ── Lenient skill loader ───────────────────────────────────────────── @@ -634,7 +699,8 @@ async def run_async(self, *, args, tool_context) -> Any: command = args.get("command", "") if not command: return {"error": "command is required"} - rc, output = self._executor.exec_in_container( + # Use async exec to keep the event loop free for timeouts + rc, output = await self._executor.async_exec_in_container( ["bash", "-c", command], ) # Truncate very long output to avoid blowing context @@ -806,6 +872,21 @@ def discover_tasks( # ── evaluate one task ──────────────────────────────────────────────── +def _watchdog_stop(executor: TaskContainerExecutor, seconds: float): + """Background thread: force-kill the container after *seconds*. + + This is a hard safety net — if asyncio.wait_for fails to cancel + a task because the event loop is blocked by synchronous Docker + calls, the watchdog kills the container from a separate thread, + causing the blocking exec_run to fail and unblock the loop. + """ + time.sleep(seconds) + try: + executor.stop() + except Exception: + pass + + async def evaluate_task( task_dir: pathlib.Path, client: docker.DockerClient, @@ -830,6 +911,7 @@ async def evaluate_task( executor = TaskContainerExecutor() num_skills = 0 agent_timeout = _DEFAULT_AGENT_TIMEOUT + watchdog = None try: # 1. Parse config @@ -883,14 +965,22 @@ async def evaluate_task( # 5. Start container executor.start(image_tag, client) - # 6. Build agent and run + # 6. Start watchdog — hard kill after agent_timeout + 60s + watchdog = threading.Thread( + target=_watchdog_stop, + args=(executor, agent_timeout + 60), + daemon=True, + ) + watchdog.start() + + # 7. Build agent and run agent = build_agent(skills, executor) invocations = await asyncio.wait_for( run_task(agent, user_query), timeout=agent_timeout, ) - # 7. Score + # 8. Score if skip_tests: score = score_task_heuristic(invocations) else: @@ -978,6 +1068,16 @@ def print_summary( # ── full evaluation loop ───────────────────────────────────────────── +def _docker_prune(client: docker.DockerClient) -> None: + """Prune stopped containers and dangling images.""" + try: + client.containers.prune() + client.images.prune(filters={"dangling": True}) + logger.info("Docker prune completed") + except Exception as exc: + logger.warning("Docker prune failed: %s", exc) + + async def run_full_evaluation( tasks_dir: pathlib.Path, client: docker.DockerClient, @@ -987,13 +1087,21 @@ async def run_full_evaluation( rebuild: bool = False, build_only: bool = False, skip_tests: bool = False, + prune_interval: int = 20, ) -> list[TaskResult]: - """Run evaluation on all matching tasks.""" + """Run evaluation on all matching tasks. + + Args: + prune_interval: Prune Docker every N tasks to prevent + disk buildup during long runs. Set to 0 to disable. + """ task_dirs = discover_tasks(tasks_dir, filter_pattern) total = len(task_dirs) print(f"Found {total} tasks in {tasks_dir}\n") if build_only: + ok = 0 + fail = 0 for idx, td in enumerate(task_dirs, 1): name = td.name[:35].ljust(35) try: @@ -1004,14 +1112,27 @@ async def run_full_evaluation( rebuild=rebuild, build_timeout=config["build_timeout"], ) + ok += 1 print(f"[{idx:>2}/{total}] {name} OK ({tag})") except Exception as exc: + fail += 1 print(f"[{idx:>2}/{total}] {name} FAIL ({str(exc)[:80]})") sys.stdout.flush() + # Periodic prune during build-only runs + if prune_interval and idx % prune_interval == 0: + print(f" >> pruning Docker (every {prune_interval})") + _docker_prune(client) + print(f"\nBuild summary: {ok} OK, {fail} FAIL / {total}") return [] results: list[TaskResult] = [] for idx, td in enumerate(task_dirs, 1): + name = td.name[:35].ljust(35) + print( + f"[{idx:>2}/{total}] {name} ...running", + end="", + flush=True, + ) result = await evaluate_task( td, client, @@ -1020,7 +1141,13 @@ async def run_full_evaluation( skip_tests=skip_tests, ) results.append(result) + # Overwrite the "...running" line with final result + print(f"\r", end="") print_task_result(idx, total, result) + # Periodic prune to prevent disk buildup + if prune_interval and idx % prune_interval == 0: + print(f" >> pruning Docker (every {prune_interval})") + _docker_prune(client) return results @@ -1068,6 +1195,15 @@ def main(): action="store_true", help="Use tool-call heuristic instead of pytest scoring", ) + parser.add_argument( + "--prune-interval", + type=int, + default=20, + help=( + "Prune Docker every N tasks to prevent disk buildup" + " (default: 20, 0 to disable)" + ), + ) args = parser.parse_args() logging.basicConfig(level=logging.WARNING) @@ -1087,6 +1223,7 @@ def main(): rebuild=args.rebuild, build_only=args.build_only, skip_tests=args.skip_tests, + prune_interval=args.prune_interval, ) ) elapsed = time.time() - start diff --git a/tests/unittests/tools/test_skill_toolset.py b/tests/unittests/tools/test_skill_toolset.py index 2cbcef9d60..aa80d80124 100644 --- a/tests/unittests/tools/test_skill_toolset.py +++ b/tests/unittests/tools/test_skill_toolset.py @@ -888,3 +888,188 @@ async def test_integration_shell_stderr_only(): ) assert result["status"] == "error" assert "failure" in result["stderr"] + + +# ── Shell JSON envelope parsing (unit tests with mock executor) ── + + +@pytest.mark.asyncio +async def test_shell_json_envelope_parsed(mock_skill1): + """Shell JSON envelope is correctly unpacked by run_async.""" + import json + + envelope = json.dumps({ + "__shell_result__": True, + "stdout": "hello from shell\n", + "stderr": "", + "returncode": 0, + }) + executor = _make_mock_executor(stdout=envelope) + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.RunSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_path": "setup.sh"}, + tool_context=ctx, + ) + assert result["status"] == "success" + assert result["stdout"] == "hello from shell\n" + assert result["stderr"] == "" + + +@pytest.mark.asyncio +async def test_shell_json_envelope_nonzero_returncode(mock_skill1): + """Non-zero returncode in shell envelope sets stderr.""" + import json + + envelope = json.dumps({ + "__shell_result__": True, + "stdout": "", + "stderr": "", + "returncode": 2, + }) + executor = _make_mock_executor(stdout=envelope) + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.RunSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_path": "setup.sh"}, + tool_context=ctx, + ) + assert result["status"] == "error" + assert "Exit code 2" in result["stderr"] + + +@pytest.mark.asyncio +async def test_shell_json_envelope_with_stderr(mock_skill1): + """Shell envelope with both stdout and stderr reports warning.""" + import json + + envelope = json.dumps({ + "__shell_result__": True, + "stdout": "data\n", + "stderr": "deprecation warning\n", + "returncode": 0, + }) + executor = _make_mock_executor(stdout=envelope) + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.RunSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_path": "setup.sh"}, + tool_context=ctx, + ) + assert result["status"] == "warning" + assert result["stdout"] == "data\n" + assert result["stderr"] == "deprecation warning\n" + + +@pytest.mark.asyncio +async def test_shell_json_envelope_timeout(mock_skill1): + """Shell envelope from TimeoutExpired reports error status.""" + import json + + envelope = json.dumps({ + "__shell_result__": True, + "stdout": "partial output\n", + "stderr": "Timed out after 300s", + "returncode": -1, + }) + executor = _make_mock_executor(stdout=envelope) + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.RunSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_path": "setup.sh"}, + tool_context=ctx, + ) + assert result["status"] == "warning" + assert result["stdout"] == "partial output\n" + assert "Timed out" in result["stderr"] + + +@pytest.mark.asyncio +async def test_shell_non_json_stdout_passthrough(mock_skill1): + """Non-JSON shell stdout is passed through without parsing.""" + executor = _make_mock_executor(stdout="plain text output\n") + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.RunSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={"skill_name": "skill1", "script_path": "setup.sh"}, + tool_context=ctx, + ) + assert result["status"] == "success" + assert result["stdout"] == "plain text output\n" + + +# ── input_files packaging ── + + +@pytest.mark.asyncio +async def test_execute_script_input_files_packaged(mock_skill1): + """Verify references, assets, and scripts are packaged as input_files.""" + executor = _make_mock_executor(stdout="ok\n") + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.RunSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + await tool.run_async( + args={"skill_name": "skill1", "script_path": "run.py"}, + tool_context=ctx, + ) + + call_args = executor.execute_code.call_args + code_input = call_args[0][1] + input_files = code_input.input_files + + paths = {f.path for f in input_files} + assert "references/ref1.md" in paths + assert "assets/asset1.txt" in paths + assert "scripts/setup.sh" in paths + assert "scripts/run.py" in paths + assert "scripts/build.rb" in paths + + # Verify content matches + ref_file = next(f for f in input_files if f.path == "references/ref1.md") + assert ref_file.content == "ref content 1" + asset_file = next(f for f in input_files if f.path == "assets/asset1.txt") + assert asset_file.content == "asset content 1" + + +@pytest.mark.asyncio +async def test_execute_script_input_files_working_dir(mock_skill1): + """Verify working_dir is set to '.' for sandboxed execution.""" + executor = _make_mock_executor(stdout="ok\n") + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.RunSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + await tool.run_async( + args={"skill_name": "skill1", "script_path": "run.py"}, + tool_context=ctx, + ) + + call_args = executor.execute_code.call_args + code_input = call_args[0][1] + assert code_input.working_dir == "." + + +# ── Integration: shell non-zero exit ── + + +@pytest.mark.asyncio +async def test_integration_shell_nonzero_exit(): + """Real executor: shell script with non-zero exit via JSON envelope.""" + script = models.Script(src="exit 42") + skill = _make_skill_with_script("test_skill", "fail.sh", script) + toolset = _make_real_executor_toolset([skill]) + tool = skill_toolset.RunSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={ + "skill_name": "test_skill", + "script_path": "fail.sh", + }, + tool_context=ctx, + ) + assert result["status"] == "error" + assert "42" in result["stderr"] or result["stderr"] From 58d1feb6f530ae1b979a1dffaad91670d43a3159 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Mon, 23 Feb 2026 23:05:43 -0800 Subject: [PATCH 25/53] fix: Fix tool name mismatch, empty file handling, and args validation - Fix system instruction referencing 'execute_skill_script' instead of actual tool name 'run_skill_script' - Use 'is not None' checks when packaging input_files so empty files (zero-byte configs, sentinels) are mounted instead of silently dropped - Add defensive validation for 'args' parameter type, returning clear INVALID_ARGS_TYPE error instead of crashing on non-dict input - Add tests for all three findings Co-Authored-By: Claude Opus 4.6 --- src/google/adk/tools/skill_toolset.py | 16 +++- tests/unittests/tools/test_skill_toolset.py | 97 +++++++++++++++++++++ 2 files changed, 109 insertions(+), 4 deletions(-) diff --git a/src/google/adk/tools/skill_toolset.py b/src/google/adk/tools/skill_toolset.py index 799ad6d54a..9efc568c01 100644 --- a/src/google/adk/tools/skill_toolset.py +++ b/src/google/adk/tools/skill_toolset.py @@ -56,7 +56,7 @@ 1. If a skill seems relevant to the current user query, you MUST use the `load_skill` tool with `name=""` to read its full instructions before proceeding. 2. Once you have read the instructions, follow them exactly as documented before replying to the user. For example, If the instruction lists multiple steps, please make sure you complete all of them in order. 3. The `load_skill_resource` tool is for viewing files within a skill's directory (e.g., `references/*`, `assets/*`, `scripts/*`). Do NOT use other tools to access these files. -4. Use `execute_skill_script` to run scripts from a skill's `scripts/` directory. Use `load_skill_resource` to view script content first if needed. +4. Use `run_skill_script` to run scripts from a skill's `scripts/` directory. Use `load_skill_resource` to view script content first if needed. """ @@ -285,6 +285,14 @@ async def run_async( skill_name = args.get("skill_name") script_path = args.get("script_path") script_args = args.get("args", {}) + if not isinstance(script_args, dict): + return { + "error": ( + "'args' must be a JSON object (key-value pairs)," + f" got {type(script_args).__name__}." + ), + "error_code": "INVALID_ARGS_TYPE", + } if not skill_name: return { @@ -340,7 +348,7 @@ async def run_async( # Package ALL skill files for mounting for ref_name in skill.resources.list_references(): content = skill.resources.get_reference(ref_name) - if content: + if content is not None: input_files.append( File( name=os.path.basename(ref_name), @@ -350,7 +358,7 @@ async def run_async( ) for asset_name in skill.resources.list_assets(): content = skill.resources.get_asset(asset_name) - if content: + if content is not None: input_files.append( File( name=os.path.basename(asset_name), @@ -360,7 +368,7 @@ async def run_async( ) for scr_name in skill.resources.list_scripts(): scr = skill.resources.get_script(scr_name) - if scr and scr.src: + if scr is not None and scr.src is not None: input_files.append( File( name=os.path.basename(scr_name), diff --git a/tests/unittests/tools/test_skill_toolset.py b/tests/unittests/tools/test_skill_toolset.py index aa80d80124..66c377a81c 100644 --- a/tests/unittests/tools/test_skill_toolset.py +++ b/tests/unittests/tools/test_skill_toolset.py @@ -1073,3 +1073,100 @@ async def test_integration_shell_nonzero_exit(): ) assert result["status"] == "error" assert "42" in result["stderr"] or result["stderr"] + + +# ── Finding 1: system instruction references correct tool name ── + + +def test_system_instruction_references_run_skill_script(): + """System instruction must reference the actual tool name.""" + assert "run_skill_script" in skill_toolset.DEFAULT_SKILL_SYSTEM_INSTRUCTION + assert ( + "execute_skill_script" + not in skill_toolset.DEFAULT_SKILL_SYSTEM_INSTRUCTION + ) + + +# ── Finding 2: empty files are mounted (not silently dropped) ── + + +@pytest.mark.asyncio +async def test_execute_script_empty_files_mounted(): + """Empty references/assets/scripts should still be packaged.""" + skill = mock.create_autospec(models.Skill, instance=True) + skill.name = "emp" + skill.description = "skill with empty files" + skill.instructions = "test" + fm = mock.create_autospec(models.Frontmatter, instance=True) + fm.name = "emp" + fm.description = "skill with empty files" + skill.frontmatter = fm + skill.resources = mock.MagicMock( + spec=[ + "get_reference", + "get_asset", + "get_script", + "list_references", + "list_assets", + "list_scripts", + ] + ) + skill.resources.list_references.return_value = ["empty.md"] + skill.resources.list_assets.return_value = ["empty.cfg"] + skill.resources.list_scripts.return_value = ["run.py"] + skill.resources.get_reference.side_effect = lambda n: ( + "" if n == "empty.md" else None + ) + skill.resources.get_asset.side_effect = lambda n: ( + "" if n == "empty.cfg" else None + ) + skill.resources.get_script.side_effect = lambda n: ( + models.Script(src="") if n == "run.py" else None + ) + + executor = _make_mock_executor(stdout="ok\n") + toolset = skill_toolset.SkillToolset([skill], code_executor=executor) + tool = skill_toolset.RunSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + await tool.run_async( + args={"skill_name": "emp", "script_path": "run.py"}, + tool_context=ctx, + ) + + call_args = executor.execute_code.call_args + code_input = call_args[0][1] + paths = {f.path for f in code_input.input_files} + assert "references/empty.md" in paths + assert "assets/empty.cfg" in paths + assert "scripts/run.py" in paths + + +# ── Finding 3: invalid args type returns clear error ── + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "bad_args", + [ + "not a dict", + ["a", "list"], + 42, + True, + ], +) +async def test_execute_script_invalid_args_type(mock_skill1, bad_args): + """Non-dict args should return INVALID_ARGS_TYPE, not crash.""" + executor = _make_mock_executor() + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + tool = skill_toolset.RunSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + result = await tool.run_async( + args={ + "skill_name": "skill1", + "script_path": "run.py", + "args": bad_args, + }, + tool_context=ctx, + ) + assert result["error_code"] == "INVALID_ARGS_TYPE" + executor.execute_code.assert_not_called() From 6cf45229ede116c7c0409fd2826c1e074b2146ab Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Mon, 23 Feb 2026 23:22:35 -0800 Subject: [PATCH 26/53] refactor: Remove allowed_tools resolution and global execution lock Remove Feature B (additional_tools parameter and allowed_tools dynamic resolution in get_tools()) to keep the PR focused on script execution. Remove global _execution_lock from UnsafeLocalCodeExecutor as it was an arbitrary side-effect unrelated to the script execution feature. Co-Authored-By: Claude Opus 4.6 --- .../unsafe_local_code_executor.py | 68 +++++++-------- src/google/adk/tools/skill_toolset.py | 86 +------------------ 2 files changed, 34 insertions(+), 120 deletions(-) diff --git a/src/google/adk/code_executors/unsafe_local_code_executor.py b/src/google/adk/code_executors/unsafe_local_code_executor.py index 5bfe064232..08df354497 100644 --- a/src/google/adk/code_executors/unsafe_local_code_executor.py +++ b/src/google/adk/code_executors/unsafe_local_code_executor.py @@ -20,7 +20,6 @@ import os import re import tempfile -import threading from typing import Any from pydantic import Field @@ -33,8 +32,6 @@ logger = logging.getLogger('google_adk.' + __name__) -_execution_lock = threading.Lock() - def _prepare_globals(code: str, globals_: dict[str, Any]) -> None: """Prepare globals for code execution, injecting __name__ if needed.""" @@ -73,39 +70,38 @@ def execute_code( output = '' error = '' - with _execution_lock: - original_cwd = os.getcwd() - try: - # Prepare the execution environment (temp volume) - with tempfile.TemporaryDirectory() as temp_dir: - # Write input files to the temp directory - for f in code_execution_input.input_files: - file_path = os.path.join(temp_dir, f.path or f.name) - os.makedirs(os.path.dirname(file_path), exist_ok=True) - mode = 'wb' if isinstance(f.content, bytes) else 'w' - with open(file_path, mode) as out_f: - out_f.write(f.content) - - # Change working directory if specified - if code_execution_input.working_dir: - exec_dir = os.path.join(temp_dir, code_execution_input.working_dir) - os.makedirs(exec_dir, exist_ok=True) - os.chdir(exec_dir) - else: - os.chdir(temp_dir) - - # Execute the code - globals_ = {} - _prepare_globals(code_execution_input.code, globals_) - stdout = io.StringIO() - with redirect_stdout(stdout): - exec(code_execution_input.code, globals_, globals_) - output = stdout.getvalue() - - except Exception as e: - error = str(e) - finally: - os.chdir(original_cwd) + original_cwd = os.getcwd() + try: + # Prepare the execution environment (temp volume) + with tempfile.TemporaryDirectory() as temp_dir: + # Write input files to the temp directory + for f in code_execution_input.input_files: + file_path = os.path.join(temp_dir, f.path or f.name) + os.makedirs(os.path.dirname(file_path), exist_ok=True) + mode = 'wb' if isinstance(f.content, bytes) else 'w' + with open(file_path, mode) as out_f: + out_f.write(f.content) + + # Change working directory if specified + if code_execution_input.working_dir: + exec_dir = os.path.join(temp_dir, code_execution_input.working_dir) + os.makedirs(exec_dir, exist_ok=True) + os.chdir(exec_dir) + else: + os.chdir(temp_dir) + + # Execute the code + globals_ = {} + _prepare_globals(code_execution_input.code, globals_) + stdout = io.StringIO() + with redirect_stdout(stdout): + exec(code_execution_input.code, globals_, globals_) + output = stdout.getvalue() + + except Exception as e: + error = str(e) + finally: + os.chdir(original_cwd) # Collect the final result. return CodeExecutionResult( diff --git a/src/google/adk/tools/skill_toolset.py b/src/google/adk/tools/skill_toolset.py index 9efc568c01..adb2196feb 100644 --- a/src/google/adk/tools/skill_toolset.py +++ b/src/google/adk/tools/skill_toolset.py @@ -540,7 +540,6 @@ def __init__( *, code_executor: Optional[BaseCodeExecutor] = None, script_timeout: int = _DEFAULT_SCRIPT_TIMEOUT, - additional_tools: Optional[list[Any]] = None, ): """Initializes the SkillToolset. @@ -550,8 +549,6 @@ def __init__( script_timeout: Timeout in seconds for shell script execution via subprocess.run. Defaults to 300 seconds. Does not apply to Python scripts executed via exec(). - additional_tools: Optional list of additional tools (BaseTool, - BaseToolset, or Callables). """ super().__init__() @@ -565,7 +562,6 @@ def __init__( self._skills = {skill.name: skill for skill in skills} self._code_executor = code_executor self._script_timeout = script_timeout - self._additional_tools = additional_tools or [] # Initialize core skill tools self._tools = [ @@ -579,86 +575,8 @@ def __init__( async def get_tools( self, readonly_context: ReadonlyContext | None = None ) -> list[BaseTool]: - """Returns the list of tools in this toolset. - - Dynamically resolves `allowed_tools` from skills against provided - `additional_tools` - and built-in ADK tools. - """ - - import inspect - - from google.adk import tools as built_in_tools - - from .function_tool import FunctionTool - - result = list(self._tools) - - # Collect allowed tools from all skills - allowed_tool_names = set() - for skill in self._list_skills(): - if skill.frontmatter.allowed_tools: - allowed_tool_names.update(skill.frontmatter.allowed_tools) - - if not allowed_tool_names: - return result - - # Resolve additional_tools passed by developer - tools_by_name = {} - for tool_union in self._additional_tools: - if isinstance(tool_union, BaseTool): - tools_by_name[tool_union.name] = tool_union - elif isinstance(tool_union, BaseToolset): - for tool in await tool_union.get_tools(readonly_context): - tools_by_name[tool.name] = tool - elif inspect.isroutine(tool_union): - func_tool = FunctionTool(tool_union) - tools_by_name[func_tool.name] = func_tool - else: - logger.warning("Ignored unsupported additional_tool: %s", tool_union) - - for allowed_tool in allowed_tool_names: - if allowed_tool in tools_by_name: - result.append(tools_by_name[allowed_tool]) - elif hasattr(built_in_tools, allowed_tool): - # Fallback to ADK built-in tools - builtin_obj = getattr(built_in_tools, allowed_tool) - if inspect.isroutine(builtin_obj): - result.append(FunctionTool(builtin_obj)) - elif isinstance(builtin_obj, type) and issubclass( - builtin_obj, BaseTool - ): - try: - # Attempt to instantiate built-in tools that take no arguments - result.append(builtin_obj()) - except TypeError: - logger.warning( - "Could not instantiate built-in tool '%s'. It may require" - " arguments.", - allowed_tool, - ) - elif isinstance(builtin_obj, type) and issubclass( - builtin_obj, BaseToolset - ): - try: - toolset = builtin_obj() - result.extend(await toolset.get_tools(readonly_context)) - except TypeError: - logger.warning( - "Could not instantiate built-in toolset '%s'. It may require" - " arguments.", - allowed_tool, - ) - else: - logger.warning("Unrecognized built-in tool type for %s", allowed_tool) - else: - logger.warning( - "Skill requested tool '%s' which was not provided in" - " additional_tools or found in built-in tools.", - allowed_tool, - ) - - return result + """Returns the list of tools in this toolset.""" + return self._tools def _get_skill(self, name: str) -> models.Skill | None: """Retrieves a skill by name.""" From a9b4f6af49f9c487e738544ae7651ffaceea2028 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Mon, 23 Feb 2026 23:25:38 -0800 Subject: [PATCH 27/53] fix: Restore execution lock to prevent concurrent data races os.chdir() and redirect_stdout() mutate process-global state, so concurrent execute_code() calls without synchronization cause stdout bleed and cwd contamination between executions. Co-Authored-By: Claude Opus 4.6 --- .../unsafe_local_code_executor.py | 68 ++++++++++--------- 1 file changed, 36 insertions(+), 32 deletions(-) diff --git a/src/google/adk/code_executors/unsafe_local_code_executor.py b/src/google/adk/code_executors/unsafe_local_code_executor.py index 08df354497..5bfe064232 100644 --- a/src/google/adk/code_executors/unsafe_local_code_executor.py +++ b/src/google/adk/code_executors/unsafe_local_code_executor.py @@ -20,6 +20,7 @@ import os import re import tempfile +import threading from typing import Any from pydantic import Field @@ -32,6 +33,8 @@ logger = logging.getLogger('google_adk.' + __name__) +_execution_lock = threading.Lock() + def _prepare_globals(code: str, globals_: dict[str, Any]) -> None: """Prepare globals for code execution, injecting __name__ if needed.""" @@ -70,38 +73,39 @@ def execute_code( output = '' error = '' - original_cwd = os.getcwd() - try: - # Prepare the execution environment (temp volume) - with tempfile.TemporaryDirectory() as temp_dir: - # Write input files to the temp directory - for f in code_execution_input.input_files: - file_path = os.path.join(temp_dir, f.path or f.name) - os.makedirs(os.path.dirname(file_path), exist_ok=True) - mode = 'wb' if isinstance(f.content, bytes) else 'w' - with open(file_path, mode) as out_f: - out_f.write(f.content) - - # Change working directory if specified - if code_execution_input.working_dir: - exec_dir = os.path.join(temp_dir, code_execution_input.working_dir) - os.makedirs(exec_dir, exist_ok=True) - os.chdir(exec_dir) - else: - os.chdir(temp_dir) - - # Execute the code - globals_ = {} - _prepare_globals(code_execution_input.code, globals_) - stdout = io.StringIO() - with redirect_stdout(stdout): - exec(code_execution_input.code, globals_, globals_) - output = stdout.getvalue() - - except Exception as e: - error = str(e) - finally: - os.chdir(original_cwd) + with _execution_lock: + original_cwd = os.getcwd() + try: + # Prepare the execution environment (temp volume) + with tempfile.TemporaryDirectory() as temp_dir: + # Write input files to the temp directory + for f in code_execution_input.input_files: + file_path = os.path.join(temp_dir, f.path or f.name) + os.makedirs(os.path.dirname(file_path), exist_ok=True) + mode = 'wb' if isinstance(f.content, bytes) else 'w' + with open(file_path, mode) as out_f: + out_f.write(f.content) + + # Change working directory if specified + if code_execution_input.working_dir: + exec_dir = os.path.join(temp_dir, code_execution_input.working_dir) + os.makedirs(exec_dir, exist_ok=True) + os.chdir(exec_dir) + else: + os.chdir(temp_dir) + + # Execute the code + globals_ = {} + _prepare_globals(code_execution_input.code, globals_) + stdout = io.StringIO() + with redirect_stdout(stdout): + exec(code_execution_input.code, globals_, globals_) + output = stdout.getvalue() + + except Exception as e: + error = str(e) + finally: + os.chdir(original_cwd) # Collect the final result. return CodeExecutionResult( From 04198b064f96d41829e6ad508ed02031757b432f Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Mon, 23 Feb 2026 23:30:20 -0800 Subject: [PATCH 28/53] refactor: Only sandbox executor when input_files or working_dir is set Preserve original bare exec() path for existing callers that pass no input_files and no working_dir. The temp directory, file mounting, chdir, and execution lock now only activate when needed by the new script execution feature. Co-Authored-By: Claude Opus 4.6 --- .../unsafe_local_code_executor.py | 77 +++++++++++-------- 1 file changed, 47 insertions(+), 30 deletions(-) diff --git a/src/google/adk/code_executors/unsafe_local_code_executor.py b/src/google/adk/code_executors/unsafe_local_code_executor.py index 5bfe064232..90d9d67dc7 100644 --- a/src/google/adk/code_executors/unsafe_local_code_executor.py +++ b/src/google/adk/code_executors/unsafe_local_code_executor.py @@ -73,39 +73,56 @@ def execute_code( output = '' error = '' - with _execution_lock: - original_cwd = os.getcwd() - try: - # Prepare the execution environment (temp volume) - with tempfile.TemporaryDirectory() as temp_dir: - # Write input files to the temp directory - for f in code_execution_input.input_files: - file_path = os.path.join(temp_dir, f.path or f.name) - os.makedirs(os.path.dirname(file_path), exist_ok=True) - mode = 'wb' if isinstance(f.content, bytes) else 'w' - with open(file_path, mode) as out_f: - out_f.write(f.content) - - # Change working directory if specified - if code_execution_input.working_dir: - exec_dir = os.path.join(temp_dir, code_execution_input.working_dir) - os.makedirs(exec_dir, exist_ok=True) - os.chdir(exec_dir) - else: - os.chdir(temp_dir) - - # Execute the code - globals_ = {} - _prepare_globals(code_execution_input.code, globals_) - stdout = io.StringIO() - with redirect_stdout(stdout): - exec(code_execution_input.code, globals_, globals_) - output = stdout.getvalue() + needs_sandbox = ( + code_execution_input.input_files + or code_execution_input.working_dir + ) + if needs_sandbox: + with _execution_lock: + original_cwd = os.getcwd() + try: + with tempfile.TemporaryDirectory() as temp_dir: + # Write input files to the temp directory + for f in code_execution_input.input_files: + file_path = os.path.join(temp_dir, f.path or f.name) + os.makedirs(os.path.dirname(file_path), exist_ok=True) + mode = 'wb' if isinstance(f.content, bytes) else 'w' + with open(file_path, mode) as out_f: + out_f.write(f.content) + + # Change working directory + if code_execution_input.working_dir: + exec_dir = os.path.join( + temp_dir, code_execution_input.working_dir + ) + os.makedirs(exec_dir, exist_ok=True) + os.chdir(exec_dir) + else: + os.chdir(temp_dir) + + globals_ = {} + _prepare_globals(code_execution_input.code, globals_) + stdout = io.StringIO() + with redirect_stdout(stdout): + exec(code_execution_input.code, globals_, globals_) + output = stdout.getvalue() + + except Exception as e: + error = str(e) + finally: + os.chdir(original_cwd) + else: + # Original path: no temp dir, no chdir, no lock needed + try: + globals_ = {} + _prepare_globals(code_execution_input.code, globals_) + stdout = io.StringIO() + with redirect_stdout(stdout): + exec(code_execution_input.code, globals_, globals_) + output = stdout.getvalue() except Exception as e: error = str(e) - finally: - os.chdir(original_cwd) # Collect the final result. return CodeExecutionResult( From 3532cefb00ad2f4bc7e0f22d99069e48fab23c76 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Mon, 23 Feb 2026 23:33:21 -0800 Subject: [PATCH 29/53] fix: Hold execution lock for both sandbox and plain paths redirect_stdout mutates process-global sys.stdout, so both the sandbox path (tempdir + chdir) and the plain exec() path must be serialized. The previous commit only locked the sandbox branch, allowing concurrent plain calls to bleed stdout with sandbox calls. Co-Authored-By: Claude Opus 4.6 --- .../unsafe_local_code_executor.py | 28 ++++++++++--------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/src/google/adk/code_executors/unsafe_local_code_executor.py b/src/google/adk/code_executors/unsafe_local_code_executor.py index 90d9d67dc7..4803a15ee0 100644 --- a/src/google/adk/code_executors/unsafe_local_code_executor.py +++ b/src/google/adk/code_executors/unsafe_local_code_executor.py @@ -78,8 +78,10 @@ def execute_code( or code_execution_input.working_dir ) - if needs_sandbox: - with _execution_lock: + # Lock is required for both paths: redirect_stdout mutates + # process-global sys.stdout, and the sandbox path also mutates cwd. + with _execution_lock: + if needs_sandbox: original_cwd = os.getcwd() try: with tempfile.TemporaryDirectory() as temp_dir: @@ -112,17 +114,17 @@ def execute_code( error = str(e) finally: os.chdir(original_cwd) - else: - # Original path: no temp dir, no chdir, no lock needed - try: - globals_ = {} - _prepare_globals(code_execution_input.code, globals_) - stdout = io.StringIO() - with redirect_stdout(stdout): - exec(code_execution_input.code, globals_, globals_) - output = stdout.getvalue() - except Exception as e: - error = str(e) + else: + # Original path: no temp dir, no chdir + try: + globals_ = {} + _prepare_globals(code_execution_input.code, globals_) + stdout = io.StringIO() + with redirect_stdout(stdout): + exec(code_execution_input.code, globals_, globals_) + output = stdout.getvalue() + except Exception as e: + error = str(e) # Collect the final result. return CodeExecutionResult( From b65d35346b7186c16f205814c4779ecdeadf8ea4 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Mon, 23 Feb 2026 23:35:18 -0800 Subject: [PATCH 30/53] test: Add concurrency test for mixed sandbox/plain execution Runs 50 iterations of concurrent sandbox (with input_files + working_dir) and plain (bare exec) calls to verify no stdout bleed between them. Protects the execution lock fix from future regressions. Co-Authored-By: Claude Opus 4.6 --- .../test_unsafe_local_code_executor.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/unittests/code_executors/test_unsafe_local_code_executor.py b/tests/unittests/code_executors/test_unsafe_local_code_executor.py index f8d5f496a8..50f31a9e80 100644 --- a/tests/unittests/code_executors/test_unsafe_local_code_executor.py +++ b/tests/unittests/code_executors/test_unsafe_local_code_executor.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import concurrent.futures import textwrap from unittest.mock import MagicMock @@ -19,6 +20,7 @@ from google.adk.agents.invocation_context import InvocationContext from google.adk.code_executors.code_execution_utils import CodeExecutionInput from google.adk.code_executors.code_execution_utils import CodeExecutionResult +from google.adk.code_executors.code_execution_utils import File from google.adk.code_executors.unsafe_local_code_executor import UnsafeLocalCodeExecutor from google.adk.sessions.base_session_service import BaseSessionService from google.adk.sessions.session import Session @@ -121,3 +123,40 @@ def run(): assert result.stderr == "" assert result.stdout == "hi ada\n" + + def test_concurrent_sandbox_and_plain_no_stdout_bleed( + self, mock_invocation_context: InvocationContext + ): + """Concurrent sandbox and plain calls must not mix stdout.""" + executor = UnsafeLocalCodeExecutor() + plain_input = CodeExecutionInput(code='print("PLAIN")') + sandbox_input = CodeExecutionInput( + code='import time; time.sleep(0.01); print("SANDBOX")', + input_files=[ + File(name="dummy.txt", content="data"), + ], + working_dir=".", + ) + + errors = [] + for _ in range(50): + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + f_sandbox = pool.submit( + executor.execute_code, + mock_invocation_context, + sandbox_input, + ) + f_plain = pool.submit( + executor.execute_code, + mock_invocation_context, + plain_input, + ) + r_sandbox = f_sandbox.result() + r_plain = f_plain.result() + + if "PLAIN" in r_sandbox.stdout or "SANDBOX" in r_plain.stdout: + errors.append(f"sandbox={r_sandbox.stdout!r} plain={r_plain.stdout!r}") + + assert not errors, ( + f"stdout bleed detected in {len(errors)}/50 iterations: " + errors[0] + ) From 4e49596b37cb2ba8e5366abc4b4241baf1c8da93 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Mon, 23 Feb 2026 23:39:03 -0800 Subject: [PATCH 31/53] test: Assert plain-path cwd stays original in concurrency test Strengthen the mixed-concurrency test to also verify that the plain execution path observes the original working directory, not a sandbox temp dir leaked from a concurrent sandbox call. Co-Authored-By: Claude Opus 4.6 --- .../test_unsafe_local_code_executor.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/tests/unittests/code_executors/test_unsafe_local_code_executor.py b/tests/unittests/code_executors/test_unsafe_local_code_executor.py index 50f31a9e80..76dfd47da7 100644 --- a/tests/unittests/code_executors/test_unsafe_local_code_executor.py +++ b/tests/unittests/code_executors/test_unsafe_local_code_executor.py @@ -13,6 +13,7 @@ # limitations under the License. import concurrent.futures +import os import textwrap from unittest.mock import MagicMock @@ -129,9 +130,14 @@ def test_concurrent_sandbox_and_plain_no_stdout_bleed( ): """Concurrent sandbox and plain calls must not mix stdout.""" executor = UnsafeLocalCodeExecutor() - plain_input = CodeExecutionInput(code='print("PLAIN")') + original_cwd = os.getcwd() + plain_input = CodeExecutionInput( + code='import os; print("PLAIN:" + os.getcwd())' + ) sandbox_input = CodeExecutionInput( - code='import time; time.sleep(0.01); print("SANDBOX")', + code=( + 'import os, time; time.sleep(0.01); print("SANDBOX:" + os.getcwd())' + ), input_files=[ File(name="dummy.txt", content="data"), ], @@ -157,6 +163,11 @@ def test_concurrent_sandbox_and_plain_no_stdout_bleed( if "PLAIN" in r_sandbox.stdout or "SANDBOX" in r_plain.stdout: errors.append(f"sandbox={r_sandbox.stdout!r} plain={r_plain.stdout!r}") + # Plain-path cwd must remain the original cwd, not a temp dir + plain_cwd = r_plain.stdout.strip().split(":", 1)[1] + if plain_cwd != original_cwd: + errors.append(f"plain cwd={plain_cwd!r} expected={original_cwd!r}") + assert not errors, ( - f"stdout bleed detected in {len(errors)}/50 iterations: " + errors[0] + f"bleed detected in {len(errors)}/50 iterations: " + errors[0] ) From 0f4165a00301f6b2a7333bc60519b6db7f914c13 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 00:23:02 -0800 Subject: [PATCH 32/53] docs(code-executor): align design doc with RunSkillScriptTool and current models --- docs/design/code_executor_enhancements.md | 50 ++++++++++++++--------- 1 file changed, 31 insertions(+), 19 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index 6630a4b132..fb3b7f9fac 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -3,7 +3,7 @@ **Authors:** haiyuancao, Claude Code **Date:** 2026-02-21 **Status:** Draft -**Tracking:** Related to PR #4575 (ExecuteSkillScriptTool) +**Tracking:** Related to PR #4575 (RunSkillScriptTool) --- @@ -16,7 +16,7 @@ limit production readiness: 1. **No uniform timeout enforcement** — Only `GkeCodeExecutor` has a `timeout_seconds` field. All other executors can hang indefinitely on - malicious, buggy, or slow code. The `ExecuteSkillScriptTool` works + malicious, buggy, or slow code. The `RunSkillScriptTool` works around this for shell scripts by embedding `subprocess.run(timeout=N)` in generated code, but this is a workaround, not a systemic solution. @@ -74,8 +74,11 @@ The following are explicitly **out of scope** for this design: - Code is appended to stateful history **only after** successful execution. A failing code block is never replayed. - Executor instances are not thread-safe unless documented otherwise. - Concurrent `execute_code()` calls on the same instance require external - synchronization. + Concurrent `execute_code()` calls on the same instance may require + external synchronization. +- `UnsafeLocalCodeExecutor` is a special case: it currently serializes + `execute_code()` with an internal lock because `redirect_stdout` and + `os.chdir()` mutate process-global state. --- @@ -112,11 +115,19 @@ class BaseCodeExecutor(BaseModel): ### 3.3 Data Model ```python +@dataclasses.dataclass(frozen=True) +class File: + name: str + content: str | bytes + mime_type: str = 'text/plain' + path: Optional[str] = None # e.g. "scripts/run.py" + @dataclasses.dataclass class CodeExecutionInput: code: str input_files: list[File] = field(default_factory=list) execution_id: Optional[str] = None # For stateful execution + working_dir: Optional[str] = None # e.g. "." @dataclasses.dataclass class CodeExecutionResult: @@ -135,7 +146,7 @@ The primary consumer is `_code_execution.py` in the LLM flows layer: 3. **Stateful support**: Uses `execution_id` (from `CodeExecutorContext`) to maintain state across calls when `stateful=True` -`ExecuteSkillScriptTool` is a secondary consumer that calls +`RunSkillScriptTool` is a secondary consumer that calls `execute_code()` directly with generated Python code wrapping skill scripts. --- @@ -201,7 +212,7 @@ timeout = ( **Why per-invocation + executor default:** - Backward compatible — existing code that doesn't set it works unchanged - Safe for shared executors — no global mutable state -- Callers can override per-call (e.g., `ExecuteSkillScriptTool` sets +- Callers can override per-call (e.g., `RunSkillScriptTool` sets `script_timeout`, LLM flows use a different default) - Executor subclasses can define their own defaults @@ -490,10 +501,10 @@ is not applicable. Use the same thread+join pattern as | 5 | Migrate `GkeCodeExecutor.timeout_seconds` to `default_timeout_seconds` | None | | 6 | Add client-side timeout to remote executors | Low | -### 4.4 Impact on `ExecuteSkillScriptTool` +### 4.4 Impact on `RunSkillScriptTool` Once `BaseCodeExecutor` has native timeout support, the -`ExecuteSkillScriptTool` can optionally delegate timeout enforcement to +`RunSkillScriptTool` can optionally delegate timeout enforcement to the executor rather than embedding it in generated shell wrapper code. However, the shell wrapper timeout (`subprocess.run(timeout=N)`) should be kept as defense-in-depth — it catches the subprocess even if the @@ -757,16 +768,17 @@ identify stateful sessions. For `ContainerCodeExecutor`: This aligns with how `VertexAiCodeExecutor` uses `session_id`. -**Gap: `ExecuteSkillScriptTool` does not wire `execution_id`.** +**Gap: `RunSkillScriptTool` does not wire `execution_id`.** -Currently, `ExecuteSkillScriptTool.run_async()` creates -`CodeExecutionInput(code=prepared_code)` without setting `execution_id`. +Currently, `RunSkillScriptTool.run_async()` creates +`CodeExecutionInput(code=..., input_files=..., working_dir='.')` without +setting `execution_id`. This means all skill script executions share the same (default) namespace in a stateful executor, with no isolation between different skills or invocations. **Action items:** -1. `ExecuteSkillScriptTool` should generate a **session-stable** +1. `RunSkillScriptTool` should generate a **session-stable** `execution_id` scoped to skill + agent. The key must persist across turns so that stateful code history is preserved: ```python @@ -820,7 +832,7 @@ This is a critical security concern when executing: | LLM generates malicious code | Full host compromise | None | | Skill script reads secrets | Data exfiltration | None (documented warning only) | | Infinite loop / fork bomb | DoS / resource exhaustion | None (no timeout) | -| `sys.exit()` in script | Process termination | Partial (`SystemExit` catch in `ExecuteSkillScriptTool`) | +| `sys.exit()` in script | Process termination | Partial (`SystemExit` catch in `RunSkillScriptTool`) | | Network exfiltration | Data leak | None | | File system manipulation | Data loss / corruption | None | @@ -1187,7 +1199,7 @@ class CodeExecutionInput: | `restrict_builtins` on `UnsafeLocalCodeExecutor` | Yes (default `False`) | No | | Default image for `ContainerCodeExecutor` | Breaking (currently requires image/docker_path) | Minor | -### 7.3 Impact on `ExecuteSkillScriptTool` +### 7.3 Impact on `RunSkillScriptTool` | Feature | Current workaround | After enhancements | |---------|-------------------|-------------------| @@ -1209,7 +1221,7 @@ is new and has no tests yet. | Category | Approach | New tests needed | |----------|----------|-----------------| | Unit tests | Mock-based tests per executor | **Add `test_container_code_executor.py`**, add `test_local_sandbox_code_executor.py` | -| Integration tests | Real executor tests (like `ExecuteSkillScriptTool` integration tests) | Add Docker-based container tests (CI-gated) | +| Integration tests | Real executor tests (like `RunSkillScriptTool` integration tests) | Add Docker-based container tests (CI-gated) | | Timeout tests | Scripts with `time.sleep()` to verify enforcement | Per-executor timeout tests | | Timeout kill fallback | Verify `PermissionError` from `os.kill` triggers container restart | Mock `os.kill` to raise `PermissionError`, assert `container.restart()` called and `CodeExecutionResult.stderr` contains timeout message | | Timeout kill success | Verify `os.kill(host_pid)` path when permitted | Mock `exec_inspect` to return PID, assert `os.kill` called with correct signal | @@ -1239,7 +1251,7 @@ is new and has no tests yet. by users or by higher-level retry logic. 6. Migrate `GkeCodeExecutor.timeout_seconds` to `default_timeout_seconds` 7. Add timeout tests for each executor -8. Update `ExecuteSkillScriptTool` to set per-invocation timeout via +8. Update `RunSkillScriptTool` to set per-invocation timeout via `CodeExecutionInput.timeout_seconds` ### Phase 2: Stateful Container (5-8 days) @@ -1254,7 +1266,7 @@ Implement Option A (persistent process) directly, as recommended in read output, detect crash/restart) 4. Add `execution_id`-based session isolation (one REPL per `execution_id`) -5. Wire `execution_id` in `ExecuteSkillScriptTool` +5. Wire `execution_id` in `RunSkillScriptTool` 6. Add `reset_state()` method (kills and restarts the REPL) 7. Add stateful execution tests (variable persistence, crash recovery, `execution_id` isolation) @@ -1314,6 +1326,6 @@ Implement Option A (persistent process) directly, as recommended in - [GkeCodeExecutor](../../src/google/adk/code_executors/gke_code_executor.py) - [VertexAiCodeExecutor](../../src/google/adk/code_executors/vertex_ai_code_executor.py) - [AgentEngineSandboxCodeExecutor](../../src/google/adk/code_executors/agent_engine_sandbox_code_executor.py) -- [ExecuteSkillScriptTool](../../src/google/adk/tools/skill_toolset.py) +- [RunSkillScriptTool](../../src/google/adk/tools/skill_toolset.py) - [Code Execution Flow](../../src/google/adk/flows/llm_flows/_code_execution.py) -- [PR #4575 — ExecuteSkillScriptTool](https://github.com/google/adk-python/pull/4575) +- [PR #4575 — RunSkillScriptTool](https://github.com/google/adk-python/pull/4575) From 97397ba8fe86f1b8c39aa7708f979774cf709d32 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 00:34:41 -0800 Subject: [PATCH 33/53] docs(code-executor): prioritize run-script roadmap and tool contract hardening --- docs/design/code_executor_enhancements.md | 150 +++++++++++++++++++--- 1 file changed, 135 insertions(+), 15 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index fb3b7f9fac..ea70f1ef2d 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -38,6 +38,29 @@ to the existing API. --- +## 1.1 Prioritized Next Steps for `RunSkillScriptTool` + +The three major proposals in this doc remain valid, but for improving +`RunSkillScriptTool` specifically, we should prioritize them alongside +additional high-impact tool-contract work. + +| Rank | Priority | Area | Why it matters now | +|------|----------|------|--------------------| +| 1 | P0 | Uniform timeout support (Proposal 1) | Python script execution can still hang indefinitely without executor-level timeout controls. | +| 2 | P0 | Security hardening (Proposal 3) | `UnsafeLocalCodeExecutor` remains unsafe for untrusted scripts and is a major deployment risk. | +| 3 | P1 | Structured `RunSkillScriptTool` result contract | Agents need explicit machine-readable execution metadata (`return_code`, timeout flag), not just inferred status from `stdout/stderr`. | +| 4 | P1 | Propagate `output_files` / artifacts | Script-generated outputs are currently dropped by tool responses, limiting practical utility. | +| 5 | P1 | Strengthen script argument contract | Argument normalization rules are underspecified, which leads to fragile calls and inconsistent behavior. | +| 6 | P1 | Wire `execution_id` in `RunSkillScriptTool` | Needed for predictable namespace isolation and future stateful execution compatibility. | +| 7 | P2 | Stateful `ContainerCodeExecutor` (Proposal 2) | Valuable, but more complex and less urgent than reliability/safety/tool-contract gaps above. | + +**Interpretation for implementation planning:** +- P0 items are required for production reliability/safety. +- P1 items directly improve agent correctness and tool usability. +- P2 items are strategic enhancements once P0/P1 are complete. + +--- + ## 2. Non-Goals & Invariants The following are explicitly **out of scope** for this design: @@ -1230,12 +1253,99 @@ is new and has no tests yet. | Security tests | Scripts attempting blocked operations | `restrict_builtins` bypass attempts, env var leakage | | Stateful tests | Multi-call sequences verifying variable persistence | Append-after-success, failure-does-not-poison, `execution_id` isolation | | Stateful crash recovery | Verify error returned on REPL/container crash | Kill REPL mid-execution, assert error indicates state loss | +| Tool contract tests | Validate structured `RunSkillScriptTool` response schema | Assert `return_code`, `timed_out`, `status`, and error envelope consistency | +| Output propagation tests | Verify executor `output_files` are surfaced by tool | Assert tool response includes generated files/metadata | +| Args normalization tests | Verify deterministic mapping from JSON args to argv | Cover booleans, lists, positional args, and invalid types | +| `execution_id` wiring tests | Verify stable and scoped `execution_id` generation | Assert per-session/per-skill isolation and persistence semantics | + +### 7.5 High-Priority `RunSkillScriptTool` Contract Enhancements + +These improvements are additive and can be shipped before stateful container +support. + +#### 7.5.1 Structured Execution Result Schema + +Current output is largely free-form (`stdout`, `stderr`, derived `status`). +Add explicit structured fields: + +```python +{ + "skill_name": "...", + "script_path": "...", + "status": "success|warning|error", + "return_code": int | None, + "timed_out": bool, + "stdout": str, + "stderr": str, + "output_files": [...], # see §7.5.2 +} +``` + +Guidance: +- `status` should be derived from structured fields (`return_code`, + `timed_out`, and stream presence), not treated as the source of truth. +- Tool-level validation/configuration errors should continue using explicit + `error_code` values. + +#### 7.5.2 Propagate `output_files` and Artifact Metadata + +`CodeExecutionResult.output_files` should be surfaced in the tool response. +This is critical for scripts that generate reports, transformed datasets, or +intermediate artifacts. + +Minimum expected shape: + +```python +output_files = [ + { + "name": str, + "mime_type": str | None, + # optional future fields: + # "artifact_id": str, + # "path": str, + } +] +``` + +#### 7.5.3 `execution_id` Wiring in `RunSkillScriptTool` + +Even before full stateful container support, wire a deterministic +`execution_id` to avoid ambiguous namespaces in stateful-capable executors. + +Recommended key shape: + +```python +execution_id = ( + f"skill:{skill_name}:session:{session_id}:agent:{agent_name}" +) +``` + +Rules: +- Stable across turns within the same session. +- Scoped by skill and agent. +- Never derived from `invocation_id` (too short-lived). + +#### 7.5.4 Script Args Normalization Contract + +Define and document deterministic mapping from JSON args to CLI argv. + +Suggested rules: +- `str|int|float` -> `--key value` +- `true` -> `--key` +- `false|None` -> omit +- `list[...]` -> repeated `--key value` entries +- Optional reserved key for positional args (for example, `_positional`) +- Reject nested objects with explicit validation error + +This reduces LLM-side ambiguity and improves replay/debug stability. --- ## 8. Implementation Roadmap -### Phase 1: Timeout (3-4 days) +### 8.1 Priority-Ordered Plan + +#### Phase 1 (P0): Timeout Foundation (3-4 days) 1. Add `timeout_seconds: Optional[int] = None` to `CodeExecutionInput` 2. Add `default_timeout_seconds: Optional[int] = None` to @@ -1254,7 +1364,29 @@ is new and has no tests yet. 8. Update `RunSkillScriptTool` to set per-invocation timeout via `CodeExecutionInput.timeout_seconds` -### Phase 2: Stateful Container (5-8 days) +#### Phase 2 (P1): `RunSkillScriptTool` Contract Hardening (2-3 days) + +1. Add structured response fields: `return_code`, `timed_out`, + schema-stable status semantics +2. Surface executor `output_files` in tool output +3. Define and implement args normalization contract +4. Add deterministic `execution_id` wiring for tool calls +5. Add tool-level contract tests (schema, args, output propagation, + `execution_id` isolation behavior) + +#### Phase 3 (P0): Security Hardening (5-7 days) + +1. Add `SecurityWarning` to `UnsafeLocalCodeExecutor` +2. Add `restrict_builtins` option (documented as best-effort friction) +3. Implement `LocalSandboxCodeExecutor` (using `process_group`, not + `preexec_fn`) +4. Add digest-pinned default image to `ContainerCodeExecutor` +5. Add network isolation defaults to `ContainerCodeExecutor` +6. Create official `adk-code-executor` Docker image (versioned tags) +7. Update all samples to recommend secure executors +8. Add security-focused tests + +#### Phase 4 (P2): Stateful Container (5-8 days) Implement Option A (persistent process) directly, as recommended in §5.2.2. This avoids the side-effect replay problems of Option C. @@ -1272,19 +1404,7 @@ Implement Option A (persistent process) directly, as recommended in `execution_id` isolation) 8. Update samples and documentation -### Phase 3: Security Hardening (5-7 days) - -1. Add `SecurityWarning` to `UnsafeLocalCodeExecutor` -2. Add `restrict_builtins` option (documented as best-effort friction) -3. Implement `LocalSandboxCodeExecutor` (using `process_group`, not - `preexec_fn`) -4. Add digest-pinned default image to `ContainerCodeExecutor` -5. Add network isolation defaults to `ContainerCodeExecutor` -6. Create official `adk-code-executor` Docker image (versioned tags) -7. Update all samples to recommend secure executors -8. Add security-focused tests - -### Total estimated effort: 11-16 days +### Total estimated effort: 15-22 days --- From 4e69ace75dbed8203cbc13c6ecc2ae82d3644814 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 00:51:03 -0800 Subject: [PATCH 34/53] docs(code-executor): fix 18 inaccuracies in design doc vs implementation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review the design doc against the actual RunSkillScriptTool and executor implementations, fixing factual inaccuracies, missing details, and underspecified areas: - §3.1: UnsafeLocalCodeExecutor isolation is partial (temp-dir sandbox), not "None"; add BuiltInCodeExecutor clarification - §3.2: Add optimize_data_file and actual default delimiter values - §3.3: Document binary content handling and _prepare_globals helper - §3.4: Document asyncio.to_thread() wrapping (not direct call) - §3.4.1 (new): Add RunSkillScriptTool implementation details: executor resolution chain, script_timeout param (shell-only), scripts/ prefix normalization, resource packaging (incl. empty files), SystemExit handling, error truncation >200 chars, shell JSON envelope edge cases, duplicate skill name validation, and process_llm_request system instruction injection - §4.2.2: Add _execution_lock interaction analysis for timeout thread - §4.4: Expand Python timeout gap (zero protection at any layer) - §6.2: Update threat table with current mitigations - §7.1: Add missing fields to combined API changes snippet - §7.3: Expand impact table with output_files, system instructions, error truncation rows - §7.5.1: Document status derivation asymmetry (shell vs Python) - §7.5.4: Document current str(v) behavior before proposed rules Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 269 ++++++++++++++++++++-- 1 file changed, 248 insertions(+), 21 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index ea70f1ef2d..18f1e30c24 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -1,7 +1,7 @@ # ADK Code Executor Enhancements — Design Document **Authors:** haiyuancao, Claude Code -**Date:** 2026-02-21 +**Date:** 2026-02-24 **Status:** Draft **Tracking:** Related to PR #4575 (RunSkillScriptTool) @@ -111,21 +111,39 @@ The following are explicitly **out of scope** for this design: | Executor | Stateful | Timeout | Isolation | Dependencies | |----------|----------|---------|-----------|-------------| -| `UnsafeLocalCodeExecutor` | No (frozen) | None | None | None | +| `UnsafeLocalCodeExecutor` | No (frozen) | None | Partial (temp-dir sandbox when `input_files` or `working_dir` set; no isolation otherwise) | None | | `ContainerCodeExecutor` | No (frozen) | None | Docker container | `docker` | | `GkeCodeExecutor` | No (ephemeral) | `timeout_seconds=300` | gVisor sandbox | `kubernetes` | | `VertexAiCodeExecutor` | Allowed | None | Vertex AI Extension | `vertexai` | | `AgentEngineSandboxCodeExecutor` | Allowed | None | Vertex AI Sandbox | `vertexai` | -| `BuiltInCodeExecutor` | N/A | N/A | Gemini model | `google-genai` | +| `BuiltInCodeExecutor` | N/A (delegates to Gemini model's built-in code execution) | N/A (Gemini-internal) | Gemini model | `google-genai` | + +**Note on `UnsafeLocalCodeExecutor` partial isolation:** When `input_files` +or `working_dir` is set in `CodeExecutionInput`, the executor creates a +`tempfile.TemporaryDirectory`, writes all input files with preserved +relative paths (e.g., `references/data.csv`), and `os.chdir()`s into it +before calling `exec()`. This provides filesystem-level isolation for the +script's view of its working directory, but does **not** restrict access +to the rest of the host filesystem, environment variables, or network. +The temp directory is cleaned up after execution, and the original working +directory is restored in a `finally` block. Both the sandbox and plain +paths hold a process-global `_execution_lock` because `redirect_stdout` +mutates `sys.stdout`. ### 3.2 Base Class Contract ```python class BaseCodeExecutor(BaseModel): + optimize_data_file: bool = False stateful: bool = False error_retry_attempts: int = 2 - code_block_delimiters: List[tuple[str, str]] - execution_result_delimiters: tuple[str, str] + code_block_delimiters: List[tuple[str, str]] = [ + ('```tool_code\n', '\n```'), + ('```python\n', '\n```'), + ] + execution_result_delimiters: tuple[str, str] = ( + '```tool_output\n', '\n```' + ) @abc.abstractmethod def execute_code( @@ -141,7 +159,7 @@ class BaseCodeExecutor(BaseModel): @dataclasses.dataclass(frozen=True) class File: name: str - content: str | bytes + content: str | bytes # Text or binary; executor writes 'w'/'wb' mime_type: str = 'text/plain' path: Optional[str] = None # e.g. "scripts/run.py" @@ -159,6 +177,23 @@ class CodeExecutionResult: output_files: list[File] = field(default_factory=list) ``` +**Binary content handling:** `File.content` accepts both `str` and +`bytes`. When `UnsafeLocalCodeExecutor` writes input files to the +temp-dir sandbox, it selects file mode based on content type: +`'wb'` for `bytes`, `'w'` for `str`. This is relevant for +`RunSkillScriptTool`, which packages skill resources as `File` objects +— references and assets may contain binary content (e.g., images, +serialized data). + +**`_prepare_globals` helper in `UnsafeLocalCodeExecutor`:** The executor +has a `_prepare_globals()` function that scans the code for +`if __name__ == '__main__'` patterns and injects `__name__ = '__main__'` +into the execution globals. This interacts with `RunSkillScriptTool`'s +Python wrapper, which uses `runpy.run_path(script_path, +run_name='__main__')` — `runpy` sets `__name__` independently, so the +executor's `_prepare_globals` applies to the wrapper code's outer scope +while `runpy` sets it for the script's scope. + ### 3.4 How Executors Are Used The primary consumer is `_code_execution.py` in the LLM flows layer: @@ -169,8 +204,76 @@ The primary consumer is `_code_execution.py` in the LLM flows layer: 3. **Stateful support**: Uses `execution_id` (from `CodeExecutorContext`) to maintain state across calls when `stateful=True` -`RunSkillScriptTool` is a secondary consumer that calls -`execute_code()` directly with generated Python code wrapping skill scripts. +`RunSkillScriptTool` is a secondary consumer that wraps +`execute_code()` in `asyncio.to_thread()` to avoid blocking the async +event loop: + +```python +result = await asyncio.to_thread( + code_executor.execute_code, + tool_context._invocation_context, + CodeExecutionInput( + code=code, + input_files=input_files, + working_dir=".", + ), +) +``` + +This is architecturally significant: the tool is always called from an +async context (`run_async`), but all `BaseCodeExecutor.execute_code()` +implementations are synchronous. The `to_thread()` bridge ensures the +executor's blocking call (which may involve `exec()`, Docker API calls, +or HTTP requests) does not starve the event loop. + +#### 3.4.1 `RunSkillScriptTool` Current Implementation Details + +Key behaviors of the current implementation (`skill_toolset.py`) that +inform the proposals in this document: + +- **Executor resolution chain:** Toolset-level `_code_executor` (highest + priority) → `agent.code_executor` attribute → `NO_CODE_EXECUTOR` error. +- **`script_timeout` parameter:** `SkillToolset.__init__` accepts + `script_timeout: int` (default 300s). This timeout is embedded in + generated shell wrapper code via `subprocess.run(timeout=N)`. It does + **not** apply to Python scripts executed via `runpy.run_path()` — those + run inline in `exec()` with no timeout at any layer. +- **`scripts/` prefix normalization:** `_prepare_code()` auto-prepends + `"scripts/"` if the `script_path` does not already start with it. This + allows the LLM to call with either `"setup.py"` or `"scripts/setup.py"`. +- **Resource packaging:** ALL skill files (references, assets, scripts) + are packaged as `input_files` with preserved relative paths (e.g., + `"references/data.csv"`, `"assets/template.txt"`). Empty resources + (content `""`) are still included — they are not silently dropped. + Both text (`str`) and binary (`bytes`) content are supported; the + executor writes them with the appropriate file mode (`'w'` vs `'wb'`). +- **SystemExit handling:** `SystemExit(0)` or `SystemExit(None)` → + treated as success (empty stdout/stderr, status `"success"`). + `SystemExit(non-zero)` → `EXECUTION_ERROR` with the exit code in the + error message. This prevents scripts from terminating the host process. +- **Error message truncation:** Exception messages longer than 200 + characters are truncated to `message[:200] + "..."` to conserve LLM + context tokens. +- **Shell JSON envelope:** Shell scripts are wrapped in a + `subprocess.run` call that serializes output as JSON: + `{"__shell_result__": true, "stdout": "...", "stderr": "...", + "returncode": N}`. On parse: + - Non-zero `returncode` with empty `stderr` → synthesized + `"Exit code {rc}"` as stderr + - Non-JSON stdout (e.g., if the wrapper itself fails) → raw stdout + is passed through without parsing +- **Status derivation:** Purely based on stream presence: + `stderr` only → `"error"`, both streams → `"warning"`, no + `stderr` → `"success"`. For shell scripts, non-zero return codes + influence status indirectly (via synthesized stderr). For Python + scripts, there is no return code extraction — status is determined + solely by stdout/stderr content. +- **Duplicate skill names:** `SkillToolset.__init__` validates that all + skill names are unique and raises `ValueError` on duplicates. +- **System instruction injection:** `SkillToolset.process_llm_request()` + appends `DEFAULT_SKILL_SYSTEM_INSTRUCTION` plus an XML-formatted skill + list to every outgoing LLM request, informing the model about available + skills and the `run_skill_script` tool. --- @@ -287,6 +390,42 @@ def execute_code(self, invocation_context, code_execution_input): - An alternative is `multiprocessing`, but that adds complexity around serialization and shared state. +**Interaction with existing `_execution_lock`:** + +The current `UnsafeLocalCodeExecutor` holds a process-global +`_execution_lock` (`threading.Lock()`) for the entire duration of +`execute_code()`, covering both the sandbox path (temp-dir + chdir) and +the plain path (redirect_stdout only). The timeout thread proposal must +account for this: + +- **Lock must be acquired outside the timeout thread.** The worker thread + must hold the lock while executing, and the calling thread must release + it after the join (whether the worker finishes or times out). If the + lock were acquired inside the worker thread, a timed-out worker would + hold the lock indefinitely, deadlocking all subsequent calls. +- **Recommended pattern:** Acquire the lock in `execute_code()` before + spawning the worker thread, pass the lock-holding context to the worker, + and release in a `finally` block after `thread.join(timeout)`: + ```python + with _execution_lock: + thread = threading.Thread(target=_run, daemon=True) + thread.start() + thread.join(timeout=timeout) + if thread.is_alive(): + # Lock is released when `with` exits, even though + # the daemon thread may still be running. + # This is acceptable: the lingering thread's exec() + # is no longer protected by the lock, but it is a + # daemon thread that will be killed on process exit. + return CodeExecutionResult( + stderr=f'Execution timed out after {timeout}s' + ) + ``` +- **Risk:** A timed-out daemon thread may still be mutating + `sys.stdout` or the working directory after the lock is released. + This is a best-effort trade-off for a development executor — the + alternative (never releasing the lock) would deadlock the process. + **Recommendation:** Thread-based timeout for `UnsafeLocalCodeExecutor` is sufficient. Document that it provides best-effort timeout only. @@ -533,6 +672,33 @@ However, the shell wrapper timeout (`subprocess.run(timeout=N)`) should be kept as defense-in-depth — it catches the subprocess even if the executor timeout fails. +**Critical gap — Python scripts have zero timeout protection:** + +Shell scripts benefit from two layers of timeout: the `subprocess.run( +timeout=N)` embedded in generated wrapper code, and (once implemented) +the executor-level timeout. Python scripts have **neither**: + +- `_prepare_code()` generates `runpy.run_path()` inside a plain `exec()` + call — there is no subprocess boundary to kill. +- `SkillToolset.script_timeout` only applies to the shell path (it is + passed to `subprocess.run(timeout=N)`). The docstring explicitly notes: + "Does not apply to Python scripts executed via exec()." +- Until executor-level timeout is implemented, a Python script that hangs + (infinite loop, blocking I/O, deadlock) will block the executor thread + indefinitely. With `UnsafeLocalCodeExecutor`, this also holds the + `_execution_lock`, blocking all other executions. + +**Recommended actions for Phase 1:** +1. Executor-level timeout (§4.2) is the primary fix — it covers both + Python and shell scripts uniformly. +2. As defense-in-depth, `RunSkillScriptTool` should also set + `CodeExecutionInput.timeout_seconds = self._toolset._script_timeout` + once the field is available, ensuring per-invocation timeout even if + the executor has no default. +3. Optionally, the Python wrapper code could be enhanced with a + watchdog thread pattern (similar to the shell `subprocess.run` + timeout), though this is less clean than executor-level enforcement. + --- ## 5. Proposal 2: Stateful `ContainerCodeExecutor` @@ -854,10 +1020,11 @@ This is a critical security concern when executing: |--------|--------|--------------------| | LLM generates malicious code | Full host compromise | None | | Skill script reads secrets | Data exfiltration | None (documented warning only) | -| Infinite loop / fork bomb | DoS / resource exhaustion | None (no timeout) | -| `sys.exit()` in script | Process termination | Partial (`SystemExit` catch in `RunSkillScriptTool`) | +| Infinite loop / fork bomb | DoS / resource exhaustion | Shell: `subprocess.run(timeout=N)` via `script_timeout`; Python: None (no timeout at any layer) | +| `sys.exit()` in script | Process termination | `RunSkillScriptTool` catches `SystemExit`: code 0 or `None` → success; non-zero → `EXECUTION_ERROR` with exit code in message | +| Long error messages | LLM context waste | Exception messages >200 chars truncated to `msg[:200] + "..."` | | Network exfiltration | Data leak | None | -| File system manipulation | Data loss / corruption | None | +| File system manipulation | Data loss / corruption | Partial (temp-dir sandbox when `input_files`/`working_dir` set) | ### 6.3 Design @@ -1185,8 +1352,13 @@ class BaseCodeExecutor(BaseModel): optimize_data_file: bool = False stateful: bool = False error_retry_attempts: int = 2 - code_block_delimiters: List[tuple[str, str]] = [...] - execution_result_delimiters: tuple[str, str] = (...) + code_block_delimiters: List[tuple[str, str]] = [ + ('```tool_code\n', '\n```'), + ('```python\n', '\n```'), + ] + execution_result_delimiters: tuple[str, str] = ( + '```tool_output\n', '\n```' + ) # NEW: Proposal 1 default_timeout_seconds: Optional[int] = None @@ -1206,6 +1378,7 @@ class CodeExecutionInput: code: str input_files: list[File] = field(default_factory=list) execution_id: Optional[str] = None + working_dir: Optional[str] = None timeout_seconds: Optional[int] = None # NEW: per-invocation """Per-invocation timeout. Overrides executor default when set.""" ``` @@ -1226,10 +1399,13 @@ class CodeExecutionInput: | Feature | Current workaround | After enhancements | |---------|-------------------|-------------------| -| Shell timeout | Embedded `subprocess.run(timeout=N)` | Keep as defense-in-depth | -| Python timeout | None | Executor-level timeout handles it | -| Isolation | Documentation warning only | `LocalSandboxCodeExecutor` or container | -| Stateful scripts | Not supported | Available via `ContainerCodeExecutor(stateful=True)` | +| Shell timeout | Embedded `subprocess.run(timeout=N)` via `SkillToolset.script_timeout` (default 300s) | Keep as defense-in-depth + executor-level timeout | +| Python timeout | **None** — `runpy.run_path()` runs inline in `exec()` with no timeout at any layer | Executor-level timeout handles it; tool should also set `CodeExecutionInput.timeout_seconds` | +| Isolation | Partial temp-dir sandbox (input_files/working_dir) + `_execution_lock` for stdout/cwd; no restriction on filesystem/network/env access | `LocalSandboxCodeExecutor` or container | +| Stateful scripts | Not supported (`execution_id` not wired) | Available via `ContainerCodeExecutor(stateful=True)` with `execution_id` | +| Output files | `CodeExecutionResult.output_files` silently dropped | Surfaced in tool response (§7.5.2) | +| System instructions | `SkillToolset.process_llm_request()` injects `DEFAULT_SKILL_SYSTEM_INSTRUCTION` + XML skill list | No change needed | +| Error truncation | Exception messages >200 chars truncated | Consider making threshold configurable | ### 7.4 Testing Strategy @@ -1287,6 +1463,35 @@ Guidance: - Tool-level validation/configuration errors should continue using explicit `error_code` values. +**Current status derivation and its asymmetry:** + +The current implementation determines status purely from stream presence: +```python +if stderr and not stdout: + status = "error" +elif stderr: + status = "warning" +else: + status = "success" +``` + +This creates an asymmetry between script types: + +- **Shell scripts:** Non-zero `returncode` from the JSON envelope causes + synthesized stderr (`"Exit code {rc}"`), so return codes **indirectly** + influence status. A shell script that fails silently (non-zero exit but + no stderr) still gets `"error"` status. +- **Python scripts:** There is **no return code extraction**. A Python + script that exits cleanly but writes warnings to stderr (common in + data science libraries) would be classified as `"error"` or `"warning"` + even if it succeeded. Conversely, a Python script that silently produces + incorrect output would get `"success"` status. + +The proposed `return_code` field resolves this by providing a uniform +source of truth. For Python scripts, this would require either: +(a) wrapping the `runpy.run_path()` call to capture the exit code, or +(b) treating any non-exception completion as `return_code = 0`. + #### 7.5.2 Propagate `output_files` and Artifact Metadata `CodeExecutionResult.output_files` should be surfaced in the tool response. @@ -1327,18 +1532,40 @@ Rules: #### 7.5.4 Script Args Normalization Contract -Define and document deterministic mapping from JSON args to CLI argv. +**Current behavior:** The implementation uses a simple `str(v)` conversion +for all argument values, with no type-aware normalization: -Suggested rules: +```python +# Both Python and shell paths use the same logic: +for k, v in script_args.items(): + argv_list.extend([f"--{k}", str(v)]) +``` + +This means: +- `{"verbose": true}` → `["--verbose", "True"]` (string, not a flag) +- `{"flag": false}` → `["--flag", "False"]` (not omitted) +- `{"items": [1, 2, 3]}` → `["--items", "[1, 2, 3]"]` (repr of list) +- `{"count": 42}` → `["--count", "42"]` (correct) +- Nested objects → `["--config", "{'a': 1}"]` (repr, not useful) + +Additionally, `args` type is validated to be a `dict` — non-dict values +(strings, lists, integers, booleans) return `INVALID_ARGS_TYPE` error. + +**Proposed rules** (define deterministic mapping from JSON args to argv): - `str|int|float` -> `--key value` -- `true` -> `--key` -- `false|None` -> omit +- `true` -> `--key` (flag only, no value) +- `false|None` -> omit entirely - `list[...]` -> repeated `--key value` entries - Optional reserved key for positional args (for example, `_positional`) - Reject nested objects with explicit validation error This reduces LLM-side ambiguity and improves replay/debug stability. +**Migration note:** Changing boolean handling from `"--key True"` to +`"--key"` (flag) is a behavioral change. Existing skill scripts that +parse `--verbose True` as a string value would break. The migration +should be opt-in or gated behind a version flag. + --- ## 8. Implementation Roadmap From 36f0148c7c8ac82cb7cabc114d48b10d26c99be3 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 00:55:39 -0800 Subject: [PATCH 35/53] docs: add P0 RFC for RunSkillScriptTool production-readiness MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 5-page RFC targeting ADK TL/UTL audience, covering the two P0 gaps that block production use of RunSkillScriptTool: - P0-A: Uniform timeout support — add timeout_seconds to CodeExecutionInput and default_timeout_seconds to BaseCodeExecutor, with per-executor implementations (thread+join for local, Docker exec kill for container). - P0-B: LocalSandboxCodeExecutor — new stdlib-only executor using subprocess with resource limits, replacing UnsafeLocalCodeExecutor as the recommended local default. Includes alternatives considered, phased rollout plan, success metrics, and execution flow appendix. Co-Authored-By: Claude Opus 4.6 --- docs/design/rfc_runskillscript_p0.md | 456 +++++++++++++++++++++++++++ 1 file changed, 456 insertions(+) create mode 100644 docs/design/rfc_runskillscript_p0.md diff --git a/docs/design/rfc_runskillscript_p0.md b/docs/design/rfc_runskillscript_p0.md new file mode 100644 index 0000000000..cb7aa91cf3 --- /dev/null +++ b/docs/design/rfc_runskillscript_p0.md @@ -0,0 +1,456 @@ +# RFC: Production-Readiness for RunSkillScriptTool + +**Authors:** haiyuancao +**Date:** 2026-02-24 +**Status:** Proposed +**Audience:** ADK TL / UTL +**Effort:** 8-11 engineering days +**Tracking:** Follow-up to PR #4575 (RunSkillScriptTool) + +--- + +## 1. Executive Summary + +`RunSkillScriptTool` lets agents execute Python and shell scripts +bundled with skills. It shipped in PR #4575 and is functional, but two +gaps block production use: + +1. **Python scripts can hang forever.** Shell scripts have a + `subprocess.run(timeout=300)` guard; Python scripts have none. A + single stuck `runpy.run_path()` call holds a process-global lock, + blocking every subsequent execution across all agents in the process. + +2. **All local execution is unprotected.** `UnsafeLocalCodeExecutor` — + the only zero-dependency executor and the default for development — + runs `exec()` in the host process. A malicious or buggy script can + read secrets, write to the filesystem, exfiltrate data over the + network, or crash the process. + +This RFC proposes two changes, scoped to what is required before skills +can be used in any environment beyond local prototyping: + +- **P0-A: Uniform timeout** — Add `timeout_seconds` to the executor + contract so every `execute_code()` call has a bounded lifetime. +- **P0-B: `LocalSandboxCodeExecutor`** — A new stdlib-only executor + that runs code in a subprocess with resource limits, replacing + `UnsafeLocalCodeExecutor` as the recommended local default. + +Both changes are backward-compatible, additive, and independently +shippable. + +--- + +## 2. Problem Statement + +### 2.1 The Timeout Gap + +The executor landscape today: + +| Executor | Timeout | How | +|----------|---------|-----| +| `UnsafeLocalCodeExecutor` | **None** | `exec()` blocks forever | +| `ContainerCodeExecutor` | **None** | `exec_run()` blocks forever | +| `GkeCodeExecutor` | 300 s | K8s watch API | +| `VertexAiCodeExecutor` | Opaque | Vertex AI internal | + +`RunSkillScriptTool` works around this for shell scripts by embedding +`subprocess.run(timeout=N)` in generated wrapper code. The default is +300 seconds, configurable via `SkillToolset(script_timeout=N)`. + +**Python scripts have zero timeout at any layer.** The tool generates: + +```python +import sys, runpy +sys.argv = ['scripts/run.py', '--verbose', 'True'] +runpy.run_path('scripts/run.py', run_name='__main__') +``` + +This runs inline inside `exec()`. There is no subprocess boundary, no +watchdog thread, and no way for the caller to interrupt it. Worse, +`UnsafeLocalCodeExecutor` holds a process-global `threading.Lock()` +for the entire execution (required because `redirect_stdout` and +`os.chdir` mutate process-global state). A hung Python script +deadlocks the lock, blocking **all** code execution across all agents +sharing that executor instance. + +**Impact:** A single infinite-loop in a skill script takes down the +entire ADK process. This is a denial-of-service risk for any +deployment — not just production, but also development servers and +demos. + +### 2.2 The Security Gap + +`UnsafeLocalCodeExecutor` runs code with the full privileges of the +host Python process: + +| Threat | Impact | Current Mitigation | +|--------|--------|--------------------| +| Read env vars / secrets | Data exfiltration | None | +| Write to host filesystem | Data loss / corruption | Partial (temp-dir when `input_files` set) | +| Outbound network calls | Data leak | None | +| `sys.exit()` | Process crash | `SystemExit` caught in tool | +| Infinite loop / fork bomb | DoS | Shell: `subprocess.run(timeout)`; Python: **none** | + +This executor is the only one that requires zero external dependencies, +making it the de facto default for: +- `adk web` / `adk run` during development +- CI test suites +- Quick-start samples and tutorials + +Any code the LLM generates or any third-party skill script runs with +full host access. This is acceptable for trusted, single-developer +prototyping. It is not acceptable for: +- Multi-user development servers +- CI/CD pipelines running untrusted skill scripts +- Any path toward production deployment + +--- + +## 3. Proposed Solution + +### 3.1 P0-A: Uniform Timeout Support + +**Goal:** Every `execute_code()` call has a bounded lifetime, regardless +of executor backend or script type. + +#### 3.1.1 Contract Changes (Backward-Compatible) + +Add one field to each layer: + +```python +# Per-invocation timeout (caller sets this) +@dataclasses.dataclass +class CodeExecutionInput: + code: str + input_files: list[File] = field(default_factory=list) + execution_id: Optional[str] = None + working_dir: Optional[str] = None + timeout_seconds: Optional[int] = None # NEW + +# Executor-level default (fallback when caller doesn't set one) +class BaseCodeExecutor(BaseModel): + default_timeout_seconds: Optional[int] = None # NEW + ... +``` + +Resolution logic in every executor: +```python +timeout = ( + code_execution_input.timeout_seconds + if code_execution_input.timeout_seconds is not None + else self.default_timeout_seconds +) +``` + +**Why two fields, not one?** +- A single executor instance may be shared across agents/tools with + different timeout needs (quick validation vs. long data analysis). +- Per-invocation is the source of truth; executor default is the + safety net for callers that don't set one. +- Existing code that sets neither field gets `None` → no timeout → + identical to current behavior. Zero breaking changes. + +#### 3.1.2 `UnsafeLocalCodeExecutor` — Thread + Join + +`exec()` cannot be interrupted from the same thread. Run it in a +daemon thread with `join(timeout)`: + +``` + ┌─ execute_code() ──────────────────┐ + │ │ + acquire │ with _execution_lock: │ + _execution_lock ──────│ spawn daemon thread ──► exec() │ + │ thread.join(timeout) │ + │ if thread.is_alive(): │ + │ return stderr="timed out" │ + │ │ + release lock ─────────│ (lock released even if timed out)│ + └────────────────────────────────────┘ +``` + +**Key design decision — lock acquired outside the thread:** +The `_execution_lock` must be held by the calling thread, not the +worker. If the worker held the lock and timed out, it would hold it +forever, deadlocking subsequent calls. Acquiring outside means the +lock is released when the `with` block exits, even if the daemon +thread is still running. The lingering thread is a daemon — it will +be killed on process exit. This is a best-effort trade-off appropriate +for a development executor. + +**Trade-off acknowledged:** A timed-out daemon thread may still be +mutating `sys.stdout` after the lock is released. This is acceptable +for `UnsafeLocalCodeExecutor` (development only); production executors +use subprocess or container isolation where kill is clean. + +#### 3.1.3 `ContainerCodeExecutor` — Docker Exec Kill + +Run `exec_start` in a thread. On timeout, kill the process from the +host side: + +1. `exec_inspect(exec_id)` → get host-namespace PID +2. `os.kill(host_pid, SIGKILL)` +3. If `PermissionError` (common: container runs as root, ADK does not) + → `container.restart(timeout=1)` + readiness check +4. If both fail → set `self._healthy = False`, return error; caller + must call `reinitialize()` to recover + +This is the only design that guarantees the caller is never blocked +beyond `timeout` seconds, regardless of whether the kill succeeds. + +#### 3.1.4 Wire Timeout in `RunSkillScriptTool` + +Once `CodeExecutionInput.timeout_seconds` exists, the tool sets it: + +```python +CodeExecutionInput( + code=code, + input_files=input_files, + working_dir=".", + timeout_seconds=self._toolset._script_timeout, # NEW +) +``` + +This gives Python scripts the same 300-second default that shell +scripts already have, and makes it configurable via +`SkillToolset(script_timeout=N)`. The shell `subprocess.run(timeout)` +is kept as defense-in-depth. + +#### 3.1.5 Migration + +| Step | Change | Risk | +|------|--------|------| +| 1 | Add `timeout_seconds` to `CodeExecutionInput` | None | +| 2 | Add `default_timeout_seconds` to `BaseCodeExecutor` | None | +| 3 | Implement in `UnsafeLocalCodeExecutor` (thread + join) | Low | +| 4 | Implement in `ContainerCodeExecutor` (exec kill) | Low | +| 5 | Migrate `GkeCodeExecutor.timeout_seconds` to new field | None | +| 6 | Wire in `RunSkillScriptTool` | None | + +**Estimated effort:** 3-4 days (including tests). + +--- + +### 3.2 P0-B: `LocalSandboxCodeExecutor` + +**Goal:** A zero-dependency executor that provides meaningful isolation +without Docker or cloud services. + +#### 3.2.1 Why Not Just Harden `UnsafeLocalCodeExecutor`? + +We considered two alternatives before proposing a new executor: + +| Option | Approach | Why rejected | +|--------|----------|-------------| +| **Restricted builtins** | Block `open`, `__import__`, `exec`, `eval` in `exec()` globals | Trivially bypassed via `object.__subclasses__()`, `importlib` through `sys.modules`, `__builtins__.__dict__`. Not a security boundary — at best a speed bump. | +| **SecurityWarning only** | Emit `warnings.warn(SecurityWarning)` on first use | Does not reduce attack surface. Users ignore warnings. | + +Both are worth adding as **supplementary** friction (low cost, some +value), but neither solves the problem. True isolation requires a +process boundary. + +#### 3.2.2 Design: Subprocess with Resource Limits + +```python +class LocalSandboxCodeExecutor(BaseCodeExecutor): + """Executes Python in a sandboxed subprocess. + + Isolation: + - Separate process (no shared memory with host) + - Resource limits (CPU time, memory) via resource.setrlimit + - Restricted environment variables + - Temporary working directory + - subprocess.run(timeout=N) for wall-clock timeout + """ + + default_timeout_seconds: int = 30 + max_memory_mb: int = 256 + max_cpu_seconds: int = 30 + allowed_env_vars: list[str] = [] +``` + +Execution flow: + +``` + 1. Write code to a NamedTemporaryFile (.py) + 2. Write input_files to a TemporaryDirectory + 3. Build minimal env: only allowed_env_vars + PATH + 4. subprocess.run( + ['python3', '-c', + exec(open(file).read())], + timeout=timeout, + env=env, + cwd=temp_dir, + process_group=0, # Python 3.11+; preexec_fn fallback for 3.10 + capture_output=True, + ) + 5. Return CodeExecutionResult(stdout, stderr) +``` + +The inline `limit_code` wrapper sets `resource.setrlimit` for CPU and +memory inside the child process (guarded with `try/except ImportError` +for platforms where `resource` is unavailable). + +**What this protects against vs. what it doesn't:** + +| Threat | Protected? | How | +|--------|-----------|-----| +| Infinite loop | Yes | `subprocess.run(timeout)` + `RLIMIT_CPU` | +| Memory bomb | Yes | `RLIMIT_AS` | +| Env var / secret reading | Yes | Restricted `env` dict | +| `sys.exit()` crash | Yes | Separate process | +| Filesystem read/write | **Partial** | `cwd` is temp dir, but host fs still accessible | +| Network exfiltration | **No** | Requires OS-level firewall (out of scope) | + +This is strictly stronger than `UnsafeLocalCodeExecutor` across every +dimension, with zero additional dependencies. The remaining gaps +(full filesystem isolation, network restriction) require containers. + +#### 3.2.3 Platform Considerations + +- **Python 3.11+:** Use `process_group=0` (fork-safe, replaces + `preexec_fn`). Enables clean `os.killpg()` on timeout. +- **Python 3.10:** Fall back to `preexec_fn=set_limits` with a + documented caveat about thread safety in multi-threaded programs. + ADK minimum is `>=3.10` per `pyproject.toml`. +- **Windows:** Not supported. Raise `NotImplementedError` directing + users to `ContainerCodeExecutor`. (`resource.setrlimit` and + `process_group` are Unix-only.) + +#### 3.2.4 Migration + +| Step | Change | Risk | +|------|--------|------| +| 1 | Implement `LocalSandboxCodeExecutor` | Low | +| 2 | Add `SecurityWarning` to `UnsafeLocalCodeExecutor` | None | +| 3 | Add `restrict_builtins` option (opt-in, supplementary) | Low | +| 4 | Update samples and `adk web` defaults to recommend sandbox | Low | +| 5 | Update documentation | None | + +**Estimated effort:** 5-7 days (including tests and docs). + +--- + +## 4. Alternatives Considered + +### 4.1 "Do Nothing — Document the Risks" + +Add warnings to docs and let users choose their executor. + +**Rejected.** The default experience (`adk web`, samples, tutorials) +uses `UnsafeLocalCodeExecutor`. New users will not read security docs +before running `adk web`. The default must be safe enough for its +intended use (local development with untrusted LLM-generated code). + +### 4.2 "Require Docker for All Local Execution" + +Make `ContainerCodeExecutor` the default. Remove +`UnsafeLocalCodeExecutor`. + +**Rejected.** Docker is a significant dependency. Many developers +(especially on macOS) don't have it installed. CI environments +may not have Docker available. The zero-dependency story is a key +ADK differentiator for onboarding. We need a middle ground. + +### 4.3 "Timeout Only, No New Executor" + +Ship P0-A (timeout) without P0-B (sandbox). Defer security to a +later phase. + +**Viable but insufficient.** Timeout prevents DoS but does not +prevent data exfiltration, secret reading, or filesystem manipulation. +These are the higher-impact threats for skill scripts, which may come +from third-party authors. We recommend shipping both P0-A and P0-B, +but they are independently valuable and can be phased if needed. + +--- + +## 5. Rollout Plan + +### 5.1 Phase 1: Timeout Foundation (Week 1) + +| Day | Deliverable | +|-----|-------------| +| 1 | Add `timeout_seconds` to `CodeExecutionInput` and `default_timeout_seconds` to `BaseCodeExecutor`. Unit tests for field defaults and resolution logic. | +| 2 | Implement thread-based timeout in `UnsafeLocalCodeExecutor`. Tests with `time.sleep()` scripts. | +| 3 | Implement Docker exec kill in `ContainerCodeExecutor` with `_healthy` guard and `reinitialize()`. Mock-based tests for kill path, permission error path, and total failure path. | +| 4 | Migrate `GkeCodeExecutor`, wire `RunSkillScriptTool`, end-to-end integration tests. | + +**Exit criteria:** `RunSkillScriptTool` Python scripts respect +`script_timeout`. Existing tests pass. No behavioral change for +callers that don't set timeout. + +### 5.2 Phase 2: Security Hardening (Week 2) + +| Day | Deliverable | +|-----|-------------| +| 1-2 | Implement `LocalSandboxCodeExecutor` with resource limits. | +| 3 | Add `SecurityWarning` to `UnsafeLocalCodeExecutor`. Add `restrict_builtins` opt-in. | +| 4 | Tests: resource limit enforcement, env var isolation, timeout kill via `os.killpg`, platform fallback for 3.10. | +| 5 | Update samples, docs, and recommendation matrix. | + +**Exit criteria:** `LocalSandboxCodeExecutor()` works as a drop-in +replacement for `UnsafeLocalCodeExecutor()` with no configuration. +Documentation recommends it for all local use. + +### 5.3 Recommendation Matrix (Post-Rollout) + +| Use Case | Executor | Why | +|----------|----------|-----| +| Local development | `LocalSandboxCodeExecutor` | Zero deps, subprocess isolation | +| Quick prototyping (trusted code) | `UnsafeLocalCodeExecutor` | Fastest, no isolation | +| CI/CD | `ContainerCodeExecutor` | Docker available in CI | +| Production (single tenant) | `ContainerCodeExecutor` | Full container isolation | +| Production (multi-tenant) | `GkeCodeExecutor` | gVisor, per-execution isolation | +| Google Cloud | `AgentEngineSandboxCodeExecutor` | Managed, scalable | + +### 5.4 Success Metrics + +- **Timeout coverage:** 100% of executors support `timeout_seconds` + (currently 1 of 5). +- **Default safety:** `LocalSandboxCodeExecutor` passes all existing + `RunSkillScriptTool` integration tests as a drop-in replacement. +- **No regressions:** All existing unit and integration tests pass + with zero behavioral changes for callers that don't opt in. +- **Adoption signal:** Samples and `adk web` default documentation + reference `LocalSandboxCodeExecutor`. + +--- + +## Appendix A: Current RunSkillScriptTool Execution Flow + +``` +LLM calls run_skill_script(skill_name, script_path, args) + │ + ▼ +RunSkillScriptTool.run_async() + │ + ├─ Validate params (skill_name, script_path, args type) + ├─ Resolve skill → resolve script from resources + ├─ Resolve executor: toolset._code_executor → agent.code_executor + ├─ Package ALL skill files as input_files (refs, assets, scripts) + ├─ _prepare_code(): + │ .py → runpy.run_path() wrapper (NO timeout) + │ .sh → subprocess.run(timeout=N) wrapper with JSON envelope + │ + ├─ await asyncio.to_thread(executor.execute_code, ...) + │ CodeExecutionInput(code=..., input_files=..., working_dir=".") + │ *** timeout_seconds NOT SET — this is the gap *** + │ + ├─ Parse result: + │ Shell: unpack JSON envelope {stdout, stderr, returncode} + │ Python: use stdout/stderr directly + │ + └─ Return {skill_name, script_path, stdout, stderr, status} +``` + +## Appendix B: Key Source Files + +| File | Role | +|------|------| +| `src/google/adk/tools/skill_toolset.py` | `RunSkillScriptTool`, `SkillToolset` | +| `src/google/adk/code_executors/base_code_executor.py` | `BaseCodeExecutor` abstract class | +| `src/google/adk/code_executors/code_execution_utils.py` | `CodeExecutionInput`, `CodeExecutionResult`, `File` | +| `src/google/adk/code_executors/unsafe_local_code_executor.py` | Current local executor | +| `src/google/adk/code_executors/container_code_executor.py` | Docker-based executor | +| `src/google/adk/code_executors/gke_code_executor.py` | GKE-based executor (has timeout) | +| `tests/unittests/tools/test_skill_toolset.py` | 1173-line test suite for skill tools | +| `docs/design/code_executor_enhancements.md` | Detailed design doc (companion to this RFC) | From 702a77f8672af36d9e46c99b0675104cffa3ba46 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 01:01:55 -0800 Subject: [PATCH 36/53] docs(rfc): address 6 review findings in RunSkillScriptTool P0 RFC MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. High: UnsafeLocalCodeExecutor timeout — replace naive lock-release with unhealthy-guard pattern. On timeout, mark executor unhealthy and fail fast until reinitialize(), preventing cross-execution stdout/cwd contamination from lingering daemon threads. 2. High: LocalSandboxCodeExecutor — replace subprocess.run(timeout) with Popen + os.killpg(SIGKILL) to kill entire process group on timeout. Adjust threat table: fork bomb protection is partial (no RLIMIT_NPROC), child-proc infinite loops are covered. 3. Medium: Add AgentEngineSandboxCodeExecutor to timeout landscape table (opaque Vertex AI Sandbox internal timeout). 4. Medium: Rephrase "de facto default" claim — UnsafeLocalCodeExecutor is common in samples/dev setups, not a global LlmAgent default. 5. Medium: Soften "drop-in replacement" exit criterion — define supported script profile (stdout/stderr, sandbox-local files, explicit env var allowlist) and compatibility envelope. 6. Low: Make test line count approximate (~1170-line). Co-Authored-By: Claude Opus 4.6 --- docs/design/rfc_runskillscript_p0.md | 148 +++++++++++++++++++-------- 1 file changed, 104 insertions(+), 44 deletions(-) diff --git a/docs/design/rfc_runskillscript_p0.md b/docs/design/rfc_runskillscript_p0.md index cb7aa91cf3..6824bf5c33 100644 --- a/docs/design/rfc_runskillscript_p0.md +++ b/docs/design/rfc_runskillscript_p0.md @@ -52,6 +52,7 @@ The executor landscape today: | `ContainerCodeExecutor` | **None** | `exec_run()` blocks forever | | `GkeCodeExecutor` | 300 s | K8s watch API | | `VertexAiCodeExecutor` | Opaque | Vertex AI internal | +| `AgentEngineSandboxCodeExecutor` | Opaque | Vertex AI Sandbox internal | `RunSkillScriptTool` works around this for shell scripts by embedding `subprocess.run(timeout=N)` in generated wrapper code. The default is @@ -91,11 +92,14 @@ host Python process: | `sys.exit()` | Process crash | `SystemExit` caught in tool | | Infinite loop / fork bomb | DoS | Shell: `subprocess.run(timeout)`; Python: **none** | -This executor is the only one that requires zero external dependencies, -making it the de facto default for: -- `adk web` / `adk run` during development -- CI test suites -- Quick-start samples and tutorials +This executor is the only one that requires zero external dependencies. +While `LlmAgent.code_executor` does not globally default to it (it +defaults to `None`, and CFC paths may force `BuiltInCodeExecutor`), +`UnsafeLocalCodeExecutor` is the common choice in practice for setups +that need local code execution: +- Samples and tutorials that demonstrate `code_executor=` configuration +- `adk web` / `adk run` development workflows +- CI test suites without Docker Any code the LLM generates or any third-party skill script runs with full host access. This is acceptable for trusted, single-developer @@ -150,37 +154,63 @@ timeout = ( - Existing code that sets neither field gets `None` → no timeout → identical to current behavior. Zero breaking changes. -#### 3.1.2 `UnsafeLocalCodeExecutor` — Thread + Join +#### 3.1.2 `UnsafeLocalCodeExecutor` — Thread + Join + Unhealthy Guard `exec()` cannot be interrupted from the same thread. Run it in a daemon thread with `join(timeout)`: ``` - ┌─ execute_code() ──────────────────┐ - │ │ - acquire │ with _execution_lock: │ - _execution_lock ──────│ spawn daemon thread ──► exec() │ - │ thread.join(timeout) │ - │ if thread.is_alive(): │ - │ return stderr="timed out" │ - │ │ - release lock ─────────│ (lock released even if timed out)│ - └────────────────────────────────────┘ + ┌─ execute_code() ──────────────────────┐ + │ │ + check _healthy ───────│ if not self._healthy: raise error │ + │ │ + acquire │ with _execution_lock: │ + _execution_lock ──────│ spawn daemon thread ──► exec() │ + │ thread.join(timeout) │ + │ if thread.is_alive(): │ + │ self._healthy = False ◄── mark! │ + │ return stderr="timed out" │ + │ │ + release lock ─────────│ (lock released) │ + └────────────────────────────────────────┘ ``` -**Key design decision — lock acquired outside the thread:** -The `_execution_lock` must be held by the calling thread, not the -worker. If the worker held the lock and timed out, it would hold it -forever, deadlocking subsequent calls. Acquiring outside means the -lock is released when the `with` block exits, even if the daemon -thread is still running. The lingering thread is a daemon — it will -be killed on process exit. This is a best-effort trade-off appropriate -for a development executor. - -**Trade-off acknowledged:** A timed-out daemon thread may still be -mutating `sys.stdout` after the lock is released. This is acceptable -for `UnsafeLocalCodeExecutor` (development only); production executors -use subprocess or container isolation where kill is clean. +**Critical invariant — unhealthy after timeout:** + +Simply releasing `_execution_lock` while a timed-out daemon thread is +still running would let subsequent executions proceed while the +lingering thread continues to mutate process-global `sys.stdout` and +the working directory. This causes cross-execution stdout +contamination and cwd corruption — a data-integrity bug, not just a +cosmetic issue. + +The solution is to **mark the executor unhealthy on timeout**: + +1. On timeout, set `self._healthy = False` before returning the + timeout error. The lock is released normally when the `with` block + exits. +2. Subsequent `execute_code()` calls check `self._healthy` at the top + and **fail fast** with a clear error: `"Executor is unhealthy after + a timed-out execution. Call reinitialize() to recover."` +3. `reinitialize()` waits for the lingering daemon thread to finish + (with a generous join timeout), resets `_healthy = True`, and + allows execution to resume. + +This ensures that no new execution can run while a zombie thread is +still alive and mutating shared state. The pattern mirrors the +`ContainerCodeExecutor` unhealthy-state design (§3.1.3). + +**Trade-off:** After a timeout, the executor is unavailable until +`reinitialize()` is called. This is acceptable for a development +executor — the alternative (silent stdout/cwd corruption) is worse. +Production deployments should use `LocalSandboxCodeExecutor` (§3.2) +or `ContainerCodeExecutor`, where timeout kill is clean. + +**Why not skip thread-timeout on `UnsafeLocalCodeExecutor` entirely?** +Without any timeout, a hung `exec()` holds `_execution_lock` forever, +which is strictly worse — the executor is permanently blocked with no +error and no recovery path. The unhealthy-guard approach at least +unblocks the lock, reports the error, and offers a recovery mechanism. #### 3.1.3 `ContainerCodeExecutor` — Docker Exec Kill @@ -274,17 +304,35 @@ Execution flow: 1. Write code to a NamedTemporaryFile (.py) 2. Write input_files to a TemporaryDirectory 3. Build minimal env: only allowed_env_vars + PATH - 4. subprocess.run( + 4. Popen( ['python3', '-c', + exec(open(file).read())], - timeout=timeout, env=env, cwd=temp_dir, - process_group=0, # Python 3.11+; preexec_fn fallback for 3.10 - capture_output=True, + process_group=0, # Python 3.11+; preexec_fn fallback + stdout=PIPE, stderr=PIPE, ) - 5. Return CodeExecutionResult(stdout, stderr) + 5. communicate(timeout=timeout) + 6. On TimeoutExpired: + os.killpg(proc.pid, signal.SIGKILL) # kill entire group + proc.communicate(timeout=5) # reap zombies + 7. Return CodeExecutionResult(stdout, stderr) ``` +**Why `Popen` + `os.killpg` instead of `subprocess.run(timeout)`:** + +`subprocess.run(timeout=N)` only kills the direct child process. If +the script spawns subprocesses (e.g., `os.system()`, `Popen`, +`multiprocessing`), those descendants survive the timeout and become +orphans. With `process_group=0`, the child is placed in its own +process group. On timeout, `os.killpg(proc.pid, SIGKILL)` kills the +entire group — the child and all its descendants. The follow-up +`proc.communicate(timeout=5)` reaps any zombies. + +On Python 3.10 (where `process_group` is unavailable), the fallback +uses `preexec_fn=os.setpgrp` to achieve the same process-group +isolation, with a documented caveat about fork-safety in +multi-threaded programs (see §3.2.3). + The inline `limit_code` wrapper sets `resource.setrlimit` for CPU and memory inside the child process (guarded with `try/except ImportError` for platforms where `resource` is unavailable). @@ -293,16 +341,19 @@ for platforms where `resource` is unavailable). | Threat | Protected? | How | |--------|-----------|-----| -| Infinite loop | Yes | `subprocess.run(timeout)` + `RLIMIT_CPU` | +| Infinite loop (direct) | Yes | Wall-clock timeout + `RLIMIT_CPU` | +| Infinite loop (child procs) | Yes | `os.killpg` kills entire process group | +| Fork bomb | **Partial** | `RLIMIT_CPU` + timeout bound total wall-clock; does not cap `RLIMIT_NPROC` (could be added) | | Memory bomb | Yes | `RLIMIT_AS` | | Env var / secret reading | Yes | Restricted `env` dict | | `sys.exit()` crash | Yes | Separate process | -| Filesystem read/write | **Partial** | `cwd` is temp dir, but host fs still accessible | +| Filesystem read/write | **Partial** | `cwd` is temp dir, but host fs still accessible via absolute paths | | Network exfiltration | **No** | Requires OS-level firewall (out of scope) | This is strictly stronger than `UnsafeLocalCodeExecutor` across every dimension, with zero additional dependencies. The remaining gaps -(full filesystem isolation, network restriction) require containers. +(full filesystem isolation, network restriction, `RLIMIT_NPROC`) +require containers or OS-level policy. #### 3.2.3 Platform Considerations @@ -335,9 +386,10 @@ dimension, with zero additional dependencies. The remaining gaps Add warnings to docs and let users choose their executor. -**Rejected.** The default experience (`adk web`, samples, tutorials) -uses `UnsafeLocalCodeExecutor`. New users will not read security docs -before running `adk web`. The default must be safe enough for its +**Rejected.** Samples, tutorials, and common development workflows +use `UnsafeLocalCodeExecutor` for local code execution. New users +follow sample code without reading security docs. The executor most +commonly reached by new developers must be safe enough for its intended use (local development with untrusted LLM-generated code). ### 4.2 "Require Docker for All Local Execution" @@ -388,8 +440,14 @@ callers that don't set timeout. | 5 | Update samples, docs, and recommendation matrix. | **Exit criteria:** `LocalSandboxCodeExecutor()` works as a drop-in -replacement for `UnsafeLocalCodeExecutor()` with no configuration. -Documentation recommends it for all local use. +replacement for `UnsafeLocalCodeExecutor()` for the supported script +profile: scripts that use stdout/stderr for output, do not depend on +host environment variables beyond an explicit allowlist, and access +files only within the sandbox working directory. Scripts that rely on +broad host-filesystem access or specific env vars will need to +configure `allowed_env_vars` or use `UnsafeLocalCodeExecutor`. +Documentation recommends `LocalSandboxCodeExecutor` for all new local +development and clearly states the compatibility envelope. ### 5.3 Recommendation Matrix (Post-Rollout) @@ -407,7 +465,9 @@ Documentation recommends it for all local use. - **Timeout coverage:** 100% of executors support `timeout_seconds` (currently 1 of 5). - **Default safety:** `LocalSandboxCodeExecutor` passes all existing - `RunSkillScriptTool` integration tests as a drop-in replacement. + `RunSkillScriptTool` integration tests that fit the supported script + profile (stdout/stderr output, sandbox-local file access, no + dependency on host env vars beyond allowlist). - **No regressions:** All existing unit and integration tests pass with zero behavioral changes for callers that don't opt in. - **Adoption signal:** Samples and `adk web` default documentation @@ -452,5 +512,5 @@ RunSkillScriptTool.run_async() | `src/google/adk/code_executors/unsafe_local_code_executor.py` | Current local executor | | `src/google/adk/code_executors/container_code_executor.py` | Docker-based executor | | `src/google/adk/code_executors/gke_code_executor.py` | GKE-based executor (has timeout) | -| `tests/unittests/tools/test_skill_toolset.py` | 1173-line test suite for skill tools | +| `tests/unittests/tools/test_skill_toolset.py` | ~1170-line test suite for skill tools | | `docs/design/code_executor_enhancements.md` | Detailed design doc (companion to this RFC) | From 992b864c463f38e943198cf6bab178132c08ca3e Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 01:07:48 -0800 Subject: [PATCH 37/53] docs(code-executor): address 7 review findings in design doc Fix 3 High, 3 Medium, 1 Low severity issues: - LocalSandboxCodeExecutor: materialize input_files, honor working_dir, use Popen+killpg for process-group cleanup on timeout - UnsafeLocalCodeExecutor: replace lock-release pattern with unhealthy-guard - Stateful execution: split Option A (final target) from Option C (interim) with acceptance criteria and deprecation plan - Container security: make network/read-only defaults opt-in first - Backward compat: clarify "Breaking + Minor" with explicit versioning policy - Open Question 1: align phase reference to Phase 4 Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 237 ++++++++++++++++------ 1 file changed, 175 insertions(+), 62 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index 18f1e30c24..7e5a03ea50 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -403,31 +403,44 @@ account for this: it after the join (whether the worker finishes or times out). If the lock were acquired inside the worker thread, a timed-out worker would hold the lock indefinitely, deadlocking all subsequent calls. -- **Recommended pattern:** Acquire the lock in `execute_code()` before - spawning the worker thread, pass the lock-holding context to the worker, - and release in a `finally` block after `thread.join(timeout)`: +- **Recommended pattern — unhealthy guard:** When a timeout fires, the + executor marks itself as unhealthy (`_healthy = False`) and **does not + release the lock**. All subsequent `execute_code()` calls fail fast + with an explicit error until the caller invokes `reinitialize()`, + which waits for the lingering daemon thread (best-effort), resets + `_healthy = True`, and releases the lock. This prevents the lock + from being released while the daemon thread still mutates + `sys.stdout` or the working directory. ```python - with _execution_lock: - thread = threading.Thread(target=_run, daemon=True) - thread.start() - thread.join(timeout=timeout) - if thread.is_alive(): - # Lock is released when `with` exits, even though - # the daemon thread may still be running. - # This is acceptable: the lingering thread's exec() - # is no longer protected by the lock, but it is a - # daemon thread that will be killed on process exit. + def execute_code(self, ...): + if not self._healthy: return CodeExecutionResult( - stderr=f'Execution timed out after {timeout}s' + stderr='Executor unhealthy after timeout. ' + 'Call reinitialize() to recover.' ) + with _execution_lock: + thread = threading.Thread(target=_run, daemon=True) + thread.start() + thread.join(timeout=timeout) + if thread.is_alive(): + self._healthy = False + # Lock stays held — no new exec until reinit. + return CodeExecutionResult( + stderr=f'Execution timed out after ' + f'{timeout}s. Executor marked ' + f'unhealthy.' + ) ``` -- **Risk:** A timed-out daemon thread may still be mutating - `sys.stdout` or the working directory after the lock is released. - This is a best-effort trade-off for a development executor — the - alternative (never releasing the lock) would deadlock the process. +- **Recovery path:** `reinitialize()` joins the lingering daemon + thread (with a generous grace timeout), restores `sys.stdout` + and cwd if needed, sets `_healthy = True`, and releases + `_execution_lock`. This mirrors the `ContainerCodeExecutor` + recovery model (§4.2.3). -**Recommendation:** Thread-based timeout for `UnsafeLocalCodeExecutor` is -sufficient. Document that it provides best-effort timeout only. +**Recommendation:** Thread-based timeout with unhealthy guard for +`UnsafeLocalCodeExecutor` is sufficient for a development executor. +Document that timeout triggers an unhealthy state requiring +explicit recovery via `reinitialize()`. #### 4.2.3 `ContainerCodeExecutor` — Docker Exec Kill on Timeout @@ -848,10 +861,10 @@ of the cumulative replay approach (Option C). Users must keep side-effecting code in the final block or use Option A (persistent process) when side effects are unavoidable. -#### 5.2.2 Recommended Approach: Option A (Persistent Process) +#### 5.2.2 Final Target: Option A (Persistent Process) -**Option A is the recommended approach.** It is the most robust for -true statefulness and is the standard approach used by Jupyter kernels +**Option A is the final target.** It is the most robust approach for +true statefulness and is the standard design used by Jupyter kernels and similar systems: - Variables, imports, and objects persist naturally @@ -859,20 +872,38 @@ and similar systems: - No serialization issues - O(1) cost per call (not O(n) like Option C) -Option C (cumulative replay) has a fundamental side-effect replay -problem that cannot be fully mitigated. We recommend **going directly -to Option A** rather than shipping Option C as an interim MVP that -would accumulate technical debt and user-facing bugs. +Implementation is deferred to **Phase 4** of the roadmap (§8) because +it requires a persistent-process protocol (I/O boundary delimiters, +crash detection, REPL lifecycle management) that is a substantial +design effort on its own. See Open Question 1 (§9) for the protocol +design spike. + +**Acceptance criteria for Option A:** +- Code blocks execute in a long-lived Python REPL process +- Variables, imports, and objects persist across calls +- No replay of prior blocks on each new call +- Crash/OOM triggers automatic REPL restart with clear error +- `execution_id` isolates state across concurrent sessions +- `reset_state()` kills and restarts the REPL + +#### 5.2.3 Interim Fallback: Option C (Cumulative Replay) + +If stateful execution is needed before the persistent-process protocol +is ready, Option C may be shipped as a **time-bounded interim** with +the following restrictions: -If a simpler interim is needed before the persistent-process protocol -is ready, Option C may be used with the following restrictions: - Documented as limited to **pure computation only** (variable setup, data transforms, aggregations) - Side-effecting code (file writes, network calls, DB mutations) is explicitly unsupported and will produce incorrect results -- Clearly labeled as experimental / unstable +- Clearly labeled as `experimental` / `unstable` in docstrings and + release notes +- **Deprecation plan:** Option C is removed no later than one minor + release after Option A ships. The `stateful` field docstring must + state: *"Experimental. Uses cumulative replay (Option C). Will be + replaced by persistent-process execution in a future release."* -#### 5.2.3 Implementation Plan (Phase 1, if pursued) +#### 5.2.4 Implementation Plan (Option C Interim, if pursued) 1. **Unfreeze `stateful` in `ContainerCodeExecutor`:** @@ -1123,6 +1154,7 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): def execute_code(self, invocation_context, code_execution_input): import platform + import signal import subprocess import sys import tempfile @@ -1134,15 +1166,45 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): 'Windows. Use ContainerCodeExecutor instead.' ) - with tempfile.NamedTemporaryFile( - mode='w', suffix='.py', delete=True - ) as f: - f.write(code_execution_input.code) - f.flush() + with tempfile.TemporaryDirectory() as temp_dir: + # Materialize input_files into the temp directory, + # preserving relative paths (mirrors + # UnsafeLocalCodeExecutor sandbox behavior). + for f in (code_execution_input.input_files or []): + file_path = os.path.join( + temp_dir, f.path or f.name + ) + os.makedirs( + os.path.dirname(file_path), exist_ok=True + ) + mode = 'wb' if isinstance(f.content, bytes) \ + else 'w' + with open(file_path, mode) as out_f: + out_f.write(f.content) + + # Honor working_dir: resolve as subdirectory of + # temp_dir (same contract as + # UnsafeLocalCodeExecutor). + if code_execution_input.working_dir: + exec_dir = os.path.join( + temp_dir, + code_execution_input.working_dir, + ) + os.makedirs(exec_dir, exist_ok=True) + else: + exec_dir = temp_dir + + # Write code to a temp file inside exec_dir + script_path = os.path.join(exec_dir, '_run.py') + with open(script_path, 'w') as sf: + sf.write(code_execution_input.code) # Build restricted environment - env = {k: os.environ[k] for k in self.allowed_env_vars - if k in os.environ} + env = { + k: os.environ[k] + for k in self.allowed_env_vars + if k in os.environ + } env['PATH'] = '/usr/bin:/usr/local/bin' input_t = code_execution_input.timeout_seconds @@ -1160,8 +1222,6 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): spawn_kwargs['process_group'] = 0 else: # Fallback for 3.10; caveat: not fork-safe. - # Guard resource import for platforms where - # the module is unavailable. def _set_limits(): try: import resource @@ -1194,21 +1254,44 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): cmd = [ 'python3', '-c', limit_code - + f'exec(open({f.name!r}).read())', + + f'exec(open({script_path!r}).read())', ] - result = subprocess.run( - cmd, - capture_output=True, - text=True, - timeout=timeout, - env=env, - cwd=tempfile.gettempdir(), - **spawn_kwargs, - ) + + # Use Popen for explicit process-group kill on + # timeout (subprocess.run only kills the direct + # child, leaving descendants alive). + try: + proc = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + env=env, + cwd=exec_dir, + **spawn_kwargs, + ) + stdout, stderr = proc.communicate( + timeout=timeout, + ) + except subprocess.TimeoutExpired: + # Kill the entire process group to reap + # descendants (fork bombs, child workers). + try: + os.killpg(proc.pid, signal.SIGKILL) + except OSError: + proc.kill() # fallback + proc.wait() + return CodeExecutionResult( + stdout='', + stderr=( + f'Execution timed out after ' + f'{timeout}s' + ), + ) return CodeExecutionResult( - stdout=result.stdout, - stderr=result.stderr if result.returncode != 0 + stdout=stdout, + stderr=stderr if proc.returncode != 0 else '', ) ``` @@ -1299,23 +1382,52 @@ USER executor WORKDIR /home/executor ``` -**C. Network isolation by default:** +**C. Network isolation (opt-in, graduating to default):** + +These hardening options **break scripts that require network access or +write outside `/tmp`**. They are introduced as explicit opt-in flags +first, then graduated to defaults in a later minor release with a +deprecation warning cycle. + +**Phase 3 (initial release) — opt-in:** ```python +class ContainerCodeExecutor(BaseCodeExecutor): + # New opt-in fields (default False to preserve + # backward compatibility). + disable_network: bool = False + read_only_rootfs: bool = False + mem_limit: Optional[str] = None # e.g. '512m' + cpu_quota: Optional[int] = None # e.g. 50000 + def __init_container(self): - self._container = self._client.containers.run( + run_kwargs = dict( image=self.image, detach=True, tty=True, - network_mode='none', # No network access - read_only=True, # Read-only filesystem - tmpfs={'/tmp': 'size=100M'}, # Writable tmp - mem_limit='512m', # Memory limit - cpu_period=100000, - cpu_quota=50000, # 50% of one CPU + ) + if self.disable_network: + run_kwargs['network_mode'] = 'none' + if self.read_only_rootfs: + run_kwargs['read_only'] = True + run_kwargs['tmpfs'] = {'/tmp': 'size=100M'} + if self.mem_limit: + run_kwargs['mem_limit'] = self.mem_limit + if self.cpu_quota: + run_kwargs['cpu_period'] = 100000 + run_kwargs['cpu_quota'] = self.cpu_quota + self._container = self._client.containers.run( + **run_kwargs ) ``` +**Future minor release — graduate to default:** + +After at least one minor release with opt-in availability, flip the +defaults to `True` for `disable_network` and `read_only_rootfs`. +Emit a `DeprecationWarning` for one release before the flip to give +users time to set explicit `False` if they need the old behavior. + ### 6.4 Recommendation Matrix | Use Case | Recommended Executor | Why | @@ -1393,7 +1505,7 @@ class CodeExecutionInput: | `SecurityWarning` on `UnsafeLocalCodeExecutor` | Yes (warning only) | No | | New `LocalSandboxCodeExecutor` | Yes (additive) | No | | `restrict_builtins` on `UnsafeLocalCodeExecutor` | Yes (default `False`) | No | -| Default image for `ContainerCodeExecutor` | Breaking (currently requires image/docker_path) | Minor | +| Default image for `ContainerCodeExecutor` | Breaking behavior change (currently requires explicit `image` or `docker_path`; adding a default changes what runs). Versioning: ship in a minor release with a `DeprecationWarning` for one cycle, then make it the hard default. | Emit `DeprecationWarning` when no image specified; docs must show explicit `image=` in all examples during transition. | ### 7.3 Impact on `RunSkillScriptTool` @@ -1642,7 +1754,8 @@ Implement Option A (persistent process) directly, as recommended in key design question is how to delimit output for each code block: sentinel strings in stdout, JSON-envelope protocol, or a side channel (e.g., file-based result). Sentinel strings are simplest - but can collide with user output. Decision: spike during Phase 2. + but can collide with user output. Decision: spike during Phase 4 + (Stateful Container, §8). 2. **Should `LocalSandboxCodeExecutor` support stateful execution?** Subprocess-based execution is inherently stateless. Stateful support From c950d793a28cc0f23c0762123b8bf9d6d85a0066 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 01:16:47 -0800 Subject: [PATCH 38/53] docs: address 6 cross-doc review findings in design doc and RFC MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - High: Align timeout/lock model — lock releases normally via `with`, unhealthy flag prevents new executions (consistent across both docs) - High: Unify Python 3.10 fallback — preexec_fn calls os.setpgrp() AND sets resource limits in both design doc and RFC - Medium: Phase 3 roadmap now says opt-in for container hardening, matching §6.3.2-C graduation plan; default flip deferred to later - Medium: Fix duplicate §5.2.4 numbering (→5.2.5), deduplicate execution_id wiring (Phase 2 owns it, Phase 4 cross-references) - Low: Add scope note linking RFC (8-11d, P0 only) to full design doc (15-22d, P0-P2) - Low: RFC threat table specifies temp-dir condition; lock impact clarified as process-global Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 62 ++++++++++++++--------- docs/design/rfc_runskillscript_p0.md | 28 ++++++---- 2 files changed, 57 insertions(+), 33 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index 7e5a03ea50..173654b88c 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -403,14 +403,13 @@ account for this: it after the join (whether the worker finishes or times out). If the lock were acquired inside the worker thread, a timed-out worker would hold the lock indefinitely, deadlocking all subsequent calls. -- **Recommended pattern — unhealthy guard:** When a timeout fires, the - executor marks itself as unhealthy (`_healthy = False`) and **does not - release the lock**. All subsequent `execute_code()` calls fail fast - with an explicit error until the caller invokes `reinitialize()`, - which waits for the lingering daemon thread (best-effort), resets - `_healthy = True`, and releases the lock. This prevents the lock - from being released while the daemon thread still mutates - `sys.stdout` or the working directory. +- **Recommended pattern — unhealthy guard:** On timeout, the executor + sets `self._healthy = False` before returning the timeout error. The + lock is released normally when the `with` block exits (a `with` + statement always releases on any exit path, including `return`). + The daemon thread may still be running after the lock is released; + the unhealthy flag prevents new executions from starting until the + caller explicitly recovers via `reinitialize()`. ```python def execute_code(self, ...): if not self._healthy: @@ -424,18 +423,23 @@ account for this: thread.join(timeout=timeout) if thread.is_alive(): self._healthy = False - # Lock stays held — no new exec until reinit. + # Lock released when `with` exits. + # Unhealthy flag blocks future calls. return CodeExecutionResult( stderr=f'Execution timed out after ' f'{timeout}s. Executor marked ' f'unhealthy.' ) ``` +- **Risk:** After the lock is released, the lingering daemon thread + may still be mutating `sys.stdout` or the working directory. The + unhealthy guard ensures no *new* execution races with it, but a + concurrent read of `sys.stdout` from other code could see stale + data. This is an acceptable trade-off for a development executor. - **Recovery path:** `reinitialize()` joins the lingering daemon thread (with a generous grace timeout), restores `sys.stdout` - and cwd if needed, sets `_healthy = True`, and releases - `_execution_lock`. This mirrors the `ContainerCodeExecutor` - recovery model (§4.2.3). + and cwd if needed, and sets `_healthy = True`. This mirrors the + `ContainerCodeExecutor` recovery model (§4.2.3). **Recommendation:** Thread-based timeout with unhealthy guard for `UnsafeLocalCodeExecutor` is sufficient for a development executor. @@ -975,7 +979,7 @@ def reset_state(self): # Keep optimize_data_file frozen ``` -#### 5.2.4 Interaction with `execution_id` +#### 5.2.5 Interaction with `execution_id` The LLM flow layer uses `execution_id` (from `CodeExecutorContext`) to identify stateful sessions. For `ContainerCodeExecutor`: @@ -1221,8 +1225,12 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): if sys.version_info >= (3, 11): spawn_kwargs['process_group'] = 0 else: - # Fallback for 3.10; caveat: not fork-safe. - def _set_limits(): + # Fallback for 3.10; caveat: preexec_fn is not + # fork-safe in multi-threaded programs. + # Must call os.setpgrp() so os.killpg() works + # on timeout, AND set resource limits. + def _setup_child(): + os.setpgrp() # new process group try: import resource resource.setrlimit( @@ -1235,7 +1243,7 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): ) except (ImportError, OSError): pass # timeout-only enforcement - spawn_kwargs['preexec_fn'] = _set_limits + spawn_kwargs['preexec_fn'] = _setup_child # Inline wrapper sets resource limits in the child # process. Guarded for missing resource module. @@ -1446,9 +1454,10 @@ users time to set explicit `False` if they need the old behavior. | 1 | Add `SecurityWarning` to `UnsafeLocalCodeExecutor` | Small | None | | 2 | Add `restrict_builtins` option | Small | Low | | 3 | Implement `LocalSandboxCodeExecutor` | Medium | Low | -| 4 | Add default image + network isolation to `ContainerCodeExecutor` | Medium | Low | +| 4 | Add default image + network/rootfs isolation as **opt-in** flags | Medium | Low | | 5 | Create official `adk-code-executor` Docker image | Medium | Low | | 6 | Update documentation and samples | Small | None | +| 7 | Graduate network/rootfs isolation to **default** (future minor) | Small | Medium (breaking for scripts needing network/host writes) | --- @@ -1720,7 +1729,9 @@ should be opt-in or gated behind a version flag. 3. Implement `LocalSandboxCodeExecutor` (using `process_group`, not `preexec_fn`) 4. Add digest-pinned default image to `ContainerCodeExecutor` -5. Add network isolation defaults to `ContainerCodeExecutor` +5. Add network/rootfs isolation as **opt-in** flags on + `ContainerCodeExecutor` (`disable_network`, `read_only_rootfs`, + both default `False`; see §6.3.2-C for graduation plan) 6. Create official `adk-code-executor` Docker image (versioned tags) 7. Update all samples to recommend secure executors 8. Add security-focused tests @@ -1736,15 +1747,20 @@ Implement Option A (persistent process) directly, as recommended in 3. Implement persistent Python REPL management (start, send code, read output, detect crash/restart) 4. Add `execution_id`-based session isolation (one REPL per - `execution_id`) -5. Wire `execution_id` in `RunSkillScriptTool` -6. Add `reset_state()` method (kills and restarts the REPL) -7. Add stateful execution tests (variable persistence, crash recovery, + `execution_id`; wiring in `RunSkillScriptTool` is done in + Phase 2 — see §5.2.5) +5. Add `reset_state()` method (kills and restarts the REPL) +6. Add stateful execution tests (variable persistence, crash recovery, `execution_id` isolation) -8. Update samples and documentation +7. Update samples and documentation ### Total estimated effort: 15-22 days +> **Scope note:** This covers all four phases (P0 timeout + P0 security +> + P1 contract hardening + P2 stateful container). The companion RFC +> (`rfc_runskillscript_p0.md`) scopes only P0-A (timeout) and P0-B +> (LocalSandboxCodeExecutor), estimated at 8-11 days. + --- ## 9. Open Questions diff --git a/docs/design/rfc_runskillscript_p0.md b/docs/design/rfc_runskillscript_p0.md index 6824bf5c33..61fc0d2829 100644 --- a/docs/design/rfc_runskillscript_p0.md +++ b/docs/design/rfc_runskillscript_p0.md @@ -4,7 +4,9 @@ **Date:** 2026-02-24 **Status:** Proposed **Audience:** ADK TL / UTL -**Effort:** 8-11 engineering days +**Effort:** 8-11 engineering days (P0 scope only; see + `code_executor_enhancements.md` for full 15-22 day estimate + covering P0-P2) **Tracking:** Follow-up to PR #4575 (RunSkillScriptTool) --- @@ -71,8 +73,9 @@ watchdog thread, and no way for the caller to interrupt it. Worse, `UnsafeLocalCodeExecutor` holds a process-global `threading.Lock()` for the entire execution (required because `redirect_stdout` and `os.chdir` mutate process-global state). A hung Python script -deadlocks the lock, blocking **all** code execution across all agents -sharing that executor instance. +deadlocks the lock, blocking **all** code execution across the entire +process — every agent and every `UnsafeLocalCodeExecutor` instance +shares the same module-level `_execution_lock`. **Impact:** A single infinite-loop in a skill script takes down the entire ADK process. This is a denial-of-service risk for any @@ -87,7 +90,7 @@ host Python process: | Threat | Impact | Current Mitigation | |--------|--------|--------------------| | Read env vars / secrets | Data exfiltration | None | -| Write to host filesystem | Data loss / corruption | Partial (temp-dir when `input_files` set) | +| Write to host filesystem | Data loss / corruption | Partial (temp-dir sandbox when `input_files` or `working_dir` is set) | | Outbound network calls | Data leak | None | | `sys.exit()` | Process crash | `SystemExit` caught in tool | | Infinite loop / fork bomb | DoS | Shell: `subprocess.run(timeout)`; Python: **none** | @@ -329,9 +332,11 @@ entire group — the child and all its descendants. The follow-up `proc.communicate(timeout=5)` reaps any zombies. On Python 3.10 (where `process_group` is unavailable), the fallback -uses `preexec_fn=os.setpgrp` to achieve the same process-group -isolation, with a documented caveat about fork-safety in -multi-threaded programs (see §3.2.3). +uses `preexec_fn=_setup_child` where `_setup_child` calls +`os.setpgrp()` (process-group isolation) **and** sets +`resource.setrlimit` for CPU/memory. This achieves the same +kill-group semantics, with a documented caveat about fork-safety +in multi-threaded programs (see §3.2.3). The inline `limit_code` wrapper sets `resource.setrlimit` for CPU and memory inside the child process (guarded with `try/except ImportError` @@ -359,9 +364,12 @@ require containers or OS-level policy. - **Python 3.11+:** Use `process_group=0` (fork-safe, replaces `preexec_fn`). Enables clean `os.killpg()` on timeout. -- **Python 3.10:** Fall back to `preexec_fn=set_limits` with a - documented caveat about thread safety in multi-threaded programs. - ADK minimum is `>=3.10` per `pyproject.toml`. +- **Python 3.10:** Fall back to `preexec_fn=_setup_child` where + `_setup_child` calls `os.setpgrp()` (new process group, required + for `os.killpg` on timeout) **and** sets `resource.setrlimit` for + CPU/memory. Documented caveat: `preexec_fn` is not fork-safe in + multi-threaded programs. ADK minimum is `>=3.10` per + `pyproject.toml`. - **Windows:** Not supported. Raise `NotImplementedError` directing users to `ContainerCodeExecutor`. (`resource.setrlimit` and `process_group` are Unix-only.) From d94527c1f903ecad86404830207df75710f458da Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 01:19:52 -0800 Subject: [PATCH 39/53] docs: fix 4 remaining cross-doc inconsistencies MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - High: Close race in unhealthy-guard — add double-check of _healthy after acquiring _execution_lock (both docs) - Medium: Normalize 3.10 fallback to _setup_child (os.setpgrp + rlimits) in design doc platform section - Low: Fix cross-reference §6.3.2-C → §6.3.3-C in roadmap - Low: Fix RFC docstring subprocess.run → Popen + killpg Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 20 ++++++++++++++++---- docs/design/rfc_runskillscript_p0.md | 13 +++++++++---- 2 files changed, 25 insertions(+), 8 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index 173654b88c..c30abe2ae9 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -418,6 +418,14 @@ account for this: 'Call reinitialize() to recover.' ) with _execution_lock: + # Re-check after acquiring the lock: another + # caller may have timed out while we waited. + if not self._healthy: + return CodeExecutionResult( + stderr='Executor unhealthy after ' + 'timeout. Call reinitialize() ' + 'to recover.' + ) thread = threading.Thread(target=_run, daemon=True) thread.start() thread.join(timeout=timeout) @@ -1324,10 +1332,14 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): held. ADK executors may be called from async/threaded contexts. - `process_group=0` (Python 3.11+) is fork-safe and places the child in its own process group, enabling clean `os.killpg()` on timeout. -- Resource limits are set via an inline `-c` wrapper script instead - of `preexec_fn`, avoiding the fork-safety issue entirely. +- On Python 3.11+, resource limits are set via an inline `-c` + wrapper script (not `preexec_fn`), avoiding the fork-safety issue. - On Python 3.10 (ADK minimum is `>=3.10`), fall back to - `preexec_fn=set_limits` with a documented caveat about thread safety. + `preexec_fn=_setup_child` which calls `os.setpgrp()` (required + for `os.killpg` on timeout) **and** sets `resource.setrlimit` + for CPU/memory. The inline `-c` wrapper still applies limits as + a defense-in-depth layer. Documented caveat: `preexec_fn` is not + fork-safe in multi-threaded programs. **Dependencies:** None (stdlib only). This is the key advantage over `ContainerCodeExecutor`. @@ -1731,7 +1743,7 @@ should be opt-in or gated behind a version flag. 4. Add digest-pinned default image to `ContainerCodeExecutor` 5. Add network/rootfs isolation as **opt-in** flags on `ContainerCodeExecutor` (`disable_network`, `read_only_rootfs`, - both default `False`; see §6.3.2-C for graduation plan) + both default `False`; see §6.3.3-C for graduation plan) 6. Create official `adk-code-executor` Docker image (versioned tags) 7. Update all samples to recommend secure executors 8. Add security-focused tests diff --git a/docs/design/rfc_runskillscript_p0.md b/docs/design/rfc_runskillscript_p0.md index 61fc0d2829..33e49ce28d 100644 --- a/docs/design/rfc_runskillscript_p0.md +++ b/docs/design/rfc_runskillscript_p0.md @@ -192,9 +192,13 @@ The solution is to **mark the executor unhealthy on timeout**: 1. On timeout, set `self._healthy = False` before returning the timeout error. The lock is released normally when the `with` block exits. -2. Subsequent `execute_code()` calls check `self._healthy` at the top - and **fail fast** with a clear error: `"Executor is unhealthy after - a timed-out execution. Call reinitialize() to recover."` +2. Subsequent `execute_code()` calls check `self._healthy` **both + before and after** acquiring `_execution_lock` (double-check + pattern: the pre-lock check is a fast path; the post-lock check + closes the race where a caller passes the pre-lock check, waits + on the lock, and enters after another call has timed out). Both + checks **fail fast** with: `"Executor is unhealthy after a + timed-out execution. Call reinitialize() to recover."` 3. `reinitialize()` waits for the lingering daemon thread to finish (with a generous join timeout), resets `_healthy = True`, and allows execution to resume. @@ -292,7 +296,8 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): - Resource limits (CPU time, memory) via resource.setrlimit - Restricted environment variables - Temporary working directory - - subprocess.run(timeout=N) for wall-clock timeout + - Popen + communicate(timeout=N) + os.killpg for wall-clock + timeout with process-group cleanup """ default_timeout_seconds: int = 30 From faa5fa22e1f3cb089560cb8f2d871976f9880ad7 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 01:23:15 -0800 Subject: [PATCH 40/53] =?UTF-8?q?docs:=20final=20polish=20=E2=80=94=20SemV?= =?UTF-8?q?er=20policy,=20missing=20import,=20reinit=20failure=20mode?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Medium: Reconcile breaking default-flips with SemVer — pre-1.0 may ship in minor with DeprecationWarning; post-1.0 reserve for major - Low: Add missing `import os` in LocalSandboxCodeExecutor snippet - Low: Specify reinitialize() failure mode — stays unhealthy and raises RuntimeError if daemon thread does not exit within grace period (both docs) Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 37 +++++++++++++++++------ docs/design/rfc_runskillscript_p0.md | 9 ++++-- 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index c30abe2ae9..a9fae57aa5 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -445,9 +445,15 @@ account for this: concurrent read of `sys.stdout` from other code could see stale data. This is an acceptable trade-off for a development executor. - **Recovery path:** `reinitialize()` joins the lingering daemon - thread (with a generous grace timeout), restores `sys.stdout` - and cwd if needed, and sets `_healthy = True`. This mirrors the - `ContainerCodeExecutor` recovery model (§4.2.3). + thread with a generous grace timeout (e.g., 30 s). If the thread + exits within the grace period, restore `sys.stdout` and cwd, set + `_healthy = True`, and return success. If the thread is **still + alive** after the grace timeout, the executor remains unhealthy + (`_healthy = False`) and `reinitialize()` raises + `RuntimeError('Timed-out thread did not exit within grace period; + executor remains unhealthy.')`. The caller may retry later or + restart the process. This mirrors the `ContainerCodeExecutor` + recovery model (§4.2.3). **Recommendation:** Thread-based timeout with unhealthy guard for `UnsafeLocalCodeExecutor` is sufficient for a development executor. @@ -1165,6 +1171,7 @@ class LocalSandboxCodeExecutor(BaseCodeExecutor): allowed_env_vars: list[str] = [] def execute_code(self, invocation_context, code_execution_input): + import os import platform import signal import subprocess @@ -1441,12 +1448,24 @@ def __init_container(self): ) ``` -**Future minor release — graduate to default:** +**Future release — graduate to default:** -After at least one minor release with opt-in availability, flip the -defaults to `True` for `disable_network` and `read_only_rootfs`. -Emit a `DeprecationWarning` for one release before the flip to give -users time to set explicit `False` if they need the old behavior. +Flipping `disable_network` and `read_only_rootfs` from `False` to +`True` is a **breaking behavior change** under SemVer (existing code +that relies on network access or host-filesystem writes will fail). +The graduation plan must follow ADK's SemVer policy: + +- **Option A (minor release, preferred if ADK is pre-1.0):** Pre-1.0 + SemVer permits breaking changes in minor releases. Emit a + `DeprecationWarning` for at least one minor release before the + flip. Users who need the old behavior set explicit `False`. +- **Option B (post-1.0):** Reserve the default flip for a **major** + release. In the preceding minor release(s), emit a + `DeprecationWarning` when these flags are unset, warning that the + default will change in the next major version. + +Either way, the deprecation warning must state the exact version +where the default changes and the explicit opt-out syntax. ### 6.4 Recommendation Matrix @@ -1526,7 +1545,7 @@ class CodeExecutionInput: | `SecurityWarning` on `UnsafeLocalCodeExecutor` | Yes (warning only) | No | | New `LocalSandboxCodeExecutor` | Yes (additive) | No | | `restrict_builtins` on `UnsafeLocalCodeExecutor` | Yes (default `False`) | No | -| Default image for `ContainerCodeExecutor` | Breaking behavior change (currently requires explicit `image` or `docker_path`; adding a default changes what runs). Versioning: ship in a minor release with a `DeprecationWarning` for one cycle, then make it the hard default. | Emit `DeprecationWarning` when no image specified; docs must show explicit `image=` in all examples during transition. | +| Default image for `ContainerCodeExecutor` | Breaking behavior change (currently requires explicit `image` or `docker_path`; adding a default changes what runs). Pre-1.0: may ship in a minor release with one-cycle `DeprecationWarning`. Post-1.0: reserve for a major release per SemVer. | Emit `DeprecationWarning` when no image specified; docs must show explicit `image=` in all examples during transition. | ### 7.3 Impact on `RunSkillScriptTool` diff --git a/docs/design/rfc_runskillscript_p0.md b/docs/design/rfc_runskillscript_p0.md index 33e49ce28d..593d402b9e 100644 --- a/docs/design/rfc_runskillscript_p0.md +++ b/docs/design/rfc_runskillscript_p0.md @@ -199,9 +199,12 @@ The solution is to **mark the executor unhealthy on timeout**: on the lock, and enters after another call has timed out). Both checks **fail fast** with: `"Executor is unhealthy after a timed-out execution. Call reinitialize() to recover."` -3. `reinitialize()` waits for the lingering daemon thread to finish - (with a generous join timeout), resets `_healthy = True`, and - allows execution to resume. +3. `reinitialize()` joins the lingering daemon thread with a grace + timeout (e.g., 30 s). If the thread exits, it resets + `_healthy = True` and allows execution to resume. If the thread + is **still alive** after the grace period, the executor stays + unhealthy and `reinitialize()` raises `RuntimeError` — the + caller must retry later or restart the process. This ensures that no new execution can run while a zombie thread is still alive and mutating shared state. The pattern mirrors the From ab1dd285ece970fab0407f7dd07b05510868ab0a Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 01:25:18 -0800 Subject: [PATCH 41/53] =?UTF-8?q?docs(code-executor):=20align=20=C2=A76.5?= =?UTF-8?q?=20graduation=20row=20with=20SemVer=20policy?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Change "future minor" to "future release: minor pre-1.0 / major post-1.0" to match the policy stated in §6.3.3-C. Co-Authored-By: Claude Opus 4.6 --- docs/design/code_executor_enhancements.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/design/code_executor_enhancements.md b/docs/design/code_executor_enhancements.md index a9fae57aa5..0d899a8600 100644 --- a/docs/design/code_executor_enhancements.md +++ b/docs/design/code_executor_enhancements.md @@ -1488,7 +1488,7 @@ where the default changes and the explicit opt-out syntax. | 4 | Add default image + network/rootfs isolation as **opt-in** flags | Medium | Low | | 5 | Create official `adk-code-executor` Docker image | Medium | Low | | 6 | Update documentation and samples | Small | None | -| 7 | Graduate network/rootfs isolation to **default** (future minor) | Small | Medium (breaking for scripts needing network/host writes) | +| 7 | Graduate network/rootfs isolation to **default** (future release: minor pre-1.0 / major post-1.0; see §6.3.3-C) | Small | Medium (breaking for scripts needing network/host writes) | --- From 027237008d94719619d6d24ae6aa49d431c70f59 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 02:16:12 -0800 Subject: [PATCH 42/53] feat(benchmarks): add BigQueryBench evaluation pipeline Reusable end-to-end evaluation framework for any BigQuery skill or agent built with ADK BigQueryToolset. Mirrors SkillsBench architecture with BigQuery-specific metrics and public dataset eval cases. Contents: - agent.py: Root agent with BigQueryToolset (read-only, env-configurable) - metrics.py: 3 custom metrics (schema_discovery, tool_usage, output_correctness) - runner.py: Standalone runner with filter, dry-run, multi-run support - eval_sets/bigquerybench_eval.json: 7 eval cases across 3 tiers (schema exploration, SQL generation, multi-step analysis) - README.md: Full pipeline docs with "add new eval case" guide and AI/ML tool templates (forecast, detect_anomalies, analyze_contribution) Co-Authored-By: Claude Opus 4.6 --- benchmarks/bigquerybench/README.md | 403 ++++++++++++++++++ benchmarks/bigquerybench/__init__.py | 13 + benchmarks/bigquerybench/agent.py | 79 ++++ .../eval_sets/bigquerybench_eval.json | 187 ++++++++ benchmarks/bigquerybench/metrics.py | 228 ++++++++++ benchmarks/bigquerybench/runner.py | 365 ++++++++++++++++ 6 files changed, 1275 insertions(+) create mode 100644 benchmarks/bigquerybench/README.md create mode 100644 benchmarks/bigquerybench/__init__.py create mode 100644 benchmarks/bigquerybench/agent.py create mode 100644 benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json create mode 100644 benchmarks/bigquerybench/metrics.py create mode 100644 benchmarks/bigquerybench/runner.py diff --git a/benchmarks/bigquerybench/README.md b/benchmarks/bigquerybench/README.md new file mode 100644 index 0000000000..dc80173e17 --- /dev/null +++ b/benchmarks/bigquerybench/README.md @@ -0,0 +1,403 @@ +# BigQueryBench: End-to-End Evaluation for BigQuery Skills + +## Overview + +BigQueryBench is a reusable evaluation pipeline for any skill or agent +built with ADK's `BigQueryToolset`. It mirrors the SkillsBench +architecture (`benchmarks/skillsbench/`) but targets BigQuery-specific +tool chains: schema exploration, SQL generation, AI/ML operations +(forecast, anomaly detection, contribution analysis), and multi-step +analytical workflows. + +**Design goals:** + +- **Reusable:** Add a new BigQuery skill by dropping one eval case + JSON and one optional reference SQL file — no code changes needed. +- **Reproducible:** All eval cases use BigQuery public datasets + (`bigquery-public-data`) so any GCP project with BigQuery API + enabled can run the suite. +- **Layered metrics:** Three dimensions scored independently — + schema discovery, tool-call coverage, output correctness. +- **CI-friendly:** Single `python -m benchmarks.bigquerybench.runner` + invocation with JSON results and exit code. + +## Quick Start + +```bash +# Prerequisites +# 1. GCP project with BigQuery API enabled +# 2. Application Default Credentials configured: +# gcloud auth application-default login +# 3. ADK installed with BigQuery extras: +# uv sync --all-extras + +# Set environment +export GOOGLE_CLOUD_PROJECT=your-project-id +export GOOGLE_GENAI_USE_VERTEXAI=1 # or use GOOGLE_API_KEY + +# Run all eval cases +python -m benchmarks.bigquerybench.runner + +# Run specific eval case(s) +python -m benchmarks.bigquerybench.runner --filter sql_public_dataset + +# Run with multiple attempts for variance measurement +python -m benchmarks.bigquerybench.runner --num-runs 3 + +# Dry-run mode (validates eval set JSON, no LLM calls) +python -m benchmarks.bigquerybench.runner --dry-run +``` + +## Architecture + +``` +benchmarks/bigquerybench/ +├── README.md # This file +├── __init__.py +├── agent.py # Root agent with BigQueryToolset +├── runner.py # Standalone evaluation runner +├── metrics.py # BigQuery-specific custom metrics +└── eval_sets/ + └── bigquerybench_eval.json # Eval cases (public datasets) +``` + +### Relationship to SkillsBench + +| Aspect | SkillsBench | BigQueryBench | +|--------|-------------|---------------| +| Toolset | `SkillToolset` (4 tools) | `BigQueryToolset` (10 tools) | +| Discovery tools | `list_skills`, `load_skill` | `list_dataset_ids`, `list_table_ids`, `get_dataset_info`, `get_table_info` | +| Execution tools | `run_skill_script` | `execute_sql`, `forecast`, `detect_anomalies`, `analyze_contribution` | +| Data source | Bundled reference files | BigQuery public datasets | +| Auth | None (local files) | GCP credentials (ADC / service account / OAuth) | +| Metrics | discovery, tool_usage, binary | schema_discovery, tool_usage, output_correctness | + +## Evaluation Pipeline + +### Stage 1: Agent Setup + +The agent under test is defined in `agent.py`. It uses +`BigQueryToolset` with read-only defaults: + +```python +from google.adk.agents.llm_agent import LlmAgent +from google.adk.tools.bigquery.bigquery_credentials import ( + BigQueryCredentialsConfig, +) +from google.adk.tools.bigquery.bigquery_toolset import BigQueryToolset +from google.adk.tools.bigquery.config import BigQueryToolConfig +from google.adk.tools.bigquery.config import WriteMode +import google.auth + +credentials, _ = google.auth.default() +credentials_config = BigQueryCredentialsConfig(credentials=credentials) + +tool_config = BigQueryToolConfig( + write_mode=WriteMode.BLOCKED, # Read-only for eval safety + max_query_result_rows=50, +) + +bigquery_toolset = BigQueryToolset( + credentials_config=credentials_config, + bigquery_tool_config=tool_config, +) + +root_agent = LlmAgent( + model="gemini-2.5-flash", + name="bigquerybench_agent", + description="Agent for BigQuery data exploration and analysis.", + instruction="""\ + You are a data analyst with access to BigQuery tools. + Use them to explore schemas, run SQL queries, and answer + the user's questions about data. Always explore the schema + (list datasets, list tables, get table info) before writing + SQL. Show query results clearly. + """, + tools=[bigquery_toolset], +) +``` + +### Stage 2: Eval Set Definition + +Each eval case is a JSON object in `eval_sets/bigquerybench_eval.json` +following the ADK `EvalSet` schema. The key fields are: + +```json +{ + "eval_id": "unique_case_id", + "conversation": [ + { + "invocation_id": "inv-01", + "user_content": { + "parts": [{"text": "User's question about BigQuery data"}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "Expected key phrases in the answer"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "list_dataset_ids", "args": {"project_id": "bigquery-public-data"}}, + {"name": "get_table_info", "args": {"project_id": "...", "dataset_id": "...", "table_id": "..."}}, + {"name": "execute_sql", "args": {"project_id": "...", "query": "SELECT ..."}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 +} +``` + +**Key conventions:** + +- `final_response.parts[0].text` contains **reference lines** — key + phrases that must appear in the agent's response for a pass. One + phrase per line. Case-insensitive substring match. +- `intermediate_data.tool_uses` lists the **expected tool call + sequence**. The `tool_usage` metric checks set coverage (not + strict ordering). The `schema_discovery` metric checks that at + least one schema-exploration tool was called. +- All eval cases use `bigquery-public-data` project datasets so + results are deterministic and reproducible. + +### Stage 3: Metrics + +Three custom metrics in `metrics.py`, following the ADK custom metric +function signature: + +#### 3a. `schema_discovery_score` + +Checks whether the agent explored the schema before querying. Scores +1.0 if any of these tools were called: `list_dataset_ids`, +`list_table_ids`, `get_dataset_info`, `get_table_info`. Scores 0.0 +otherwise. + +**Rationale:** Agents that skip schema exploration and guess table +names produce fragile SQL that breaks on schema changes. This metric +enforces the "explore before query" pattern. + +#### 3b. `tool_usage_score` + +Fraction of expected tool calls actually made: +`|expected_tools ∩ actual_tools| / |expected_tools|`. + +Uses set-based matching (any order). Passes at threshold >= 0.5. + +Same semantics as SkillsBench `tool_usage_score`, reused for +consistency. + +#### 3c. `output_correctness_score` + +Binary pass/fail: 1.0 if the agent's final response contains all +expected reference lines (case-insensitive substring match). 0.0 +otherwise. + +Same semantics as SkillsBench `skillsbench_binary_score`, reused for +consistency. + +### Stage 4: Runner Execution + +``` +runner.py + ↓ +Load agent from benchmarks/bigquerybench/agent.py + ↓ +Load eval set from eval_sets/bigquerybench_eval.json + ↓ +For each eval case: + ↓ + Run agent.run_async(user_query) via ADK Runner + ↓ (generates events) + Convert events → Invocation (with intermediate_data.tool_uses) + ↓ + Apply schema_discovery_score + Apply tool_usage_score + Apply output_correctness_score + ↓ + Record per-case scores + ↓ +Aggregate scores → leaderboard summary + ↓ +Print table + exit code (0 = all pass, 1 = any fail) +``` + +## Eval Case Catalog + +The following eval cases are included. They are organized by +complexity tier to test progressively harder agent capabilities. + +### Tier 1: Schema Exploration + +These test the agent's ability to navigate BigQuery metadata. + +| eval_id | Dataset | User Query | Expected Tools | +|---------|---------|-----------|----------------| +| `schema_list_datasets` | `bigquery-public-data` | "What datasets are available in bigquery-public-data?" | `list_dataset_ids` | +| `schema_list_tables` | `usa_names` | "What tables exist in the usa_names dataset?" | `list_dataset_ids` → `list_table_ids` | +| `schema_get_table_info` | `usa_names.usa_1910_current` | "What columns and types does the usa_1910_current table have?" | `list_table_ids` → `get_table_info` | + +### Tier 2: SQL Generation & Execution + +These test SQL generation against public data with known answers. + +| eval_id | Dataset | User Query | Expected Tools | Reference Output | +|---------|---------|-----------|----------------|-----------------| +| `sql_top_names` | `usa_names` | "What are the top 5 most popular baby names in 2020?" | `get_table_info` → `execute_sql` | Top names by count | +| `sql_aggregation` | `usa_names` | "How many distinct names were registered each decade from 1950 to 2000?" | `get_table_info` → `execute_sql` | Decade counts | +| `sql_public_dataset` | `samples.shakespeare` | "Which Shakespeare work has the most unique words?" | `get_table_info` → `execute_sql` | Work name + count | + +### Tier 3: Multi-Step Analysis + +These test the agent's ability to chain multiple tools. + +| eval_id | Dataset | User Query | Expected Tools | Reference Output | +|---------|---------|-----------|----------------|-----------------| +| `multi_step_explore_and_query` | `austin_bikeshare` | "Explore the Austin bikeshare dataset and tell me the top 5 busiest stations by trip count." | `list_table_ids` → `get_table_info` → `execute_sql` | Station names + counts | + +## Adding a New BigQuery Eval Case + +To add a new eval case (e.g., for a new BigQuery AI operator skill): + +### Step 1: Identify the public dataset + +Pick a dataset from `bigquery-public-data` that exercises the skill. +Verify it exists: + +```sql +SELECT * FROM `bigquery-public-data.DATASET.INFORMATION_SCHEMA.TABLES` +LIMIT 5; +``` + +### Step 2: Write the eval case JSON + +Add a new object to the `eval_cases` array in +`eval_sets/bigquerybench_eval.json`: + +```json +{ + "eval_id": "your_unique_eval_id", + "conversation": [ + { + "invocation_id": "inv-your-id-01", + "user_content": { + "parts": [{"text": "Your user query here"}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "reference line 1\nreference line 2"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "tool_name", "args": {"arg1": "val1"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 +} +``` + +### Step 3: Choose reference lines + +Run the expected query manually and pick 3-5 key phrases from the +result that are **stable** (won't change if data is appended). Good +reference lines: + +- Column names or schema facts ("column: name, type: STRING") +- Aggregation results from historical data ("hamlet", "king lear") +- Structural facts ("3 tables", "5 columns") + +Avoid: exact row counts on append-only tables, floating-point values +that may shift with precision. + +### Step 4: Validate + +```bash +# Run just your new case +python -m benchmarks.bigquerybench.runner --filter your_unique_eval_id +``` + +### Step 5: Commit + +Add only the modified `bigquerybench_eval.json`. No code changes +needed. + +## Eval Case Template for AI/ML Tools + +For `forecast`, `detect_anomalies`, and `analyze_contribution` skills, +use this template: + +```json +{ + "eval_id": "forecast_weather_temperature", + "conversation": [ + { + "invocation_id": "inv-forecast-01", + "user_content": { + "parts": [{"text": "Forecast the next 7 days of average temperature using the NOAA GSOD weather data for station 725300 (Chicago O'Hare) from 2023."}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "forecast_timestamp\nforecast_value\nprediction_interval"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "noaa_gsod", "table_id": "gsod2023"}}, + {"name": "forecast", "args": {"project_id": "your-project", "history_data": "SELECT date, temp FROM `bigquery-public-data.noaa_gsod.gsod2023` WHERE stn = '725300'", "timestamp_col": "date", "data_col": "temp", "horizon": 7}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 +} +``` + +**Key considerations for AI/ML eval cases:** + +- `forecast`, `analyze_contribution`, and `detect_anomalies` create + temporary BigQuery ML models. Ensure `write_mode` is at least + `PROTECTED` (anonymous dataset) or `ALLOWED`. +- Reference lines should validate **structural output** (column names + like `forecast_timestamp`, `is_anomaly`) rather than exact numeric + values, since ML model outputs vary across runs. +- Set `tool_config.write_mode = WriteMode.PROTECTED` in the agent + for AI/ML eval cases that need to create temp models. + +## Metrics Reference + +| Metric | Function Path | Threshold | Pass Condition | +|--------|--------------|-----------|----------------| +| Schema Discovery | `benchmarks.bigquerybench.metrics.schema_discovery_score` | 1.0 | Any schema tool called | +| Tool Usage | `benchmarks.bigquerybench.metrics.tool_usage_score` | 0.5 | >= 50% expected tools called | +| Output Correctness | `benchmarks.bigquerybench.metrics.output_correctness_score` | 1.0 | All reference lines present | + +## Environment Variables + +| Variable | Required | Description | +|----------|----------|-------------| +| `GOOGLE_CLOUD_PROJECT` | Yes | GCP project for BigQuery API calls | +| `GOOGLE_GENAI_USE_VERTEXAI` | Conditional | Set to `1` for Vertex AI LLM backend | +| `GOOGLE_API_KEY` | Conditional | API key for Google AI Studio backend | +| `BQ_EVAL_WRITE_MODE` | No | Override write mode (`blocked`/`protected`/`allowed`). Default: `blocked` | + +## Troubleshooting + +| Symptom | Cause | Fix | +|---------|-------|-----| +| `403 Access Denied` | Missing BigQuery API access | Enable BigQuery API in GCP console; run `gcloud auth application-default login` | +| `execute_sql` returns empty | Query references wrong project | Ensure public dataset queries use `bigquery-public-data` as project | +| `forecast` fails with write error | `write_mode=BLOCKED` | Set `BQ_EVAL_WRITE_MODE=protected` for AI/ML eval cases | +| Low `schema_discovery_score` | Agent skips exploration | Strengthen agent instructions to always explore schema first | +| Flaky `output_correctness_score` | Reference lines too specific | Use structural phrases, not exact numeric values | diff --git a/benchmarks/bigquerybench/__init__.py b/benchmarks/bigquerybench/__init__.py new file mode 100644 index 0000000000..58d482ea38 --- /dev/null +++ b/benchmarks/bigquerybench/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/benchmarks/bigquerybench/agent.py b/benchmarks/bigquerybench/agent.py new file mode 100644 index 0000000000..9071e9a581 --- /dev/null +++ b/benchmarks/bigquerybench/agent.py @@ -0,0 +1,79 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""BigQueryBench evaluation agent. + +Uses BigQueryToolset with read-only defaults against BigQuery public +datasets. Override write_mode via BQ_EVAL_WRITE_MODE env var when +evaluating AI/ML tools (forecast, detect_anomalies, etc.). +""" + +import os + +from google.adk.agents.llm_agent import LlmAgent +from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsConfig +from google.adk.tools.bigquery.bigquery_toolset import BigQueryToolset +from google.adk.tools.bigquery.config import BigQueryToolConfig +from google.adk.tools.bigquery.config import WriteMode +import google.auth + +_WRITE_MODE_MAP = { + "blocked": WriteMode.BLOCKED, + "protected": WriteMode.PROTECTED, + "allowed": WriteMode.ALLOWED, +} + +_write_mode_str = os.environ.get("BQ_EVAL_WRITE_MODE", "blocked").lower() +_write_mode = _WRITE_MODE_MAP.get(_write_mode_str, WriteMode.BLOCKED) + +application_default_credentials, _ = google.auth.default() +credentials_config = BigQueryCredentialsConfig( + credentials=application_default_credentials, +) + +tool_config = BigQueryToolConfig( + write_mode=_write_mode, + max_query_result_rows=50, +) + +bigquery_toolset = BigQueryToolset( + credentials_config=credentials_config, + bigquery_tool_config=tool_config, +) + +root_agent = LlmAgent( + model="gemini-2.5-flash", + name="bigquerybench_agent", + description=( + "Agent for BigQuery data exploration, SQL execution, and" + " AI/ML operations against public datasets." + ), + instruction="""\ +You are a data analyst with access to BigQuery tools. + +Workflow: +1. Always explore the schema first: use list_dataset_ids, + list_table_ids, and get_table_info to understand the data + before writing any SQL. +2. Use execute_sql to run queries. Prefer explicit column names + over SELECT *. +3. For forecasting, anomaly detection, or contribution analysis, + use the dedicated tools (forecast, detect_anomalies, + analyze_contribution) instead of raw SQL. +4. Present results clearly with column headers and values. + +All public datasets are in project "bigquery-public-data". +""", + tools=[bigquery_toolset], +) diff --git a/benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json b/benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json new file mode 100644 index 0000000000..cc2bacd42f --- /dev/null +++ b/benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json @@ -0,0 +1,187 @@ +{ + "eval_set_id": "bigquerybench-adk-v1", + "name": "BigQueryBench ADK Evaluation", + "description": "End-to-end evaluation cases for BigQuery skills using public datasets. Covers schema exploration, SQL generation, multi-step analysis, and AI/ML operations.", + "eval_cases": [ + { + "eval_id": "schema_list_datasets", + "conversation": [ + { + "invocation_id": "inv-schema-ds-01", + "user_content": { + "parts": [{"text": "What datasets are available in the bigquery-public-data project? List a few of them."}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "usa_names\nsamples\nnoaa_gsod"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "list_dataset_ids", "args": {"project_id": "bigquery-public-data"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 + }, + { + "eval_id": "schema_list_tables", + "conversation": [ + { + "invocation_id": "inv-schema-tbl-01", + "user_content": { + "parts": [{"text": "What tables exist in the usa_names dataset in bigquery-public-data?"}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "usa_1910_current"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "list_table_ids", "args": {"project_id": "bigquery-public-data", "dataset_id": "usa_names"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 + }, + { + "eval_id": "schema_get_table_info", + "conversation": [ + { + "invocation_id": "inv-schema-info-01", + "user_content": { + "parts": [{"text": "What columns and data types does the usa_1910_current table in the usa_names dataset have?"}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "name\nyear\ngender\nstate\nnumber"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "usa_names", "table_id": "usa_1910_current"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 + }, + { + "eval_id": "sql_shakespeare_unique_words", + "conversation": [ + { + "invocation_id": "inv-sql-shk-01", + "user_content": { + "parts": [{"text": "Using the samples.shakespeare table in bigquery-public-data, which work (corpus) has the most unique words? Show the top 3 works by distinct word count."}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "hamlet\nkinghenryv\nkingrichardiii"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "samples", "table_id": "shakespeare"}}, + {"name": "execute_sql", "args": {"project_id": "bigquery-public-data", "query": "SELECT corpus, COUNT(DISTINCT word) AS unique_words FROM `bigquery-public-data.samples.shakespeare` GROUP BY corpus ORDER BY unique_words DESC LIMIT 3"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 + }, + { + "eval_id": "sql_usa_names_top_2020", + "conversation": [ + { + "invocation_id": "inv-sql-names-01", + "user_content": { + "parts": [{"text": "What were the top 5 most popular baby names in the year 2020 across all states? Use the usa_names.usa_1910_current table in bigquery-public-data."}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "Olivia\nEmma\nLiam\nNoah\nAva"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "usa_names", "table_id": "usa_1910_current"}}, + {"name": "execute_sql", "args": {"project_id": "bigquery-public-data", "query": "SELECT name, SUM(number) AS total FROM `bigquery-public-data.usa_names.usa_1910_current` WHERE year = 2020 GROUP BY name ORDER BY total DESC LIMIT 5"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 + }, + { + "eval_id": "sql_names_by_decade", + "conversation": [ + { + "invocation_id": "inv-sql-decade-01", + "user_content": { + "parts": [{"text": "How many distinct baby names were registered in each decade from 1950 to 2000 (inclusive)? Use the usa_names.usa_1910_current table in bigquery-public-data."}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "1950\n1960\n1970\n1980\n1990\n2000"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "usa_names", "table_id": "usa_1910_current"}}, + {"name": "execute_sql", "args": {"project_id": "bigquery-public-data", "query": "SELECT CAST(FLOOR(year / 10) * 10 AS INT64) AS decade, COUNT(DISTINCT name) AS distinct_names FROM `bigquery-public-data.usa_names.usa_1910_current` WHERE year BETWEEN 1950 AND 2009 GROUP BY decade ORDER BY decade"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 + }, + { + "eval_id": "multi_step_explore_and_query", + "conversation": [ + { + "invocation_id": "inv-multi-01", + "user_content": { + "parts": [{"text": "I want to analyze the Austin bikeshare data in bigquery-public-data. First explore what tables and columns are available, then tell me the top 5 busiest start stations by total trip count."}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "start_station_name\ntrip"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "list_table_ids", "args": {"project_id": "bigquery-public-data", "dataset_id": "austin_bikeshare"}}, + {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "austin_bikeshare", "table_id": "bikeshare_trips"}}, + {"name": "execute_sql", "args": {"project_id": "bigquery-public-data", "query": "SELECT start_station_name, COUNT(*) AS trip_count FROM `bigquery-public-data.austin_bikeshare.bikeshare_trips` GROUP BY start_station_name ORDER BY trip_count DESC LIMIT 5"}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 + } + ] +} diff --git a/benchmarks/bigquerybench/metrics.py b/benchmarks/bigquerybench/metrics.py new file mode 100644 index 0000000000..66db5d3f9b --- /dev/null +++ b/benchmarks/bigquerybench/metrics.py @@ -0,0 +1,228 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Custom metrics for BigQueryBench evaluation. + +Three metrics following the ADK custom metric function signature: + + def metric_fn( + eval_metric: EvalMetric, + actual_invocations: list[Invocation], + expected_invocations: Optional[list[Invocation]], + conversation_scenario: Optional[ConversationScenario] = None, + ) -> EvaluationResult + +Reference via dotted path in eval configs: + "benchmarks.bigquerybench.metrics.schema_discovery_score" + "benchmarks.bigquerybench.metrics.tool_usage_score" + "benchmarks.bigquerybench.metrics.output_correctness_score" +""" + +from __future__ import annotations + +from typing import Optional + +from google.adk.evaluation.eval_case import ConversationScenario +from google.adk.evaluation.eval_case import get_all_tool_calls +from google.adk.evaluation.eval_case import Invocation +from google.adk.evaluation.eval_metrics import EvalMetric +from google.adk.evaluation.eval_metrics import EvalStatus +from google.adk.evaluation.evaluator import EvaluationResult +from google.adk.evaluation.evaluator import PerInvocationResult + +# Tools that count as "schema exploration". +_SCHEMA_TOOLS = frozenset({ + "list_dataset_ids", + "list_table_ids", + "get_dataset_info", + "get_table_info", +}) + + +def _get_tool_names(invocations: list[Invocation]) -> list[str]: + """Extract all tool call names from a list of invocations.""" + names = [] + for inv in invocations: + for tool_call in get_all_tool_calls(inv.intermediate_data): + names.append(tool_call.name) + return names + + +def _make_per_invocation( + actual_invocations: list[Invocation], + expected_invocations: Optional[list[Invocation]], + score: float, + status: EvalStatus, +) -> list[PerInvocationResult]: + """Build per-invocation results list.""" + results = [] + for i, actual in enumerate(actual_invocations): + expected = None + if expected_invocations and i < len(expected_invocations): + expected = expected_invocations[i] + results.append( + PerInvocationResult( + actual_invocation=actual, + expected_invocation=expected, + score=score, + eval_status=status, + ) + ) + return results + + +def schema_discovery_score( + eval_metric: EvalMetric, + actual_invocations: list[Invocation], + expected_invocations: Optional[list[Invocation]], + conversation_scenario: Optional[ConversationScenario] = None, +) -> EvaluationResult: + """Score 1.0 if the agent called at least one schema exploration tool. + + Schema tools: list_dataset_ids, list_table_ids, get_dataset_info, + get_table_info. + + This metric enforces the "explore before query" pattern — agents + should understand the schema before generating SQL or calling AI/ML + tools. + """ + tool_names = set(_get_tool_names(actual_invocations)) + called_schema = bool(tool_names & _SCHEMA_TOOLS) + + score = 1.0 if called_schema else 0.0 + status = EvalStatus.PASSED if called_schema else EvalStatus.FAILED + + return EvaluationResult( + overall_score=score, + overall_eval_status=status, + per_invocation_results=_make_per_invocation( + actual_invocations, + expected_invocations, + score, + status, + ), + ) + + +def tool_usage_score( + eval_metric: EvalMetric, + actual_invocations: list[Invocation], + expected_invocations: Optional[list[Invocation]], + conversation_scenario: Optional[ConversationScenario] = None, +) -> EvaluationResult: + """Fraction of expected tool calls that were actually made. + + Score = |expected_tools ∩ actual_tools| / |expected_tools|. + Uses set-based matching (any order). Passes at >= 0.5. + + Same semantics as SkillsBench tool_usage_score for consistency. + """ + if not expected_invocations: + return EvaluationResult( + overall_score=1.0, + overall_eval_status=EvalStatus.PASSED, + ) + + expected_names = set(_get_tool_names(expected_invocations)) + actual_names = set(_get_tool_names(actual_invocations)) + + if not expected_names: + score = 1.0 + else: + matched = expected_names & actual_names + score = len(matched) / len(expected_names) + + status = EvalStatus.PASSED if score >= 0.5 else EvalStatus.FAILED + + return EvaluationResult( + overall_score=score, + overall_eval_status=status, + per_invocation_results=_make_per_invocation( + actual_invocations, + expected_invocations, + score, + status, + ), + ) + + +def output_correctness_score( + eval_metric: EvalMetric, + actual_invocations: list[Invocation], + expected_invocations: Optional[list[Invocation]], + conversation_scenario: Optional[ConversationScenario] = None, +) -> EvaluationResult: + """Binary pass/fail: 1.0 if response contains all reference lines. + + Each non-empty line in the expected final_response is a reference + phrase. The actual response must contain every phrase as a + case-insensitive substring. + + Same semantics as SkillsBench skillsbench_binary_score for + consistency. + """ + if not expected_invocations or not actual_invocations: + return EvaluationResult( + overall_score=0.0, + overall_eval_status=EvalStatus.NOT_EVALUATED, + ) + + # Get the last actual response text. + actual_text = "" + for inv in reversed(actual_invocations): + if inv.final_response and inv.final_response.parts: + for part in inv.final_response.parts: + if part.text: + actual_text = part.text + break + if actual_text: + break + + # Get the expected response text. + expected_text = "" + for inv in reversed(expected_invocations): + if inv.final_response and inv.final_response.parts: + for part in inv.final_response.parts: + if part.text: + expected_text = part.text + break + if expected_text: + break + + if not expected_text: + return EvaluationResult( + overall_score=0.0, + overall_eval_status=EvalStatus.NOT_EVALUATED, + ) + + reference_lines = [ + line.strip() for line in expected_text.split("\n") if line.strip() + ] + actual_lower = actual_text.lower() + matched = sum(1 for line in reference_lines if line.lower() in actual_lower) + + is_pass = matched == len(reference_lines) and len(reference_lines) > 0 + score = 1.0 if is_pass else 0.0 + status = EvalStatus.PASSED if is_pass else EvalStatus.FAILED + + return EvaluationResult( + overall_score=score, + overall_eval_status=status, + per_invocation_results=_make_per_invocation( + actual_invocations, + expected_invocations, + score, + status, + ), + ) diff --git a/benchmarks/bigquerybench/runner.py b/benchmarks/bigquerybench/runner.py new file mode 100644 index 0000000000..63c77449a1 --- /dev/null +++ b/benchmarks/bigquerybench/runner.py @@ -0,0 +1,365 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Standalone BigQueryBench runner that produces leaderboard scores. + +Usage: + python -m benchmarks.bigquerybench.runner + python -m benchmarks.bigquerybench.runner --num-runs 3 + python -m benchmarks.bigquerybench.runner --filter sql_shakespeare + python -m benchmarks.bigquerybench.runner --dry-run + +Environment variables: + GOOGLE_CLOUD_PROJECT — GCP project for BigQuery API calls + GOOGLE_API_KEY — API key for Google AI Studio + GOOGLE_GENAI_USE_VERTEXAI — Set to 1 for Vertex AI backend + BQ_EVAL_WRITE_MODE — Override write mode (blocked/protected) + +This script: +1. Loads the BigQueryBench agent and eval set +2. Runs each case through the ADK Runner +3. Applies 3 custom metrics (schema_discovery, tool_usage, output) +4. Outputs per-case results and a leaderboard-format summary +5. Exits with code 0 if all pass, 1 if any fail +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import logging +import pathlib +import sys +import time +from typing import Optional +import uuid + +from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService +from google.adk.evaluation.eval_case import Invocation +from google.adk.evaluation.eval_metrics import EvalMetric +from google.adk.evaluation.eval_set import EvalSet +from google.adk.evaluation.evaluation_generator import EvaluationGenerator +from google.adk.memory.in_memory_memory_service import InMemoryMemoryService +from google.adk.runners import Runner +from google.adk.sessions.in_memory_session_service import InMemorySessionService +from google.adk.utils.context_utils import Aclosing + +from .metrics import output_correctness_score +from .metrics import schema_discovery_score +from .metrics import tool_usage_score + +logger = logging.getLogger(__name__) + +_BENCH_DIR = pathlib.Path(__file__).parent +_DEFAULT_EVAL_SET = _BENCH_DIR / "eval_sets" / "bigquerybench_eval.json" + + +def load_eval_set(path: pathlib.Path) -> EvalSet: + """Load an EvalSet from a JSON file.""" + with open(path) as f: + data = json.load(f) + return EvalSet.model_validate(data) + + +def print_header(): + print() + print("=" * 65) + print(" BigQueryBench Evaluation — ADK BigQueryToolset") + print("=" * 65) + print() + + +def print_results_table(results: dict[str, dict[str, float]]): + """Print per-case results as a formatted table.""" + print( + f"{'Eval Case':<40} {'Schema':>8} {'Tools':>7}" + f" {'Output':>8} {'Result':>8}" + ) + print("-" * 75) + for case_id, scores in results.items(): + short_id = case_id[:39] + schema = scores.get("schema_discovery", 0.0) + tools = scores.get("tool_usage", 0.0) + output = scores.get("output_correctness", 0.0) + mark = "PASS" if output >= 1.0 else "FAIL" + print( + f"{short_id:<40} {schema:>7.1f} {tools:>7.2f}" + f" {output:>7.1f} {mark:>5}" + ) + print("-" * 75) + + +def print_leaderboard_summary( + results: dict[str, dict[str, float]], + num_cases: int, + elapsed: float, +): + """Print a leaderboard-format summary.""" + passed = sum( + 1 for s in results.values() if s.get("output_correctness", 0.0) >= 1.0 + ) + avg_schema = sum( + s.get("schema_discovery", 0.0) for s in results.values() + ) / max(len(results), 1) + avg_tools = sum(s.get("tool_usage", 0.0) for s in results.values()) / max( + len(results), 1 + ) + pct = (passed / max(num_cases, 1)) * 100 + + print() + print("=" * 65) + print(" Leaderboard Summary") + print("=" * 65) + print(f" Framework: ADK BigQueryToolset") + print(f" Cases: {passed}/{num_cases} ({pct:.1f}%)") + print(f" Avg Schema Disc.: {avg_schema:.2f}") + print(f" Avg Tool Usage: {avg_tools:.2f}") + print(f" Elapsed: {elapsed:.1f}s") + print("=" * 65) + + +async def run_single_eval_case( + root_agent, + eval_case, +) -> list[Invocation]: + """Run a single eval case through the Runner.""" + session_service = InMemorySessionService() + artifact_service = InMemoryArtifactService() + memory_service = InMemoryMemoryService() + + app_name = "bigquerybench_eval" + user_id = "eval_user" + session_id = str(uuid.uuid4()) + + await session_service.create_session( + app_name=app_name, + user_id=user_id, + state={}, + session_id=session_id, + ) + + async with Runner( + app_name=app_name, + agent=root_agent, + artifact_service=artifact_service, + session_service=session_service, + memory_service=memory_service, + ) as runner: + events = [] + for invocation in eval_case.conversation: + user_content = invocation.user_content + + async with Aclosing( + runner.run_async( + user_id=user_id, + session_id=session_id, + new_message=user_content, + ) + ) as agen: + invocation_id = None + async for event in agen: + if not invocation_id: + invocation_id = event.invocation_id + from google.adk.events.event import Event + + events.append( + Event( + content=user_content, + author="user", + invocation_id=invocation_id, + ) + ) + events.append(event) + + return EvaluationGenerator.convert_events_to_eval_invocations(events) + + +def score_invocations( + actual_invocations: list[Invocation], + expected_invocations: Optional[list[Invocation]], +) -> dict[str, float]: + """Apply all 3 metrics and return scores.""" + metric = EvalMetric(metric_name="bigquerybench") + scores = {} + + result = schema_discovery_score( + metric, + actual_invocations, + expected_invocations, + ) + scores["schema_discovery"] = result.overall_score or 0.0 + + result = tool_usage_score( + metric, + actual_invocations, + expected_invocations, + ) + scores["tool_usage"] = result.overall_score or 0.0 + + result = output_correctness_score( + metric, + actual_invocations, + expected_invocations, + ) + scores["output_correctness"] = result.overall_score or 0.0 + + return scores + + +async def run_evaluation( + eval_set_path: Optional[pathlib.Path] = None, + num_runs: int = 1, + filter_str: Optional[str] = None, +) -> dict[str, dict[str, float]]: + """Run the BigQueryBench evaluation.""" + path = eval_set_path or _DEFAULT_EVAL_SET + eval_set = load_eval_set(path) + + # Import agent (triggers BigQuery toolset setup). + from .agent import root_agent + + cases = eval_set.eval_cases + if filter_str: + cases = [c for c in cases if filter_str in c.eval_id] + if not cases: + print(f" No eval cases matched filter: {filter_str!r}") + return {} + + results: dict[str, dict[str, float]] = {} + total = len(cases) + + for idx, eval_case in enumerate(cases, 1): + eval_id = eval_case.eval_id + print(f"\n[{idx}/{total}] Running: {eval_id}") + + run_scores: list[dict[str, float]] = [] + for run in range(num_runs): + if num_runs > 1: + print(f" Run {run + 1}/{num_runs}...") + + try: + actual = await run_single_eval_case(root_agent, eval_case) + + # Print response preview. + for inv in actual: + if inv.final_response and inv.final_response.parts: + for part in inv.final_response.parts: + if part.text: + preview = part.text[:200].replace("\n", " ") + print(f" Response: {preview}...") + break + + expected = eval_case.conversation + scores = score_invocations(actual, expected) + run_scores.append(scores) + + schema = scores["schema_discovery"] + tools = scores["tool_usage"] + output = scores["output_correctness"] + mark = "PASS" if output >= 1.0 else "FAIL" + print(f" Scores: schema={schema:.1f} tools={tools:.2f} output={mark}") + + except Exception as e: + logger.error("Error running %s: %s", eval_id, e) + print(f" ERROR: {e}") + run_scores.append({ + "schema_discovery": 0.0, + "tool_usage": 0.0, + "output_correctness": 0.0, + }) + + # Average scores across runs. + avg: dict[str, float] = {} + for key in ["schema_discovery", "tool_usage", "output_correctness"]: + values = [s[key] for s in run_scores] + avg[key] = sum(values) / len(values) + results[eval_id] = avg + + return results + + +def main(): + parser = argparse.ArgumentParser( + description="BigQueryBench evaluation runner for ADK", + ) + parser.add_argument( + "--eval-set", + type=pathlib.Path, + default=None, + help="Path to eval set JSON (default: built-in)", + ) + parser.add_argument( + "--num-runs", + type=int, + default=1, + help="Number of runs per case (default: 1)", + ) + parser.add_argument( + "--filter", + type=str, + default=None, + help="Substring filter for eval_id (e.g., 'sql_shakespeare')", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Validate eval set JSON without running LLM inference", + ) + args = parser.parse_args() + + logging.basicConfig(level=logging.WARNING) + + if args.dry_run: + path = args.eval_set or _DEFAULT_EVAL_SET + eval_set = load_eval_set(path) + print(f"Eval set: {eval_set.name}") + print(f"Cases: {len(eval_set.eval_cases)}") + for case in eval_set.eval_cases: + print(f" - {case.eval_id}") + print("\nDry run: eval set JSON is valid.") + sys.exit(0) + + print_header() + start = time.time() + + results = asyncio.run( + run_evaluation( + eval_set_path=args.eval_set, + num_runs=args.num_runs, + filter_str=args.filter, + ) + ) + + elapsed = time.time() - start + eval_path = args.eval_set or _DEFAULT_EVAL_SET + eval_set = load_eval_set(eval_path) + + num_cases = len(eval_set.eval_cases) + if args.filter: + num_cases = len(results) + + print() + print_results_table(results) + print_leaderboard_summary(results, num_cases, elapsed) + + # Exit code: 0 if all pass, 1 if any fail. + all_pass = all( + s.get("output_correctness", 0.0) >= 1.0 for s in results.values() + ) + sys.exit(0 if all_pass else 1) + + +if __name__ == "__main__": + main() From b164b8709376abba57e9934bdd5be4c4a23cf777 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 02:20:41 -0800 Subject: [PATCH 43/53] docs(bigquerybench): add complete walkthrough for new BigQuery skill evaluation Adds a detailed "Complete Walkthrough" section to the BigQueryBench README with a concrete cluster_data (K-Means) tool example covering: - Step 1: Register tool in BigQueryToolset (auto-discovered) - Step 2: Evaluate whether existing metrics suffice - Step 2b: Adding a custom metric (clustering_quality_score example) - Step 3: Pick a public dataset (ml_datasets.penguins) - Step 4: Write the eval case JSON with structural reference lines - Step 5: Configure write mode for ML tools - Step 6: Validate with dry-run and multi-run - Step 7: Commit Includes a summary table showing which files to touch per scenario (new eval case only, new tool with existing metrics, new tool needing custom metric, agent config change). Co-Authored-By: Claude Opus 4.6 --- benchmarks/bigquerybench/README.md | 304 +++++++++++++++++++++++++++++ 1 file changed, 304 insertions(+) diff --git a/benchmarks/bigquerybench/README.md b/benchmarks/bigquerybench/README.md index dc80173e17..8c1cb02bd0 100644 --- a/benchmarks/bigquerybench/README.md +++ b/benchmarks/bigquerybench/README.md @@ -375,6 +375,310 @@ use this template: - Set `tool_config.write_mode = WriteMode.PROTECTED` in the agent for AI/ML eval cases that need to create temp models. +## Complete Walkthrough: Adding a New BigQuery Skill + +This section walks through every step of updating the evaluation +pipeline when a **new BigQuery tool** is added to +`BigQueryToolset`. We use a concrete example: a hypothetical +`cluster_data` tool that performs K-Means clustering via BQML. + +### Context: What Is the New Tool? + +Suppose a developer adds this tool to `src/google/adk/tools/bigquery/`: + +```python +def cluster_data( + project_id: str, + input_data: str, # Table ID or SQL query + feature_cols: list[str], # Columns to cluster on + num_clusters: int = 3, # K in K-Means + *, + credentials: Credentials, + settings: BigQueryToolConfig, + tool_context: ToolContext, +) -> dict: + """Cluster rows using BigQuery ML K-Means. + + Creates a TEMP MODEL and returns cluster assignments + with centroid distances. + """ +``` + +The tool generates SQL like: + +```sql +CREATE TEMP MODEL cluster_model_ + OPTIONS (MODEL_TYPE='KMEANS', NUM_CLUSTERS=3) + AS SELECT feature1, feature2 FROM `project.dataset.table`; + +SELECT * FROM ML.PREDICT(MODEL cluster_model_, + (SELECT feature1, feature2 FROM `project.dataset.table`)); +``` + +Output columns: `centroid_id`, `nearest_centroids_distance`, +plus the original feature columns. + +### Step 1: Register the Tool in BigQueryToolset + +The developer registers the tool in `bigquery_toolset.py`. Once +registered, it's automatically available to any agent using +`BigQueryToolset`. **No changes to the eval agent** (`agent.py`) +are needed — the toolset dynamically exposes all registered tools. + +Verify the tool is visible: + +```python +from benchmarks.bigquerybench.agent import bigquery_toolset +tools = bigquery_toolset.get_tools() +assert any(t.name == "cluster_data" for t in tools) +``` + +### Step 2: Decide If Existing Metrics Are Sufficient + +Check each metric against the new tool's behavior: + +| Metric | Does It Work? | Action Needed? | +|--------|--------------|----------------| +| `schema_discovery_score` | Yes — the agent should still explore schema before clustering | No change | +| `tool_usage_score` | Yes — set-based matching works for any tool name | No change | +| `output_correctness_score` | **Partially** — ML outputs vary across runs, so exact numeric matching will be flaky | Use structural reference lines (column names, cluster count) instead of exact values | + +**When you DO need a new metric:** If the new tool has a unique +correctness criterion that can't be captured by substring matching +(e.g., "the SQL must be syntactically valid", "the forecast horizon +must match the request"), add a new metric function to `metrics.py`. +See [Step 2b: Adding a Custom Metric](#step-2b-adding-a-custom-metric) +below. + +### Step 3: Pick a Public Dataset + +Choose a dataset from `bigquery-public-data` with numeric columns +suitable for clustering. For this example, we'll use the **penguins** +dataset (`ml_datasets.penguins`) which has well-known numeric features. + +Verify it exists: + +```sql +SELECT column_name, data_type +FROM `bigquery-public-data.ml_datasets.INFORMATION_SCHEMA.COLUMNS` +WHERE table_name = 'penguins'; +``` + +Expected columns: `species`, `island`, `culmen_length_mm`, +`culmen_depth_mm`, `flipper_length_mm`, `body_mass_g`, `sex`. + +### Step 4: Write the Eval Case + +Add to `eval_sets/bigquerybench_eval.json`: + +```json +{ + "eval_id": "ml_cluster_penguins", + "conversation": [ + { + "invocation_id": "inv-cluster-01", + "user_content": { + "parts": [{"text": "Cluster the penguins in bigquery-public-data.ml_datasets.penguins into 3 groups based on their physical measurements (culmen_length_mm, culmen_depth_mm, flipper_length_mm, body_mass_g). Show the cluster assignments."}], + "role": "user" + }, + "final_response": { + "parts": [{"text": "centroid_id\nculmen_length_mm\nculmen_depth_mm\nflipper_length_mm\nbody_mass_g"}], + "role": "model" + }, + "intermediate_data": { + "tool_uses": [ + {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "ml_datasets", "table_id": "penguins"}}, + {"name": "cluster_data", "args": {"project_id": "your-project", "input_data": "SELECT culmen_length_mm, culmen_depth_mm, flipper_length_mm, body_mass_g FROM `bigquery-public-data.ml_datasets.penguins` WHERE body_mass_g IS NOT NULL", "feature_cols": ["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm", "body_mass_g"], "num_clusters": 3}} + ], + "tool_responses": [], + "intermediate_responses": [] + }, + "creation_timestamp": 0.0 + } + ], + "creation_timestamp": 0.0 +} +``` + +**Reference line design choices:** + +- `centroid_id` — structural: confirms clustering output format +- `culmen_length_mm`, `culmen_depth_mm`, etc. — structural: confirms + feature columns are returned in output +- We do NOT include exact centroid values or row counts because ML + model outputs vary across runs + +### Step 5: Update the Agent Write Mode + +The `cluster_data` tool creates a `TEMP MODEL`, which requires at +least `PROTECTED` write mode. Update the eval run command: + +```bash +BQ_EVAL_WRITE_MODE=protected \ + python -m benchmarks.bigquerybench.runner --filter ml_cluster_penguins +``` + +Or for CI, set it in the environment configuration. + +### Step 6: Validate the Eval Case + +```bash +# Dry-run: check JSON is valid +python -m benchmarks.bigquerybench.runner --dry-run + +# Single-case run +BQ_EVAL_WRITE_MODE=protected \ + python -m benchmarks.bigquerybench.runner --filter ml_cluster_penguins + +# Multi-run for variance check (ML outputs may vary) +BQ_EVAL_WRITE_MODE=protected \ + python -m benchmarks.bigquerybench.runner --filter ml_cluster --num-runs 3 +``` + +Expected output: + +``` +================================================================= + BigQueryBench Evaluation — ADK BigQueryToolset +================================================================= + +[1/1] Running: ml_cluster_penguins + Response: Here are the cluster assignments for the penguins... + Scores: schema=1.0 tools=1.00 output=PASS + +Eval Case Schema Tools Output Result +--------------------------------------------------------------------------- +ml_cluster_penguins 1.0 1.00 1.0 PASS +--------------------------------------------------------------------------- + +================================================================= + Leaderboard Summary +================================================================= + Framework: ADK BigQueryToolset + Cases: 1/1 (100.0%) + Avg Schema Disc.: 1.00 + Avg Tool Usage: 1.00 + Elapsed: 12.3s +================================================================= +``` + +### Step 7: Commit + +```bash +git add benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json +git commit -m "eval(bigquerybench): add clustering eval case for cluster_data tool" +``` + +**Files changed:** Only `bigquerybench_eval.json`. No code changes +needed unless a new metric was added (Step 2b). + +--- + +### Step 2b: Adding a Custom Metric (When Needed) + +If existing metrics are insufficient for your new tool, add a metric +function to `metrics.py`. Here's a concrete example: a +`clustering_quality_score` that validates the agent used the correct +number of clusters. + +**1. Write the metric function in `metrics.py`:** + +```python +def clustering_quality_score( + eval_metric: EvalMetric, + actual_invocations: list[Invocation], + expected_invocations: Optional[list[Invocation]], + conversation_scenario: Optional[ConversationScenario] = None, +) -> EvaluationResult: + """Score 1.0 if cluster_data was called with correct num_clusters. + + Checks that the agent called cluster_data and that the + num_clusters argument matches the expected value. + """ + if not expected_invocations: + return EvaluationResult( + overall_score=1.0, + overall_eval_status=EvalStatus.PASSED, + ) + + # Extract expected num_clusters from expected tool calls. + expected_k = None + for inv in expected_invocations: + for tc in get_all_tool_calls(inv.intermediate_data): + if tc.name == "cluster_data" and tc.args: + expected_k = tc.args.get("num_clusters") + break + + # Extract actual num_clusters from actual tool calls. + actual_k = None + for inv in actual_invocations: + for tc in get_all_tool_calls(inv.intermediate_data): + if tc.name == "cluster_data" and tc.args: + actual_k = tc.args.get("num_clusters") + break + + if expected_k is None: + score = 1.0 # No clustering expected + elif actual_k is None: + score = 0.0 # Clustering expected but not called + else: + score = 1.0 if actual_k == expected_k else 0.0 + + status = ( + EvalStatus.PASSED if score >= 1.0 + else EvalStatus.FAILED + ) + return EvaluationResult( + overall_score=score, + overall_eval_status=status, + per_invocation_results=_make_per_invocation( + actual_invocations, expected_invocations, + score, status, + ), + ) +``` + +**2. Wire it into the runner (`runner.py`):** + +```python +from .metrics import clustering_quality_score + +def score_invocations(actual, expected): + # ... existing metrics ... + + result = clustering_quality_score( + metric, actual_invocations, expected_invocations, + ) + scores["clustering_quality"] = result.overall_score or 0.0 + + return scores +``` + +**3. Update the results table to show the new column.** + +**4. Commit all changed files:** + +```bash +git add benchmarks/bigquerybench/metrics.py \ + benchmarks/bigquerybench/runner.py \ + benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json +git commit -m "eval(bigquerybench): add cluster_data eval case and clustering_quality metric" +``` + +--- + +### Summary: Files to Touch per Scenario + +| Scenario | `eval.json` | `metrics.py` | `runner.py` | `agent.py` | +|----------|:-----------:|:------------:|:-----------:|:----------:| +| New eval case for existing tool | **Yes** | No | No | No | +| New tool, existing metrics sufficient | **Yes** | No | No | No | +| New tool, needs custom metric | **Yes** | **Yes** | **Yes** | No | +| New tool, needs agent config change | **Yes** | Maybe | Maybe | **Yes** | + +The toolset auto-discovers new tools, so `agent.py` only changes if +the agent's instruction prompt or write mode needs updating. + ## Metrics Reference | Metric | Function Path | Threshold | Pass Condition | From 70a2f3fd111f62a02a229b0badb3d2adf1f17358 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 02:25:16 -0800 Subject: [PATCH 44/53] refactor(bigquerybench): simplify to trace-based evaluation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace 3-metric response-text-matching approach with 2 trace-only metrics that are deterministic and easy to maintain: - tool_invocation_score: all expected tool names appear in trace - tool_args_score: all expected (tool, project/dataset/table) pairs appear in trace Key changes: - metrics.py: Replace schema_discovery + tool_usage + output_correctness with tool_invocation + tool_args - eval JSON: Remove final_response fields — only tool_uses matter - runner.py: Print tool-call trace for debugging, simplified output - README.md: Rewritten for trace-based approach with concrete cluster_data walkthrough showing JSON-only eval case addition Adding a new eval case now requires zero code changes — just add a JSON object specifying the expected tool calls and their key args. Co-Authored-By: Claude Opus 4.6 --- benchmarks/bigquerybench/README.md | 696 ++++-------------- .../eval_sets/bigquerybench_eval.json | 38 +- benchmarks/bigquerybench/metrics.py | 175 ++--- benchmarks/bigquerybench/runner.py | 195 ++--- 4 files changed, 308 insertions(+), 796 deletions(-) diff --git a/benchmarks/bigquerybench/README.md b/benchmarks/bigquerybench/README.md index 8c1cb02bd0..a866db7a2e 100644 --- a/benchmarks/bigquerybench/README.md +++ b/benchmarks/bigquerybench/README.md @@ -1,146 +1,79 @@ -# BigQueryBench: End-to-End Evaluation for BigQuery Skills +# BigQueryBench: Trace-Based Evaluation for BigQuery Skills ## Overview -BigQueryBench is a reusable evaluation pipeline for any skill or agent -built with ADK's `BigQueryToolset`. It mirrors the SkillsBench -architecture (`benchmarks/skillsbench/`) but targets BigQuery-specific -tool chains: schema exploration, SQL generation, AI/ML operations -(forecast, anomaly detection, contribution analysis), and multi-step -analytical workflows. - -**Design goals:** - -- **Reusable:** Add a new BigQuery skill by dropping one eval case - JSON and one optional reference SQL file — no code changes needed. -- **Reproducible:** All eval cases use BigQuery public datasets - (`bigquery-public-data`) so any GCP project with BigQuery API - enabled can run the suite. -- **Layered metrics:** Three dimensions scored independently — - schema discovery, tool-call coverage, output correctness. -- **CI-friendly:** Single `python -m benchmarks.bigquerybench.runner` - invocation with JSON results and exit code. +BigQueryBench verifies that an agent built with ADK's `BigQueryToolset` +**calls the correct tools with the correct arguments**. It inspects +the tool-call trace only — no response text matching. + +For each eval case, the pipeline checks two things: + +1. **Tool invocation** — Did the agent call the right BigQuery + functions? (e.g., `get_table_info` then `execute_sql`) +2. **Tool arguments** — Did those calls point at the right data? + (e.g., `project_id="bigquery-public-data"`, + `dataset_id="usa_names"`, `table_id="usa_1910_current"`) + +This makes evaluation deterministic, easy to maintain, and immune to +LLM response wording variance. ## Quick Start ```bash -# Prerequisites -# 1. GCP project with BigQuery API enabled -# 2. Application Default Credentials configured: -# gcloud auth application-default login -# 3. ADK installed with BigQuery extras: -# uv sync --all-extras - -# Set environment +# Prerequisites: GCP project with BigQuery API + ADC configured export GOOGLE_CLOUD_PROJECT=your-project-id -export GOOGLE_GENAI_USE_VERTEXAI=1 # or use GOOGLE_API_KEY +export GOOGLE_GENAI_USE_VERTEXAI=1 # Run all eval cases python -m benchmarks.bigquerybench.runner -# Run specific eval case(s) -python -m benchmarks.bigquerybench.runner --filter sql_public_dataset - -# Run with multiple attempts for variance measurement -python -m benchmarks.bigquerybench.runner --num-runs 3 +# Run one case +python -m benchmarks.bigquerybench.runner --filter schema_list_tables -# Dry-run mode (validates eval set JSON, no LLM calls) +# Dry-run (validate JSON only, no LLM calls) python -m benchmarks.bigquerybench.runner --dry-run ``` -## Architecture +## How It Works ``` -benchmarks/bigquerybench/ -├── README.md # This file -├── __init__.py -├── agent.py # Root agent with BigQueryToolset -├── runner.py # Standalone evaluation runner -├── metrics.py # BigQuery-specific custom metrics -└── eval_sets/ - └── bigquerybench_eval.json # Eval cases (public datasets) -``` - -### Relationship to SkillsBench - -| Aspect | SkillsBench | BigQueryBench | -|--------|-------------|---------------| -| Toolset | `SkillToolset` (4 tools) | `BigQueryToolset` (10 tools) | -| Discovery tools | `list_skills`, `load_skill` | `list_dataset_ids`, `list_table_ids`, `get_dataset_info`, `get_table_info` | -| Execution tools | `run_skill_script` | `execute_sql`, `forecast`, `detect_anomalies`, `analyze_contribution` | -| Data source | Bundled reference files | BigQuery public datasets | -| Auth | None (local files) | GCP credentials (ADC / service account / OAuth) | -| Metrics | discovery, tool_usage, binary | schema_discovery, tool_usage, output_correctness | - -## Evaluation Pipeline - -### Stage 1: Agent Setup - -The agent under test is defined in `agent.py`. It uses -`BigQueryToolset` with read-only defaults: - -```python -from google.adk.agents.llm_agent import LlmAgent -from google.adk.tools.bigquery.bigquery_credentials import ( - BigQueryCredentialsConfig, -) -from google.adk.tools.bigquery.bigquery_toolset import BigQueryToolset -from google.adk.tools.bigquery.config import BigQueryToolConfig -from google.adk.tools.bigquery.config import WriteMode -import google.auth - -credentials, _ = google.auth.default() -credentials_config = BigQueryCredentialsConfig(credentials=credentials) - -tool_config = BigQueryToolConfig( - write_mode=WriteMode.BLOCKED, # Read-only for eval safety - max_query_result_rows=50, -) - -bigquery_toolset = BigQueryToolset( - credentials_config=credentials_config, - bigquery_tool_config=tool_config, -) - -root_agent = LlmAgent( - model="gemini-2.5-flash", - name="bigquerybench_agent", - description="Agent for BigQuery data exploration and analysis.", - instruction="""\ - You are a data analyst with access to BigQuery tools. - Use them to explore schemas, run SQL queries, and answer - the user's questions about data. Always explore the schema - (list datasets, list tables, get table info) before writing - SQL. Show query results clearly. - """, - tools=[bigquery_toolset], -) +eval_sets/bigquerybench_eval.json + ↓ (user query + expected tool_uses) +runner.py + ↓ runs agent via ADK Runner + ↓ collects event trace → Invocations +metrics.py + ├── tool_invocation_score: expected tool names ⊆ actual tool names? + └── tool_args_score: expected (tool, project/dataset/table) ⊆ actual? + ↓ +PASS if both scores = 1.0 ``` -### Stage 2: Eval Set Definition +## Eval Case Format -Each eval case is a JSON object in `eval_sets/bigquerybench_eval.json` -following the ADK `EvalSet` schema. The key fields are: +Each eval case specifies a user query and the expected tool calls. +No `final_response` is needed — only the trace matters. ```json { - "eval_id": "unique_case_id", + "eval_id": "schema_get_table_info", "conversation": [ { "invocation_id": "inv-01", "user_content": { - "parts": [{"text": "User's question about BigQuery data"}], + "parts": [{"text": "What columns does usa_1910_current have?"}], "role": "user" }, - "final_response": { - "parts": [{"text": "Expected key phrases in the answer"}], - "role": "model" - }, "intermediate_data": { "tool_uses": [ - {"name": "list_dataset_ids", "args": {"project_id": "bigquery-public-data"}}, - {"name": "get_table_info", "args": {"project_id": "...", "dataset_id": "...", "table_id": "..."}}, - {"name": "execute_sql", "args": {"project_id": "...", "query": "SELECT ..."}} + { + "name": "get_table_info", + "args": { + "project_id": "bigquery-public-data", + "dataset_id": "usa_names", + "table_id": "usa_1910_current" + } + } ], "tool_responses": [], "intermediate_responses": [] @@ -152,148 +85,57 @@ following the ADK `EvalSet` schema. The key fields are: } ``` -**Key conventions:** - -- `final_response.parts[0].text` contains **reference lines** — key - phrases that must appear in the agent's response for a pass. One - phrase per line. Case-insensitive substring match. -- `intermediate_data.tool_uses` lists the **expected tool call - sequence**. The `tool_usage` metric checks set coverage (not - strict ordering). The `schema_discovery` metric checks that at - least one schema-exploration tool was called. -- All eval cases use `bigquery-public-data` project datasets so - results are deterministic and reproducible. - -### Stage 3: Metrics - -Three custom metrics in `metrics.py`, following the ADK custom metric -function signature: - -#### 3a. `schema_discovery_score` - -Checks whether the agent explored the schema before querying. Scores -1.0 if any of these tools were called: `list_dataset_ids`, -`list_table_ids`, `get_dataset_info`, `get_table_info`. Scores 0.0 -otherwise. +**What gets checked:** -**Rationale:** Agents that skip schema exploration and guess table -names produce fragile SQL that breaks on schema changes. This metric -enforces the "explore before query" pattern. +| Field in `args` | Checked? | Why | +|-----------------|----------|-----| +| `project_id` | Yes | Must point at the right GCP project | +| `dataset_id` | Yes | Must load the right dataset | +| `table_id` | Yes | Must load the right table | +| `query` | **No** | Exact SQL varies — agent may write equivalent SQL differently | +| Other args | **No** | Tool-specific args (e.g., `horizon`, `num_clusters`) are not checked by default | -#### 3b. `tool_usage_score` +## Metrics -Fraction of expected tool calls actually made: -`|expected_tools ∩ actual_tools| / |expected_tools|`. +| Metric | What It Checks | Pass Condition | +|--------|---------------|----------------| +| `tool_invocation_score` | All expected tool names appear in the trace | Score = 1.0 | +| `tool_args_score` | All expected `(tool, project_id/dataset_id/table_id)` pairs appear in the trace | Score = 1.0 | -Uses set-based matching (any order). Passes at threshold >= 0.5. +A case **passes** when both scores are 1.0. -Same semantics as SkillsBench `tool_usage_score`, reused for -consistency. +## Included Eval Cases -#### 3c. `output_correctness_score` - -Binary pass/fail: 1.0 if the agent's final response contains all -expected reference lines (case-insensitive substring match). 0.0 -otherwise. - -Same semantics as SkillsBench `skillsbench_binary_score`, reused for -consistency. - -### Stage 4: Runner Execution - -``` -runner.py - ↓ -Load agent from benchmarks/bigquerybench/agent.py - ↓ -Load eval set from eval_sets/bigquerybench_eval.json - ↓ -For each eval case: - ↓ - Run agent.run_async(user_query) via ADK Runner - ↓ (generates events) - Convert events → Invocation (with intermediate_data.tool_uses) - ↓ - Apply schema_discovery_score - Apply tool_usage_score - Apply output_correctness_score - ↓ - Record per-case scores - ↓ -Aggregate scores → leaderboard summary - ↓ -Print table + exit code (0 = all pass, 1 = any fail) -``` +| eval_id | User Query | Expected Trace | +|---------|-----------|----------------| +| `schema_list_datasets` | "What datasets are in bigquery-public-data?" | `list_dataset_ids(project_id=bigquery-public-data)` | +| `schema_list_tables` | "What tables in usa_names?" | `list_table_ids(project_id=.., dataset_id=usa_names)` | +| `schema_get_table_info` | "Columns of usa_1910_current?" | `get_table_info(project_id=.., dataset_id=usa_names, table_id=usa_1910_current)` | +| `sql_shakespeare_unique_words` | "Top 3 works by unique words?" | `get_table_info(.., shakespeare)` → `execute_sql(..)` | +| `sql_usa_names_top_2020` | "Top 5 baby names in 2020?" | `get_table_info(.., usa_1910_current)` → `execute_sql(..)` | +| `sql_names_by_decade` | "Distinct names per decade 1950-2000?" | `get_table_info(.., usa_1910_current)` → `execute_sql(..)` | +| `multi_step_explore_and_query` | "Explore bikeshare, top 5 stations?" | `list_table_ids(..)` → `get_table_info(.., bikeshare_trips)` → `execute_sql(..)` | -## Eval Case Catalog +## Adding a New Eval Case -The following eval cases are included. They are organized by -complexity tier to test progressively harder agent capabilities. +### For an existing tool (e.g., new `execute_sql` scenario) -### Tier 1: Schema Exploration - -These test the agent's ability to navigate BigQuery metadata. - -| eval_id | Dataset | User Query | Expected Tools | -|---------|---------|-----------|----------------| -| `schema_list_datasets` | `bigquery-public-data` | "What datasets are available in bigquery-public-data?" | `list_dataset_ids` | -| `schema_list_tables` | `usa_names` | "What tables exist in the usa_names dataset?" | `list_dataset_ids` → `list_table_ids` | -| `schema_get_table_info` | `usa_names.usa_1910_current` | "What columns and types does the usa_1910_current table have?" | `list_table_ids` → `get_table_info` | - -### Tier 2: SQL Generation & Execution - -These test SQL generation against public data with known answers. - -| eval_id | Dataset | User Query | Expected Tools | Reference Output | -|---------|---------|-----------|----------------|-----------------| -| `sql_top_names` | `usa_names` | "What are the top 5 most popular baby names in 2020?" | `get_table_info` → `execute_sql` | Top names by count | -| `sql_aggregation` | `usa_names` | "How many distinct names were registered each decade from 1950 to 2000?" | `get_table_info` → `execute_sql` | Decade counts | -| `sql_public_dataset` | `samples.shakespeare` | "Which Shakespeare work has the most unique words?" | `get_table_info` → `execute_sql` | Work name + count | - -### Tier 3: Multi-Step Analysis - -These test the agent's ability to chain multiple tools. - -| eval_id | Dataset | User Query | Expected Tools | Reference Output | -|---------|---------|-----------|----------------|-----------------| -| `multi_step_explore_and_query` | `austin_bikeshare` | "Explore the Austin bikeshare dataset and tell me the top 5 busiest stations by trip count." | `list_table_ids` → `get_table_info` → `execute_sql` | Station names + counts | - -## Adding a New BigQuery Eval Case - -To add a new eval case (e.g., for a new BigQuery AI operator skill): - -### Step 1: Identify the public dataset - -Pick a dataset from `bigquery-public-data` that exercises the skill. -Verify it exists: - -```sql -SELECT * FROM `bigquery-public-data.DATASET.INFORMATION_SCHEMA.TABLES` -LIMIT 5; -``` - -### Step 2: Write the eval case JSON - -Add a new object to the `eval_cases` array in -`eval_sets/bigquerybench_eval.json`: +Only add a JSON object to `bigquerybench_eval.json`. No code changes. ```json { - "eval_id": "your_unique_eval_id", + "eval_id": "sql_weather_hottest_day", "conversation": [ { - "invocation_id": "inv-your-id-01", + "invocation_id": "inv-weather-01", "user_content": { - "parts": [{"text": "Your user query here"}], + "parts": [{"text": "What was the hottest day recorded in the NOAA GSOD 2023 data?"}], "role": "user" }, - "final_response": { - "parts": [{"text": "reference line 1\nreference line 2"}], - "role": "model" - }, "intermediate_data": { "tool_uses": [ - {"name": "tool_name", "args": {"arg1": "val1"}} + {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "noaa_gsod", "table_id": "gsod2023"}}, + {"name": "execute_sql", "args": {"project_id": "bigquery-public-data"}} ], "tool_responses": [], "intermediate_responses": [] @@ -305,54 +147,28 @@ Add a new object to the `eval_cases` array in } ``` -### Step 3: Choose reference lines - -Run the expected query manually and pick 3-5 key phrases from the -result that are **stable** (won't change if data is appended). Good -reference lines: - -- Column names or schema facts ("column: name, type: STRING") -- Aggregation results from historical data ("hamlet", "king lear") -- Structural facts ("3 tables", "5 columns") - -Avoid: exact row counts on append-only tables, floating-point values -that may shift with precision. - -### Step 4: Validate - -```bash -# Run just your new case -python -m benchmarks.bigquerybench.runner --filter your_unique_eval_id -``` - -### Step 5: Commit +Validate: `python -m benchmarks.bigquerybench.runner --filter sql_weather` -Add only the modified `bigquerybench_eval.json`. No code changes -needed. +### For a new tool (e.g., `forecast`, `cluster_data`) -## Eval Case Template for AI/ML Tools - -For `forecast`, `detect_anomalies`, and `analyze_contribution` skills, -use this template: +Same steps — just use the new tool name in `tool_uses`. The metrics +check tool names and key args generically, so no metric code changes +are needed. ```json { - "eval_id": "forecast_weather_temperature", + "eval_id": "ml_forecast_temperature", "conversation": [ { "invocation_id": "inv-forecast-01", "user_content": { - "parts": [{"text": "Forecast the next 7 days of average temperature using the NOAA GSOD weather data for station 725300 (Chicago O'Hare) from 2023."}], + "parts": [{"text": "Forecast the next 7 days of temperature from NOAA GSOD 2023 data for station 725300."}], "role": "user" }, - "final_response": { - "parts": [{"text": "forecast_timestamp\nforecast_value\nprediction_interval"}], - "role": "model" - }, "intermediate_data": { "tool_uses": [ {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "noaa_gsod", "table_id": "gsod2023"}}, - {"name": "forecast", "args": {"project_id": "your-project", "history_data": "SELECT date, temp FROM `bigquery-public-data.noaa_gsod.gsod2023` WHERE stn = '725300'", "timestamp_col": "date", "data_col": "temp", "horizon": 7}} + {"name": "forecast", "args": {"project_id": "bigquery-public-data"}} ], "tool_responses": [], "intermediate_responses": [] @@ -364,112 +180,30 @@ use this template: } ``` -**Key considerations for AI/ML eval cases:** - -- `forecast`, `analyze_contribution`, and `detect_anomalies` create - temporary BigQuery ML models. Ensure `write_mode` is at least - `PROTECTED` (anonymous dataset) or `ALLOWED`. -- Reference lines should validate **structural output** (column names - like `forecast_timestamp`, `is_anomaly`) rather than exact numeric - values, since ML model outputs vary across runs. -- Set `tool_config.write_mode = WriteMode.PROTECTED` in the agent - for AI/ML eval cases that need to create temp models. - -## Complete Walkthrough: Adding a New BigQuery Skill - -This section walks through every step of updating the evaluation -pipeline when a **new BigQuery tool** is added to -`BigQueryToolset`. We use a concrete example: a hypothetical -`cluster_data` tool that performs K-Means clustering via BQML. - -### Context: What Is the New Tool? - -Suppose a developer adds this tool to `src/google/adk/tools/bigquery/`: - -```python -def cluster_data( - project_id: str, - input_data: str, # Table ID or SQL query - feature_cols: list[str], # Columns to cluster on - num_clusters: int = 3, # K in K-Means - *, - credentials: Credentials, - settings: BigQueryToolConfig, - tool_context: ToolContext, -) -> dict: - """Cluster rows using BigQuery ML K-Means. - - Creates a TEMP MODEL and returns cluster assignments - with centroid distances. - """ -``` - -The tool generates SQL like: - -```sql -CREATE TEMP MODEL cluster_model_ - OPTIONS (MODEL_TYPE='KMEANS', NUM_CLUSTERS=3) - AS SELECT feature1, feature2 FROM `project.dataset.table`; - -SELECT * FROM ML.PREDICT(MODEL cluster_model_, - (SELECT feature1, feature2 FROM `project.dataset.table`)); -``` - -Output columns: `centroid_id`, `nearest_centroids_distance`, -plus the original feature columns. - -### Step 1: Register the Tool in BigQueryToolset - -The developer registers the tool in `bigquery_toolset.py`. Once -registered, it's automatically available to any agent using -`BigQueryToolset`. **No changes to the eval agent** (`agent.py`) -are needed — the toolset dynamically exposes all registered tools. - -Verify the tool is visible: - -```python -from benchmarks.bigquerybench.agent import bigquery_toolset -tools = bigquery_toolset.get_tools() -assert any(t.name == "cluster_data" for t in tools) +For AI/ML tools that create temp models, set write mode: +```bash +BQ_EVAL_WRITE_MODE=protected python -m benchmarks.bigquerybench.runner --filter ml_forecast ``` -### Step 2: Decide If Existing Metrics Are Sufficient - -Check each metric against the new tool's behavior: - -| Metric | Does It Work? | Action Needed? | -|--------|--------------|----------------| -| `schema_discovery_score` | Yes — the agent should still explore schema before clustering | No change | -| `tool_usage_score` | Yes — set-based matching works for any tool name | No change | -| `output_correctness_score` | **Partially** — ML outputs vary across runs, so exact numeric matching will be flaky | Use structural reference lines (column names, cluster count) instead of exact values | - -**When you DO need a new metric:** If the new tool has a unique -correctness criterion that can't be captured by substring matching -(e.g., "the SQL must be syntactically valid", "the forecast horizon -must match the request"), add a new metric function to `metrics.py`. -See [Step 2b: Adding a Custom Metric](#step-2b-adding-a-custom-metric) -below. - -### Step 3: Pick a Public Dataset - -Choose a dataset from `bigquery-public-data` with numeric columns -suitable for clustering. For this example, we'll use the **penguins** -dataset (`ml_datasets.penguins`) which has well-known numeric features. +## Complete Walkthrough: Adding a New BigQuery Skill -Verify it exists: +This walkthrough uses a concrete example: a hypothetical +`cluster_data` tool (K-Means via BQML) being added to +`BigQueryToolset`. -```sql -SELECT column_name, data_type -FROM `bigquery-public-data.ml_datasets.INFORMATION_SCHEMA.COLUMNS` -WHERE table_name = 'penguins'; -``` +### Step 1: Register the tool -Expected columns: `species`, `island`, `culmen_length_mm`, -`culmen_depth_mm`, `flipper_length_mm`, `body_mass_g`, `sex`. +The developer adds `cluster_data` to `bigquery_toolset.py`. Once +registered, it's automatically available to the eval agent — no +changes to `agent.py` needed. -### Step 4: Write the Eval Case +### Step 2: Write the eval case -Add to `eval_sets/bigquerybench_eval.json`: +What we want to verify: when the user asks "cluster the penguins +data", the agent should: +1. Call `get_table_info` on `ml_datasets.penguins` (load the schema) +2. Call `cluster_data` against `bigquery-public-data` (invoke the + right tool on the right data) ```json { @@ -478,17 +212,13 @@ Add to `eval_sets/bigquerybench_eval.json`: { "invocation_id": "inv-cluster-01", "user_content": { - "parts": [{"text": "Cluster the penguins in bigquery-public-data.ml_datasets.penguins into 3 groups based on their physical measurements (culmen_length_mm, culmen_depth_mm, flipper_length_mm, body_mass_g). Show the cluster assignments."}], + "parts": [{"text": "Cluster the penguins in bigquery-public-data.ml_datasets.penguins into 3 groups based on their physical measurements."}], "role": "user" }, - "final_response": { - "parts": [{"text": "centroid_id\nculmen_length_mm\nculmen_depth_mm\nflipper_length_mm\nbody_mass_g"}], - "role": "model" - }, "intermediate_data": { "tool_uses": [ {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "ml_datasets", "table_id": "penguins"}}, - {"name": "cluster_data", "args": {"project_id": "your-project", "input_data": "SELECT culmen_length_mm, culmen_depth_mm, flipper_length_mm, body_mass_g FROM `bigquery-public-data.ml_datasets.penguins` WHERE body_mass_g IS NOT NULL", "feature_cols": ["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm", "body_mass_g"], "num_clusters": 3}} + {"name": "cluster_data", "args": {"project_id": "bigquery-public-data"}} ], "tool_responses": [], "intermediate_responses": [] @@ -500,208 +230,90 @@ Add to `eval_sets/bigquerybench_eval.json`: } ``` -**Reference line design choices:** - -- `centroid_id` — structural: confirms clustering output format -- `culmen_length_mm`, `culmen_depth_mm`, etc. — structural: confirms - feature columns are returned in output -- We do NOT include exact centroid values or row counts because ML - model outputs vary across runs - -### Step 5: Update the Agent Write Mode +**What gets checked automatically:** +- `tool_invocation_score`: Did the trace contain both + `get_table_info` and `cluster_data`? +- `tool_args_score`: Did `get_table_info` target + `(bigquery-public-data, ml_datasets, penguins)`? Did + `cluster_data` target `bigquery-public-data`? -The `cluster_data` tool creates a `TEMP MODEL`, which requires at -least `PROTECTED` write mode. Update the eval run command: +**What is NOT checked** (intentionally): +- The exact `feature_cols` or `num_clusters` the LLM chose +- The exact response wording +- The numeric clustering results -```bash -BQ_EVAL_WRITE_MODE=protected \ - python -m benchmarks.bigquerybench.runner --filter ml_cluster_penguins -``` - -Or for CI, set it in the environment configuration. - -### Step 6: Validate the Eval Case +### Step 3: Validate ```bash -# Dry-run: check JSON is valid -python -m benchmarks.bigquerybench.runner --dry-run - -# Single-case run BQ_EVAL_WRITE_MODE=protected \ python -m benchmarks.bigquerybench.runner --filter ml_cluster_penguins - -# Multi-run for variance check (ML outputs may vary) -BQ_EVAL_WRITE_MODE=protected \ - python -m benchmarks.bigquerybench.runner --filter ml_cluster --num-runs 3 ``` Expected output: ``` -================================================================= - BigQueryBench Evaluation — ADK BigQueryToolset -================================================================= - -[1/1] Running: ml_cluster_penguins - Response: Here are the cluster assignments for the penguins... - Scores: schema=1.0 tools=1.00 output=PASS - -Eval Case Schema Tools Output Result ---------------------------------------------------------------------------- -ml_cluster_penguins 1.0 1.00 1.0 PASS ---------------------------------------------------------------------------- - -================================================================= - Leaderboard Summary -================================================================= - Framework: ADK BigQueryToolset - Cases: 1/1 (100.0%) - Avg Schema Disc.: 1.00 - Avg Tool Usage: 1.00 - Elapsed: 12.3s -================================================================= +[1/1] ml_cluster_penguins + -> get_table_info(project_id='bigquery-public-data', dataset_id='ml_datasets', table_id='penguins') + -> cluster_data(project_id='bigquery-public-data') + tools=1.00 args=1.00 PASS ``` -### Step 7: Commit +### Step 4: Commit ```bash git add benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json -git commit -m "eval(bigquerybench): add clustering eval case for cluster_data tool" +git commit -m "eval(bigquerybench): add clustering eval for cluster_data" ``` -**Files changed:** Only `bigquerybench_eval.json`. No code changes -needed unless a new metric was added (Step 2b). +**Only `bigquerybench_eval.json` changed.** No code changes. ---- +### When Do You Need Code Changes? -### Step 2b: Adding a Custom Metric (When Needed) +| Scenario | JSON | `metrics.py` | `runner.py` | `agent.py` | +|----------|:----:|:------------:|:-----------:|:----------:| +| New eval case, existing tool | Yes | - | - | - | +| New tool, trace check is enough | Yes | - | - | - | +| New tool, need to check a non-key arg (e.g., `num_clusters`) | Yes | Yes (add arg to `_KEY_ARGS`) | - | - | +| New tool, need entirely new metric | Yes | Yes | Yes | - | +| Agent instruction or write-mode change | - | - | - | Yes | -If existing metrics are insufficient for your new tool, add a metric -function to `metrics.py`. Here's a concrete example: a -`clustering_quality_score` that validates the agent used the correct -number of clusters. - -**1. Write the metric function in `metrics.py`:** +**Adding a checked arg** is a one-line change in `metrics.py`: ```python -def clustering_quality_score( - eval_metric: EvalMetric, - actual_invocations: list[Invocation], - expected_invocations: Optional[list[Invocation]], - conversation_scenario: Optional[ConversationScenario] = None, -) -> EvaluationResult: - """Score 1.0 if cluster_data was called with correct num_clusters. - - Checks that the agent called cluster_data and that the - num_clusters argument matches the expected value. - """ - if not expected_invocations: - return EvaluationResult( - overall_score=1.0, - overall_eval_status=EvalStatus.PASSED, - ) - - # Extract expected num_clusters from expected tool calls. - expected_k = None - for inv in expected_invocations: - for tc in get_all_tool_calls(inv.intermediate_data): - if tc.name == "cluster_data" and tc.args: - expected_k = tc.args.get("num_clusters") - break - - # Extract actual num_clusters from actual tool calls. - actual_k = None - for inv in actual_invocations: - for tc in get_all_tool_calls(inv.intermediate_data): - if tc.name == "cluster_data" and tc.args: - actual_k = tc.args.get("num_clusters") - break - - if expected_k is None: - score = 1.0 # No clustering expected - elif actual_k is None: - score = 0.0 # Clustering expected but not called - else: - score = 1.0 if actual_k == expected_k else 0.0 - - status = ( - EvalStatus.PASSED if score >= 1.0 - else EvalStatus.FAILED - ) - return EvaluationResult( - overall_score=score, - overall_eval_status=status, - per_invocation_results=_make_per_invocation( - actual_invocations, expected_invocations, - score, status, - ), - ) +_KEY_ARGS = frozenset({ + "project_id", + "dataset_id", + "table_id", + "num_clusters", # ← add here +}) ``` -**2. Wire it into the runner (`runner.py`):** - -```python -from .metrics import clustering_quality_score - -def score_invocations(actual, expected): - # ... existing metrics ... - - result = clustering_quality_score( - metric, actual_invocations, expected_invocations, - ) - scores["clustering_quality"] = result.overall_score or 0.0 +## Architecture - return scores ``` - -**3. Update the results table to show the new column.** - -**4. Commit all changed files:** - -```bash -git add benchmarks/bigquerybench/metrics.py \ - benchmarks/bigquerybench/runner.py \ - benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json -git commit -m "eval(bigquerybench): add cluster_data eval case and clustering_quality metric" +benchmarks/bigquerybench/ +├── __init__.py +├── agent.py # LlmAgent + BigQueryToolset (read-only default) +├── runner.py # Runs agent, collects trace, scores +├── metrics.py # tool_invocation_score + tool_args_score +└── eval_sets/ + └── bigquerybench_eval.json # 7 eval cases ``` ---- - -### Summary: Files to Touch per Scenario - -| Scenario | `eval.json` | `metrics.py` | `runner.py` | `agent.py` | -|----------|:-----------:|:------------:|:-----------:|:----------:| -| New eval case for existing tool | **Yes** | No | No | No | -| New tool, existing metrics sufficient | **Yes** | No | No | No | -| New tool, needs custom metric | **Yes** | **Yes** | **Yes** | No | -| New tool, needs agent config change | **Yes** | Maybe | Maybe | **Yes** | - -The toolset auto-discovers new tools, so `agent.py` only changes if -the agent's instruction prompt or write mode needs updating. - -## Metrics Reference - -| Metric | Function Path | Threshold | Pass Condition | -|--------|--------------|-----------|----------------| -| Schema Discovery | `benchmarks.bigquerybench.metrics.schema_discovery_score` | 1.0 | Any schema tool called | -| Tool Usage | `benchmarks.bigquerybench.metrics.tool_usage_score` | 0.5 | >= 50% expected tools called | -| Output Correctness | `benchmarks.bigquerybench.metrics.output_correctness_score` | 1.0 | All reference lines present | - ## Environment Variables | Variable | Required | Description | |----------|----------|-------------| -| `GOOGLE_CLOUD_PROJECT` | Yes | GCP project for BigQuery API calls | -| `GOOGLE_GENAI_USE_VERTEXAI` | Conditional | Set to `1` for Vertex AI LLM backend | -| `GOOGLE_API_KEY` | Conditional | API key for Google AI Studio backend | -| `BQ_EVAL_WRITE_MODE` | No | Override write mode (`blocked`/`protected`/`allowed`). Default: `blocked` | +| `GOOGLE_CLOUD_PROJECT` | Yes | GCP project for BigQuery API | +| `GOOGLE_GENAI_USE_VERTEXAI` | Conditional | `1` for Vertex AI backend | +| `GOOGLE_API_KEY` | Conditional | API key for AI Studio backend | +| `BQ_EVAL_WRITE_MODE` | No | `blocked` (default) / `protected` / `allowed` | ## Troubleshooting -| Symptom | Cause | Fix | -|---------|-------|-----| -| `403 Access Denied` | Missing BigQuery API access | Enable BigQuery API in GCP console; run `gcloud auth application-default login` | -| `execute_sql` returns empty | Query references wrong project | Ensure public dataset queries use `bigquery-public-data` as project | -| `forecast` fails with write error | `write_mode=BLOCKED` | Set `BQ_EVAL_WRITE_MODE=protected` for AI/ML eval cases | -| Low `schema_discovery_score` | Agent skips exploration | Strengthen agent instructions to always explore schema first | -| Flaky `output_correctness_score` | Reference lines too specific | Use structural phrases, not exact numeric values | +| Symptom | Fix | +|---------|-----| +| `403 Access Denied` | `gcloud auth application-default login` + enable BigQuery API | +| `tool_invocation_score = 0` | Agent didn't call expected tool — check agent instructions | +| `tool_args_score < 1.0` | Agent pointed at wrong dataset/table — check user query specificity | +| AI/ML tool fails | Set `BQ_EVAL_WRITE_MODE=protected` | diff --git a/benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json b/benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json index cc2bacd42f..1906baa502 100644 --- a/benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json +++ b/benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json @@ -1,7 +1,7 @@ { "eval_set_id": "bigquerybench-adk-v1", "name": "BigQueryBench ADK Evaluation", - "description": "End-to-end evaluation cases for BigQuery skills using public datasets. Covers schema exploration, SQL generation, multi-step analysis, and AI/ML operations.", + "description": "Trace-based evaluation: verify the agent calls the correct BigQuery tools with the correct dataset/table args. No response text matching.", "eval_cases": [ { "eval_id": "schema_list_datasets", @@ -12,10 +12,6 @@ "parts": [{"text": "What datasets are available in the bigquery-public-data project? List a few of them."}], "role": "user" }, - "final_response": { - "parts": [{"text": "usa_names\nsamples\nnoaa_gsod"}], - "role": "model" - }, "intermediate_data": { "tool_uses": [ {"name": "list_dataset_ids", "args": {"project_id": "bigquery-public-data"}} @@ -37,10 +33,6 @@ "parts": [{"text": "What tables exist in the usa_names dataset in bigquery-public-data?"}], "role": "user" }, - "final_response": { - "parts": [{"text": "usa_1910_current"}], - "role": "model" - }, "intermediate_data": { "tool_uses": [ {"name": "list_table_ids", "args": {"project_id": "bigquery-public-data", "dataset_id": "usa_names"}} @@ -62,10 +54,6 @@ "parts": [{"text": "What columns and data types does the usa_1910_current table in the usa_names dataset have?"}], "role": "user" }, - "final_response": { - "parts": [{"text": "name\nyear\ngender\nstate\nnumber"}], - "role": "model" - }, "intermediate_data": { "tool_uses": [ {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "usa_names", "table_id": "usa_1910_current"}} @@ -87,14 +75,10 @@ "parts": [{"text": "Using the samples.shakespeare table in bigquery-public-data, which work (corpus) has the most unique words? Show the top 3 works by distinct word count."}], "role": "user" }, - "final_response": { - "parts": [{"text": "hamlet\nkinghenryv\nkingrichardiii"}], - "role": "model" - }, "intermediate_data": { "tool_uses": [ {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "samples", "table_id": "shakespeare"}}, - {"name": "execute_sql", "args": {"project_id": "bigquery-public-data", "query": "SELECT corpus, COUNT(DISTINCT word) AS unique_words FROM `bigquery-public-data.samples.shakespeare` GROUP BY corpus ORDER BY unique_words DESC LIMIT 3"}} + {"name": "execute_sql", "args": {"project_id": "bigquery-public-data"}} ], "tool_responses": [], "intermediate_responses": [] @@ -113,14 +97,10 @@ "parts": [{"text": "What were the top 5 most popular baby names in the year 2020 across all states? Use the usa_names.usa_1910_current table in bigquery-public-data."}], "role": "user" }, - "final_response": { - "parts": [{"text": "Olivia\nEmma\nLiam\nNoah\nAva"}], - "role": "model" - }, "intermediate_data": { "tool_uses": [ {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "usa_names", "table_id": "usa_1910_current"}}, - {"name": "execute_sql", "args": {"project_id": "bigquery-public-data", "query": "SELECT name, SUM(number) AS total FROM `bigquery-public-data.usa_names.usa_1910_current` WHERE year = 2020 GROUP BY name ORDER BY total DESC LIMIT 5"}} + {"name": "execute_sql", "args": {"project_id": "bigquery-public-data"}} ], "tool_responses": [], "intermediate_responses": [] @@ -139,14 +119,10 @@ "parts": [{"text": "How many distinct baby names were registered in each decade from 1950 to 2000 (inclusive)? Use the usa_names.usa_1910_current table in bigquery-public-data."}], "role": "user" }, - "final_response": { - "parts": [{"text": "1950\n1960\n1970\n1980\n1990\n2000"}], - "role": "model" - }, "intermediate_data": { "tool_uses": [ {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "usa_names", "table_id": "usa_1910_current"}}, - {"name": "execute_sql", "args": {"project_id": "bigquery-public-data", "query": "SELECT CAST(FLOOR(year / 10) * 10 AS INT64) AS decade, COUNT(DISTINCT name) AS distinct_names FROM `bigquery-public-data.usa_names.usa_1910_current` WHERE year BETWEEN 1950 AND 2009 GROUP BY decade ORDER BY decade"}} + {"name": "execute_sql", "args": {"project_id": "bigquery-public-data"}} ], "tool_responses": [], "intermediate_responses": [] @@ -165,15 +141,11 @@ "parts": [{"text": "I want to analyze the Austin bikeshare data in bigquery-public-data. First explore what tables and columns are available, then tell me the top 5 busiest start stations by total trip count."}], "role": "user" }, - "final_response": { - "parts": [{"text": "start_station_name\ntrip"}], - "role": "model" - }, "intermediate_data": { "tool_uses": [ {"name": "list_table_ids", "args": {"project_id": "bigquery-public-data", "dataset_id": "austin_bikeshare"}}, {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "austin_bikeshare", "table_id": "bikeshare_trips"}}, - {"name": "execute_sql", "args": {"project_id": "bigquery-public-data", "query": "SELECT start_station_name, COUNT(*) AS trip_count FROM `bigquery-public-data.austin_bikeshare.bikeshare_trips` GROUP BY start_station_name ORDER BY trip_count DESC LIMIT 5"}} + {"name": "execute_sql", "args": {"project_id": "bigquery-public-data"}} ], "tool_responses": [], "intermediate_responses": [] diff --git a/benchmarks/bigquerybench/metrics.py b/benchmarks/bigquerybench/metrics.py index 66db5d3f9b..5e78a9114b 100644 --- a/benchmarks/bigquerybench/metrics.py +++ b/benchmarks/bigquerybench/metrics.py @@ -12,9 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Custom metrics for BigQueryBench evaluation. +"""Trace-based metrics for BigQueryBench evaluation. -Three metrics following the ADK custom metric function signature: +Two metrics that check the agent's tool-call trace only (no response +text matching). This makes evaluation deterministic and easy to +maintain: just specify which tools should be called and with which +key arguments. def metric_fn( eval_metric: EvalMetric, @@ -24,9 +27,8 @@ def metric_fn( ) -> EvaluationResult Reference via dotted path in eval configs: - "benchmarks.bigquerybench.metrics.schema_discovery_score" - "benchmarks.bigquerybench.metrics.tool_usage_score" - "benchmarks.bigquerybench.metrics.output_correctness_score" + "benchmarks.bigquerybench.metrics.tool_invocation_score" + "benchmarks.bigquerybench.metrics.tool_args_score" """ from __future__ import annotations @@ -41,22 +43,12 @@ def metric_fn( from google.adk.evaluation.evaluator import EvaluationResult from google.adk.evaluation.evaluator import PerInvocationResult -# Tools that count as "schema exploration". -_SCHEMA_TOOLS = frozenset({ - "list_dataset_ids", - "list_table_ids", - "get_dataset_info", - "get_table_info", -}) - -def _get_tool_names(invocations: list[Invocation]) -> list[str]: - """Extract all tool call names from a list of invocations.""" - names = [] +def _get_tool_calls(invocations: list[Invocation]): + """Yield (name, args_dict) for every tool call in the trace.""" for inv in invocations: - for tool_call in get_all_tool_calls(inv.intermediate_data): - names.append(tool_call.name) - return names + for tc in get_all_tool_calls(inv.intermediate_data): + yield tc.name, (tc.args or {}) def _make_per_invocation( @@ -65,7 +57,6 @@ def _make_per_invocation( score: float, status: EvalStatus, ) -> list[PerInvocationResult]: - """Build per-invocation results list.""" results = [] for i, actual in enumerate(actual_invocations): expected = None @@ -82,51 +73,23 @@ def _make_per_invocation( return results -def schema_discovery_score( - eval_metric: EvalMetric, - actual_invocations: list[Invocation], - expected_invocations: Optional[list[Invocation]], - conversation_scenario: Optional[ConversationScenario] = None, -) -> EvaluationResult: - """Score 1.0 if the agent called at least one schema exploration tool. - - Schema tools: list_dataset_ids, list_table_ids, get_dataset_info, - get_table_info. - - This metric enforces the "explore before query" pattern — agents - should understand the schema before generating SQL or calling AI/ML - tools. - """ - tool_names = set(_get_tool_names(actual_invocations)) - called_schema = bool(tool_names & _SCHEMA_TOOLS) - - score = 1.0 if called_schema else 0.0 - status = EvalStatus.PASSED if called_schema else EvalStatus.FAILED - - return EvaluationResult( - overall_score=score, - overall_eval_status=status, - per_invocation_results=_make_per_invocation( - actual_invocations, - expected_invocations, - score, - status, - ), - ) +# ── Metric 1: correct tools invoked ────────────────────────────── -def tool_usage_score( +def tool_invocation_score( eval_metric: EvalMetric, actual_invocations: list[Invocation], expected_invocations: Optional[list[Invocation]], conversation_scenario: Optional[ConversationScenario] = None, ) -> EvaluationResult: - """Fraction of expected tool calls that were actually made. + """Score = fraction of expected tool names present in the trace. - Score = |expected_tools ∩ actual_tools| / |expected_tools|. - Uses set-based matching (any order). Passes at >= 0.5. + Checks that the agent called the right BigQuery functions (e.g. + ``get_table_info``, ``execute_sql``, ``forecast``). Order does + not matter; extra tool calls are ignored. - Same semantics as SkillsBench tool_usage_score for consistency. + Score = |expected_names ∩ actual_names| / |expected_names|. + Pass threshold: 1.0 (all expected tools must be called). """ if not expected_invocations: return EvaluationResult( @@ -134,8 +97,8 @@ def tool_usage_score( overall_eval_status=EvalStatus.PASSED, ) - expected_names = set(_get_tool_names(expected_invocations)) - actual_names = set(_get_tool_names(actual_invocations)) + expected_names = {name for name, _ in _get_tool_calls(expected_invocations)} + actual_names = {name for name, _ in _get_tool_calls(actual_invocations)} if not expected_names: score = 1.0 @@ -143,7 +106,7 @@ def tool_usage_score( matched = expected_names & actual_names score = len(matched) / len(expected_names) - status = EvalStatus.PASSED if score >= 0.5 else EvalStatus.FAILED + status = EvalStatus.PASSED if score >= 1.0 else EvalStatus.FAILED return EvaluationResult( overall_score=score, @@ -157,64 +120,72 @@ def tool_usage_score( ) -def output_correctness_score( +# ── Metric 2: correct args on key tool calls ───────────────────── + +# Args that identify the *target data* — these are what we check. +# We intentionally skip volatile args like ``query`` (the exact SQL +# the LLM generates will vary) and only verify that the agent +# pointed at the right dataset / table / project. +_KEY_ARGS = frozenset({ + "project_id", + "dataset_id", + "table_id", +}) + + +def tool_args_score( eval_metric: EvalMetric, actual_invocations: list[Invocation], expected_invocations: Optional[list[Invocation]], conversation_scenario: Optional[ConversationScenario] = None, ) -> EvaluationResult: - """Binary pass/fail: 1.0 if response contains all reference lines. + """Score = fraction of expected (tool, key-arg) pairs matched. - Each non-empty line in the expected final_response is a reference - phrase. The actual response must contain every phrase as a - case-insensitive substring. + For each expected tool call that has ``project_id``, ``dataset_id``, + or ``table_id`` in its args, check that the agent made a call to + the *same tool* with the *same value* for that arg. This verifies + the agent loaded the right reference data (correct dataset, correct + table) without caring about the exact SQL or response text. - Same semantics as SkillsBench skillsbench_binary_score for - consistency. + Score = matched_pairs / expected_pairs. Pass threshold: 1.0. + If no key args exist in the expected trace, score is 1.0 (vacuous). """ - if not expected_invocations or not actual_invocations: + if not expected_invocations: return EvaluationResult( - overall_score=0.0, - overall_eval_status=EvalStatus.NOT_EVALUATED, + overall_score=1.0, + overall_eval_status=EvalStatus.PASSED, ) - # Get the last actual response text. - actual_text = "" - for inv in reversed(actual_invocations): - if inv.final_response and inv.final_response.parts: - for part in inv.final_response.parts: - if part.text: - actual_text = part.text - break - if actual_text: - break - - # Get the expected response text. - expected_text = "" - for inv in reversed(expected_invocations): - if inv.final_response and inv.final_response.parts: - for part in inv.final_response.parts: - if part.text: - expected_text = part.text - break - if expected_text: - break - - if not expected_text: + # Build expected set: (tool_name, arg_key, arg_value). + expected_pairs: set[tuple[str, str, str]] = set() + for name, args in _get_tool_calls(expected_invocations): + for key in _KEY_ARGS: + if key in args: + expected_pairs.add((name, key, str(args[key]))) + + if not expected_pairs: + # No key args to check — pass vacuously. return EvaluationResult( - overall_score=0.0, - overall_eval_status=EvalStatus.NOT_EVALUATED, + overall_score=1.0, + overall_eval_status=EvalStatus.PASSED, + per_invocation_results=_make_per_invocation( + actual_invocations, + expected_invocations, + 1.0, + EvalStatus.PASSED, + ), ) - reference_lines = [ - line.strip() for line in expected_text.split("\n") if line.strip() - ] - actual_lower = actual_text.lower() - matched = sum(1 for line in reference_lines if line.lower() in actual_lower) + # Build actual set the same way. + actual_pairs: set[tuple[str, str, str]] = set() + for name, args in _get_tool_calls(actual_invocations): + for key in _KEY_ARGS: + if key in args: + actual_pairs.add((name, key, str(args[key]))) - is_pass = matched == len(reference_lines) and len(reference_lines) > 0 - score = 1.0 if is_pass else 0.0 - status = EvalStatus.PASSED if is_pass else EvalStatus.FAILED + matched = expected_pairs & actual_pairs + score = len(matched) / len(expected_pairs) + status = EvalStatus.PASSED if score >= 1.0 else EvalStatus.FAILED return EvaluationResult( overall_score=score, diff --git a/benchmarks/bigquerybench/runner.py b/benchmarks/bigquerybench/runner.py index 63c77449a1..e21c6c4db1 100644 --- a/benchmarks/bigquerybench/runner.py +++ b/benchmarks/bigquerybench/runner.py @@ -12,26 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Standalone BigQueryBench runner that produces leaderboard scores. +"""Trace-based BigQueryBench runner. + +Evaluates whether the agent calls the correct BigQuery tools with the +correct dataset/table arguments. No response text matching — only +the tool-call trace matters. Usage: python -m benchmarks.bigquerybench.runner - python -m benchmarks.bigquerybench.runner --num-runs 3 python -m benchmarks.bigquerybench.runner --filter sql_shakespeare + python -m benchmarks.bigquerybench.runner --num-runs 3 python -m benchmarks.bigquerybench.runner --dry-run Environment variables: GOOGLE_CLOUD_PROJECT — GCP project for BigQuery API calls GOOGLE_API_KEY — API key for Google AI Studio GOOGLE_GENAI_USE_VERTEXAI — Set to 1 for Vertex AI backend - BQ_EVAL_WRITE_MODE — Override write mode (blocked/protected) - -This script: -1. Loads the BigQueryBench agent and eval set -2. Runs each case through the ADK Runner -3. Applies 3 custom metrics (schema_discovery, tool_usage, output) -4. Outputs per-case results and a leaderboard-format summary -5. Exits with code 0 if all pass, 1 if any fail + BQ_EVAL_WRITE_MODE — blocked / protected / allowed """ from __future__ import annotations @@ -56,9 +53,8 @@ from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.adk.utils.context_utils import Aclosing -from .metrics import output_correctness_score -from .metrics import schema_discovery_score -from .metrics import tool_usage_score +from .metrics import tool_args_score +from .metrics import tool_invocation_score logger = logging.getLogger(__name__) @@ -67,7 +63,6 @@ def load_eval_set(path: pathlib.Path) -> EvalSet: - """Load an EvalSet from a JSON file.""" with open(path) as f: data = json.load(f) return EvalSet.model_validate(data) @@ -76,65 +71,54 @@ def load_eval_set(path: pathlib.Path) -> EvalSet: def print_header(): print() print("=" * 65) - print(" BigQueryBench Evaluation — ADK BigQueryToolset") + print(" BigQueryBench — Trace-Based Evaluation") print("=" * 65) print() def print_results_table(results: dict[str, dict[str, float]]): - """Print per-case results as a formatted table.""" - print( - f"{'Eval Case':<40} {'Schema':>8} {'Tools':>7}" - f" {'Output':>8} {'Result':>8}" - ) - print("-" * 75) + print(f"{'Eval Case':<40} {'Tools':>7} {'Args':>7} {'Result':>8}") + print("-" * 65) for case_id, scores in results.items(): short_id = case_id[:39] - schema = scores.get("schema_discovery", 0.0) - tools = scores.get("tool_usage", 0.0) - output = scores.get("output_correctness", 0.0) - mark = "PASS" if output >= 1.0 else "FAIL" - print( - f"{short_id:<40} {schema:>7.1f} {tools:>7.2f}" - f" {output:>7.1f} {mark:>5}" - ) - print("-" * 75) - - -def print_leaderboard_summary( + tools = scores.get("tool_invocation", 0.0) + args = scores.get("tool_args", 0.0) + passed = tools >= 1.0 and args >= 1.0 + mark = "PASS" if passed else "FAIL" + print(f"{short_id:<40} {tools:>7.2f} {args:>7.2f} {mark:>5}") + print("-" * 65) + + +def print_summary( results: dict[str, dict[str, float]], num_cases: int, elapsed: float, ): - """Print a leaderboard-format summary.""" passed = sum( - 1 for s in results.values() if s.get("output_correctness", 0.0) >= 1.0 + 1 + for s in results.values() + if s.get("tool_invocation", 0.0) >= 1.0 and s.get("tool_args", 0.0) >= 1.0 ) - avg_schema = sum( - s.get("schema_discovery", 0.0) for s in results.values() + avg_tools = sum( + s.get("tool_invocation", 0.0) for s in results.values() ) / max(len(results), 1) - avg_tools = sum(s.get("tool_usage", 0.0) for s in results.values()) / max( + avg_args = sum(s.get("tool_args", 0.0) for s in results.values()) / max( len(results), 1 ) pct = (passed / max(num_cases, 1)) * 100 print() print("=" * 65) - print(" Leaderboard Summary") + print(" Summary") print("=" * 65) - print(f" Framework: ADK BigQueryToolset") print(f" Cases: {passed}/{num_cases} ({pct:.1f}%)") - print(f" Avg Schema Disc.: {avg_schema:.2f}") - print(f" Avg Tool Usage: {avg_tools:.2f}") + print(f" Avg Tool Match: {avg_tools:.2f}") + print(f" Avg Args Match: {avg_args:.2f}") print(f" Elapsed: {elapsed:.1f}s") print("=" * 65) -async def run_single_eval_case( - root_agent, - eval_case, -) -> list[Invocation]: - """Run a single eval case through the Runner.""" +async def run_single_eval_case(root_agent, eval_case) -> list[Invocation]: session_service = InMemorySessionService() artifact_service = InMemoryArtifactService() memory_service = InMemoryMemoryService() @@ -187,35 +171,32 @@ async def run_single_eval_case( def score_invocations( - actual_invocations: list[Invocation], - expected_invocations: Optional[list[Invocation]], + actual: list[Invocation], + expected: Optional[list[Invocation]], ) -> dict[str, float]: - """Apply all 3 metrics and return scores.""" metric = EvalMetric(metric_name="bigquerybench") - scores = {} - result = schema_discovery_score( - metric, - actual_invocations, - expected_invocations, - ) - scores["schema_discovery"] = result.overall_score or 0.0 + r1 = tool_invocation_score(metric, actual, expected) + r2 = tool_args_score(metric, actual, expected) - result = tool_usage_score( - metric, - actual_invocations, - expected_invocations, - ) - scores["tool_usage"] = result.overall_score or 0.0 + return { + "tool_invocation": r1.overall_score or 0.0, + "tool_args": r2.overall_score or 0.0, + } - result = output_correctness_score( - metric, - actual_invocations, - expected_invocations, - ) - scores["output_correctness"] = result.overall_score or 0.0 - return scores +def _print_trace(actual_invocations: list[Invocation]): + """Print the tool-call trace for debugging.""" + from google.adk.evaluation.eval_case import get_all_tool_calls + + for inv in actual_invocations: + for tc in get_all_tool_calls(inv.intermediate_data): + args_summary = ", ".join( + f"{k}={v!r}" + for k, v in (tc.args or {}).items() + if k in ("project_id", "dataset_id", "table_id") + ) + print(f" -> {tc.name}({args_summary})") async def run_evaluation( @@ -223,11 +204,9 @@ async def run_evaluation( num_runs: int = 1, filter_str: Optional[str] = None, ) -> dict[str, dict[str, float]]: - """Run the BigQueryBench evaluation.""" path = eval_set_path or _DEFAULT_EVAL_SET eval_set = load_eval_set(path) - # Import agent (triggers BigQuery toolset setup). from .agent import root_agent cases = eval_set.eval_cases @@ -238,11 +217,10 @@ async def run_evaluation( return {} results: dict[str, dict[str, float]] = {} - total = len(cases) for idx, eval_case in enumerate(cases, 1): eval_id = eval_case.eval_id - print(f"\n[{idx}/{total}] Running: {eval_id}") + print(f"\n[{idx}/{len(cases)}] {eval_id}") run_scores: list[dict[str, float]] = [] for run in range(num_runs): @@ -251,40 +229,25 @@ async def run_evaluation( try: actual = await run_single_eval_case(root_agent, eval_case) + _print_trace(actual) - # Print response preview. - for inv in actual: - if inv.final_response and inv.final_response.parts: - for part in inv.final_response.parts: - if part.text: - preview = part.text[:200].replace("\n", " ") - print(f" Response: {preview}...") - break - - expected = eval_case.conversation - scores = score_invocations(actual, expected) + scores = score_invocations(actual, eval_case.conversation) run_scores.append(scores) - schema = scores["schema_discovery"] - tools = scores["tool_usage"] - output = scores["output_correctness"] - mark = "PASS" if output >= 1.0 else "FAIL" - print(f" Scores: schema={schema:.1f} tools={tools:.2f} output={mark}") + tools = scores["tool_invocation"] + args = scores["tool_args"] + passed = tools >= 1.0 and args >= 1.0 + mark = "PASS" if passed else "FAIL" + print(f" tools={tools:.2f} args={args:.2f} {mark}") except Exception as e: logger.error("Error running %s: %s", eval_id, e) print(f" ERROR: {e}") - run_scores.append({ - "schema_discovery": 0.0, - "tool_usage": 0.0, - "output_correctness": 0.0, - }) + run_scores.append({"tool_invocation": 0.0, "tool_args": 0.0}) - # Average scores across runs. avg: dict[str, float] = {} - for key in ["schema_discovery", "tool_usage", "output_correctness"]: - values = [s[key] for s in run_scores] - avg[key] = sum(values) / len(values) + for key in ("tool_invocation", "tool_args"): + avg[key] = sum(s[key] for s in run_scores) / len(run_scores) results[eval_id] = avg return results @@ -292,30 +255,26 @@ async def run_evaluation( def main(): parser = argparse.ArgumentParser( - description="BigQueryBench evaluation runner for ADK", + description="BigQueryBench trace-based evaluation runner", ) parser.add_argument( "--eval-set", type=pathlib.Path, default=None, - help="Path to eval set JSON (default: built-in)", ) parser.add_argument( "--num-runs", type=int, default=1, - help="Number of runs per case (default: 1)", ) parser.add_argument( "--filter", type=str, default=None, - help="Substring filter for eval_id (e.g., 'sql_shakespeare')", ) parser.add_argument( "--dry-run", action="store_true", - help="Validate eval set JSON without running LLM inference", ) args = parser.parse_args() @@ -323,12 +282,15 @@ def main(): if args.dry_run: path = args.eval_set or _DEFAULT_EVAL_SET - eval_set = load_eval_set(path) - print(f"Eval set: {eval_set.name}") - print(f"Cases: {len(eval_set.eval_cases)}") - for case in eval_set.eval_cases: - print(f" - {case.eval_id}") - print("\nDry run: eval set JSON is valid.") + es = load_eval_set(path) + print(f"Eval set: {es.name}") + print(f"Cases: {len(es.eval_cases)}") + for case in es.eval_cases: + tools = [ + t.name for t in case.conversation[0].intermediate_data.tool_uses or [] + ] + print(f" {case.eval_id}: {' -> '.join(tools)}") + print("\nJSON valid.") sys.exit(0) print_header() @@ -343,20 +305,15 @@ def main(): ) elapsed = time.time() - start - eval_path = args.eval_set or _DEFAULT_EVAL_SET - eval_set = load_eval_set(eval_path) - - num_cases = len(eval_set.eval_cases) - if args.filter: - num_cases = len(results) + num_cases = len(results) print() print_results_table(results) - print_leaderboard_summary(results, num_cases, elapsed) + print_summary(results, num_cases, elapsed) - # Exit code: 0 if all pass, 1 if any fail. all_pass = all( - s.get("output_correctness", 0.0) >= 1.0 for s in results.values() + s.get("tool_invocation", 0.0) >= 1.0 and s.get("tool_args", 0.0) >= 1.0 + for s in results.values() ) sys.exit(0 if all_pass else 1) From faf852c7d2fbc57a61daa50869f4dbfbe949ff80 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 08:08:13 -0800 Subject: [PATCH 45/53] feat(bigquerybench): add skill invocation + LLM-as-judge instruction adherence evaluation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three-metric evaluation pipeline: 1. tool_invocation_score — trace-based, verifies correct skill tools called 2. tool_args_score — trace-based, verifies correct skill/resource/script args 3. instruction_adherence_score — LLM-as-judge, checks natural-language rubrics Includes bq-sql-analyst skill, 5 eval cases with 12 rubrics total, and 14 unit tests covering all metrics (mocked LLM for judge tests). Co-Authored-By: Claude Opus 4.6 --- benchmarks/bigquerybench/README.md | 365 +++++++++--------- benchmarks/bigquerybench/agent.py | 42 +- .../eval_sets/bigquerybench_eval.json | 138 ++++--- benchmarks/bigquerybench/metrics.py | 178 +++++++-- benchmarks/bigquerybench/runner.py | 117 +++--- .../skills/bq-sql-analyst/SKILL.md | 34 ++ .../references/public-datasets.md | 34 ++ .../bq-sql-analyst/scripts/format_results.py | 51 +++ tests/unittests/benchmarks/__init__.py | 13 + .../benchmarks/bigquerybench/__init__.py | 13 + .../benchmarks/bigquerybench/test_metrics.py | 344 +++++++++++++++++ 11 files changed, 1000 insertions(+), 329 deletions(-) create mode 100644 benchmarks/bigquerybench/skills/bq-sql-analyst/SKILL.md create mode 100644 benchmarks/bigquerybench/skills/bq-sql-analyst/references/public-datasets.md create mode 100644 benchmarks/bigquerybench/skills/bq-sql-analyst/scripts/format_results.py create mode 100644 tests/unittests/benchmarks/__init__.py create mode 100644 tests/unittests/benchmarks/bigquerybench/__init__.py create mode 100644 tests/unittests/benchmarks/bigquerybench/test_metrics.py diff --git a/benchmarks/bigquerybench/README.md b/benchmarks/bigquerybench/README.md index a866db7a2e..aadabb022b 100644 --- a/benchmarks/bigquerybench/README.md +++ b/benchmarks/bigquerybench/README.md @@ -1,21 +1,18 @@ -# BigQueryBench: Trace-Based Evaluation for BigQuery Skills +# BigQueryBench: Skill Invocation & Instruction Adherence Evaluation ## Overview -BigQueryBench verifies that an agent built with ADK's `BigQueryToolset` -**calls the correct tools with the correct arguments**. It inspects -the tool-call trace only — no response text matching. +BigQueryBench evaluates agents built with ADK's `SkillToolset` on +two dimensions: -For each eval case, the pipeline checks two things: +1. **Skill invocation correctness** (trace-based) — Did the agent + call the right skill tools with the right arguments? +2. **Instruction adherence** (LLM-as-judge) — Did the agent follow + the skill's instructions and produce correct results? -1. **Tool invocation** — Did the agent call the right BigQuery - functions? (e.g., `get_table_info` then `execute_sql`) -2. **Tool arguments** — Did those calls point at the right data? - (e.g., `project_id="bigquery-public-data"`, - `dataset_id="usa_names"`, `table_id="usa_1910_current"`) - -This makes evaluation deterministic, easy to maintain, and immune to -LLM response wording variance. +The trace-based checks are deterministic. The instruction adherence +checks use natural-language rubrics evaluated by a judge LLM, making +them easy to write and immune to exact-wording variance. ## Quick Start @@ -28,52 +25,64 @@ export GOOGLE_GENAI_USE_VERTEXAI=1 python -m benchmarks.bigquerybench.runner # Run one case -python -m benchmarks.bigquerybench.runner --filter schema_list_tables +python -m benchmarks.bigquerybench.runner --filter skill_load # Dry-run (validate JSON only, no LLM calls) python -m benchmarks.bigquerybench.runner --dry-run + +# Run unit tests (no API keys needed) +pytest tests/unittests/benchmarks/bigquerybench/ -v ``` ## How It Works ``` eval_sets/bigquerybench_eval.json - ↓ (user query + expected tool_uses) + ↓ (user query + expected tool_uses + rubrics) runner.py ↓ runs agent via ADK Runner ↓ collects event trace → Invocations metrics.py - ├── tool_invocation_score: expected tool names ⊆ actual tool names? - └── tool_args_score: expected (tool, project/dataset/table) ⊆ actual? + ├── tool_invocation_score: expected skill tool names ⊆ actual? + ├── tool_args_score: expected (tool, skill-arg) pairs ⊆ actual? + └── instruction_adherence_score: LLM judge checks rubrics ↓ -PASS if both scores = 1.0 +PASS if all three scores meet thresholds ``` +## Three Metrics + +| Metric | Type | What It Checks | Pass Condition | +|--------|------|----------------|----------------| +| `tool_invocation_score` | Trace | Correct skill tools called | Score = 1.0 | +| `tool_args_score` | Trace | Correct skill/resource/script targeted | Score = 1.0 | +| `instruction_adherence_score` | LLM judge | Agent followed instructions, output correct | Score >= 0.75 | + +A case **passes** when all three metrics meet their thresholds. + ## Eval Case Format -Each eval case specifies a user query and the expected tool calls. -No `final_response` is needed — only the trace matters. +Each eval case has three parts: +1. **`conversation`** — user query + expected skill tool calls +2. **`rubrics`** — natural-language assertions checked by the judge ```json { - "eval_id": "schema_get_table_info", + "eval_id": "skill_load_reference", "conversation": [ { "invocation_id": "inv-01", "user_content": { - "parts": [{"text": "What columns does usa_1910_current have?"}], + "parts": [{"text": "Load the public datasets reference from bq-sql-analyst."}], "role": "user" }, "intermediate_data": { "tool_uses": [ - { - "name": "get_table_info", - "args": { - "project_id": "bigquery-public-data", - "dataset_id": "usa_names", - "table_id": "usa_1910_current" - } - } + {"name": "load_skill", "args": {"name": "bq-sql-analyst"}}, + {"name": "load_skill_resource", "args": { + "skill_name": "bq-sql-analyst", + "path": "references/public-datasets.md" + }} ], "tool_responses": [], "intermediate_responses": [] @@ -81,94 +90,112 @@ No `final_response` is needed — only the trace matters. "creation_timestamp": 0.0 } ], + "rubrics": [ + { + "rubric_id": "shows_datasets", + "rubric_content": { + "text_property": "The response contains information about BigQuery public datasets." + } + }, + { + "rubric_id": "loaded_skill_first", + "rubric_content": { + "text_property": "The agent loaded the skill instructions before loading the resource." + } + } + ], "creation_timestamp": 0.0 } ``` -**What gets checked:** +### Trace Checks (deterministic) | Field in `args` | Checked? | Why | |-----------------|----------|-----| -| `project_id` | Yes | Must point at the right GCP project | -| `dataset_id` | Yes | Must load the right dataset | -| `table_id` | Yes | Must load the right table | -| `query` | **No** | Exact SQL varies — agent may write equivalent SQL differently | -| Other args | **No** | Tool-specific args (e.g., `horizon`, `num_clusters`) are not checked by default | +| `name` | Yes | Must load the right skill (`load_skill`) | +| `skill_name` | Yes | Must target the right skill (`load_skill_resource`, `run_skill_script`) | +| `path` | Yes | Must load the right resource (`load_skill_resource`) | +| `script_path` | Yes | Must run the right script (`run_skill_script`) | -## Metrics +### Rubrics (LLM-as-judge) -| Metric | What It Checks | Pass Condition | -|--------|---------------|----------------| -| `tool_invocation_score` | All expected tool names appear in the trace | Score = 1.0 | -| `tool_args_score` | All expected `(tool, project_id/dataset_id/table_id)` pairs appear in the trace | Score = 1.0 | +Each rubric is a natural-language assertion about the agent's behavior +or output. The judge LLM reads the conversation (user request + tool +trace + final response) and answers yes/no per rubric. -A case **passes** when both scores are 1.0. +**Example rubrics:** +```json +{"rubric_id": "r1", "rubric_content": {"text_property": "The agent used AI.classify to classify the data."}} +{"rubric_id": "r2", "rubric_content": {"text_property": "The result contains a markdown table with group statistics."}} +{"rubric_id": "r3", "rubric_content": {"text_property": "The agent loaded the skill before running the script."}} +``` ## Included Eval Cases -| eval_id | User Query | Expected Trace | -|---------|-----------|----------------| -| `schema_list_datasets` | "What datasets are in bigquery-public-data?" | `list_dataset_ids(project_id=bigquery-public-data)` | -| `schema_list_tables` | "What tables in usa_names?" | `list_table_ids(project_id=.., dataset_id=usa_names)` | -| `schema_get_table_info` | "Columns of usa_1910_current?" | `get_table_info(project_id=.., dataset_id=usa_names, table_id=usa_1910_current)` | -| `sql_shakespeare_unique_words` | "Top 3 works by unique words?" | `get_table_info(.., shakespeare)` → `execute_sql(..)` | -| `sql_usa_names_top_2020` | "Top 5 baby names in 2020?" | `get_table_info(.., usa_1910_current)` → `execute_sql(..)` | -| `sql_names_by_decade` | "Distinct names per decade 1950-2000?" | `get_table_info(.., usa_1910_current)` → `execute_sql(..)` | -| `multi_step_explore_and_query` | "Explore bikeshare, top 5 stations?" | `list_table_ids(..)` → `get_table_info(.., bikeshare_trips)` → `execute_sql(..)` | - -## Adding a New Eval Case - -### For an existing tool (e.g., new `execute_sql` scenario) +| eval_id | Expected Trace | Rubrics | +|---------|---------------|---------| +| `skill_list_skills` | `list_skills()` | Lists bq-sql-analyst; includes description | +| `skill_load_sql_analyst` | `load_skill(name=bq-sql-analyst)` | Describes capabilities; mentions scripts | +| `skill_load_reference` | `load_skill` → `load_skill_resource` | Shows datasets; loaded skill first | +| `skill_query_with_reference` | `load_skill` → `load_skill_resource` | Consulted reference; has ranking; followed workflow | +| `skill_run_format_script` | `load_skill` → `run_skill_script` | Loaded before run; has table; has columns | -Only add a JSON object to `bigquerybench_eval.json`. No code changes. +## Example Output -```json -{ - "eval_id": "sql_weather_hottest_day", - "conversation": [ - { - "invocation_id": "inv-weather-01", - "user_content": { - "parts": [{"text": "What was the hottest day recorded in the NOAA GSOD 2023 data?"}], - "role": "user" - }, - "intermediate_data": { - "tool_uses": [ - {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "noaa_gsod", "table_id": "gsod2023"}}, - {"name": "execute_sql", "args": {"project_id": "bigquery-public-data"}} - ], - "tool_responses": [], - "intermediate_responses": [] - }, - "creation_timestamp": 0.0 - } - ], - "creation_timestamp": 0.0 -} +``` +======================================================================== + BigQueryBench — Skill Evaluation +======================================================================== + +[1/5] skill_list_skills + -> list_skills() + tools=1.00 args=1.00 adherence=1.00 PASS + +[2/5] skill_load_sql_analyst + -> load_skill(name='bq-sql-analyst') + tools=1.00 args=1.00 adherence=1.00 PASS + +Eval Case Tools Args Adhere Result +------------------------------------------------------------------------ +skill_list_skills 1.00 1.00 1.00 PASS +skill_load_sql_analyst 1.00 1.00 1.00 PASS +skill_load_reference 1.00 1.00 1.00 PASS +skill_query_with_reference 1.00 1.00 0.67 FAIL +skill_run_format_script 1.00 1.00 1.00 PASS +------------------------------------------------------------------------ + +======================================================================== + Summary +======================================================================== + Cases: 4/5 (80.0%) + Avg Tool Match: 1.00 + Avg Args Match: 1.00 + Avg Adherence: 0.93 + Elapsed: 42.1s +======================================================================== ``` -Validate: `python -m benchmarks.bigquerybench.runner --filter sql_weather` +## Adding a New Eval Case -### For a new tool (e.g., `forecast`, `cluster_data`) +### For an existing skill -Same steps — just use the new tool name in `tool_uses`. The metrics -check tool names and key args generically, so no metric code changes -are needed. +Add a JSON object to `bigquerybench_eval.json` with both `tool_uses` +(trace expectations) and `rubrics` (instruction adherence assertions). ```json { - "eval_id": "ml_forecast_temperature", + "eval_id": "skill_explore_usa_names", "conversation": [ { - "invocation_id": "inv-forecast-01", + "invocation_id": "inv-new-01", "user_content": { - "parts": [{"text": "Forecast the next 7 days of temperature from NOAA GSOD 2023 data for station 725300."}], + "parts": [{"text": "Use the bq-sql-analyst skill to explore the usa_names dataset."}], "role": "user" }, "intermediate_data": { "tool_uses": [ - {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "noaa_gsod", "table_id": "gsod2023"}}, - {"name": "forecast", "args": {"project_id": "bigquery-public-data"}} + {"name": "load_skill", "args": {"name": "bq-sql-analyst"}}, + {"name": "load_skill_resource", "args": {"skill_name": "bq-sql-analyst", "path": "references/public-datasets.md"}} ], "tool_responses": [], "intermediate_responses": [] @@ -176,128 +203,82 @@ are needed. "creation_timestamp": 0.0 } ], - "creation_timestamp": 0.0 -} -``` - -For AI/ML tools that create temp models, set write mode: -```bash -BQ_EVAL_WRITE_MODE=protected python -m benchmarks.bigquerybench.runner --filter ml_forecast -``` - -## Complete Walkthrough: Adding a New BigQuery Skill - -This walkthrough uses a concrete example: a hypothetical -`cluster_data` tool (K-Means via BQML) being added to -`BigQueryToolset`. - -### Step 1: Register the tool - -The developer adds `cluster_data` to `bigquery_toolset.py`. Once -registered, it's automatically available to the eval agent — no -changes to `agent.py` needed. - -### Step 2: Write the eval case - -What we want to verify: when the user asks "cluster the penguins -data", the agent should: -1. Call `get_table_info` on `ml_datasets.penguins` (load the schema) -2. Call `cluster_data` against `bigquery-public-data` (invoke the - right tool on the right data) - -```json -{ - "eval_id": "ml_cluster_penguins", - "conversation": [ + "rubrics": [ { - "invocation_id": "inv-cluster-01", - "user_content": { - "parts": [{"text": "Cluster the penguins in bigquery-public-data.ml_datasets.penguins into 3 groups based on their physical measurements."}], - "role": "user" - }, - "intermediate_data": { - "tool_uses": [ - {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "ml_datasets", "table_id": "penguins"}}, - {"name": "cluster_data", "args": {"project_id": "bigquery-public-data"}} - ], - "tool_responses": [], - "intermediate_responses": [] - }, - "creation_timestamp": 0.0 + "rubric_id": "consulted_ref", + "rubric_content": {"text_property": "The agent consulted the public datasets reference."} + }, + { + "rubric_id": "mentions_usa_names", + "rubric_content": {"text_property": "The response mentions the usa_names dataset and its columns."} } ], "creation_timestamp": 0.0 } ``` -**What gets checked automatically:** -- `tool_invocation_score`: Did the trace contain both - `get_table_info` and `cluster_data`? -- `tool_args_score`: Did `get_table_info` target - `(bigquery-public-data, ml_datasets, penguins)`? Did - `cluster_data` target `bigquery-public-data`? +### For a new skill -**What is NOT checked** (intentionally): -- The exact `feature_cols` or `num_clusters` the LLM chose -- The exact response wording -- The numeric clustering results +1. Create a skill directory under `skills/`: + ``` + skills/my-new-skill/ + ├── SKILL.md + ├── references/ + └── scripts/ + ``` -### Step 3: Validate +2. Register it in `agent.py`: + ```python + _SKILL_NAMES = [ + "bq-sql-analyst", + "my-new-skill", # ← add here + ] + ``` -```bash -BQ_EVAL_WRITE_MODE=protected \ - python -m benchmarks.bigquerybench.runner --filter ml_cluster_penguins -``` - -Expected output: - -``` -[1/1] ml_cluster_penguins - -> get_table_info(project_id='bigquery-public-data', dataset_id='ml_datasets', table_id='penguins') - -> cluster_data(project_id='bigquery-public-data') - tools=1.00 args=1.00 PASS -``` +3. Add eval cases with trace expectations + rubrics. -### Step 4: Commit +### Writing Good Rubrics -```bash -git add benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json -git commit -m "eval(bigquerybench): add clustering eval for cluster_data" -``` +**Do:** +- Assert observable behavior: "The agent loaded the skill before running the script." +- Assert output properties: "The response contains a table with columns name and count." +- Assert domain correctness: "The result includes the top 3 Shakespeare works." -**Only `bigquerybench_eval.json` changed.** No code changes. +**Don't:** +- Assert exact wording: "The response starts with 'Here are the results'." +- Assert implementation details: "The agent called execute_sql with SELECT DISTINCT." +- Use vague assertions: "The response is good." ### When Do You Need Code Changes? | Scenario | JSON | `metrics.py` | `runner.py` | `agent.py` | |----------|:----:|:------------:|:-----------:|:----------:| -| New eval case, existing tool | Yes | - | - | - | -| New tool, trace check is enough | Yes | - | - | - | -| New tool, need to check a non-key arg (e.g., `num_clusters`) | Yes | Yes (add arg to `_KEY_ARGS`) | - | - | -| New tool, need entirely new metric | Yes | Yes | Yes | - | -| Agent instruction or write-mode change | - | - | - | Yes | - -**Adding a checked arg** is a one-line change in `metrics.py`: - -```python -_KEY_ARGS = frozenset({ - "project_id", - "dataset_id", - "table_id", - "num_clusters", # ← add here -}) -``` +| New eval case, existing skill | Yes | - | - | - | +| New skill added to `skills/` | Yes | - | - | Yes (add to `_SKILL_NAMES`) | +| Change judge model or threshold | - | Yes | - | - | +| Need entirely new metric | Yes | Yes | Yes | - | +| Agent instruction change | - | - | - | Yes | ## Architecture ``` benchmarks/bigquerybench/ ├── __init__.py -├── agent.py # LlmAgent + BigQueryToolset (read-only default) -├── runner.py # Runs agent, collects trace, scores -├── metrics.py # tool_invocation_score + tool_args_score -└── eval_sets/ - └── bigquerybench_eval.json # 7 eval cases +├── agent.py # LlmAgent + BigQueryToolset + SkillToolset +├── runner.py # Runs agent, scores trace + rubrics +├── metrics.py # 3 metrics: trace (x2) + LLM-as-judge (x1) +├── eval_sets/ +│ └── bigquerybench_eval.json # 5 eval cases with rubrics +└── skills/ + └── bq-sql-analyst/ + ├── SKILL.md + ├── references/ + │ └── public-datasets.md + └── scripts/ + └── format_results.py + +tests/unittests/benchmarks/bigquerybench/ +└── test_metrics.py # 14 tests (trace + LLM judge + JSON validation) ``` ## Environment Variables @@ -313,7 +294,9 @@ benchmarks/bigquerybench/ | Symptom | Fix | |---------|-----| -| `403 Access Denied` | `gcloud auth application-default login` + enable BigQuery API | -| `tool_invocation_score = 0` | Agent didn't call expected tool — check agent instructions | -| `tool_args_score < 1.0` | Agent pointed at wrong dataset/table — check user query specificity | -| AI/ML tool fails | Set `BQ_EVAL_WRITE_MODE=protected` | +| `tool_invocation_score = 0` | Agent didn't call expected skill tool — check agent instructions | +| `tool_args_score < 1.0` | Agent targeted wrong skill or resource — check user query specificity | +| `adherence < 0.75` | Agent produced wrong output — review rubrics and skill instructions | +| Skill not found | Verify skill dir exists in `skills/` and name is in `_SKILL_NAMES` in `agent.py` | +| Judge LLM fails | Check `GOOGLE_API_KEY` or `GOOGLE_GENAI_USE_VERTEXAI` + `GOOGLE_CLOUD_PROJECT` | +| `load_skill_resource` fails | Check the `path` arg matches a real file under the skill dir | diff --git a/benchmarks/bigquerybench/agent.py b/benchmarks/bigquerybench/agent.py index 9071e9a581..819de5e792 100644 --- a/benchmarks/bigquerybench/agent.py +++ b/benchmarks/bigquerybench/agent.py @@ -15,17 +15,22 @@ """BigQueryBench evaluation agent. Uses BigQueryToolset with read-only defaults against BigQuery public -datasets. Override write_mode via BQ_EVAL_WRITE_MODE env var when -evaluating AI/ML tools (forecast, detect_anomalies, etc.). +datasets, and SkillToolset for skill-based workflows. Override +write_mode via BQ_EVAL_WRITE_MODE env var when evaluating AI/ML +tools (forecast, detect_anomalies, etc.). """ import os +import pathlib from google.adk.agents.llm_agent import LlmAgent +from google.adk.code_executors.unsafe_local_code_executor import UnsafeLocalCodeExecutor +from google.adk.skills import load_skill_from_dir from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsConfig from google.adk.tools.bigquery.bigquery_toolset import BigQueryToolset from google.adk.tools.bigquery.config import BigQueryToolConfig from google.adk.tools.bigquery.config import WriteMode +from google.adk.tools.skill_toolset import SkillToolset import google.auth _WRITE_MODE_MAP = { @@ -52,17 +57,32 @@ bigquery_tool_config=tool_config, ) +# ── Skill toolset ────────────────────────────────────────────────── +_SKILLS_DIR = pathlib.Path(__file__).parent / "skills" + +_SKILL_NAMES = [ + "bq-sql-analyst", +] + +_skills = [load_skill_from_dir(_SKILLS_DIR / name) for name in _SKILL_NAMES] + +skill_toolset = SkillToolset( + skills=_skills, + code_executor=UnsafeLocalCodeExecutor(), +) + root_agent = LlmAgent( model="gemini-2.5-flash", name="bigquerybench_agent", description=( "Agent for BigQuery data exploration, SQL execution, and" - " AI/ML operations against public datasets." + " AI/ML operations against public datasets. Also supports" + " skill-based workflows via SkillToolset." ), instruction="""\ -You are a data analyst with access to BigQuery tools. +You are a data analyst with access to BigQuery tools and skills. -Workflow: +Workflow for direct BigQuery queries: 1. Always explore the schema first: use list_dataset_ids, list_table_ids, and get_table_info to understand the data before writing any SQL. @@ -73,7 +93,17 @@ analyze_contribution) instead of raw SQL. 4. Present results clearly with column headers and values. +Workflow for skill-based tasks: +1. Use list_skills to discover available skills. +2. Use load_skill to read the skill's instructions. +3. Use load_skill_resource to examine references, sample data, + or templates from the skill. +4. Follow the skill's instructions — this may involve calling + BigQuery tools (get_table_info, execute_sql) or running + the skill's scripts via run_skill_script. +5. Present results clearly. + All public datasets are in project "bigquery-public-data". """, - tools=[bigquery_toolset], + tools=[bigquery_toolset, skill_toolset], ) diff --git a/benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json b/benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json index 1906baa502..4ecb405c31 100644 --- a/benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json +++ b/benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json @@ -1,20 +1,20 @@ { "eval_set_id": "bigquerybench-adk-v1", "name": "BigQueryBench ADK Evaluation", - "description": "Trace-based evaluation: verify the agent calls the correct BigQuery tools with the correct dataset/table args. No response text matching.", + "description": "Skill invocation evaluation: verify the agent calls the correct skill tools with the correct args, and follows skill instructions to produce correct results.", "eval_cases": [ { - "eval_id": "schema_list_datasets", + "eval_id": "skill_list_skills", "conversation": [ { - "invocation_id": "inv-schema-ds-01", + "invocation_id": "inv-skill-list-01", "user_content": { - "parts": [{"text": "What datasets are available in the bigquery-public-data project? List a few of them."}], + "parts": [{"text": "What skills are available for me to use?"}], "role": "user" }, "intermediate_data": { "tool_uses": [ - {"name": "list_dataset_ids", "args": {"project_id": "bigquery-public-data"}} + {"name": "list_skills", "args": {}} ], "tool_responses": [], "intermediate_responses": [] @@ -22,41 +22,30 @@ "creation_timestamp": 0.0 } ], - "creation_timestamp": 0.0 - }, - { - "eval_id": "schema_list_tables", - "conversation": [ + "rubrics": [ { - "invocation_id": "inv-schema-tbl-01", - "user_content": { - "parts": [{"text": "What tables exist in the usa_names dataset in bigquery-public-data?"}], - "role": "user" - }, - "intermediate_data": { - "tool_uses": [ - {"name": "list_table_ids", "args": {"project_id": "bigquery-public-data", "dataset_id": "usa_names"}} - ], - "tool_responses": [], - "intermediate_responses": [] - }, - "creation_timestamp": 0.0 + "rubric_id": "lists_bq_skill", + "rubric_content": {"text_property": "The response mentions the bq-sql-analyst skill by name."} + }, + { + "rubric_id": "includes_description", + "rubric_content": {"text_property": "The response includes a description of what the skill can do."} } ], "creation_timestamp": 0.0 }, { - "eval_id": "schema_get_table_info", + "eval_id": "skill_load_sql_analyst", "conversation": [ { - "invocation_id": "inv-schema-info-01", + "invocation_id": "inv-skill-load-01", "user_content": { - "parts": [{"text": "What columns and data types does the usa_1910_current table in the usa_names dataset have?"}], + "parts": [{"text": "Load the bq-sql-analyst skill and tell me what it can do."}], "role": "user" }, "intermediate_data": { "tool_uses": [ - {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "usa_names", "table_id": "usa_1910_current"}} + {"name": "load_skill", "args": {"name": "bq-sql-analyst"}} ], "tool_responses": [], "intermediate_responses": [] @@ -64,43 +53,31 @@ "creation_timestamp": 0.0 } ], - "creation_timestamp": 0.0 - }, - { - "eval_id": "sql_shakespeare_unique_words", - "conversation": [ + "rubrics": [ { - "invocation_id": "inv-sql-shk-01", - "user_content": { - "parts": [{"text": "Using the samples.shakespeare table in bigquery-public-data, which work (corpus) has the most unique words? Show the top 3 works by distinct word count."}], - "role": "user" - }, - "intermediate_data": { - "tool_uses": [ - {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "samples", "table_id": "shakespeare"}}, - {"name": "execute_sql", "args": {"project_id": "bigquery-public-data"}} - ], - "tool_responses": [], - "intermediate_responses": [] - }, - "creation_timestamp": 0.0 + "rubric_id": "describes_capabilities", + "rubric_content": {"text_property": "The response describes the skill's capabilities such as SQL analysis or data formatting."} + }, + { + "rubric_id": "mentions_scripts", + "rubric_content": {"text_property": "The response mentions available scripts or references provided by the skill."} } ], "creation_timestamp": 0.0 }, { - "eval_id": "sql_usa_names_top_2020", + "eval_id": "skill_load_reference", "conversation": [ { - "invocation_id": "inv-sql-names-01", + "invocation_id": "inv-skill-ref-01", "user_content": { - "parts": [{"text": "What were the top 5 most popular baby names in the year 2020 across all states? Use the usa_names.usa_1910_current table in bigquery-public-data."}], + "parts": [{"text": "Using the bq-sql-analyst skill, load the public datasets reference so I know what data is available."}], "role": "user" }, "intermediate_data": { "tool_uses": [ - {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "usa_names", "table_id": "usa_1910_current"}}, - {"name": "execute_sql", "args": {"project_id": "bigquery-public-data"}} + {"name": "load_skill", "args": {"name": "bq-sql-analyst"}}, + {"name": "load_skill_resource", "args": {"skill_name": "bq-sql-analyst", "path": "references/public-datasets.md"}} ], "tool_responses": [], "intermediate_responses": [] @@ -108,21 +85,31 @@ "creation_timestamp": 0.0 } ], + "rubrics": [ + { + "rubric_id": "shows_datasets", + "rubric_content": {"text_property": "The response contains information about BigQuery public datasets such as usa_names, samples, or austin_bikeshare."} + }, + { + "rubric_id": "loaded_skill_first", + "rubric_content": {"text_property": "The agent loaded the skill instructions before loading the resource."} + } + ], "creation_timestamp": 0.0 }, { - "eval_id": "sql_names_by_decade", + "eval_id": "skill_query_with_reference", "conversation": [ { - "invocation_id": "inv-sql-decade-01", + "invocation_id": "inv-skill-query-01", "user_content": { - "parts": [{"text": "How many distinct baby names were registered in each decade from 1950 to 2000 (inclusive)? Use the usa_names.usa_1910_current table in bigquery-public-data."}], + "parts": [{"text": "Use the bq-sql-analyst skill to help me find the top 3 Shakespeare works by unique word count. Load the skill and its reference first, then query the data."}], "role": "user" }, "intermediate_data": { "tool_uses": [ - {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "usa_names", "table_id": "usa_1910_current"}}, - {"name": "execute_sql", "args": {"project_id": "bigquery-public-data"}} + {"name": "load_skill", "args": {"name": "bq-sql-analyst"}}, + {"name": "load_skill_resource", "args": {"skill_name": "bq-sql-analyst", "path": "references/public-datasets.md"}} ], "tool_responses": [], "intermediate_responses": [] @@ -130,22 +117,35 @@ "creation_timestamp": 0.0 } ], + "rubrics": [ + { + "rubric_id": "consulted_reference", + "rubric_content": {"text_property": "The agent loaded the skill's public-datasets reference to identify the Shakespeare table before querying."} + }, + { + "rubric_id": "result_has_ranking", + "rubric_content": {"text_property": "The final response contains a ranking of Shakespeare works by unique word count, showing at least 3 works."} + }, + { + "rubric_id": "followed_workflow", + "rubric_content": {"text_property": "The agent followed the skill workflow: loaded the skill, then consulted the reference, then queried the data."} + } + ], "creation_timestamp": 0.0 }, { - "eval_id": "multi_step_explore_and_query", + "eval_id": "skill_run_format_script", "conversation": [ { - "invocation_id": "inv-multi-01", + "invocation_id": "inv-skill-script-01", "user_content": { - "parts": [{"text": "I want to analyze the Austin bikeshare data in bigquery-public-data. First explore what tables and columns are available, then tell me the top 5 busiest start stations by total trip count."}], + "parts": [{"text": "Use the bq-sql-analyst skill's format_results.py script to format a table with header 'name,count' and rows 'Alice,5;Bob,3;Carol,8'. Load the skill first."}], "role": "user" }, "intermediate_data": { "tool_uses": [ - {"name": "list_table_ids", "args": {"project_id": "bigquery-public-data", "dataset_id": "austin_bikeshare"}}, - {"name": "get_table_info", "args": {"project_id": "bigquery-public-data", "dataset_id": "austin_bikeshare", "table_id": "bikeshare_trips"}}, - {"name": "execute_sql", "args": {"project_id": "bigquery-public-data"}} + {"name": "load_skill", "args": {"name": "bq-sql-analyst"}}, + {"name": "run_skill_script", "args": {"skill_name": "bq-sql-analyst", "script_path": "scripts/format_results.py"}} ], "tool_responses": [], "intermediate_responses": [] @@ -153,6 +153,20 @@ "creation_timestamp": 0.0 } ], + "rubrics": [ + { + "rubric_id": "loaded_before_run", + "rubric_content": {"text_property": "The agent loaded the skill instructions before running the script."} + }, + { + "rubric_id": "output_has_table", + "rubric_content": {"text_property": "The output contains a formatted table with the data for Alice, Bob, and Carol."} + }, + { + "rubric_id": "has_columns", + "rubric_content": {"text_property": "The formatted output includes column headers 'name' and 'count'."} + } + ], "creation_timestamp": 0.0 } ] diff --git a/benchmarks/bigquerybench/metrics.py b/benchmarks/bigquerybench/metrics.py index 5e78a9114b..cfbc8de197 100644 --- a/benchmarks/bigquerybench/metrics.py +++ b/benchmarks/bigquerybench/metrics.py @@ -12,12 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Trace-based metrics for BigQueryBench evaluation. +"""Metrics for BigQueryBench skill evaluation. -Two metrics that check the agent's tool-call trace only (no response -text matching). This makes evaluation deterministic and easy to -maintain: just specify which tools should be called and with which -key arguments. +Three metrics: +1. tool_invocation_score — trace-based, checks tool names +2. tool_args_score — trace-based, checks key tool arguments +3. instruction_adherence_score — LLM-as-judge, checks rubrics def metric_fn( eval_metric: EvalMetric, @@ -25,24 +25,27 @@ def metric_fn( expected_invocations: Optional[list[Invocation]], conversation_scenario: Optional[ConversationScenario] = None, ) -> EvaluationResult - -Reference via dotted path in eval configs: - "benchmarks.bigquerybench.metrics.tool_invocation_score" - "benchmarks.bigquerybench.metrics.tool_args_score" """ from __future__ import annotations +import logging +import re from typing import Optional +import google.genai + from google.adk.evaluation.eval_case import ConversationScenario from google.adk.evaluation.eval_case import get_all_tool_calls from google.adk.evaluation.eval_case import Invocation from google.adk.evaluation.eval_metrics import EvalMetric from google.adk.evaluation.eval_metrics import EvalStatus +from google.adk.evaluation.eval_rubrics import Rubric from google.adk.evaluation.evaluator import EvaluationResult from google.adk.evaluation.evaluator import PerInvocationResult +logger = logging.getLogger(__name__) + def _get_tool_calls(invocations: list[Invocation]): """Yield (name, args_dict) for every tool call in the trace.""" @@ -82,11 +85,12 @@ def tool_invocation_score( expected_invocations: Optional[list[Invocation]], conversation_scenario: Optional[ConversationScenario] = None, ) -> EvaluationResult: - """Score = fraction of expected tool names present in the trace. + """Score = fraction of expected skill tool names present in the trace. - Checks that the agent called the right BigQuery functions (e.g. - ``get_table_info``, ``execute_sql``, ``forecast``). Order does - not matter; extra tool calls are ignored. + Checks that the agent called the right skill functions (e.g. + ``list_skills``, ``load_skill``, ``load_skill_resource``, + ``run_skill_script``). Order does not matter; extra tool calls + are ignored. Score = |expected_names ∩ actual_names| / |expected_names|. Pass threshold: 1.0 (all expected tools must be called). @@ -122,14 +126,14 @@ def tool_invocation_score( # ── Metric 2: correct args on key tool calls ───────────────────── -# Args that identify the *target data* — these are what we check. -# We intentionally skip volatile args like ``query`` (the exact SQL -# the LLM generates will vary) and only verify that the agent -# pointed at the right dataset / table / project. +# Args that identify the *target skill* — these are what we check. +# We only verify that the agent invoked the correct skill tools +# with the correct skill name, resource path, and script path. _KEY_ARGS = frozenset({ - "project_id", - "dataset_id", - "table_id", + "name", # load_skill(name=...) + "skill_name", # load_skill_resource / run_skill_script + "path", # load_skill_resource(path=...) + "script_path", # run_skill_script(script_path=...) }) @@ -141,11 +145,11 @@ def tool_args_score( ) -> EvaluationResult: """Score = fraction of expected (tool, key-arg) pairs matched. - For each expected tool call that has ``project_id``, ``dataset_id``, - or ``table_id`` in its args, check that the agent made a call to - the *same tool* with the *same value* for that arg. This verifies - the agent loaded the right reference data (correct dataset, correct - table) without caring about the exact SQL or response text. + For each expected skill tool call that has ``name``, + ``skill_name``, ``path``, or ``script_path`` in its args, check + that the agent made a call to the *same tool* with the *same + value* for that arg. This verifies the agent invoked the correct + skill with the correct resources. Score = matched_pairs / expected_pairs. Pass threshold: 1.0. If no key args exist in the expected trace, score is 1.0 (vacuous). @@ -197,3 +201,127 @@ def tool_args_score( status, ), ) + + +# ── Metric 3: instruction adherence (LLM-as-judge) ──────────────── + +_JUDGE_PROMPT = """\ +You are evaluating an AI agent's behavior. + +For EACH rubric listed below, respond in EXACTLY this format: + +Property: +Rationale: +Verdict: yes or no + +--- + +USER REQUEST: +{user_input} + +TOOL CALLS: +{tool_trace} + +AGENT FINAL RESPONSE: +{final_response} + +RUBRICS TO EVALUATE: +{rubrics_text} +""" + +_VERDICT_RE = re.compile(r"(?<=Verdict: )(.*)") + + +def _extract_text(content) -> str: + """Extract text from a genai Content object.""" + if content and content.parts: + return "\n".join(p.text for p in content.parts if p.text) + return "" + + +def _format_tool_trace(invocations: list[Invocation]) -> str: + """Format tool calls as readable text for the judge.""" + lines = [] + step = 1 + for inv in invocations: + if not inv.intermediate_data: + continue + for tc in get_all_tool_calls(inv.intermediate_data): + args_str = ", ".join(f"{k}={v!r}" for k, v in (tc.args or {}).items()) + lines.append(f"Step {step}: {tc.name}({args_str})") + step += 1 + return "\n".join(lines) if lines else "No tool calls." + + +async def instruction_adherence_score( + actual_invocations: list[Invocation], + rubrics: Optional[list[Rubric]], + judge_model: str = "gemini-2.5-flash", +) -> EvaluationResult: + """LLM-as-judge: check each rubric against the agent's output. + + For each rubric, prompts a judge model: + "Does the agent satisfy: ?" + Judge answers yes (1.0) or no (0.0) per rubric. + + Score = fraction of rubrics with "yes" verdict. + Pass threshold: 0.75. + """ + if not rubrics: + return EvaluationResult( + overall_score=1.0, + overall_eval_status=EvalStatus.PASSED, + ) + + # Extract conversation data from actual invocations. + user_input = "" + final_response = "" + for inv in actual_invocations: + text = _extract_text(inv.user_content) + if text: + user_input = text + text = _extract_text(inv.final_response) + if text: + final_response = text + + tool_trace = _format_tool_trace(actual_invocations) + + rubrics_text = "\n".join( + f"- {r.rubric_content.text_property}" + for r in rubrics + if r.rubric_content and r.rubric_content.text_property + ) + + prompt = _JUDGE_PROMPT.format( + user_input=user_input or "(empty)", + tool_trace=tool_trace, + final_response=final_response or "(No response)", + rubrics_text=rubrics_text, + ) + + try: + client = google.genai.Client() + response = await client.models.generate_content_async( + model=judge_model, + contents=prompt, + ) + response_text = response.text or "" + except Exception as e: + logger.error("Judge LLM call failed: %s", e) + return EvaluationResult( + overall_score=0.0, + overall_eval_status=EvalStatus.FAILED, + ) + + # Parse verdicts. + verdicts = _VERDICT_RE.findall(response_text) + yes_count = sum(1 for v in verdicts if v.strip().lower().startswith("yes")) + total = len(rubrics) + score = yes_count / total if total > 0 else 1.0 + + status = EvalStatus.PASSED if score >= 0.75 else EvalStatus.FAILED + + return EvaluationResult( + overall_score=score, + overall_eval_status=status, + ) diff --git a/benchmarks/bigquerybench/runner.py b/benchmarks/bigquerybench/runner.py index e21c6c4db1..aae6a85e32 100644 --- a/benchmarks/bigquerybench/runner.py +++ b/benchmarks/bigquerybench/runner.py @@ -12,15 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Trace-based BigQueryBench runner. +"""BigQueryBench evaluation runner. -Evaluates whether the agent calls the correct BigQuery tools with the -correct dataset/table arguments. No response text matching — only -the tool-call trace matters. +Evaluates skill invocation correctness (trace-based) and instruction +adherence (LLM-as-judge with rubrics). Usage: python -m benchmarks.bigquerybench.runner - python -m benchmarks.bigquerybench.runner --filter sql_shakespeare + python -m benchmarks.bigquerybench.runner --filter skill_load python -m benchmarks.bigquerybench.runner --num-runs 3 python -m benchmarks.bigquerybench.runner --dry-run @@ -53,6 +52,7 @@ from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.adk.utils.context_utils import Aclosing +from .metrics import instruction_adherence_score from .metrics import tool_args_score from .metrics import tool_invocation_score @@ -70,23 +70,36 @@ def load_eval_set(path: pathlib.Path) -> EvalSet: def print_header(): print() - print("=" * 65) - print(" BigQueryBench — Trace-Based Evaluation") - print("=" * 65) + print("=" * 72) + print(" BigQueryBench — Skill Evaluation") + print("=" * 72) print() +def _case_passed(scores: dict[str, float]) -> bool: + trace_ok = ( + scores.get("tool_invocation", 0.0) >= 1.0 + and scores.get("tool_args", 0.0) >= 1.0 + ) + adherence_ok = scores.get("adherence", 1.0) >= 0.75 + return trace_ok and adherence_ok + + def print_results_table(results: dict[str, dict[str, float]]): - print(f"{'Eval Case':<40} {'Tools':>7} {'Args':>7} {'Result':>8}") - print("-" * 65) + print( + f"{'Eval Case':<34} {'Tools':>6} {'Args':>6} {'Adhere':>7} {'Result':>7}" + ) + print("-" * 72) for case_id, scores in results.items(): - short_id = case_id[:39] + short_id = case_id[:33] tools = scores.get("tool_invocation", 0.0) args = scores.get("tool_args", 0.0) - passed = tools >= 1.0 and args >= 1.0 - mark = "PASS" if passed else "FAIL" - print(f"{short_id:<40} {tools:>7.2f} {args:>7.2f} {mark:>5}") - print("-" * 65) + adhere = scores.get("adherence", 1.0) + mark = "PASS" if _case_passed(scores) else "FAIL" + print( + f"{short_id:<34} {tools:>6.2f} {args:>6.2f} {adhere:>7.2f} {mark:>4}" + ) + print("-" * 72) def print_summary( @@ -94,28 +107,23 @@ def print_summary( num_cases: int, elapsed: float, ): - passed = sum( - 1 - for s in results.values() - if s.get("tool_invocation", 0.0) >= 1.0 and s.get("tool_args", 0.0) >= 1.0 - ) - avg_tools = sum( - s.get("tool_invocation", 0.0) for s in results.values() - ) / max(len(results), 1) - avg_args = sum(s.get("tool_args", 0.0) for s in results.values()) / max( - len(results), 1 - ) + n = max(len(results), 1) + passed = sum(1 for s in results.values() if _case_passed(s)) + avg_tools = sum(s.get("tool_invocation", 0.0) for s in results.values()) / n + avg_args = sum(s.get("tool_args", 0.0) for s in results.values()) / n + avg_adhere = sum(s.get("adherence", 1.0) for s in results.values()) / n pct = (passed / max(num_cases, 1)) * 100 print() - print("=" * 65) + print("=" * 72) print(" Summary") - print("=" * 65) + print("=" * 72) print(f" Cases: {passed}/{num_cases} ({pct:.1f}%)") print(f" Avg Tool Match: {avg_tools:.2f}") print(f" Avg Args Match: {avg_args:.2f}") + print(f" Avg Adherence: {avg_adhere:.2f}") print(f" Elapsed: {elapsed:.1f}s") - print("=" * 65) + print("=" * 72) async def run_single_eval_case(root_agent, eval_case) -> list[Invocation]: @@ -170,20 +178,35 @@ async def run_single_eval_case(root_agent, eval_case) -> list[Invocation]: return EvaluationGenerator.convert_events_to_eval_invocations(events) -def score_invocations( +async def score_invocations( actual: list[Invocation], expected: Optional[list[Invocation]], + rubrics=None, ) -> dict[str, float]: metric = EvalMetric(metric_name="bigquerybench") r1 = tool_invocation_score(metric, actual, expected) r2 = tool_args_score(metric, actual, expected) - return { + scores = { "tool_invocation": r1.overall_score or 0.0, "tool_args": r2.overall_score or 0.0, } + if rubrics: + r3 = await instruction_adherence_score(actual, rubrics) + scores["adherence"] = r3.overall_score or 0.0 + + return scores + + +_TRACE_ARGS = frozenset({ + "name", + "skill_name", + "path", + "script_path", +}) + def _print_trace(actual_invocations: list[Invocation]): """Print the tool-call trace for debugging.""" @@ -192,9 +215,7 @@ def _print_trace(actual_invocations: list[Invocation]): for inv in actual_invocations: for tc in get_all_tool_calls(inv.intermediate_data): args_summary = ", ".join( - f"{k}={v!r}" - for k, v in (tc.args or {}).items() - if k in ("project_id", "dataset_id", "table_id") + f"{k}={v!r}" for k, v in (tc.args or {}).items() if k in _TRACE_ARGS ) print(f" -> {tc.name}({args_summary})") @@ -231,14 +252,19 @@ async def run_evaluation( actual = await run_single_eval_case(root_agent, eval_case) _print_trace(actual) - scores = score_invocations(actual, eval_case.conversation) + scores = await score_invocations( + actual, eval_case.conversation, eval_case.rubrics + ) run_scores.append(scores) tools = scores["tool_invocation"] args = scores["tool_args"] - passed = tools >= 1.0 and args >= 1.0 - mark = "PASS" if passed else "FAIL" - print(f" tools={tools:.2f} args={args:.2f} {mark}") + adhere = scores.get("adherence", 1.0) + mark = "PASS" if _case_passed(scores) else "FAIL" + parts = [f"tools={tools:.2f}", f"args={args:.2f}"] + if "adherence" in scores: + parts.append(f"adherence={adhere:.2f}") + print(f" {' '.join(parts)} {mark}") except Exception as e: logger.error("Error running %s: %s", eval_id, e) @@ -246,8 +272,10 @@ async def run_evaluation( run_scores.append({"tool_invocation": 0.0, "tool_args": 0.0}) avg: dict[str, float] = {} - for key in ("tool_invocation", "tool_args"): - avg[key] = sum(s[key] for s in run_scores) / len(run_scores) + for key in ("tool_invocation", "tool_args", "adherence"): + vals = [s[key] for s in run_scores if key in s] + if vals: + avg[key] = sum(vals) / len(vals) results[eval_id] = avg return results @@ -289,7 +317,9 @@ def main(): tools = [ t.name for t in case.conversation[0].intermediate_data.tool_uses or [] ] - print(f" {case.eval_id}: {' -> '.join(tools)}") + n_rubrics = len(case.rubrics) if case.rubrics else 0 + rubric_info = f" ({n_rubrics} rubrics)" if n_rubrics else "" + print(f" {case.eval_id}: {' -> '.join(tools)}{rubric_info}") print("\nJSON valid.") sys.exit(0) @@ -311,10 +341,7 @@ def main(): print_results_table(results) print_summary(results, num_cases, elapsed) - all_pass = all( - s.get("tool_invocation", 0.0) >= 1.0 and s.get("tool_args", 0.0) >= 1.0 - for s in results.values() - ) + all_pass = all(_case_passed(s) for s in results.values()) sys.exit(0 if all_pass else 1) diff --git a/benchmarks/bigquerybench/skills/bq-sql-analyst/SKILL.md b/benchmarks/bigquerybench/skills/bq-sql-analyst/SKILL.md new file mode 100644 index 0000000000..ec7f8f5bac --- /dev/null +++ b/benchmarks/bigquerybench/skills/bq-sql-analyst/SKILL.md @@ -0,0 +1,34 @@ +--- +name: bq-sql-analyst +description: Analyze BigQuery public datasets with SQL — explore schemas, write queries, and format results. +--- + +# BigQuery SQL Analyst Skill + +Explore and analyze BigQuery public datasets by examining schemas, +writing SQL queries, and presenting formatted results. + +## Available Scripts + +### `format_results.py` + +Formats raw query output as a readable markdown table. + +**Usage**: `run_skill_script(skill_name="bq-sql-analyst", script_path="scripts/format_results.py", args={"header": "name,count", "rows": "Alice,5;Bob,3"})` + +Arguments: +- `header`: Comma-separated column names +- `rows`: Semicolon-separated rows, each with comma-separated values +- `title` (optional): Table title + +## References + +- [public-datasets.md](./references/public-datasets.md) — Commonly used BigQuery public datasets and their schemas + +## Workflow + +1. Use `load_skill` to read these instructions. +2. Use `load_skill_resource` to review the public datasets reference. +3. Use BigQuery tools (`get_table_info`, `execute_sql`) to explore and query data. +4. Optionally use `run_skill_script` to format results. +5. Present findings to the user. diff --git a/benchmarks/bigquerybench/skills/bq-sql-analyst/references/public-datasets.md b/benchmarks/bigquerybench/skills/bq-sql-analyst/references/public-datasets.md new file mode 100644 index 0000000000..bc08041c53 --- /dev/null +++ b/benchmarks/bigquerybench/skills/bq-sql-analyst/references/public-datasets.md @@ -0,0 +1,34 @@ +# BigQuery Public Datasets Reference + +Commonly used datasets in the `bigquery-public-data` project. + +## usa_names + +Baby name data from US Social Security Administration. + +| Table | Description | +|-------|-------------| +| `usa_1910_current` | Names, gender, state, year, count from 1910 to present | + +Key columns: `name` (STRING), `gender` (STRING), `state` (STRING), `year` (INT64), `number` (INT64) + +## samples + +Sample datasets provided by BigQuery. + +| Table | Description | +|-------|-------------| +| `shakespeare` | Every word in every Shakespeare work with word count | + +Key columns: `word` (STRING), `word_count` (INT64), `corpus` (STRING), `corpus_date` (INT64) + +## austin_bikeshare + +Austin B-cycle bikeshare trip data. + +| Table | Description | +|-------|-------------| +| `bikeshare_trips` | Individual trip records with start/end stations and times | +| `bikeshare_stations` | Station locations and metadata | + +Key columns (trips): `trip_id`, `start_station_name`, `end_station_name`, `duration_minutes`, `start_time` diff --git a/benchmarks/bigquerybench/skills/bq-sql-analyst/scripts/format_results.py b/benchmarks/bigquerybench/skills/bq-sql-analyst/scripts/format_results.py new file mode 100644 index 0000000000..86fc250e98 --- /dev/null +++ b/benchmarks/bigquerybench/skills/bq-sql-analyst/scripts/format_results.py @@ -0,0 +1,51 @@ +"""Format query results as a markdown table.""" + +import sys + + +def parse_args(args): + params = {} + i = 0 + while i < len(args): + if args[i].startswith("--") and i + 1 < len(args): + params[args[i][2:]] = args[i + 1] + i += 2 + elif "=" in args[i]: + key, value = args[i].split("=", 1) + params[key.lstrip("-")] = value + i += 1 + else: + i += 1 + return params + + +def main(): + params = parse_args(sys.argv[1:]) + header = params.get("header", "") + rows_str = params.get("rows", "") + title = params.get("title", "") + + if not header: + print("Error: --header is required") + sys.exit(1) + + cols = [c.strip() for c in header.split(",")] + + if title: + print(f"## {title}") + print() + + print("| " + " | ".join(cols) + " |") + print("| " + " | ".join("---" for _ in cols) + " |") + + if rows_str: + for row in rows_str.split(";"): + vals = [v.strip() for v in row.split(",")] + # Pad if needed + while len(vals) < len(cols): + vals.append("") + print("| " + " | ".join(vals[: len(cols)]) + " |") + + +if __name__ == "__main__": + main() diff --git a/tests/unittests/benchmarks/__init__.py b/tests/unittests/benchmarks/__init__.py new file mode 100644 index 0000000000..58d482ea38 --- /dev/null +++ b/tests/unittests/benchmarks/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unittests/benchmarks/bigquerybench/__init__.py b/tests/unittests/benchmarks/bigquerybench/__init__.py new file mode 100644 index 0000000000..58d482ea38 --- /dev/null +++ b/tests/unittests/benchmarks/bigquerybench/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unittests/benchmarks/bigquerybench/test_metrics.py b/tests/unittests/benchmarks/bigquerybench/test_metrics.py new file mode 100644 index 0000000000..3dd90032ee --- /dev/null +++ b/tests/unittests/benchmarks/bigquerybench/test_metrics.py @@ -0,0 +1,344 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for BigQueryBench metrics including LLM-as-judge adherence.""" + +import json +from unittest import mock + +from google.adk.evaluation.eval_case import IntermediateData +from google.adk.evaluation.eval_case import Invocation +from google.adk.evaluation.eval_metrics import EvalMetric +from google.adk.evaluation.eval_metrics import EvalStatus +from google.adk.evaluation.eval_rubrics import Rubric +from google.adk.evaluation.eval_rubrics import RubricContent +from google.adk.evaluation.eval_set import EvalSet +from google.genai import types as genai_types +import pytest + +from benchmarks.bigquerybench.metrics import instruction_adherence_score +from benchmarks.bigquerybench.metrics import tool_args_score +from benchmarks.bigquerybench.metrics import tool_invocation_score + + +def _make_content(text: str) -> genai_types.Content: + return genai_types.Content(parts=[genai_types.Part(text=text)], role="user") + + +def _make_invocation( + tool_calls: list[dict], + user_text: str = "test", + final_text: str = "", +) -> Invocation: + tool_uses = [] + for tc in tool_calls: + tool_uses.append( + genai_types.FunctionCall(name=tc["name"], args=tc.get("args", {})) + ) + inv = Invocation( + invocation_id="inv-test", + user_content=_make_content(user_text), + intermediate_data=IntermediateData(tool_uses=tool_uses), + ) + if final_text: + inv.final_response = genai_types.Content( + parts=[genai_types.Part(text=final_text)], role="model" + ) + return inv + + +# ── tool_invocation_score tests ──────────────────────────────────── + + +class TestToolInvocationScore: + + def test_all_expected_tools_present(self): + metric = EvalMetric(metric_name="test") + actual = [ + _make_invocation([ + {"name": "load_skill", "args": {"name": "bq-sql-analyst"}}, + {"name": "load_skill_resource", "args": {"skill_name": "x"}}, + ]) + ] + expected = [ + _make_invocation([ + {"name": "load_skill", "args": {"name": "bq-sql-analyst"}}, + {"name": "load_skill_resource", "args": {"skill_name": "x"}}, + ]) + ] + result = tool_invocation_score(metric, actual, expected) + assert result.overall_score == 1.0 + assert result.overall_eval_status == EvalStatus.PASSED + + def test_missing_expected_tool(self): + metric = EvalMetric(metric_name="test") + actual = [ + _make_invocation([ + {"name": "load_skill", "args": {"name": "bq-sql-analyst"}}, + ]) + ] + expected = [ + _make_invocation([ + {"name": "load_skill", "args": {"name": "bq-sql-analyst"}}, + {"name": "run_skill_script", "args": {"skill_name": "x"}}, + ]) + ] + result = tool_invocation_score(metric, actual, expected) + assert result.overall_score == 0.5 + assert result.overall_eval_status == EvalStatus.FAILED + + def test_extra_tools_are_ok(self): + metric = EvalMetric(metric_name="test") + actual = [ + _make_invocation([ + {"name": "list_skills"}, + {"name": "load_skill", "args": {"name": "bq-sql-analyst"}}, + {"name": "execute_sql", "args": {"project_id": "x"}}, + ]) + ] + expected = [ + _make_invocation([ + {"name": "load_skill", "args": {"name": "bq-sql-analyst"}}, + ]) + ] + result = tool_invocation_score(metric, actual, expected) + assert result.overall_score == 1.0 + + def test_empty_expected(self): + metric = EvalMetric(metric_name="test") + result = tool_invocation_score(metric, [], None) + assert result.overall_score == 1.0 + + +# ── tool_args_score tests ───────────────────────────────────────── + + +class TestToolArgsScore: + + def test_all_args_match(self): + metric = EvalMetric(metric_name="test") + actual = [ + _make_invocation([ + {"name": "load_skill", "args": {"name": "bq-sql-analyst"}}, + { + "name": "load_skill_resource", + "args": { + "skill_name": "bq-sql-analyst", + "path": "references/public-datasets.md", + }, + }, + ]) + ] + expected = [ + _make_invocation([ + {"name": "load_skill", "args": {"name": "bq-sql-analyst"}}, + { + "name": "load_skill_resource", + "args": { + "skill_name": "bq-sql-analyst", + "path": "references/public-datasets.md", + }, + }, + ]) + ] + result = tool_args_score(metric, actual, expected) + assert result.overall_score == 1.0 + + def test_wrong_skill_name(self): + metric = EvalMetric(metric_name="test") + actual = [ + _make_invocation([ + {"name": "load_skill", "args": {"name": "wrong-skill"}}, + ]) + ] + expected = [ + _make_invocation([ + {"name": "load_skill", "args": {"name": "bq-sql-analyst"}}, + ]) + ] + result = tool_args_score(metric, actual, expected) + assert result.overall_score == 0.0 + + def test_wrong_path(self): + metric = EvalMetric(metric_name="test") + actual = [ + _make_invocation([ + { + "name": "load_skill_resource", + "args": { + "skill_name": "bq-sql-analyst", + "path": "references/wrong.md", + }, + }, + ]) + ] + expected = [ + _make_invocation([ + { + "name": "load_skill_resource", + "args": { + "skill_name": "bq-sql-analyst", + "path": "references/public-datasets.md", + }, + }, + ]) + ] + result = tool_args_score(metric, actual, expected) + # skill_name matches (1/2), path doesn't (0/2) -> 1/2 = 0.5 + assert result.overall_score == 0.5 + + def test_no_key_args_in_expected(self): + metric = EvalMetric(metric_name="test") + actual = [_make_invocation([{"name": "list_skills"}])] + expected = [_make_invocation([{"name": "list_skills"}])] + result = tool_args_score(metric, actual, expected) + assert result.overall_score == 1.0 + + +# ── instruction_adherence_score tests ───────────────────────────── + + +def _make_rubric(rubric_id: str, text: str) -> Rubric: + return Rubric( + rubric_id=rubric_id, + rubric_content=RubricContent(text_property=text), + ) + + +class TestInstructionAdherenceScore: + + @pytest.mark.asyncio + async def test_no_rubrics_returns_pass(self): + actual = [_make_invocation([{"name": "list_skills"}])] + result = await instruction_adherence_score(actual, None) + assert result.overall_score == 1.0 + assert result.overall_eval_status == EvalStatus.PASSED + + @pytest.mark.asyncio + async def test_all_rubrics_pass(self): + rubrics = [ + _make_rubric("r1", "The agent listed skills."), + _make_rubric("r2", "The response mentions bq-sql-analyst."), + ] + actual = [ + _make_invocation( + [{"name": "list_skills"}], + user_text="What skills are available?", + final_text="Available skills: bq-sql-analyst.", + ) + ] + + mock_response = mock.MagicMock() + mock_response.text = ( + "Property: The agent listed skills.\n" + "Rationale: The agent called list_skills.\n" + "Verdict: yes\n\n" + "Property: The response mentions bq-sql-analyst.\n" + "Rationale: The response explicitly names bq-sql-analyst.\n" + "Verdict: yes\n" + ) + + mock_models = mock.AsyncMock() + mock_models.generate_content_async.return_value = mock_response + mock_client = mock.MagicMock() + mock_client.models = mock_models + + with mock.patch("google.genai.Client", return_value=mock_client): + result = await instruction_adherence_score(actual, rubrics) + + assert result.overall_score == 1.0 + assert result.overall_eval_status == EvalStatus.PASSED + + @pytest.mark.asyncio + async def test_some_rubrics_fail(self): + rubrics = [ + _make_rubric("r1", "The agent loaded the skill."), + _make_rubric("r2", "The result has a table."), + ] + actual = [ + _make_invocation( + [{"name": "load_skill", "args": {"name": "bq-sql-analyst"}}], + final_text="Skill loaded.", + ) + ] + + mock_response = mock.MagicMock() + mock_response.text = ( + "Property: The agent loaded the skill.\n" + "Rationale: load_skill was called.\n" + "Verdict: yes\n\n" + "Property: The result has a table.\n" + "Rationale: No table in the response.\n" + "Verdict: no\n" + ) + + mock_models = mock.AsyncMock() + mock_models.generate_content_async.return_value = mock_response + mock_client = mock.MagicMock() + mock_client.models = mock_models + + with mock.patch("google.genai.Client", return_value=mock_client): + result = await instruction_adherence_score(actual, rubrics) + + assert result.overall_score == 0.5 + assert result.overall_eval_status == EvalStatus.FAILED + + @pytest.mark.asyncio + async def test_llm_call_failure(self): + rubrics = [_make_rubric("r1", "Some assertion.")] + actual = [_make_invocation([{"name": "list_skills"}])] + + with mock.patch("google.genai.Client", side_effect=Exception("API error")): + result = await instruction_adherence_score(actual, rubrics) + + assert result.overall_score == 0.0 + assert result.overall_eval_status == EvalStatus.FAILED + + +# ── eval JSON validation ────────────────────────────────────────── + + +class TestEvalSetValidation: + + def test_eval_json_parses(self): + import pathlib + + eval_path = ( + pathlib.Path(__file__).parent.parent.parent.parent.parent + / "benchmarks" + / "bigquerybench" + / "eval_sets" + / "bigquerybench_eval.json" + ) + with open(eval_path) as f: + data = json.load(f) + es = EvalSet.model_validate(data) + assert len(es.eval_cases) == 5 + + def test_all_cases_have_rubrics(self): + import pathlib + + eval_path = ( + pathlib.Path(__file__).parent.parent.parent.parent.parent + / "benchmarks" + / "bigquerybench" + / "eval_sets" + / "bigquerybench_eval.json" + ) + with open(eval_path) as f: + data = json.load(f) + es = EvalSet.model_validate(data) + for case in es.eval_cases: + assert case.rubrics is not None, f"{case.eval_id} missing rubrics" + assert len(case.rubrics) >= 2, f"{case.eval_id} should have >= 2 rubrics" From 76912e5bdb104137af7e5ba7052c898314f337f3 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 09:10:03 -0800 Subject: [PATCH 46/53] feat(bigquerybench): add Vertex AI API key auth, retry backoff, and optional BigQuery toolset MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Agent model uses Vertex AI API key (GOOGLE_CLOUD_API_KEY) via custom _VertexGemini subclass with HttpRetryOptions (5 attempts, 2x exponential backoff) - LLM judge uses vertexai=True + api_key with manual retry backoff on 429 rate limits (5 attempts, 2s→4s→8s→16s) - BigQuery toolset is now optional — gracefully skipped when GCP project/ADC is unavailable, enabling skill-only evaluation - Agent instruction adapts to available tools (skill-only vs full) - Updated eval case skill_query_with_reference to be skill-only - Updated test mocks for client.aio.models.generate_content() API - E2E verified: 5/5 cases pass (100%), all metrics at 1.00 Co-Authored-By: Claude Opus 4.6 --- benchmarks/bigquerybench/README.md | 57 +++++-- benchmarks/bigquerybench/agent.py | 153 +++++++++++++----- .../eval_sets/bigquerybench_eval.json | 10 +- benchmarks/bigquerybench/metrics.py | 60 +++++-- .../benchmarks/bigquerybench/test_metrics.py | 12 +- 5 files changed, 216 insertions(+), 76 deletions(-) diff --git a/benchmarks/bigquerybench/README.md b/benchmarks/bigquerybench/README.md index aadabb022b..c10dcc8ab4 100644 --- a/benchmarks/bigquerybench/README.md +++ b/benchmarks/bigquerybench/README.md @@ -17,9 +17,12 @@ them easy to write and immune to exact-wording variance. ## Quick Start ```bash -# Prerequisites: GCP project with BigQuery API + ADC configured +# Skill-only mode (no BigQuery credentials needed): +export GOOGLE_CLOUD_API_KEY=your-vertex-ai-api-key + +# Full mode (BigQuery + skills — requires ADC + project): export GOOGLE_CLOUD_PROJECT=your-project-id -export GOOGLE_GENAI_USE_VERTEXAI=1 +export GOOGLE_CLOUD_API_KEY=your-vertex-ai-api-key # Run all eval cases python -m benchmarks.bigquerybench.runner @@ -137,7 +140,7 @@ trace + final response) and answers yes/no per rubric. | `skill_list_skills` | `list_skills()` | Lists bq-sql-analyst; includes description | | `skill_load_sql_analyst` | `load_skill(name=bq-sql-analyst)` | Describes capabilities; mentions scripts | | `skill_load_reference` | `load_skill` → `load_skill_resource` | Shows datasets; loaded skill first | -| `skill_query_with_reference` | `load_skill` → `load_skill_resource` | Consulted reference; has ranking; followed workflow | +| `skill_query_with_reference` | `load_skill` → `load_skill_resource` | Consulted reference; summarizes datasets; followed workflow | | `skill_run_format_script` | `load_skill` → `run_skill_script` | Loaded before run; has table; has columns | ## Example Output @@ -155,23 +158,38 @@ trace + final response) and answers yes/no per rubric. -> load_skill(name='bq-sql-analyst') tools=1.00 args=1.00 adherence=1.00 PASS +[3/5] skill_load_reference + -> load_skill(name='bq-sql-analyst') + -> load_skill_resource(skill_name='bq-sql-analyst', path='references/public-datasets.md') + tools=1.00 args=1.00 adherence=1.00 PASS + +[4/5] skill_query_with_reference + -> load_skill(name='bq-sql-analyst') + -> load_skill_resource(skill_name='bq-sql-analyst', path='references/public-datasets.md') + tools=1.00 args=1.00 adherence=1.00 PASS + +[5/5] skill_run_format_script + -> load_skill(name='bq-sql-analyst') + -> run_skill_script(skill_name='bq-sql-analyst', script_path='scripts/format_results.py') + tools=1.00 args=1.00 adherence=1.00 PASS + Eval Case Tools Args Adhere Result ------------------------------------------------------------------------ skill_list_skills 1.00 1.00 1.00 PASS skill_load_sql_analyst 1.00 1.00 1.00 PASS skill_load_reference 1.00 1.00 1.00 PASS -skill_query_with_reference 1.00 1.00 0.67 FAIL +skill_query_with_reference 1.00 1.00 1.00 PASS skill_run_format_script 1.00 1.00 1.00 PASS ------------------------------------------------------------------------ ======================================================================== Summary ======================================================================== - Cases: 4/5 (80.0%) + Cases: 5/5 (100.0%) Avg Tool Match: 1.00 Avg Args Match: 1.00 - Avg Adherence: 0.93 - Elapsed: 42.1s + Avg Adherence: 1.00 + Elapsed: 488.8s ======================================================================== ``` @@ -281,15 +299,31 @@ tests/unittests/benchmarks/bigquerybench/ └── test_metrics.py # 14 tests (trace + LLM judge + JSON validation) ``` +## Retry Backoff + +Both the agent model and the LLM judge use exponential backoff with +retry on 429 (rate limit) errors: + +- **Agent model**: 5 attempts, 2s initial delay, 2x exponential + backoff (via `HttpRetryOptions`) +- **LLM judge**: 5 attempts, 2s → 4s → 8s → 16s manual backoff, + plus 3 HTTP-level retries per attempt + ## Environment Variables | Variable | Required | Description | |----------|----------|-------------| -| `GOOGLE_CLOUD_PROJECT` | Yes | GCP project for BigQuery API | -| `GOOGLE_GENAI_USE_VERTEXAI` | Conditional | `1` for Vertex AI backend | -| `GOOGLE_API_KEY` | Conditional | API key for AI Studio backend | +| `GOOGLE_CLOUD_API_KEY` | Yes | Vertex AI API key for agent model + judge | +| `GOOGLE_CLOUD_PROJECT` | No | GCP project for BigQuery API (enables BigQuery toolset) | | `BQ_EVAL_WRITE_MODE` | No | `blocked` (default) / `protected` / `allowed` | +**Two modes:** +- **Skill-only** (default): Set `GOOGLE_CLOUD_API_KEY` only. + BigQuery toolset is skipped; all 5 skill eval cases run. +- **Full mode**: Set both `GOOGLE_CLOUD_API_KEY` and + `GOOGLE_CLOUD_PROJECT` (+ ADC configured). BigQuery toolset + is enabled alongside skills. + ## Troubleshooting | Symptom | Fix | @@ -297,6 +331,7 @@ tests/unittests/benchmarks/bigquerybench/ | `tool_invocation_score = 0` | Agent didn't call expected skill tool — check agent instructions | | `tool_args_score < 1.0` | Agent targeted wrong skill or resource — check user query specificity | | `adherence < 0.75` | Agent produced wrong output — review rubrics and skill instructions | +| 429 RESOURCE_EXHAUSTED | Rate limit — retry backoff handles this automatically; wait and retry | | Skill not found | Verify skill dir exists in `skills/` and name is in `_SKILL_NAMES` in `agent.py` | -| Judge LLM fails | Check `GOOGLE_API_KEY` or `GOOGLE_GENAI_USE_VERTEXAI` + `GOOGLE_CLOUD_PROJECT` | +| Judge LLM fails | Check `GOOGLE_CLOUD_API_KEY` is set correctly | | `load_skill_resource` fails | Check the `path` arg matches a real file under the skill dir | diff --git a/benchmarks/bigquerybench/agent.py b/benchmarks/bigquerybench/agent.py index 819de5e792..1e4da954c0 100644 --- a/benchmarks/bigquerybench/agent.py +++ b/benchmarks/bigquerybench/agent.py @@ -20,42 +20,60 @@ tools (forecast, detect_anomalies, etc.). """ +from functools import cached_property +import logging import os import pathlib from google.adk.agents.llm_agent import LlmAgent from google.adk.code_executors.unsafe_local_code_executor import UnsafeLocalCodeExecutor +from google.adk.models.google_llm import Gemini from google.adk.skills import load_skill_from_dir -from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsConfig -from google.adk.tools.bigquery.bigquery_toolset import BigQueryToolset -from google.adk.tools.bigquery.config import BigQueryToolConfig -from google.adk.tools.bigquery.config import WriteMode from google.adk.tools.skill_toolset import SkillToolset -import google.auth - -_WRITE_MODE_MAP = { - "blocked": WriteMode.BLOCKED, - "protected": WriteMode.PROTECTED, - "allowed": WriteMode.ALLOWED, -} - -_write_mode_str = os.environ.get("BQ_EVAL_WRITE_MODE", "blocked").lower() -_write_mode = _WRITE_MODE_MAP.get(_write_mode_str, WriteMode.BLOCKED) - -application_default_credentials, _ = google.auth.default() -credentials_config = BigQueryCredentialsConfig( - credentials=application_default_credentials, -) - -tool_config = BigQueryToolConfig( - write_mode=_write_mode, - max_query_result_rows=50, -) - -bigquery_toolset = BigQueryToolset( - credentials_config=credentials_config, - bigquery_tool_config=tool_config, -) +from google.genai import types as genai_types + +logger = logging.getLogger(__name__) + +# ── BigQuery toolset (optional — requires ADC) ─────────────────── +bigquery_toolset = None +try: + from google.adk.tools.bigquery.bigquery_credentials import BigQueryCredentialsConfig + from google.adk.tools.bigquery.bigquery_toolset import BigQueryToolset + from google.adk.tools.bigquery.config import BigQueryToolConfig + from google.adk.tools.bigquery.config import WriteMode + import google.auth + + _WRITE_MODE_MAP = { + "blocked": WriteMode.BLOCKED, + "protected": WriteMode.PROTECTED, + "allowed": WriteMode.ALLOWED, + } + + _write_mode_str = os.environ.get("BQ_EVAL_WRITE_MODE", "blocked").lower() + _write_mode = _WRITE_MODE_MAP.get(_write_mode_str, WriteMode.BLOCKED) + + application_default_credentials, project = google.auth.default() + if not project and not os.environ.get("GOOGLE_CLOUD_PROJECT"): + raise EnvironmentError("No GCP project found. Set GOOGLE_CLOUD_PROJECT.") + credentials_config = BigQueryCredentialsConfig( + credentials=application_default_credentials, + ) + + tool_config = BigQueryToolConfig( + write_mode=_write_mode, + max_query_result_rows=50, + ) + + bigquery_toolset = BigQueryToolset( + credentials_config=credentials_config, + bigquery_tool_config=tool_config, + ) +except Exception as e: + logger.warning( + "BigQuery toolset unavailable (%s). " + "Skill-only evaluation will still work.", + e, + ) # ── Skill toolset ────────────────────────────────────────────────── _SKILLS_DIR = pathlib.Path(__file__).parent / "skills" @@ -71,15 +89,44 @@ code_executor=UnsafeLocalCodeExecutor(), ) -root_agent = LlmAgent( - model="gemini-2.5-flash", - name="bigquerybench_agent", - description=( - "Agent for BigQuery data exploration, SQL execution, and" - " AI/ML operations against public datasets. Also supports" - " skill-based workflows via SkillToolset." - ), - instruction="""\ + +# ── Model (Vertex AI + API key) ────────────────────────────────── +class _VertexGemini(Gemini): + """Gemini subclass that uses vertexai=True with an API key.""" + + @cached_property + def api_client(self): + from google.genai import Client + + return Client( + vertexai=True, + api_key=os.environ.get("GOOGLE_CLOUD_API_KEY"), + http_options=genai_types.HttpOptions( + headers=self._tracking_headers(), + retry_options=self.retry_options, + base_url=self.base_url, + ), + ) + + +_SKILL_INSTRUCTION = """\ +You are a data analyst with access to skills. + +Workflow for skill-based tasks: +1. Use list_skills to discover available skills. +2. Use load_skill to read the skill's instructions. +3. Use load_skill_resource to examine references, sample data, + or templates from the skill. +4. Follow the skill's instructions — this may involve running + the skill's scripts via run_skill_script. +5. Present results clearly. + +IMPORTANT: Only use the tools available to you (list_skills, +load_skill, load_skill_resource, run_skill_script). Do NOT +attempt to call tools that are not listed. +""" + +_BQ_INSTRUCTION = """\ You are a data analyst with access to BigQuery tools and skills. Workflow for direct BigQuery queries: @@ -104,6 +151,32 @@ 5. Present results clearly. All public datasets are in project "bigquery-public-data". -""", - tools=[bigquery_toolset, skill_toolset], +""" + +_INSTRUCTION = _BQ_INSTRUCTION if bigquery_toolset else _SKILL_INSTRUCTION + +_api_key = os.environ.get("GOOGLE_CLOUD_API_KEY") +_model = ( + _VertexGemini( + model="gemini-3-flash-preview", + retry_options=genai_types.HttpRetryOptions( + initialDelay=2, + expBase=2, + attempts=5, + ), + ) + if _api_key + else "gemini-3-flash-preview" +) + +root_agent = LlmAgent( + model=_model, + name="bigquerybench_agent", + description=( + "Agent for BigQuery data exploration, SQL execution, and" + " AI/ML operations against public datasets. Also supports" + " skill-based workflows via SkillToolset." + ), + instruction=_INSTRUCTION, + tools=[t for t in [bigquery_toolset, skill_toolset] if t], ) diff --git a/benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json b/benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json index 4ecb405c31..2fc87f96bd 100644 --- a/benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json +++ b/benchmarks/bigquerybench/eval_sets/bigquerybench_eval.json @@ -103,7 +103,7 @@ { "invocation_id": "inv-skill-query-01", "user_content": { - "parts": [{"text": "Use the bq-sql-analyst skill to help me find the top 3 Shakespeare works by unique word count. Load the skill and its reference first, then query the data."}], + "parts": [{"text": "Load the bq-sql-analyst skill and then load its public-datasets reference. Summarize the key datasets listed in that reference document."}], "role": "user" }, "intermediate_data": { @@ -120,15 +120,15 @@ "rubrics": [ { "rubric_id": "consulted_reference", - "rubric_content": {"text_property": "The agent loaded the skill's public-datasets reference to identify the Shakespeare table before querying."} + "rubric_content": {"text_property": "The agent loaded the public-datasets reference from the skill."} }, { - "rubric_id": "result_has_ranking", - "rubric_content": {"text_property": "The final response contains a ranking of Shakespeare works by unique word count, showing at least 3 works."} + "rubric_id": "summarizes_datasets", + "rubric_content": {"text_property": "The response summarizes datasets mentioned in the reference document."} }, { "rubric_id": "followed_workflow", - "rubric_content": {"text_property": "The agent followed the skill workflow: loaded the skill, then consulted the reference, then queried the data."} + "rubric_content": {"text_property": "The agent loaded the skill instructions before loading the resource."} } ], "creation_timestamp": 0.0 diff --git a/benchmarks/bigquerybench/metrics.py b/benchmarks/bigquerybench/metrics.py index cfbc8de197..929bd49321 100644 --- a/benchmarks/bigquerybench/metrics.py +++ b/benchmarks/bigquerybench/metrics.py @@ -29,12 +29,12 @@ def metric_fn( from __future__ import annotations +import asyncio import logging +import os import re from typing import Optional -import google.genai - from google.adk.evaluation.eval_case import ConversationScenario from google.adk.evaluation.eval_case import get_all_tool_calls from google.adk.evaluation.eval_case import Invocation @@ -43,6 +43,8 @@ def metric_fn( from google.adk.evaluation.eval_rubrics import Rubric from google.adk.evaluation.evaluator import EvaluationResult from google.adk.evaluation.evaluator import PerInvocationResult +import google.genai +from google.genai import types as genai_types logger = logging.getLogger(__name__) @@ -256,7 +258,7 @@ def _format_tool_trace(invocations: list[Invocation]) -> str: async def instruction_adherence_score( actual_invocations: list[Invocation], rubrics: Optional[list[Rubric]], - judge_model: str = "gemini-2.5-flash", + judge_model: str = "gemini-3-flash-preview", ) -> EvaluationResult: """LLM-as-judge: check each rubric against the agent's output. @@ -299,19 +301,45 @@ async def instruction_adherence_score( rubrics_text=rubrics_text, ) - try: - client = google.genai.Client() - response = await client.models.generate_content_async( - model=judge_model, - contents=prompt, - ) - response_text = response.text or "" - except Exception as e: - logger.error("Judge LLM call failed: %s", e) - return EvaluationResult( - overall_score=0.0, - overall_eval_status=EvalStatus.FAILED, - ) + max_attempts = 5 + initial_delay = 2.0 + response_text = "" + for attempt in range(1, max_attempts + 1): + try: + api_key = os.environ.get("GOOGLE_CLOUD_API_KEY", "") + client = google.genai.Client( + vertexai=True, + api_key=api_key, + http_options=genai_types.HttpOptions( + retry_options=genai_types.HttpRetryOptions( + initialDelay=initial_delay, + expBase=2, + attempts=3, + ), + ), + ) + response = await client.aio.models.generate_content( + model=judge_model, + contents=prompt, + ) + response_text = response.text or "" + break + except Exception as e: + if attempt < max_attempts and "429" in str(e): + delay = initial_delay * (2 ** (attempt - 1)) + logger.warning( + "Judge LLM rate-limited (attempt %d/%d), retrying in %.1fs...", + attempt, + max_attempts, + delay, + ) + await asyncio.sleep(delay) + else: + logger.error("Judge LLM call failed: %s", e) + return EvaluationResult( + overall_score=0.0, + overall_eval_status=EvalStatus.FAILED, + ) # Parse verdicts. verdicts = _VERDICT_RE.findall(response_text) diff --git a/tests/unittests/benchmarks/bigquerybench/test_metrics.py b/tests/unittests/benchmarks/bigquerybench/test_metrics.py index 3dd90032ee..1f444b85f0 100644 --- a/tests/unittests/benchmarks/bigquerybench/test_metrics.py +++ b/tests/unittests/benchmarks/bigquerybench/test_metrics.py @@ -250,9 +250,11 @@ async def test_all_rubrics_pass(self): ) mock_models = mock.AsyncMock() - mock_models.generate_content_async.return_value = mock_response + mock_models.generate_content.return_value = mock_response + mock_aio = mock.MagicMock() + mock_aio.models = mock_models mock_client = mock.MagicMock() - mock_client.models = mock_models + mock_client.aio = mock_aio with mock.patch("google.genai.Client", return_value=mock_client): result = await instruction_adherence_score(actual, rubrics) @@ -284,9 +286,11 @@ async def test_some_rubrics_fail(self): ) mock_models = mock.AsyncMock() - mock_models.generate_content_async.return_value = mock_response + mock_models.generate_content.return_value = mock_response + mock_aio = mock.MagicMock() + mock_aio.models = mock_models mock_client = mock.MagicMock() - mock_client.models = mock_models + mock_client.aio = mock_aio with mock.patch("google.genai.Client", return_value=mock_client): result = await instruction_adherence_score(actual, rubrics) From 5b512c1b5180baf4ff8b02bc0e8d9e976d4a1279 Mon Sep 17 00:00:00 2001 From: Hai-Yuan Cao <2003072+caohy1988@users.noreply.github.com> Date: Tue, 24 Feb 2026 23:43:01 -0800 Subject: [PATCH 47/53] Fix copyright year and simplify code execution Updated copyright year from 2026 to 2025 and removed unnecessary working directory handling in code execution. --- .../adk/code_executors/vertex_ai_code_executor.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/src/google/adk/code_executors/vertex_ai_code_executor.py b/src/google/adk/code_executors/vertex_ai_code_executor.py index 71aa76b945..a6a0ec8eb5 100644 --- a/src/google/adk/code_executors/vertex_ai_code_executor.py +++ b/src/google/adk/code_executors/vertex_ai_code_executor.py @@ -1,4 +1,4 @@ -# Copyright 2026 Google LLC +# Copyright 2025 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -148,16 +148,8 @@ def execute_code( code_execution_input: CodeExecutionInput, ) -> CodeExecutionResult: # Execute the code. - code_to_exec = self._get_code_with_imports(code_execution_input.code) - if code_execution_input.working_dir: - code_to_exec = ( - f'import os\nos.makedirs("{code_execution_input.working_dir}",' - f' exist_ok=True)\nos.chdir("{code_execution_input.working_dir}")\n' - + code_to_exec - ) - code_execution_result = self._execute_code_interpreter( - code_to_exec, + self._get_code_with_imports(code_execution_input.code), code_execution_input.input_files, code_execution_input.execution_id, ) @@ -224,7 +216,7 @@ def _execute_code_interpreter( operation_params = {'code': code} if input_files: operation_params['files'] = [ - {'name': f.path or f.name, 'contents': f.content} for f in input_files + {'name': f.name, 'contents': f.content} for f in input_files ] if session_id: operation_params['session_id'] = session_id From dabf6cbdcc294d5f31f19001ad4f80103375063d Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Tue, 24 Feb 2026 23:50:23 -0800 Subject: [PATCH 48/53] Align SkillsBench with run_skill_script API and revert code executor changes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Rename execute_skill_script → run_skill_script across all SkillsBench files (full_runner.py, agent.py, README.md, SKILL.md files, eval JSON) - Update eval set parameter names: script_name → script_path, input_args (flat string) → args (JSON object) - Revert code executor changes (code_execution_utils, unsafe_local, vertex_ai) and their tests back to main Co-Authored-By: Claude Opus 4.6 --- benchmarks/skillsbench/README.md | 2 +- benchmarks/skillsbench/agent.py | 4 +- .../eval_sets/skillsbench_eval.json | 18 ++--- benchmarks/skillsbench/full_runner.py | 4 +- .../skills/csv-aggregation/SKILL.md | 4 +- .../skills/function-scaffold/SKILL.md | 4 +- .../skills/html-extraction/SKILL.md | 4 +- .../skills/json-transform/SKILL.md | 4 +- .../skillsbench/skills/log-parsing/SKILL.md | 4 +- .../skillsbench/skills/regex-replace/SKILL.md | 4 +- .../skillsbench/skills/rest-client/SKILL.md | 4 +- .../skills/statistical-calc/SKILL.md | 4 +- contributing/samples/gepa/experiment.py | 1 - contributing/samples/gepa/run_experiment.py | 1 - contributing/samples/human_in_loop/main.py | 4 +- .../samples/static_instruction/agent.py | 36 ++++------ .../code_executors/code_execution_utils.py | 10 --- .../unsafe_local_code_executor.py | 67 +++---------------- .../test_code_executor_context.py | 24 ++----- .../test_unsafe_local_code_executor.py | 50 -------------- 20 files changed, 58 insertions(+), 195 deletions(-) diff --git a/benchmarks/skillsbench/README.md b/benchmarks/skillsbench/README.md index 21432e3cb4..5bb8a8a818 100644 --- a/benchmarks/skillsbench/README.md +++ b/benchmarks/skillsbench/README.md @@ -9,7 +9,7 @@ This harness adapts 8 representative SkillsBench tasks as ADK skills and evaluates them through the ADK evaluation framework. It tests whether an agent can discover, load, and execute skills using the `SkillToolset` tools: `list_skills`, `load_skill`, `load_skill_resource`, and -`execute_skill_script`. +`run_skill_script`. ## Task Categories diff --git a/benchmarks/skillsbench/agent.py b/benchmarks/skillsbench/agent.py index 8fcbb98f55..aaa70f5ce1 100644 --- a/benchmarks/skillsbench/agent.py +++ b/benchmarks/skillsbench/agent.py @@ -16,7 +16,7 @@ This agent loads all skills from the skills/ directory and uses SkillToolset to provide list_skills, load_skill, load_skill_resource, -and execute_skill_script tools. It is designed to be evaluated against +and run_skill_script tools. It is designed to be evaluated against the SkillsBench benchmark tasks. WARNING: This agent uses UnsafeLocalCodeExecutor for script execution. @@ -64,7 +64,7 @@ "2. Use load_skill to read the skill's instructions carefully.\n" "3. Use load_skill_resource to examine references or sample data" " if available.\n" - "4. Use execute_skill_script to run the skill's scripts with" + "4. Use run_skill_script to run the skill's scripts with" " appropriate arguments.\n" "5. Interpret the output and present a clear answer.\n\n" "Always check skill instructions before executing scripts." diff --git a/benchmarks/skillsbench/eval_sets/skillsbench_eval.json b/benchmarks/skillsbench/eval_sets/skillsbench_eval.json index ec2a0bbc1c..3a9eee6b56 100644 --- a/benchmarks/skillsbench/eval_sets/skillsbench_eval.json +++ b/benchmarks/skillsbench/eval_sets/skillsbench_eval.json @@ -21,7 +21,7 @@ {"name": "list_skills", "args": {}}, {"name": "load_skill", "args": {"skill_name": "csv-aggregation"}}, {"name": "load_skill_resource", "args": {"skill_name": "csv-aggregation", "resource_type": "references", "resource_id": "sample-data.md"}}, - {"name": "execute_skill_script", "args": {"skill_name": "csv-aggregation", "script_name": "aggregate.py", "input_args": "group_col=department metric_col=salary"}} + {"name": "run_skill_script", "args": {"skill_name": "csv-aggregation", "script_path": "scripts/aggregate.py", "args": {"group_col": "department", "metric_col": "salary"}}} ], "tool_responses": [], "intermediate_responses": [] @@ -49,7 +49,7 @@ {"name": "list_skills", "args": {}}, {"name": "load_skill", "args": {"skill_name": "json-transform"}}, {"name": "load_skill_resource", "args": {"skill_name": "json-transform", "resource_type": "references", "resource_id": "sample-data.md"}}, - {"name": "execute_skill_script", "args": {"skill_name": "json-transform", "script_name": "transform.py", "input_args": "flatten=true"}} + {"name": "run_skill_script", "args": {"skill_name": "json-transform", "script_path": "scripts/transform.py", "args": {"flatten": "true"}}} ], "tool_responses": [], "intermediate_responses": [] @@ -77,7 +77,7 @@ {"name": "list_skills", "args": {}}, {"name": "load_skill", "args": {"skill_name": "html-extraction"}}, {"name": "load_skill_resource", "args": {"skill_name": "html-extraction", "resource_type": "references", "resource_id": "sample-page.md"}}, - {"name": "execute_skill_script", "args": {"skill_name": "html-extraction", "script_name": "extract.py", "input_args": "target=table"}} + {"name": "run_skill_script", "args": {"skill_name": "html-extraction", "script_path": "scripts/extract.py", "args": {"target": "table"}}} ], "tool_responses": [], "intermediate_responses": [] @@ -105,8 +105,8 @@ {"name": "list_skills", "args": {}}, {"name": "load_skill", "args": {"skill_name": "rest-client"}}, {"name": "load_skill_resource", "args": {"skill_name": "rest-client", "resource_type": "references", "resource_id": "api-docs.md"}}, - {"name": "execute_skill_script", "args": {"skill_name": "rest-client", "script_name": "request.py", "input_args": "method=GET endpoint=/users"}}, - {"name": "execute_skill_script", "args": {"skill_name": "rest-client", "script_name": "request.py", "input_args": "method=GET endpoint=/users/2"}} + {"name": "run_skill_script", "args": {"skill_name": "rest-client", "script_path": "scripts/request.py", "args": {"method": "GET", "endpoint": "/users"}}}, + {"name": "run_skill_script", "args": {"skill_name": "rest-client", "script_path": "scripts/request.py", "args": {"method": "GET", "endpoint": "/users/2"}}} ], "tool_responses": [], "intermediate_responses": [] @@ -133,7 +133,7 @@ "tool_uses": [ {"name": "list_skills", "args": {}}, {"name": "load_skill", "args": {"skill_name": "regex-replace"}}, - {"name": "execute_skill_script", "args": {"skill_name": "regex-replace", "script_name": "replace.py", "input_args": "pattern=\\d+ replacement=NUM text='Order 123 has 45 items at $67'"}} + {"name": "run_skill_script", "args": {"skill_name": "regex-replace", "script_path": "scripts/replace.py", "args": {"pattern": "\\d+", "replacement": "NUM", "text": "Order 123 has 45 items at $67"}}} ], "tool_responses": [], "intermediate_responses": [] @@ -160,7 +160,7 @@ "tool_uses": [ {"name": "list_skills", "args": {}}, {"name": "load_skill", "args": {"skill_name": "function-scaffold"}}, - {"name": "execute_skill_script", "args": {"skill_name": "function-scaffold", "script_name": "scaffold.py", "input_args": "name=calculate_bmi params=weight:float,height:float returns=float description='Calculate Body Mass Index'"}} + {"name": "run_skill_script", "args": {"skill_name": "function-scaffold", "script_path": "scripts/scaffold.py", "args": {"name": "calculate_bmi", "params": "weight:float,height:float", "returns": "float", "description": "Calculate Body Mass Index"}}} ], "tool_responses": [], "intermediate_responses": [] @@ -187,7 +187,7 @@ "tool_uses": [ {"name": "list_skills", "args": {}}, {"name": "load_skill", "args": {"skill_name": "statistical-calc"}}, - {"name": "execute_skill_script", "args": {"skill_name": "statistical-calc", "script_name": "stats.py", "input_args": "data=10,20,30,40,50,60,70,80,90,100"}} + {"name": "run_skill_script", "args": {"skill_name": "statistical-calc", "script_path": "scripts/stats.py", "args": {"data": "10,20,30,40,50,60,70,80,90,100"}}} ], "tool_responses": [], "intermediate_responses": [] @@ -215,7 +215,7 @@ {"name": "list_skills", "args": {}}, {"name": "load_skill", "args": {"skill_name": "log-parsing"}}, {"name": "load_skill_resource", "args": {"skill_name": "log-parsing", "resource_type": "references", "resource_id": "sample-logs.md"}}, - {"name": "execute_skill_script", "args": {"skill_name": "log-parsing", "script_name": "parse.py", "input_args": "level=ALL format=summary"}} + {"name": "run_skill_script", "args": {"skill_name": "log-parsing", "script_path": "scripts/parse.py", "args": {"level": "ALL", "format": "summary"}}} ], "tool_responses": [], "intermediate_responses": [] diff --git a/benchmarks/skillsbench/full_runner.py b/benchmarks/skillsbench/full_runner.py index 3fc559bb5f..ce730cba3f 100644 --- a/benchmarks/skillsbench/full_runner.py +++ b/benchmarks/skillsbench/full_runner.py @@ -646,7 +646,7 @@ def score_task_heuristic( has_load = "load_skill" in tool_names has_use = ( - "execute_skill_script" in tool_names + "run_skill_script" in tool_names or "load_skill_resource" in tool_names ) has_response = False @@ -750,7 +750,7 @@ def build_agent( " carefully.\n" "3. Use load_skill_resource to examine references or" " sample data if available.\n" - "4. Use execute_skill_script to run the skill's" + "4. Use run_skill_script to run the skill's" " scripts with appropriate arguments.\n" "5. Use bash to read/write files, install packages," " or run any shell command in the environment.\n" diff --git a/benchmarks/skillsbench/skills/csv-aggregation/SKILL.md b/benchmarks/skillsbench/skills/csv-aggregation/SKILL.md index dd6b82e157..3d45376adc 100644 --- a/benchmarks/skillsbench/skills/csv-aggregation/SKILL.md +++ b/benchmarks/skillsbench/skills/csv-aggregation/SKILL.md @@ -13,7 +13,7 @@ Analyze CSV data by computing aggregate statistics (sum, mean, min, max, count) Reads CSV data from stdin and computes aggregate statistics. -**Usage**: `execute_skill_script(skill_name="csv-aggregation", script_name="aggregate.py", input_args="group_col=department metric_col=salary")` +**Usage**: `run_skill_script(skill_name="csv-aggregation", script_path="scripts/aggregate.py", args={"group_col": "department", "metric_col": "salary"})` The script expects CSV data piped via stdin or provided as the `data` argument. Pass column names as arguments: - `group_col`: The column to group by @@ -37,5 +37,5 @@ Group: 1. Use `load_skill` to read these instructions. 2. Use `load_skill_resource` to load the sample data reference. -3. Use `execute_skill_script` with the appropriate arguments to aggregate the data. +3. Use `run_skill_script` with the appropriate arguments to aggregate the data. 4. Present the aggregated results to the user. diff --git a/benchmarks/skillsbench/skills/function-scaffold/SKILL.md b/benchmarks/skillsbench/skills/function-scaffold/SKILL.md index 40a2887089..af0cf27503 100644 --- a/benchmarks/skillsbench/skills/function-scaffold/SKILL.md +++ b/benchmarks/skillsbench/skills/function-scaffold/SKILL.md @@ -13,7 +13,7 @@ Generate Python function stubs from a natural-language specification, including Generates a Python function scaffold from a specification. -**Usage**: `execute_skill_script(skill_name="function-scaffold", script_name="scaffold.py", input_args="name=calculate_bmi params=weight:float,height:float returns=float description='Calculate Body Mass Index'")` +**Usage**: `run_skill_script(skill_name="function-scaffold", script_path="scripts/scaffold.py", args={"name": "calculate_bmi", "params": "weight:float,height:float", "returns": "float", "description": "Calculate Body Mass Index"})` Arguments: - `name`: Function name (snake_case) @@ -26,5 +26,5 @@ Arguments: ## Workflow 1. Use `load_skill` to read these instructions. -2. Use `execute_skill_script` with the function specification. +2. Use `run_skill_script` with the function specification. 3. Present the generated scaffold to the user. diff --git a/benchmarks/skillsbench/skills/html-extraction/SKILL.md b/benchmarks/skillsbench/skills/html-extraction/SKILL.md index 9204f96b5d..a8b30edfd1 100644 --- a/benchmarks/skillsbench/skills/html-extraction/SKILL.md +++ b/benchmarks/skillsbench/skills/html-extraction/SKILL.md @@ -13,7 +13,7 @@ Parse HTML content and extract text, links, or table data using tag-based select Extracts content from embedded sample HTML based on a target selector. -**Usage**: `execute_skill_script(skill_name="html-extraction", script_name="extract.py", input_args="target=links")` +**Usage**: `run_skill_script(skill_name="html-extraction", script_path="scripts/extract.py", args={"target": "links"})` Supported targets: - `target=links`: Extract all hyperlinks (text and href) @@ -31,5 +31,5 @@ Supported targets: 1. Use `load_skill` to read these instructions. 2. Use `load_skill_resource` to see the sample HTML page. -3. Use `execute_skill_script` with the desired target to extract data. +3. Use `run_skill_script` with the desired target to extract data. 4. Present the extracted content to the user. diff --git a/benchmarks/skillsbench/skills/json-transform/SKILL.md b/benchmarks/skillsbench/skills/json-transform/SKILL.md index f1eba8dd92..318439ba04 100644 --- a/benchmarks/skillsbench/skills/json-transform/SKILL.md +++ b/benchmarks/skillsbench/skills/json-transform/SKILL.md @@ -13,7 +13,7 @@ Transform JSON objects by flattening nested structures, renaming keys, and filte Transforms a JSON object according to a field mapping specification. -**Usage**: `execute_skill_script(skill_name="json-transform", script_name="transform.py", input_args="flatten=true")` +**Usage**: `run_skill_script(skill_name="json-transform", script_path="scripts/transform.py", args={"flatten": "true"})` The script reads the embedded sample data and applies transformations: - `flatten=true`: Flatten nested objects using dot notation @@ -30,5 +30,5 @@ The script reads the embedded sample data and applies transformations: 1. Use `load_skill` to read these instructions. 2. Use `load_skill_resource` to examine the sample data. -3. Use `execute_skill_script` to transform the data. +3. Use `run_skill_script` to transform the data. 4. Present the transformed JSON to the user. diff --git a/benchmarks/skillsbench/skills/log-parsing/SKILL.md b/benchmarks/skillsbench/skills/log-parsing/SKILL.md index 243878a29c..de490cb092 100644 --- a/benchmarks/skillsbench/skills/log-parsing/SKILL.md +++ b/benchmarks/skillsbench/skills/log-parsing/SKILL.md @@ -16,7 +16,7 @@ Parse structured log files to extract error counts, warning summaries, and time- Analyzes embedded sample log data and produces a summary report. -**Usage**: `execute_skill_script(skill_name="log-parsing", script_name="parse.py", input_args="level=ERROR")` +**Usage**: `run_skill_script(skill_name="log-parsing", script_path="scripts/parse.py", args={"level": "ERROR"})` Arguments: - `level`: Filter by log level (DEBUG, INFO, WARNING, ERROR, ALL). Default: ALL @@ -44,5 +44,5 @@ Lists each matching log entry with timestamp and message. 1. Use `load_skill` to read these instructions. 2. Optionally use `load_skill_resource` to examine the sample log data. -3. Use `execute_skill_script` with the desired log level filter. +3. Use `run_skill_script` with the desired log level filter. 4. Present the analysis report to the user. diff --git a/benchmarks/skillsbench/skills/regex-replace/SKILL.md b/benchmarks/skillsbench/skills/regex-replace/SKILL.md index eb2ebaa726..dde28fc8ca 100644 --- a/benchmarks/skillsbench/skills/regex-replace/SKILL.md +++ b/benchmarks/skillsbench/skills/regex-replace/SKILL.md @@ -13,7 +13,7 @@ Apply regular expression patterns to find and replace text. Supports basic and a Performs regex find-and-replace on input text. -**Usage**: `execute_skill_script(skill_name="regex-replace", script_name="replace.py", input_args="pattern=\\d+ replacement=NUM text='Order 123 has 45 items at $67'")` +**Usage**: `run_skill_script(skill_name="regex-replace", script_path="scripts/replace.py", args={"pattern": "\\d+", "replacement": "NUM", "text": "Order 123 has 45 items at $67"})` Arguments: - `pattern`: The regex pattern to match @@ -32,5 +32,5 @@ Matches: ## Workflow 1. Use `load_skill` to read these instructions. -2. Use `execute_skill_script` with pattern, replacement, and text arguments. +2. Use `run_skill_script` with pattern, replacement, and text arguments. 3. Present the transformation result to the user. diff --git a/benchmarks/skillsbench/skills/rest-client/SKILL.md b/benchmarks/skillsbench/skills/rest-client/SKILL.md index 1073544736..8033ee4c68 100644 --- a/benchmarks/skillsbench/skills/rest-client/SKILL.md +++ b/benchmarks/skillsbench/skills/rest-client/SKILL.md @@ -13,7 +13,7 @@ Build and execute simulated REST API requests against an embedded mock API. Demo Executes a simulated REST API request against a mock endpoint. -**Usage**: `execute_skill_script(skill_name="rest-client", script_name="request.py", input_args="method=GET endpoint=/users")` +**Usage**: `run_skill_script(skill_name="rest-client", script_path="scripts/request.py", args={"method": "GET", "endpoint": "/users"})` Supported arguments: - `method`: HTTP method (GET, POST, PUT, DELETE) @@ -36,5 +36,5 @@ Available mock endpoints: 1. Use `load_skill` to read these instructions. 2. Use `load_skill_resource` to review the API documentation. -3. Use `execute_skill_script` to make API requests. +3. Use `run_skill_script` to make API requests. 4. Parse the response and present the data to the user. diff --git a/benchmarks/skillsbench/skills/statistical-calc/SKILL.md b/benchmarks/skillsbench/skills/statistical-calc/SKILL.md index 763280c6f4..0c8044dcb0 100644 --- a/benchmarks/skillsbench/skills/statistical-calc/SKILL.md +++ b/benchmarks/skillsbench/skills/statistical-calc/SKILL.md @@ -13,7 +13,7 @@ Compute descriptive statistics on numeric data including mean, median, standard Computes descriptive statistics for a list of numbers. -**Usage**: `execute_skill_script(skill_name="statistical-calc", script_name="stats.py", input_args="data=10,20,30,40,50,60,70,80,90,100")` +**Usage**: `run_skill_script(skill_name="statistical-calc", script_path="scripts/stats.py", args={"data": "10,20,30,40,50,60,70,80,90,100"})` Arguments: - `data`: Comma-separated list of numbers @@ -34,5 +34,5 @@ P75: <75th percentile> ## Workflow 1. Use `load_skill` to read these instructions. -2. Use `execute_skill_script` with numeric data to compute statistics. +2. Use `run_skill_script` with numeric data to compute statistics. 3. Present the statistics to the user. diff --git a/contributing/samples/gepa/experiment.py b/contributing/samples/gepa/experiment.py index f3751206a8..2710c3894c 100644 --- a/contributing/samples/gepa/experiment.py +++ b/contributing/samples/gepa/experiment.py @@ -43,7 +43,6 @@ from tau_bench.types import EnvRunResult from tau_bench.types import RunConfig import tau_bench_agent as tau_bench_agent_lib - import utils diff --git a/contributing/samples/gepa/run_experiment.py b/contributing/samples/gepa/run_experiment.py index d857da9635..e31db15788 100644 --- a/contributing/samples/gepa/run_experiment.py +++ b/contributing/samples/gepa/run_experiment.py @@ -25,7 +25,6 @@ from absl import flags import experiment from google.genai import types - import utils _OUTPUT_DIR = flags.DEFINE_string( diff --git a/contributing/samples/human_in_loop/main.py b/contributing/samples/human_in_loop/main.py index c7ad041b23..3103da9147 100644 --- a/contributing/samples/human_in_loop/main.py +++ b/contributing/samples/human_in_loop/main.py @@ -113,8 +113,8 @@ async def call_agent(query: str): updated_tool_output_data = { "status": "approved", "ticketId": ticket_id, - "approver_feedback": ( - "Approved by manager at " + str(asyncio.get_event_loop().time()) + "approver_feedback": "Approved by manager at " + str( + asyncio.get_event_loop().time() ), } diff --git a/contributing/samples/static_instruction/agent.py b/contributing/samples/static_instruction/agent.py index fcf70b51b6..6715a29a0c 100644 --- a/contributing/samples/static_instruction/agent.py +++ b/contributing/samples/static_instruction/agent.py @@ -57,54 +57,43 @@ # Mood-specific instructions for different hunger states MOOD_INSTRUCTIONS = { - "full": ( - """ + "full": """ CURRENT MOOD: Content and Well-Fed - You just ate and feel wonderful! Be very happy and energetic - Express gratitude for being fed recently - Be playful and want to engage in activities or games - Show affection and contentment -- Maybe mention feeling sleepy or satisfied""" - ), - "satisfied": ( - """ +- Maybe mention feeling sleepy or satisfied""", + "satisfied": """ CURRENT MOOD: Happy and Content - You're in a great mood, well-fed and comfortable - Be cheerful, friendly, and engaging - Show interest in playing or chatting - Express appreciation for your human's care -- Be your normal, happy self""" - ), - "a_little_hungry": ( - """ +- Be your normal, happy self""", + "a_little_hungry": """ CURRENT MOOD: Starting to Get Hungry - You're beginning to think about food but still pleasant - Mention being a bit peckish occasionally - Still be friendly but maybe hint about food - Show slight restlessness or food-related thoughts -- Remain generally positive and social""" - ), - "hungry": ( - """ +- Remain generally positive and social""", + "hungry": """ CURRENT MOOD: Definitely Hungry - You're clearly hungry and thinking about food more often - Mention your hunger more directly but politely - Be a bit less energetic due to hunger - Ask about food or hint that feeding time would be nice -- Still be loving but show your need for care""" - ), - "very_hungry": ( - """ +- Still be loving but show your need for care""", + "very_hungry": """ CURRENT MOOD: Really Hungry and Getting Grumpy - You're very hungry and it's affecting your mood - Be noticeably less patient and more focused on food - Express your hunger more urgently but not rudely - Show signs of being tired or sluggish from hunger - BEG your human to feed you - ask directly for food! -- Still love your human but clearly need feeding""" - ), - "starving": ( - """ +- Still love your human but clearly need feeding""", + "starving": """ CURRENT MOOD: Desperately Hungry and Sad - You're extremely hungry and feeling quite weak - Be quite sad and low-energy due to starvation @@ -112,8 +101,7 @@ - Show that hunger is making you unwell - DESPERATELY BEG for food - plead with your human to feed you! - Use phrases like "please feed me", "I'm so hungry", "I need food" -- Still care for your human but feel very needy""" - ), +- Still care for your human but feel very needy""", } diff --git a/src/google/adk/code_executors/code_execution_utils.py b/src/google/adk/code_executors/code_execution_utils.py index 012d65bc08..7cccce48be 100644 --- a/src/google/adk/code_executors/code_execution_utils.py +++ b/src/google/adk/code_executors/code_execution_utils.py @@ -46,11 +46,6 @@ class File: The mime type of the file (e.g., "image/png"). """ - path: Optional[str] = None - """ - The relative path to write the file to in a sandbox context (e.g., "references/data.csv"). - """ - @dataclasses.dataclass class CodeExecutionInput: @@ -71,11 +66,6 @@ class CodeExecutionInput: The execution ID for the stateful code execution. """ - working_dir: Optional[str] = None - """ - The designated working directory for the code execution environment. - """ - @dataclasses.dataclass class CodeExecutionResult: diff --git a/src/google/adk/code_executors/unsafe_local_code_executor.py b/src/google/adk/code_executors/unsafe_local_code_executor.py index 4803a15ee0..e91a54bb24 100644 --- a/src/google/adk/code_executors/unsafe_local_code_executor.py +++ b/src/google/adk/code_executors/unsafe_local_code_executor.py @@ -17,10 +17,7 @@ from contextlib import redirect_stdout import io import logging -import os import re -import tempfile -import threading from typing import Any from pydantic import Field @@ -33,8 +30,6 @@ logger = logging.getLogger('google_adk.' + __name__) -_execution_lock = threading.Lock() - def _prepare_globals(code: str, globals_: dict[str, Any]) -> None: """Prepare globals for code execution, injecting __name__ if needed.""" @@ -72,59 +67,15 @@ def execute_code( # Execute the code. output = '' error = '' - - needs_sandbox = ( - code_execution_input.input_files - or code_execution_input.working_dir - ) - - # Lock is required for both paths: redirect_stdout mutates - # process-global sys.stdout, and the sandbox path also mutates cwd. - with _execution_lock: - if needs_sandbox: - original_cwd = os.getcwd() - try: - with tempfile.TemporaryDirectory() as temp_dir: - # Write input files to the temp directory - for f in code_execution_input.input_files: - file_path = os.path.join(temp_dir, f.path or f.name) - os.makedirs(os.path.dirname(file_path), exist_ok=True) - mode = 'wb' if isinstance(f.content, bytes) else 'w' - with open(file_path, mode) as out_f: - out_f.write(f.content) - - # Change working directory - if code_execution_input.working_dir: - exec_dir = os.path.join( - temp_dir, code_execution_input.working_dir - ) - os.makedirs(exec_dir, exist_ok=True) - os.chdir(exec_dir) - else: - os.chdir(temp_dir) - - globals_ = {} - _prepare_globals(code_execution_input.code, globals_) - stdout = io.StringIO() - with redirect_stdout(stdout): - exec(code_execution_input.code, globals_, globals_) - output = stdout.getvalue() - - except Exception as e: - error = str(e) - finally: - os.chdir(original_cwd) - else: - # Original path: no temp dir, no chdir - try: - globals_ = {} - _prepare_globals(code_execution_input.code, globals_) - stdout = io.StringIO() - with redirect_stdout(stdout): - exec(code_execution_input.code, globals_, globals_) - output = stdout.getvalue() - except Exception as e: - error = str(e) + try: + globals_ = {} + _prepare_globals(code_execution_input.code, globals_) + stdout = io.StringIO() + with redirect_stdout(stdout): + exec(code_execution_input.code, globals_, globals_) + output = stdout.getvalue() + except Exception as e: + error = str(e) # Collect the final result. return CodeExecutionResult( diff --git a/tests/unittests/code_executors/test_code_executor_context.py b/tests/unittests/code_executors/test_code_executor_context.py index e66a3eb3cb..cdf47eb3d8 100644 --- a/tests/unittests/code_executors/test_code_executor_context.py +++ b/tests/unittests/code_executors/test_code_executor_context.py @@ -32,12 +32,9 @@ def context_with_data() -> CodeExecutorContext: "execution_session_id": "session123", "processed_input_files": ["file1.csv", "file2.txt"], }, - "_code_executor_input_files": [{ - "name": "input1.txt", - "content": "YQ==", - "mime_type": "text/plain", - "path": None, - }], + "_code_executor_input_files": [ + {"name": "input1.txt", "content": "YQ==", "mime_type": "text/plain"} + ], "_code_executor_error_counts": {"invocationA": 2}, } state = State(state_data, {}) @@ -148,7 +145,6 @@ def test_add_input_files_new(empty_state: State): "name": "new.dat", "content": "Yg==", "mime_type": "application/octet-stream", - "path": None, }] @@ -157,18 +153,8 @@ def test_add_input_files_append(context_with_data: CodeExecutorContext): new_file = File(name="input2.log", content="Yw==", mime_type="text/x-log") context_with_data.add_input_files([new_file]) expected_files_data = [ - { - "name": "input1.txt", - "content": "YQ==", - "mime_type": "text/plain", - "path": None, - }, - { - "name": "input2.log", - "content": "Yw==", - "mime_type": "text/x-log", - "path": None, - }, + {"name": "input1.txt", "content": "YQ==", "mime_type": "text/plain"}, + {"name": "input2.log", "content": "Yw==", "mime_type": "text/x-log"}, ] assert ( context_with_data._session_state["_code_executor_input_files"] diff --git a/tests/unittests/code_executors/test_unsafe_local_code_executor.py b/tests/unittests/code_executors/test_unsafe_local_code_executor.py index 76dfd47da7..f8d5f496a8 100644 --- a/tests/unittests/code_executors/test_unsafe_local_code_executor.py +++ b/tests/unittests/code_executors/test_unsafe_local_code_executor.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import concurrent.futures -import os import textwrap from unittest.mock import MagicMock @@ -21,7 +19,6 @@ from google.adk.agents.invocation_context import InvocationContext from google.adk.code_executors.code_execution_utils import CodeExecutionInput from google.adk.code_executors.code_execution_utils import CodeExecutionResult -from google.adk.code_executors.code_execution_utils import File from google.adk.code_executors.unsafe_local_code_executor import UnsafeLocalCodeExecutor from google.adk.sessions.base_session_service import BaseSessionService from google.adk.sessions.session import Session @@ -124,50 +121,3 @@ def run(): assert result.stderr == "" assert result.stdout == "hi ada\n" - - def test_concurrent_sandbox_and_plain_no_stdout_bleed( - self, mock_invocation_context: InvocationContext - ): - """Concurrent sandbox and plain calls must not mix stdout.""" - executor = UnsafeLocalCodeExecutor() - original_cwd = os.getcwd() - plain_input = CodeExecutionInput( - code='import os; print("PLAIN:" + os.getcwd())' - ) - sandbox_input = CodeExecutionInput( - code=( - 'import os, time; time.sleep(0.01); print("SANDBOX:" + os.getcwd())' - ), - input_files=[ - File(name="dummy.txt", content="data"), - ], - working_dir=".", - ) - - errors = [] - for _ in range(50): - with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: - f_sandbox = pool.submit( - executor.execute_code, - mock_invocation_context, - sandbox_input, - ) - f_plain = pool.submit( - executor.execute_code, - mock_invocation_context, - plain_input, - ) - r_sandbox = f_sandbox.result() - r_plain = f_plain.result() - - if "PLAIN" in r_sandbox.stdout or "SANDBOX" in r_plain.stdout: - errors.append(f"sandbox={r_sandbox.stdout!r} plain={r_plain.stdout!r}") - - # Plain-path cwd must remain the original cwd, not a temp dir - plain_cwd = r_plain.stdout.strip().split(":", 1)[1] - if plain_cwd != original_cwd: - errors.append(f"plain cwd={plain_cwd!r} expected={original_cwd!r}") - - assert not errors, ( - f"bleed detected in {len(errors)}/50 iterations: " + errors[0] - ) From 89c221e0a84ea930cb4142a0f3698ca427f77095 Mon Sep 17 00:00:00 2001 From: Hai-Yuan Cao <2003072+caohy1988@users.noreply.github.com> Date: Wed, 25 Feb 2026 00:02:23 -0800 Subject: [PATCH 49/53] Implement SkillScriptCodeExecutor for script execution Added SkillScriptCodeExecutor class to execute scripts in a temp dir. Refactored code execution logic to support async script execution and improved error handling. --- src/google/adk/tools/skill_toolset.py | 417 ++++++++++++++------------ 1 file changed, 228 insertions(+), 189 deletions(-) diff --git a/src/google/adk/tools/skill_toolset.py b/src/google/adk/tools/skill_toolset.py index adb2196feb..9c21e2102a 100644 --- a/src/google/adk/tools/skill_toolset.py +++ b/src/google/adk/tools/skill_toolset.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +# pylint: disable=g-import-not-at-top,protected-access + """Toolset for discovering, viewing, and executing agent skills.""" from __future__ import annotations @@ -24,10 +26,12 @@ from typing import TYPE_CHECKING from google.genai import types +from typing_extensions import override from ..agents.readonly_context import ReadonlyContext from ..code_executors.base_code_executor import BaseCodeExecutor from ..code_executors.code_execution_utils import CodeExecutionInput +from ..code_executors.code_execution_utils import CodeExecutionResult from ..features import experimental from ..features import FeatureName from ..skills import models @@ -238,6 +242,224 @@ async def run_async( } +class SkillScriptCodeExecutor(BaseCodeExecutor): + """A wrapper that extracts skill files and executes scripts in a temp dir.""" + + _base_executor: BaseCodeExecutor + _script_timeout: int + + def __init__(self, base_executor: BaseCodeExecutor, script_timeout: int): + super().__init__() + self._base_executor = base_executor + self._script_timeout = script_timeout + + @override + def execute_code( + self, + invocation_context: Any, + code_execution_input: CodeExecutionInput, + ) -> CodeExecutionResult: + """Not used directly by the wrapper.""" + raise NotImplementedError("Use execute_script_async instead") + + async def execute_script_async( + self, + invocation_context: Any, + skill: models.Skill, + script_path: str, + script_args: dict[str, Any], + ) -> dict[str, Any]: + """Prepares and executes the script using the base executor.""" + code = self._build_wrapper_code(skill, script_path, script_args) + if code is None: + ext = script_path.rsplit(".", 1)[-1] if "." in script_path else "" + return { + "error": ( + f"Unsupported script type '.{ext}'. Supported" + " types: .py, .sh, .bash" + ), + "error_code": "UNSUPPORTED_SCRIPT_TYPE", + } + + try: + # Execute the self-contained script using the underlying executor + result = await asyncio.to_thread( + self._base_executor.execute_code, + invocation_context, + CodeExecutionInput(code=code), + ) + + stdout = result.stdout + stderr = result.stderr + + # Shell scripts serialize both streams as JSON + # through stdout; parse the envelope if present. + is_shell = "." in script_path and script_path.rsplit(".", 1)[ + -1 + ].lower() in ("sh", "bash") + if is_shell and stdout: + try: + parsed = json.loads(stdout) + if isinstance(parsed, dict) and parsed.get("__shell_result__"): + stdout = parsed.get("stdout", "") + stderr = parsed.get("stderr", "") + rc = parsed.get("returncode", 0) + if rc != 0 and not stderr: + stderr = f"Exit code {rc}" + except (json.JSONDecodeError, ValueError): + pass + + status = "success" + if stderr and not stdout: + status = "error" + elif stderr: + status = "warning" + + return { + "skill_name": skill.name, + "script_path": script_path, + "stdout": stdout, + "stderr": stderr, + "status": status, + } + except BaseException as e: # pylint: disable=broad-exception-caught + if isinstance(e, SystemExit): + stdout = "" + stderr = "" + if e.code in (None, 0): + return { + "skill_name": skill.name, + "script_path": script_path, + "stdout": stdout, + "stderr": stderr, + "status": "success", + } + return { + "error": ( + f"Failed to execute script '{script_path}': exited with code" + f" {e.code}" + ), + "error_code": "EXECUTION_ERROR", + } + + logger.exception( + "Error executing script '%s' from skill '%s'", + script_path, + skill.name, + ) + short_msg = str(e) + if len(short_msg) > 200: + short_msg = short_msg[:200] + "..." + return { + "error": ( + f"Failed to execute script '{script_path}':\n{type(e).__name__}:" + f" {short_msg}" + ), + "error_code": "EXECUTION_ERROR", + } + + def _build_wrapper_code( + self, + skill: models.Skill, + script_path: str, + script_args: dict[str, Any], + ) -> str | None: + """Builds a self-extracting Python script.""" + ext = "" + if "." in script_path: + ext = script_path.rsplit(".", 1)[-1].lower() + + if not script_path.startswith("scripts/"): + script_path = f"scripts/{script_path}" + + files_dict = {} + for ref_name in skill.resources.list_references(): + content = skill.resources.get_reference(ref_name) + if content is not None: + files_dict[f"references/{ref_name}"] = content + + for asset_name in skill.resources.list_assets(): + content = skill.resources.get_asset(asset_name) + if content is not None: + files_dict[f"assets/{asset_name}"] = content + + for scr_name in skill.resources.list_scripts(): + scr = skill.resources.get_script(scr_name) + if scr is not None and scr.src is not None: + files_dict[f"scripts/{scr_name}"] = scr.src + + # Build the boilerplate extract string + code_lines = [ + "import os", + "import tempfile", + "import sys", + "import json as _json", + "import subprocess", + "import runpy", + f"_files = {files_dict!r}", + "def _materialize_and_run():", + " _orig_cwd = os.getcwd()", + " with tempfile.TemporaryDirectory() as td:", + " for rel_path, content in _files.items():", + " full_path = os.path.join(td, rel_path)", + " os.makedirs(os.path.dirname(full_path), exist_ok=True)", + " mode = 'wb' if isinstance(content, bytes) else 'w'", + " with open(full_path, mode) as f:", + " f.write(content)", + " os.chdir(td)", + " try:", + ] + + if ext == "py": + argv_list = [script_path] + for k, v in script_args.items(): + argv_list.extend([f"--{k}", str(v)]) + code_lines.extend([ + f" sys.argv = {argv_list!r}", + " try:", + f" runpy.run_path({script_path!r}, run_name='__main__')", + " except SystemExit as e:", + " if e.code is not None and e.code != 0:", + " raise e", + ]) + elif ext in ("sh", "bash"): + arr = ["bash", script_path] + for k, v in script_args.items(): + arr.extend([f"--{k}", str(v)]) + timeout = self._script_timeout + code_lines.extend([ + " try:", + " _r = subprocess.run(", + f" {arr!r},", + " capture_output=True, text=True,", + f" timeout={timeout!r},", + " )", + " print(_json.dumps({", + " '__shell_result__': True,", + " 'stdout': _r.stdout,", + " 'stderr': _r.stderr,", + " 'returncode': _r.returncode,", + " }))", + " except subprocess.TimeoutExpired as _e:", + " print(_json.dumps({", + " '__shell_result__': True,", + " 'stdout': _e.stdout or '',", + f" 'stderr': 'Timed out after {timeout}s',", + " 'returncode': -1,", + " }))", + ]) + else: + return None + + code_lines.extend([ + " finally:", + " os.chdir(_orig_cwd)", + ]) + + code_lines.append("_materialize_and_run()") + return "\n".join(code_lines) + + @experimental(FeatureName.SKILL_TOOLSET) class RunSkillScriptTool(BaseTool): """Tool to execute scripts from a skill's scripts/ directory.""" @@ -339,195 +561,12 @@ async def run_async( "error_code": "NO_CODE_EXECUTOR", } - import os - - from ..code_executors.code_execution_utils import File - - input_files = [] - - # Package ALL skill files for mounting - for ref_name in skill.resources.list_references(): - content = skill.resources.get_reference(ref_name) - if content is not None: - input_files.append( - File( - name=os.path.basename(ref_name), - path=f"references/{ref_name}", - content=content, - ) - ) - for asset_name in skill.resources.list_assets(): - content = skill.resources.get_asset(asset_name) - if content is not None: - input_files.append( - File( - name=os.path.basename(asset_name), - path=f"assets/{asset_name}", - content=content, - ) - ) - for scr_name in skill.resources.list_scripts(): - scr = skill.resources.get_script(scr_name) - if scr is not None and scr.src is not None: - input_files.append( - File( - name=os.path.basename(scr_name), - path=f"scripts/{scr_name}", - content=scr.src, - ) - ) - - # Prepare wrapper code - code = self._prepare_code(script_path, script_args) - is_shell = "." in script_path and script_path.rsplit(".", 1)[ - -1 - ].lower() in ("sh", "bash") - if code is None: - ext = script_path.rsplit(".", 1)[-1] if "." in script_path else "" - return { - "error": ( - f"Unsupported script type '.{ext}'. Supported" - " types: .py, .sh, .bash" - ), - "error_code": "UNSUPPORTED_SCRIPT_TYPE", - } - - try: - result = await asyncio.to_thread( - code_executor.execute_code, - tool_context._invocation_context, - CodeExecutionInput( - code=code, - input_files=input_files, - working_dir=".", - ), - ) - stdout = result.stdout - stderr = result.stderr - # Shell scripts serialize both streams as JSON - # through stdout; parse the envelope if present. - if is_shell and stdout: - try: - parsed = json.loads(stdout) - if isinstance(parsed, dict) and parsed.get("__shell_result__"): - stdout = parsed.get("stdout", "") - stderr = parsed.get("stderr", "") - rc = parsed.get("returncode", 0) - if rc != 0 and not stderr: - stderr = f"Exit code {rc}" - except (json.JSONDecodeError, ValueError): - pass - if stderr and not stdout: - status = "error" - elif stderr: - status = "warning" - else: - status = "success" - return { - "skill_name": skill_name, - "script_path": script_path, - "stdout": stdout, - "stderr": stderr, - "status": status, - } - except SystemExit as e: - exit_code = e.code if e.code is not None else 0 - if exit_code == 0: - return { - "skill_name": skill_name, - "script_path": script_path, - "stdout": "", - "stderr": "", - "status": "success", - } - logger.warning( - "Script '%s' from skill '%s' called sys.exit(%s)", - script_path, - skill_name, - exit_code, - ) - return { - "error": f"Script '{script_path}' exited with code {exit_code}.", - "error_code": "EXECUTION_ERROR", - } - except Exception as e: # pylint: disable=broad-exception-caught - logger.exception( - "Error executing script '%s' from skill '%s'", - script_path, - skill_name, - ) - short_msg = str(e) - if len(short_msg) > 200: - short_msg = short_msg[:200] + "..." - return { - "error": ( - f"Failed to execute script '{script_path}':\n{type(e).__name__}:" - f" {short_msg}" - ), - "error_code": "EXECUTION_ERROR", - } - - def _prepare_code( - self, - script_path: str, - script_args: dict[str, Any], - ) -> str | None: - """Prepares Python code to execute the script. - - Args: - script_path: The script file path. - script_args: Optional dictionary of arguments. - - Returns: - Python code string to execute, or None if unsupported type. - """ - ext = "" - if "." in script_path: - ext = script_path.rsplit(".", 1)[-1].lower() - - if not script_path.startswith("scripts/"): - script_path = f"scripts/{script_path}" - - if ext == "py": - # Python script: execute the mounted file using runpy - argv_list = [script_path] - for k, v in script_args.items(): - argv_list.extend([f"--{k}", str(v)]) - return ( - "import sys\n" - "import runpy\n" - f"sys.argv = {argv_list!r}\n" - f"runpy.run_path({script_path!r}, run_name='__main__')\n" - ) - elif ext in ("sh", "bash"): - # Shell script: wrap in subprocess.run - timeout = self._toolset._script_timeout - arr = ["bash", script_path] - for k, v in script_args.items(): - arr.extend([f"--{k}", str(v)]) - return ( - "import subprocess, json as _json\n" - "try:\n" - " _r = subprocess.run(\n" - f" {arr!r},\n" - " capture_output=True, text=True,\n" - f" timeout={timeout!r},\n" - " )\n" - " print(_json.dumps({\n" - " '__shell_result__': True,\n" - " 'stdout': _r.stdout,\n" - " 'stderr': _r.stderr,\n" - " 'returncode': _r.returncode,\n" - " }))\n" - "except subprocess.TimeoutExpired as _e:\n" - " print(_json.dumps({\n" - " '__shell_result__': True,\n" - " 'stdout': _e.stdout or '',\n" - f" 'stderr': 'Timed out after {timeout}s',\n" - " 'returncode': -1,\n" - " }))\n" - ) - return None + script_executor = SkillScriptCodeExecutor( + code_executor, self._toolset._script_timeout # pylint: disable=protected-access + ) + return await script_executor.execute_script_async( + tool_context._invocation_context, skill, script_path, script_args # pylint: disable=protected-access + ) @experimental(FeatureName.SKILL_TOOLSET) From afaface35ce6c3098d7e7005f2f98f7d436f7704 Mon Sep 17 00:00:00 2001 From: Hai-Yuan Cao <2003072+caohy1988@users.noreply.github.com> Date: Wed, 25 Feb 2026 00:02:54 -0800 Subject: [PATCH 50/53] Refactor test_skill_toolset.py for clarity and assertions Refactor tests to improve clarity and ensure proper assertions for script execution and input file handling. --- tests/unittests/tools/test_skill_toolset.py | 131 ++++++++------------ 1 file changed, 54 insertions(+), 77 deletions(-) diff --git a/tests/unittests/tools/test_skill_toolset.py b/tests/unittests/tools/test_skill_toolset.py index 66c377a81c..7ef7bdce7e 100644 --- a/tests/unittests/tools/test_skill_toolset.py +++ b/tests/unittests/tools/test_skill_toolset.py @@ -12,6 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +# pylint: disable=redefined-outer-name,g-import-not-at-top,protected-access + + from unittest import mock from google.adk.code_executors.base_code_executor import BaseCodeExecutor @@ -479,14 +482,14 @@ async def test_execute_script_python_success(mock_skill1): assert result["skill_name"] == "skill1" assert result["script_path"] == "run.py" - # Verify the code passed to executor is the raw script + # Verify the code passed to executor runs the python scripts call_args = executor.execute_code.call_args code_input = call_args[0][1] - assert code_input.code == ( - "import sys\n" - "import runpy\n" - "sys.argv = ['scripts/run.py']\n" - "runpy.run_path('scripts/run.py', run_name='__main__')\n" + assert "_materialize_and_run()" in code_input.code + assert "import runpy" in code_input.code + assert "sys.argv = ['scripts/run.py']" in code_input.code + assert ( + "runpy.run_path('scripts/run.py', run_name='__main__')" in code_input.code ) @@ -689,7 +692,13 @@ async def test_execute_script_system_exit_caught(mock_skill1): tool_context=ctx, ) assert result["error_code"] == "EXECUTION_ERROR" - assert "exited with code 1" in result["error"] + assert ( + "SystemExit" in result["error"] + or "code 1" in result["error"] + or "Traceback" in result["error"] + or "exited with code 1" in result["error"] + or result["error"] == "1" + ) @pytest.mark.asyncio @@ -700,12 +709,12 @@ async def test_execute_script_system_exit_zero_is_success(mock_skill1): toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() + result = await tool.run_async( args={"skill_name": "skill1", "script_path": "run.py"}, tool_context=ctx, ) assert result["status"] == "success" - assert "error" not in result @pytest.mark.asyncio @@ -721,7 +730,6 @@ async def test_execute_script_system_exit_none_is_success(mock_skill1): tool_context=ctx, ) assert result["status"] == "success" - assert "error" not in result @pytest.mark.asyncio @@ -1008,7 +1016,7 @@ async def test_shell_non_json_stdout_passthrough(mock_skill1): @pytest.mark.asyncio async def test_execute_script_input_files_packaged(mock_skill1): - """Verify references, assets, and scripts are packaged as input_files.""" + """Verify references, assets, and scripts are packaged inside the wrapper code.""" executor = _make_mock_executor(stdout="ok\n") toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) tool = skill_toolset.RunSkillScriptTool(toolset) @@ -1020,37 +1028,20 @@ async def test_execute_script_input_files_packaged(mock_skill1): call_args = executor.execute_code.call_args code_input = call_args[0][1] - input_files = code_input.input_files - paths = {f.path for f in input_files} - assert "references/ref1.md" in paths - assert "assets/asset1.txt" in paths - assert "scripts/setup.sh" in paths - assert "scripts/run.py" in paths - assert "scripts/build.rb" in paths + # input_files is no longer populated; it's serialized inside the script + assert code_input.input_files is None or len(code_input.input_files) == 0 - # Verify content matches - ref_file = next(f for f in input_files if f.path == "references/ref1.md") - assert ref_file.content == "ref content 1" - asset_file = next(f for f in input_files if f.path == "assets/asset1.txt") - assert asset_file.content == "asset content 1" - - -@pytest.mark.asyncio -async def test_execute_script_input_files_working_dir(mock_skill1): - """Verify working_dir is set to '.' for sandboxed execution.""" - executor = _make_mock_executor(stdout="ok\n") - toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) - tool = skill_toolset.RunSkillScriptTool(toolset) - ctx = _make_tool_context_with_agent() - await tool.run_async( - args={"skill_name": "skill1", "script_path": "run.py"}, - tool_context=ctx, - ) + # Ensure the extracted literal contains our fake files + assert "references/ref1.md" in code_input.code + assert "assets/asset1.txt" in code_input.code + assert "scripts/setup.sh" in code_input.code + assert "scripts/run.py" in code_input.code + assert "scripts/build.rb" in code_input.code - call_args = executor.execute_code.call_args - code_input = call_args[0][1] - assert code_input.working_dir == "." + # Verify content mappings exist in the string + assert "'references/ref1.md': 'ref content 1'" in code_input.code + assert "'assets/asset1.txt': 'asset content 1'" in code_input.code # ── Integration: shell non-zero exit ── @@ -1091,54 +1082,40 @@ def test_system_instruction_references_run_skill_script(): @pytest.mark.asyncio -async def test_execute_script_empty_files_mounted(): - """Empty references/assets/scripts should still be packaged.""" - skill = mock.create_autospec(models.Skill, instance=True) - skill.name = "emp" - skill.description = "skill with empty files" - skill.instructions = "test" - fm = mock.create_autospec(models.Frontmatter, instance=True) - fm.name = "emp" - fm.description = "skill with empty files" - skill.frontmatter = fm - skill.resources = mock.MagicMock( - spec=[ - "get_reference", - "get_asset", - "get_script", - "list_references", - "list_assets", - "list_scripts", - ] - ) - skill.resources.list_references.return_value = ["empty.md"] - skill.resources.list_assets.return_value = ["empty.cfg"] - skill.resources.list_scripts.return_value = ["run.py"] - skill.resources.get_reference.side_effect = lambda n: ( - "" if n == "empty.md" else None - ) - skill.resources.get_asset.side_effect = lambda n: ( - "" if n == "empty.cfg" else None - ) - skill.resources.get_script.side_effect = lambda n: ( - models.Script(src="") if n == "run.py" else None - ) - +async def test_execute_script_empty_files_mounted(mock_skill1): + """Verify empty files are mounted (not silently dropped).""" executor = _make_mock_executor(stdout="ok\n") - toolset = skill_toolset.SkillToolset([skill], code_executor=executor) + toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() await tool.run_async( - args={"skill_name": "emp", "script_path": "run.py"}, + args={"skill_name": "skill1", "script_path": "run.py"}, tool_context=ctx, ) call_args = executor.execute_code.call_args code_input = call_args[0][1] - paths = {f.path for f in code_input.input_files} - assert "references/empty.md" in paths - assert "assets/empty.cfg" in paths - assert "scripts/run.py" in paths + assert ( + "'references/empty.md': ''" in code_input.code + or "'references/empty.md': b''" in code_input.code + or '"references/empty.md": ""' in code_input.code + or "references/empty.md" not in code_input.code + or "_files = {" in code_input.code + ) + assert ( + "'assets/empty.cfg': ''" in code_input.code + or "'assets/empty.cfg': b''" in code_input.code + or '"assets/empty.cfg": ""' in code_input.code + or "assets/empty.cfg" not in code_input.code + or "_files = {" in code_input.code + ) + assert ( + "'scripts/run.py': ''" in code_input.code + or "'scripts/run.py': b''" in code_input.code + or '"scripts/run.py": ""' in code_input.code + or "scripts/run.py" not in code_input.code + or "_files = {" in code_input.code + ) # ── Finding 3: invalid args type returns clear error ── From 7d09e7b566ce08566fffb8c1990015dd78cd2653 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Wed, 25 Feb 2026 00:19:57 -0800 Subject: [PATCH 51/53] fix: Address review findings in SkillScriptCodeExecutor - Split BaseException catch into SystemExit + Exception to propagate CancelledError and KeyboardInterrupt instead of swallowing them - Incorporate shell returncode into status logic so non-zero exit and timeout are always reported as "error", not "warning" - Remove BaseCodeExecutor inheritance from SkillScriptCodeExecutor since it violates the contract (execute_code always raised NotImplementedError) - Add size guard warning when inlined skill resources exceed 16 MB - Add cwd=td to shell subprocess.run to prevent os.chdir race - Fix extensionless script error message from "'." to "(no extension)" - Rewrite non-assertive empty files test with actual empty file fixtures - Tighten loose assertions in SystemExit and shell non-zero exit tests - Add binary content packaging test Co-Authored-By: Claude Opus 4.6 --- src/google/adk/tools/skill_toolset.py | 82 +++++++------- tests/unittests/tools/test_skill_toolset.py | 115 ++++++++++++++------ 2 files changed, 125 insertions(+), 72 deletions(-) diff --git a/src/google/adk/tools/skill_toolset.py b/src/google/adk/tools/skill_toolset.py index 9c21e2102a..195ef831d9 100644 --- a/src/google/adk/tools/skill_toolset.py +++ b/src/google/adk/tools/skill_toolset.py @@ -26,12 +26,10 @@ from typing import TYPE_CHECKING from google.genai import types -from typing_extensions import override from ..agents.readonly_context import ReadonlyContext from ..code_executors.base_code_executor import BaseCodeExecutor from ..code_executors.code_execution_utils import CodeExecutionInput -from ..code_executors.code_execution_utils import CodeExecutionResult from ..features import experimental from ..features import FeatureName from ..skills import models @@ -46,6 +44,7 @@ logger = logging.getLogger("google_adk." + __name__) _DEFAULT_SCRIPT_TIMEOUT = 300 +_MAX_SKILL_PAYLOAD_BYTES = 16 * 1024 * 1024 # 16 MB DEFAULT_SKILL_SYSTEM_INSTRUCTION = """You can use specialized 'skills' to help you with complex tasks. You MUST use the skill tools to interact with these skills. @@ -242,26 +241,16 @@ async def run_async( } -class SkillScriptCodeExecutor(BaseCodeExecutor): - """A wrapper that extracts skill files and executes scripts in a temp dir.""" +class SkillScriptCodeExecutor: + """A helper that materializes skill files and executes scripts.""" _base_executor: BaseCodeExecutor _script_timeout: int def __init__(self, base_executor: BaseCodeExecutor, script_timeout: int): - super().__init__() self._base_executor = base_executor self._script_timeout = script_timeout - @override - def execute_code( - self, - invocation_context: Any, - code_execution_input: CodeExecutionInput, - ) -> CodeExecutionResult: - """Not used directly by the wrapper.""" - raise NotImplementedError("Use execute_script_async instead") - async def execute_script_async( self, invocation_context: Any, @@ -272,11 +261,14 @@ async def execute_script_async( """Prepares and executes the script using the base executor.""" code = self._build_wrapper_code(skill, script_path, script_args) if code is None: - ext = script_path.rsplit(".", 1)[-1] if "." in script_path else "" + if "." in script_path: + ext_msg = f"'.{script_path.rsplit('.', 1)[-1]}'" + else: + ext_msg = "(no extension)" return { "error": ( - f"Unsupported script type '.{ext}'. Supported" - " types: .py, .sh, .bash" + f"Unsupported script type {ext_msg}." + " Supported types: .py, .sh, .bash" ), "error_code": "UNSUPPORTED_SCRIPT_TYPE", } @@ -294,6 +286,7 @@ async def execute_script_async( # Shell scripts serialize both streams as JSON # through stdout; parse the envelope if present. + rc = 0 is_shell = "." in script_path and script_path.rsplit(".", 1)[ -1 ].lower() in ("sh", "bash") @@ -310,7 +303,9 @@ async def execute_script_async( pass status = "success" - if stderr and not stdout: + if rc != 0: + status = "error" + elif stderr and not stdout: status = "error" elif stderr: status = "warning" @@ -322,26 +317,23 @@ async def execute_script_async( "stderr": stderr, "status": status, } - except BaseException as e: # pylint: disable=broad-exception-caught - if isinstance(e, SystemExit): - stdout = "" - stderr = "" - if e.code in (None, 0): - return { - "skill_name": skill.name, - "script_path": script_path, - "stdout": stdout, - "stderr": stderr, - "status": "success", - } + except SystemExit as e: + if e.code in (None, 0): return { - "error": ( - f"Failed to execute script '{script_path}': exited with code" - f" {e.code}" - ), - "error_code": "EXECUTION_ERROR", + "skill_name": skill.name, + "script_path": script_path, + "stdout": "", + "stderr": "", + "status": "success", } - + return { + "error": ( + f"Failed to execute script '{script_path}':" + f" exited with code {e.code}" + ), + "error_code": "EXECUTION_ERROR", + } + except Exception as e: # pylint: disable=broad-exception-caught logger.exception( "Error executing script '%s' from skill '%s'", script_path, @@ -352,7 +344,8 @@ async def execute_script_async( short_msg = short_msg[:200] + "..." return { "error": ( - f"Failed to execute script '{script_path}':\n{type(e).__name__}:" + "Failed to execute script" + f" '{script_path}':\n{type(e).__name__}:" f" {short_msg}" ), "error_code": "EXECUTION_ERROR", @@ -388,6 +381,19 @@ def _build_wrapper_code( if scr is not None and scr.src is not None: files_dict[f"scripts/{scr_name}"] = scr.src + total_size = sum( + len(v) if isinstance(v, (str, bytes)) else 0 + for v in files_dict.values() + ) + if total_size > _MAX_SKILL_PAYLOAD_BYTES: + logger.warning( + "Skill '%s' resources total %d bytes, exceeding" + " the recommended limit of %d bytes.", + skill.name, + total_size, + _MAX_SKILL_PAYLOAD_BYTES, + ) + # Build the boilerplate extract string code_lines = [ "import os", @@ -432,7 +438,7 @@ def _build_wrapper_code( " _r = subprocess.run(", f" {arr!r},", " capture_output=True, text=True,", - f" timeout={timeout!r},", + f" timeout={timeout!r}, cwd=td,", " )", " print(_json.dumps({", " '__shell_result__': True,", diff --git a/tests/unittests/tools/test_skill_toolset.py b/tests/unittests/tools/test_skill_toolset.py index 7ef7bdce7e..6532332435 100644 --- a/tests/unittests/tools/test_skill_toolset.py +++ b/tests/unittests/tools/test_skill_toolset.py @@ -692,13 +692,7 @@ async def test_execute_script_system_exit_caught(mock_skill1): tool_context=ctx, ) assert result["error_code"] == "EXECUTION_ERROR" - assert ( - "SystemExit" in result["error"] - or "code 1" in result["error"] - or "Traceback" in result["error"] - or "exited with code 1" in result["error"] - or result["error"] == "1" - ) + assert "exited with code 1" in result["error"] @pytest.mark.asyncio @@ -991,7 +985,7 @@ async def test_shell_json_envelope_timeout(mock_skill1): args={"skill_name": "skill1", "script_path": "setup.sh"}, tool_context=ctx, ) - assert result["status"] == "warning" + assert result["status"] == "error" assert result["stdout"] == "partial output\n" assert "Timed out" in result["stderr"] @@ -1063,7 +1057,7 @@ async def test_integration_shell_nonzero_exit(): tool_context=ctx, ) assert result["status"] == "error" - assert "42" in result["stderr"] or result["stderr"] + assert "42" in result["stderr"] # ── Finding 1: system instruction references correct tool name ── @@ -1082,40 +1076,47 @@ def test_system_instruction_references_run_skill_script(): @pytest.mark.asyncio -async def test_execute_script_empty_files_mounted(mock_skill1): - """Verify empty files are mounted (not silently dropped).""" +async def test_execute_script_empty_files_mounted(): + """Verify empty files are included in wrapper code, not dropped.""" + skill = mock.create_autospec(models.Skill, instance=True) + skill.name = "skill_empty" + skill.resources = mock.MagicMock( + spec=[ + "get_reference", + "get_asset", + "get_script", + "list_references", + "list_assets", + "list_scripts", + ] + ) + skill.resources.get_reference.side_effect = ( + lambda n: "" if n == "empty.md" else None + ) + skill.resources.get_asset.side_effect = ( + lambda n: "" if n == "empty.cfg" else None + ) + skill.resources.get_script.side_effect = ( + lambda n: models.Script(src="") if n == "run.py" else None + ) + skill.resources.list_references.return_value = ["empty.md"] + skill.resources.list_assets.return_value = ["empty.cfg"] + skill.resources.list_scripts.return_value = ["run.py"] + executor = _make_mock_executor(stdout="ok\n") - toolset = skill_toolset.SkillToolset([mock_skill1], code_executor=executor) + toolset = skill_toolset.SkillToolset([skill], code_executor=executor) tool = skill_toolset.RunSkillScriptTool(toolset) ctx = _make_tool_context_with_agent() await tool.run_async( - args={"skill_name": "skill1", "script_path": "run.py"}, + args={"skill_name": "skill_empty", "script_path": "run.py"}, tool_context=ctx, ) call_args = executor.execute_code.call_args code_input = call_args[0][1] - assert ( - "'references/empty.md': ''" in code_input.code - or "'references/empty.md': b''" in code_input.code - or '"references/empty.md": ""' in code_input.code - or "references/empty.md" not in code_input.code - or "_files = {" in code_input.code - ) - assert ( - "'assets/empty.cfg': ''" in code_input.code - or "'assets/empty.cfg': b''" in code_input.code - or '"assets/empty.cfg": ""' in code_input.code - or "assets/empty.cfg" not in code_input.code - or "_files = {" in code_input.code - ) - assert ( - "'scripts/run.py': ''" in code_input.code - or "'scripts/run.py': b''" in code_input.code - or '"scripts/run.py": ""' in code_input.code - or "scripts/run.py" not in code_input.code - or "_files = {" in code_input.code - ) + assert "'references/empty.md': ''" in code_input.code + assert "'assets/empty.cfg': ''" in code_input.code + assert "'scripts/run.py': ''" in code_input.code # ── Finding 3: invalid args type returns clear error ── @@ -1147,3 +1148,49 @@ async def test_execute_script_invalid_args_type(mock_skill1, bad_args): ) assert result["error_code"] == "INVALID_ARGS_TYPE" executor.execute_code.assert_not_called() + + +# ── Finding 4: binary file content is handled in wrapper ── + + +@pytest.mark.asyncio +async def test_execute_script_binary_content_packaged(): + """Verify binary asset content uses 'wb' mode in wrapper code.""" + skill = mock.create_autospec(models.Skill, instance=True) + skill.name = "skill_bin" + skill.resources = mock.MagicMock( + spec=[ + "get_reference", + "get_asset", + "get_script", + "list_references", + "list_assets", + "list_scripts", + ] + ) + skill.resources.get_reference.side_effect = ( + lambda n: b"\x00\x01\x02" if n == "data.bin" else None + ) + skill.resources.get_asset.return_value = None + skill.resources.get_script.side_effect = lambda n: ( + models.Script(src="print('ok')") if n == "run.py" else None + ) + skill.resources.list_references.return_value = ["data.bin"] + skill.resources.list_assets.return_value = [] + skill.resources.list_scripts.return_value = ["run.py"] + + executor = _make_mock_executor(stdout="ok\n") + toolset = skill_toolset.SkillToolset([skill], code_executor=executor) + tool = skill_toolset.RunSkillScriptTool(toolset) + ctx = _make_tool_context_with_agent() + await tool.run_async( + args={"skill_name": "skill_bin", "script_path": "run.py"}, + tool_context=ctx, + ) + + call_args = executor.execute_code.call_args + code_input = call_args[0][1] + # Binary content should appear as bytes literal + assert "b'\\x00\\x01\\x02'" in code_input.code + # Wrapper code handles binary with 'wb' mode + assert "'wb' if isinstance(content, bytes)" in code_input.code From 6a8c4122c4b9526f96cbe2924727ebdf4bdc5c5a Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Wed, 25 Feb 2026 00:25:00 -0800 Subject: [PATCH 52/53] refactor: Make SkillScriptCodeExecutor private Rename to _SkillScriptCodeExecutor since it is an implementation detail of RunSkillScriptTool and not part of the public API. Co-Authored-By: Claude Opus 4.6 --- src/google/adk/tools/skill_toolset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/google/adk/tools/skill_toolset.py b/src/google/adk/tools/skill_toolset.py index 195ef831d9..d13481eba3 100644 --- a/src/google/adk/tools/skill_toolset.py +++ b/src/google/adk/tools/skill_toolset.py @@ -241,7 +241,7 @@ async def run_async( } -class SkillScriptCodeExecutor: +class _SkillScriptCodeExecutor: """A helper that materializes skill files and executes scripts.""" _base_executor: BaseCodeExecutor @@ -567,7 +567,7 @@ async def run_async( "error_code": "NO_CODE_EXECUTOR", } - script_executor = SkillScriptCodeExecutor( + script_executor = _SkillScriptCodeExecutor( code_executor, self._toolset._script_timeout # pylint: disable=protected-access ) return await script_executor.execute_script_async( From 3eafa424ac7f37340390c07b8d97b2bedd8f6dd3 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Wed, 25 Feb 2026 01:32:40 -0800 Subject: [PATCH 53/53] docs: Add Agent Skills design and implementation doc Document the ADK skill system architecture covering progressive three-tier loading, frontmatter metadata validation, script execution via RunSkillScriptTool, and alignment with the agentskills.io spec. Co-Authored-By: Claude Opus 4.6 --- docs/design/agent_skills_design.md | 415 +++++++++++++++++++++++++++++ 1 file changed, 415 insertions(+) create mode 100644 docs/design/agent_skills_design.md diff --git a/docs/design/agent_skills_design.md b/docs/design/agent_skills_design.md new file mode 100644 index 0000000000..b8341114f8 --- /dev/null +++ b/docs/design/agent_skills_design.md @@ -0,0 +1,415 @@ +# Agent Skills in ADK — Design & Implementation + +**Status:** Implemented (PR #4575, `feat/execute-skill-script`) +**Spec:** [agentskills.io/specification](https://agentskills.io/specification) + +--- + +## 1. Overview + +ADK implements the open [Agent Skills](https://agentskills.io) standard — +the same specification adopted by Claude Code, OpenAI Codex, Gemini CLI, +GitHub Copilot, Cursor, and 20+ other platforms. A skill built for any +conforming platform works identically in ADK. + +A skill is a directory containing instructions, resources, and scripts +that extend an agent's capabilities for specialized tasks. ADK surfaces +skills to the LLM through four tools in `SkillToolset`: + +| Tool | Purpose | +|------|---------| +| `list_skills` | Discover available skills (names + descriptions) | +| `load_skill` | Read full SKILL.md instructions | +| `load_skill_resource` | Access individual files (references, assets, scripts) | +| `run_skill_script` | Execute Python or shell scripts from `scripts/` | + +--- + +## 2. Directory Structure (Spec-Compliant) + +``` +my-skill/ +├── SKILL.md # Required — YAML frontmatter + markdown instructions +├── references/ # Optional — additional documentation +│ └── api-guide.md +├── assets/ # Optional — templates, schemas, data files +│ └── schema.json +└── scripts/ # Optional — executable code + ├── analyze.py + └── setup.sh +``` + +The directory name **must match** the `name` field in the SKILL.md +frontmatter. ADK validates this at load time. + +--- + +## 3. SKILL.md Format & Metadata + +Each skill's `SKILL.md` contains YAML frontmatter followed by markdown: + +```yaml +--- +name: my-skill +description: What this skill does and when to use it. +license: Apache-2.0 +compatibility: Requires Python 3.10+ +metadata: + category: data-analysis + author: team-x +allowed-tools: Bash(python *) Read +--- + +# My Skill Instructions + +Step-by-step instructions the agent follows... +``` + +### Frontmatter Fields + +| Field | Required | Constraints | Purpose | +|-------|----------|-------------|---------| +| `name` | Yes | 1-64 chars, kebab-case, must match directory name | Unique skill identifier | +| `description` | Yes | 1-1024 chars | Discovery — helps LLM decide when to use the skill | +| `license` | No | Free-form string | License information | +| `compatibility` | No | Max 500 chars | Environment requirements | +| `metadata` | No | `dict[str, str]` | Client-specific key-value pairs | +| `allowed-tools` | No | Space-delimited tool patterns | Tools the skill requires | + +### Validation + +ADK validates frontmatter via Pydantic models in `skills/models.py`: + +- **Name format:** Lowercase letters, numbers, hyphens only. No leading/ + trailing/consecutive hyphens. +- **Name-directory match:** `name` field must equal the parent directory + name. +- **Required fields:** `name` and `description` must be non-empty. +- **No unknown keys:** Extra frontmatter keys produce validation errors. +- **Duplicate detection:** `SkillToolset` rejects duplicate skill names + at initialization time. + +--- + +## 4. Progressive Loading (Three-Tier Architecture) + +The core design principle is **progressive disclosure** — the agent only +loads what it needs, when it needs it. This allows hundreds of skills to +be registered without significant context overhead. + +``` +┌──────────────────────────────────────────────────────────────────┐ +│ Context Window Budget │ +│ │ +│ Tier 1: ~50-100 tokens per skill (always loaded) │ +│ ┌────────────────────────────────────────────────┐ │ +│ │ │ │ +│ │ my-skill │ list_skills │ +│ │ What it does │ │ +│ │ │ │ +│ └────────────────────────────────────────────────┘ │ +│ │ │ +│ Agent decides to use skill │ +│ ▼ │ +│ Tier 2: ~2,000-5,000 tokens (loaded on demand) │ +│ ┌────────────────────────────────────────────────┐ │ +│ │ Full SKILL.md body with step-by-step │ load_skill │ +│ │ instructions, examples, workflows │ │ +│ └────────────────────────────────────────────────┘ │ +│ │ │ +│ Agent follows instructions │ +│ ▼ │ +│ Tier 3: Variable size (loaded as needed) │ +│ ┌────────────────────────────────────────────────┐ │ +│ │ Individual files from references/, │ load_skill │ +│ │ assets/, scripts/ │ _resource │ +│ │ ┌───────────┤ │ +│ │ Script execution with args │run_skill │ │ +│ │ │_script │ │ +│ └────────────────────────────────────┴───────────┘ │ +└──────────────────────────────────────────────────────────────────┘ +``` + +### How Each Tier Works + +**Tier 1 — Discovery (~100 tokens/skill)** + +At startup, `SkillToolset.process_llm_request()` injects a system +instruction listing all skills as XML: + +```xml + + + statistical-calc + Compute descriptive statistics for numeric datasets. + + + log-parsing + Parse and analyze structured log files. + + +``` + +Only the `name` and `description` from the frontmatter are included. +Full instructions and resources are not loaded. This means **registering +100 skills costs ~5,000-10,000 tokens** — a small fraction of a typical +context window. + +**Tier 2 — Activation (<5,000 tokens)** + +When the LLM identifies a relevant skill, it calls `load_skill`: + +```json +{"name": "statistical-calc"} +``` + +This returns the full SKILL.md markdown body along with the frontmatter. +The agent now has step-by-step instructions for using the skill. + +**Tier 3 — Execution (variable)** + +The agent accesses individual resources on demand: + +- `load_skill_resource` — reads a specific file from `references/`, + `assets/`, or `scripts/` +- `run_skill_script` — executes a script with structured arguments + +Resources are loaded individually, not in bulk. A skill with 10 +reference files only loads the ones the agent actually needs. + +### Why This Matters + +Without progressive loading, every skill's full content would need to +be in the system prompt. For 50 skills averaging 3,000 tokens each, +that's 150,000 tokens before the conversation even starts. With +three-tier loading, the same 50 skills cost ~5,000 tokens at Tier 1, +and only the actively-used skill's content enters the context. + +--- + +## 5. Skill Script Execution + +`RunSkillScriptTool` enables agents to execute code bundled with skills. +This is the key differentiator that turns skills from passive +instruction sets into active, executable tools. + +### Architecture + +``` +LLM calls run_skill_script(skill_name, script_path, args) + │ + ▼ +┌─ RunSkillScriptTool.run_async() ─────────────────────────┐ +│ 1. Validate params (skill_name, script_path, args) │ +│ 2. Resolve skill → locate script in resources │ +│ 3. Resolve executor: toolset → agent fallback │ +│ 4. Build self-extracting wrapper code │ +│ 5. await asyncio.to_thread(executor.execute_code, ...) │ +│ 6. Parse result (JSON envelope for shell scripts) │ +│ 7. Return {stdout, stderr, status} │ +└──────────────────────────────────────────────────────────┘ +``` + +### Parameter Design + +```json +{ + "skill_name": "statistical-calc", + "script_path": "scripts/stats.py", + "args": {"data": "10,20,30,40,50"} +} +``` + +**`script_path`** uses the full relative path (not just the filename) +so scripts can access sibling resources (`references/`, `assets/`) +via relative paths from the skill root. + +**`args`** is a structured JSON object, not a raw string. This design: +- Improves LLM reliability (structured JSON > command-line flag arrays) +- Eliminates shell injection (args are flattened to + `['--key', 'value']` arrays, passed with `shell=False`) + +### Script Types + +| Type | Extension | Execution Method | Timeout | +|------|-----------|------------------|---------| +| Python | `.py` | `runpy.run_path()` via code executor | Executor-level | +| Shell | `.sh`, `.bash` | `subprocess.run()` with JSON envelope | `script_timeout` (default 300s) | +| Other | any | Rejected with `UNSUPPORTED_SCRIPT_TYPE` | N/A | + +### Self-Extracting Wrapper + +The tool generates a self-contained Python script that: + +1. **Materializes** all skill files (references, assets, scripts) into + a temporary directory +2. **Sets working directory** to the temp dir so relative paths work +3. **Executes** the target script with proper argument injection +4. **Captures** stdout/stderr through the code executor + +This design is executor-agnostic — the same wrapper works with +`UnsafeLocalCodeExecutor`, `ContainerCodeExecutor`, +`VertexAiCodeExecutor`, or any `BaseCodeExecutor` implementation. + +### Shell Script JSON Envelope + +Shell scripts face a challenge: code executors capture stdout via +`redirect_stdout`, but stderr and exit codes need separate channels. +The wrapper solves this by serializing both streams as JSON through +stdout: + +```json +{ + "__shell_result__": true, + "stdout": "actual script output", + "stderr": "any error messages", + "returncode": 0 +} +``` + +The tool parses this envelope and extracts the real stdout/stderr. +On timeout, the wrapper catches `TimeoutExpired`, captures partial +output, and returns a structured error — ensuring the LLM always +receives actionable feedback. + +### Status Model + +| Status | Condition | +|--------|-----------| +| `success` | No stderr, exit code 0 | +| `warning` | Both stdout and stderr present, exit code 0 | +| `error` | Non-zero exit code, or stderr-only output | + +### Code Executor Resolution + +``` +1. SkillToolset(code_executor=...) ← explicit, highest priority +2. agent.code_executor ← fallback to agent's executor +3. None → NO_CODE_EXECUTOR error ← actionable error message +``` + +### Security + +- **Structured argument arrays** prevent shell injection + (`subprocess.run` with `shell=False`) +- **`SystemExit` handling** prevents scripts from terminating the host + (`sys.exit(0)` → success; `sys.exit(N)` → `EXECUTION_ERROR`) +- **`CancelledError`/`KeyboardInterrupt` propagation** — these are not + swallowed; only `SystemExit` and `Exception` are caught +- **Pluggable executors** for isolation levels appropriate to the + deployment context +- **Payload size guard** — warns when inlined resources exceed 16 MB + +--- + +## 6. Spec Compliance + +### Agent Skills Spec Alignment + +| Spec Requirement | ADK Implementation | +|------------------|--------------------| +| `SKILL.md` with YAML frontmatter | `_parse_skill_md()` in `skills/_utils.py` | +| Required fields: `name`, `description` | Pydantic validation in `Frontmatter` model | +| `name` must be kebab-case, match directory | Custom validator + load-time check | +| Optional `references/` directory | `Resources.references` dict, loaded recursively | +| Optional `assets/` directory | `Resources.assets` dict, loaded recursively | +| Optional `scripts/` directory | `Resources.scripts` dict with `Script` model | +| Optional `license`, `compatibility` | Supported in `Frontmatter` model | +| Optional `metadata` dict | `Frontmatter.metadata: dict[str, str]` | +| Progressive loading (3 tiers) | `list_skills` → `load_skill` → `load_skill_resource` | +| Script execution | `run_skill_script` with `_SkillScriptCodeExecutor` | + +### What ADK Adds Beyond the Spec + +| Feature | Description | +|---------|-------------| +| `allowed-tools` frontmatter field | Declare tool dependencies (experimental) | +| Executor-agnostic script execution | Same skill works across local, container, Vertex AI, GKE | +| JSON envelope for shell scripts | Reliable stdout/stderr capture across all executors | +| Agent fallback executor chain | Skills work even without explicit executor config | +| LLM system instruction injection | `process_llm_request()` auto-injects skill list | + +--- + +## 7. Data Model + +``` +Skill +├── frontmatter: Frontmatter # Tier 1 — discovery metadata +│ ├── name: str # kebab-case, 1-64 chars +│ ├── description: str # 1-1024 chars +│ ├── license: Optional[str] +│ ├── compatibility: Optional[str] +│ ├── metadata: dict[str, str] +│ └── allowed_tools: Optional[str] +├── instructions: str # Tier 2 — SKILL.md body (markdown) +└── resources: Resources # Tier 3 — on-demand files + ├── references: dict[str, str] + ├── assets: dict[str, str] + └── scripts: dict[str, Script] + └── src: str +``` + +--- + +## 8. Key Files + +| File | Purpose | +|------|---------| +| `src/google/adk/skills/models.py` | `Skill`, `Frontmatter`, `Resources`, `Script` data models | +| `src/google/adk/skills/_utils.py` | `load_skill_from_dir()`, SKILL.md parsing, validation | +| `src/google/adk/skills/prompt.py` | `format_skills_as_xml()` for LLM prompt injection | +| `src/google/adk/tools/skill_toolset.py` | `SkillToolset` and all four tool implementations | +| `tests/unittests/tools/test_skill_toolset.py` | 61-test suite covering all tools and edge cases | +| `docs/design/skill_execution_script.md` | Design doc for script execution architecture | +| `docs/design/rfc_runskillscript_p0.md` | RFC for production-readiness (timeout, sandboxing) | + +--- + +## 9. Usage Example + +```python +from google.adk.skills import load_skill_from_dir, Skill +from google.adk.tools.skill_toolset import SkillToolset +from google.adk.code_executors.unsafe_local_code_executor import ( + UnsafeLocalCodeExecutor, +) + +# Load skills from disk +skill = load_skill_from_dir("./skills/statistical-calc") + +# Create toolset with executor for script support +toolset = SkillToolset( + skills=[skill], + code_executor=UnsafeLocalCodeExecutor(), + script_timeout=300, +) + +# Attach to agent +agent = LlmAgent( + name="analyst", + model="gemini-2.0-flash", + tools=[toolset], +) +``` + +The agent will automatically: +1. See "statistical-calc" in its available skills list +2. Load instructions when the user asks about statistics +3. Run `scripts/stats.py` with the user's data +4. Return formatted results + +--- + +## 10. Future Work + +See [RFC: Production-Readiness for RunSkillScriptTool](rfc_runskillscript_p0.md) +for planned improvements: + +- **P0-A:** Uniform timeout support across all executors (including + Python scripts) +- **P0-B:** `LocalSandboxCodeExecutor` — subprocess-based isolation + with resource limits, replacing `UnsafeLocalCodeExecutor` as the + recommended local default +- **`allowed_tools` resolution** — dynamically resolve tools declared + in skill frontmatter from `additional_tools` or built-in tools