Skip to content

Commit df9ac26

Browse files
committed
fix(skills): support aliased registry IDs and address review feedback
- Add _registry_skill_aliases map to resolve pinned skills by registry ID when frontmatter name differs, avoiding re-downloads in load_skill. - Filter out pinned skills in search_skills matching either alias or frontmatter name. - Remove redundant prefetch calls from ListSkillsTool and SearchSkillsTool. - Expand registry_skills docstring explaining cache lifecycle and naming. - Demonstrate registry_skills in GCP skill registry sample agent. - Add unit tests for LAZY list_skills, no-refetch caching, frontmatter name mismatches, and aliased search filtering.
1 parent 4c469f8 commit df9ac26

3 files changed

Lines changed: 156 additions & 7 deletions

File tree

contributing/samples/integrations/gcp_skill_registry_agent/agent.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,11 @@
2424
)
2525

2626
# Initialize SkillToolset with registry
27-
skill_toolset = SkillToolset(skills=[], registry=registry)
27+
skill_toolset = SkillToolset(
28+
skills=[],
29+
registry=registry,
30+
registry_skills=["your-pinned-skill-id"],
31+
)
2832

2933
root_agent = Agent(
3034
model="gemini-2.5-flash",

src/google/adk/tools/skill_toolset.py

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -273,7 +273,6 @@ def _get_declaration(self) -> types.FunctionDeclaration | None:
273273
async def run_async(
274274
self, *, args: dict[str, Any], tool_context: ToolContext
275275
) -> Any:
276-
await self._toolset.prefetch()
277276
skills = self._toolset._list_skills()
278277
return prompt.format_skills_as_xml(skills)
279278

@@ -316,7 +315,6 @@ def _get_declaration(self) -> types.FunctionDeclaration | None:
316315
async def run_async(
317316
self, *, args: dict[str, Any], tool_context: ToolContext
318317
) -> Any:
319-
await self._toolset.prefetch()
320318
query = args.get("query")
321319
if not query:
322320
return {
@@ -327,7 +325,10 @@ async def run_async(
327325
results = await self._toolset._registry.search_skills(query=query)
328326
formatted_results = []
329327
for r in results:
330-
if r.name in self._toolset._skills:
328+
if (
329+
r.name in self._toolset._skills
330+
or r.name in self._toolset._registry_skill_aliases
331+
):
331332
logger.warning(
332333
"Skill naming conflict: skill '%s' already exists locally."
333334
" Registry skill is filtered.",
@@ -1379,8 +1380,12 @@ def __init__(
13791380
Args:
13801381
skills: List of skills to register.
13811382
registry: Optional skill registry for dynamic loading.
1382-
registry_skills: Optional list of skill names in the registry to pin and
1383-
fetch into the local catalog.
1383+
registry_skills: Optional list of skill names in the registry to pin.
1384+
Pinned skills are fetched once per process on first use, then appear in
1385+
`list_skills` / the EAGER catalog and are served locally without further
1386+
registry calls; unpinned registry skills are still reachable via
1387+
`search_skills`. The skill is stored under the name in the archive's
1388+
frontmatter, which may differ from the registry resource name.
13841389
code_executor: Optional code executor for script execution.
13851390
environment: Optional environment for executing scripts.
13861391
skills_folder: Optional absolute path where skills are stored in the
@@ -1417,6 +1422,7 @@ def __init__(
14171422
)
14181423
self._registry_skills_loaded = False
14191424
self._fetched_registry_skills: set[str] = set()
1425+
self._registry_skill_aliases: dict[str, str] = {}
14201426
self._registry_skills_lock: asyncio.Lock | None = None
14211427
self._code_executor = code_executor
14221428
self._env = environment
@@ -1536,6 +1542,7 @@ async def prefetch(self) -> None:
15361542
)
15371543
else:
15381544
self._skills[skill.name] = skill
1545+
self._registry_skill_aliases[name] = skill.name
15391546
self._fetched_registry_skills.add(name)
15401547

15411548
if len(self._fetched_registry_skills) == len(self._registry_skills):
@@ -1649,6 +1656,10 @@ async def _get_or_fetch_skill(
16491656
if skill:
16501657
return skill
16511658

1659+
if aliased_name := self._registry_skill_aliases.get(skill_name):
1660+
if skill := self._get_skill(aliased_name):
1661+
return skill
1662+
16521663
if not self._registry:
16531664
return None
16541665

tests/unittests/tools/test_skill_toolset.py

Lines changed: 135 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3821,7 +3821,8 @@ async def test_skill_toolset_search_skills_filters_pinned_registry_skill(
38213821
toolset = skill_toolset.SkillToolset(
38223822
registry=mock_registry, registry_skills=["skill1"]
38233823
)
3824-
tool = skill_toolset.SearchSkillsTool(toolset)
3824+
tools = await toolset.get_tools()
3825+
tool = next(t for t in tools if isinstance(t, skill_toolset.SearchSkillsTool))
38253826

38263827
result = await tool.run_async(
38273828
args={"query": "test"}, tool_context=tool_context_instance
@@ -3859,3 +3860,136 @@ async def test_skill_toolset_prefetch_concurrent_calls(
38593860
)
38603861

38613862
mock_registry.get_skill.assert_called_once_with(name="skill1")
3863+
3864+
3865+
@pytest.mark.asyncio
3866+
async def test_skill_toolset_list_skills_tool_includes_pinned_registry_skill(
3867+
mock_registry, mock_skill1, tool_context_instance
3868+
):
3869+
"""LAZY mode: the list_skills tool output must contain pinned skills."""
3870+
mock_registry.get_skill.return_value = mock_skill1
3871+
toolset = skill_toolset.SkillToolset(
3872+
registry=mock_registry, registry_skills=["skill1"]
3873+
)
3874+
tools = await toolset.get_tools()
3875+
tool = next(t for t in tools if isinstance(t, skill_toolset.ListSkillsTool))
3876+
3877+
result = await tool.run_async(args={}, tool_context=tool_context_instance)
3878+
3879+
assert "<available_skills>" in result
3880+
assert "skill1" in result
3881+
mock_registry.get_skill.assert_called_once_with(name="skill1")
3882+
3883+
3884+
@pytest.mark.asyncio
3885+
async def test_skill_toolset_load_skill_does_not_refetch_pinned_registry_skill(
3886+
mock_registry, mock_skill1, tool_context_instance
3887+
):
3888+
"""After prefetch, load_skill for a pinned name is served locally."""
3889+
mock_registry.get_skill.return_value = mock_skill1
3890+
toolset = skill_toolset.SkillToolset(
3891+
registry=mock_registry, registry_skills=["skill1"]
3892+
)
3893+
tool_context_instance.state.get.return_value = None
3894+
3895+
await toolset.prefetch()
3896+
assert mock_registry.get_skill.call_count == 1
3897+
3898+
tool = skill_toolset.LoadSkillTool(toolset)
3899+
# Two invocations: the per-invocation cache must not be what serves this.
3900+
for invocation_id in ("inv-1", "inv-2"):
3901+
tool_context_instance.invocation_id = invocation_id
3902+
result = await tool.run_async(
3903+
args={"skill_name": "skill1"}, tool_context=tool_context_instance
3904+
)
3905+
assert result["skill_name"] == "skill1"
3906+
3907+
assert mock_registry.get_skill.call_count == 1
3908+
3909+
3910+
@pytest.mark.asyncio
3911+
async def test_skill_toolset_prefetch_stores_skill_under_frontmatter_name(
3912+
mock_registry, tool_context_instance
3913+
):
3914+
"""A pinned name is marked fetched even when the archive's frontmatter name
3915+
differs; the skill is stored under the frontmatter name, and can be resolved
3916+
by either name without refetching."""
3917+
fetched = mock.create_autospec(models.Skill, instance=True)
3918+
fetched.name = "bar"
3919+
fetched.instructions = "Instructions for bar"
3920+
fetched.frontmatter = mock.create_autospec(models.Frontmatter, instance=True)
3921+
fetched.frontmatter.metadata = {}
3922+
mock_registry.get_skill.return_value = fetched
3923+
toolset = skill_toolset.SkillToolset(
3924+
registry=mock_registry, registry_skills=["foo"]
3925+
)
3926+
3927+
await toolset.prefetch()
3928+
3929+
mock_registry.get_skill.assert_called_once_with(name="foo")
3930+
assert toolset._get_skill("bar") is fetched
3931+
assert toolset._get_skill("foo") is None
3932+
assert toolset._registry_skills_loaded is True
3933+
3934+
# Resolving via _get_or_fetch_skill with the requested registry id "foo"
3935+
# uses the alias map and does not call registry again.
3936+
resolved = await toolset._get_or_fetch_skill("foo")
3937+
assert resolved is fetched
3938+
assert mock_registry.get_skill.call_count == 1
3939+
3940+
# Also via LoadSkillTool using the registry id "foo"
3941+
tool_context_instance.state.get.return_value = None
3942+
load_tool = skill_toolset.LoadSkillTool(toolset)
3943+
res = await load_tool.run_async(
3944+
args={"skill_name": "foo"}, tool_context=tool_context_instance
3945+
)
3946+
assert res["skill_name"] == "foo"
3947+
assert res["instructions"] == "Instructions for bar"
3948+
assert mock_registry.get_skill.call_count == 1
3949+
3950+
# No refetch on the next prefetch.
3951+
await toolset.prefetch()
3952+
assert mock_registry.get_skill.call_count == 1
3953+
3954+
3955+
@pytest.mark.asyncio
3956+
async def test_skill_toolset_search_skills_filters_aliased_registry_skill(
3957+
mock_registry, tool_context_instance
3958+
):
3959+
"""SearchSkillsTool filters out pinned skills by both registry ID and frontmatter name."""
3960+
fetched = mock.create_autospec(models.Skill, instance=True)
3961+
fetched.name = "bar"
3962+
mock_registry.get_skill.return_value = fetched
3963+
3964+
mock_frontmatter_reg = mock.create_autospec(models.Frontmatter, instance=True)
3965+
mock_frontmatter_reg.name = "foo"
3966+
mock_frontmatter_reg.model_dump.return_value = {"name": "foo"}
3967+
3968+
mock_frontmatter_fm = mock.create_autospec(models.Frontmatter, instance=True)
3969+
mock_frontmatter_fm.name = "bar"
3970+
mock_frontmatter_fm.model_dump.return_value = {"name": "bar"}
3971+
3972+
mock_frontmatter_other = mock.create_autospec(
3973+
models.Frontmatter, instance=True
3974+
)
3975+
mock_frontmatter_other.name = "other"
3976+
mock_frontmatter_other.model_dump.return_value = {"name": "other"}
3977+
3978+
mock_registry.search_skills.return_value = [
3979+
mock_frontmatter_reg,
3980+
mock_frontmatter_fm,
3981+
mock_frontmatter_other,
3982+
]
3983+
3984+
toolset = skill_toolset.SkillToolset(
3985+
registry=mock_registry, registry_skills=["foo"]
3986+
)
3987+
tools = await toolset.get_tools()
3988+
tool = next(t for t in tools if isinstance(t, skill_toolset.SearchSkillsTool))
3989+
3990+
result = await tool.run_async(
3991+
args={"query": "test"}, tool_context=tool_context_instance
3992+
)
3993+
3994+
# Both "foo" (aliased registry id) and "bar" (stored frontmatter name) filtered
3995+
assert result == [{"name": "other"}]

0 commit comments

Comments
 (0)