diff --git a/tests/cli/test_cli_harness_open_source_defaults.py b/tests/cli/test_cli_harness_open_source_defaults.py index e28b46d42..952afa970 100644 --- a/tests/cli/test_cli_harness_open_source_defaults.py +++ b/tests/cli/test_cli_harness_open_source_defaults.py @@ -43,11 +43,10 @@ def test_harness_dockerfile_uses_accelerated_source_with_official_fallback() -> in cli_harness._DOCKERFILE ) assert "https://github.com/volcengine/veadk-python.git" in cli_harness._DOCKERFILE - # The harness advertises `runtime: codex` in harness.yaml and honours a - # per-request runtime override, so the image must carry the codex extra or - # every such request fails with an ImportError on an already-deployed - # runtime. - assert '"./src[harness,codex]"' in cli_harness._DOCKERFILE + # The generated image is the shared HarnessApp runtime: it must carry the + # optional backend extras used by request-level resources, plus codex for + # runtime overrides. + assert '"./src[extensions,database,harness,codex]"' in cli_harness._DOCKERFILE old_package_path = "packages/" + "agentkit" + "-harness-python" assert old_package_path not in cli_harness._DOCKERFILE diff --git a/tests/cli/test_frontend_trace.py b/tests/cli/test_frontend_trace.py index 99f3e10f1..397c5eb7d 100644 --- a/tests/cli/test_frontend_trace.py +++ b/tests/cli/test_frontend_trace.py @@ -15,7 +15,11 @@ from __future__ import annotations +import importlib import json +import sys +from collections.abc import Iterator +from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace from typing import cast @@ -55,14 +59,33 @@ def _write_agent_app(tmp_path: Path, app_name: str, source: str) -> None: (app_dir / "agent.py").write_text(source, encoding="utf-8") -def _build_adk_web_client(tmp_path: Path) -> TestClient: +_MISSING = object() + + +@contextmanager +def _build_adk_web_client(tmp_path: Path) -> Iterator[TestClient]: from google.adk.cli.fast_api import get_fast_api_app from veadk.utils.patches import patch_adk_build_graph_serialization + original_sys_path = list(sys.path) + app_names = [path.name for path in tmp_path.iterdir() if path.is_dir()] + module_names = {name for name in app_names} + module_names.update(f"{name}.agent" for name in app_names) + original_modules = {name: sys.modules.get(name, _MISSING) for name in module_names} patch_adk_build_graph_serialization() - app = get_fast_api_app(agents_dir=str(tmp_path), web=True) - return TestClient(app) + try: + app = get_fast_api_app(agents_dir=str(tmp_path), web=True) + with TestClient(app) as client: + yield client + finally: + sys.path[:] = original_sys_path + for name, module in original_modules.items(): + if module is _MISSING: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + importlib.invalidate_caches() def test_session_trace_route_returns_json_spans() -> None: @@ -177,7 +200,8 @@ def hello() -> str: """, ) - response = _build_adk_web_client(tmp_path).get("/dev/apps/demo_agent/build_graph") + with _build_adk_web_client(tmp_path) as client: + response = client.get("/dev/apps/demo_agent/build_graph") assert response.status_code == 200 payload = response.json() @@ -216,7 +240,8 @@ def test_build_graph_serializes_nested_veadk_agent_models(tmp_path: Path) -> Non """, ) - response = _build_adk_web_client(tmp_path).get("/dev/apps/nested_agent/build_graph") + with _build_adk_web_client(tmp_path) as client: + response = client.get("/dev/apps/nested_agent/build_graph") assert response.status_code == 200 payload = response.json() diff --git a/tests/cloud/test_harness_app_contract.py b/tests/cloud/test_harness_app_contract.py index 8e2f26770..145a64fb1 100644 --- a/tests/cloud/test_harness_app_contract.py +++ b/tests/cloud/test_harness_app_contract.py @@ -24,28 +24,61 @@ import time, so it is intentionally left out to keep these tests offline. """ +import json +import os from pathlib import Path +from types import SimpleNamespace +import pytest +from google.genai import types +from pydantic import ValidationError + +from veadk import Agent +from veadk.cloud.harness_app.agentkit_resources import ( + AgentKitMcpRouterResolver, + AgentKitResourceResolver, +) +from veadk.cloud.harness_app.env_mapping import to_runtime_env from veadk.cloud.harness_app.types import ( + HarnessAgentConfigRequest, + HarnessBuiltinTool, HarnessCompactionMetric, HarnessConfig, + HarnessCreateSessionRequest, HarnessEnhanceOverrides, + HarnessMcpServer, HarnessOverrides, HarnessPluginMetrics, + HarnessRegistryOverride, + HarnessResourceOverride, HarnessResponseMetrics, + HarnessSelectedSkill, InvokeHarnessRequest, InvokeHarnessResponse, LlmUsageMetrics, RunAgentRequest, ) -from veadk.cloud.harness_app.env_mapping import to_runtime_env from veadk.cloud.harness_app.utils import ( + ResourceResolutionError, agent_name_from_harness, config_from_env, + harness_overrides_from_env, + init_harness_agent, + merge_harness_overrides, + normalize_harness_overrides, + set_harness_mcp_router_resolver, + set_harness_resource_resolver, + spawn_harness_agent, + spawn_harness_run_agent, split_csv, ) from veadk.consts import DEFAULT_MODEL_AGENT_NAME +from veadk.knowledgebase import KnowledgeBase +from veadk.memory.long_term_memory import LongTermMemory +from veadk.memory.save_session_callback import save_session_to_long_term_memory from veadk.prompts.agent_default_prompt import DEFAULT_INSTRUCTION +from veadk.tools.builtin_tools.create_agent.sources.cloud import CloudCredentials +from veadk.tools.builtin_tools.load_knowledgebase import LoadKnowledgebaseTool def _fields(model) -> dict: @@ -53,42 +86,226 @@ def _fields(model) -> dict: return dict(model.model_fields) +def _after_agent_callbacks(agent: Agent) -> list[object]: + callback = agent.after_agent_callback + if callback is None: + return [] + if isinstance(callback, list): + return callback + return [callback] + + +@pytest.fixture(autouse=True) +def _offline_embedding_key(monkeypatch): + monkeypatch.setenv("MODEL_EMBEDDING_API_KEY", "test-embedding-key") + set_harness_resource_resolver(None) + set_harness_mcp_router_resolver(None) + yield + set_harness_resource_resolver(None) + set_harness_mcp_router_resolver(None) + + class TestHarnessOverrides: def test_fields(self): assert set(_fields(HarnessOverrides)) == { "model_name", "tools", + "builtin_tools", + "mcp_router_id", "skills", + "selected_skills", + "mcp", "system_prompt", "runtime", "registry_space_id", "registry_endpoint", "registry_region", "registry_top_k", + "registry", + "knowledgebase", + "longterm_memory", + "temperature", + "top_p", + "max_tokens", + "presence_penalty", + "frequency_penalty", + "penalty", + "max_llm_calls", } def test_defaults(self): fields = _fields(HarnessOverrides) assert fields["model_name"].default == DEFAULT_MODEL_AGENT_NAME assert fields["tools"].default == "" + assert HarnessOverrides().builtin_tools == [] + assert fields["mcp_router_id"].default == "" assert fields["skills"].default == "" + assert HarnessOverrides().selected_skills == [] + assert HarnessOverrides().mcp == [] assert fields["system_prompt"].default == "You are a helpful assistant." assert fields["runtime"].default == "adk" assert fields["registry_space_id"].default == "" assert fields["registry_endpoint"].default == "" assert fields["registry_region"].default == "" assert fields["registry_top_k"].default == 3 + assert fields["registry"].default is None + assert fields["knowledgebase"].default is None + assert fields["longterm_memory"].default is None + assert fields["temperature"].default is None + assert fields["top_p"].default is None + assert fields["max_tokens"].default is None + assert fields["presence_penalty"].default is None + assert fields["frequency_penalty"].default is None + assert fields["penalty"].default is None + assert fields["max_llm_calls"].default is None + + def test_max_llm_calls_has_only_minimum_limit(self): + assert HarnessOverrides(max_llm_calls=20).max_llm_calls == 20 + assert HarnessConfig(max_llm_calls=20).max_llm_calls == 20 + assert RunAgentRequest(user_id="u1", session_id="s1", max_llm_calls=20) + with pytest.raises(ValidationError): + HarnessOverrides(max_llm_calls=0) + with pytest.raises(ValidationError): + HarnessConfig(max_llm_calls=0) + with pytest.raises(ValidationError): + RunAgentRequest(user_id="u1", session_id="s1", max_llm_calls=0) - def test_tools_and_skills_are_csv_strings(self): - # The server splits these with split_csv(); they must stay plain strings, - # not lists, to keep the CLI/curl pass-through contract. + def test_legacy_tools_and_skills_are_still_csv_strings(self): h = HarnessOverrides() assert isinstance(h.tools, str) assert isinstance(h.skills, str) + def test_structured_tools_skills_and_mcp_are_accepted(self): + h = HarnessOverrides.model_validate( + { + "mcp_router_id": "mt-1", + "builtin_tools": [ + {"id": "web_search"}, + {"id": "run_code", "config": {"tool_id": "t-1"}}, + ], + "selected_skills": [ + {"source": "skillhub", "slug": "team/reporting"}, + { + "source": "skillspace", + "space_id": "ss-1", + "skill_id": "skill-1", + }, + ], + "mcp": [{"name": "db", "url": "http://db.test/mcp"}], + } + ) + + assert h.builtin_tools == [ + HarnessBuiltinTool(id="web_search"), + HarnessBuiltinTool(id="run_code", config={"tool_id": "t-1"}), + ] + assert h.mcp_router_id == "mt-1" + assert h.selected_skills[1] == HarnessSelectedSkill( + source="skillspace", + skill_space_id="ss-1", + skill_id="skill-1", + ) + assert h.mcp == [HarnessMcpServer(name="db", server_url="http://db.test/mcp")] + + def test_structured_registry_is_accepted_and_normalized(self): + h = HarnessOverrides.model_validate( + { + "registry": { + "space_id": "space-1", + "region": "cn-beijing", + "top_k": 5, + } + } + ) + + assert h.registry == HarnessRegistryOverride( + space_id="space-1", + region="cn-beijing", + top_k=5, + ) + + normalized = normalize_harness_overrides(h) + + assert normalized.registry is None + assert normalized.model_dump(mode="json", exclude_unset=True) == { + "registry_space_id": "space-1", + "registry_region": "cn-beijing", + "registry_top_k": 5, + } + + def test_legacy_csv_fields_normalize_to_structured_fields(self): + h = normalize_harness_overrides( + HarnessOverrides( + tools="web_search, run_code", + mcp_router_id="mt-1", + skills="team/reporting, space:ss-1", + ) + ) + + assert h.model_dump(mode="json", exclude_unset=True) == { + "builtin_tools": [ + {"id": "web_search", "config": {}}, + {"id": "run_code", "config": {}}, + {"id": "mcp_router", "config": {"mcp_router_id": "mt-1"}}, + ], + "mcp_router_id": "mt-1", + "selected_skills": [ + {"source": "skillhub", "slug": "team/reporting"}, + {"source": "skillspace", "skill_space_id": "ss-1"}, + ], + } + + def test_merge_harness_overrides_applies_session_then_current(self): + merged = merge_harness_overrides( + { + "model_name": "session-model", + "tools": "web_search", + "system_prompt": "session prompt", + "temperature": 0.2, + }, + HarnessOverrides( + builtin_tools=[{"id": "link_reader"}], + system_prompt="current prompt", + temperature=None, + top_p=0.8, + ), + ) + + assert merged.model_dump(mode="json", exclude_unset=True) == { + "model_name": "session-model", + "builtin_tools": [{"id": "link_reader", "config": {}}], + "system_prompt": "current prompt", + "temperature": None, + "top_p": 0.8, + } + + def test_long_term_memory_alias_is_accepted(self): + h = HarnessOverrides.model_validate( + { + "long_term_memory": { + "type": "local", + "id": "ltm-1", + "config": {"index": "ltm-index"}, + } + } + ) + + assert h.longterm_memory == HarnessResourceOverride( + type="local", id="ltm-1", config={"index": "ltm-index"} + ) + assert h.model_fields_set == {"longterm_memory"} + + def test_misspelled_knowledgebase_alias_is_accepted(self): + h = HarnessOverrides.model_validate( + {"konwledgebase": {"type": "local", "id": "kb-1"}} + ) + + assert h.knowledgebase == HarnessResourceOverride(type="local", id="kb-1") + assert h.model_fields_set == {"knowledgebase"} + def test_every_field_has_a_description(self): - # Descriptions are the single source of truth for the generated - # `veadk harness invoke` flags, so each field must carry one. + # Descriptions feed FastAPI schemas and the subset of CLI flags that are + # still generated from this model, so each field must carry one. for name, field in _fields(HarnessOverrides).items(): assert field.description, f"{name} is missing a description" @@ -120,6 +337,7 @@ def test_component_defaults(self): assert fields["knowledgebase_type"].default == "" assert fields["longterm_memory_type"].default == "" assert fields["shortterm_memory_type"].default == "local" + assert fields["max_llm_calls"].default == 10 assert fields["structured_tool_calls"].default is False assert fields["include_tools_every_turn"].default is True assert fields["registry_type"].default == "" @@ -157,11 +375,116 @@ def test_tool_calling_yaml_maps_to_runtime_env(self): { "structured_tool_calls": True, "include_tools_every_turn": True, + "mcp_router_id": "mt-1", } ) assert envs["STRUCTURED_TOOL_CALLS"] == "true" assert envs["INCLUDE_TOOLS_EVERY_TURN"] == "true" + assert envs["MCP_ROUTER_ID"] == "mt-1" + assert envs["TOOLS"] == "mcp_router" + + def test_doc_harness_block_maps_to_runtime_env(self): + envs = to_runtime_env( + { + "harness": { + "model_name": "doc-model", + "temperature": 0.3, + "top_p": 0.9, + "max_llm_calls": 8, + "builtin_tools": [ + { + "id": "run_code", + "config": { + "tool_id": "t-script-1", + "region": "cn-beijing", + }, + }, + { + "id": "mcp_router", + "config": { + "url": "http://router.test/mcp", + "api_key": "router-token", + }, + }, + ], + "selected_skills": [ + {"source": "skillhub", "slug": "team/reporting"}, + ], + "mcp": [{"name": "db", "server_url": "http://db.test/mcp"}], + "knowledgebase": { + "type": "viking", + "config": { + "index": "kb-viking-index", + "app_name": "kb-viking-index", + "project": "default", + "region": "cn-beijing", + "resource_id": "resource-xxx", + }, + }, + "longterm_memory": { + "type": "mem0", + "config": { + "index": "memory-index", + "app_name": "memory-index", + "api_key": "mem0-token", + "base_url": "https://api.mem0.ai", + }, + }, + } + } + ) + + assert envs["MODEL_AGENT_NAME"] == "doc-model" + assert envs["MODEL_NAME"] == "doc-model" + assert envs["MODEL_AGENT_TEMPERATURE"] == "0.3" + assert envs["MODEL_AGENT_TOP_P"] == "0.9" + assert envs["MAX_LLM_CALLS"] == "8" + assert envs["TOOLS"] == "run_code,mcp_router" + assert envs["AGENTKIT_TOOL_ID_SCRIPT"] == "t-script-1" + assert envs["AGENTKIT_TOOL_REGION"] == "cn-beijing" + assert envs["TOOL_MCP_ROUTER_URL"] == "http://router.test/mcp" + assert envs["TOOL_MCP_ROUTER_API_KEY"] == "router-token" + assert json.loads(envs["SELECTED_SKILLS_JSON"]) == [ + {"source": "skillhub", "slug": "team/reporting"} + ] + assert json.loads(envs["MCP_SERVERS_JSON"]) == [ + {"name": "db", "server_url": "http://db.test/mcp"} + ] + assert envs["KNOWLEDGEBASE_TYPE"] == "viking" + assert envs["DATABASE_VIKING_PROJECT"] == "default" + assert envs["DATABASE_VIKING_REGION"] == "cn-beijing" + assert envs["DATABASE_VIKING_RESOURCE_ID"] == "resource-xxx" + assert json.loads(envs["KNOWLEDGEBASE_CONFIG_JSON"]) == { + "index": "kb-viking-index", + "app_name": "kb-viking-index", + "project": "default", + "region": "cn-beijing", + "resource_id": "resource-xxx", + } + assert envs["LONG_TERM_MEMORY_TYPE"] == "mem0" + assert envs["DATABASE_MEM0_API_KEY"] == "mem0-token" + assert envs["DATABASE_MEM0_BASE_URL"] == "https://api.mem0.ai" + assert json.loads(envs["LONG_TERM_MEMORY_CONFIG_JSON"]) == { + "index": "memory-index", + "app_name": "memory-index", + "api_key": "mem0-token", + "base_url": "https://api.mem0.ai", + } + + def test_cli_style_model_yaml_maps_current_and_legacy_model_env(self): + envs = to_runtime_env( + { + "model": {"name": "cli-model"}, + "temperature": 0.2, + "top_p": 0.7, + } + ) + + assert envs["MODEL_AGENT_NAME"] == "cli-model" + assert envs["MODEL_NAME"] == "cli-model" + assert envs["MODEL_AGENT_TEMPERATURE"] == "0.2" + assert envs["MODEL_AGENT_TOP_P"] == "0.7" def test_config_from_env_reads_registry_fields(self, monkeypatch): monkeypatch.setenv("REGISTRY_TYPE", "agentkit_a2a") @@ -185,6 +508,131 @@ def test_config_from_env_reads_tool_calling_fields(self, monkeypatch): assert config.structured_tool_calls is True assert config.include_tools_every_turn is False + def test_config_from_env_reads_agentkit_runtime_fields(self, monkeypatch): + monkeypatch.setenv("MODEL_AGENT_NAME", "model-agent") + monkeypatch.setenv("MODEL_NAME", "legacy-model") + monkeypatch.setenv("MODEL_AGENT_TEMPERATURE", "0.2") + monkeypatch.setenv("MODEL_AGENT_TOP_P", "0.8") + monkeypatch.setenv( + "SELECTED_SKILLS_JSON", + '{"selected_skills":[{"source":"skillhub","slug":"team/reporting"}]}', + ) + monkeypatch.setenv( + "MCP_SERVERS_JSON", + '[{"name":"db","server_url":"http://db.test/mcp","bear_token":"tok"}]', + ) + monkeypatch.setenv("MCP_ROUTER_ID", "mt-1") + monkeypatch.setenv("KNOWLEDGEBASE_TYPE", "local") + monkeypatch.setenv("KNOWLEDGEBASE_ID", "kb-1") + monkeypatch.setenv( + "KNOWLEDGEBASE_CONFIG_JSON", + '{"index":"kb-index","top_k":7}', + ) + monkeypatch.setenv("LONG_TERM_MEMORY_TYPE", "local") + monkeypatch.setenv("LONG_TERM_MEMORY_ID", "mem-1") + monkeypatch.setenv( + "LONG_TERM_MEMORY_CONFIG_JSON", + '{"index":"memory-index"}', + ) + + config = config_from_env() + + assert config.model_name == "model-agent" + assert config.temperature == 0.2 + assert config.top_p == 0.8 + assert config.selected_skills == [ + HarnessSelectedSkill(source="skillhub", slug="team/reporting") + ] + assert config.mcp == [ + HarnessMcpServer( + name="db", + server_url="http://db.test/mcp", + bear_token="tok", + ) + ] + assert config.mcp_router_id == "mt-1" + assert config.knowledgebase == HarnessResourceOverride( + type="local", + id="kb-1", + config={"index": "kb-index", "top_k": 7}, + ) + assert config.longterm_memory == HarnessResourceOverride( + type="local", + id="mem-1", + config={"index": "memory-index"}, + ) + + def test_harness_overrides_from_env_uses_run_sse_harness_shape(self, monkeypatch): + monkeypatch.setenv("MODEL_AGENT_NAME", "model-agent") + monkeypatch.setenv("MODEL_AGENT_TEMPERATURE", "0.2") + monkeypatch.setenv("MODEL_AGENT_TOP_P", "0.8") + monkeypatch.setenv("MAX_LLM_CALLS", "4") + monkeypatch.setenv("TOOLS", "web_search") + monkeypatch.setenv("MCP_ROUTER_ID", "mt-1") + monkeypatch.setenv( + "SELECTED_SKILLS_JSON", + '[{"source":"skillhub","slug":"team/reporting"}]', + ) + monkeypatch.setenv("KNOWLEDGEBASE_TYPE", "viking") + monkeypatch.setenv("KNOWLEDGEBASE_ID", "kb-1") + monkeypatch.setenv( + "KNOWLEDGEBASE_CONFIG_JSON", + '{"index":"secret-kb-index","api_key":"secret"}', + ) + monkeypatch.setenv("LONG_TERM_MEMORY_TYPE", "mem0") + monkeypatch.setenv("LONG_TERM_MEMORY_ID", "memory-1") + monkeypatch.setenv( + "LONG_TERM_MEMORY_CONFIG_JSON", + '{"index":"secret-memory-index","api_key":"secret"}', + ) + + config = harness_overrides_from_env() + + assert config.model_name == "model-agent" + assert config.temperature == 0.2 + assert config.top_p == 0.8 + assert config.max_llm_calls == 4 + assert config.builtin_tools == [ + HarnessBuiltinTool(id="web_search"), + HarnessBuiltinTool(id="mcp_router", config={"mcp_router_id": "mt-1"}), + ] + assert config.selected_skills == [ + HarnessSelectedSkill(source="skillhub", slug="team/reporting") + ] + assert config.knowledgebase == HarnessResourceOverride( + type="viking", + id="kb-1", + config={"index": "secret-kb-index", "api_key": "secret"}, + ) + assert config.longterm_memory == HarnessResourceOverride( + type="mem0", + id="memory-1", + config={"index": "secret-memory-index", "api_key": "secret"}, + ) + + def test_resource_yaml_ids_map_to_runtime_env(self): + envs = to_runtime_env( + { + "konwledgebase": { + "type": "viking", + "_id": "kb-1", + "project": "default", + }, + "long_term_memory": { + "type": "mem0", + "id": "mem-1", + "base_url": "https://api.mem0.ai", + }, + } + ) + + assert envs["KNOWLEDGEBASE_TYPE"] == "viking" + assert envs["KNOWLEDGEBASE_ID"] == "kb-1" + assert envs["DATABASE_VIKING_PROJECT"] == "default" + assert envs["LONG_TERM_MEMORY_TYPE"] == "mem0" + assert envs["LONG_TERM_MEMORY_ID"] == "mem-1" + assert envs["DATABASE_MEM0_BASE_URL"] == "https://api.mem0.ai" + def test_registry_overrides_remount_registry_tools(self): source = Path("veadk/cloud/harness_app/utils.py").read_text() @@ -217,8 +665,768 @@ def test_registry_request_auth_is_bound_to_run_agent_config(self): assert "use_registry_tip_token" not in registry_source assert "use_registry_tip_token" not in app_source + def test_spawn_mounts_registry_tools_from_structured_registry(self): + base = Agent(model_name="base-model", model_api_key="test-key") + + cloned = spawn_harness_agent( + base, + HarnessOverrides.model_validate( + { + "registry": { + "space_id": "space-test", + "region": "cn-beijing", + "top_k": 5, + } + } + ), + ) + + registry_config = cloned._veadk_a2a_registry_config + assert registry_config.space_id == "space-test" + assert registry_config.region == "cn-beijing" + assert registry_config.top_k == 5 + assert { + "a2a_registry_search_agent_cards", + "a2a_registry_task_create", + "a2a_registry_task_poll", + }.issubset({getattr(tool, "__name__", "") for tool in cloned.tools}) + + def test_spawn_applies_sampling_overrides_to_clone_only(self): + base = Agent( + model_name="base-model", + model_api_key="test-key", + generate_content_config=types.GenerateContentConfig(temperature=0.1), + ) + + cloned = spawn_harness_agent( + base, + HarnessOverrides( + top_p=0.9, + max_tokens=128, + penalty=0.2, + presence_penalty=0.3, + ), + ) + + assert base.generate_content_config is not None + assert base.generate_content_config.temperature == 0.1 + assert base.generate_content_config.top_p is None + assert cloned.generate_content_config.temperature == 0.1 + assert cloned.generate_content_config.top_p == 0.9 + assert cloned.generate_content_config.max_output_tokens == 128 + assert cloned.generate_content_config.presence_penalty == 0.3 + assert cloned.generate_content_config.frequency_penalty == 0.2 + + def test_spawn_applies_resource_overrides_to_clone_only(self): + base_kb = KnowledgeBase(backend="local", app_name="base-app") + base_memory = LongTermMemory(backend="local", app_name="base-app") + base = Agent( + model_name="base-model", + model_api_key="test-key", + knowledgebase=base_kb, + long_term_memory=base_memory, + ) + + cloned = spawn_harness_agent( + base, + HarnessOverrides( + knowledgebase={ + "type": "local", + "config": {"index": "request-kb"}, + }, + longterm_memory={ + "type": "local", + "config": {"index": "request-memory"}, + }, + ), + app_name="request-app", + ) + + base_kb_tools = [ + tool for tool in base.tools if isinstance(tool, LoadKnowledgebaseTool) + ] + cloned_kb_tools = [ + tool for tool in cloned.tools if isinstance(tool, LoadKnowledgebaseTool) + ] + + assert base.knowledgebase is base_kb + assert base.long_term_memory is base_memory + assert base.knowledgebase.index == "base-app" + assert base.long_term_memory.index == "base-app" + assert len(base_kb_tools) == 1 + assert base_kb_tools[0].knowledgebase.index == "base-app" + + assert cloned.knowledgebase is not base.knowledgebase + assert cloned.long_term_memory is not base.long_term_memory + assert cloned.knowledgebase.index == "request-kb" + assert cloned.long_term_memory.index == "request-memory" + assert cloned.auto_save_session is True + assert save_session_to_long_term_memory in _after_agent_callbacks(cloned) + assert len(cloned_kb_tools) == 1 + assert cloned_kb_tools[0].knowledgebase.index == "request-kb" + assert ( + sum(getattr(tool, "name", None) == "load_memory" for tool in cloned.tools) + == 1 + ) + + def test_spawn_removes_auto_save_when_longterm_memory_is_cleared(self): + base_memory = LongTermMemory(backend="local", app_name="base-app") + + def custom_callback(*_args): + return None + + base = Agent( + model_name="base-model", + model_api_key="test-key", + long_term_memory=base_memory, + auto_save_session=True, + after_agent_callback=custom_callback, + ) + + cloned = spawn_harness_agent( + base, + HarnessOverrides.model_validate({"longterm_memory": None}), + app_name="request-app", + ) + + assert base.long_term_memory is base_memory + assert base.auto_save_session is True + assert save_session_to_long_term_memory in _after_agent_callbacks(base) + assert cloned.long_term_memory is None + assert cloned.auto_save_session is False + assert _after_agent_callbacks(cloned) == [custom_callback] + assert not any( + getattr(tool, "name", None) == "load_memory" for tool in cloned.tools + ) + + def test_resource_id_without_resolver_falls_back_to_index_and_app_name(self): + base = Agent(model_name="base-model", model_api_key="test-key") + + cloned = spawn_harness_agent( + base, + HarnessOverrides( + knowledgebase={"type": "local", "id": "kb-id"}, + longterm_memory={"type": "local", "id": "memory-id"}, + ), + app_name="request-app", + ) + + assert cloned.knowledgebase.index == "kb-id" + assert cloned.knowledgebase.app_name == "kb-id" + assert cloned.long_term_memory.index == "memory-id" + assert cloned.long_term_memory.app_name == "memory-id" + + def test_resource_resolver_merges_control_plane_config(self): + base = Agent(model_name="base-model", model_api_key="test-key") + calls = [] + + def resolver(kind, resource): + calls.append((kind, resource.type, resource.id)) + if kind == "knowledgebase": + return {"type": "local", "index": "resolved-kb", "top_k": 4} + return HarnessResourceOverride( + type="local", + config={"index": "resolved-memory", "top_k": 2}, + ) + + set_harness_resource_resolver(resolver) + try: + cloned = spawn_harness_agent( + base, + HarnessOverrides( + knowledgebase={ + "type": "local", + "id": "kb-id", + "config": {"top_k": 7}, + }, + longterm_memory={"type": "local", "id": "memory-id"}, + ), + app_name="request-app", + ) + finally: + set_harness_resource_resolver(None) + + assert calls == [ + ("knowledgebase", "local", "kb-id"), + ("longterm_memory", "local", "memory-id"), + ] + assert cloned.knowledgebase.index == "resolved-kb" + assert cloned.knowledgebase.top_k == 7 + assert cloned.long_term_memory.index == "resolved-memory" + assert cloned.long_term_memory.top_k == 2 + + def test_resource_resolver_can_supply_missing_resource_type(self): + base = Agent(model_name="base-model", model_api_key="test-key") + calls = [] + + def resolver(kind, resource): + calls.append((kind, resource.type, resource.id)) + if kind == "knowledgebase": + return HarnessResourceOverride( + type="local", + config={"index": "resolved-kb"}, + ) + return HarnessResourceOverride( + type="local", + config={"index": "resolved-memory"}, + ) + + set_harness_resource_resolver(resolver) + try: + cloned = spawn_harness_agent( + base, + HarnessOverrides( + knowledgebase={"id": "kb-id"}, + longterm_memory={"id": "memory-id"}, + ), + app_name="request-app", + ) + finally: + set_harness_resource_resolver(None) + + assert calls == [ + ("knowledgebase", "", "kb-id"), + ("longterm_memory", "", "memory-id"), + ] + assert cloned.knowledgebase.index == "resolved-kb" + assert cloned.long_term_memory.index == "resolved-memory" + + def test_init_harness_agent_resolves_env_resource_ids_without_types( + self, monkeypatch + ): + calls = [] + + def resolver(kind, resource): + calls.append((kind, resource.type, resource.id)) + if kind == "knowledgebase": + return HarnessResourceOverride( + type="local", + config={"index": "env-kb"}, + ) + return HarnessResourceOverride( + type="local", + config={"index": "env-memory"}, + ) + + def mcp_router_resolver(mcp_router_id, config): + calls.append(("mcp_router", "", mcp_router_id)) + return { + "url": "http://router.test/mcp", + "api_key": "router-token", + } + + monkeypatch.setenv("MODEL_AGENT_NAME", "env-agent") + monkeypatch.setenv("MCP_ROUTER_ID", "mt-1") + monkeypatch.setenv("KNOWLEDGEBASE_ID", "kb-id") + monkeypatch.setenv("LONG_TERM_MEMORY_ID", "memory-id") + set_harness_resource_resolver(resolver) + set_harness_mcp_router_resolver(mcp_router_resolver) + try: + agent, _memory = init_harness_agent() + finally: + set_harness_resource_resolver(None) + set_harness_mcp_router_resolver(None) + + assert calls == [ + ("mcp_router", "", "mt-1"), + ("knowledgebase", "", "kb-id"), + ("longterm_memory", "", "memory-id"), + ] + mcp_router = next( + tool + for tool in agent.tools + if getattr(tool, "_veadk_harness_builtin_tool_id", "") == "mcp_router" + ) + assert mcp_router._connection_params.url == "http://router.test/mcp" + assert mcp_router._connection_params.headers == { + "Authorization": "Bearer router-token" + } + assert agent.knowledgebase.index == "env-kb" + assert agent.long_term_memory.index == "env-memory" + assert agent.auto_save_session is True + assert save_session_to_long_term_memory in _after_agent_callbacks(agent) + + def test_resource_resolver_missing_id_fails_clearly(self): + base = Agent(model_name="base-model", model_api_key="test-key") + set_harness_resource_resolver(lambda *_args: None) + try: + with pytest.raises( + ResourceResolutionError, + match="No runtime config found for knowledgebase resource 'kb-id'", + ): + spawn_harness_agent( + base, + HarnessOverrides(knowledgebase={"type": "local", "id": "kb-id"}), + app_name="request-app", + ) + finally: + set_harness_resource_resolver(None) + + def test_agentkit_resolver_fetches_mem0_connection_info(self): + memory_requests = [] + + class MemoryClient: + def get_memory_collection(self, request): + memory_requests.append(("get", request.memory_id)) + return SimpleNamespace( + memory_id=request.memory_id, + name="my_agent_memory", + provider_collection_id="ak-my_agent_memory", + provider_type="MEM0", + region="cn-beijing", + project_name="default", + ) + + def get_memory_connection_info(self, request): + memory_requests.append(("connection", request.memory_id)) + return SimpleNamespace( + memory_id=request.memory_id, + provider_collection_id="ak-my_agent_memory", + provider_type="MEM0", + connection_infos=[ + SimpleNamespace( + status="Ready", + addr_type="Public", + auth_key="mem0-api-key", + base_url="https://mem0.example.com", + ) + ], + ) + + resolver = AgentKitResourceResolver( + region="cn-beijing", + credential_resolver=lambda: CloudCredentials("ak", "sk", "sts"), + memory_client_factory=lambda credentials, region: MemoryClient(), + ) + + resolved = resolver( + "longterm_memory", + HarnessResourceOverride(type="", id="mem-1"), + ) + + assert memory_requests == [("get", "mem-1"), ("connection", "mem-1")] + assert resolved == HarnessResourceOverride( + type="mem0", + id="mem-1", + config={ + "index": "ak-my_agent_memory", + "app_name": "ak-my_agent_memory", + "api_key": "mem0-api-key", + "base_url": "https://mem0.example.com", + "project_id": "ak-my_agent_memory", + }, + ) + + def test_agentkit_resolver_fetches_viking_knowledge_connection_info(self): + knowledge_requests = [] + + class KnowledgeClient: + def get_knowledge_base(self, request): + knowledge_requests.append(("get", request.knowledge_id)) + return SimpleNamespace( + knowledge_id=request.knowledge_id, + name="my_travel_knowledge", + description="travel knowledge", + provider_knowledge_id="kb-travel-001", + provider_type="VIKINGDB_KNOWLEDGE", + region="cn-beijing", + project_name="default", + ) + + def get_knowledge_connection_info(self, request): + knowledge_requests.append(("connection", request.knowledge_id)) + return SimpleNamespace( + knowledge_id=request.knowledge_id, + provider_knowledge_id="kb-travel-001", + provider_type="VIKINGDB_KNOWLEDGE", + connection_infos=[ + SimpleNamespace( + status="Ready", + addr_type="Public", + auth_type="STS", + auth_key=json.dumps( + { + "AccessKeyId": "temporary-ak", + "SecretAccessKey": "temporary-sk", + "SessionToken": "temporary-token", + } + ), + base_url="https://knowledge.example.com", + region="cn-beijing", + ) + ], + ) + + resolver = AgentKitResourceResolver( + region="cn-beijing", + credential_resolver=lambda: CloudCredentials("fallback-ak", "fallback-sk"), + knowledge_client_factory=lambda credentials, region: KnowledgeClient(), + ) + + resolved = resolver( + "knowledgebase", + HarnessResourceOverride( + type="", + id="kb-1", + config={"top_k": 7}, + ), + ) + + assert knowledge_requests == [("get", "kb-1"), ("connection", "kb-1")] + assert resolved.type == "viking" + assert resolved.id == "kb-1" + assert resolved.config == { + "name": "my_travel_knowledge", + "description": "travel knowledge", + "index": "my_travel_knowledge", + "app_name": "my_travel_knowledge", + "resource_id": "kb-travel-001", + "region": "cn-beijing", + "volcengine_project": "default", + "volcengine_access_key": "temporary-ak", + "volcengine_secret_key": "temporary-sk", + "session_token": "temporary-token", + "cloud_provider": "volcengine", + "base_url": "https://knowledge.example.com", + "host": "knowledge.example.com", + "schema": "https", + "top_k": 7, + } + + def test_agentkit_resolver_without_credentials_allows_config_fallback(self): + resolver = AgentKitResourceResolver( + credential_resolver=lambda: None, + ) + + resolved = resolver( + "knowledgebase", + HarnessResourceOverride( + type="viking", + id="kb-1", + config={"index": "explicit-index"}, + ), + ) + + assert resolved is None + + def test_agentkit_mcp_router_resolver_fetches_toolset_connection_info(self): + mcp_requests = [] + + class MCPClient: + def get_mcp_toolset(self, request): + mcp_requests.append(request.mcp_toolset_id) + return SimpleNamespace( + mcp_toolset=SimpleNamespace( + mcp_toolset_id=request.mcp_toolset_id, + name="router", + path="/mcp", + network_configurations=[ + SimpleNamespace( + network_type="Public", + endpoint="https://router.example.com", + ) + ], + authorizer_configuration=SimpleNamespace( + authorizer_type="ApiKey", + authorizer=SimpleNamespace( + key_auth=SimpleNamespace( + api_keys=[ + SimpleNamespace( + key="router-token", + name="default", + ) + ] + ) + ), + ), + ) + ) + + resolver = AgentKitMcpRouterResolver( + region="cn-beijing", + credential_resolver=lambda: CloudCredentials("ak", "sk"), + mcp_client_factory=lambda credentials, region: MCPClient(), + ) + + resolved = resolver("mt-1") + + assert mcp_requests == ["mt-1"] + assert resolved == { + "mcp_router_id": "mt-1", + "url": "https://router.example.com/mcp", + "api_key": "router-token", + "name": "router", + } + + def test_spawn_mounts_mcp_router_from_id(self): + base = Agent(model_name="base-model", model_api_key="test-key") + calls = [] + + def resolver(mcp_router_id, config): + calls.append((mcp_router_id, config)) + return { + "url": "http://router.test/mcp", + "api_key": "router-token", + } + + set_harness_mcp_router_resolver(resolver) + try: + cloned = spawn_harness_agent( + base, + HarnessOverrides(mcp_router_id="mt-1"), + ) + finally: + set_harness_mcp_router_resolver(None) + + mcp_router = next( + tool + for tool in cloned.tools + if getattr(tool, "_veadk_harness_builtin_tool_id", "") == "mcp_router" + ) + + assert calls == [("mt-1", {"mcp_router_id": "mt-1"})] + assert mcp_router._connection_params.url == "http://router.test/mcp" + assert mcp_router._connection_params.headers == { + "Authorization": "Bearer router-token" + } + + def test_spawn_replaces_builtin_tools(self, monkeypatch): + from veadk.cloud.harness_app import utils + + def base_web_search(): + return "base-web" + + def custom_tool(): + return "custom" + + def fake_get_builtin_tool(name): + def tool(): + return name + + tool.__name__ = name + return tool + + base_web_search.__name__ = "web_search" + custom_tool.__name__ = "custom_tool" + monkeypatch.setattr(utils, "get_builtin_tool", fake_get_builtin_tool) + base = Agent(model_name="base-model", model_api_key="test-key") + base.tools = [base_web_search, custom_tool] + + cloned = spawn_harness_agent( + base, + HarnessOverrides(builtin_tools=[{"id": "link_reader"}]), + ) + + tool_names = [getattr(tool, "__name__", "") for tool in cloned.tools] + assert tool_names == ["custom_tool", "link_reader"] + + def test_spawn_replaces_skills(self, monkeypatch): + from veadk.cloud.harness_app import utils + + class FakeSkillToolset: + def __init__(self, names): + self.names = names + + old_skill_toolset = FakeSkillToolset(["old"]) + + def custom_tool(): + return "custom" + + def fake_build_skill_toolset(skill_ids, download_dir=None): + return FakeSkillToolset(skill_ids) if skill_ids else None + + custom_tool.__name__ = "custom_tool" + monkeypatch.setattr(utils, "SkillToolset", FakeSkillToolset) + monkeypatch.setattr(utils, "build_skill_toolset", fake_build_skill_toolset) + base = Agent(model_name="base-model", model_api_key="test-key") + base.tools = [old_skill_toolset, custom_tool] + + cloned = spawn_harness_agent( + base, + HarnessOverrides( + selected_skills=[{"source": "skillhub", "slug": "team/new-skill"}] + ), + ) + + assert custom_tool in cloned.tools + skill_toolsets = [ + tool for tool in cloned.tools if isinstance(tool, FakeSkillToolset) + ] + assert old_skill_toolset not in cloned.tools + assert len(skill_toolsets) == 1 + assert skill_toolsets[0].names == ["team/new-skill"] + + def test_spawn_clears_builtin_tools_and_skills_with_empty_overrides( + self, monkeypatch + ): + from veadk.cloud.harness_app import utils + + class FakeSkillToolset: + pass + + def base_web_search(): + return "base-web" + + base_web_search.__name__ = "web_search" + monkeypatch.setattr(utils, "SkillToolset", FakeSkillToolset) + base = Agent(model_name="base-model", model_api_key="test-key") + base.tools = [base_web_search, FakeSkillToolset()] + + cloned = spawn_harness_agent( + base, + HarnessOverrides.model_validate( + { + "tools": "", + "selected_skills": [], + } + ), + ) + + assert cloned.tools == [] + + def test_spawn_applies_agentkit_structured_runtime_tools(self, monkeypatch): + from veadk.cloud.harness_app import utils + + def fake_get_builtin_tool(name): + def tool(): + return { + "tool_id": os.environ.get("AGENTKIT_TOOL_ID_SCRIPT"), + "region": os.environ.get("AGENTKIT_TOOL_REGION"), + } + + tool.__name__ = name + return tool + + monkeypatch.setattr(utils, "get_builtin_tool", fake_get_builtin_tool) + base = Agent(model_name="base-model", model_api_key="test-key") + cloned = spawn_harness_agent( + base, + HarnessOverrides( + builtin_tools=[ + { + "id": "run_code", + "config": { + "tool_id": "t-script-1", + "region": "cn-beijing", + }, + }, + { + "id": "mcp_router", + "config": { + "url": "http://router.test/mcp", + "api_key": "router-token", + }, + }, + ], + mcp=[ + { + "name": "db", + "server_url": "http://db.test/mcp", + "bear_token": "db-token", + } + ], + ), + ) + + run_code = next( + tool + for tool in cloned.tools + if getattr(tool, "_veadk_harness_builtin_tool_id", "") == "run_code" + ) + mcp_router = next( + tool + for tool in cloned.tools + if getattr(tool, "_veadk_harness_builtin_tool_id", "") == "mcp_router" + ) + mcp = next( + tool + for tool in cloned.tools + if getattr(tool, "_veadk_harness_mcp_server", "") == "db" + ) + + assert run_code() == { + "tool_id": "t-script-1", + "region": "cn-beijing", + } + assert mcp_router._connection_params.url == "http://router.test/mcp" + assert mcp_router._connection_params.headers == { + "Authorization": "Bearer router-token" + } + assert mcp._connection_params.url == "http://db.test/mcp" + assert mcp._connection_params.headers == {"Authorization": "Bearer db-token"} + + def test_spawn_run_agent_merges_session_and_current_overrides(self): + base = Agent( + model_name="base-model", + model_api_key="test-key", + instruction="base prompt", + generate_content_config=types.GenerateContentConfig(temperature=0.1), + ) + + cloned = spawn_harness_run_agent( + base, + "hello", + session_overrides={ + "model_name": "session-model", + "system_prompt": "session prompt", + "temperature": 0.2, + }, + current_overrides={ + "system_prompt": "current prompt", + "top_p": 0.8, + }, + ) + + assert base.model_name == "base-model" + assert base.instruction == "base prompt" + assert base.generate_content_config.temperature == 0.1 + assert cloned.model_name == "session-model" + assert cloned.instruction == "current prompt" + assert cloned.generate_content_config.temperature == 0.2 + assert cloned.generate_content_config.top_p == 0.8 + class TestRequestResponseSchemas: + def test_create_session_request_accepts_agentkit_id(self): + request = HarnessCreateSessionRequest.model_validate( + { + "id": "session-1", + "state": {"foo": "bar"}, + "events": [], + } + ) + + assert request.id == "session-1" + assert request.session_id is None + assert request.state == {"foo": "bar"} + assert request.events == [] + + def test_create_session_request_accepts_adk_session_id_alias(self): + request = HarnessCreateSessionRequest.model_validate({"sessionId": "session-1"}) + + assert request.id is None + assert request.session_id == "session-1" + + def test_get_agent_config_request_accepts_camel_case(self): + request = HarnessAgentConfigRequest.model_validate( + { + "appName": "test_agent", + "userId": "test_user", + "sessionId": "session-1", + } + ) + + assert request.app_name == "test_agent" + assert request.user_id == "test_user" + assert request.session_id == "session-1" + + def test_get_agent_config_request_defaults_user_and_session(self): + request = HarnessAgentConfigRequest.model_validate({}) + + assert request.app_name is None + assert request.user_id == "default" + assert request.session_id == "default" + def test_run_agent_request_fields(self): assert set(_fields(RunAgentRequest)) == { "user_id", @@ -239,6 +1447,7 @@ def test_invoke_request_fields(self): "prompt", "harness_name", "harness", + "harness_merge", "harness_enhance", "run_agent_request", } diff --git a/tests/cloud/test_harness_app_http.py b/tests/cloud/test_harness_app_http.py index bd1dfab95..08a3ad612 100644 --- a/tests/cloud/test_harness_app_http.py +++ b/tests/cloud/test_harness_app_http.py @@ -13,10 +13,31 @@ # limitations under the License. import importlib +import types +import uuid +import pytest +from fastapi.responses import PlainTextResponse from fastapi.testclient import TestClient +def _first_run_sse_endpoint(harness_module): + for route in harness_module.app.router.routes: + if getattr(route, "path", None) == "/run_sse" and "POST" in getattr( + route, "methods", set() + ): + return route.endpoint + raise AssertionError("run_sse route not found") + + +def _closure_cell(function, name): + freevars = function.__code__.co_freevars + closure = function.__closure__ or () + if name not in freevars: + raise AssertionError(f"{function.__name__} has no closure cell {name}") + return closure[freevars.index(name)] + + def test_harness_app_exposes_agent_info(monkeypatch): monkeypatch.setenv("MODEL_AGENT_API_KEY", "test-api-key") monkeypatch.setenv("MODEL_NAME", "test-model") @@ -76,3 +97,620 @@ def test_harness_app_disables_bff_tool_host_by_default(monkeypatch): getattr(route, "path", None) == "/harness/run_sse" for route in harness_module.app.router.routes ) + + +def test_harness_session_create_accepts_id_and_get_agent_config(monkeypatch): + monkeypatch.setenv("MODEL_AGENT_API_KEY", "test-api-key") + monkeypatch.setenv("MODEL_NAME", "test-model") + monkeypatch.setenv("HARNESS_NAME", "test-harness") + + harness_module = importlib.import_module("veadk.cloud.harness_app.app") + + with TestClient(harness_module.app) as client: + app_name = client.get("/list-apps").json()[0] + user_id = f"user-{uuid.uuid4()}" + session_id = f"session-{uuid.uuid4()}" + created = client.post( + f"/apps/{app_name}/users/{user_id}/sessions", + json={"id": session_id}, + ) + + assert created.status_code == 200 + assert created.json()["id"] == session_id + + config = client.get( + "/get_agent_config", + params={ + "app_name": app_name, + "user_id": user_id, + "session_id": session_id, + }, + ) + + assert config.status_code == 200 + body = config.json() + assert body["app_name"] == app_name + assert body["user_id"] == user_id + assert body["session_id"] == session_id + assert body["harness"]["model_name"] == "test-model" + assert body["harness"]["runtime"] == "adk" + assert body["harness"]["max_llm_calls"] == 10 + + camel_config = client.get( + "/get_agent_config", + params={ + "appName": app_name, + "userId": user_id, + "sessionId": session_id, + }, + ) + + assert camel_config.status_code == 200 + assert camel_config.json()["harness"]["model_name"] == "test-model" + + default_config = client.get("/get_agent_config") + + assert default_config.status_code == 200 + assert default_config.json()["user_id"] == "default" + assert default_config.json()["session_id"] == "default" + assert default_config.json()["harness"]["model_name"] == "test-model" + + +def test_run_sse_harness_does_not_persist_across_same_session(monkeypatch): + monkeypatch.setenv("MODEL_AGENT_API_KEY", "test-api-key") + monkeypatch.setenv("MODEL_NAME", "test-model") + monkeypatch.setenv("HARNESS_NAME", "test-harness") + + harness_module = importlib.import_module("veadk.cloud.harness_app.app") + seen = [] + + async def fake_run_sse_events( + self, req, tip_token="", auth_header="", plugins=None + ): + seen.append( + req.harness.model_dump(mode="json", exclude_unset=True) + if req.harness is not None + else None + ) + yield "data: [DONE]\n\n" + + monkeypatch.setattr(harness_module, "has_a2a_registry_config", lambda agent: True) + monkeypatch.setattr( + harness_module.harness_app, + "_run_sse_events", + types.MethodType(fake_run_sse_events, harness_module.harness_app), + ) + + with TestClient(harness_module.app) as client: + app_name = client.get("/list-apps").json()[0] + user_id = f"user-{uuid.uuid4()}" + session_id = f"session-{uuid.uuid4()}" + client.post( + f"/apps/{app_name}/users/{user_id}/sessions", + json={"id": session_id}, + ) + + first = client.post( + "/run_sse", + json={ + "app_name": app_name, + "user_id": user_id, + "session_id": session_id, + "streaming": True, + "new_message": { + "role": "user", + "parts": [{"text": "hello"}], + }, + "harness": { + "knowledgebase": { + "type": "local", + "id": "kb-1", + }, + }, + }, + ) + second = client.post( + "/run_sse", + json={ + "app_name": app_name, + "user_id": user_id, + "session_id": session_id, + "streaming": True, + "new_message": { + "role": "user", + "parts": [{"text": "hello again"}], + }, + "harness": { + "longterm_memory": { + "type": "local", + "id": "memory-1", + }, + }, + }, + ) + third = client.post( + "/run_sse", + json={ + "app_name": app_name, + "user_id": user_id, + "session_id": session_id, + "streaming": True, + "new_message": { + "role": "user", + "parts": [{"text": "hello from default config"}], + }, + }, + ) + config = client.post( + "/get_agent_config", + json={ + "app_name": app_name, + "user_id": user_id, + "session_id": session_id, + }, + ) + + assert first.status_code == 200 + assert second.status_code == 200 + assert third.status_code == 200 + assert seen == [ + { + "knowledgebase": { + "type": "local", + "id": "kb-1", + }, + }, + { + "longterm_memory": { + "type": "local", + "id": "memory-1", + }, + }, + None, + ] + assert config.status_code == 200 + returned_harness = config.json()["harness"] + assert returned_harness["runtime"] == "adk" + assert returned_harness["max_llm_calls"] == 10 + assert "knowledgebase" not in returned_harness + assert "longterm_memory" not in returned_harness + + +def test_run_sse_harness_merge_uses_default_config_not_session_config(monkeypatch): + monkeypatch.setenv("MODEL_AGENT_API_KEY", "test-api-key") + monkeypatch.setenv("MODEL_NAME", "test-model") + monkeypatch.setenv("HARNESS_NAME", "test-harness") + + harness_module = importlib.import_module("veadk.cloud.harness_app.app") + monkeypatch.setattr( + harness_module.harness_app, + "default_harness_config", + { + "model_name": "default-model", + "knowledgebase": { + "type": "local", + "id": "default-kb", + }, + "temperature": 0.2, + }, + ) + seen = [] + + async def fake_run_sse_events( + self, req, tip_token="", auth_header="", plugins=None + ): + seen.append(req.harness.model_dump(mode="json", exclude_unset=True)) + yield "data: [DONE]\n\n" + + monkeypatch.setattr( + harness_module.harness_app, + "_run_sse_events", + types.MethodType(fake_run_sse_events, harness_module.harness_app), + ) + + with TestClient(harness_module.app) as client: + app_name = client.get("/list-apps").json()[0] + user_id = f"user-{uuid.uuid4()}" + session_id = f"session-{uuid.uuid4()}" + client.post( + f"/apps/{app_name}/users/{user_id}/sessions", + json={"id": session_id}, + ) + + first = client.post( + "/run_sse", + json={ + "app_name": app_name, + "user_id": user_id, + "session_id": session_id, + "streaming": True, + "new_message": { + "role": "user", + "parts": [{"text": "session kb should not be reused"}], + }, + "harness": { + "knowledgebase": { + "type": "local", + "id": "session-kb", + }, + }, + }, + ) + second = client.post( + "/run_sse", + json={ + "app_name": app_name, + "user_id": user_id, + "session_id": session_id, + "streaming": True, + "harness_merge": True, + "new_message": { + "role": "user", + "parts": [{"text": "merge with default only"}], + }, + "harness": { + "longterm_memory": { + "type": "local", + "id": "memory-1", + }, + }, + }, + ) + + assert first.status_code == 200 + assert second.status_code == 200 + assert seen == [ + { + "knowledgebase": { + "type": "local", + "id": "session-kb", + }, + }, + { + "model_name": "default-model", + "knowledgebase": { + "type": "local", + "id": "default-kb", + }, + "temperature": 0.2, + "longterm_memory": { + "type": "local", + "id": "memory-1", + }, + }, + ] + + +@pytest.mark.parametrize( + ( + "payload_extra", + "headers", + "body_plugin_names", + "header_plugin_names", + "default_plugin_names", + "expected_plugin_names", + ), + [ + ( + {"harness_enhance": {"enabled": True, "components": "compactor"}}, + {"x-test-harness": "1"}, + ["body-plugin"], + ["header-plugin"], + ["default-plugin"], + ["body-plugin"], + ), + ( + {}, + {"x-test-harness": "1"}, + [], + ["header-plugin"], + ["default-plugin"], + ["header-plugin"], + ), + ( + {}, + {}, + [], + [], + ["default-plugin"], + ["default-plugin"], + ), + ], +) +def test_run_sse_uses_harness_plugins_from_body_headers_or_default( + monkeypatch, + payload_extra, + headers, + body_plugin_names, + header_plugin_names, + default_plugin_names, + expected_plugin_names, +): + monkeypatch.setenv("MODEL_AGENT_API_KEY", "test-api-key") + monkeypatch.setenv("MODEL_NAME", "test-model") + monkeypatch.setenv("HARNESS_NAME", "test-harness") + + harness_module = importlib.import_module("veadk.cloud.harness_app.app") + body_plugins = [types.SimpleNamespace(name=name) for name in body_plugin_names] + header_plugins = [types.SimpleNamespace(name=name) for name in header_plugin_names] + default_plugins = [ + types.SimpleNamespace(name=name) for name in default_plugin_names + ] + seen = [] + + def fake_plugins_from_enhance(enhance): + return body_plugins if enhance is not None and enhance.enabled else [] + + def fake_plugins_from_headers(request_headers): + return header_plugins if request_headers.get("x-test-harness") else [] + + async def fake_run_sse_events( + self, req, tip_token="", auth_header="", plugins=None + ): + seen.append([getattr(plugin, "name", "") for plugin in plugins or []]) + yield "data: [DONE]\n\n" + + monkeypatch.setattr(harness_module, "has_a2a_registry_config", lambda agent: False) + monkeypatch.setattr( + harness_module, + "build_harness_plugins_from_enhance", + fake_plugins_from_enhance, + ) + monkeypatch.setattr( + harness_module, + "build_harness_plugins_from_headers", + fake_plugins_from_headers, + ) + monkeypatch.setattr(harness_module.harness_app, "plugins", default_plugins) + monkeypatch.setattr( + harness_module.harness_app, + "_run_sse_events", + types.MethodType(fake_run_sse_events, harness_module.harness_app), + ) + + with TestClient(harness_module.app) as client: + app_name = client.get("/list-apps").json()[0] + payload = { + "app_name": app_name, + "user_id": f"user-{uuid.uuid4()}", + "session_id": f"session-{uuid.uuid4()}", + "streaming": True, + "new_message": { + "role": "user", + "parts": [{"text": "hello with plugins"}], + }, + **payload_extra, + } + response = client.post("/run_sse", json=payload, headers=headers) + + assert response.status_code == 200 + assert seen == [expected_plugin_names] + + +def test_get_agent_config_redacts_resource_configs(monkeypatch): + monkeypatch.setenv("MODEL_AGENT_API_KEY", "test-api-key") + monkeypatch.setenv("MODEL_NAME", "test-model") + monkeypatch.setenv("HARNESS_NAME", "test-harness") + + harness_module = importlib.import_module("veadk.cloud.harness_app.app") + monkeypatch.setattr( + harness_module.harness_app, + "default_harness_config", + { + "model_name": "test-model", + "knowledgebase": { + "type": "viking", + "id": "kb-1", + "config": {"index": "secret-kb-index", "api_key": "secret"}, + }, + "longterm_memory": { + "type": "mem0", + "id": "memory-1", + "config": {"index": "secret-memory-index", "api_key": "secret"}, + }, + }, + ) + + with TestClient(harness_module.app) as client: + response = client.get("/get_agent_config") + + assert response.status_code == 200 + assert response.json()["harness"]["knowledgebase"] == {"id": "kb-1"} + assert response.json()["harness"]["longterm_memory"] == {"id": "memory-1"} + + +def test_run_sse_without_harness_or_session_config_delegates_to_adk(monkeypatch): + monkeypatch.setenv("MODEL_AGENT_API_KEY", "test-api-key") + monkeypatch.setenv("MODEL_NAME", "test-model") + monkeypatch.setenv("HARNESS_NAME", "test-harness") + + harness_module = importlib.import_module("veadk.cloud.harness_app.app") + endpoint = _first_run_sse_endpoint(harness_module) + adk_run_sse_cell = _closure_cell(endpoint, "adk_run_sse") + original_adk_run_sse = adk_run_sse_cell.cell_contents + delegated = [] + + async def fake_adk_run_sse(req): + delegated.append( + { + "app_name": req.app_name, + "user_id": req.user_id, + "session_id": req.session_id, + "harness": req.harness, + } + ) + return PlainTextResponse("delegated") + + async def fail_run_sse_events( + self, req, tip_token="", auth_header="", plugins=None + ): + raise AssertionError("run_sse should delegate to the ADK handler") + yield "data: unreachable\n\n" + + monkeypatch.setattr(harness_module, "has_a2a_registry_config", lambda agent: False) + monkeypatch.setattr( + harness_module.harness_app, + "_run_sse_events", + types.MethodType(fail_run_sse_events, harness_module.harness_app), + ) + adk_run_sse_cell.cell_contents = fake_adk_run_sse + try: + with TestClient(harness_module.app) as client: + app_name = client.get("/list-apps").json()[0] + response = client.post( + "/run_sse", + json={ + "app_name": app_name, + "user_id": f"user-{uuid.uuid4()}", + "session_id": f"session-{uuid.uuid4()}", + "streaming": True, + "new_message": { + "role": "user", + "parts": [{"text": "hello without overrides"}], + }, + }, + ) + finally: + adk_run_sse_cell.cell_contents = original_adk_run_sse + + assert response.status_code == 200 + assert response.text == "delegated" + assert len(delegated) == 1 + assert delegated[0]["app_name"] == app_name + assert delegated[0]["harness"] is None + + +def test_harness_invoke_uses_harness_max_llm_calls(monkeypatch): + monkeypatch.setenv("MODEL_AGENT_API_KEY", "test-api-key") + monkeypatch.setenv("MODEL_NAME", "test-model") + monkeypatch.setenv("HARNESS_NAME", "test-harness") + + harness_module = importlib.import_module("veadk.cloud.harness_app.app") + seen = [] + + class FakeRunner: + def __init__(self, **kwargs): + self.kwargs = kwargs + + async def run(self, **kwargs): + run_config = kwargs["run_config"] + seen.append(getattr(run_config, "max_llm_calls", None)) + return "ok" + + monkeypatch.setattr(harness_module, "Runner", FakeRunner) + monkeypatch.setattr( + harness_module, + "spawn_harness_run_agent", + lambda agent, *_args, **_kwargs: agent, + ) + + with TestClient(harness_module.app) as client: + response = client.post( + "/harness/invoke", + json={ + "prompt": "hello", + "harness_name": "test-harness", + "harness": {"model_name": "model-b", "max_llm_calls": 7}, + "run_agent_request": { + "user_id": f"user-{uuid.uuid4()}", + "session_id": f"session-{uuid.uuid4()}", + }, + }, + ) + + assert response.status_code == 200 + assert response.json()["output"] == "ok" + assert seen == [7] + + +def test_harness_invoke_harness_merge_uses_default_config(monkeypatch): + monkeypatch.setenv("MODEL_AGENT_API_KEY", "test-api-key") + monkeypatch.setenv("MODEL_NAME", "test-model") + monkeypatch.setenv("HARNESS_NAME", "test-harness") + + harness_module = importlib.import_module("veadk.cloud.harness_app.app") + monkeypatch.setattr( + harness_module.harness_app, + "default_harness_config", + { + "model_name": "default-model", + "knowledgebase": { + "type": "local", + "id": "default-kb", + }, + }, + ) + seen = [] + + class FakeRunner: + def __init__(self, **kwargs): + self.kwargs = kwargs + + async def run(self, **kwargs): + return "ok" + + def fake_spawn(agent, prompt, overrides, *_args, **_kwargs): + seen.append(overrides.model_dump(mode="json", exclude_unset=True)) + return agent + + monkeypatch.setattr(harness_module, "Runner", FakeRunner) + monkeypatch.setattr(harness_module, "spawn_harness_run_agent", fake_spawn) + + with TestClient(harness_module.app) as client: + first = client.post( + "/harness/invoke", + json={ + "prompt": "hello", + "harness_name": "test-harness", + "harness": { + "longterm_memory": { + "type": "local", + "id": "memory-1", + }, + }, + "run_agent_request": { + "user_id": f"user-{uuid.uuid4()}", + "session_id": f"session-{uuid.uuid4()}", + }, + }, + ) + second = client.post( + "/harness/invoke", + json={ + "prompt": "hello", + "harness_name": "test-harness", + "harness_merge": True, + "harness": { + "longterm_memory": { + "type": "local", + "id": "memory-1", + }, + }, + "run_agent_request": { + "user_id": f"user-{uuid.uuid4()}", + "session_id": f"session-{uuid.uuid4()}", + }, + }, + ) + + assert first.status_code == 200 + assert second.status_code == 200 + assert seen == [ + { + "longterm_memory": { + "type": "local", + "id": "memory-1", + }, + }, + { + "model_name": "default-model", + "knowledgebase": { + "type": "local", + "id": "default-kb", + }, + "longterm_memory": { + "type": "local", + "id": "memory-1", + }, + }, + ] diff --git a/tests/cloud/test_harness_enhance_env.py b/tests/cloud/test_harness_enhance_env.py index ac2f5b25c..83526873c 100644 --- a/tests/cloud/test_harness_enhance_env.py +++ b/tests/cloud/test_harness_enhance_env.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json + from veadk.cloud.harness_app.env_mapping import to_runtime_env @@ -43,3 +45,31 @@ def test_harness_enhance_config_flattens_to_runtime_env(): assert env["HARNESS_MAX_CONTEXT_CHARS"] == "12000" assert env["HARNESS_MAX_TOOL_RESULT_CHARS"] == "3000" assert env["HARNESS_VERIFIER_MODE"] == "observe" + + +def test_structured_skills_and_mcp_map_to_json_runtime_env(): + env = to_runtime_env( + { + "selected_skills": [ + {"source": "skillhub", "slug": "team/reporting"}, + ], + "mcp": [ + { + "name": "db", + "server_url": "http://db.test/mcp", + "bear_token": "secret", + }, + ], + } + ) + + assert json.loads(env["SELECTED_SKILLS_JSON"]) == [ + {"source": "skillhub", "slug": "team/reporting"} + ] + assert json.loads(env["MCP_SERVERS_JSON"]) == [ + { + "name": "db", + "server_url": "http://db.test/mcp", + "bear_token": "secret", + } + ] diff --git a/tests/cloud/test_harness_skill_download.py b/tests/cloud/test_harness_skill_download.py index 6be1a6965..c38f8fd2e 100644 --- a/tests/cloud/test_harness_skill_download.py +++ b/tests/cloud/test_harness_skill_download.py @@ -152,36 +152,39 @@ def fail_skillhub_download(*args, **kwargs): assert isinstance(toolset._code_executor, UnsafeLocalCodeExecutor) -def test_incremental_skills_preserve_existing_code_executor(monkeypatch, tmp_path): +def test_replace_skills_replaces_existing_skill_toolset(monkeypatch, tmp_path): existing_dir = tmp_path / "existing" incoming_dir = tmp_path / "incoming" _write_adk_skill(existing_dir, name="existing") _write_adk_skill(incoming_dir, name="incoming") - executor = UnsafeLocalCodeExecutor() existing_toolset = SkillToolset( skills=[load_skill_from_dir(existing_dir)], - code_executor=executor, + code_executor=UnsafeLocalCodeExecutor(), ) incoming_toolset = SkillToolset( skills=[load_skill_from_dir(incoming_dir)], code_executor=UnsafeLocalCodeExecutor(), ) - agent = SimpleNamespace(tools=[existing_toolset]) + marker_tool = object() + agent = SimpleNamespace(tools=[marker_tool, existing_toolset]) + calls: list[tuple[list[str], object]] = [] + + def fake_build_skill_toolset(skill_ids, download_dir=None): + calls.append((skill_ids, download_dir)) + return incoming_toolset + monkeypatch.setattr( utils, "build_skill_toolset", - lambda skill_ids, download_dir=None: incoming_toolset, + fake_build_skill_toolset, ) - utils._add_incremental_skills(agent, ["incoming"]) + utils._replace_skills(agent, ["incoming"], download_dir=tmp_path) - merged_toolset = agent.tools[0] - assert [skill.name for skill in merged_toolset._list_skills()] == [ - "existing", - "incoming", - ] - assert merged_toolset._code_executor is executor + assert calls == [(["incoming"], tmp_path)] + assert agent.tools == [marker_tool, incoming_toolset] + assert [skill.name for skill in incoming_toolset._list_skills()] == ["incoming"] @pytest.mark.asyncio diff --git a/veadk/cli/cli_harness.py b/veadk/cli/cli_harness.py index 537717579..e29b1be8e 100644 --- a/veadk/cli/cli_harness.py +++ b/veadk/cli/cli_harness.py @@ -233,7 +233,7 @@ done; \\ test -d src/veadk RUN uv pip install --system --index-url https://mirrors.aliyun.com/pypi/simple/ \\ - "./src[harness,codex]" fastapi "uvicorn[standard]" + "./src[extensions,database,harness,codex]" fastapi "uvicorn[standard]" EXPOSE 8000 CMD ["python", "-m", "uvicorn", "veadk.cloud.harness_app.app:app", "--host", "0.0.0.0", "--port", "8000"] """ @@ -293,7 +293,6 @@ @click.group() def harness() -> None: """Create, configure, and deploy a VeADK harness server.""" - pass @harness.command("create") @@ -383,17 +382,34 @@ def _connection_options(func): return func +_HTTP_ONLY_OVERRIDE_FIELDS = { + "knowledgebase", + "longterm_memory", + "temperature", + "top_p", + "max_tokens", + "presence_penalty", + "frequency_penalty", + "penalty", + "max_llm_calls", +} + + +def _hide_from_cli_override_flags(name: str) -> bool: + return name.startswith("registry_") or name in _HTTP_ONLY_OVERRIDE_FIELDS + + def _override_options(func): """Attach a ``--flag`` for every :class:`HarnessOverrides` field. Shared by ``add`` and ``invoke`` so their model / tools / skills / system-prompt / runtime flags stay identical and in sync with the model. - ``registry_*`` overrides are accepted by the HTTP API for AgentKit, but are - intentionally hidden from the VeADK CLI. Each exposed flag defaults to - ``None`` (unset → not applied). + Some HTTP-only overrides use nested objects or have dedicated CLI flags; + those are intentionally hidden from the VeADK CLI. Each exposed flag + defaults to ``None`` (unset → not applied). """ for name, field in reversed(list(HarnessOverrides.model_fields.items())): - if name.startswith("registry_"): + if _hide_from_cli_override_flags(name): continue option: dict = { "default": None, @@ -565,14 +581,15 @@ def show(path: str) -> None: click.echo("") click.secho("Overridable at invoke time:", fg="green", bold=True) for name, field in HarnessOverrides.model_fields.items(): - if name.startswith("registry_"): + if _hide_from_cli_override_flags(name): continue flag = "--" + name.replace("_", "-") click.echo(f" {flag}: {field.description or name}") click.echo("") click.echo( "Override per call via `veadk harness invoke ... --`. " - "Memory, knowledgebase, and registry are not exposed as VeADK CLI overrides." + "HTTP-only overrides such as memory, knowledgebase, sampling params, " + "and registry are not exposed as VeADK CLI overrides." ) diff --git a/veadk/cloud/harness_app/Dockerfile b/veadk/cloud/harness_app/Dockerfile index a127a5f90..d0ea59d7f 100644 --- a/veadk/cloud/harness_app/Dockerfile +++ b/veadk/cloud/harness_app/Dockerfile @@ -18,9 +18,10 @@ WORKDIR /app # `[extensions]` pulls llama-index / redis / opensearch, required when the # KNOWLEDGEBASE_TYPE or LONG_TERM_MEMORY_TYPE env vars enable those components. +# `[database]` adds optional long-term memory backends such as mem0. # (Viking / MySQL / PostgreSQL backends are already in the base dependencies.) RUN apt-get update && apt-get install -y --no-install-recommends git && \ - pip3 install --no-cache-dir "veadk-python[extensions] @ git+https://github.com/volcengine/veadk-python.git" && \ + pip3 install --no-cache-dir "veadk-python[extensions,database,harness] @ git+https://github.com/volcengine/veadk-python.git" && \ apt-get purge -y git && apt-get autoremove -y && \ apt-get clean && rm -rf /var/lib/apt/lists/* diff --git a/veadk/cloud/harness_app/agentkit_resources.py b/veadk/cloud/harness_app/agentkit_resources.py new file mode 100644 index 000000000..c7410ca84 --- /dev/null +++ b/veadk/cloud/harness_app/agentkit_resources.py @@ -0,0 +1,603 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# 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. + +"""Resolve AgentKit control-plane resource ids into VeADK runtime config.""" + +from __future__ import annotations + +import json +import re +from collections.abc import Callable +from typing import Any +from urllib.parse import urlsplit + +from veadk.cloud.harness_app.types import HarnessResourceOverride +from veadk.tools.builtin_tools.create_agent.sources.cloud import ( + CloudCredentials, + default_agentkit_region, + resolve_cloud_credentials, +) +from veadk.utils.cloud_provider import cloud_provider_from_env +from veadk.utils.logger import get_logger + +logger = get_logger(__name__) + +CredentialsResolver = Callable[[], CloudCredentials | None] +MemoryClientFactory = Callable[[CloudCredentials, str], Any] +KnowledgeClientFactory = Callable[[CloudCredentials, str], Any] +McpClientFactory = Callable[[CloudCredentials, str], Any] + +_SUPPORTED_AUTH_TYPES = {"", "aksk", "sts", "temporaryaksk", "temporarycredentials"} + + +class AgentKitResourceResolver: + """Resolve memory/knowledge resource ids with the AgentKit SDK clients.""" + + def __init__( + self, + *, + region: str | None = None, + credential_resolver: CredentialsResolver | None = None, + memory_client_factory: MemoryClientFactory | None = None, + knowledge_client_factory: KnowledgeClientFactory | None = None, + ) -> None: + self.region = region or default_agentkit_region() + self._credential_resolver = credential_resolver or resolve_cloud_credentials + self._memory_client_factory = ( + memory_client_factory or _default_memory_client_factory + ) + self._knowledge_client_factory = ( + knowledge_client_factory or _default_knowledge_client_factory + ) + + def __call__( + self, + kind: str, + resource: HarnessResourceOverride, + ) -> HarnessResourceOverride | None: + if not resource.id: + return resource + + credentials = self._credential_resolver() + if credentials is None: + logger.warning( + "AgentKit credentials are unavailable; cannot resolve %s id=%s.", + kind, + resource.id, + ) + return None + + if kind == "knowledgebase": + return self._resolve_knowledgebase(resource, credentials) + if kind == "longterm_memory": + return self._resolve_longterm_memory(resource, credentials) + raise ValueError(f"Unsupported Harness resource kind: {kind}") + + def _resolve_knowledgebase( + self, + resource: HarnessResourceOverride, + credentials: CloudCredentials, + ) -> HarnessResourceOverride: + from agentkit.sdk.knowledge import types + + client = self._knowledge_client_factory(credentials, self.region) + knowledge_id = resource.id or "" + detail = client.get_knowledge_base( + types.GetKnowledgeBaseRequest(KnowledgeId=knowledge_id) + ) + connection = client.get_knowledge_connection_info( + types.GetKnowledgeConnectionInfoRequest(KnowledgeId=knowledge_id) + ) + info = _preferred_connection_info(connection.connection_infos or []) + provider_type = _text( + getattr(connection, "provider_type", "") + or getattr(detail, "provider_type", "") + ) + backend = _knowledge_backend_type(provider_type) + _validate_requested_type(resource.type, backend, "knowledgebase") + config = _resolved_knowledge_config( + detail=detail, + connection=connection, + info=info, + credentials=credentials, + control_plane_region=self.region, + ) + config.update(resource.config or {}) + return HarnessResourceOverride( + type=backend, + id=resource.id, + config=config, + ) + + def _resolve_longterm_memory( + self, + resource: HarnessResourceOverride, + credentials: CloudCredentials, + ) -> HarnessResourceOverride: + from agentkit.sdk.memory import types + + client = self._memory_client_factory(credentials, self.region) + memory_id = resource.id or "" + detail = client.get_memory_collection( + types.GetMemoryCollectionRequest(MemoryId=memory_id) + ) + connection = client.get_memory_connection_info( + types.GetMemoryConnectionInfoRequest(MemoryId=memory_id) + ) + info = _preferred_connection_info(connection.connection_infos or []) + provider_type = _text( + getattr(connection, "provider_type", "") + or getattr(detail, "provider_type", "") + ) + backend = _memory_backend_type(provider_type) + _validate_requested_type(resource.type, backend, "longterm_memory") + config = _resolved_memory_config( + backend=backend, + detail=detail, + connection=connection, + info=info, + credentials=credentials, + control_plane_region=self.region, + ) + config.update(resource.config or {}) + return HarnessResourceOverride( + type=backend, + id=resource.id, + config=config, + ) + + +class AgentKitMcpRouterResolver: + """Resolve an AgentKit MCP toolset id into mcp_router runtime config.""" + + def __init__( + self, + *, + region: str | None = None, + credential_resolver: CredentialsResolver | None = None, + mcp_client_factory: McpClientFactory | None = None, + ) -> None: + self.region = region or default_agentkit_region() + self._credential_resolver = credential_resolver or resolve_cloud_credentials + self._mcp_client_factory = mcp_client_factory or _default_mcp_client_factory + + def __call__( + self, + mcp_router_id: str, + config: dict[str, Any] | None = None, + ) -> dict[str, Any] | None: + credentials = self._credential_resolver() + if credentials is None: + logger.warning( + "AgentKit credentials are unavailable; cannot resolve MCP router id=%s.", + mcp_router_id, + ) + return None + + from agentkit.sdk.mcp import types + + client = self._mcp_client_factory(credentials, self.region) + response = client.get_mcp_toolset( + types.GetMCPToolsetRequest(MCPToolsetId=mcp_router_id) + ) + toolset = getattr(response, "mcp_toolset", None) + if toolset is None: + raise ValueError(f"AgentKit MCP toolset '{mcp_router_id}' was not found.") + + url = _mcp_toolset_url(toolset) + api_key = _mcp_toolset_api_key(toolset) + resolved = { + "mcp_router_id": mcp_router_id, + "url": url, + "api_key": api_key, + "name": _text(getattr(toolset, "name", "")), + } + resolved = { + key: value for key, value in resolved.items() if value not in {None, ""} + } + resolved.update(config or {}) + return resolved + + +def default_agentkit_resource_resolver() -> AgentKitResourceResolver: + """Build the default resolver used by HarnessApp at runtime.""" + + return AgentKitResourceResolver() + + +def default_agentkit_mcp_router_resolver() -> AgentKitMcpRouterResolver: + """Build the default MCP router resolver used by HarnessApp at runtime.""" + + return AgentKitMcpRouterResolver() + + +def _default_memory_client_factory(credentials: CloudCredentials, region: str) -> Any: + from agentkit.platform.context import default_cloud_provider + from agentkit.sdk.memory.client import AgentkitMemoryClient + + with default_cloud_provider(cloud_provider_from_env()): + return AgentkitMemoryClient( + access_key=credentials.access_key, + secret_key=credentials.secret_key, + session_token=credentials.session_token, + region=region, + ) + + +def _default_knowledge_client_factory( + credentials: CloudCredentials, region: str +) -> Any: + from agentkit.platform.context import default_cloud_provider + from agentkit.sdk.knowledge.client import AgentkitKnowledgeClient + + with default_cloud_provider(cloud_provider_from_env()): + return AgentkitKnowledgeClient( + access_key=credentials.access_key, + secret_key=credentials.secret_key, + session_token=credentials.session_token, + region=region, + ) + + +def _default_mcp_client_factory(credentials: CloudCredentials, region: str) -> Any: + from agentkit.platform.context import default_cloud_provider + from agentkit.sdk.mcp.client import AgentkitMCPClient + + with default_cloud_provider(cloud_provider_from_env()): + return AgentkitMCPClient( + access_key=credentials.access_key, + secret_key=credentials.secret_key, + session_token=credentials.session_token, + region=region, + ) + + +def _preferred_connection_info(infos: list[Any]) -> Any: + info = next( + ( + item + for item in infos + if _text(getattr(item, "status", "")).casefold() + in {"", "ready", "available", "active"} + and _text(getattr(item, "addr_type", "")).casefold() + in {"", "public", "internet"} + ), + infos[0] if infos else None, + ) + if info is None: + raise ValueError("AgentKit returned no provider connection info.") + return info + + +def _knowledge_backend_type(provider_type: str) -> str: + normalized = _provider_type(provider_type) + if normalized == "vikingdbknowledge": + return "viking" + raise ValueError(f"Unsupported AgentKit knowledge provider type: {provider_type}") + + +def _memory_backend_type(provider_type: str) -> str: + normalized = _provider_type(provider_type) + if normalized == "mem0": + return "mem0" + if normalized == "vikingdbmemory": + return "viking" + raise ValueError(f"Unsupported AgentKit memory provider type: {provider_type}") + + +def _provider_type(value: str) -> str: + return re.sub(r"[^a-z0-9]", "", value.casefold()) + + +def _validate_requested_type( + requested_type: str, resolved_type: str, kind: str +) -> None: + requested = requested_type.strip().casefold() + if requested and requested != resolved_type: + raise ValueError( + f"AgentKit {kind} provider resolved to '{resolved_type}', " + f"but request specified '{requested_type}'." + ) + + +def _resolved_knowledge_config( + *, + detail: Any, + connection: Any, + info: Any, + credentials: CloudCredentials, + control_plane_region: str, +) -> dict[str, Any]: + provider_id = _text( + getattr(connection, "provider_knowledge_id", "") + or getattr(detail, "provider_knowledge_id", "") + ) + knowledge_id = _text(getattr(detail, "knowledge_id", "")) + name = _text(getattr(detail, "name", "")) + region = ( + _text(getattr(info, "region", "")) + or _text(getattr(detail, "region", "")) + or control_plane_region + ) + project_name = _text(getattr(detail, "project_name", "")) or "default" + access_key, secret_key, session_token = _connection_credentials( + auth_type=_text(getattr(info, "auth_type", "")), + auth_key=_text(getattr(info, "auth_key", "")), + extra_config=_text(getattr(info, "extra_config", "")), + fallback=credentials, + ) + base_url_config = _base_url_config(_text(getattr(info, "base_url", ""))) + resource_id = ( + provider_id + if provider_id.startswith("kb-") + else (knowledge_id if knowledge_id.startswith("kb-") else "") + ) + index = _viking_index(provider_id, name or knowledge_id) + config = { + "name": name or knowledge_id, + "description": _text(getattr(detail, "description", "")), + "index": index, + "app_name": index, + "resource_id": resource_id, + "region": region, + "volcengine_project": project_name, + "volcengine_access_key": access_key, + "volcengine_secret_key": secret_key, + "session_token": session_token or "", + "cloud_provider": cloud_provider_from_env(), + **base_url_config, + } + return {key: value for key, value in config.items() if value not in {None, ""}} + + +def _resolved_memory_config( + *, + backend: str, + detail: Any, + connection: Any, + info: Any, + credentials: CloudCredentials, + control_plane_region: str, +) -> dict[str, Any]: + provider_id = _text( + getattr(connection, "provider_collection_id", "") + or getattr(detail, "provider_collection_id", "") + ) + memory_id = _text(getattr(detail, "memory_id", "")) + name = _text(getattr(detail, "name", "")) + region = ( + _text(getattr(info, "region", "")) + or _text(getattr(detail, "region", "")) + or control_plane_region + ) + project_name = _text(getattr(detail, "project_name", "")) or "default" + if backend == "mem0": + return _resolved_mem0_config( + provider_id=provider_id, + memory_id=memory_id, + name=name, + info=info, + ) + access_key, secret_key, session_token = _connection_credentials( + auth_type=_text(getattr(info, "auth_type", "")), + auth_key=_text(getattr(info, "auth_key", "")), + extra_config=_text(getattr(info, "extra_config", "")), + fallback=credentials, + ) + index = _viking_index(provider_id, name or memory_id) + config = { + "index": index, + "app_name": index, + "region": region, + "volcengine_project": project_name, + "volcengine_access_key": access_key, + "volcengine_secret_key": secret_key, + "session_token": session_token or "", + "cloud_provider": cloud_provider_from_env(), + } + return {key: value for key, value in config.items() if value not in {None, ""}} + + +def _resolved_mem0_config( + *, + provider_id: str, + memory_id: str, + name: str, + info: Any, +) -> dict[str, Any]: + extra = _json_object(_text(getattr(info, "extra_config", ""))) + auth = _json_object(_text(getattr(info, "auth_key", ""))) + api_key = ( + _credential_value(auth, {"apikey", "token", "authkey"}) + or _credential_value(extra, {"apikey", "token", "authkey"}) + or _text(getattr(info, "auth_key", "")) + ) + base_url = _text(getattr(info, "base_url", "")) or _string_value( + extra, "base_url", "BaseUrl", "host", "Host" + ) + index = provider_id or name or memory_id + config = { + "index": index, + "app_name": index, + "api_key": api_key, + "base_url": base_url, + "project_id": provider_id, + } + return {key: value for key, value in config.items() if value not in {None, ""}} + + +def _connection_credentials( + *, + auth_type: str, + auth_key: str, + extra_config: str, + fallback: CloudCredentials, +) -> tuple[str, str, str]: + normalized_auth_type = _provider_type(auth_type) + if not auth_key and normalized_auth_type in {"", "aksk"}: + return fallback.access_key, fallback.secret_key, fallback.session_token + if normalized_auth_type not in _SUPPORTED_AUTH_TYPES: + raise ValueError(f"Unsupported AgentKit provider auth type: {auth_type}") + + combined: dict[str, Any] = {} + for raw in (auth_key, extra_config): + if not raw: + continue + combined.update(_json_object_or_error(raw)) + + access_key = _credential_value(combined, {"ak", "accesskey", "accesskeyid"}) + secret_key = _credential_value(combined, {"sk", "secretkey", "secretaccesskey"}) + session_token = _credential_value( + combined, + {"sessiontoken", "securitytoken", "ststoken"}, + ) + if not access_key or not secret_key: + return fallback.access_key, fallback.secret_key, fallback.session_token + if ( + normalized_auth_type in {"sts", "temporaryaksk", "temporarycredentials"} + and not session_token + ): + raise ValueError("AgentKit provider temporary credentials lack STS token.") + return access_key, secret_key, session_token + + +def _credential_value(payload: dict[str, Any], aliases: set[str]) -> str: + pending: list[tuple[dict[str, Any], int]] = [(payload, 0)] + while pending: + current, depth = pending.pop() + for key, value in current.items(): + normalized = _provider_type(str(key)) + if normalized in aliases and isinstance(value, str) and value.strip(): + return value.strip() + if depth < 2 and isinstance(value, dict): + pending.append((value, depth + 1)) + return "" + + +def _json_object(raw: str) -> dict[str, Any]: + raw = raw.strip() + if not raw: + return {} + try: + value = json.loads(raw) + except json.JSONDecodeError: + return {} + return value if isinstance(value, dict) else {} + + +def _json_object_or_error(raw: str) -> dict[str, Any]: + raw = raw.strip() + if not raw: + return {} + try: + value = json.loads(raw) + except json.JSONDecodeError as e: + raise ValueError("AgentKit provider auth info is not valid JSON.") from e + if not isinstance(value, dict): + raise TypeError("AgentKit provider auth info must be a JSON object.") + return value + + +def _string_value(payload: dict[str, Any], *keys: str) -> str: + for key in keys: + value = payload.get(key) + if isinstance(value, str) and value.strip(): + return value.strip() + return "" + + +def _text(value: Any) -> str: + return str(value or "").strip() + + +def _viking_index(provider_id: str, fallback: str) -> str: + provider_id = provider_id.strip() + if re.fullmatch(r"[A-Za-z][A-Za-z0-9_]{0,127}", provider_id): + return provider_id + return fallback + + +def _base_url_config(base_url: str) -> dict[str, str]: + if not base_url: + return {} + parsed = urlsplit(base_url) + host = parsed.netloc or parsed.path.split("/", 1)[0] + return { + "base_url": base_url, + "host": host, + "schema": parsed.scheme or "https", + } + + +def _mcp_toolset_url(toolset: Any) -> str: + path = _text(getattr(toolset, "path", "")) + endpoint = _preferred_mcp_endpoint( + getattr(toolset, "network_configurations", None) or [] + ) + if not endpoint: + for service in getattr(toolset, "mcp_services", None) or []: + endpoint = _preferred_mcp_endpoint( + getattr(service, "network_configurations", None) or [] + ) + path = path or _text(getattr(service, "path", "")) + if endpoint: + break + if not endpoint: + raise ValueError("AgentKit MCP toolset has no network endpoint.") + return _join_url_path(endpoint, path) + + +def _preferred_mcp_endpoint(networks: list[Any]) -> str: + selected = next( + ( + network + for network in networks + if _text(getattr(network, "network_type", "")).casefold() + in {"", "public", "internet"} + and _text(getattr(network, "endpoint", "")) + ), + networks[0] if networks else None, + ) + return _text(getattr(selected, "endpoint", "")) if selected else "" + + +def _join_url_path(endpoint: str, path: str) -> str: + endpoint = endpoint.strip() + path = path.strip() + if not path: + return endpoint + if not path.startswith("/"): + path = f"/{path}" + if endpoint.endswith(path): + return endpoint + return f"{endpoint.rstrip('/')}{path}" + + +def _mcp_toolset_api_key(toolset: Any) -> str: + authorizer_config = getattr(toolset, "authorizer_configuration", None) + auth_type = _provider_type(getattr(authorizer_config, "authorizer_type", "")) + if auth_type and auth_type not in {"apikey", "keyauth"}: + raise ValueError( + "Unsupported AgentKit MCP router authorizer type: " + f"{getattr(authorizer_config, 'authorizer_type', '')}" + ) + authorizer = getattr(authorizer_config, "authorizer", None) + key_auth = getattr(authorizer, "key_auth", None) + for item in getattr(key_auth, "api_keys", None) or []: + key = _text(getattr(item, "key", "")) + if key: + return key + if auth_type in {"apikey", "keyauth"}: + raise ValueError("AgentKit MCP router ApiKey authorizer has no api key.") + return "" diff --git a/veadk/cloud/harness_app/app.py b/veadk/cloud/harness_app/app.py index 75493ef28..fdbb4d629 100644 --- a/veadk/cloud/harness_app/app.py +++ b/veadk/cloud/harness_app/app.py @@ -41,7 +41,7 @@ from pathlib import Path from typing import Any -from fastapi import FastAPI, Request +from fastapi import FastAPI, HTTPException, Query, Request from fastapi.responses import StreamingResponse from google.adk.agents import RunConfig from google.adk.agents.base_agent import BaseAgent @@ -56,8 +56,10 @@ LocalEvalSetResultsManager, ) from google.adk.evaluation.local_eval_sets_manager import LocalEvalSetsManager +from google.adk.events import Event from google.adk.memory.in_memory_memory_service import InMemoryMemoryService from google.adk.plugins import BasePlugin +from google.adk.sessions import Session from google.adk.utils.context_utils import Aclosing from typing_extensions import override @@ -75,7 +77,10 @@ ) from veadk.cloud.harness_app.metrics import HarnessLlmUsagePlugin from veadk.cloud.harness_app.types import ( + HarnessAgentConfigRequest, HarnessCompactionMetric, + HarnessCreateSessionRequest, + HarnessEnhanceOverrides, HarnessOverrides, HarnessPluginMetrics, HarnessResponseMetrics, @@ -85,7 +90,9 @@ from veadk.cloud.harness_app.utils import ( SkillLoadError, ToolLoadError, + harness_overrides_from_env, has_a2a_registry_config, + merge_harness_overrides, spawn_harness_run_agent, ) from veadk.integrations.agentkit.app import ( @@ -100,10 +107,10 @@ logger = get_logger(__name__) HARNESS_NAME = os.getenv("HARNESS_NAME", "default") -# Optional harness default max LLM calls per run, from harness.yaml (overridable -# per invocation). Unset -> falls through to ADK RunConfig's own default. +# Harness default max LLM calls per run, from harness.yaml (overridable per +# invocation). Unset uses the documented default. DEFAULT_MAX_LLM_CALLS = ( - int(os.environ["MAX_LLM_CALLS"]) if os.environ.get("MAX_LLM_CALLS") else None + int(os.environ["MAX_LLM_CALLS"]) if os.environ.get("MAX_LLM_CALLS") else 10 ) RETURN_LLM_USAGE = os.getenv("HARNESS_APP_RETURN_LLM_USAGE", "").lower() in { "1", @@ -111,6 +118,8 @@ "yes", "on", } +_RESOURCE_HARNESS_FIELDS = ("knowledgebase", "longterm_memory") +_RESOURCE_ID_FIELDS = ("id", "_id", "resource_id", "index", "app_name") def _content_text(content: Any) -> str: @@ -123,6 +132,45 @@ def _content_text(content: Any) -> str: return "\n".join(texts) +def _harness_config_delta(overrides: HarnessOverrides) -> dict[str, Any]: + return merge_harness_overrides(overrides).model_dump( + mode="json", + exclude_unset=True, + exclude_none=False, + ) + + +def _merge_harness_config( + current: dict[str, Any] | HarnessOverrides | None, + overrides: dict[str, Any] | HarnessOverrides | None, +) -> dict[str, Any]: + return _harness_config_delta(merge_harness_overrides(current, overrides)) + + +def _public_harness_config(config: dict[str, Any]) -> dict[str, Any]: + public_config = dict(config) + for field in _RESOURCE_HARNESS_FIELDS: + value = public_config.get(field) + if value is None: + public_config.pop(field, None) + continue + if isinstance(value, dict): + resource_id = _public_resource_id(value) + public_config[field] = {"id": resource_id} if resource_id else {} + return public_config + + +def _public_resource_id(value: dict[str, Any]) -> str: + config = value.get("config") + candidates = [value.get(field) for field in _RESOURCE_ID_FIELDS] + if isinstance(config, dict): + candidates.extend(config.get(field) for field in _RESOURCE_ID_FIELDS) + for candidate in candidates: + if candidate not in {None, ""}: + return str(candidate) + return "" + + class _HarnessAgentLoader(BaseAgentLoader): """Serve the single env-built harness agent to the ADK web server. @@ -165,6 +213,8 @@ class HarnessRunAgentRequest(RunAgentRequest): """ harness: HarnessOverrides | None = None + harness_merge: bool = False + harness_enhance: HarnessEnhanceOverrides | None = None class HarnessApp: @@ -181,6 +231,9 @@ def __init__( self.harness_name = harness_name self.max_llm_calls = max_llm_calls self.return_llm_usage = RETURN_LLM_USAGE + self.default_harness_config = _harness_config_delta( + harness_overrides_from_env() + ) self.plugins = build_harness_plugins_from_runtime_env() self.runner = Runner( agent=agent, @@ -233,16 +286,137 @@ async def lifespan(app: FastAPI): self._mount_run_sse_override() self.app.mount("/", self._a2a_app) + def _promote_route(self, endpoint: Any) -> None: + routes = self.app.router.routes + for i, route in enumerate(routes): + if getattr(route, "endpoint", None) is endpoint: + routes.insert(0, routes.pop(i)) + return + + async def _get_session_or_404( + self, + app_name: str, + user_id: str, + session_id: str, + ) -> Session: + session = await self.short_term_memory.session_service.get_session( + app_name=app_name, + user_id=user_id, + session_id=session_id, + ) + if not session: + raise HTTPException(status_code=404, detail="Session not found") + return session + + def _effective_harness_config( + self, + overrides: HarnessOverrides, + *, + harness_merge: bool = False, + ) -> HarnessOverrides: + if harness_merge: + return HarnessOverrides.model_validate( + _merge_harness_config(self.default_harness_config, overrides) + ) + return HarnessOverrides.model_validate(_harness_config_delta(overrides)) + + async def _agent_config_response( + self, + app_name: str, + user_id: str, + session_id: str, + ) -> dict[str, Any]: + return { + "app_name": app_name, + "user_id": user_id, + "session_id": session_id, + "harness": _public_harness_config(self.default_harness_config), + } + def mount(self): + @self.app.post( + "/apps/{app_name}/users/{user_id}/sessions", + response_model_exclude_none=True, + ) + async def create_harness_session( + app_name: str, + user_id: str, + req: HarnessCreateSessionRequest | None = None, + ) -> Session: + session_id = None + state = None + events = None + if req is not None: + session_id = req.id or req.session_id or None + state = req.state + events = req.events + + session = await self._server._create_session( + app_name=app_name, + user_id=user_id, + state=state, + session_id=session_id, + ) + for event_data in events or []: + event = Event.model_validate(event_data) + await self.short_term_memory.session_service.append_event( + session=session, + event=event, + ) + return session + + self._promote_route(create_harness_session) + + @self.app.get("/get_agent_config") + async def get_agent_config( + user_id: str | None = Query(default="default"), + session_id: str | None = Query(default="default"), + app_name: str | None = Query(default=None), + user_id_camel: str | None = Query(default=None, alias="userId"), + session_id_camel: str | None = Query(default=None, alias="sessionId"), + app_name_camel: str | None = Query(default=None, alias="appName"), + ) -> dict[str, Any]: + resolved_user_id = user_id_camel or user_id or "default" + resolved_session_id = session_id_camel or session_id or "default" + return await self._agent_config_response( + app_name or app_name_camel or self.harness_name, + resolved_user_id, + resolved_session_id, + ) + + @self.app.post("/get_agent_config") + async def post_get_agent_config( + req: HarnessAgentConfigRequest, + ) -> dict[str, Any]: + return await self._agent_config_response( + req.app_name or self.harness_name, + req.user_id, + req.session_id, + ) + @self.app.post("/harness/invoke") async def invoke_harness( request: InvokeHarnessRequest, http_request: Request, ) -> InvokeHarnessResponse: - # max LLM calls: per-call override, else the harness default; if - # neither is set, fall through to ADK RunConfig's own default. + effective_harness = ( + self._effective_harness_config( + request.harness, + harness_merge=request.harness_merge, + ) + if request.harness is not None + else None + ) + # max LLM calls: per-call override, then request harness override, + # then the app default from env/config. max_llm_calls = ( - request.run_agent_request.max_llm_calls or self.max_llm_calls + request.run_agent_request.max_llm_calls + or ( + effective_harness.max_llm_calls + if effective_harness and effective_harness.max_llm_calls is not None + else None + ) + or self.max_llm_calls ) run_config = ( RunConfig(max_llm_calls=max_llm_calls) @@ -277,9 +451,9 @@ async def invoke_harness( or bool(header_plugins) or usage_plugin is not None ) - if request.harness is not None: + if effective_harness is not None: logger.info( - f"Applying once-time harness override: {request.harness}" + f"Applying once-time harness override: {effective_harness}" ) # The override clones the base agent and may download incremental # skills into a temp dir; the skill files are read from disk while @@ -291,8 +465,9 @@ async def invoke_harness( agent = spawn_harness_run_agent( self.agent, request.prompt, - request.harness, + effective_harness, download_dir=Path(work_dir), + app_name=self.harness_name, registry_tip_token=tip_token, registry_authorization=auth_header, ) @@ -313,6 +488,7 @@ async def invoke_harness( run_agent = spawn_harness_run_agent( self.agent, request.prompt, + app_name=self.harness_name, registry_tip_token=tip_token, registry_authorization=auth_header, ) @@ -390,8 +566,20 @@ def _mount_run_sse_override(self): @self.app.post("/run_sse") async def run_sse(req: HarnessRunAgentRequest, http_request: Request): + app_name = req.app_name or self.harness_name + req.app_name = app_name + if req.harness is not None: + req.harness = self._effective_harness_config( + req.harness, + harness_merge=req.harness_merge, + ) + header_plugins = build_harness_plugins_from_headers(http_request.headers) + body_plugins = build_harness_plugins_from_enhance(req.harness_enhance) + harness_plugins = body_plugins or header_plugins or self.plugins + self._reset_plugin_diagnostics(harness_plugins) if ( req.harness is None + and not harness_plugins and not has_a2a_registry_config(self.agent) and adk_run_sse is not None ): @@ -400,30 +588,40 @@ async def run_sse(req: HarnessRunAgentRequest, http_request: Request): tip_token = registry_tip_token_from_headers(http_request.headers) auth_header = registry_authorization_from_headers(http_request.headers) return StreamingResponse( - self._run_sse_events(req, tip_token, auth_header), + self._run_sse_events( + req, + tip_token, + auth_header, + plugins=self._plugins_for_run(harness_plugins, None), + ), media_type="text/event-stream", ) # Move ours to the front so it wins (Starlette matches the first route), # without deleting the default we delegate to. - routes = self.app.router.routes - for i, r in enumerate(routes): - if getattr(r, "path", None) == "/run_sse" and ( - getattr(r, "endpoint", None) is run_sse - ): - routes.insert(0, routes.pop(i)) - break + self._promote_route(run_sse) async def _run_sse_events( self, req: "HarnessRunAgentRequest", tip_token: str = "", auth_header: str = "", + plugins: list[BasePlugin] | None = None, ): """Yield SSE ``data:`` lines for a run, spawning the agent on override.""" - run_config = RunConfig( - streaming_mode=StreamingMode.SSE if req.streaming else StreamingMode.NONE + run_config_kwargs: dict[str, Any] = { + "streaming_mode": StreamingMode.SSE if req.streaming else StreamingMode.NONE + } + if req.custom_metadata: + run_config_kwargs["custom_metadata"] = req.custom_metadata + max_llm_calls = ( + req.harness.max_llm_calls + if req.harness and req.harness.max_llm_calls is not None + else self.max_llm_calls ) + if max_llm_calls is not None: + run_config_kwargs["max_llm_calls"] = max_llm_calls + run_config = RunConfig(**run_config_kwargs) work_dir_ctx = None prompt = _content_text(req.new_message) try: @@ -438,6 +636,7 @@ async def _run_sse_events( prompt, req.harness, download_dir=Path(work_dir_ctx.name), + app_name=req.app_name, registry_tip_token=tip_token, registry_authorization=auth_header, ) @@ -449,6 +648,7 @@ async def _run_sse_events( agent = spawn_harness_run_agent( self.agent, prompt, + app_name=req.app_name, registry_tip_token=tip_token, registry_authorization=auth_header, ) @@ -459,6 +659,7 @@ async def _run_sse_events( agent=agent, short_term_memory=self.short_term_memory, app_name=req.app_name, + plugins=plugins or None, ) # Be self-sufficient: create the session if the caller did not. if not await runner.session_service.get_session( @@ -477,15 +678,34 @@ async def _run_sse_events( user_id=req.user_id, session_id=req.session_id, new_message=req.new_message, + state_delta=req.state_delta, run_config=run_config, + invocation_id=req.invocation_id, ) ) as agen: async for event in agen: - yield ( - "data: " - + event.model_dump_json(exclude_none=True, by_alias=True) - + "\n\n" - ) + events_to_stream = [event] + if ( + not req.function_call_event_id + and event.actions.artifact_delta + and event.content + and event.content.parts + ): + content_event = event.model_copy(deep=True) + content_event.actions.artifact_delta = {} + artifact_event = event.model_copy(deep=True) + artifact_event.content = None + events_to_stream = [content_event, artifact_event] + + for event_to_stream in events_to_stream: + yield ( + "data: " + + event_to_stream.model_dump_json( + exclude_none=True, + by_alias=True, + ) + + "\n\n" + ) except Exception as e: logger.exception("run_sse failed") yield f"data: {json.dumps({'error': str(e)})}\n\n" diff --git a/veadk/cloud/harness_app/env_mapping.py b/veadk/cloud/harness_app/env_mapping.py index 7b3d5d639..c51dc3c7a 100644 --- a/veadk/cloud/harness_app/env_mapping.py +++ b/veadk/cloud/harness_app/env_mapping.py @@ -44,6 +44,7 @@ and a Viking long-term memory). """ +import json from typing import Any from veadk.utils.misc import flatten_dict @@ -56,6 +57,50 @@ "short_term_memory": "SHORT_TERM_MEMORY_TYPE", } +COMPONENT_ID_ENV: dict[str, str] = { + "knowledgebase": "KNOWLEDGEBASE_ID", + "long_term_memory": "LONG_TERM_MEMORY_ID", +} + +COMPONENT_CONFIG_ENV: dict[str, str] = { + "knowledgebase": "KNOWLEDGEBASE_CONFIG_JSON", + "long_term_memory": "LONG_TERM_MEMORY_CONFIG_JSON", +} + +COMPONENT_ALIASES: dict[str, tuple[str, ...]] = { + "knowledgebase": ("konwledgebase",), + "long_term_memory": ("longterm_memory",), +} + +HARNESS_ROOT_KEY = "harness" + +MODEL_ENV_ALIASES: dict[str, tuple[str, ...]] = { + "model_name": ("MODEL_AGENT_NAME", "MODEL_NAME"), + "temperature": ("MODEL_AGENT_TEMPERATURE",), + "top_p": ("MODEL_AGENT_TOP_P",), + "max_llm_calls": ("MAX_LLM_CALLS",), +} + +TOOL_CONFIG_ENV: dict[str, dict[str, str]] = { + "run_code": { + "tool_id": "AGENTKIT_TOOL_ID_SCRIPT", + "region": "AGENTKIT_TOOL_REGION", + }, + "coding": { + "tool_id": "AGENTKIT_TOOL_ID_OPENCODE", + "region": "AGENTKIT_TOOL_REGION", + }, + "mcp_router": { + "id": "MCP_ROUTER_ID", + "_id": "MCP_ROUTER_ID", + "mcp_router_id": "MCP_ROUTER_ID", + "mcp_toolset_id": "MCP_ROUTER_ID", + "url": "TOOL_MCP_ROUTER_URL", + "api_key": "TOOL_MCP_ROUTER_API_KEY", + "apikey": "TOOL_MCP_ROUTER_API_KEY", + }, +} + # Backend ``type`` -> {harness connection param: VeADK env var}. Mirrors the # pydantic-settings env prefixes in :mod:`veadk.configs.database_configs`; # credentials map to the shared top-level ``VOLCENGINE_*`` vars. Backends with no @@ -64,6 +109,7 @@ "viking": { "project": "DATABASE_VIKING_PROJECT", "region": "DATABASE_VIKING_REGION", + "resource_id": "DATABASE_VIKING_RESOURCE_ID", "access_key": "VOLCENGINE_ACCESS_KEY", "secret_key": "VOLCENGINE_SECRET_KEY", }, @@ -153,27 +199,53 @@ def to_runtime_env(spec: dict[str, Any]) -> dict[str, str]: """Convert a parsed ``harness.yaml`` into the VeADK runtime env var dict. Empty values are skipped (VeADK falls back to its own defaults). An unknown - backend ``type`` or connection param raises ``ValueError`` (fast-fail on a - typo rather than silently dropping config). + backend ``type`` raises ``ValueError``. Backend-specific params are also + exported as ``DATABASE_*`` env vars when supported; the full component config + is preserved as JSON for resource-level runtime construction. """ + spec = _runtime_spec(spec) env: dict[str, str] = {} # Non-component fields: reuse VeADK's flatten_dict (same as config.yaml). # The `auth` block is excluded too: it configures the runtime's gateway # authorizer at deploy time (custom_jwt), not the container environment. + structured_json_fields = { + "mcp": "MCP_SERVERS_JSON", + "selected_skills": "SELECTED_SKILLS_JSON", + } + special_fields = { + "builtin_tools", + "mcp_router_id", + "mcp_toolset_id", + "model", + *MODEL_ENV_ALIASES, + } rest = { - k: v for k, v in spec.items() if k not in COMPONENT_TYPE_ENV and k != "auth" + k: v + for k, v in spec.items() + if k not in COMPONENT_TYPE_ENV + and k not in _component_aliases() + and k not in structured_json_fields + and k not in special_fields + and k != "auth" } for key, value in flatten_dict(rest).items(): if _is_empty(value): continue env[key.upper()] = _stringify(value) + _add_model_envs(env, spec) + _add_builtin_tool_envs(env, spec.get("builtin_tools")) + _add_mcp_router_envs(env, spec) + for key, env_name in structured_json_fields.items(): + value = spec.get(key) + if not _is_empty(value): + env[env_name] = json.dumps(value, ensure_ascii=False) _add_harness_enhance_aliases(env, spec) # Component sections: `type` selector + backend-specific connection params. for component, type_env in COMPONENT_TYPE_ENV.items(): - section: dict[str, Any] = spec.get(component) or {} + section: dict[str, Any] = _component_section(spec, component) if _is_empty(section.get("type")): continue backend = str(section["type"]) @@ -185,20 +257,142 @@ def to_runtime_env(spec: dict[str, Any]) -> dict[str, str]: f"Unknown backend type '{backend}' for '{component}'. " f"Known: {sorted(BACKEND_ENV)}" ) + resource_config = _component_resource_config(section) + config_env = COMPONENT_CONFIG_ENV.get(component) + if config_env and resource_config: + env[config_env] = json.dumps(resource_config, ensure_ascii=False) + for param, value in section.items(): if param == "type" or _is_empty(value): continue + if param in {"id", "_id"}: + id_env = COMPONENT_ID_ENV.get(component) + if id_env is None: + raise ValueError( + f"Param '{param}' is not supported for {component}." + ) + env[id_env] = _stringify(value) + continue env_name = params.get(param) if env_name is None: - raise ValueError( - f"Unknown param '{param}' for {component} backend '{backend}'. " - f"Known: {sorted(params)}" - ) + continue env[env_name] = _stringify(value) return env +def _runtime_spec(spec: dict[str, Any]) -> dict[str, Any]: + harness_section = spec.get(HARNESS_ROOT_KEY) + if not isinstance(harness_section, dict): + return spec + outer = {k: v for k, v in spec.items() if k != HARNESS_ROOT_KEY} + return {**outer, **harness_section} + + +def _add_model_envs(env: dict[str, str], spec: dict[str, Any]) -> None: + model_section = spec.get("model") + if isinstance(model_section, dict) and not _is_empty(model_section.get("name")): + _set_env_aliases(env, MODEL_ENV_ALIASES["model_name"], model_section["name"]) + for key, env_names in MODEL_ENV_ALIASES.items(): + value = spec.get(key) + if not _is_empty(value): + _set_env_aliases(env, env_names, value) + + +def _set_env_aliases( + env: dict[str, str], + env_names: tuple[str, ...], + value: Any, +) -> None: + string_value = _stringify(value) + for env_name in env_names: + env[env_name] = string_value + + +def _add_builtin_tool_envs(env: dict[str, str], value: Any) -> None: + entries = _builtin_tool_entries(value) + if not entries: + return + env["TOOLS"] = ",".join(entry["id"] for entry in entries) + for entry in entries: + _add_builtin_tool_config_envs(env, entry["id"], entry.get("config") or {}) + + +def _builtin_tool_entries(value: Any) -> list[dict[str, Any]]: + if _is_empty(value): + return [] + raw_entries = value if isinstance(value, list) else [value] + entries: list[dict[str, Any]] = [] + for raw in raw_entries: + if isinstance(raw, str): + tool_id = raw.strip() + config = {} + elif isinstance(raw, dict): + tool_id = str(raw.get("id") or "").strip() + config = raw.get("config") if isinstance(raw.get("config"), dict) else {} + else: + continue + if tool_id: + entries.append({"id": tool_id, "config": config}) + return entries + + +def _add_builtin_tool_config_envs( + env: dict[str, str], + tool_id: str, + config: dict[str, Any], +) -> None: + mapping = TOOL_CONFIG_ENV.get(tool_id) + if not mapping: + return + for key, env_name in mapping.items(): + value = config.get(key) + if not _is_empty(value): + env[env_name] = _stringify(value) + + +def _add_mcp_router_envs(env: dict[str, str], spec: dict[str, Any]) -> None: + mcp_router_id = spec.get("mcp_router_id") or spec.get("mcp_toolset_id") + if _is_empty(mcp_router_id): + return + env["MCP_ROUTER_ID"] = _stringify(mcp_router_id) + tools = [tool for tool in env.get("TOOLS", "").split(",") if tool] + if "mcp_router" not in tools: + tools.append("mcp_router") + env["TOOLS"] = ",".join(tools) + + +def _component_aliases() -> set[str]: + return {alias for aliases in COMPONENT_ALIASES.values() for alias in aliases} + + +def _component_section(spec: dict[str, Any], component: str) -> dict[str, Any]: + section = spec.get(component) + if isinstance(section, dict): + return _flatten_component_config(section) + for alias in COMPONENT_ALIASES.get(component, ()): + alias_section = spec.get(alias) + if isinstance(alias_section, dict): + return _flatten_component_config(alias_section) + return {} + + +def _flatten_component_config(section: dict[str, Any]) -> dict[str, Any]: + config = section.get("config") + flat = {key: value for key, value in section.items() if key != "config"} + if isinstance(config, dict): + return {**config, **flat} + return flat + + +def _component_resource_config(section: dict[str, Any]) -> dict[str, Any]: + return { + key: value + for key, value in section.items() + if key not in {"type", "id", "_id"} and not _is_empty(value) + } + + def _add_harness_enhance_aliases(env: dict[str, str], spec: dict[str, Any]) -> None: """Expose harness_enhance fields under the SDK's generic env names too.""" diff --git a/veadk/cloud/harness_app/types.py b/veadk/cloud/harness_app/types.py index 9e0de7192..3cc772320 100644 --- a/veadk/cloud/harness_app/types.py +++ b/veadk/cloud/harness_app/types.py @@ -17,23 +17,130 @@ The parameters split into two groups: * :class:`HarnessOverrides` — the subset that may be overridden per invocation - (model, prompt, tools, skills, runtime). + (model, prompt, tools, skills, runtime, resources, and sampling params). * :class:`HarnessConfig` — the full set fixed at agent creation time. It extends - the overridable params with the knowledge base and memory components, which are - bound when the agent is built and therefore **cannot** be overridden per request. + the overridable params with additional creation-time settings. ``tools`` and ``skills`` are comma-separated strings (e.g. ``"web_search,web_fetch"``). Skills may be SkillHub names/slugs or skills-center refs prefixed with ``space:``. """ -from typing import Literal +from typing import Any, Literal -from pydantic import BaseModel, Field +from pydantic import AliasChoices, BaseModel, ConfigDict, Field from veadk.consts import DEFAULT_MODEL_AGENT_NAME from veadk.prompts.agent_default_prompt import DEFAULT_DESCRIPTION, DEFAULT_INSTRUCTION +class HarnessBuiltinTool(BaseModel): + """Built-in tool selected by AgentKit runtime config.""" + + id: str = Field(description="Built-in tool id.") + config: dict[str, Any] = Field( + default_factory=dict, + description="Per-tool runtime config.", + ) + + +class HarnessMcpServer(BaseModel): + """Streamable HTTP MCP server selected by AgentKit runtime config.""" + + model_config = ConfigDict(populate_by_name=True) + + name: str = Field(default="", description="MCP server display name.") + server_url: str = Field( + default="", + validation_alias=AliasChoices("server_url", "url"), + description="Streamable HTTP MCP server URL.", + ) + bear_token: str = Field( + default="", + validation_alias=AliasChoices("bear_token", "bearer_token", "api_key"), + description="Optional bearer token for the MCP server.", + ) + + +class HarnessSelectedSkill(BaseModel): + """Skill selected by AgentKit runtime config.""" + + model_config = ConfigDict(populate_by_name=True) + + source: Literal["skillhub", "skillspace"] = Field( + default="skillhub", + description="Skill source.", + ) + slug: str | None = Field( + default=None, + description="SkillHub slug, e.g. namespace/owner/reporting-skill.", + ) + namespace: str | None = Field( + default=None, + description="Optional SkillHub namespace.", + ) + skill_space_id: str | None = Field( + default=None, + validation_alias=AliasChoices("skill_space_id", "space_id"), + description="SkillSpace id.", + ) + skill_id: str | None = Field(default=None, description="Skill id.") + version: str | None = Field(default=None, description="Skill version.") + region: str | None = Field(default=None, description="Skill region.") + + +class HarnessResourceOverride(BaseModel): + """Request-level knowledge or memory backend override.""" + + model_config = ConfigDict(populate_by_name=True) + + type: str = Field( + default="", + description="Backend type for this request-level resource override.", + ) + id: str | None = Field( + default=None, + validation_alias=AliasChoices("id", "_id"), + description="Resource id from AgentKit control plane.", + ) + config: dict[str, Any] = Field( + default_factory=dict, + description=( + "Backend or component configuration for this request-level resource." + ), + ) + + +class HarnessRegistryOverride(BaseModel): + """Request-level AgentKit A2A registry override.""" + + model_config = ConfigDict(populate_by_name=True) + + type: Literal["agentkit_a2a"] = Field( + default="agentkit_a2a", + description="AgentKit registry backend type.", + ) + space_id: str = Field( + default="", + validation_alias=AliasChoices("space_id", "registry_space_id"), + description="AgentKit A2A registry space id.", + ) + endpoint: str = Field( + default="", + validation_alias=AliasChoices("endpoint", "registry_endpoint"), + description="AgentKit A2A registry OpenAPI endpoint.", + ) + region: str = Field( + default="", + validation_alias=AliasChoices("region", "registry_region"), + description="AgentKit A2A registry OpenAPI region.", + ) + top_k: int = Field( + default=3, + validation_alias=AliasChoices("top_k", "registry_top_k"), + description="Number of A2A AgentCards to retrieve.", + ) + + class HarnessOverrides(BaseModel): """Harness parameters that may be overridden on a per-invocation basis. @@ -42,6 +149,8 @@ class HarnessOverrides(BaseModel): AgentKit's harness invoke API but intentionally hidden from the VeADK CLI. """ + model_config = ConfigDict(populate_by_name=True) + model_name: str = Field( default=DEFAULT_MODEL_AGENT_NAME, description="Reasoning model name." ) @@ -49,6 +158,15 @@ class HarnessOverrides(BaseModel): default="", description="Comma-separated built-in tool names, e.g. web_search,web_fetch.", ) + builtin_tools: list[HarnessBuiltinTool] = Field( + default_factory=list, + description="Structured built-in tools selected by AgentKit.", + ) + mcp_router_id: str = Field( + default="", + validation_alias=AliasChoices("mcp_router_id", "mcp_toolset_id"), + description="AgentKit MCP toolset id used by the built-in mcp_router tool.", + ) skills: str = Field( default="", description=( @@ -56,6 +174,14 @@ class HarnessOverrides(BaseModel): "e.g. data-visualization-cloud,space:ss-xxx." ), ) + selected_skills: list[HarnessSelectedSkill] = Field( + default_factory=list, + description="Structured skills selected by AgentKit.", + ) + mcp: list[HarnessMcpServer] = Field( + default_factory=list, + description="Streamable HTTP MCP servers selected by AgentKit.", + ) system_prompt: str = Field( default="You are a helpful assistant.", description="System prompt / instruction.", @@ -75,14 +201,53 @@ class HarnessOverrides(BaseModel): registry_top_k: int = Field( default=3, description="Override the number of A2A AgentCards to retrieve." ) + registry: HarnessRegistryOverride | None = Field( + default=None, + description="Structured AgentKit A2A registry override.", + ) + knowledgebase: HarnessResourceOverride | None = Field( + default=None, + validation_alias=AliasChoices("knowledgebase", "konwledgebase"), + description="Request-level knowledge base override.", + ) + longterm_memory: HarnessResourceOverride | None = Field( + default=None, + validation_alias=AliasChoices("longterm_memory", "long_term_memory"), + description="Request-level long-term memory override.", + ) + temperature: float | None = Field( + default=None, description="Request-level model temperature." + ) + top_p: float | None = Field(default=None, description="Request-level model top_p.") + max_tokens: int | None = Field( + default=None, + description="Request-level max output tokens.", + ) + presence_penalty: float | None = Field( + default=None, description="Request-level presence penalty." + ) + frequency_penalty: float | None = Field( + default=None, description="Request-level frequency penalty." + ) + penalty: float | None = Field( + default=None, + description=( + "Compatibility penalty applied to presence and frequency penalties " + "when those fields are not set." + ), + ) + max_llm_calls: int | None = Field( + default=None, + ge=1, + description="Request-level max LLM calls.", + ) class HarnessConfig(HarnessOverrides): """Full harness parameters fixed when the agent is created. - Extends :class:`HarnessOverrides` with the knowledge base and memory - backends. These are wired into the agent at build time and cannot be changed - per request, so they are intentionally absent from :class:`HarnessOverrides`. + Extends :class:`HarnessOverrides` with creation-time defaults such as the + short-term memory backend and registry runtime settings. An empty backend string means the component is disabled (not created). """ @@ -95,8 +260,9 @@ class HarnessConfig(HarnessOverrides): shortterm_memory_type: str = Field(default="local") runtime: Literal["adk", "codex"] = Field(default="adk") max_llm_calls: int | None = Field( - default=None, - description="Default max LLM calls per run; unset follows ADK RunConfig's default. Overridable per invocation.", + default=10, + ge=1, + description="Default max LLM calls per run. Overridable per invocation.", ) structured_tool_calls: bool = Field(default=False) include_tools_every_turn: bool = Field(default=True) @@ -136,10 +302,45 @@ class RunAgentRequest(BaseModel): session_id: str max_llm_calls: int | None = Field( default=None, + ge=1, description="Override max LLM calls for this single call (falls back to the harness default, then ADK's).", ) +class HarnessCreateSessionRequest(BaseModel): + """Session creation payload accepted by the AgentKit harness runtime.""" + + model_config = ConfigDict(populate_by_name=True) + + id: str | None = Field( + default=None, + description="Session id from AgentKit playground. Empty means create one.", + ) + session_id: str | None = Field( + default=None, + alias="sessionId", + description="ADK-compatible session id alias.", + ) + state: dict[str, Any] | None = Field( + default=None, + description="Initial session state.", + ) + events: list[dict[str, Any]] | None = Field( + default=None, + description="Initial session events.", + ) + + +class HarnessAgentConfigRequest(BaseModel): + """Request body for fetching the latest session-scoped harness config.""" + + model_config = ConfigDict(populate_by_name=True) + + app_name: str | None = Field(default=None, alias="appName") + user_id: str = Field(default="default", alias="userId") + session_id: str = Field(default="default", alias="sessionId") + + class LlmUsageMetrics(BaseModel): """Aggregated model usage for one HarnessApp invocation.""" @@ -193,9 +394,16 @@ class InvokeHarnessRequest(BaseModel): prompt: str harness_name: str # When present, a once-time override applied on top of the served agent for - # this single call. Only the fields actually set are applied; memory and the - # knowledge base are never overridable (absent from HarnessOverrides). + # this single call. Only the fields actually set are applied. harness: HarnessOverrides | None = None + harness_merge: bool = Field( + default=False, + description=( + "When true, merge the request harness with the default harness " + "configuration before running. When false, the request harness fully " + "replaces the default configurable overlay." + ), + ) harness_enhance: HarnessEnhanceOverrides | None = None run_agent_request: RunAgentRequest diff --git a/veadk/cloud/harness_app/utils.py b/veadk/cloud/harness_app/utils.py index c06f9c22a..d960f4414 100644 --- a/veadk/cloud/harness_app/utils.py +++ b/veadk/cloud/harness_app/utils.py @@ -20,18 +20,23 @@ :class:`HarnessConfig` and builds the long-lived agent, downloading its skills from the skill hub and mounting them as an ADK skill toolset. * :func:`spawn_harness_agent` — temporary, one-off creation that clones the base - agent and applies a per-request override (incremental tools/skills on top). + agent and applies a per-request override (configured tools/skills replace the + base harness selection). * :func:`spawn_harness_run_agent` — per-turn clone that also attaches dynamic registry-discovered remote A2A tools for the current user message. """ +import inspect import io +import json import os import re import shutil import tempfile import zipfile +from collections.abc import Callable, Mapping from dataclasses import replace +from functools import wraps from pathlib import Path from typing import Any from urllib.parse import urlencode @@ -41,15 +46,25 @@ from google.adk.code_executors import UnsafeLocalCodeExecutor from google.adk.skills import load_skill_from_dir from google.adk.tools.skill_toolset import SkillToolset +from google.genai import types from veadk import Agent -from veadk.cloud.harness_app.types import HarnessConfig, HarnessOverrides +from veadk.cloud.harness_app.types import ( + HarnessBuiltinTool, + HarnessConfig, + HarnessMcpServer, + HarnessOverrides, + HarnessRegistryOverride, + HarnessResourceOverride, + HarnessSelectedSkill, +) from veadk.knowledgebase import KnowledgeBase from veadk.memory.long_term_memory import LongTermMemory from veadk.memory.short_term_memory import ShortTermMemory from veadk.skills.materializer import materialize_remote_skill from veadk.skills.utils import _load_skills_from_space_id from veadk.tools import get_builtin_tool, list_builtin_tools +from veadk.tools.builtin_tools.load_knowledgebase import LoadKnowledgebaseTool from veadk.utils.logger import get_logger logger = get_logger(__name__) @@ -66,21 +81,38 @@ "registry_region", "registry_top_k", } +_SAMPLING_OVERRIDE_FIELDS = { + "temperature", + "top_p", + "max_tokens", + "presence_penalty", + "frequency_penalty", + "penalty", +} +_KNOWLEDGEBASE_TOOL_NAMES = {"load_knowledgebase", "load_kb_queries"} +_LONGTERM_MEMORY_TOOL_NAMES = {"load_memory"} _SKILL_CENTER_SPACE_PREFIX = "space:" __all__ = [ "HarnessConfig", "HarnessOverrides", - "split_csv", - "agent_name_from_harness", - "build_skill_toolset", + "HarnessResourceResolver", + "ResourceResolutionError", "SkillLoadError", "ToolLoadError", + "agent_name_from_harness", + "build_skill_toolset", "config_from_env", + "harness_overrides_from_env", + "has_a2a_registry_config", "init_harness_agent", + "merge_harness_overrides", + "normalize_harness_overrides", + "set_harness_mcp_router_resolver", + "set_harness_resource_resolver", "spawn_harness_agent", "spawn_harness_run_agent", - "has_a2a_registry_config", + "split_csv", ] @@ -104,6 +136,71 @@ def _load_builtin_tool(name: str) -> Any: ) from e +class ResourceResolutionError(RuntimeError): + """A control-plane resource id could not be resolved into runtime config.""" + + +HarnessResourceResolver = Callable[ + [str, HarnessResourceOverride], + HarnessResourceOverride | Mapping[str, Any] | None, +] +McpRouterResolver = Callable[[str, dict[str, Any]], Mapping[str, Any] | None] + +_resource_resolver: HarnessResourceResolver | None = None +_mcp_router_resolver: McpRouterResolver | None = None + + +def set_harness_resource_resolver( + resolver: HarnessResourceResolver | None, +) -> None: + """Register a resolver for AgentKit resource ids. + + The resolver receives a resource kind (``"knowledgebase"`` or + ``"longterm_memory"``) and the request/env resource override. It may return + another :class:`HarnessResourceOverride` or a mapping with ``type``, ``id`` + and ``config`` keys. Flat mappings are treated as ``config`` for convenience. + Passing ``None`` clears the resolver and keeps the built-in id-as-index + fallback. + """ + global _resource_resolver + _resource_resolver = resolver + + +def set_harness_mcp_router_resolver(resolver: McpRouterResolver | None) -> None: + """Register a resolver for AgentKit MCP toolset ids.""" + + global _mcp_router_resolver + _mcp_router_resolver = resolver + + +def _ensure_default_resource_resolver() -> None: + global _resource_resolver + if _resource_resolver is not None: + return + try: + from veadk.cloud.harness_app.agentkit_resources import ( + default_agentkit_resource_resolver, + ) + except ImportError as e: + logger.warning("AgentKit resource resolver is unavailable: %s", e) + return + _resource_resolver = default_agentkit_resource_resolver() + + +def _ensure_default_mcp_router_resolver() -> None: + global _mcp_router_resolver + if _mcp_router_resolver is not None: + return + try: + from veadk.cloud.harness_app.agentkit_resources import ( + default_agentkit_mcp_router_resolver, + ) + except ImportError as e: + logger.warning("AgentKit MCP router resolver is unavailable: %s", e) + return + _mcp_router_resolver = default_agentkit_mcp_router_resolver() + + # Skill hub download endpoint. A skill name in a harness is the path after # `/download/`, e.g. "namespace/owner/skill-name". SKILL_HUB_DOWNLOAD_URL = os.getenv( @@ -117,28 +214,47 @@ def _load_builtin_tool(name: str) -> Any: # populated via its "name" alias. Only variables that are set are passed, so the # model's own defaults apply to everything else. _ENV_FIELDS = { - "model_name": "MODEL_NAME", - "tools": "TOOLS", - "skills": "SKILLS", - "system_prompt": "SYSTEM_PROMPT", - "description": "DESCRIPTION", - "runtime": "RUNTIME", - "structured_tool_calls": "STRUCTURED_TOOL_CALLS", - "include_tools_every_turn": "INCLUDE_TOOLS_EVERY_TURN", - "name": "HARNESS_NAME", - "knowledgebase_type": "KNOWLEDGEBASE_TYPE", - "longterm_memory_type": "LONG_TERM_MEMORY_TYPE", - "shortterm_memory_type": "SHORT_TERM_MEMORY_TYPE", - "max_llm_calls": "MAX_LLM_CALLS", - "registry_type": "REGISTRY_TYPE", - "registry_space_id": "REGISTRY_SPACE_ID", - "registry_endpoint": "REGISTRY_ENDPOINT", - "registry_version": "REGISTRY_VERSION", - "registry_service_name": "REGISTRY_SERVICE_NAME", - "registry_region": "REGISTRY_REGION", - "registry_top_k": "REGISTRY_TOP_K", - "registry_timeout_ms": "REGISTRY_TIMEOUT_MS", - "registry_poll_interval_ms": "REGISTRY_POLL_INTERVAL_MS", + "model_name": ("MODEL_AGENT_NAME", "MODEL_NAME"), + "tools": ("TOOLS",), + "mcp_router_id": ("MCP_ROUTER_ID", "MCP_TOOLSET_ID"), + "skills": ("SKILLS",), + "system_prompt": ("SYSTEM_PROMPT",), + "description": ("DESCRIPTION",), + "runtime": ("RUNTIME",), + "temperature": ("MODEL_AGENT_TEMPERATURE",), + "top_p": ("MODEL_AGENT_TOP_P",), + "structured_tool_calls": ("STRUCTURED_TOOL_CALLS",), + "include_tools_every_turn": ("INCLUDE_TOOLS_EVERY_TURN",), + "name": ("HARNESS_NAME",), + "knowledgebase_type": ("KNOWLEDGEBASE_TYPE",), + "longterm_memory_type": ("LONG_TERM_MEMORY_TYPE",), + "shortterm_memory_type": ("SHORT_TERM_MEMORY_TYPE",), + "max_llm_calls": ("MAX_LLM_CALLS",), + "registry_type": ("REGISTRY_TYPE",), + "registry_space_id": ("REGISTRY_SPACE_ID",), + "registry_endpoint": ("REGISTRY_ENDPOINT",), + "registry_version": ("REGISTRY_VERSION",), + "registry_service_name": ("REGISTRY_SERVICE_NAME",), + "registry_region": ("REGISTRY_REGION",), + "registry_top_k": ("REGISTRY_TOP_K",), + "registry_timeout_ms": ("REGISTRY_TIMEOUT_MS",), + "registry_poll_interval_ms": ("REGISTRY_POLL_INTERVAL_MS",), +} +_HARNESS_BUILTIN_TOOL_ID_ATTR = "_veadk_harness_builtin_tool_id" +_HARNESS_MCP_SERVER_ATTR = "_veadk_harness_mcp_server" +_RUN_CODE_TOOL_ENVS = { + "run_code": "AGENTKIT_TOOL_ID_SCRIPT", + "coding": "AGENTKIT_TOOL_ID_OPENCODE", +} +_RESOURCE_DIRECT_FIELDS = { + "name", + "description", + "top_k", + "app_name", + "index", + "enable_profile", + "query_with_user_profile", + "user_id", } @@ -150,6 +266,215 @@ def split_csv(value: str) -> list[str]: return [item.strip() for item in value.split(",") if item.strip()] +def _env_value(*names: str) -> str | None: + for name in names: + value = os.environ.get(name) + if value is not None: + return value + return None + + +def _json_env(name: str, key: str) -> Any: + raw = os.environ.get(name) + if not raw or not raw.strip(): + return None + try: + value = json.loads(raw) + except json.JSONDecodeError as e: + raise ValueError(f"Invalid JSON in {name}: {e}") from e + if isinstance(value, dict): + return value.get(key) + return value + + +def _json_object_env(name: str) -> dict[str, Any]: + raw = os.environ.get(name) + if not raw or not raw.strip(): + return {} + try: + value = json.loads(raw) + except json.JSONDecodeError as e: + raise ValueError(f"Invalid JSON in {name}: {e}") from e + if not isinstance(value, dict): + raise TypeError(f"{name} must be a JSON object.") + return value + + +def _builtin_tool_from_name(name: str) -> HarnessBuiltinTool: + return HarnessBuiltinTool(id=name) + + +def _skill_from_legacy_ref(skill: str) -> HarnessSelectedSkill: + if _is_skill_center_space_ref(skill): + return HarnessSelectedSkill( + source="skillspace", + skill_space_id=_skill_center_space_id(skill), + ) + return HarnessSelectedSkill(source="skillhub", slug=skill) + + +def _selected_skill_ref(skill: HarnessSelectedSkill) -> str: + if skill.source == "skillspace": + return f"{_SKILL_CENTER_SPACE_PREFIX}{skill.skill_space_id or ''}".strip() + return (skill.slug or "").strip() + + +def _dedupe_builtin_tools( + entries: list[HarnessBuiltinTool], +) -> list[HarnessBuiltinTool]: + ordered: dict[str, HarnessBuiltinTool] = {} + for entry in entries: + tool_id = entry.id.strip() + if not tool_id: + continue + ordered[tool_id] = entry.model_copy(update={"id": tool_id}) + return list(ordered.values()) + + +def _dedupe_selected_skills( + entries: list[HarnessSelectedSkill], +) -> list[HarnessSelectedSkill]: + ordered: dict[str, HarnessSelectedSkill] = {} + for entry in entries: + ref = _selected_skill_ref(entry) + if ref: + ordered[ref] = entry + return list(ordered.values()) + + +def _builtin_tool_entries( + config: HarnessOverrides, + *, + only_set: bool, +) -> list[HarnessBuiltinTool]: + set_fields = config.model_fields_set + entries: list[HarnessBuiltinTool] = [] + if (not only_set or "tools" in set_fields) and config.tools: + entries.extend( + _builtin_tool_from_name(name) for name in split_csv(config.tools) + ) + if (not only_set or "builtin_tools" in set_fields) and config.builtin_tools: + entries.extend(config.builtin_tools) + if (not only_set or "mcp_router_id" in set_fields) and config.mcp_router_id: + entries = _add_mcp_router_id_entry(entries, config.mcp_router_id) + return _dedupe_builtin_tools(entries) + + +def _add_mcp_router_id_entry( + entries: list[HarnessBuiltinTool], + mcp_router_id: str, +) -> list[HarnessBuiltinTool]: + updated: list[HarnessBuiltinTool] = [] + found = False + for entry in entries: + if entry.id.strip() == "mcp_router": + config = {"mcp_router_id": mcp_router_id, **(entry.config or {})} + updated.append(entry.model_copy(update={"config": config})) + found = True + else: + updated.append(entry) + if not found: + updated.append( + HarnessBuiltinTool( + id="mcp_router", + config={"mcp_router_id": mcp_router_id}, + ) + ) + return updated + + +def _selected_skill_entries( + config: HarnessOverrides, + *, + only_set: bool, +) -> list[HarnessSelectedSkill]: + set_fields = config.model_fields_set + entries: list[HarnessSelectedSkill] = [] + if (not only_set or "skills" in set_fields) and config.skills: + entries.extend( + _skill_from_legacy_ref(skill) for skill in split_csv(config.skills) + ) + if (not only_set or "selected_skills" in set_fields) and config.selected_skills: + entries.extend(config.selected_skills) + return _dedupe_selected_skills(entries) + + +def normalize_harness_overrides(overrides: HarnessOverrides) -> HarnessOverrides: + """Return a canonical override using AgentKit's structured field names.""" + data = overrides.model_dump( + mode="json", + exclude_unset=True, + exclude_none=False, + ) + if "registry" in overrides.model_fields_set: + data.update(_registry_override_delta(overrides.registry)) + data.pop("registry", None) + builtin_tools = _builtin_tool_entries(overrides, only_set=True) + if builtin_tools or "tools" in overrides.model_fields_set: + data["builtin_tools"] = [ + item.model_dump(mode="json", exclude_none=True) for item in builtin_tools + ] + data.pop("tools", None) + selected_skills = _selected_skill_entries(overrides, only_set=True) + if selected_skills or "skills" in overrides.model_fields_set: + data["selected_skills"] = [ + item.model_dump(mode="json", exclude_none=True) for item in selected_skills + ] + data.pop("skills", None) + return HarnessOverrides.model_validate(data) + + +def _registry_override_delta( + registry: HarnessRegistryOverride | None, +) -> dict[str, Any]: + if registry is None: + return {} + + updates: dict[str, Any] = {} + if "space_id" in registry.model_fields_set: + updates["registry_space_id"] = registry.space_id + if "endpoint" in registry.model_fields_set: + updates["registry_endpoint"] = registry.endpoint + if "region" in registry.model_fields_set: + updates["registry_region"] = registry.region + if "top_k" in registry.model_fields_set: + updates["registry_top_k"] = registry.top_k + return updates + + +def merge_harness_overrides( + *overrides_list: HarnessOverrides | Mapping[str, Any] | None, +) -> HarnessOverrides: + """Merge Harness override layers in order, keeping AgentKit field names. + + This represents the runtime precedence used by HarnessApp: + ``env/base agent < session latest < current request``. The env/base layer is + already materialized into ``base_agent``; this helper merges the request + overlay layers without injecting unset defaults. + """ + merged: dict[str, Any] = {} + for overrides in overrides_list: + if overrides is None: + continue + if isinstance(overrides, HarnessOverrides): + normalized = normalize_harness_overrides(overrides) + else: + normalized = normalize_harness_overrides( + HarnessOverrides.model_validate(overrides) + ) + delta = normalized.model_dump( + mode="json", + exclude_unset=True, + exclude_none=False, + ) + if "builtin_tools" in delta: + merged.pop("tools", None) + if "selected_skills" in delta: + merged.pop("skills", None) + merged.update(delta) + return HarnessOverrides.model_validate(merged) + + def agent_name_from_harness(harness_name: str) -> str: """Derive a valid ADK agent name from the harness name. @@ -245,7 +570,7 @@ def _download_skill_response(name: str) -> httpx.Response: def _looks_like_zip(content: bytes) -> bool: - return content.startswith(b"PK\x03\x04") or content.startswith(b"PK\x05\x06") + return content.startswith((b"PK\x03\x04", b"PK\x05\x06")) def _resolve_skill_download_name(name: str) -> str | None: @@ -259,7 +584,7 @@ def _resolve_skill_download_name(name: str) -> str | None: if response.status_code != 200: return None data = response.json() - except Exception: + except (httpx.HTTPError, ValueError): return None for item in _skill_search_items(data): @@ -381,13 +706,248 @@ def build_skill_toolset( def config_from_env() -> HarnessConfig: """Parse the environment into a :class:`HarnessConfig` (validated by pydantic).""" kwargs: dict[str, Any] = { - field: os.environ[env] - for field, env in _ENV_FIELDS.items() - if env in os.environ + field: value + for field, env_names in _ENV_FIELDS.items() + if (value := _env_value(*env_names)) is not None } + selected_skills = _json_env("SELECTED_SKILLS_JSON", "selected_skills") + if selected_skills is not None: + kwargs["selected_skills"] = selected_skills + mcp_servers = _json_env("MCP_SERVERS_JSON", "mcp") + if mcp_servers is not None: + kwargs["mcp"] = mcp_servers + knowledgebase_id = os.environ.get("KNOWLEDGEBASE_ID") + knowledgebase_config = _json_object_env("KNOWLEDGEBASE_CONFIG_JSON") + if kwargs.get("knowledgebase_type") or knowledgebase_id: + kwargs["knowledgebase"] = { + "type": kwargs.get("knowledgebase_type", ""), + "id": knowledgebase_id, + "config": knowledgebase_config, + } + longterm_memory_id = os.environ.get("LONG_TERM_MEMORY_ID") + longterm_memory_config = _json_object_env("LONG_TERM_MEMORY_CONFIG_JSON") + if kwargs.get("longterm_memory_type") or longterm_memory_id: + kwargs["longterm_memory"] = { + "type": kwargs.get("longterm_memory_type", ""), + "id": longterm_memory_id, + "config": longterm_memory_config, + } return HarnessConfig(**kwargs) +def harness_overrides_from_env() -> HarnessOverrides: + """Expose startup env config using the same shape as ``run_sse.harness``. + + ``HarnessConfig`` includes creation-time fields such as app name and memory + backend selectors. The runtime config endpoint only speaks the request-time + ``HarnessOverrides`` contract, so this projects the startup env into that + shape while preserving the model defaults for frontend initialization. + """ + + config = config_from_env() + data = config.model_dump( + mode="json", + include=set(HarnessOverrides.model_fields), + exclude_none=True, + ) + return normalize_harness_overrides(HarnessOverrides.model_validate(data)) + + +def _with_temporary_env(tool: Any, env: dict[str, str]) -> Any: + if not env: + return tool + + if inspect.iscoroutinefunction(tool): + + @wraps(tool) + async def async_wrapped(*args, **kwargs): + old = {key: os.environ.get(key) for key in env} + os.environ.update(env) + try: + return await tool(*args, **kwargs) + finally: + _restore_env(old) + + return async_wrapped + + @wraps(tool) + def wrapped(*args, **kwargs): + old = {key: os.environ.get(key) for key in env} + os.environ.update(env) + try: + return tool(*args, **kwargs) + finally: + _restore_env(old) + + return wrapped + + +def _restore_env(values: dict[str, str | None]) -> None: + for key, value in values.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +def _configured_builtin_tool(entry: HarnessBuiltinTool) -> Any: + name = entry.id.strip() + config = dict(entry.config or {}) + if name == "mcp_router": + tool = _mcp_router_tool(config) + else: + tool = _load_builtin_tool(name) + env = _tool_env_overrides(name, config) + tool = _with_temporary_env(tool, env) + setattr(tool, _HARNESS_BUILTIN_TOOL_ID_ATTR, name) + if config: + tool._veadk_harness_tool_config = config + return tool + + +def _tool_env_overrides(name: str, config: dict[str, Any]) -> dict[str, str]: + env: dict[str, str] = {} + tool_id_env = _RUN_CODE_TOOL_ENVS.get(name) + tool_id = config.get("tool_id") + if tool_id_env and tool_id: + env[tool_id_env] = str(tool_id) + region = config.get("region") + if region: + env["AGENTKIT_TOOL_REGION"] = str(region) + return env + + +def _mcp_router_tool(config: dict[str, Any]) -> Any: + from google.adk.tools.mcp_tool.mcp_session_manager import ( + StreamableHTTPConnectionParams, + ) + from google.adk.tools.mcp_tool.mcp_toolset import McpToolset + + config = _resolve_mcp_router_config(config) + url = _config_or_env(config, "url", "url_env", "TOOL_MCP_ROUTER_URL") + api_key = _config_or_env( + config, + "api_key", + "api_key_env", + "TOOL_MCP_ROUTER_API_KEY", + ) or _config_or_env(config, "apikey", "apikey_env", "TOOL_MCP_ROUTER_API_KEY") + if not url: + raise ToolLoadError("Tool 'mcp_router' requires url or TOOL_MCP_ROUTER_URL.") + headers = {"Authorization": f"Bearer {api_key}"} if api_key else None + return McpToolset( + connection_params=StreamableHTTPConnectionParams(url=url, headers=headers) + ) + + +def _resolve_mcp_router_config(config: dict[str, Any]) -> dict[str, Any]: + mcp_router_id = ( + config.get("mcp_router_id") + or config.get("mcp_toolset_id") + or config.get("id") + or config.get("_id") + or os.environ.get("MCP_ROUTER_ID") + or os.environ.get("MCP_TOOLSET_ID") + ) + if not mcp_router_id: + return config + if _mcp_router_resolver is None: + if not _config_or_env(config, "url", "url_env", "TOOL_MCP_ROUTER_URL"): + logger.warning( + "No Harness MCP router resolver configured for id=%s; " + "falling back to explicit url/api_key config.", + mcp_router_id, + ) + return config + try: + resolved = _mcp_router_resolver(str(mcp_router_id), config) + except Exception as e: + raise ToolLoadError( + f"Failed to resolve mcp_router_id '{mcp_router_id}': {e}" + ) from e + if resolved is None: + return config + merged = dict(resolved) + merged.update(config) + return merged + + +def _config_or_env( + config: dict[str, Any], + value_key: str, + env_key: str, + default_env: str, +) -> str: + value = config.get(value_key) + if value: + return str(value) + env_name = str(config.get(env_key) or default_env) + return os.environ.get(env_name, "") + + +def _build_builtin_tools(entries: list[HarnessBuiltinTool]) -> list[Any]: + return [_configured_builtin_tool(entry) for entry in _dedupe_builtin_tools(entries)] + + +def _add_or_replace_builtin_tools( + agent: Agent, + entries: list[HarnessBuiltinTool], +) -> None: + entries = _dedupe_builtin_tools(entries) + requested = {entry.id for entry in entries} + agent.tools = [ + tool for tool in agent.tools if _builtin_tool_id(tool) not in requested + ] + agent.tools.extend(_build_builtin_tools(entries)) + + +def _builtin_tool_id(tool: Any) -> str | None: + return getattr(tool, _HARNESS_BUILTIN_TOOL_ID_ATTR, None) or _tool_name(tool) + + +def _mcp_server_key(server: HarnessMcpServer) -> str: + return server.name.strip() or server.server_url.strip() + + +def _build_mcp_toolset(server: HarnessMcpServer) -> Any: + from google.adk.tools.mcp_tool.mcp_session_manager import ( + StreamableHTTPConnectionParams, + ) + from google.adk.tools.mcp_tool.mcp_toolset import McpToolset + + url = server.server_url.strip() + if not url: + raise ToolLoadError("MCP server requires server_url.") + headers = ( + {"Authorization": f"Bearer {server.bear_token}"} if server.bear_token else None + ) + toolset = McpToolset( + connection_params=StreamableHTTPConnectionParams(url=url, headers=headers) + ) + setattr(toolset, _HARNESS_MCP_SERVER_ATTR, _mcp_server_key(server)) + return toolset + + +def _build_mcp_toolsets(servers: list[HarnessMcpServer]) -> list[Any]: + return [_build_mcp_toolset(server) for server in servers] + + +def _remove_harness_mcp_toolsets(agent: Agent) -> None: + agent.tools = [ + tool for tool in agent.tools if not getattr(tool, _HARNESS_MCP_SERVER_ATTR, "") + ] + + +def _add_mcp_toolsets(agent: Agent, servers: list[HarnessMcpServer]) -> None: + if not servers: + return + agent.tools.extend(_build_mcp_toolsets(servers)) + + +def _selected_skill_refs(skills: list[HarnessSelectedSkill]) -> list[str]: + refs = [_selected_skill_ref(skill) for skill in skills] + return [ref for ref in refs if ref] + + def _assemble_agent(config: HarnessConfig) -> tuple[Agent, ShortTermMemory]: """Build an agent and its short-term memory from a :class:`HarnessConfig`. @@ -396,14 +956,15 @@ def _assemble_agent(config: HarnessConfig) -> tuple[Agent, ShortTermMemory]: base / long-term memory. Backend values are validated by each component's pydantic model (fast-fail on an unknown value). """ - tools = [_load_builtin_tool(name) for name in split_csv(config.tools)] + tools = _build_builtin_tools(_builtin_tool_entries(config, only_set=False)) - skills = split_csv(config.skills) + skills = _selected_skill_refs(_selected_skill_entries(config, only_set=False)) if skills: logger.info(f"Loading skills {skills} for harness.") skill_toolset = build_skill_toolset(skills) if skill_toolset is not None: tools.append(skill_toolset) + tools.extend(_build_mcp_toolsets(config.mcp)) registry_config = None if config.registry_type: @@ -426,25 +987,50 @@ def _assemble_agent(config: HarnessConfig) -> tuple[Agent, ShortTermMemory]: tools.extend(build_a2a_registry_tools(registry_config)) knowledgebase = None - if config.knowledgebase_type: + knowledgebase_override = config.knowledgebase + if knowledgebase_override is None and config.knowledgebase_type: + knowledgebase_override = HarnessResourceOverride( + type=config.knowledgebase_type, + ) + if knowledgebase_override: + knowledgebase_override = _resolve_resource_override( + "knowledgebase", knowledgebase_override + ) + if knowledgebase_override and knowledgebase_override.type: logger.info( - f"Initializing knowledge base: backend={config.knowledgebase_type} " + f"Initializing knowledge base: backend={knowledgebase_override.type} " f"index={config.app_name}" ) knowledgebase = KnowledgeBase( - backend=config.knowledgebase_type, # type: ignore[arg-type] - app_name=config.app_name, + backend=knowledgebase_override.type, # type: ignore[arg-type] + **_request_resource_config( + "knowledgebase", knowledgebase_override, config.app_name, resolve=False + ), ) long_term_memory = None - if config.longterm_memory_type: + longterm_memory_override = config.longterm_memory + if longterm_memory_override is None and config.longterm_memory_type: + longterm_memory_override = HarnessResourceOverride( + type=config.longterm_memory_type, + ) + if longterm_memory_override: + longterm_memory_override = _resolve_resource_override( + "longterm_memory", longterm_memory_override + ) + if longterm_memory_override and longterm_memory_override.type: logger.info( - f"Initializing long-term memory: backend={config.longterm_memory_type} " + f"Initializing long-term memory: backend={longterm_memory_override.type} " f"index={config.app_name}" ) long_term_memory = LongTermMemory( - backend=config.longterm_memory_type, # type: ignore[arg-type] - app_name=config.app_name, + backend=longterm_memory_override.type, # type: ignore[arg-type] + **_request_resource_config( + "longterm_memory", + longterm_memory_override, + config.app_name, + resolve=False, + ), ) logger.info( @@ -465,8 +1051,10 @@ def _assemble_agent(config: HarnessConfig) -> tuple[Agent, ShortTermMemory]: enable_responses_cache=not config.include_tools_every_turn, knowledgebase=knowledgebase, long_term_memory=long_term_memory, + auto_save_session=long_term_memory is not None, short_term_memory=short_term_memory, ) + _apply_sampling_overrides(agent, config) if registry_config is not None: setattr(agent, _REGISTRY_CONFIG_ATTR, registry_config) return agent, short_term_memory @@ -479,6 +1067,8 @@ def init_harness_agent() -> tuple[Agent, ShortTermMemory]: A ``(agent, short_term_memory)`` tuple. The short-term memory is returned separately so the server can share the same instance with its ``Runner``. """ + _ensure_default_resource_resolver() + _ensure_default_mcp_router_resolver() return _assemble_agent(config_from_env()) @@ -487,62 +1077,309 @@ def _tool_name(tool: Any) -> str | None: return getattr(tool, "__name__", None) or getattr(tool, "name", None) -def _add_incremental_tools(agent: Agent, tool_names: list[str]) -> None: - """Append the requested built-in tools, skipping ones already on the agent.""" - existing = {name for tool in agent.tools if (name := _tool_name(tool))} - for name in tool_names: - if name in existing: - logger.info(f"Tool '{name}' already on the agent; skipping.") - continue - agent.tools.append(_load_builtin_tool(name)) - existing.add(name) +def _replace_builtin_tools( + agent: Agent, + entries: list[HarnessBuiltinTool], +) -> None: + """Replace the harness-selected built-in tools with ``entries``.""" + + agent.tools = [tool for tool in agent.tools if not _is_harness_builtin_tool(tool)] + agent.tools.extend(_build_builtin_tools(entries)) + + +def _is_harness_builtin_tool(tool: Any) -> bool: + if getattr(tool, _HARNESS_BUILTIN_TOOL_ID_ATTR, None): + return True + name = _tool_name(tool) + return bool(name and name in set(list_builtin_tools())) -def _add_incremental_skills( +def _replace_skills( agent: Agent, skill_ids: list[str], download_dir: Path | None = None ) -> None: - """Mount the requested skills, skipping ones whose name is already loaded. + """Replace the harness-selected skill toolset with ``skill_ids``.""" - Skills already present are dropped (deduped by skill name). Any genuinely new - skills are merged into the agent's existing :class:`SkillToolset` so the agent - keeps a single toolset (two would expose duplicate ``list_skills``/``load_skill`` - tools); if the agent has none yet, a new toolset is mounted. ``download_dir`` - is where the skills are downloaded (cleaned up by the caller after the run). - """ + agent.tools = [tool for tool in agent.tools if not isinstance(tool, SkillToolset)] toolset = build_skill_toolset(skill_ids, download_dir=download_dir) - if toolset is None: - return - new_skills = toolset._list_skills() - - existing_toolset = next( - (tool for tool in agent.tools if isinstance(tool, SkillToolset)), None - ) - if existing_toolset is None: + if toolset is not None: agent.tools.append(toolset) - return - existing_skills = existing_toolset._list_skills() - existing_names = {skill.name for skill in existing_skills} - new_skills = [skill for skill in new_skills if skill.name not in existing_names] - if not new_skills: - logger.info("All requested skills already loaded; skipping.") - return - agent.tools.remove(existing_toolset) - agent.tools.append( - SkillToolset( - skills=existing_skills + new_skills, - code_executor=(existing_toolset._code_executor or toolset._code_executor), +def _remove_a2a_registry_tools(agent: Agent) -> None: + agent.tools = [ + tool for tool in agent.tools if _tool_name(tool) not in _REGISTRY_TOOL_NAMES + ] + + +def _request_resource_config( + kind: str, + resource: HarnessResourceOverride, + app_name: str | None, + *, + resolve: bool = True, +) -> dict[str, Any]: + if resolve: + resource = _resolve_resource_override(kind, resource) + config = dict(resource.config or {}) + configured_backend_config = config.pop("backend_config", None) + if resource.id: + config.setdefault("index", resource.id) + config.setdefault("app_name", resource.id) + elif app_name: + config.setdefault("app_name", app_name) + config.pop("type", None) + config.pop("backend", None) + direct = { + key: value for key, value in config.items() if key in _RESOURCE_DIRECT_FIELDS + } + backend_config = { + key: value + for key, value in config.items() + if key not in _RESOURCE_DIRECT_FIELDS + } + if isinstance(configured_backend_config, Mapping): + backend_config = {**configured_backend_config, **backend_config} + elif configured_backend_config: + backend_config["backend_config"] = configured_backend_config + if backend_config: + if "index" in direct: + backend_config.setdefault("index", direct["index"]) + if "app_name" in direct: + backend_config.setdefault("app_name", direct["app_name"]) + direct["backend_config"] = backend_config + return direct + + +def _resolve_resource_override( + kind: str, + resource: HarnessResourceOverride, +) -> HarnessResourceOverride: + if not resource.id: + return resource + if _resource_resolver is None: + if not resource.config: + logger.warning( + "No Harness resource resolver configured for %s id=%s; " + "falling back to id as index/app_name.", + kind, + resource.id, + ) + return resource + try: + resolved = _resource_resolver(kind, resource) + except Exception as e: + raise ResourceResolutionError( + f"Failed to resolve {kind} resource '{resource.id}': {e}" + ) from e + if resolved is None: + if resource.config: + logger.warning( + "Harness resource resolver returned no config for %s id=%s; " + "using the explicit request config.", + kind, + resource.id, + ) + return resource + raise ResourceResolutionError( + f"No runtime config found for {kind} resource '{resource.id}'. " + "Provide config explicitly or register a Harness resource resolver." ) + return _merge_resolved_resource(resource, resolved) + + +def _merge_resolved_resource( + requested: HarnessResourceOverride, + resolved: HarnessResourceOverride | Mapping[str, Any], +) -> HarnessResourceOverride: + resolved_resource = _coerce_resolved_resource(requested, resolved) + config = dict(resolved_resource.config or {}) + config.update(requested.config or {}) + return HarnessResourceOverride( + type=requested.type or resolved_resource.type, + id=requested.id or resolved_resource.id, + config=config, ) -def _remove_a2a_registry_tools(agent: Agent) -> None: +def _coerce_resolved_resource( + requested: HarnessResourceOverride, + resolved: HarnessResourceOverride | Mapping[str, Any], +) -> HarnessResourceOverride: + if isinstance(resolved, HarnessResourceOverride): + return resolved + raw = dict(resolved) + override_keys = {"type", "id", "_id", "config"} + override_data = {key: raw[key] for key in override_keys if key in raw} + flat_config = {key: value for key, value in raw.items() if key not in override_keys} + if override_data: + resolved_resource = HarnessResourceOverride.model_validate(override_data) + if flat_config: + config = dict(resolved_resource.config or {}) + config.update(flat_config) + resolved_resource = resolved_resource.model_copy(update={"config": config}) + return resolved_resource + return HarnessResourceOverride( + type=requested.type, + id=requested.id, + config=flat_config, + ) + + +def _remove_knowledgebase_tools(agent: Agent) -> None: agent.tools = [ - tool for tool in agent.tools if _tool_name(tool) not in _REGISTRY_TOOL_NAMES + tool + for tool in agent.tools + if not isinstance(tool, LoadKnowledgebaseTool) + and _tool_name(tool) not in _KNOWLEDGEBASE_TOOL_NAMES ] +def _mount_knowledgebase_tools(agent: Agent) -> None: + if not agent.knowledgebase: + return + + agent.tools.append(LoadKnowledgebaseTool(knowledgebase=agent.knowledgebase)) + if agent.knowledgebase.enable_profile: + from veadk.tools.builtin_tools.load_kb_queries import load_kb_queries + + agent.tools.append(load_kb_queries) + + +def _remove_longterm_memory_tools(agent: Agent) -> None: + agent.tools = [ + tool + for tool in agent.tools + if _tool_name(tool) not in _LONGTERM_MEMORY_TOOL_NAMES + ] + + +def _mount_longterm_memory_tools(agent: Agent) -> None: + if agent.long_term_memory is None: + return + + from google.adk.tools.load_memory_tool import LoadMemoryTool + + load_memory_tool = LoadMemoryTool() + if hasattr(load_memory_tool, "custom_metadata"): + if not load_memory_tool.custom_metadata: + load_memory_tool.custom_metadata = {} + load_memory_tool.custom_metadata["backend"] = agent.long_term_memory.backend + agent.tools.append(load_memory_tool) + + +def _set_longterm_memory_auto_save(agent: Agent, enabled: bool) -> None: + from veadk.memory.save_session_callback import save_session_to_long_term_memory + + agent.auto_save_session = enabled + callback = getattr(agent, "after_agent_callback", None) + + if enabled: + if callback is None: + agent.after_agent_callback = save_session_to_long_term_memory + elif isinstance(callback, list): + if save_session_to_long_term_memory not in callback: + callback.append(save_session_to_long_term_memory) + elif callback is not save_session_to_long_term_memory: + agent.after_agent_callback = [callback, save_session_to_long_term_memory] + return + + if callback is save_session_to_long_term_memory: + agent.after_agent_callback = None + elif isinstance(callback, list): + callbacks = [ + item for item in callback if item is not save_session_to_long_term_memory + ] + if not callbacks: + agent.after_agent_callback = None + elif len(callbacks) == 1: + agent.after_agent_callback = callbacks[0] + else: + agent.after_agent_callback = callbacks + + +def _apply_resource_overrides( + agent: Agent, + overrides: HarnessOverrides, + app_name: str | None, +) -> None: + set_fields = overrides.model_fields_set + + if "knowledgebase" in set_fields: + _remove_knowledgebase_tools(agent) + knowledgebase_override = overrides.knowledgebase + if knowledgebase_override: + knowledgebase_override = _resolve_resource_override( + "knowledgebase", knowledgebase_override + ) + if knowledgebase_override and knowledgebase_override.type: + agent.knowledgebase = KnowledgeBase( + backend=knowledgebase_override.type, # type: ignore[arg-type] + **_request_resource_config( + "knowledgebase", knowledgebase_override, app_name, resolve=False + ), + ) + _mount_knowledgebase_tools(agent) + else: + agent.knowledgebase = None + + if "longterm_memory" in set_fields: + _remove_longterm_memory_tools(agent) + longterm_memory_override = overrides.longterm_memory + if longterm_memory_override: + longterm_memory_override = _resolve_resource_override( + "longterm_memory", longterm_memory_override + ) + if longterm_memory_override and longterm_memory_override.type: + agent.long_term_memory = LongTermMemory( + backend=longterm_memory_override.type, # type: ignore[arg-type] + **_request_resource_config( + "longterm_memory", + longterm_memory_override, + app_name, + resolve=False, + ), + ) + _mount_longterm_memory_tools(agent) + _set_longterm_memory_auto_save(agent, True) + else: + agent.long_term_memory = None + _set_longterm_memory_auto_save(agent, False) + + +def _apply_sampling_overrides(agent: Agent, overrides: HarnessOverrides) -> None: + set_fields = overrides.model_fields_set + if not (_SAMPLING_OVERRIDE_FIELDS & set_fields): + return + + updates: dict[str, Any] = {} + if "temperature" in set_fields and overrides.temperature is not None: + updates["temperature"] = overrides.temperature + if "top_p" in set_fields and overrides.top_p is not None: + updates["top_p"] = overrides.top_p + if "max_tokens" in set_fields and overrides.max_tokens is not None: + updates["max_output_tokens"] = overrides.max_tokens + if "presence_penalty" in set_fields and overrides.presence_penalty is not None: + updates["presence_penalty"] = overrides.presence_penalty + if "frequency_penalty" in set_fields and overrides.frequency_penalty is not None: + updates["frequency_penalty"] = overrides.frequency_penalty + if "penalty" in set_fields and overrides.penalty is not None: + if "presence_penalty" not in updates: + updates["presence_penalty"] = overrides.penalty + if "frequency_penalty" not in updates: + updates["frequency_penalty"] = overrides.penalty + + if not updates: + return + + base_config = getattr(agent, "generate_content_config", None) + generate_content_config = ( + base_config.model_copy(deep=True) + if base_config is not None + else types.GenerateContentConfig() + ) + agent.generate_content_config = generate_content_config.model_copy(update=updates) + + def _apply_registry_overrides( agent: Agent, base_config, @@ -627,20 +1464,22 @@ def _add_dynamic_a2a_agent_tools(agent: Agent, prompt: str) -> None: def spawn_harness_agent( - base_agent: Agent, overrides: HarnessOverrides, download_dir: Path | None = None + base_agent: Agent, + overrides: HarnessOverrides, + download_dir: Path | None = None, + app_name: str | None = None, ) -> Agent: """Clone the base agent for a one-off invocation and apply per-request overrides. - Uses ADK's :meth:`~google.adk.agents.base_agent.BaseAgent.clone`, so the clone - inherits the base agent's knowledge base and memory — these are never - overridable. Only the fields the request actually set are applied: ``model_name``, - ``system_prompt`` and ``runtime`` replace the base value, while ``tools`` and - ``skills`` are mounted *incrementally* — anything already on the agent (same - tool name / skill name) is skipped, so only the delta is added. + Uses ADK's :meth:`~google.adk.agents.base_agent.BaseAgent.clone`, then applies + only the fields the request actually set. ``model_name``, ``system_prompt``, + ``runtime``, sampling params, tools, skills, knowledge base, and long-term + memory replace the clone's value. - ``download_dir`` is where any incremental skills are downloaded; the caller - owns it and should remove it once the invocation finishes. + ``download_dir`` is where any skills are downloaded; the caller owns it and + should remove it once the invocation finishes. """ + overrides = normalize_harness_overrides(overrides) set_fields = overrides.model_fields_set update: dict[str, Any] = {} @@ -653,12 +1492,25 @@ def spawn_harness_agent( if "model_name" in set_fields: cloned.update_model(overrides.model_name) - if "tools" in set_fields: - _add_incremental_tools(cloned, split_csv(overrides.tools)) + if "builtin_tools" in set_fields: + _replace_builtin_tools( + cloned, + _builtin_tool_entries(overrides, only_set=True), + ) + + if "selected_skills" in set_fields: + _replace_skills( + cloned, + _selected_skill_refs(_selected_skill_entries(overrides, only_set=True)), + download_dir, + ) - if "skills" in set_fields: - _add_incremental_skills(cloned, split_csv(overrides.skills), download_dir) + if "mcp" in set_fields: + _remove_harness_mcp_toolsets(cloned) + _add_mcp_toolsets(cloned, overrides.mcp) + _apply_sampling_overrides(cloned, overrides) + _apply_resource_overrides(cloned, overrides, app_name) _apply_registry_overrides( cloned, getattr(base_agent, _REGISTRY_CONFIG_ATTR, None), @@ -673,13 +1525,28 @@ def spawn_harness_run_agent( prompt: str, overrides: HarnessOverrides | None = None, download_dir: Path | None = None, + app_name: str | None = None, registry_tip_token: str = "", registry_authorization: str = "", + session_overrides: HarnessOverrides | Mapping[str, Any] | None = None, + current_overrides: HarnessOverrides | Mapping[str, Any] | None = None, ) -> Agent: """Clone a harness agent for one run and attach per-turn dynamic tools.""" + if session_overrides is not None or current_overrides is not None: + overrides = merge_harness_overrides( + overrides, + session_overrides, + current_overrides, + ) + if overrides is not None: - cloned = spawn_harness_agent(base_agent, overrides, download_dir=download_dir) + cloned = spawn_harness_agent( + base_agent, + overrides, + download_dir=download_dir, + app_name=app_name, + ) else: cloned = base_agent.clone(update={}) diff --git a/veadk/tools/builtin_tools/mcp_router.py b/veadk/tools/builtin_tools/mcp_router.py index 3be329333..7ae9cc503 100644 --- a/veadk/tools/builtin_tools/mcp_router.py +++ b/veadk/tools/builtin_tools/mcp_router.py @@ -12,18 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset +import os -from veadk.config import getenv from google.adk.tools.mcp_tool.mcp_session_manager import ( StreamableHTTPConnectionParams, ) +from google.adk.tools.mcp_tool.mcp_toolset import McpToolset -url = getenv("TOOL_MCP_ROUTER_URL") -api_key = getenv("TOOL_MCP_ROUTER_API_KEY") +url = os.getenv("TOOL_MCP_ROUTER_URL", "") +api_key = os.getenv("TOOL_MCP_ROUTER_API_KEY", "") -mcp_router = MCPToolset( +mcp_router = McpToolset( connection_params=StreamableHTTPConnectionParams( - url=url, headers={"Authorization": f"Bearer {api_key}"} - ), + url=url, + headers={"Authorization": f"Bearer {api_key}"} if api_key else None, + ) )