Skip to content

Commit 1fa2a62

Browse files
committed
feat(tools): add rolling window for activated skills in SkillToolset
Add max_activated_skills parameter to SkillToolset that enforces a rolling window on the number of concurrently activated skills per agent. When the limit is exceeded, the oldest activated skills (and their associated additional tools) are evicted from session state. This addresses unbounded tool context growth during long-running sessions where many skills get loaded but only recent ones are relevant. Key changes: - Add max_activated_skills param to SkillToolset.__init__ - Enforce LRU-style eviction in LoadSkillTool.run_async: re-loading an already active skill promotes it to most-recent position - Always persist state (even for re-activations) to maintain ordering - Default is None (no limit) for backward compatibility
1 parent baf7efb commit 1fa2a62

2 files changed

Lines changed: 356 additions & 1 deletion

File tree

src/google/adk/tools/skill_toolset.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -241,7 +241,17 @@ async def run_async(
241241
activated_skills = list(tool_context.state.get(state_key) or [])
242242
if skill_name not in activated_skills:
243243
activated_skills.append(skill_name)
244-
tool_context.state[state_key] = activated_skills
244+
else:
245+
# Move to the end (most recently used) for rolling window.
246+
activated_skills.remove(skill_name)
247+
activated_skills.append(skill_name)
248+
249+
# Enforce rolling window: evict oldest skills if limit is exceeded.
250+
max_skills = self._toolset._max_activated_skills
251+
if max_skills is not None and len(activated_skills) > max_skills:
252+
activated_skills = activated_skills[-max_skills:]
253+
254+
tool_context.state[state_key] = activated_skills
245255

246256
return {
247257
"skill_name": skill_name,
@@ -886,6 +896,7 @@ def __init__(
886896
code_executor: BaseCodeExecutor | None = None,
887897
script_timeout: int = _DEFAULT_SCRIPT_TIMEOUT,
888898
additional_tools: list[ToolUnion] | None = None,
899+
max_activated_skills: int | None = None,
889900
):
890901
"""Initializes the SkillToolset.
891902
@@ -898,6 +909,10 @@ def __init__(
898909
scripts executed via exec().
899910
additional_tools: Optional list of `BaseTool` or `BaseToolset` instances
900911
to be made available to the agent when certain skills are activated.
912+
max_activated_skills: Optional maximum number of skills to keep activated
913+
simultaneously. When set, only the most recently activated skills are
914+
retained (rolling window). Older skills and their associated tools are
915+
evicted from the session state. Defaults to None (no limit).
901916
"""
902917
super().__init__()
903918

@@ -914,6 +929,7 @@ def __init__(
914929
self._registry = registry
915930
self._code_executor = code_executor
916931
self._script_timeout = script_timeout
932+
self._max_activated_skills = max_activated_skills
917933
# Needed for mid-turn reloading of skill tools.
918934
self._use_invocation_cache = False
919935
# Cache fetched remote skill definitions per turn to reduce requests to registry

tests/unittests/tools/test_skill_toolset.py

Lines changed: 339 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1925,3 +1925,342 @@ async def test_close_cancels_futures_and_clears_cache():
19251925
assert fut1.cancelled()
19261926
assert not fut2.cancelled() # Done futures shouldn't/can't be cancelled
19271927
assert not toolset._fetched_skill_cache
1928+
1929+
1930+
# --- Tests for max_activated_skills (rolling window) ---
1931+
1932+
1933+
@pytest.mark.asyncio
1934+
async def test_max_activated_skills_evicts_oldest(
1935+
mock_skill1, mock_skill2, tool_context_instance
1936+
):
1937+
"""When max_activated_skills is exceeded, the oldest skill is evicted."""
1938+
toolset = skill_toolset.SkillToolset(
1939+
[mock_skill1, mock_skill2], max_activated_skills=1
1940+
)
1941+
tool = skill_toolset.LoadSkillTool(toolset)
1942+
1943+
state_key = "_adk_activated_skill_test_agent"
1944+
1945+
# Load skill1
1946+
tool_context_instance.state.get.return_value = None
1947+
await tool.run_async(
1948+
args={"skill_name": "skill1"}, tool_context=tool_context_instance
1949+
)
1950+
tool_context_instance.state.__setitem__.assert_called_with(
1951+
state_key, ["skill1"]
1952+
)
1953+
1954+
# Load skill2 — skill1 should be evicted
1955+
tool_context_instance.state.get.return_value = ["skill1"]
1956+
await tool.run_async(
1957+
args={"skill_name": "skill2"}, tool_context=tool_context_instance
1958+
)
1959+
tool_context_instance.state.__setitem__.assert_called_with(
1960+
state_key, ["skill2"]
1961+
)
1962+
1963+
1964+
@pytest.mark.asyncio
1965+
async def test_max_activated_skills_keeps_recent_n(
1966+
tool_context_instance,
1967+
):
1968+
"""Rolling window keeps the N most recently activated skills."""
1969+
skills = []
1970+
for i in range(5):
1971+
skill = mock.create_autospec(models.Skill, instance=True)
1972+
skill.name = f"skill_{i}"
1973+
skill.instructions = f"instructions for skill_{i}"
1974+
frontmatter = mock.create_autospec(models.Frontmatter, instance=True)
1975+
frontmatter.name = f"skill_{i}"
1976+
frontmatter.metadata = {}
1977+
frontmatter.model_dump.return_value = {"name": f"skill_{i}"}
1978+
skill.frontmatter = frontmatter
1979+
skills.append(skill)
1980+
1981+
toolset = skill_toolset.SkillToolset(skills, max_activated_skills=3)
1982+
tool = skill_toolset.LoadSkillTool(toolset)
1983+
1984+
state_key = "_adk_activated_skill_test_agent"
1985+
1986+
# Sequentially load skills 0 through 4
1987+
current_state = []
1988+
for i in range(5):
1989+
tool_context_instance.state.get.return_value = list(current_state)
1990+
await tool.run_async(
1991+
args={"skill_name": f"skill_{i}"}, tool_context=tool_context_instance
1992+
)
1993+
# Capture what was set
1994+
call_args = tool_context_instance.state.__setitem__.call_args
1995+
current_state = call_args[0][1]
1996+
1997+
# After loading 5 skills with max=3, only the last 3 should remain
1998+
assert current_state == ["skill_2", "skill_3", "skill_4"]
1999+
2000+
2001+
@pytest.mark.asyncio
2002+
async def test_max_activated_skills_reloading_moves_to_end(
2003+
mock_skill1, mock_skill2, tool_context_instance
2004+
):
2005+
"""Re-loading an already active skill moves it to the end (most recent)."""
2006+
# Create a third skill
2007+
mock_skill3 = mock.create_autospec(models.Skill, instance=True)
2008+
mock_skill3.name = "skill3"
2009+
mock_skill3.instructions = "instructions for skill3"
2010+
frontmatter3 = mock.create_autospec(models.Frontmatter, instance=True)
2011+
frontmatter3.name = "skill3"
2012+
frontmatter3.metadata = {}
2013+
frontmatter3.model_dump.return_value = {"name": "skill3"}
2014+
mock_skill3.frontmatter = frontmatter3
2015+
2016+
toolset = skill_toolset.SkillToolset(
2017+
[mock_skill1, mock_skill2, mock_skill3], max_activated_skills=3
2018+
)
2019+
tool = skill_toolset.LoadSkillTool(toolset)
2020+
2021+
state_key = "_adk_activated_skill_test_agent"
2022+
2023+
# Start with all 3 loaded
2024+
tool_context_instance.state.get.return_value = [
2025+
"skill1",
2026+
"skill2",
2027+
"skill3",
2028+
]
2029+
2030+
# Reload skill1 — should move to end
2031+
await tool.run_async(
2032+
args={"skill_name": "skill1"}, tool_context=tool_context_instance
2033+
)
2034+
call_args = tool_context_instance.state.__setitem__.call_args
2035+
assert call_args[0][1] == ["skill2", "skill3", "skill1"]
2036+
2037+
2038+
@pytest.mark.asyncio
2039+
async def test_no_max_activated_skills_unlimited(
2040+
tool_context_instance,
2041+
):
2042+
"""Without max_activated_skills, all skills are retained (no eviction)."""
2043+
skills = []
2044+
for i in range(10):
2045+
skill = mock.create_autospec(models.Skill, instance=True)
2046+
skill.name = f"skill_{i}"
2047+
skill.instructions = f"instructions for skill_{i}"
2048+
frontmatter = mock.create_autospec(models.Frontmatter, instance=True)
2049+
frontmatter.name = f"skill_{i}"
2050+
frontmatter.metadata = {}
2051+
frontmatter.model_dump.return_value = {"name": f"skill_{i}"}
2052+
skill.frontmatter = frontmatter
2053+
skills.append(skill)
2054+
2055+
toolset = skill_toolset.SkillToolset(skills) # No max_activated_skills
2056+
tool = skill_toolset.LoadSkillTool(toolset)
2057+
2058+
state_key = "_adk_activated_skill_test_agent"
2059+
2060+
current_state = []
2061+
for i in range(10):
2062+
tool_context_instance.state.get.return_value = list(current_state)
2063+
await tool.run_async(
2064+
args={"skill_name": f"skill_{i}"}, tool_context=tool_context_instance
2065+
)
2066+
call_args = tool_context_instance.state.__setitem__.call_args
2067+
current_state = call_args[0][1]
2068+
2069+
# All 10 skills should be retained
2070+
assert len(current_state) == 10
2071+
assert current_state == [f"skill_{i}" for i in range(10)]
2072+
2073+
2074+
@pytest.mark.asyncio
2075+
async def test_evicted_skill_tools_not_resolved():
2076+
"""After a skill is evicted, its additional tools are no longer resolved."""
2077+
# Create two skills, each declaring different additional tools
2078+
skill_a = mock.create_autospec(models.Skill, instance=True)
2079+
skill_a.name = "skill_a"
2080+
skill_a.instructions = "instructions A"
2081+
frontmatter_a = mock.create_autospec(models.Frontmatter, instance=True)
2082+
frontmatter_a.name = "skill_a"
2083+
frontmatter_a.metadata = {"adk_additional_tools": ["tool_alpha"]}
2084+
frontmatter_a.model_dump.return_value = {"name": "skill_a"}
2085+
skill_a.frontmatter = frontmatter_a
2086+
2087+
skill_b = mock.create_autospec(models.Skill, instance=True)
2088+
skill_b.name = "skill_b"
2089+
skill_b.instructions = "instructions B"
2090+
frontmatter_b = mock.create_autospec(models.Frontmatter, instance=True)
2091+
frontmatter_b.name = "skill_b"
2092+
frontmatter_b.metadata = {"adk_additional_tools": ["tool_beta"]}
2093+
frontmatter_b.model_dump.return_value = {"name": "skill_b"}
2094+
skill_b.frontmatter = frontmatter_b
2095+
2096+
# Create mock additional tools
2097+
tool_alpha = mock.create_autospec(skill_toolset.BaseTool, instance=True)
2098+
tool_alpha.name = "tool_alpha"
2099+
tool_beta = mock.create_autospec(skill_toolset.BaseTool, instance=True)
2100+
tool_beta.name = "tool_beta"
2101+
2102+
toolset = skill_toolset.SkillToolset(
2103+
[skill_a, skill_b],
2104+
additional_tools=[tool_alpha, tool_beta],
2105+
max_activated_skills=1,
2106+
)
2107+
2108+
# Simulate: skill_a was activated, then skill_b was activated (evicting A)
2109+
readonly_context = mock.create_autospec(ReadonlyContext, instance=True)
2110+
readonly_context.agent_name = "test_agent"
2111+
readonly_context.invocation_id = "inv-1"
2112+
# After eviction, only skill_b is in state
2113+
readonly_context.state.get.return_value = ["skill_b"]
2114+
2115+
tools = await toolset.get_tools(readonly_context=readonly_context)
2116+
tool_names = {t.name for t in tools}
2117+
2118+
# tool_beta should be present (skill_b is active)
2119+
assert "tool_beta" in tool_names
2120+
# tool_alpha should NOT be present (skill_a was evicted)
2121+
assert "tool_alpha" not in tool_names
2122+
2123+
2124+
@pytest.mark.asyncio
2125+
async def test_evicted_tool_call_raises_value_error():
2126+
"""Simulates the LLM calling a tool from an evicted skill.
2127+
2128+
This mirrors what happens in the real flow: after skill eviction, the
2129+
tool is removed from tools_dict, so _get_tool raises ValueError.
2130+
This test validates that the tools_dict correctly excludes evicted tools.
2131+
"""
2132+
# Skill with additional tool
2133+
skill_a = mock.create_autospec(models.Skill, instance=True)
2134+
skill_a.name = "skill_a"
2135+
skill_a.instructions = "instructions A"
2136+
frontmatter_a = mock.create_autospec(models.Frontmatter, instance=True)
2137+
frontmatter_a.name = "skill_a"
2138+
frontmatter_a.metadata = {"adk_additional_tools": ["expensive_api_tool"]}
2139+
frontmatter_a.model_dump.return_value = {"name": "skill_a"}
2140+
skill_a.frontmatter = frontmatter_a
2141+
2142+
skill_b = mock.create_autospec(models.Skill, instance=True)
2143+
skill_b.name = "skill_b"
2144+
skill_b.instructions = "instructions B"
2145+
frontmatter_b = mock.create_autospec(models.Frontmatter, instance=True)
2146+
frontmatter_b.name = "skill_b"
2147+
frontmatter_b.metadata = {}
2148+
frontmatter_b.model_dump.return_value = {"name": "skill_b"}
2149+
skill_b.frontmatter = frontmatter_b
2150+
2151+
expensive_tool = mock.create_autospec(skill_toolset.BaseTool, instance=True)
2152+
expensive_tool.name = "expensive_api_tool"
2153+
2154+
toolset = skill_toolset.SkillToolset(
2155+
[skill_a, skill_b],
2156+
additional_tools=[expensive_tool],
2157+
max_activated_skills=1,
2158+
)
2159+
2160+
# After eviction: only skill_b remains
2161+
readonly_context = mock.create_autospec(ReadonlyContext, instance=True)
2162+
readonly_context.agent_name = "test_agent"
2163+
readonly_context.invocation_id = "inv-2"
2164+
readonly_context.state.get.return_value = ["skill_b"]
2165+
2166+
tools = await toolset.get_tools(readonly_context=readonly_context)
2167+
tools_dict = {t.name: t for t in tools}
2168+
2169+
# The LLM might still try to call "expensive_api_tool" from history
2170+
assert "expensive_api_tool" not in tools_dict
2171+
2172+
# Simulate what _get_tool in functions.py would do
2173+
from google.genai import types as genai_types
2174+
2175+
fake_function_call = mock.MagicMock()
2176+
fake_function_call.name = "expensive_api_tool"
2177+
2178+
# This mirrors the behavior in flows/llm_flows/functions.py:_get_tool
2179+
with pytest.raises(ValueError, match="Tool 'expensive_api_tool' not found"):
2180+
from google.adk.flows.llm_flows.functions import _get_tool
2181+
2182+
_get_tool(fake_function_call, tools_dict)
2183+
2184+
2185+
@pytest.mark.asyncio
2186+
async def test_rolling_window_full_lifecycle_with_tool_resolution():
2187+
"""End-to-end: load skills, evict, verify tool availability at each step."""
2188+
skills = []
2189+
additional_tools = []
2190+
for i in range(4):
2191+
skill = mock.create_autospec(models.Skill, instance=True)
2192+
skill.name = f"skill_{i}"
2193+
skill.instructions = f"instructions {i}"
2194+
frontmatter = mock.create_autospec(models.Frontmatter, instance=True)
2195+
frontmatter.name = f"skill_{i}"
2196+
frontmatter.metadata = {"adk_additional_tools": [f"tool_{i}"]}
2197+
frontmatter.model_dump.return_value = {"name": f"skill_{i}"}
2198+
skill.frontmatter = frontmatter
2199+
skills.append(skill)
2200+
2201+
tool = mock.create_autospec(skill_toolset.BaseTool, instance=True)
2202+
tool.name = f"tool_{i}"
2203+
additional_tools.append(tool)
2204+
2205+
toolset = skill_toolset.SkillToolset(
2206+
skills,
2207+
additional_tools=additional_tools,
2208+
max_activated_skills=2,
2209+
)
2210+
2211+
load_tool = skill_toolset.LoadSkillTool(toolset)
2212+
ctx = mock.create_autospec(skill_toolset.ToolContext, instance=True)
2213+
ctx._invocation_context = mock.MagicMock()
2214+
ctx.agent_name = "test_agent"
2215+
ctx.invocation_id = "inv-1"
2216+
2217+
readonly_context = mock.create_autospec(ReadonlyContext, instance=True)
2218+
readonly_context.agent_name = "test_agent"
2219+
readonly_context.invocation_id = "inv-1"
2220+
2221+
# Step 1: Load skill_0 and skill_1
2222+
current_state = []
2223+
for i in range(2):
2224+
ctx.state.get.return_value = list(current_state)
2225+
await load_tool.run_async(
2226+
args={"skill_name": f"skill_{i}"}, tool_context=ctx
2227+
)
2228+
current_state = ctx.state.__setitem__.call_args[0][1]
2229+
2230+
assert current_state == ["skill_0", "skill_1"]
2231+
2232+
readonly_context.state.get.return_value = list(current_state)
2233+
tools = await toolset.get_tools(readonly_context=readonly_context)
2234+
tool_names = {t.name for t in tools}
2235+
assert "tool_0" in tool_names
2236+
assert "tool_1" in tool_names
2237+
assert "tool_2" not in tool_names
2238+
2239+
# Step 2: Load skill_2 — should evict skill_0
2240+
ctx.state.get.return_value = list(current_state)
2241+
await load_tool.run_async(args={"skill_name": "skill_2"}, tool_context=ctx)
2242+
current_state = ctx.state.__setitem__.call_args[0][1]
2243+
2244+
assert current_state == ["skill_1", "skill_2"]
2245+
2246+
readonly_context.state.get.return_value = list(current_state)
2247+
tools = await toolset.get_tools(readonly_context=readonly_context)
2248+
tool_names = {t.name for t in tools}
2249+
assert "tool_0" not in tool_names # evicted!
2250+
assert "tool_1" in tool_names
2251+
assert "tool_2" in tool_names
2252+
2253+
# Step 3: Load skill_3 — should evict skill_1
2254+
ctx.state.get.return_value = list(current_state)
2255+
await load_tool.run_async(args={"skill_name": "skill_3"}, tool_context=ctx)
2256+
current_state = ctx.state.__setitem__.call_args[0][1]
2257+
2258+
assert current_state == ["skill_2", "skill_3"]
2259+
2260+
readonly_context.state.get.return_value = list(current_state)
2261+
tools = await toolset.get_tools(readonly_context=readonly_context)
2262+
tool_names = {t.name for t in tools}
2263+
assert "tool_0" not in tool_names
2264+
assert "tool_1" not in tool_names # evicted!
2265+
assert "tool_2" in tool_names
2266+
assert "tool_3" in tool_names

0 commit comments

Comments
 (0)