From 5e38fa22c4f87376b1ce05843a4e117ddf98b104 Mon Sep 17 00:00:00 2001 From: Million <15158090088@163.com> Date: Wed, 19 Aug 2026 16:53:47 +0800 Subject: [PATCH] feat(viking): tool filtering, error diagnostics, and URI-prefix instructions - viking_search/find accept multiple target_uri prefixes and default to the full allowed_uri_prefixes allowlist (previously only searched the first) - get_instructions() renders a dynamic Allowed URI Prefixes block so the model passes target_uri directly and skips discovery probing - all 15 viking_* tools include the exception type in error returns so empty-message failures like ReadTimeout('') remain diagnosable - VikingCapabilityConfig gains enabled_tools/disabled_tools whitelist/ blacklist filtering (mutually exclusive) to disable slow semantic-search tools --- .../2026-08-18-viking-allowed-uri-prefixes.md | 18 ++- ...026-08-19-viking-tool-error-diagnostics.md | 18 +++ .../2026-08-19-viking-tool-filtering.md | 26 ++++ .../capabilities/viking/__init__.py | 21 ++- .../capabilities/viking/instructions.py | 33 +++++ src/wolfharness/capabilities/viking/tools.py | 87 +++++++---- src/wolfharness_config/capabilities.py | 28 +++- tests/capabilities/viking/test_viking.py | 138 +++++++++++++++--- .../viking/test_viking_integration.py | 86 ++++++++++- 9 files changed, 396 insertions(+), 59 deletions(-) create mode 100644 changelog/unreleased/2026-08-19-viking-tool-error-diagnostics.md create mode 100644 changelog/unreleased/2026-08-19-viking-tool-filtering.md diff --git a/changelog/unreleased/2026-08-18-viking-allowed-uri-prefixes.md b/changelog/unreleased/2026-08-18-viking-allowed-uri-prefixes.md index bbda3370a..3fa8111d0 100644 --- a/changelog/unreleased/2026-08-18-viking-allowed-uri-prefixes.md +++ b/changelog/unreleased/2026-08-18-viking-allowed-uri-prefixes.md @@ -19,4 +19,20 @@ single subtree such as `viking://resources/wiki/` without also allowing Empty list (the default) preserves unrestricted behavior for backward compatibility. `viking_search`/`viking_find` without a `target_uri` -automatically scope to the first allowed prefix when a allowlist is set. \ No newline at end of file +automatically scope to the configured allowlist when one is set (see below). + +## Search scoping covers all allowed prefixes (2026-08-19) + +When a `target_uri` is omitted, `viking_search`/`viking_find` previously +scoped to only the **first** allowed prefix, silently dropping results from +the other allowed trees. `target_uri` now accepts a list, and both tools pass +**every** allowed prefix to the SDK — the server searches each tree and the +result set covers the full allowlist. Explicit `str` targets are still +validated against the allowlist as before; explicit `list` targets are +validated element-wise. + +Alongside this, `get_instructions()` now renders a dynamic **Allowed URI +Prefixes** section listing the exact prefixes when the allowlist is +configured. The model sees the scoping boundary up front, so it can pass the +most specific `target_uri` directly and skip discovery probing (`viking_ls`), +which also avoids the slower whole-allowlist search. \ No newline at end of file diff --git a/changelog/unreleased/2026-08-19-viking-tool-error-diagnostics.md b/changelog/unreleased/2026-08-19-viking-tool-error-diagnostics.md new file mode 100644 index 000000000..8053ac5dc --- /dev/null +++ b/changelog/unreleased/2026-08-19-viking-tool-error-diagnostics.md @@ -0,0 +1,18 @@ +# Viking tool errors now include the exception type + +All 15 `viking_*` tools previously rendered failures as +`viking_search error: {e}`, relying on `str(e)` for the diagnostic text. +Some exceptions carry an empty message — notably `httpx.ReadTimeout('')` +when a slow knowledge-graph search exceeds the configured timeout — which +produced a useless `viking_search error:` with no trailing context and no +indication of what went wrong. + +Tool error returns now follow `viking_search error ({ExcType}): {e}`, so +an empty-message timeout renders as `viking_search error (ReadTimeout):` +and the failure class is always identifiable even when the message is +blank. This applies uniformly to all 15 tools (search, find, recall, +grep, glob, ls, read, expand, write, edit, mkdir, add_resource, forget, +link, set_tags). + +Also adds a regression test asserting that empty-message exceptions still +surface their exception type in the tool return value. \ No newline at end of file diff --git a/changelog/unreleased/2026-08-19-viking-tool-filtering.md b/changelog/unreleased/2026-08-19-viking-tool-filtering.md new file mode 100644 index 000000000..5675eb4d5 --- /dev/null +++ b/changelog/unreleased/2026-08-19-viking-tool-filtering.md @@ -0,0 +1,26 @@ +# Viking tool filtering via enabled_tools / disabled_tools + +`VikingCapabilityConfig` gains two mutually-exclusive tool-filter fields, +mirroring `StdioMCPServerConfig`: + +- `enabled_tools` — whitelist: only these `viking_*` tools are exposed. +- `disabled_tools` — blacklist: these tools are excluded from the exposed set. + +The filter applies in `build_tools()` after mode-based assembly, so it works +for every mode (`retrieve`/`write`/`graph`/`all`) and propagates to both the +toolset and the OpenCode `/experimental/tool` listing via `get_tools()`. + +Motivation: the knowledge-graph semantic search backend +(`viking_search`/`viking_find`) can be slow on large resource trees (40-60s+ +per query), exceeding typical client timeouts. Operators can now disable +just those tools while keeping the deterministic ones: + +```yaml +capabilities: + - type: viking + mode: retrieve + disabled_tools: ["viking_search", "viking_find"] +``` + +Specifying both `enabled_tools` and `disabled_tools` raises a validation +error (mutually exclusive), matching the MCP server config pattern. \ No newline at end of file diff --git a/src/wolfharness/capabilities/viking/__init__.py b/src/wolfharness/capabilities/viking/__init__.py index 55fefd7b7..30ac49cfc 100644 --- a/src/wolfharness/capabilities/viking/__init__.py +++ b/src/wolfharness/capabilities/viking/__init__.py @@ -217,6 +217,14 @@ class VikingCapability(AbstractCapability[Any]): """When True (default), profile injection runs only on the first turn of a session. When False, injection runs on every ``before_model_request`` call (not recommended — expensive and static).""" + enabled_tools: list[str] | None = None + """If set, only these tools are exposed (whitelist). Mutually exclusive + with ``disabled_tools``.""" + disabled_tools: list[str] | None = None + """Tools to exclude from the exposed set (blacklist). Mutually exclusive + with ``enabled_tools``. For example, disable a slow semantic-search + backend while keeping deterministic tools: + ``["viking_search", "viking_find"]``.""" compaction_enabled: bool = False """When True, archive old conversation messages to Viking before context overflow. Disabled by default.""" @@ -545,12 +553,21 @@ async def for_run(self, ctx: RunContext[Any]) -> VikingCapability: def get_instructions(self) -> str | None: """Return the Viking workflow instructions. + When ``allowed_uri_prefixes`` is configured, appends a dynamic + block listing the allowed prefixes so the model can pass a + ``target_uri`` and skip discovery probing. + Returns: The instruction string from ``instructions.py``. """ - from wolfharness.capabilities.viking.instructions import _VIKING_INSTRUCTIONS + from wolfharness.capabilities.viking.instructions import ( + _VIKING_INSTRUCTIONS, + format_allowed_prefixes_block, + ) - return _VIKING_INSTRUCTIONS + if not self.allowed_uri_prefixes: + return _VIKING_INSTRUCTIONS + return _VIKING_INSTRUCTIONS + format_allowed_prefixes_block(self.allowed_uri_prefixes) def get_toolset(self) -> AgentToolset[Any] | None: """Build a ``FunctionToolset`` from tools filtered by ``self.mode``. diff --git a/src/wolfharness/capabilities/viking/instructions.py b/src/wolfharness/capabilities/viking/instructions.py index 00317c3ae..3d14d08ab 100644 --- a/src/wolfharness/capabilities/viking/instructions.py +++ b/src/wolfharness/capabilities/viking/instructions.py @@ -11,6 +11,39 @@ from __future__ import annotations +from typing import TYPE_CHECKING + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +def format_allowed_prefixes_block(prefixes: Sequence[str]) -> str: + """Build the allowed-prefix instruction block. + + Appended to the base instructions when ``allowed_uri_prefixes`` is + configured. Rendered dynamically at ``get_instructions()`` time so the + model sees the exact prefixes this session may access — it can then pass + the most specific prefix as ``target_uri`` and skip discovery probing. + + Args: + prefixes: The configured allowed URI prefixes. + + Returns: + A markdown block listing the prefixes and the search guidance. + """ + allowed = "\n".join(f" - `{p}`" for p in prefixes) + return ( + "\n### Allowed URI Prefixes (this session)\n" + "Your Viking access is restricted to these prefixes:\n" + f"{allowed}\n" + "When calling `viking_search` / `viking_find`, ALWAYS pass the most " + "specific matching prefix as `target_uri` — it scopes the search and " + "is faster. Omitting it searches all allowed prefixes, which is slower. " + "Use `viking_ls` on these prefixes to see what is in scope. Do not " + "probe outside them — you will get access errors." + ) + _VIKING_INSTRUCTIONS = """\ ## Viking Knowledge Graph Tools diff --git a/src/wolfharness/capabilities/viking/tools.py b/src/wolfharness/capabilities/viking/tools.py index 3654d5b6a..1fe1f7a7b 100644 --- a/src/wolfharness/capabilities/viking/tools.py +++ b/src/wolfharness/capabilities/viking/tools.py @@ -93,7 +93,7 @@ async def viking_search( limit: int = 10, min_score: float = 0.35, level: list[int] | None = None, - target_uri: str = "", + target_uri: str | list[str] = "", ) -> ToolReturn: """Search the Viking knowledge graph semantically. @@ -105,7 +105,8 @@ async def viking_search( limit: Maximum number of results to return. min_score: Minimum relevance score (0.0 to 1.0). level: Filter by content level (e.g. [0, 1, 2] for L0-L2). - target_uri: Restrict search to a specific URI subtree. + target_uri: Restrict search to specific URI subtrees — a + single ``viking://`` URI or a list of them. Returns: Formatted search results grouped by context type. @@ -115,12 +116,22 @@ async def viking_search( sid = _get_session_id(ctx) sdk_filter: dict[str, Any] | None = {"level": level} if level else None if cap.allowed_uri_prefixes: - if target_uri: - err = cap._check_uri_allowed(target_uri, tool_name="viking_search") - if err: - return ToolReturn(return_value=err) + if isinstance(target_uri, str): + if target_uri: + err = cap._check_uri_allowed(target_uri, tool_name="viking_search") + if err: + return ToolReturn(return_value=err) else: - target_uri = cap.allowed_uri_prefixes[0] + for u in target_uri: + err = cap._check_uri_allowed(u, tool_name="viking_search") + if err: + return ToolReturn(return_value=err) + if not target_uri: + # SDK target_uri accepts a list — the server searches + # every allowed prefix. The old default used only the + # first prefix, silently dropping the other allowed + # trees. + target_uri = cap.allowed_uri_prefixes result = await client.search( query, target_uri=target_uri, @@ -131,7 +142,7 @@ async def viking_search( ) return ToolReturn(return_value=format_search_results(result)) except Exception as e: - return ToolReturn(return_value=f"viking_search error: {e}") + return ToolReturn(return_value=f"viking_search error ({type(e).__name__}): {e}") async def viking_find( ctx: RunContext[Any], @@ -139,7 +150,7 @@ async def viking_find( limit: int = 10, min_score: float = 0.35, level: list[int] | None = None, - target_uri: str = "", + target_uri: str | list[str] = "", ) -> ToolReturn: """Find content in Viking, deduplicating results. @@ -151,7 +162,8 @@ async def viking_find( limit: Maximum number of results to return. min_score: Minimum relevance score (0.0 to 1.0). level: Filter by content level (e.g. [0, 1, 2] for L0-L2). - target_uri: Restrict search to a specific URI subtree. + target_uri: Restrict search to specific URI subtrees — a + single ``viking://`` URI or a list of them. Returns: Formatted search results grouped by context type. @@ -160,12 +172,20 @@ async def viking_find( client = await cap._ensure_client() sdk_filter: dict[str, Any] | None = {"level": level} if level else None if cap.allowed_uri_prefixes: - if target_uri: - err = cap._check_uri_allowed(target_uri, tool_name="viking_find") - if err: - return ToolReturn(return_value=err) + if isinstance(target_uri, str): + if target_uri: + err = cap._check_uri_allowed(target_uri, tool_name="viking_find") + if err: + return ToolReturn(return_value=err) else: - target_uri = cap.allowed_uri_prefixes[0] + for u in target_uri: + err = cap._check_uri_allowed(u, tool_name="viking_find") + if err: + return ToolReturn(return_value=err) + if not target_uri: + # See viking_search: SDK target_uri accepts a list so we + # search every allowed prefix, not just the first. + target_uri = cap.allowed_uri_prefixes result = await client.find( query, target_uri=target_uri, @@ -175,7 +195,7 @@ async def viking_find( ) return ToolReturn(return_value=format_search_results(result)) except Exception as e: - return ToolReturn(return_value=f"viking_find error: {e}") + return ToolReturn(return_value=f"viking_find error ({type(e).__name__}): {e}") async def viking_recall( ctx: RunContext[Any], @@ -243,7 +263,7 @@ async def viking_recall( merged = "\n\n".join(sections) return ToolReturn(return_value=truncate_text(merged, max_chars)) except Exception as e: - return ToolReturn(return_value=f"viking_recall error: {e}") + return ToolReturn(return_value=f"viking_recall error ({type(e).__name__}): {e}") async def viking_grep( ctx: RunContext[Any], @@ -300,7 +320,7 @@ async def _grep_one(p: str) -> list[dict[str, Any]]: return ToolReturn(return_value=format_grep_results(all_matches, patterns)) except Exception as e: - return ToolReturn(return_value=f"viking_grep error: {e}") + return ToolReturn(return_value=f"viking_grep error ({type(e).__name__}): {e}") async def viking_glob( ctx: RunContext[Any], @@ -333,7 +353,7 @@ async def viking_glob( uris = result if isinstance(result, list) else [] return ToolReturn(return_value=format_glob_results([str(u) for u in uris], pattern)) except Exception as e: - return ToolReturn(return_value=f"viking_glob error: {e}") + return ToolReturn(return_value=f"viking_glob error ({type(e).__name__}): {e}") async def viking_ls( ctx: RunContext[Any], @@ -416,7 +436,7 @@ async def _safe_abstract(entry_uri: str) -> str: return ToolReturn(return_value=format_ls_entries(entry_list)) except Exception as e: - return ToolReturn(return_value=f"viking_ls error: {e}") + return ToolReturn(return_value=f"viking_ls error ({type(e).__name__}): {e}") async def viking_read( ctx: RunContext[Any], @@ -510,7 +530,7 @@ async def viking_read( ) return ToolReturn(return_value="\n\n".join(sections)) except Exception as e: - return ToolReturn(return_value=f"viking_read error: {e}") + return ToolReturn(return_value=f"viking_read error ({type(e).__name__}): {e}") async def viking_expand( ctx: RunContext[Any], @@ -539,7 +559,7 @@ async def viking_expand( return_value=str(content) if content else "No content found at URI." ) except Exception as e: - return ToolReturn(return_value=f"viking_expand error: {e}") + return ToolReturn(return_value=f"viking_expand error ({type(e).__name__}): {e}") retrieve_tools: list[Callable[..., Awaitable[ToolReturn]]] = [ viking_search, @@ -618,7 +638,7 @@ async def viking_write( return_value=f"Wrote {len(content)} chars to {uri} (mode={mode})." ) except Exception as e: - return ToolReturn(return_value=f"viking_write error: {e}") + return ToolReturn(return_value=f"viking_write error ({type(e).__name__}): {e}") async def viking_edit( ctx: RunContext[Any], @@ -669,7 +689,7 @@ async def viking_edit( await client.write(uri, modified, mode="replace") return ToolReturn(return_value=f"Replaced {count} occurrence(s) in {uri}.") except Exception as e: - return ToolReturn(return_value=f"viking_edit error: {e}") + return ToolReturn(return_value=f"viking_edit error ({type(e).__name__}): {e}") async def viking_mkdir( ctx: RunContext[Any], @@ -692,7 +712,7 @@ async def viking_mkdir( await client.mkdir(uri, description=description) return ToolReturn(return_value=f"Created directory {uri}.") except Exception as e: - return ToolReturn(return_value=f"viking_mkdir error: {e}") + return ToolReturn(return_value=f"viking_mkdir error ({type(e).__name__}): {e}") async def viking_add_resource( ctx: RunContext[Any], @@ -736,7 +756,9 @@ async def viking_add_resource( ) return ToolReturn(return_value=f"Added resource {path} to Viking. Result: {result}") except Exception as e: - return ToolReturn(return_value=f"viking_add_resource error: {e}") + return ToolReturn( + return_value=f"viking_add_resource error ({type(e).__name__}): {e}" + ) async def viking_forget( ctx: RunContext[Any], @@ -759,7 +781,7 @@ async def viking_forget( await client.rm(uri, recursive=recursive) return ToolReturn(return_value=f"Removed {uri}.") except Exception as e: - return ToolReturn(return_value=f"viking_forget error: {e}") + return ToolReturn(return_value=f"viking_forget error ({type(e).__name__}): {e}") write_tools: list[Callable[..., Awaitable[ToolReturn]]] = [ viking_write, @@ -809,7 +831,7 @@ async def viking_link( return_value=f"Linked {from_uri} -> {', '.join(targets)} (reason: {reason!r})." ) except Exception as e: - return ToolReturn(return_value=f"viking_link error: {e}") + return ToolReturn(return_value=f"viking_link error ({type(e).__name__}): {e}") async def viking_set_tags( ctx: RunContext[Any], @@ -834,11 +856,18 @@ async def viking_set_tags( await client.set_tags(uri, tags, mode="replace", recursive=recursive) return ToolReturn(return_value=f"Set {len(tags)} tag(s) on {uri}.") except Exception as e: - return ToolReturn(return_value=f"viking_set_tags error: {e}") + return ToolReturn(return_value=f"viking_set_tags error ({type(e).__name__}): {e}") graph_tools: list[Callable[..., Awaitable[ToolReturn]]] = [viking_set_tags] if cap.enable_link: graph_tools.append(viking_link) tools.extend(graph_tools) + if cap.enabled_tools is not None: + names = {getattr(fn, "__name__", "") for fn in tools} + allowed = {n for n in cap.enabled_tools if n in names} + tools = [fn for fn in tools if getattr(fn, "__name__", "") in allowed] + elif cap.disabled_tools is not None: + tools = [fn for fn in tools if getattr(fn, "__name__", "") not in set(cap.disabled_tools)] + return tools diff --git a/src/wolfharness_config/capabilities.py b/src/wolfharness_config/capabilities.py index 193890dc3..884098e1f 100644 --- a/src/wolfharness_config/capabilities.py +++ b/src/wolfharness_config/capabilities.py @@ -16,9 +16,9 @@ from __future__ import annotations -from typing import Annotated, Any, Literal +from typing import Annotated, Any, Literal, Self -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator KNOWN_CAPABILITY_TYPES: frozenset[str] = frozenset({ @@ -323,6 +323,30 @@ class VikingCapabilityConfig(BaseModel): """When True (default), profile injection runs only on the first turn of a session (message count <= 2). When False, injection runs on every before_model_request call where _profile_injected is False.""" + enabled_tools: list[str] | None = Field( + default=None, + examples=[["viking_ls", "viking_read", "viking_grep"]], + title="Enabled tools", + ) + """If set, only these tools will be available (whitelist). + Mutually exclusive with disabled_tools.""" + + disabled_tools: list[str] | None = Field( + default=None, + examples=[["viking_search", "viking_find"]], + title="Disabled tools", + ) + """Tools to exclude from this capability (blacklist). Mutually exclusive + with enabled_tools. For example, disable a slow knowledge-graph semantic + search backend while keeping the deterministic tools: + ``disabled_tools: ["viking_search", "viking_find"]``.""" + + @model_validator(mode="after") + def _validate_tool_filters(self) -> Self: + """Validate that enabled_tools and disabled_tools are mutually exclusive.""" + if self.enabled_tools is not None and self.disabled_tools is not None: + raise ValueError("Cannot specify both 'enabled_tools' and 'disabled_tools'") + return self # --------------------------------------------------------------------------- diff --git a/tests/capabilities/viking/test_viking.py b/tests/capabilities/viking/test_viking.py index 789aff1d8..01319c74a 100644 --- a/tests/capabilities/viking/test_viking.py +++ b/tests/capabilities/viking/test_viking.py @@ -1827,7 +1827,7 @@ async def test_search_error(self, viking_cap: VikingCapability, mock_client: Asy tools = build_tools(viking_cap) ctx = _make_ctx() result = await _get_tool(tools, "viking_search")(ctx, query="test") - assert "viking_search error: connection failed" in result.return_value + assert "viking_search error (RuntimeError): connection failed" in result.return_value @pytest.mark.asyncio async def test_find_error(self, viking_cap: VikingCapability, mock_client: AsyncMock) -> None: @@ -1835,7 +1835,7 @@ async def test_find_error(self, viking_cap: VikingCapability, mock_client: Async tools = build_tools(viking_cap) ctx = _make_ctx() result = await _get_tool(tools, "viking_find")(ctx, query="test") - assert "viking_find error: timeout" in result.return_value + assert "viking_find error (RuntimeError): timeout" in result.return_value @pytest.mark.asyncio async def test_recall_error(self, viking_cap: VikingCapability, mock_client: AsyncMock) -> None: @@ -1843,7 +1843,7 @@ async def test_recall_error(self, viking_cap: VikingCapability, mock_client: Asy tools = build_tools(viking_cap) ctx = _make_ctx() result = await _get_tool(tools, "viking_recall")(ctx, query="test") - assert "viking_recall error: server error" in result.return_value + assert "viking_recall error (RuntimeError): server error" in result.return_value @pytest.mark.asyncio async def test_grep_error(self, viking_cap: VikingCapability, mock_client: AsyncMock) -> None: @@ -1861,7 +1861,7 @@ async def test_glob_error(self, viking_cap: VikingCapability, mock_client: Async tools = build_tools(viking_cap) ctx = _make_ctx() result = await _get_tool(tools, "viking_glob")(ctx, pattern="**/*.md") - assert "viking_glob error: error" in result.return_value + assert "viking_glob error (RuntimeError): error" in result.return_value @pytest.mark.asyncio async def test_ls_error(self, viking_cap: VikingCapability, mock_client: AsyncMock) -> None: @@ -1869,7 +1869,7 @@ async def test_ls_error(self, viking_cap: VikingCapability, mock_client: AsyncMo tools = build_tools(viking_cap) ctx = _make_ctx() result = await _get_tool(tools, "viking_ls")(ctx, uri="viking://missing/") - assert "viking_ls error: not found" in result.return_value + assert "viking_ls error (RuntimeError): not found" in result.return_value @pytest.mark.asyncio async def test_read_error(self, viking_cap: VikingCapability, mock_client: AsyncMock) -> None: @@ -1877,7 +1877,7 @@ async def test_read_error(self, viking_cap: VikingCapability, mock_client: Async tools = build_tools(viking_cap) ctx = _make_ctx() result = await _get_tool(tools, "viking_read")(ctx, uris="viking://secret.md") - assert "viking_read error: permission denied" in result.return_value + assert "viking_read error (RuntimeError): permission denied" in result.return_value @pytest.mark.asyncio async def test_write_error(self, viking_cap: VikingCapability, mock_client: AsyncMock) -> None: @@ -1885,7 +1885,7 @@ async def test_write_error(self, viking_cap: VikingCapability, mock_client: Asyn tools = build_tools(viking_cap) ctx = _make_ctx() result = await _get_tool(tools, "viking_write")(ctx, uri="viking://doc.md", content="data") - assert "viking_write error: disk full" in result.return_value + assert "viking_write error (RuntimeError): disk full" in result.return_value @pytest.mark.asyncio async def test_edit_error(self, viking_cap: VikingCapability, mock_client: AsyncMock) -> None: @@ -1895,7 +1895,7 @@ async def test_edit_error(self, viking_cap: VikingCapability, mock_client: Async result = await _get_tool(tools, "viking_edit")( ctx, uri="viking://doc.md", old_string="a", new_string="b" ) - assert "viking_edit error: network error" in result.return_value + assert "viking_edit error (RuntimeError): network error" in result.return_value @pytest.mark.asyncio async def test_mkdir_error(self, viking_cap: VikingCapability, mock_client: AsyncMock) -> None: @@ -1903,7 +1903,7 @@ async def test_mkdir_error(self, viking_cap: VikingCapability, mock_client: Asyn tools = build_tools(viking_cap) ctx = _make_ctx() result = await _get_tool(tools, "viking_mkdir")(ctx, uri="viking://exists/") - assert "viking_mkdir error: exists" in result.return_value + assert "viking_mkdir error (RuntimeError): exists" in result.return_value @pytest.mark.asyncio async def test_add_resource_error( @@ -1913,7 +1913,7 @@ async def test_add_resource_error( tools = build_tools(viking_cap) ctx = _make_ctx() result = await _get_tool(tools, "viking_add_resource")(ctx, path="/bad/path") - assert "viking_add_resource error: invalid path" in result.return_value + assert "viking_add_resource error (RuntimeError): invalid path" in result.return_value @pytest.mark.asyncio async def test_forget_error(self, viking_cap: VikingCapability, mock_client: AsyncMock) -> None: @@ -1921,7 +1921,7 @@ async def test_forget_error(self, viking_cap: VikingCapability, mock_client: Asy tools = build_tools(viking_cap) ctx = _make_ctx() result = await _get_tool(tools, "viking_forget")(ctx, uri="viking://protected.md") - assert "viking_forget error: protected" in result.return_value + assert "viking_forget error (RuntimeError): protected" in result.return_value @pytest.mark.asyncio async def test_link_error(self, viking_cap: VikingCapability, mock_client: AsyncMock) -> None: @@ -1931,7 +1931,7 @@ async def test_link_error(self, viking_cap: VikingCapability, mock_client: Async result = await _get_tool(tools, "viking_link")( ctx, from_uri="viking://a.md", to_uris="viking://b.md" ) - assert "viking_link error: cycle detected" in result.return_value + assert "viking_link error (RuntimeError): cycle detected" in result.return_value @pytest.mark.asyncio async def test_set_tags_error( @@ -1941,7 +1941,7 @@ async def test_set_tags_error( tools = build_tools(viking_cap) ctx = _make_ctx() result = await _get_tool(tools, "viking_set_tags")(ctx, uri="viking://doc.md", tags=["bad"]) - assert "viking_set_tags error: invalid tag" in result.return_value + assert "viking_set_tags error (RuntimeError): invalid tag" in result.return_value @pytest.mark.asyncio async def test_ensure_client_lazy_init(self) -> None: @@ -2235,6 +2235,36 @@ def test_all_mode_16_tools_with_flags(self) -> None: tools = build_tools(cap) assert len(tools) == 16 + def test_disabled_tools_excludes_search_find(self) -> None: + """disabled_tools blacklist removes viking_search and viking_find.""" + cap = VikingCapability( + mode="retrieve", + disabled_tools=["viking_search", "viking_find"], + ) + cap._client = AsyncMock() + tools = build_tools(cap) + names = {t.__name__ for t in tools} + assert "viking_search" not in names + assert "viking_find" not in names + assert "viking_read" in names + assert "viking_grep" in names + + def test_enabled_tools_whitelist(self) -> None: + """enabled_tools whitelist keeps only the listed tools.""" + cap = VikingCapability(mode="retrieve", enabled_tools=["viking_ls", "viking_read"]) + cap._client = AsyncMock() + tools = build_tools(cap) + names = {t.__name__ for t in tools} + assert names == {"viking_ls", "viking_read"} + + def test_enabled_tools_unknown_names_ignored(self) -> None: + """enabled_tools entries that don't match any tool are ignored.""" + cap = VikingCapability(mode="retrieve", enabled_tools=["viking_ls", "nope_tool"]) + cap._client = AsyncMock() + tools = build_tools(cap) + names = {t.__name__ for t in tools} + assert names == {"viking_ls"} + def test_get_toolset_retrieve(self) -> None: """get_toolset() returns a FunctionToolset with 7 tools for retrieve mode (default).""" from pydantic_ai.toolsets import FunctionToolset @@ -2358,6 +2388,31 @@ def test_instructions_consistent_across_modes(self) -> None: cap_retrieve = VikingCapability(mode="retrieve") assert cap_all.get_instructions() == cap_retrieve.get_instructions() + def test_instructions_no_prefix_block_when_unrestricted(self) -> None: + """get_instructions() omits the allowed-prefix block when unrestricted.""" + cap = VikingCapability(mode="all") + instructions = cap.get_instructions() + assert instructions is not None + assert "Allowed URI Prefixes" not in instructions + + def test_instructions_include_prefix_block_when_restricted(self) -> None: + """get_instructions() lists the allowed prefixes. + + So the model can pass a target_uri and skip discovery probing. + """ + cap = VikingCapability( + mode="all", + allowed_uri_prefixes=[ + "viking://resources/wiki/", + "viking://resources/raw/", + ], + ) + instructions = cap.get_instructions() + assert instructions is not None + assert "Allowed URI Prefixes" in instructions + assert "viking://resources/wiki/" in instructions + assert "viking://resources/raw/" in instructions + def test_on_change_returns_none(self) -> None: """on_change() returns None.""" cap = VikingCapability(mode="all") @@ -2963,7 +3018,7 @@ async def test_viking_read_abstract_error( ctx = _make_ctx() result = await read_tool(ctx, uris="viking://doc.md", level="abstract") - assert "viking_read error: not available" in result.return_value + assert "viking_read error (RuntimeError): not available" in result.return_value class TestTieredLoadingReadResource: @@ -4515,7 +4570,11 @@ async def test_viking_write_blocks_outside_prefix(self, mock_client: AsyncMock) @pytest.mark.asyncio async def test_viking_search_defaults_to_first_prefix(self, mock_client: AsyncMock) -> None: - """viking_search without target_uri scopes to the first allowed prefix.""" + """viking_search passes the single allowed prefix as a one-element list. + + A list is the SDK's multi-prefix scoping contract, so a single + prefix is passed as a one-element list rather than a bare string. + """ cap = VikingCapability(mode="retrieve", allowed_uri_prefixes=["viking://resources/wiki/"]) cap._client = mock_client tools = build_tools(cap) @@ -4525,7 +4584,52 @@ async def test_viking_search_defaults_to_first_prefix(self, mock_client: AsyncMo await search_tool(ctx, query="hydraulic") kwargs = mock_client.search.call_args.kwargs - assert kwargs["target_uri"] == "viking://resources/wiki/" + assert kwargs["target_uri"] == ["viking://resources/wiki/"] + + @pytest.mark.asyncio + async def test_viking_search_multi_prefix_defaults_to_all(self, mock_client: AsyncMock) -> None: + """viking_search without target_uri passes ALL allowed prefixes to the SDK. + + The SDK's target_uri accepts a list and the server searches each + prefix — the old behavior of silently using only the first prefix + dropped results from the other allowed trees. + """ + cap = VikingCapability( + mode="retrieve", + allowed_uri_prefixes=["viking://resources/wiki/", "viking://resources/raw/"], + ) + cap._client = mock_client + tools = build_tools(cap) + search_tool = _get_tool(tools, "viking_search") + + ctx = _make_ctx() + await search_tool(ctx, query="hydraulic") + + kwargs = mock_client.search.call_args.kwargs + assert kwargs["target_uri"] == [ + "viking://resources/wiki/", + "viking://resources/raw/", + ] + + @pytest.mark.asyncio + async def test_viking_find_multi_prefix_defaults_to_all(self, mock_client: AsyncMock) -> None: + """viking_find without target_uri passes ALL allowed prefixes to the SDK.""" + cap = VikingCapability( + mode="retrieve", + allowed_uri_prefixes=["viking://resources/wiki/", "viking://resources/raw/"], + ) + cap._client = mock_client + tools = build_tools(cap) + find_tool = _get_tool(tools, "viking_find") + + ctx = _make_ctx() + await find_tool(ctx, query="hydraulic") + + kwargs = mock_client.find.call_args.kwargs + assert kwargs["target_uri"] == [ + "viking://resources/wiki/", + "viking://resources/raw/", + ] @pytest.mark.asyncio async def test_viking_search_blocks_outside_target(self, mock_client: AsyncMock) -> None: @@ -5147,7 +5251,7 @@ async def test_viking_expand_tool_error(self, mock_client: AsyncMock) -> None: ctx = _make_ctx() result = await expand_tool(ctx, uri="viking://missing.md") - assert "viking_expand error: not found" in result.return_value + assert "viking_expand error (RuntimeError): not found" in result.return_value @pytest.mark.asyncio async def test_viking_expand_tool_empty_content(self, mock_client: AsyncMock) -> None: diff --git a/tests/capabilities/viking/test_viking_integration.py b/tests/capabilities/viking/test_viking_integration.py index bbf8f4092..643afb6e7 100644 --- a/tests/capabilities/viking/test_viking_integration.py +++ b/tests/capabilities/viking/test_viking_integration.py @@ -139,6 +139,49 @@ def test_yaml_config_loading_viking_retrieve() -> None: assert cfg.mode == "retrieve" +def test_yaml_config_loading_viking_disabled_tools() -> None: + """YAML config with disabled_tools excludes the listed tools.""" + from wolfharness import AgentsManifest + + yaml_str = """ +agents: + test_agent: + type: native + model: test + capabilities: + - type: viking + mode: retrieve + disabled_tools: ["viking_search", "viking_find"] +""" + d = yamling.load_yaml(yaml_str, verify_type=dict) + manifest = AgentsManifest.model_validate(d) + cfg = manifest.agents["test_agent"].capabilities[0] + assert isinstance(cfg, VikingCapabilityConfig) + assert cfg.disabled_tools == ["viking_search", "viking_find"] + cap = build_capability(cfg) + from wolfharness.capabilities.viking.tools import build_tools + + names = {t.__name__ for t in build_tools(cap)} + assert "viking_search" not in names + assert "viking_find" not in names + assert "viking_read" in names + + +def test_viking_config_enabled_disabled_mutually_exclusive() -> None: + """enabled_tools and disabled_tools together raise a validation error.""" + import pytest + + from wolfharness_config.capabilities import VikingCapabilityConfig + + with pytest.raises(ValueError, match="Cannot specify both"): + VikingCapabilityConfig( + type="viking", + mode="retrieve", + enabled_tools=["viking_ls"], + disabled_tools=["viking_search"], + ) + + def test_yaml_config_loading_viking_with_fields() -> None: """YAML config with all fields populated parses correctly.""" from wolfharness import AgentsManifest @@ -363,7 +406,7 @@ async def test_network_error_graceful() -> None: ctx.deps.session_id = "test" result = await search_tool(ctx, query="test") - assert "viking_search error: network down" in result.return_value + assert "viking_search error (ConnectionError): network down" in result.return_value assert isinstance(result, ToolReturn) @@ -385,7 +428,7 @@ async def test_invalid_uri_graceful() -> None: ctx.deps.session_id = "test" result = await read_tool(ctx, uris="not-a-valid-uri") - assert "viking_read error: invalid URI format" in result.return_value + assert "viking_read error (ValueError): invalid URI format" in result.return_value @pytest.mark.asyncio @@ -406,7 +449,7 @@ async def test_permission_error_graceful() -> None: ctx.deps.session_id = "test" result = await write_tool(ctx, uri="viking://protected/doc.md", content="data") - assert "viking_write error: access denied" in result.return_value + assert "viking_write error (PermissionError): access denied" in result.return_value @pytest.mark.asyncio @@ -427,7 +470,33 @@ async def test_timeout_error_graceful() -> None: ctx.deps.session_id = "test" result = await search_tool(ctx, query="slow query") - assert "viking_search error: request timed out" in result.return_value + assert "viking_search error (TimeoutError): request timed out" in result.return_value + + +@pytest.mark.asyncio +async def test_empty_message_error_graceful() -> None: + """Exceptions with empty messages still surface the exception type. + + e.g. httpcore raises ``ReadTimeout('')`` — an empty-message exception + that would otherwise render as ``viking_search error:`` with no context. + """ + cap = VikingCapability(mode="all") + mock_client = AsyncMock() + mock_client.search = AsyncMock(side_effect=RuntimeError("")) + cap._client = mock_client + + tools = build_tools(cap) + search_tool = next(t for t in tools if t.__name__ == "viking_search") + + from unittest.mock import MagicMock + + ctx = MagicMock() + ctx.deps = MagicMock() + ctx.deps.session_id = "test" + + result = await search_tool(ctx, query="slow query") + assert "viking_search error (RuntimeError): " in result.return_value + assert result.return_value.strip().endswith("(RuntimeError):") @pytest.mark.asyncio @@ -448,7 +517,7 @@ async def test_generic_exception_graceful() -> None: ctx.deps.session_id = "test" result = await ls_tool(ctx, uri="viking://broken/") - assert "viking_ls error: unexpected error" in result.return_value + assert "viking_ls error (Exception): unexpected error" in result.return_value @pytest.mark.asyncio @@ -748,11 +817,12 @@ async def test_l2_allowed_uri_prefixes_allow_read_in_scope() -> None: @pytest.mark.asyncio async def test_l2_allowed_uri_prefixes_search_scoped_when_no_target() -> None: - """L2: viking_search defaults target_uri to the first allowed prefix. + """L2: viking_search scopes to the allowed prefix when no target is given. Given: allowed_uri_prefixes configured, no target_uri passed. When: viking_search is called. - Then: the SDK search receives target_uri equal to the first allowed prefix. + Then: the SDK search receives target_uri as a one-element list of the + allowed prefixes (list is the SDK's multi-prefix contract). """ client = _make_mock_client() cfg = VikingCapabilityConfig( @@ -766,7 +836,7 @@ async def test_l2_allowed_uri_prefixes_search_scoped_when_no_target() -> None: ctx = _make_run_context() await search_tool(ctx, query="hydraulic") - assert client.search.call_args.kwargs["target_uri"] == "viking://resources/wiki/" + assert client.search.call_args.kwargs["target_uri"] == ["viking://resources/wiki/"] @pytest.mark.asyncio