From d2f5f5bd767c75b065318d39cf6a974a8183e79d Mon Sep 17 00:00:00 2001 From: Kiki McArron Date: Thu, 10 Sep 2026 16:05:28 +0200 Subject: [PATCH 1/9] Implement parse_extra_resources function --- statek/extra_resources.py | 80 +++++++++++++++++++++++++++++++++++ tests/test_extra_resources.py | 75 ++++++++++++++++++++++++++++++++ 2 files changed, 155 insertions(+) create mode 100644 statek/extra_resources.py create mode 100644 tests/test_extra_resources.py diff --git a/statek/extra_resources.py b/statek/extra_resources.py new file mode 100644 index 0000000..93baae1 --- /dev/null +++ b/statek/extra_resources.py @@ -0,0 +1,80 @@ +# Copyright 2026 Statek authors +# +# 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. + +"""Parse incremental extra resource declarations for difficulty levels L/M/H.""" + +import re +from typing import TypeAlias + + +ExtraResources: TypeAlias = tuple[list[str] | None, list[str] | None, list[str] | None] + +_RESOURCE_NAME = r"[^\s(),:{}\[\]'\"\\]+" +_RESOURCE_GROUP = re.compile( + rf"(?P[LMH]+)\s*:\s*" + rf"(?P{_RESOURCE_NAME}(?:\s*,\s*(?!{_RESOURCE_NAME}\s*:){_RESOURCE_NAME})*)", + re.IGNORECASE, +) + + +def parse_extra_resources(input: str) -> ExtraResources: # pylint: disable=redefined-builtin + """Parse prompt metadata into incremental resource lists for L, M and H. + + Round parentheses around individual groups are optional, for example + ``(LM:assistant,researcher),(H:helpdesk)`` or + ``LM:assistant,researcher,H:helpdesk``. Selectors are case-insensitive; + resource names are case-sensitive unquoted tokens. Empty input returns + three None values. + + Each resource is stored only at its lowest declared difficulty, since it + is inherited at higher levels. Order within that lowest-level bucket is + retained. Parsing neither resolves agents nor executes expressions. + + Raises: + ValueError: If the declaration has invalid selectors, names or grouping. + """ + remaining = input.strip() + buckets: list[list[str]] = [[], [], []] + while remaining: + parenthesized = remaining.startswith("(") + if parenthesized: + remaining = remaining[1:].lstrip() + match = _RESOURCE_GROUP.match(remaining) + if match is None: + raise ValueError(f"Invalid EXTRA_RESOURCES declaration near {remaining!r}") + + level = min("LMH".index(label) for label in match["levels"].upper()) + buckets[level].extend(name.strip() for name in match["names"].split(",")) + remaining = remaining[match.end():].lstrip() + if parenthesized: + if not remaining.startswith(")"): + raise ValueError("Invalid EXTRA_RESOURCES: expected closing ')' after group") + remaining = remaining[1:].lstrip() + if remaining: + if not remaining.startswith(","): + raise ValueError(f"Invalid EXTRA_RESOURCES separator near {remaining!r}") + remaining = remaining[1:].lstrip() + if not remaining: + raise ValueError("Invalid EXTRA_RESOURCES: trailing comma") + + seen: set[str] = set() + result: list[list[str] | None] = [] + for bucket in buckets: + additions = [] + for name in bucket: + if name not in seen: + seen.add(name) + additions.append(name) + result.append(additions or None) + return result[0], result[1], result[2] diff --git a/tests/test_extra_resources.py b/tests/test_extra_resources.py new file mode 100644 index 0000000..7694d24 --- /dev/null +++ b/tests/test_extra_resources.py @@ -0,0 +1,75 @@ +"""Tests for incremental difficulty-specific extra resource declarations.""" + +import pytest + +from statek.extra_resources import parse_extra_resources + + +@pytest.mark.parametrize( + ("definition", "expected"), + [ + ("", (None, None, None)), + (" \n\t ", (None, None, None)), + ("(H:assistant)", (None, None, ["assistant"])), + ("H:assistant", (None, None, ["assistant"])), + ("LM:assistant", (["assistant"], None, None)), + ("LH:assistant", (["assistant"], None, None)), + ("MH:assistant", (None, ["assistant"], None)), + ("HML:assistant", (["assistant"], None, None)), + ("LL:assistant", (["assistant"], None, None)), + (" ( lm : Assistant , researcher ) , ( h : helpdesk ) ", + (["Assistant", "researcher"], None, ["helpdesk"])), + ("L:a,a,M:a,b,H:b,c,c", (["a"], ["b"], ["c"])), + ("H:a,b,M:a,L:b", (["b"], ["a"], None)), + ("H:a,L:b,a", (["b", "a"], None, None)), + ("(H:a),(L:a)", (["a"], None, None)), + ("L:a,L:b", (["a", "b"], None, None)), + ("L:A,a", (["A", "a"], None, None)), + ("M:team.assistant,help-desk", (None, ["team.assistant", "help-desk"], None)), + ], +) +def test_parse_extra_resources(definition, expected): + """Resources occur once, at the lowest declared difficulty.""" + assert parse_extra_resources(definition) == expected + + +@pytest.mark.parametrize( + "definition", + [ + "(LM:schedule_assistant,researcher),(H:helpdesk)", + "LM:schedule_assistant,researcher,H:helpdesk", + "(LM:schedule_assistant,researcher),H:helpdesk", + "LM:schedule_assistant,researcher,(H:helpdesk)", + ], +) +def test_parse_extra_resources_optional_group_parentheses(definition): + """Parenthesized, bare and mixed forms have identical semantics.""" + assert parse_extra_resources(definition) == ( + ["schedule_assistant", "researcher"], None, ["helpdesk"] + ) + + +@pytest.mark.parametrize( + "definition", + [ + "Q:a", "LQ:a", ":a", "a", "a,L:b", "L:", "()", "( )", + "L:a,,b", "L:a,", ",L:a", "(L:a,)", "(L:a", "L:a)", + "((L:a))", "(L:a)(H:b)", "(L:a),b", "(L:a,H:b)", + "L:a:b", "L:a,Q:b", "L:a,(Q:b)", "L:a H:b", "L:a b", + "{L:a}", "[L:a]", "L:'a'", 'L:"a"', "L:a\\b", + "H:assistant()", "L:__import__('os')", "(L:a)garbage", + ], +) +def test_parse_extra_resources_rejects_malformed_definitions(definition): + """Invalid groups and trailing input must not be silently discarded.""" + with pytest.raises(ValueError, match="EXTRA_RESOURCES"): + parse_extra_resources(definition) + + +def test_parse_extra_resources_results_own_their_lists(): + """Mutating one parse result cannot affect another result or level.""" + first = parse_extra_resources("L:a,M:b") + second = parse_extra_resources("L:a,M:b") + first[0].append("changed") + assert first[1] == ["b"] + assert second == (["a"], ["b"], None) From 8fd389ff3d5e36cd5a90f50d9143ecc7988013eb Mon Sep 17 00:00:00 2001 From: Kiki McArron Date: Sat, 12 Sep 2026 09:57:51 +0200 Subject: [PATCH 2/9] Issues 1499 and 1500 --- statek/agents/agent.py | 6 +- statek/agents/dialog_agent.py | 4 + statek/executors/job.py | 21 ++- statek/executors/utils.py | 15 ++ statek/extra_resources.py | 21 ++- statek/task.py | 6 + tests/test_agentic_loop.py | 3 + tests/test_extra_resources_integration.py | 184 ++++++++++++++++++++++ 8 files changed, 255 insertions(+), 5 deletions(-) create mode 100644 tests/test_extra_resources_integration.py diff --git a/statek/agents/agent.py b/statek/agents/agent.py index 96fdeef..0e20fbb 100644 --- a/statek/agents/agent.py +++ b/statek/agents/agent.py @@ -29,6 +29,7 @@ ) from statek.executors.post_processor import effective_post_processing from statek.provider_config import ProviderConfig +from statek.extra_resources import ExtraResources from statek.prompt_config import ( SystemPrompt, SystemPromptData, @@ -486,7 +487,7 @@ def update_warmup_def(self, unparsed_warmup_def: Optional[str]) -> bool: self._X__ref_locals = None return True - def create_job_def( + def create_job_def( # pylint: disable=too-many-arguments,too-many-positional-arguments self, tools: Optional[List[Callable]] = None, warmup_code: WarmupCodeInput = None, @@ -494,6 +495,7 @@ def create_job_def( locale=None, post_processing=None, provider_config: Optional[ProviderConfig] = None, + extra_resources: ExtraResources | None = None, **kwargs ) -> JobDef: # pylint: disable=unused-argument @@ -510,6 +512,7 @@ def create_job_def( locale: optional locale for job execution post_processing: optional resolved post-processor or sequence of post-processors for this job definition + extra_resources: parsed prompt resources transported by internal callers kwargs: job specific parameters for prompt formatting (i.e. job_params) Returns: @@ -541,6 +544,7 @@ def create_job_def( locale=locale, post_processing=resolved_post_processing, provider_config=provider_config, + extra_resources=extra_resources, ) def _combine_warmup_code(self, warmup_code): diff --git a/statek/agents/dialog_agent.py b/statek/agents/dialog_agent.py index 1b9a540..060c5cf 100644 --- a/statek/agents/dialog_agent.py +++ b/statek/agents/dialog_agent.py @@ -21,6 +21,7 @@ import dbzero as db0 from statek.agents.agent import SupervisedAgent from statek.chat_style import ChatStyle +from statek.extra_resources import ExtraResources if TYPE_CHECKING: from statek.executors.job import Job @@ -238,6 +239,7 @@ def create_job_def( # pylint: disable=no-member,arguments-renamed,too-many-argu locale=None, post_processing=None, provider_config=None, + extra_resources: ExtraResources | None = None, **kwargs ): """Create a JobDef with a dialog chat style. @@ -254,6 +256,7 @@ def create_job_def( # pylint: disable=no-member,arguments-renamed,too-many-argu chat_style: chat style override; if omitted, falls back to prompt metadata locale: optional locale for job execution post_processing: optional resolved post-processor or sequence of post-processors + extra_resources: parsed prompt resources forwarded to the parent **kwargs: job-specific parameters forwarded to the parent """ if chat_style is None and self._metadata and 'CHAT_STYLE' in self._metadata: @@ -265,6 +268,7 @@ def create_job_def( # pylint: disable=no-member,arguments-renamed,too-many-argu locale=locale, post_processing=post_processing, provider_config=provider_config, + extra_resources=extra_resources, **kwargs ) job_def.set_chat_style(chat_style or ChatStyle.MD_DIALOG) # pylint: disable=no-member diff --git a/statek/executors/job.py b/statek/executors/job.py index 1ffaf23..923e553 100644 --- a/statek/executors/job.py +++ b/statek/executors/job.py @@ -35,6 +35,9 @@ import dbzero as db0 from dbzero import memo, enum from statek.pyenv import Error, ErrorKind, PyEnv +from statek.extra_resources import ( + ExtraResources, extra_resources_from_metadata, extra_resources_identity, normalize_extra_resources, +) from statek.executors.llm_usage import LLM_Usage from statek.executors.chat_log_item import ( ChatLogItem, @@ -356,6 +359,7 @@ def _job_def_identity_hash( chat_style, post_processing, provider_config=None, + extra_resources=None, ) -> str: payload = ( warmup_code, @@ -366,6 +370,7 @@ def _job_def_identity_hash( chat_style, post_processing_identity(post_processing), provider_config_identity(provider_config), + extra_resources_identity(extra_resources), ) encoded = str(payload).encode("utf-8") return hashlib.sha1(encoded).hexdigest()[:4] @@ -380,10 +385,11 @@ def _job_def_identity_tag( chat_style, post_processing=None, provider_config=None, + extra_resources=None, ) -> str: return ( f"{_JOBDEF_HASH_TAG_PREFIX}" - f"{_job_def_identity_hash(warmup_code, model_family, model, job_params, locale, chat_style, post_processing, provider_config)}" + f"{_job_def_identity_hash(warmup_code, model_family, model, job_params, locale, chat_style, post_processing, provider_config, extra_resources)}" ) @@ -398,6 +404,7 @@ def job_def_identity_tag_for_job_def(job_def: "JobDef") -> str: getattr(job_def, "_chat_style", None), job_def.post_processing, job_def.provider_config, + job_def.extra_resources, ) @@ -421,9 +428,11 @@ class JobDef: # Optional locale for language-specific behaviour locale: Optional["StatekLocale"] = None # Optional post-processing sequence applied after LLM responses - post_processing: PostProcessingInput = None + post_processing: Optional[PostProcessingInput] = None # Durable provider configuration snapshot used to create this job definition. provider_config: Optional[ProviderConfig] = None + # Incremental prompt resource additions at L/M/H, kept out of LLM metadata. + extra_resources: Optional[ExtraResources] = None def __post_init__(self): if self.metadata is None: @@ -433,6 +442,14 @@ def __post_init__(self): else {} ) + resources = self.extra_resources + if resources is None: + resources = extra_resources_from_metadata(self.metadata) + resources = normalize_extra_resources(resources) + self.extra_resources = resources if any(resources) else None + if "EXTRA_RESOURCES" in self.metadata: + self.metadata = {key: value for key, value in self.metadata.items() if key != "EXTRA_RESOURCES"} + metadata_model = self.metadata.get('MODEL') if self.metadata else None if metadata_model is None: role = self.agent.role if self.agent is not None else '' diff --git a/statek/executors/utils.py b/statek/executors/utils.py index 5d7a2c3..80ead5d 100644 --- a/statek/executors/utils.py +++ b/statek/executors/utils.py @@ -26,6 +26,7 @@ import dbzero as db0 from statek.exceptions import FutureError, LLM_HarnessError +from statek.extra_resources import extra_resources_from_metadata, extra_resources_identity from statek.future import FutureResult @@ -1523,12 +1524,15 @@ def find_existing_job_def( chat_style: object = _MATCH_UNSET, post_processing: object = None, provider_config: Optional[ProviderConfig] = None, + extra_resources: object = _MATCH_UNSET, ) -> Optional[JobDef]: """Find an existing JobDef matching the given agent and warmup_code. Args: agent: the Agent the JobDef must be associated with warmup_code: the warmup code to match (compared after parsing) + extra_resources: parsed resource criterion; omitted means no filter, + while None or an empty triple matches no additional resources Returns: The first matching JobDef, or None if not found @@ -1553,15 +1557,19 @@ def _matches(job_def): return False if not provider_configs_match(job_def.provider_config, provider_config): return False + if extra_resources is not _MATCH_UNSET and extra_resources_identity(job_def.extra_resources) != resource_identity: + return False return True parsed = parse_warmup_code(warmup_code) + resource_identity = extra_resources_identity(extra_resources) if extra_resources is not _MATCH_UNSET else None agent_tag = db0.as_tag(agent) if ( model is not None and job_params is not _MATCH_UNSET and locale is not _MATCH_UNSET and chat_style is not _MATCH_UNSET + and extra_resources is not _MATCH_UNSET ): resolved_model_family = model_family or ensure_model_name(model).model_family lookup_tag = _job_def_identity_tag( @@ -1573,6 +1581,7 @@ def _matches(job_def): chat_style, post_processing, provider_config, + extra_resources, ) for job_def in db0.find(JobDef, agent_tag, lookup_tag): if _matches(job_def): @@ -1725,6 +1734,7 @@ def task_queue_size() -> int: # Reuse an existing matching job definition or create a new one model_family, model_to_use = _resolve_job_def_model(agent, provider) provider_config = resolve_settings_provider_config(get_statek_settings()) + extra_resources = extra_resources_from_metadata(agent._metadata) # pylint: disable=protected-access job_def = find_existing_job_def( agent, warmup_code, @@ -1734,6 +1744,7 @@ def task_queue_size() -> int: locale=None, chat_style=None, provider_config=provider_config, + extra_resources=extra_resources, ) if job_def: # Clear any previous errors on the job definition they might'be been fixed after process restart @@ -1747,6 +1758,7 @@ def task_queue_size() -> int: job_params=None, warmup_code=parsed_warmup_code, provider_config=provider_config, + extra_resources=extra_resources, ) start_jobs_func = _make_start_jobs_func(agent, job_def, task_queue_size_func, provider) @@ -1791,6 +1803,7 @@ async def run_agentic_fleet( model_family, model_to_use = _resolve_job_def_model(agent, provider) provider_config = resolve_settings_provider_config(get_statek_settings()) + extra_resources = extra_resources_from_metadata(agent._metadata) # pylint: disable=protected-access job_def = find_existing_job_def( agent, warmup_code, @@ -1800,6 +1813,7 @@ async def run_agentic_fleet( locale=None, chat_style=None, provider_config=provider_config, + extra_resources=extra_resources, ) if job_def: job_def.clear_errors() @@ -1812,6 +1826,7 @@ async def run_agentic_fleet( job_params=None, warmup_code=parsed_warmup_code, provider_config=provider_config, + extra_resources=extra_resources, ) start_jobs_funcs.append(_make_start_jobs_func(agent, job_def, task_queue_size_func, provider)) diff --git a/statek/extra_resources.py b/statek/extra_resources.py index 93baae1..5cfa1e0 100644 --- a/statek/extra_resources.py +++ b/statek/extra_resources.py @@ -15,6 +15,7 @@ """Parse incremental extra resource declarations for difficulty levels L/M/H.""" import re +from collections.abc import Mapping from typing import TypeAlias @@ -68,13 +69,29 @@ def parse_extra_resources(input: str) -> ExtraResources: # pylint: disable=rede if not remaining: raise ValueError("Invalid EXTRA_RESOURCES: trailing comma") + return normalize_extra_resources((buckets[0], buckets[1], buckets[2])) + + +def normalize_extra_resources(resources: ExtraResources | None) -> ExtraResources: + """Copy resource lists and retain each name only at its lowest level.""" seen: set[str] = set() result: list[list[str] | None] = [] - for bucket in buckets: + for bucket in resources or (None, None, None): additions = [] - for name in bucket: + for name in bucket or (): if name not in seen: seen.add(name) additions.append(name) result.append(additions or None) return result[0], result[1], result[2] + + +def extra_resources_from_metadata(metadata: Mapping[str, str] | None) -> ExtraResources: + """Prepare prompt resources once for both definition lookup and creation.""" + declaration = metadata.get("EXTRA_RESOURCES") if metadata else None + return parse_extra_resources(declaration) if declaration is not None else (None, None, None) + + +def extra_resources_identity(resources: ExtraResources | None) -> tuple[tuple[str, ...], ...]: + """Return the same ordered value identity for Python and persisted lists.""" + return tuple(tuple(bucket or ()) for bucket in normalize_extra_resources(resources)) diff --git a/statek/task.py b/statek/task.py index 4df80ed..ec2c184 100644 --- a/statek/task.py +++ b/statek/task.py @@ -20,6 +20,7 @@ from typing import Any, Dict, Iterable, List, Optional, Tuple, Type import dbzero as db0 from .exceptions import FutureError +from .extra_resources import ExtraResources, extra_resources_from_metadata from .future import FutureResult, temporal from .system import tool from .agents.agent import Agent, SupervisedAgent @@ -60,6 +61,7 @@ def _find_reusable_job_def( post_processing, kwargs: Dict[str, Any], provider_config: Optional[ProviderConfig], + extra_resources: ExtraResources, ): """Find an existing JobDef matching create_new_job's reusable definition.""" from statek.executors.utils import find_existing_job_def # pylint: disable=import-outside-toplevel @@ -79,6 +81,7 @@ def _find_reusable_job_def( chat_style=_dialog_chat_style(agent, kwargs) if isinstance(agent, DialogAgent) else None, post_processing=resolved_post_processing, provider_config=provider_config, + extra_resources=extra_resources, ) @@ -447,6 +450,7 @@ def create_new_job( # pylint: disable=too-many-arguments,too-many-positional-ar shared_vars = shared_vars or {} effective_locale = _resolve_child_locale(parent_job, locale) provider_config = resolve_settings_provider_config(_get_statek_settings()) + extra_resources = extra_resources_from_metadata(agent._metadata) # pylint: disable=protected-access job_def = _find_reusable_job_def( agent, warmup_code, @@ -454,6 +458,7 @@ def create_new_job( # pylint: disable=too-many-arguments,too-many-positional-ar post_processing, kwargs, provider_config, + extra_resources, ) if job_def is None: job_def = agent.create_job_def( @@ -461,6 +466,7 @@ def create_new_job( # pylint: disable=too-many-arguments,too-many-positional-ar locale=effective_locale, post_processing=post_processing, provider_config=provider_config, + extra_resources=extra_resources, **kwargs, ) diff --git a/tests/test_agentic_loop.py b/tests/test_agentic_loop.py index 5e14a18..0a0eaf6 100644 --- a/tests/test_agentic_loop.py +++ b/tests/test_agentic_loop.py @@ -234,6 +234,7 @@ def recording_find(*args, **kwargs): job_params=None, locale=None, chat_style=None, + extra_resources=None, ) assert result is job_def @@ -262,6 +263,7 @@ def test_hash_collision_still_uses_exact_comparison( job_params=None, locale=None, chat_style=None, + extra_resources=None, ) assert result is matching @@ -294,6 +296,7 @@ def test_provider_config_hash_collision_still_uses_exact_comparison( job_params=None, locale=None, chat_style=None, + extra_resources=None, provider_config=requested_config, ) diff --git a/tests/test_extra_resources_integration.py b/tests/test_extra_resources_integration.py new file mode 100644 index 0000000..1201af6 --- /dev/null +++ b/tests/test_extra_resources_integration.py @@ -0,0 +1,184 @@ +"""Prompt resource storage and definition reuse across job creation paths.""" +# pylint: disable=protected-access,no-member + +from pathlib import Path + +import dbzero as db0 +import pytest + +from statek.agents.dialog_agent import DialogAgent +from statek.executors.job import Job, JobDef, job_def_identity_tag_for_job_def +from statek.executors.utils import ( + AgentLoopDef, find_existing_job_def, run_agentic_fleet, run_agentic_loop, +) +from statek.extra_resources import parse_extra_resources +from statek.prompt_config import PromptDef, parse_prompt_file, update_prompt_config +from statek.task import create_new_job +from tests.conftest import DB0_DIR + + +def _send_message(body: str, media: str | None = None) -> None: # pylint: disable=unused-argument + """Accept dialog messages without external effects.""" + + +@pytest.mark.parametrize("declaration", ["H:helpdesk", "(H:helpdesk)"]) +def test_extra_resources_prompt_to_job_def( + supervised_agent, tmp_path: Path, declaration: str, +) -> None: + """Raw declarations remain on the agent, not in job execution metadata.""" + prompt = tmp_path / "test.md" + prompt.write_text( + f"# MODEL: test-model\n# TEMPERATURE: 0.3\n# EXTRA_RESOURCES: {declaration}\n" + "# System Prompt\nYou are an assistant.\n", encoding="utf-8", + ) + prompt_def = parse_prompt_file(prompt) + update_prompt_config({"test": prompt_def}, agents=[supervised_agent]) + definition = supervised_agent.create_job_def(user_question="question") + assert definition.extra_resources == (None, None, ["helpdesk"]) + assert "EXTRA_RESOURCES" not in definition.metadata + assert definition.metadata["TEMPERATURE"] == "0.3" + assert supervised_agent._metadata["EXTRA_RESOURCES"] == declaration + assert definition.job_params == {"user_question": "question"} + assert "EXTRA_RESOURCES" not in Job(job_def=definition)._build_request_data()["metadata"] + + +def test_extra_resources_direct_constructor_and_snapshot(supervised_agent) -> None: + """Direct constructors isolate resource metadata without altering the source.""" + metadata = {"MODEL": "test-model", "EXTRA_RESOURCES": "LM:a,H:a,b"} + first = JobDef(agent=supervised_agent, metadata=metadata) + second = JobDef(agent=supervised_agent, metadata=metadata) + assert first.extra_resources == (["a"], None, ["b"]) + assert "EXTRA_RESOURCES" not in first.metadata + assert metadata["EXTRA_RESOURCES"] == "LM:a,H:a,b" + first.extra_resources[0].append("changed") + assert second.extra_resources == (["a"], None, ["b"]) + + +def test_extra_resources_internal_factory_parameter(supervised_agent) -> None: + """Already parsed configuration is transported separately from job_params.""" + resources = parse_extra_resources("M:a,H:b") + definition = supervised_agent.create_job_def(extra_resources=resources, question="hello") + assert definition.extra_resources == resources + assert definition.job_params == {"question": "hello"} + resources[1].append("changed") + assert definition.extra_resources == (None, ["a"], ["b"]) + + +def test_extra_resources_absent_preserves_metadata_sharing(supervised_agent) -> None: + definition = supervised_agent.create_job_def() + assert definition.extra_resources is None + assert definition.metadata is supervised_agent._metadata + + +def test_extra_resources_malformed_metadata_fails_before_job_creation(supervised_agent) -> None: + supervised_agent.update_metadata({"MODEL": "test-model", "EXTRA_RESOURCES": "Q:a"}) + with pytest.raises(ValueError, match="EXTRA_RESOURCES"): + create_new_job(supervised_agent) + assert len(db0.find(Job, db0.as_tag(supervised_agent))) == 0 + + +def test_extra_resources_reuse_and_prompt_changes(supervised_agent) -> None: + supervised_agent.update_metadata({"MODEL": "test-model", "EXTRA_RESOURCES": "LM:a,H:b"}) + first = create_new_job(supervised_agent) + supervised_agent.update_metadata({"MODEL": "test-model", "EXTRA_RESOURCES": "(L:a),(H:b,b)"}) + assert create_new_job(supervised_agent).job_def is first.job_def + update_prompt_config({"test": PromptDef( + system=None, metadata={"MODEL": "test-model", "EXTRA_RESOURCES": "M:a,H:b"}, + )}, agents=[supervised_agent]) + second = create_new_job(supervised_agent) + assert second.job_def is not first.job_def + assert first.job_def.extra_resources == (["a"], None, ["b"]) + update_prompt_config({"test": PromptDef( + system=None, metadata={"MODEL": "test-model"}, + )}, agents=[supervised_agent]) + third = create_new_job(supervised_agent) + assert third.job_def.extra_resources is None + assert third.job_def is not first.job_def + assert len(db0.find(JobDef, db0.as_tag(supervised_agent))) == 3 + + +def test_extra_resources_lookup_filters_and_hash_collision(supervised_agent, monkeypatch) -> None: + """Exact comparison protects resource identity even with identical short hashes.""" + monkeypatch.setattr("statek.executors.job._job_def_identity_hash", lambda *args: "same") + first = supervised_agent.create_job_def(extra_resources=parse_extra_resources("L:a")) + second = supervised_agent.create_job_def(extra_resources=parse_extra_resources("M:a")) + assert job_def_identity_tag_for_job_def(first) == job_def_identity_tag_for_job_def(second) + assert find_existing_job_def( + supervised_agent, None, extra_resources=(None, ["a"], None), + ) is second + assert find_existing_job_def(supervised_agent, None, extra_resources=None) is None + assert find_existing_job_def(supervised_agent, None) in (first, second) + assert find_existing_job_def( + supervised_agent, None, model="test-model", job_params=None, locale=None, + chat_style=None, + ) in (first, second) + assert find_existing_job_def( + supervised_agent, None, model="test-model", job_params=None, locale=None, + chat_style=None, extra_resources=(None, ["a"], None), + ) is second + assert find_existing_job_def( + supervised_agent, None, model="test-model", job_params=None, locale=None, + chat_style=None, extra_resources=(["missing"], None, None), + ) is None + + +def test_extra_resources_dialog_factory(db0_fixture) -> None: # pylint: disable=unused-argument + agent = DialogAgent( + role="dialog", send_message=_send_message, + _metadata={"MODEL": "test-model", "EXTRA_RESOURCES": "H:a"}, + ) + first = create_new_job(agent) + assert first.job_def.extra_resources == (None, None, ["a"]) + assert first.job_def.job_params is None + assert create_new_job(agent).job_def is first.job_def + + +def test_extra_resources_persistence_and_ordered_identity(supervised_agent) -> None: + definition = supervised_agent.create_job_def(extra_resources=parse_extra_resources("L:a,b,H:c")) + identifier = db0.uuid(definition) + agent_identifier = db0.uuid(supervised_agent) + db0.close() + db0.init(DB0_DIR, read_write=True) + db0.open("test_prefix", "rw") + restored = db0.fetch(identifier) + agent = db0.fetch(agent_identifier) + assert restored.extra_resources == (["a", "b"], None, ["c"]) + assert find_existing_job_def(agent, None, extra_resources=(["a", "b"], None, ["c"])) is restored + assert find_existing_job_def(agent, None, extra_resources=(["b", "a"], None, ["c"])) is None + + +def test_extra_resources_empty_and_redundant_internal_values(supervised_agent) -> None: + empty = supervised_agent.create_job_def(extra_resources=([], None, [])) + assert empty.extra_resources is None + assert find_existing_job_def( + supervised_agent, None, extra_resources=(None, None, None), + ) is empty + definition = supervised_agent.create_job_def( + extra_resources=(["a", "a"], ["a", "b"], ["b"]), + ) + assert definition.extra_resources == (["a"], ["b"], None) + assert find_existing_job_def( + supervised_agent, None, extra_resources=(["a"], ["a", "b"], None), + ) is definition + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fleet", [False, True]) +async def test_extra_resources_loop_and_fleet(supervised_agent, fleet: bool) -> None: + supervised_agent.update_metadata({"MODEL": "test-model", "EXTRA_RESOURCES": "H:a"}) + for _ in range(2): + if fleet: + await run_agentic_fleet( + [AgentLoopDef(supervised_agent, None, lambda: 0)], + queue_prefixes=[db0.get_current_prefix().name], auto_terminate=True, + ) + else: + await run_agentic_loop( + supervised_agent, None, lambda: 0, + queue_prefixes=[db0.get_current_prefix().name], auto_terminate=True, + ) + definitions = tuple(db0.find(JobDef, db0.as_tag(supervised_agent))) + assert len(definitions) == 1 + assert definitions[0].extra_resources == (None, None, ["a"]) + assert "EXTRA_RESOURCES" not in definitions[0].metadata + assert supervised_agent._metadata["EXTRA_RESOURCES"] == "H:a" From 40e2ab4ee67fdf8623777e5008c8652c7be5bf4a Mon Sep 17 00:00:00 2001 From: Kiki McArron Date: Tue, 15 Sep 2026 10:06:50 +0200 Subject: [PATCH 3/9] Implement & test the 'Capacities Expansion' functionality - part 1 --- statek/executors/job.py | 40 ++++++ tests/test_extra_resources_runtime.py | 189 ++++++++++++++++++++++++++ 2 files changed, 229 insertions(+) create mode 100644 tests/test_extra_resources_runtime.py diff --git a/statek/executors/job.py b/statek/executors/job.py index 923e553..0515055 100644 --- a/statek/executors/job.py +++ b/statek/executors/job.py @@ -88,6 +88,7 @@ ) if TYPE_CHECKING: + from statek.agents.agent import Agent from statek.agents.dialog_agent import Reminder from statek.task import SubTaskHandler @@ -1477,6 +1478,45 @@ def get_request_data(self, turn_num: int) -> Dict[str, Any]: push_log=history_push_log, ) + def get_resource_agents(self) -> List["Agent"]: + """Return the receiver and configured donors active at this job's difficulty. + + Donors are resolved by role in the receiver's prefix, independently of + running jobs. Lower-level declarations remain active at higher levels; + donor declarations are not expanded recursively. The returned list is + job-local and does not modify any agent's configuration. + + Raises: + ValueError: If an activated role has no configured agent. + """ + receiver = self.job_def.agent + resources = getattr(self.job_def, "extra_resources", None) + if not resources: + return [receiver] + + level = tuple(TaskDifficulty.values()).index(self.get_current_difficulty()) + roles = dict.fromkeys( + role for index in range(level + 1) for role in resources[index] or () + if role != receiver.role + ) + if not roles: + return [receiver] + + # Agent imports JobDef from this module, so defer the runtime import. + from statek.agents.agent import Agent # pylint: disable=import-outside-toplevel + + prefix = db0.get_prefix_of(receiver).name + configured = {agent.role: agent for agent in db0.find(Agent, prefix=prefix)} + result = [receiver] + for role in roles: + if role not in configured: + raise ValueError( + f"EXTRA_RESOURCES for {receiver.role!r}: agent {role!r} " + f"is not configured in prefix {prefix!r}" + ) + result.append(configured[role]) + return result + def get_current_model(self) -> str: """Return the concrete model configured for the job's current difficulty.""" metadata = self.job_def.metadata or {} diff --git a/tests/test_extra_resources_runtime.py b/tests/test_extra_resources_runtime.py new file mode 100644 index 0000000..dc37d33 --- /dev/null +++ b/tests/test_extra_resources_runtime.py @@ -0,0 +1,189 @@ +"""Job-local capability donors resolve from configuration, not conversations.""" +# pylint: disable=protected-access,no-member,redefined-outer-name + +from collections.abc import Callable + +import dbzero as db0 +import pytest + +from statek.agents.agent import Agent, SupervisedAgent +from statek.executors.job import Job, JobDef +from statek.extra_resources import parse_extra_resources + + +@pytest.fixture +def resource_agents(agent_factory: Callable[..., Agent]) -> dict[str, Agent]: + """Configure donors in a different order from their activation order.""" + return { + role: agent_factory(role=role) + for role in ("receiver", "sounding_board", "schedule_assistant", "preferences_assistant") + } + + +@pytest.mark.parametrize("difficulty,count", [("L", 2), ("M", 3), ("H", 4)]) +def test_resource_agents_inherit_in_declaration_order( + resource_agents: dict[str, Agent], difficulty: str, count: int, +) -> None: + """Each difficulty includes cumulative donors in declaration order without creating jobs.""" + definition = JobDef( + agent=resource_agents["receiver"], + metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": difficulty}, + extra_resources=parse_extra_resources( + "L:preferences_assistant,M:schedule_assistant,H:sounding_board" + ), + ) + job = Job(job_def=definition) + expected = [resource_agents[role] for role in ( + "receiver", "preferences_assistant", "schedule_assistant", "sounding_board" + )] + + assert job.get_resource_agents() == expected[:count] + assert len(db0.find(Job)) == 1 + + +def test_resource_agents_without_declaration_return_receiver( + resource_agents: dict[str, Agent], +) -> None: + """A job without extra resources uses only its receiving agent's configuration.""" + receiver = resource_agents["receiver"] + job = Job(job_def=JobDef(agent=receiver)) + + assert job.get_resource_agents() == [receiver] + + +def test_resource_agents_deduplicate_receiver_and_donors( + resource_agents: dict[str, Agent], +) -> None: + """Repeated donor and self-references produce only one entry per configured role.""" + receiver = resource_agents["receiver"] + job = Job(job_def=JobDef( + agent=receiver, + metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": "H"}, + extra_resources=parse_extra_resources( + "L:receiver,preferences_assistant,M:receiver,preferences_assistant" + ), + )) + + assert job.get_resource_agents() == [receiver, resource_agents["preferences_assistant"]] + + +def test_resource_agents_do_not_expand_donor_declarations( + resource_agents: dict[str, Agent], +) -> None: + """Only explicitly named donors are included, not their own extra resources.""" + donor = resource_agents["preferences_assistant"] + donor.update_metadata({"MODEL": "test-model", "EXTRA_RESOURCES": "L:missing"}) + receiver = resource_agents["receiver"] + job = Job(job_def=JobDef( + agent=receiver, extra_resources=parse_extra_resources("L:preferences_assistant"), + )) + + assert job.get_resource_agents() == [receiver, donor] + + +def test_resource_agents_follow_job_escalation_without_shared_mutation( + resource_agents: dict[str, Agent], +) -> None: + """Panic expands one job's resources without changing other jobs or shared configuration.""" + receiver = resource_agents["receiver"] + definition = JobDef( + agent=receiver, + metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": "L"}, + extra_resources=parse_extra_resources("M:schedule_assistant,H:sounding_board"), + ) + first = Job(job_def=definition) + second = Job(job_def=definition) + + assert first.get_resource_agents() == [receiver] + first.panic() + assert first.get_resource_agents() == [receiver, resource_agents["schedule_assistant"]] + first.panic() + expanded = first.get_resource_agents() + assert expanded == [ + receiver, resource_agents["schedule_assistant"], resource_agents["sounding_board"] + ] + expanded.clear() + + assert len(first.get_resource_agents()) == 3 + assert second.get_resource_agents() == [receiver] + assert definition.extra_resources == (None, ["schedule_assistant"], ["sounding_board"]) + assert receiver._tools == [] + assert receiver._metadata == {"MODEL": "test-model"} + + +def test_resource_agents_ignore_running_donor_jobs(resource_agents: dict[str, Agent]) -> None: + """Creating donor jobs does not affect which configured agent supplies resources.""" + receiver = resource_agents["receiver"] + donor = resource_agents["schedule_assistant"] + job = Job(job_def=JobDef( + agent=receiver, extra_resources=parse_extra_resources("L:schedule_assistant"), + )) + expected = job.get_resource_agents() + donor_definition = JobDef(agent=donor) + for _ in range(5): + Job(job_def=donor_definition) + + assert job.get_resource_agents() == expected == [receiver, donor] + assert len(db0.find(Job)) == 6 + + +def test_resource_agents_resolve_configured_subclasses( + resource_agents: dict[str, Agent], +) -> None: + """Persisted Agent subclasses can supply resources just like base Agent configurations.""" + receiver = resource_agents["receiver"] + donor = SupervisedAgent( + role="supervised_donor", _system_prompt=None, _tools=[], _metadata={"MODEL": "test-model"}, + ) + job = Job(job_def=JobDef( + agent=receiver, extra_resources=parse_extra_resources("L:supervised_donor"), + )) + + assert job.get_resource_agents() == [receiver, donor] + + +def test_resource_agents_missing_definition_fails_only_when_activated( + resource_agents: dict[str, Agent], +) -> None: + """A missing donor raises an error only after its configured difficulty is reached.""" + receiver = resource_agents["receiver"] + job = Job(job_def=JobDef( + agent=receiver, + metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": "L"}, + extra_resources=parse_extra_resources("M:missing_agent"), + )) + + assert job.get_resource_agents() == [receiver] + job.panic() + with pytest.raises(ValueError, match="EXTRA_RESOURCES.*missing_agent"): + job.get_resource_agents() + + +def test_resource_agents_use_receiver_prefix_not_current_prefix( + resource_agents: dict[str, Agent], +) -> None: + """Lookup uses the receiver's prefix without changing the caller's current prefix.""" + receiver = resource_agents["receiver"] + donor = resource_agents["schedule_assistant"] + job = Job(job_def=JobDef( + agent=receiver, extra_resources=parse_extra_resources("L:schedule_assistant"), + )) + db0.open("unrelated_config", "rw") + Agent(role="schedule_assistant", _system_prompt=None, _tools=[]) + + assert job.get_resource_agents() == [receiver, donor] + assert db0.get_current_prefix().name == "unrelated_config" + + +def test_resource_agents_do_not_resolve_missing_role_from_another_prefix( + resource_agents: dict[str, Agent], +) -> None: + """A donor in another prefix cannot satisfy a missing local configuration.""" + job = Job(job_def=JobDef( + agent=resource_agents["receiver"], extra_resources=parse_extra_resources("L:other_donor"), + )) + db0.open("unrelated_config", "rw") + Agent(role="other_donor", _system_prompt=None, _tools=[]) + + with pytest.raises(ValueError, match="EXTRA_RESOURCES.*other_donor"): + job.get_resource_agents() From d49a0ac0f760ecf7b25ce7464eeb559dde901f6c Mon Sep 17 00:00:00 2001 From: Kiki McArron Date: Wed, 16 Sep 2026 09:53:03 +0200 Subject: [PATCH 4/9] Fix - BUG] LLM_Harness increases token limit by only one token per completed conversation turn --- statek/llm_harness.py | 6 +++--- tests/test_llm_harness.py | 34 +++++++++++++++++++++------------- 2 files changed, 24 insertions(+), 16 deletions(-) diff --git a/statek/llm_harness.py b/statek/llm_harness.py index 3d4ea3d..22dbd19 100644 --- a/statek/llm_harness.py +++ b/statek/llm_harness.py @@ -37,7 +37,7 @@ class LLM_Harness: max_token_usage: Total allowed maximum token usage (None for unlimited) limit_extension_per_completion: Fraction by which limits are extended after each job completion. Formula: - effective = base + (1.0 + extension) * num_completions + effective = base * (1.0 + extension * num_completions) """ def __init__(self, max_turns: Optional[int], max_exceptions: int, # pylint: disable=too-many-arguments,too-many-positional-arguments @@ -53,7 +53,7 @@ def _effective_limit(self, base_limit: Optional[int], num_completions: Optional[int]) -> Optional[float]: """Compute the effective limit after job completions. - Formula: base_limit + (1.0 + limit_extension_per_completion) * num_completions + Formula: base_limit * (1.0 + limit_extension_per_completion * num_completions) Returns None when base_limit is None (unlimited), or base_limit unchanged when num_completions is None (not tracked). @@ -62,7 +62,7 @@ def _effective_limit(self, base_limit: Optional[int], return None if num_completions is None: return base_limit - return base_limit + (1.0 + self.limit_extension_per_completion) * num_completions + return base_limit * (1.0 + self.limit_extension_per_completion * num_completions) def _check_exception_limits(self, job, num_completions): """Check exception-count limits shared by before- and after-step checks.""" diff --git a/tests/test_llm_harness.py b/tests/test_llm_harness.py index caee7a5..99d8e48 100644 --- a/tests/test_llm_harness.py +++ b/tests/test_llm_harness.py @@ -94,7 +94,7 @@ def test_max_consecutive_exceptions_checked_after_step(self): class TestEffectiveLimit: # pylint: disable=protected-access """Tests for _effective_limit calculation. - Formula: limit = base_limit + (1.0 + extension) * num_completions + Formula: limit = base_limit * (1.0 + extension * num_completions) """ def test_none_base_returns_none(self): @@ -118,26 +118,34 @@ def test_zero_completions_returns_base(self): limit_extension_per_completion=0.5) assert harness._effective_limit(5, 0) == 5 - def test_one_completion_extends_by_one_plus_fraction(self): - """1 completion with extension=0.5 adds 1.5 to base.""" + def test_one_completion_extends_by_base_fraction(self): + """1 completion with extension=0.5 adds half the base.""" harness = LLM_Harness(max_turns=5, max_exceptions=3, max_consecutive_exceptions=1, max_token_usage=None, limit_extension_per_completion=0.5) - assert harness._effective_limit(5, 1) == 6.5 + assert harness._effective_limit(5, 1) == 7.5 def test_two_completions(self): - """2 completions with extension=0.5 adds 3.0 to base.""" + """2 completions with extension=0.5 double the base.""" harness = LLM_Harness(max_turns=5, max_exceptions=3, max_consecutive_exceptions=1, max_token_usage=None, limit_extension_per_completion=0.5) - assert harness._effective_limit(5, 2) == 8.0 + assert harness._effective_limit(5, 2) == 10.0 def test_zero_extension(self): - """Extension=0.0 means each completion adds exactly 1.0.""" + """Zero extension preserves the base despite completions.""" harness = LLM_Harness(max_turns=10, max_exceptions=3, max_consecutive_exceptions=1, max_token_usage=None, limit_extension_per_completion=0.0) - assert harness._effective_limit(10, 3) == 13.0 + assert harness._effective_limit(10, 3) == 10.0 + + @pytest.mark.parametrize( + "completions, expected", [(0, 50000), (1, 55000), (2, 60000), (3, 65000)] + ) + def test_token_limit_grows_by_base_fraction(self, completions: int, expected: int) -> None: + """Each completion adds ten percent of the original token budget.""" + harness = LLM_Harness(None, 3, 1, 50000, limit_extension_per_completion=0.1) + assert harness._effective_limit(50000, completions) == pytest.approx(expected) class TestLimitExtensionBeforeStep: @@ -165,10 +173,10 @@ def test_exceptions_extended_after_completions(self): harness.check_before_step(_make_job(exception_count=5, num_completions=1)) def test_consecutive_exceptions_extended(self): - """max_consecutive_exceptions=1, extension=0.0, 2 completions → effective=3.0.""" + """max_consecutive_exceptions=1, extension=1.0, 2 completions → effective=3.0.""" harness = LLM_Harness(max_turns=None, max_exceptions=100, max_consecutive_exceptions=1, max_token_usage=None, - limit_extension_per_completion=0.0) + limit_extension_per_completion=1.0) harness.check_before_step(_make_job(max_consecutive_exceptions=3, num_completions=2)) with pytest.raises(LLM_HarnessError, match="consecutive"): harness.check_before_step(_make_job(max_consecutive_exceptions=4, num_completions=2)) @@ -187,13 +195,13 @@ class TestLimitExtensionAfterStep: """check_after_step uses extended limits when num_completions is set.""" def test_token_usage_extended_after_completions(self): - """max_token_usage=1000, extension=0.5, 2 completions → effective=1003.0.""" + """max_token_usage=1000, extension=0.5, 2 completions → effective=2000.0.""" harness = LLM_Harness(max_turns=None, max_exceptions=100, max_consecutive_exceptions=100, max_token_usage=1000, limit_extension_per_completion=0.5) - harness.check_after_step(_make_job(approx_token_usage=1003, num_completions=2)) + harness.check_after_step(_make_job(approx_token_usage=2000, num_completions=2)) with pytest.raises(LLM_HarnessError, match="token"): - harness.check_after_step(_make_job(approx_token_usage=1004, num_completions=2)) + harness.check_after_step(_make_job(approx_token_usage=2001, num_completions=2)) def test_exceptions_extended_after_step(self): """max_exceptions=3, extension=0.5, 1 completion → effective=4.5.""" From 7b3c31ecb9772a440c32b40889213f8563cc0905 Mon Sep 17 00:00:00 2001 From: Kiki McArron Date: Wed, 16 Sep 2026 11:20:59 +0200 Subject: [PATCH 5/9] Implement & test the 'Capacities Expansion' functionality - part 2 --- statek/agents/agent.py | 41 +++--- statek/executors/job.py | 22 +++- statek/executors/utils.py | 35 ++--- tests/test_extra_resources_runtime.py | 183 ++++++++++++++++++++++++++ tests/test_llm_api.py | 15 ++- 5 files changed, 253 insertions(+), 43 deletions(-) diff --git a/statek/agents/agent.py b/statek/agents/agent.py index 0e20fbb..a02fbf8 100644 --- a/statek/agents/agent.py +++ b/statek/agents/agent.py @@ -210,7 +210,9 @@ def update_description(self, new_description: str) -> bool: self._description = new_description return True - def _expand_tool_placeholders(self, text: str) -> Optional[str]: + def _expand_tool_placeholders( + self, text: str, available_tools: Optional[Sequence[Callable]] = None, + ) -> Optional[str]: """Expand {tools}, {brief_tools}, {detailed_tools} placeholders in text. Lines starting with '#' before a placeholder are embedded as a block comment. @@ -225,10 +227,10 @@ def _expand_tool_placeholders(self, text: str) -> Optional[str]: for name, brief, py_syntax in placeholders: pattern = re.compile(rf'^(\s*#\s*)\{{{name}\}}', re.MULTILINE) if pattern.search(text): - tools_str = self._format_tools(brief, py_syntax) + tools_str = self._format_tools(brief, py_syntax, available_tools) text = pattern.sub(block_comment(tools_str), text) elif f'{{{name}}}' in text: - tools_str = self._format_tools(brief, py_syntax) + tools_str = self._format_tools(brief, py_syntax, available_tools) text = text.replace(f'{{{name}}}', tools_str) return text @@ -236,6 +238,7 @@ def system_prompt( self, task_difficulty: TaskDifficulty, job_params: Dict = None, + available_tools: Optional[Sequence[Callable]] = None, **kwargs, ) -> str: """ @@ -263,7 +266,9 @@ def system_prompt( result = format_system_prompt( self._system_prompt, task_difficulty, - prompt_part_formatter=self._expand_tool_placeholders, + prompt_part_formatter=lambda text: self._expand_tool_placeholders( + text, available_tools, + ), ) format_ctx = {} if job_params: @@ -278,29 +283,23 @@ def system_prompt( return result.format_map(format_ctx) return result - def _format_tools(self, brief: bool, py_syntax: bool) -> str: + def _format_tools( + self, brief: bool, py_syntax: bool, + available_tools: Optional[Sequence[Callable]] = None, + ) -> str: """Format all tools with the specified settings.""" agent_name = _get_class_name(self) - formatted = [] def inner_format_tool(fn: Callable) -> str: parsed = parse_tool_docstring(fn) return format_docstring(parsed, brief=brief, py_syntax=py_syntax, agent=agent_name) + tools = self.tools if available_tools is None else available_tools formatted = [ inner_format_tool(fn) - for fn in self._tools + for fn in tools if not getattr(fn, "tool_hidden", False) ] - # also process tools specified by name - if self._tools_by_name: - for tool_name in self._tools_by_name: - fn = self.context.get(tool_name) - if fn is None: - raise ValueError(f'Missing tool defined by name "{tool_name}"') - if getattr(fn, "tool_hidden", False): - continue - formatted.append(inner_format_tool(fn)) return '\n\n'.join(formatted) @@ -328,14 +327,20 @@ def get_adapter(self, name: str) -> Optional[Any]: return None @property - def all_tools(self) -> List[Callable]: - """Return all tools assigned to this agent (both regular and internal).""" + def tools(self) -> List[Callable]: + """Return the agent's LLM-visible direct and name-registered tools.""" result = list(self._tools) if self._tools_by_name: for tool_name in self._tools_by_name: fn = self.context.get(tool_name) if fn is not None: result.append(fn) + return result + + @property + def all_tools(self) -> List[Callable]: + """Return all tools assigned to this agent (both regular and internal).""" + result = self.tools if self._internal_tools: result.extend(self._internal_tools) if self._internal_tools_by_name: diff --git a/statek/executors/job.py b/statek/executors/job.py index 0515055..a9a0db1 100644 --- a/statek/executors/job.py +++ b/statek/executors/job.py @@ -774,6 +774,7 @@ def system_prompt(self, difficulty: Optional[TaskDifficulty] = None) -> str: return self.job_def.agent.system_prompt( task_difficulty=difficulty or self.get_current_difficulty(), job_params=self.job_def.job_params, + available_tools=self.get_resource_tools(include_internal=False), ) def _ensure_ext_ref_storage(self): @@ -1418,7 +1419,7 @@ def _build_request_data( "system_prompt": system_prompt, "model": model, "metadata": metadata, - "available_tools": self.job_def.agent.all_tools, + "available_tools": self.get_resource_tools(include_internal=False), "provider_config": self.job_def.provider_config, } if temperature is not None: @@ -1517,6 +1518,25 @@ def get_resource_agents(self) -> List["Agent"]: result.append(configured[role]) return result + def get_resource_tools(self, include_internal: bool = True) -> List[Callable]: + """Return active receiver and donor tools, keeping the first tool per name. + + Receiver tools precede donor tools, so a donor cannot replace an + existing receiver binding. Internal tools are included for execution by + default but can be omitted from LLM-facing prompt documentation. + """ + result = [] + seen: set[str] = set() + for agent in self.get_resource_agents(): + tools = agent.all_tools if include_internal else agent.tools + for tool_fn in tools: + name = getattr(tool_fn, "__name__", "") + if name in seen: + continue + seen.add(name) + result.append(tool_fn) + return result + def get_current_model(self) -> str: """Return the concrete model configured for the job's current difficulty.""" metadata = self.job_def.metadata or {} diff --git a/statek/executors/utils.py b/statek/executors/utils.py index 80ead5d..b0442bf 100644 --- a/statek/executors/utils.py +++ b/statek/executors/utils.py @@ -266,9 +266,7 @@ def _setup_execution_context(job: Job, global_context: dict, local_context: dict # Merge agent's private context if available if job.job_def.agent is not None and job.job_def.agent.context is not None: global_context.update(job.job_def.agent.context) - all_direct_tools = list(job.job_def.agent._tools) - if job.job_def.agent._internal_tools: - all_direct_tools.extend(job.job_def.agent._internal_tools) + all_direct_tools = job.get_resource_tools() # Include system tools from the global registry (not stored on agents) from statek.system import find_tools # pylint: disable=import-outside-toplevel agent_tool_names = {t.__name__ for t in all_direct_tools} @@ -401,17 +399,12 @@ def _is_hidden_tool_call(call_spec: CallSpec, job: Job, policy) -> bool: if call_spec.func_name in policy.blocked_tools: return True - agent = job.job_def.agent - if agent is None: - return False - - for tool_fn in agent._tools: # pylint: disable=protected-access + for tool_fn in job.get_resource_tools(): if (getattr(tool_fn, "__name__", None) == call_spec.func_name and getattr(tool_fn, "tool_hidden", False)): return True - context_fn = agent.context.get(call_spec.func_name) - return bool(callable(context_fn) and getattr(context_fn, "tool_hidden", False)) + return False def _compile_exec_node(node: ast.AST, filename: str): @@ -851,20 +844,18 @@ def _private_print(*args, sep=' ', end='\n', **kwargs): # Re-wrap the original tool with a combined context (global + local) so that # _bind_by_name / find_locals can resolve string arguments to actual objects # (e.g. docstr(what='some_tool') → the actual callable, not the string). - agent = job.job_def.agent original_tool = None - if agent: - for t in agent._tools: # pylint: disable=protected-access - if t.__name__ == call_spec.func_name: - original_tool = t + for tool_fn in job.get_resource_tools(): + if tool_fn.__name__ == call_spec.func_name: + original_tool = tool_fn + break + # Also search system tools from the global registry + if original_tool is None: + from statek.system import find_tools # pylint: disable=import-outside-toplevel + for tool_fn in find_tools("SYSTEM", include_hidden=True): + if tool_fn.__name__ == call_spec.func_name: + original_tool = tool_fn break - # Also search system tools from the global registry - if original_tool is None: - from statek.system import find_tools # pylint: disable=import-outside-toplevel - for t in find_tools("SYSTEM", include_hidden=True): - if t.__name__ == call_spec.func_name: - original_tool = t - break if original_tool is not None: combined_ctx = {**global_context, **local_context} diff --git a/tests/test_extra_resources_runtime.py b/tests/test_extra_resources_runtime.py index dc37d33..f2d2f8b 100644 --- a/tests/test_extra_resources_runtime.py +++ b/tests/test_extra_resources_runtime.py @@ -8,7 +8,50 @@ from statek.agents.agent import Agent, SupervisedAgent from statek.executors.job import Job, JobDef +from statek.executors.utils import exec_step, exec_tool from statek.extra_resources import parse_extra_resources +from statek.llm_api import select_request_tools +from statek.prompt_config import make_system_prompt +from statek.system import tool +from statek.utils import CallSpec, get_current_agent, get_current_job + + +@tool +def receiver_action(value: str, **kwargs) -> str: # pylint: disable=unused-argument + """Return a value from the receiving agent. + + Args: + value: Value to return. + """ + return f"receiver:{value}" + + +@tool +def preference_action(value: str, **kwargs) -> str: # pylint: disable=unused-argument + """Return a preference action with the active execution identity. + + Args: + value: Value to return with the active identity. + """ + agent = get_current_agent() + job = get_current_job() + return f"{agent.role}:{job.job_def.agent.role}:{value}" + + +@tool +def schedule_action(value: str, **kwargs) -> str: # pylint: disable=unused-argument + """Return a schedule action value. + + Args: + value: Value to return. + """ + return f"schedule:{value}" + + +@tool +def _preference_internal(value: str, **kwargs) -> str: # pylint: disable=unused-argument + """Return an internal preference helper value.""" + return f"internal:{value}" @pytest.fixture @@ -20,6 +63,26 @@ def resource_agents(agent_factory: Callable[..., Agent]) -> dict[str, Agent]: } +@pytest.fixture +def resource_tool_job(resource_agents: dict[str, Agent]) -> Job: + """Create a low-difficulty job with tool-bearing configured donors.""" + receiver = resource_agents["receiver"] + receiver._system_prompt = make_system_prompt("Receiver rules only.\n{tools}") + receiver._tools.append(receiver_action) + preferences = resource_agents["preferences_assistant"] + preferences._system_prompt = make_system_prompt("Donor rules must not be merged.") + preferences._tools.append(preference_action) + preferences._internal_tools.append(_preference_internal) + resource_agents["schedule_assistant"]._tools.append(schedule_action) + return Job(job_def=JobDef( + agent=receiver, + metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": "L"}, + extra_resources=parse_extra_resources( + "L:preferences_assistant,M:schedule_assistant,H:sounding_board" + ), + )) + + @pytest.mark.parametrize("difficulty,count", [("L", 2), ("M", 3), ("H", 4)]) def test_resource_agents_inherit_in_declaration_order( resource_agents: dict[str, Agent], difficulty: str, count: int, @@ -187,3 +250,123 @@ def test_resource_agents_do_not_resolve_missing_role_from_another_prefix( with pytest.raises(ValueError, match="EXTRA_RESOURCES.*other_donor"): job.get_resource_agents() + + +def test_resource_tools_expand_request_and_receiver_prompt(resource_tool_job: Job) -> None: + """Active donor tools appear in requests and receiver tool placeholders only.""" + request = resource_tool_job.get_next_request() + names = [tool_fn.__name__ for tool_fn in request["available_tools"]] + + assert names[:2] == ["receiver_action", "preference_action"] + assert "_preference_internal" not in names + assert "schedule_action" not in names + assert "receiver_action" in request["system_prompt"] + assert "preference_action" in request["system_prompt"] + assert "_preference_internal" not in request["system_prompt"] + assert "schedule_action" not in request["system_prompt"] + assert "Donor rules must not be merged." not in request["system_prompt"] + + +def test_resource_tools_are_available_to_formal_provider_scope(resource_tool_job: Job) -> None: + """An inherited action can be explicitly selected for a formal provider tool call.""" + request = resource_tool_job.get_next_request() + selected = select_request_tools( + metadata={"LLM_TOOLS_SCOPE": "ALL"}, + available_tools=request["available_tools"], + ) + names = [tool_fn.__name__ for tool_fn in selected] + + assert "preference_action" in names + assert "_preference_internal" not in names + + +def test_resource_tools_follow_panic_without_mutating_sibling_job( + resource_tool_job: Job, +) -> None: + """Panic adds the next donor's tools only to the escalating job.""" + sibling = Job(job_def=resource_tool_job.job_def) + + resource_tool_job.panic() + + expanded = [tool_fn.__name__ for tool_fn in resource_tool_job.get_resource_tools()] + unchanged = [tool_fn.__name__ for tool_fn in sibling.get_resource_tools()] + assert "schedule_action" in expanded + assert "schedule_action" not in unchanged + assert resource_tool_job.job_def.agent._tools == [receiver_action] + + +@pytest.mark.asyncio +async def test_exec_step_calls_inherited_action_in_receiver_context( + resource_tool_job: Job, +) -> None: + """Python execution can call a donor action while retaining receiver job identity.""" + await exec_step("result = preference_action('saved')", resource_tool_job) + + assert resource_tool_job.py_env.local_state["result"] == "receiver:receiver:saved" + + +@pytest.mark.asyncio +async def test_exec_tool_calls_inherited_action_in_receiver_context( + resource_tool_job: Job, +) -> None: + """Formal tool calls resolve donor actions against the receiving execution context.""" + result, error = await exec_tool( + CallSpec( + id="inherited-action", func_name="preference_action", + args=[], kwargs={"value": "saved"}, + ), + resource_tool_job, + ) + + assert error is None + assert result == "receiver:receiver:saved" + + +@pytest.mark.asyncio +async def test_exec_step_calls_inherited_dynamic_and_internal_tools( + resource_tool_job: Job, resource_agents: dict[str, Agent], +) -> None: + """Named and internal donor tools execute but internal tools stay out of the prompt.""" + donor = resource_agents["preferences_assistant"] + + def dynamic_preference(value: str, **kwargs) -> str: # pylint: disable=unused-argument + """Return a dynamic preference result. + + Args: + value: Value to return with the active agent role. + """ + return f"dynamic:{get_current_agent().role}:{value}" + + donor.context["dynamic_preference"] = dynamic_preference + donor.append_tool("dynamic_preference") + + await exec_step( + "dynamic_result = dynamic_preference('ok')\ninternal_result = _preference_internal('ok')", + resource_tool_job, + ) + + assert resource_tool_job.py_env.local_state["dynamic_result"] == "dynamic:receiver:ok" + assert resource_tool_job.py_env.local_state["internal_result"] == "internal:ok" + prompt = resource_tool_job.system_prompt() + assert "dynamic_preference" in prompt + assert "_preference_internal" not in prompt + + +@pytest.mark.asyncio +async def test_receiver_tool_name_takes_precedence_over_donor_tool( + resource_tool_job: Job, resource_agents: dict[str, Agent], +) -> None: + """Existing first-name-wins semantics preserve the receiver's same-named action.""" + donor = resource_agents["preferences_assistant"] + + def donor_receiver_action(value: str, **kwargs) -> str: # pylint: disable=unused-argument + return f"donor:{value}" + + donor_receiver_action.__name__ = "receiver_action" + donor.context["receiver_action"] = donor_receiver_action + donor.append_tool("receiver_action") + + tools = resource_tool_job.get_resource_tools() + assert [tool_fn.__name__ for tool_fn in tools].count("receiver_action") == 1 + await exec_step("collision_result = receiver_action('ok')", resource_tool_job) + assert resource_tool_job.py_env.local_state["collision_result"] == "receiver:ok" diff --git a/tests/test_llm_api.py b/tests/test_llm_api.py index b04145e..46abe02 100644 --- a/tests/test_llm_api.py +++ b/tests/test_llm_api.py @@ -1674,10 +1674,21 @@ def test_get_next_request_available_tools_is_list(self, job_factory): assert isinstance(request["available_tools"], list) def test_get_next_request_available_tools_matches_agent(self, job_factory, agent): - """available_tools in the request matches the agent's all_tools.""" + """available_tools in the request matches the agent's LLM-visible tools.""" job = job_factory() request = job.get_next_request() - assert set(request["available_tools"]) == set(job.job_def.agent.all_tools) + assert set(request["available_tools"]) == set(job.job_def.agent.tools) + + def test_get_next_request_skips_unresolved_dynamic_tool(self, job_factory, agent): + """A volatile named tool missing from context does not block a normal request.""" + agent.append_tool("not_initialized") + job = job_factory() + + request = job.get_next_request() + + assert "not_initialized" not in { + tool_fn.__name__ for tool_fn in request["available_tools"] + } # --------------------------------------------------------------------------- From 62c5af55e626ef0981a112cb58ba2558724a8cf3 Mon Sep 17 00:00:00 2001 From: Kiki McArron Date: Wed, 16 Sep 2026 11:48:06 +0200 Subject: [PATCH 6/9] Implement & test the 'Capacities Expansion' functionality - part 3 --- statek/agents/agent.py | 32 ++++++--- statek/agents/list_of_documents.py | 16 +++-- statek/agents/list_of_examples.py | 77 ++++++++++++++++----- statek/document.py | 24 +++++-- statek/executors/job.py | 40 +++++++++-- tests/test_extra_resources_runtime.py | 37 +++++++++- tests/test_find_document.py | 11 +++ tests/test_job.py | 58 ++++++++++++++++ tests/test_list_of_documents.py | 26 ++++++- tests/test_show_document.py | 14 ++++ tests/test_show_example.py | 97 ++++++++++++++++++++++++++- 11 files changed, 387 insertions(+), 45 deletions(-) diff --git a/statek/agents/agent.py b/statek/agents/agent.py index a02fbf8..a3506bb 100644 --- a/statek/agents/agent.py +++ b/statek/agents/agent.py @@ -17,7 +17,7 @@ import re from typing import Any, List, Callable, Dict, Optional, Sequence, Union import dbzero as db0 -from statek.utils import block_comment, get_current_agent, _get_class_name +from statek.utils import block_comment, get_current_agent, get_current_job, _get_class_name from statek.system import tool from statek.docstring import parse_tool_docstring, format_docstring from statek.utils import CodeBlock @@ -51,9 +51,7 @@ def list_of_examples(start_index: int = 0, limit: int = 10, **kwargs): # pylint limit: Maximum number of examples to show (default: 10). """ from statek.agents.list_of_examples import list_of_examples as _impl # pylint: disable=import-outside-toplevel - agent = get_current_agent() - agent_name = agent.role if agent else None - _impl(agent_name, start_index, limit) + _impl(_get_resource_agent_names(), start_index, limit) @tool(system=True) @@ -70,10 +68,11 @@ def list_of_documents(topic=None, start_index: int = 0, limit: int = 25, **kwarg """ from statek.agents.list_of_documents import list_of_documents as _impl # pylint: disable=import-outside-toplevel from statek.settings import get_statek_settings # pylint: disable=import-outside-toplevel - agent = get_current_agent() - agent_name = agent.role if agent else None documents_dir = get_statek_settings().documents_dir - _impl(agent_name, documents_dir, topic=topic, start_index=start_index, limit=limit) + _impl( + _get_resource_agent_names(), documents_dir, + topic=topic, start_index=start_index, limit=limit, + ) @tool(system=True) @@ -90,10 +89,11 @@ def show_document(key, topic=None, start_from: int = 0, limit: int = 50, **kwarg """ from statek.agents.list_of_documents import show_document as _impl # pylint: disable=import-outside-toplevel from statek.settings import get_statek_settings # pylint: disable=import-outside-toplevel - agent = get_current_agent() - agent_name = agent.role if agent else None documents_dir = get_statek_settings().documents_dir - _impl(agent_name, documents_dir, key=key, topic=topic, start_from=start_from, limit=limit) + _impl( + _get_resource_agent_names(), documents_dir, + key=key, topic=topic, start_from=start_from, limit=limit, + ) @tool(system=True) @@ -107,8 +107,18 @@ def show_example(example_id: Optional[int] = None, **kwargs): # pylint: disable If not provided, uses default_example_id from the local context. """ from statek.agents.list_of_examples import show_example as _impl # pylint: disable=import-outside-toplevel + _impl(_get_resource_agent_names(), example_id) + + +def _get_resource_agent_names() -> List[str]: + """Return current receiver and active donor roles for resource lookup.""" agent = get_current_agent() - _impl(agent.role if agent else None, example_id) + if agent is None: + return [] + job = get_current_job() + if job is None: + return [agent.role] + return [resource_agent.role for resource_agent in job.get_resource_agents()] @db0.memo diff --git a/statek/agents/list_of_documents.py b/statek/agents/list_of_documents.py index 7e4990d..c50e8f0 100644 --- a/statek/agents/list_of_documents.py +++ b/statek/agents/list_of_documents.py @@ -22,9 +22,14 @@ # pylint: disable=unused-argument -from typing import Optional, Union - -from statek.document import load_documents, find_topic, find_document +from typing import Optional, Sequence, Union + +from statek.document import ( + find_document, + find_topic, + load_documents, + validate_unique_document_ids, +) from statek.settings import get_statek_settings from statek.system import tool from statek.utils import perm_ctx_set, perm_ctx_get @@ -36,7 +41,7 @@ def _get_documents_dir() -> Optional[str]: @tool(system=True) -def list_of_documents(agent_name: str, documents_dir: str, +def list_of_documents(agent_name: str | Sequence[str], documents_dir: str, topic: Optional[Union[int, str]] = None, start_index: int = 0, limit: int = 25, **kwargs): # pylint: disable=unused-argument """Lists available topics or documents within a topic. @@ -86,6 +91,7 @@ def _list_documents(topic_key, agent_name, all_topics, start_index, limit): perm_ctx_set(last_topic_id=matched.ord_no) docs = [d for d in matched.documents if d.match_audience(agent_name)] + validate_unique_document_ids(docs) total = len(docs) print(f"# Document ID: Document name ({total} total)") for d in docs[start_index:start_index + limit]: @@ -111,7 +117,7 @@ def _resolve_topic(topic_key, agent_name, all_topics): @tool(system=True) -def show_document(agent_name: str, documents_dir: str, # pylint: disable=too-many-positional-arguments,too-many-arguments +def show_document(agent_name: str | Sequence[str], documents_dir: str, # pylint: disable=too-many-positional-arguments,too-many-arguments key: Union[int, str], topic: Optional[Union[int, str]] = None, start_from: int = 0, limit: int = 50, **kwargs): """Shows the contents of a specific document. diff --git a/statek/agents/list_of_examples.py b/statek/agents/list_of_examples.py index 5679cc8..0bf2aa2 100644 --- a/statek/agents/list_of_examples.py +++ b/statek/agents/list_of_examples.py @@ -24,9 +24,10 @@ """ import os -from typing import Optional +from dataclasses import dataclass +from typing import Optional, Sequence -from statek.executors.example import load_examples, format_example +from statek.executors.example import Example, load_examples, format_example from statek.settings import get_statek_settings from statek.system import tool from statek.task_difficulty import TaskDifficulty @@ -38,6 +39,39 @@ def _get_examples_dir() -> Optional[str]: return get_statek_settings().examples_dir +@dataclass(frozen=True) +class SourcedExample: + """An example paired with its configured agent role and local index.""" + + agent_name: str + example_id: int + example: Example + + +def _agent_names(agent_name: str | Sequence[str]) -> list[str]: + """Normalize one or more role names while preserving their first occurrence.""" + names = [agent_name] if isinstance(agent_name, str) else list(agent_name) + return list(dict.fromkeys(name for name in names if name)) + + +def get_sourced_examples(agent_name: str | Sequence[str]) -> list[SourcedExample]: + """Return a fresh receiver-first combined view of examples and their sources.""" + examples_dir = _get_examples_dir() + if not examples_dir: + return [] + + result = [] + for role in _agent_names(agent_name): + path = os.path.join(examples_dir, role) + if not os.path.isdir(path): + continue + result.extend( + SourcedExample(role, example_id, example) + for example_id, example in enumerate(load_examples(path)) + ) + return result + + def _get_example(agent_name: str, example_id: int, logs: Optional[list[str]] = None): """Return an example by agent/id, optionally collecting user-facing lookup messages.""" examples_dir = _get_examples_dir() @@ -90,7 +124,9 @@ def get_example_difficulty(agent_name: str, example_id: int) -> Optional[TaskDif @tool(system=True) -def list_of_examples(agent_name: str, start_index: int = 0, limit: int = 10, **kwargs): # pylint: disable=unused-argument +def list_of_examples( # pylint: disable=unused-argument + agent_name: str | Sequence[str], start_index: int = 0, limit: int = 10, **kwargs, +): """Lists available examples for a given agent. Results are printed as a numbered list (index: name). @@ -107,19 +143,21 @@ def list_of_examples(agent_name: str, start_index: int = 0, limit: int = 10, **k list_of_examples(agent_name="coordinator") list_of_examples(agent_name="information_retriever", start_index=10, limit=5) """ - names = get_example_names(agent_name) - if not names: + sourced_examples = get_sourced_examples(agent_name) + if not sourced_examples: print("# No examples found") return - total = len(names) + total = len(sourced_examples) print(f"# Example ID: Example name ({total} total)") - for i, name in enumerate(names[start_index:start_index + limit]): + for i, sourced in enumerate(sourced_examples[start_index:start_index + limit]): idx = start_index + i - print(f"{idx}: {name}") + print(f"{idx}: {sourced.example.example_metadata.get('name', '')}") @tool(system=True) -def show_example(agent_name: str, example_id: Optional[int] = None, **kwargs): # pylint: disable=unused-argument +def show_example( # pylint: disable=unused-argument + agent_name: str | Sequence[str], example_id: Optional[int] = None, **kwargs, +): """Shows a specific example for a given agent. Prints the example content formatted using the current chat style setting. @@ -146,17 +184,26 @@ def show_example(agent_name: str, example_id: Optional[int] = None, **kwargs): except (TypeError, ValueError): print("# Example not found") return - lookup_logs = [] - example = _get_example(agent_name, example_id, logs=lookup_logs) - if example is None: - for message in lookup_logs: - print(message) + sourced_examples = get_sourced_examples(agent_name) + if example_id < 0 or example_id >= len(sourced_examples): + if sourced_examples: + print(f"# Example {example_id} not found (total: {len(sourced_examples)})") + else: + print("# No examples found") return + sourced = sourced_examples[example_id] + example = sourced.example settings = get_statek_settings() style = settings.examples_style or settings.chat_style name = example.example_metadata.get("name", "") try: - perm_ctx_set(last_example_id=example_id) + perm_ctx_set( + last_example_id=example_id, + last_example_source={ + "agent_name": sourced.agent_name, + "example_id": sourced.example_id, + }, + ) except RuntimeError: pass if settings.xml_box_example: diff --git a/statek/document.py b/statek/document.py index 6c92024..9a461c9 100644 --- a/statek/document.py +++ b/statek/document.py @@ -18,7 +18,7 @@ from dataclasses import dataclass, field from difflib import SequenceMatcher from functools import lru_cache -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional, Sequence, Union _REQUIRED_KEYS = ("ord_no", "topic", "title") @@ -32,7 +32,7 @@ class Document: """Body of text - organized by lines""" body: List[str] = field(default_factory=list) - def match_audience(self, agent_name: str) -> bool: + def match_audience(self, agent_name: str | Sequence[str]) -> bool: """Check if a specific agent is the target audience. If "audience" is not specified, the document targets all agents. @@ -40,7 +40,8 @@ def match_audience(self, agent_name: str) -> bool: audience = self.document_metadata.get("audience") if audience is None: return True - return agent_name in audience + agent_names = (agent_name,) if isinstance(agent_name, str) else agent_name + return any(name in audience for name in agent_names) @dataclass @@ -50,7 +51,7 @@ class Topic: topic: str documents: List[Document] = field(default_factory=list) - def count(self, agent_name: str) -> int: + def count(self, agent_name: str | Sequence[str]) -> int: """Count documents accessible to a given agent. Args: @@ -62,6 +63,16 @@ def count(self, agent_name: str) -> int: return sum(1 for doc in self.documents if doc.match_audience(agent_name)) +def validate_unique_document_ids(documents: Sequence[Document]) -> None: + """Reject an accessible document view containing an ambiguous numeric ID.""" + seen: set[int] = set() + for document in documents: + document_id = document.document_metadata["ord_no"] + if document_id in seen: + raise ValueError(f"Ambiguous document ID '{document_id}'") + seen.add(document_id) + + def parse_document(document: str) -> Document: """Parse a document string into a Document instance. @@ -147,7 +158,7 @@ def load_documents(path: str) -> List[Topic]: return list(topics_map.values()) -def find_topic(key: Union[int, str], agent_name: str, +def find_topic(key: Union[int, str], agent_name: str | Sequence[str], all_topics: List[Topic]) -> Optional[Topic]: """Find a topic by index, name, or name fragment. @@ -192,7 +203,7 @@ def find_topic(key: Union[int, str], agent_name: str, raise ValueError(f"Ambiguous topic '{key}', matches: {listing}") -def find_document(key: Union[int, str], agent_name: str, +def find_document(key: Union[int, str], agent_name: str | Sequence[str], topic: Topic) -> Optional[Document]: """Find a document within a topic by index, title, or title fragment. @@ -212,6 +223,7 @@ def find_document(key: Union[int, str], agent_name: str, document titles and indexes. """ accessible = [d for d in topic.documents if d.match_audience(agent_name)] + validate_unique_document_ids(accessible) if isinstance(key, int): for doc in accessible: diff --git a/statek/executors/job.py b/statek/executors/job.py index a9a0db1..780efc8 100644 --- a/statek/executors/job.py +++ b/statek/executors/job.py @@ -1565,6 +1565,36 @@ def _get_last_example_id(self) -> Optional[int]: if "default_example_id" in local_state: return local_state["default_example_id"] return None + + def _get_last_example_source(self) -> Optional[Tuple[str, int]]: + """Return a validated source role/local ID for the last shown example.""" + perm_ctx = self.py_env.perm_ctx or {} + source = perm_ctx.get("last_example_source") + if source is None and get_current_job() is self: + source = perm_ctx_get("last_example_source", None) + if not hasattr(source, "get"): + return None + agent_name = source.get("agent_name") + example_id = source.get("example_id") + if not isinstance(agent_name, str) or not isinstance(example_id, int): + return None + + baseline = self.__last_difficulty or _get_static_task_difficulty(self.job_def.metadata) + level = tuple(TaskDifficulty.values()).index(baseline) + allowed_roles = {self.job_def.agent.role} + resources = getattr(self.job_def, "extra_resources", None) + if resources: + allowed_roles.update( + role for index in range(level + 1) for role in resources[index] or () + ) + if agent_name not in allowed_roles: + return None + + from statek.agents.agent import Agent # pylint: disable=import-outside-toplevel + prefix = db0.get_prefix_of(self.job_def.agent).name + if not any(agent.role == agent_name for agent in db0.find(Agent, prefix=prefix)): + return None + return agent_name, example_id def get_current_difficulty(self) -> TaskDifficulty: """Return the current task difficulty for this job. @@ -1575,10 +1605,12 @@ def get_current_difficulty(self) -> TaskDifficulty: """ last_example_id = self._get_last_example_id() - if last_example_id is not None: - agent = self.job_def.agent if self.job_def is not None else None - agent_name = agent.role if agent is not None else None - difficulty = _get_example_difficulty_for_job(agent_name, last_example_id) + if last_example_id is not None: + source = self._get_last_example_source() + if source is None: + agent = self.job_def.agent if self.job_def is not None else None + source = (agent.role if agent is not None else None, last_example_id) + difficulty = _get_example_difficulty_for_job(*source) # update the dynamically resolved difficulty if difficulty is not None: if self.__last_difficulty is None: diff --git a/tests/test_extra_resources_runtime.py b/tests/test_extra_resources_runtime.py index f2d2f8b..d809aae 100644 --- a/tests/test_extra_resources_runtime.py +++ b/tests/test_extra_resources_runtime.py @@ -2,11 +2,19 @@ # pylint: disable=protected-access,no-member,redefined-outer-name from collections.abc import Callable +from unittest.mock import patch import dbzero as db0 import pytest -from statek.agents.agent import Agent, SupervisedAgent +from statek.agents.agent import ( + Agent, + SupervisedAgent, + list_of_documents as list_documents_wrapper, + list_of_examples as list_examples_wrapper, + show_document as show_document_wrapper, + show_example as show_example_wrapper, +) from statek.executors.job import Job, JobDef from statek.executors.utils import exec_step, exec_tool from statek.extra_resources import parse_extra_resources @@ -370,3 +378,30 @@ def donor_receiver_action(value: str, **kwargs) -> str: # pylint: disable=unuse assert [tool_fn.__name__ for tool_fn in tools].count("receiver_action") == 1 await exec_step("collision_result = receiver_action('ok')", resource_tool_job) assert resource_tool_job.py_env.local_state["collision_result"] == "receiver:ok" + + +def test_resource_wrappers_pass_active_agent_roles(resource_tool_job: Job) -> None: + """Example and document wrappers pass receiver-first active roles to their helpers.""" + expected_roles = ["receiver", "preferences_assistant"] + with ( + patch("statek.agents.agent.get_current_agent", return_value=resource_tool_job.agent), + patch("statek.agents.agent.get_current_job", return_value=resource_tool_job), + patch("statek.agents.list_of_examples.list_of_examples") as list_examples, + patch("statek.agents.list_of_examples.show_example") as show_example, + patch("statek.agents.list_of_documents.list_of_documents") as list_documents, + patch("statek.agents.list_of_documents.show_document") as show_document, + ): + list_examples_wrapper(start_index=2, limit=3) + show_example_wrapper(example_id=4) + list_documents_wrapper(topic="Guide", start_index=5, limit=6) + show_document_wrapper(key=7, topic="Guide", start_from=8, limit=9) + + list_examples.assert_called_once_with(expected_roles, 2, 3) + show_example.assert_called_once_with(expected_roles, 4) + list_documents.assert_called_once_with( + expected_roles, list_documents.call_args.args[1], topic="Guide", start_index=5, limit=6, + ) + show_document.assert_called_once_with( + expected_roles, show_document.call_args.args[1], + key=7, topic="Guide", start_from=8, limit=9, + ) diff --git a/tests/test_find_document.py b/tests/test_find_document.py index 776fa12..dd98f11 100644 --- a/tests/test_find_document.py +++ b/tests/test_find_document.py @@ -48,3 +48,14 @@ def test_find_exact_title_takes_priority(): def test_find_excludes_inaccessible_docs(): # agent_c can only see "Overview" (no audience restriction) assert find_document("Details", "agent_c", TOPIC) is None + + +def test_find_rejects_duplicate_ids_exposed_by_combined_audiences(): + """Combined resource audiences cannot silently select one of two identical IDs.""" + topic = Topic(ord_no=0, topic="T", documents=[ + _doc(1, "Agent A", audience=["agent_a"]), + _doc(1, "Agent B", audience=["agent_b"]), + ]) + + with pytest.raises(ValueError, match="Ambiguous document ID '1'"): + find_document(1, ["agent_a", "agent_b"], topic) diff --git a/tests/test_job.py b/tests/test_job.py index 565c23e..c604e32 100644 --- a/tests/test_job.py +++ b/tests/test_job.py @@ -860,6 +860,64 @@ def test_get_current_difficulty_uses_example_difficulty(job_def_factory): assert job._Job__last_difficulty == TaskDifficulty.high # pylint: disable=protected-access +def test_get_current_difficulty_uses_source_aware_donor_example( + job_def_factory, agent_factory, +): + """A persisted example source resolves difficulty against its donor's local ID.""" + donor = agent_factory(role="donor") + job = Job( + job_def=job_def_factory(extra_resources=(["donor"], None, None)), + job_status=JobStatus.READY, # pylint: disable=no-member + ) + job.py_env.local_state["_PERM_CTX"] = { + "last_example_id": 7, + "last_example_source": {"agent_name": "donor", "example_id": 1}, + } + + with patch( + "statek.executors.job._get_example_difficulty_for_job", + return_value=TaskDifficulty.high, + ) as mock_get_example_difficulty: + assert job.get_current_difficulty() == TaskDifficulty.high + + mock_get_example_difficulty.assert_called_once_with("donor", 1) + assert donor in job.get_resource_agents() + + +def test_get_current_difficulty_malformed_source_uses_legacy_receiver_id(job_def_factory): + """Malformed persisted source data falls back to the legacy receiver-local ID.""" + job = Job(job_def=job_def_factory(), job_status=JobStatus.READY) # pylint: disable=no-member + job.py_env.local_state["_PERM_CTX"] = { + "last_example_id": 3, + "last_example_source": {"agent_name": "donor", "example_id": "invalid"}, + } + + with patch( + "statek.executors.job._get_example_difficulty_for_job", + return_value=TaskDifficulty.medium, + ) as mock_get_example_difficulty: + assert job.get_current_difficulty() == TaskDifficulty.medium + + mock_get_example_difficulty.assert_called_once_with("test", 3) + + +def test_get_current_difficulty_unconfigured_source_uses_legacy_receiver_id(job_def_factory): + """A forged source outside configured resources cannot influence job difficulty.""" + job = Job(job_def=job_def_factory(), job_status=JobStatus.READY) # pylint: disable=no-member + job.py_env.local_state["_PERM_CTX"] = { + "last_example_id": 3, + "last_example_source": {"agent_name": "../foreign", "example_id": 0}, + } + + with patch( + "statek.executors.job._get_example_difficulty_for_job", + return_value=TaskDifficulty.medium, + ) as mock_get_example_difficulty: + assert job.get_current_difficulty() == TaskDifficulty.medium + + mock_get_example_difficulty.assert_called_once_with("test", 3) + + def test_get_current_difficulty_uses_registered_job_perm_ctx_example_id(job_def_factory): """The registered job's PyEnv context supplies the last example ID.""" job_def = job_def_factory( diff --git a/tests/test_list_of_documents.py b/tests/test_list_of_documents.py index 6af3bf8..5d17a4c 100644 --- a/tests/test_list_of_documents.py +++ b/tests/test_list_of_documents.py @@ -30,6 +30,13 @@ # title: Tool Guide Body.""" +DOC_PREFS_2 = """\ +# ord_no: 2 +# topic: Preferences +# audience: agent_b +# title: Donor Details +Body.""" + @pytest.fixture(autouse=True) def clear_cache(): @@ -42,6 +49,7 @@ def clear_cache(): def docs_dir(temp_dir): _write(temp_dir, "prefs0.txt", DOC_PREFS_0) _write(temp_dir, "prefs1.txt", DOC_PREFS_1) + _write(temp_dir, "prefs2.txt", DOC_PREFS_2) _write(temp_dir, "tools0.txt", DOC_TOOLS_0) return temp_dir @@ -70,7 +78,23 @@ def test_list_documents_filters_by_audience(docs_dir, capsys): run_with_statek_job(job, lambda: _impl("agent_b", docs_dir, topic="Preferences")) out = capsys.readouterr().out assert "Overview" in out - assert "Details" not in out + assert "1: Details" not in out + + +def test_list_documents_unions_resource_audiences_without_duplicate_public_docs( + docs_dir, capsys, +): + """Multiple active roles see their union and each shared document only once.""" + job = StatekContextJob() + run_with_statek_job( + job, + lambda: _impl(["agent_a", "agent_b"], docs_dir, topic="Preferences"), + ) + out = capsys.readouterr().out + + assert out.count("Overview") == 1 + assert "Details" in out + assert "Donor Details" in out def test_list_documents_sets_last_topic_id(docs_dir): diff --git a/tests/test_show_document.py b/tests/test_show_document.py index 0a4e9c2..fe0ac8d 100644 --- a/tests/test_show_document.py +++ b/tests/test_show_document.py @@ -106,6 +106,20 @@ def test_show_document_fuzzy_match_below_threshold(docs_dir, capsys): assert "not found" in out +def test_show_document_accepts_any_active_resource_audience(docs_dir, capsys): + """A donor-restricted document is visible through the combined audience.""" + job = StatekContextJob() + run_with_statek_job( + job, + lambda: show_document( + ["receiver", "agent_a"], docs_dir, key="Secret Guide", topic="Guide", + ), + ) + out = capsys.readouterr().out + + assert "Restricted content." in out + + def _write(directory, filename, content): with open(os.path.join(directory, filename), "w", encoding="utf-8") as f: f.write(content) diff --git a/tests/test_show_example.py b/tests/test_show_example.py index cb3e66f..f86368f 100644 --- a/tests/test_show_example.py +++ b/tests/test_show_example.py @@ -9,9 +9,15 @@ from tests.conftest import StatekContextJob, run_with_statek_job from statek.agents.agent import Agent -from statek.agents.list_of_examples import get_example_difficulty, show_example +from statek.agents.list_of_examples import ( + get_example_difficulty, + list_of_examples, + show_example, +) from statek.executors.example import load_examples -from statek.executors.job import TaskDifficulty +from statek.executors.job import Job, JobDef, TaskDifficulty +from statek.extra_resources import parse_extra_resources +from statek.prompt_config import make_system_prompt from statek.settings import StatekSettings, ChatStyle, get_statek_settings @@ -399,3 +405,90 @@ def test_get_example_difficulty_returns_none_when_missing(examples_dir): settings = _settings(examples_dir=examples_dir) with patch("statek.agents.list_of_examples.get_statek_settings", return_value=settings): assert get_example_difficulty("myagent", 0) is None + + +@pytest.fixture +def expanded_examples_dir(temp_dir): + """Create receiver and donor examples with overlapping local indexes.""" + examples = { + "receiver": [("Receiver zero", "low"), ("Receiver one", "medium")], + "preferences_assistant": [("Preference zero", "high")], + "schedule_assistant": [("Schedule zero", "medium")], + } + for role, role_examples in examples.items(): + role_dir = os.path.join(temp_dir, role) + os.makedirs(role_dir, exist_ok=True) + for index, (name, difficulty) in enumerate(role_examples): + content = ( + f"# seq_id: {index}\n# name: {name}\n# difficulty: {difficulty}\n" + f"```python\nprint({name!r})\n```\n" + ) + path = os.path.join(role_dir, f"example-{index}.md") + with open(path, "w", encoding="utf-8") as file: + file.write(content) + return temp_dir + + +@pytest.mark.usefixtures("db0_fixture") +def test_list_of_examples_combines_roles_with_stable_public_ids( + capsys, expanded_examples_dir, +): + """Combined examples keep receiver-first IDs and append newly active donors.""" + settings = _settings(examples_dir=expanded_examples_dir) + with patch("statek.agents.list_of_examples.get_statek_settings", return_value=settings): + list_of_examples(["receiver", "preferences_assistant"], start_index=1, limit=2) + first = capsys.readouterr().out + + with patch("statek.agents.list_of_examples.get_statek_settings", return_value=settings): + list_of_examples( + ["receiver", "preferences_assistant", "schedule_assistant"], + start_index=0, + limit=10, + ) + expanded = capsys.readouterr().out + + assert "1: Receiver one" in first + assert "2: Preference zero" in first + assert "0: Receiver zero" in expanded + assert "1: Receiver one" in expanded + assert "2: Preference zero" in expanded + assert "3: Schedule zero" in expanded + + +@pytest.mark.usefixtures("db0_fixture") +def test_show_example_records_donor_source_and_raises_job_difficulty( + capsys, expanded_examples_dir, +): + """A donor example stores its local source and drives difficulty from that donor.""" + receiver = Agent( + role="receiver", _system_prompt=make_system_prompt("Receiver"), + _tools=[], _metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": "L"}, + ) + donor = Agent( + role="preferences_assistant", _system_prompt=make_system_prompt("Preferences"), + _tools=[], _metadata={"MODEL": "test-model"}, + ) + job = Job(job_def=JobDef( + agent=receiver, + metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": "L"}, + extra_resources=parse_extra_resources("L:preferences_assistant"), + )) + settings = _settings( + examples_dir=expanded_examples_dir, + chat_style=ChatStyle.CONSOLE, # pylint: disable=no-member + ) + + with patch("statek.agents.list_of_examples.get_statek_settings", return_value=settings): + run_with_statek_job( + job, + lambda: show_example(["receiver", "preferences_assistant"], example_id=2), + ) + + capsys.readouterr() + assert job.py_env.local_state["_PERM_CTX"]["last_example_id"] == 2 + assert job.py_env.local_state["_PERM_CTX"]["last_example_source"] == { + "agent_name": "preferences_assistant", "example_id": 0, + } + with patch("statek.agents.list_of_examples.get_statek_settings", return_value=settings): + assert job.get_current_difficulty() == TaskDifficulty.high # pylint: disable=no-member + assert donor in job.get_resource_agents() From 2cf5032b8ce3a3365926fb9c7e88d3847e847fd0 Mon Sep 17 00:00:00 2001 From: Kiki McArron Date: Wed, 16 Sep 2026 12:55:59 +0200 Subject: [PATCH 7/9] Implement & test the 'Capacities Expansion' functionality - part 4 --- statek/executors/chat_log_item.py | 3 + statek/executors/job.py | 86 +++++++++++++++++------ statek/executors/utils.py | 19 ++++- tests/test_extra_resources_integration.py | 68 +++++++++++++++++- tests/test_extra_resources_runtime.py | 54 +++++++++++++- tests/test_job.py | 79 +++++++++++++++++++++ tests/test_run_job_step.py | 29 ++++++++ 7 files changed, 312 insertions(+), 26 deletions(-) diff --git a/statek/executors/chat_log_item.py b/statek/executors/chat_log_item.py index 797c2e2..327184b 100644 --- a/statek/executors/chat_log_item.py +++ b/statek/executors/chat_log_item.py @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, Union import dbzero as db0 +from statek.task_difficulty import TaskDifficulty from statek.utils import CodeBlock if TYPE_CHECKING: @@ -90,6 +91,8 @@ class LLM_LogItem(ChatLogItem): llm_resp: Union[str, CodeBlock] = None # Opaque provider continuation material replayed only by its matching formatter. llm_reasoning_payload: Optional[Any] = None + # Difficulty used to build this request; None identifies legacy log items. + request_difficulty: Optional[TaskDifficulty] = None @db0.memo(no_default_tags=True) diff --git a/statek/executors/job.py b/statek/executors/job.py index 780efc8..c992263 100644 --- a/statek/executors/job.py +++ b/statek/executors/job.py @@ -771,10 +771,16 @@ def system_prompt(self, difficulty: Optional[TaskDifficulty] = None) -> str: """Return the agent system prompt formatted for this job's current difficulty.""" if self.job_def is None or self.job_def.agent is None: return "" + resolved_difficulty = ( + difficulty if difficulty is not None else self.get_current_difficulty() + ) return self.job_def.agent.system_prompt( - task_difficulty=difficulty or self.get_current_difficulty(), + task_difficulty=resolved_difficulty, job_params=self.job_def.job_params, - available_tools=self.get_resource_tools(include_internal=False), + available_tools=self.get_resource_tools( + include_internal=False, + difficulty=resolved_difficulty, + ), ) def _ensure_ext_ref_storage(self): @@ -1395,12 +1401,16 @@ def _build_request_data( chat_log: Optional[Sequence[Union[str, ChatLogItem, UserLogItem]]] = None, console: Optional[Sequence[str]] = None, push_log: Optional[Dict[int, Union[str, List[str]]]] = None, + difficulty: Optional[TaskDifficulty] = None, ) -> Dict[str, Any]: """Build request params for the provided append-only job state.""" metadata = dict(self.job_def.metadata or {}) - model = LLM_API.require_model(self.get_current_model()) + resolved_difficulty = ( + difficulty if difficulty is not None else self.get_current_difficulty() + ) + model = LLM_API.require_model(self.get_current_model(difficulty=resolved_difficulty)) temperature = LLM_API.parse_temperature(metadata.get("TEMPERATURE")) - system_prompt = self.system_prompt() + system_prompt = self.system_prompt(difficulty=resolved_difficulty) if ( self.job_def.locale is not None @@ -1419,7 +1429,10 @@ def _build_request_data( "system_prompt": system_prompt, "model": model, "metadata": metadata, - "available_tools": self.get_resource_tools(include_internal=False), + "available_tools": self.get_resource_tools( + include_internal=False, + difficulty=resolved_difficulty, + ), "provider_config": self.job_def.provider_config, } if temperature is not None: @@ -1430,7 +1443,10 @@ def _build_request_data( return request_params - def get_next_request(self) -> Dict[str, Any]: + def get_next_request( + self, + difficulty: Optional[TaskDifficulty] = None, + ) -> Dict[str, Any]: """Build the parameter dict consumed by ``LLM_API.process_request``. ``chat_history`` is the lazy generator returned by @@ -1443,7 +1459,7 @@ def get_next_request(self) -> Dict[str, Any]: ``available_tools`` (list), and optionally ``chat_style``, ``temperature``, and ``provider_config``. """ - return self._build_request_data() + return self._build_request_data(difficulty=difficulty) def get_request_data(self, turn_num: int) -> Dict[str, Any]: """Reconstruct the request params used for a historical LLM turn.""" @@ -1451,6 +1467,7 @@ def get_request_data(self, turn_num: int) -> Dict[str, Any]: raise IndexError(f"turn_num out of range: {turn_num}") target_idx = None + target_item = None console_limit = 0 llm_turn_num = 0 for idx, item in enumerate(self.chat_log): @@ -1458,6 +1475,7 @@ def get_request_data(self, turn_num: int) -> Dict[str, Any]: continue if llm_turn_num == turn_num: target_idx = idx + target_item = item console_limit = item.console_pos break llm_turn_num += 1 @@ -1477,9 +1495,13 @@ def get_request_data(self, turn_num: int) -> Dict[str, Any]: chat_log=history_chat_log, console=history_console, push_log=history_push_log, + difficulty=getattr(target_item, "request_difficulty", None), ) - def get_resource_agents(self) -> List["Agent"]: + def get_resource_agents( + self, + difficulty: Optional[TaskDifficulty] = None, + ) -> List["Agent"]: """Return the receiver and configured donors active at this job's difficulty. Donors are resolved by role in the receiver's prefix, independently of @@ -1495,7 +1517,10 @@ def get_resource_agents(self) -> List["Agent"]: if not resources: return [receiver] - level = tuple(TaskDifficulty.values()).index(self.get_current_difficulty()) + resolved_difficulty = ( + difficulty if difficulty is not None else self.get_current_difficulty() + ) + level = tuple(TaskDifficulty.values()).index(resolved_difficulty) roles = dict.fromkeys( role for index in range(level + 1) for role in resources[index] or () if role != receiver.role @@ -1518,7 +1543,11 @@ def get_resource_agents(self) -> List["Agent"]: result.append(configured[role]) return result - def get_resource_tools(self, include_internal: bool = True) -> List[Callable]: + def get_resource_tools( + self, + include_internal: bool = True, + difficulty: Optional[TaskDifficulty] = None, + ) -> List[Callable]: """Return active receiver and donor tools, keeping the first tool per name. Receiver tools precede donor tools, so a donor cannot replace an @@ -1527,7 +1556,7 @@ def get_resource_tools(self, include_internal: bool = True) -> List[Callable]: """ result = [] seen: set[str] = set() - for agent in self.get_resource_agents(): + for agent in self.get_resource_agents(difficulty=difficulty): tools = agent.all_tools if include_internal else agent.tools for tool_fn in tools: name = getattr(tool_fn, "__name__", "") @@ -1537,15 +1566,17 @@ def get_resource_tools(self, include_internal: bool = True) -> List[Callable]: result.append(tool_fn) return result - def get_current_model(self) -> str: + def get_current_model(self, difficulty: Optional[TaskDifficulty] = None) -> str: """Return the concrete model configured for the job's current difficulty.""" metadata = self.job_def.metadata or {} model_config = metadata.get("MODEL") if not _is_model_mapping(model_config): return str(model_config) if model_config is not None else None - difficulty = self.get_current_difficulty() - return model_config[difficulty] + resolved_difficulty = ( + difficulty if difficulty is not None else self.get_current_difficulty() + ) + return model_config[resolved_difficulty] def _get_last_example_id(self) -> Optional[int]: perm_ctx = self.py_env.perm_ctx or {} @@ -1579,7 +1610,8 @@ def _get_last_example_source(self) -> Optional[Tuple[str, int]]: if not isinstance(agent_name, str) or not isinstance(example_id, int): return None - baseline = self.__last_difficulty or _get_static_task_difficulty(self.job_def.metadata) + baseline = getattr(self, "_Job__last_difficulty", None) + baseline = baseline or _get_static_task_difficulty(self.job_def.metadata) level = tuple(TaskDifficulty.values()).index(baseline) allowed_roles = {self.job_def.agent.role} resources = getattr(self.job_def, "extra_resources", None) @@ -1603,6 +1635,7 @@ def get_current_difficulty(self) -> TaskDifficulty: metadata, then settings. Example difficulty is dynamic: it updates the stored value only when it would keep or raise the current difficulty. """ + last_difficulty = getattr(self, "_Job__last_difficulty", None) last_example_id = self._get_last_example_id() if last_example_id is not None: @@ -1613,16 +1646,18 @@ def get_current_difficulty(self) -> TaskDifficulty: difficulty = _get_example_difficulty_for_job(*source) # update the dynamically resolved difficulty if difficulty is not None: - if self.__last_difficulty is None: + if last_difficulty is None: self.__last_difficulty = difficulty + last_difficulty = difficulty else: - max_difficulty = max_task_difficulty(self.__last_difficulty, difficulty) + max_difficulty = max_task_difficulty(last_difficulty, difficulty) # only update if changed - if self.__last_difficulty != max_difficulty: - self.__last_difficulty = max_difficulty + if last_difficulty != max_difficulty: + self.__last_difficulty = max_difficulty + last_difficulty = max_difficulty - if self.__last_difficulty is not None: - return self.__last_difficulty + if last_difficulty is not None: + return last_difficulty return _get_static_task_difficulty(self.job_def.metadata) @@ -1631,7 +1666,12 @@ def panic(self): current_difficulty = self.get_current_difficulty() self.__last_difficulty = _next_task_difficulty(current_difficulty) - def append_chat_log(self, request: Dict, llm_resp: LLM_Response): + def append_chat_log( + self, + request: Dict, + llm_resp: LLM_Response, + request_difficulty: Optional[TaskDifficulty] = None, + ): """ Register the LLM response in the Job's chat_log container. @@ -1643,6 +1683,7 @@ def append_chat_log(self, request: Dict, llm_resp: LLM_Response): Args: request: The original request parameters (compatible with LLM_API.process_request) llm_resp: The LLM's response as an LLM_Response object + request_difficulty: Difficulty captured when the request was built. The console_pos is set to len(console), marking the position past the end of the current console output. @@ -1686,6 +1727,7 @@ def append_chat_log(self, request: Dict, llm_resp: LLM_Response): console_pos=len(self.py_env.console) if self.py_env.console else 0, llm_resp=stored_resp, llm_reasoning_payload=step_data.reasoning_payload, + request_difficulty=request_difficulty, ) self.chat_log.append(chat_item) diff --git a/statek/executors/utils.py b/statek/executors/utils.py index b0442bf..2bbe638 100644 --- a/statek/executors/utils.py +++ b/statek/executors/utils.py @@ -460,6 +460,8 @@ def _resolve_output_future(result): value = asyncio.get_running_loop().run_until_complete(value) return value + tracks_resource_expansion = bool(getattr(job.job_def, "extra_resources", None)) + execution_difficulty = job.get_current_difficulty() if tracks_resource_expansion else None with _setup_execution_context(job, global_context, local_context, print_fn=print_fn): if policy is not None: try: @@ -512,6 +514,14 @@ def _is_print_call(node): try: is_expression = isinstance(node, ast.Expr) + current_difficulty = ( + job.get_current_difficulty() if tracks_resource_expansion else None + ) + if tracks_resource_expansion and current_difficulty != execution_difficulty: + for tool in job.get_resource_tools(): + global_context[tool.__name__] = inject_context(tool, global_context) + execution_difficulty = current_difficulty + global_context.update(local_context) sync_local = _MirrorDict( local_context, @@ -1188,7 +1198,8 @@ def _cli_error_append(cli_idx, text): # Step 12: Get next request parameters — log pending console batch first _log_pending_console(job) - request = job.get_next_request() + request_difficulty = job.get_current_difficulty() + request = job.get_next_request(difficulty=request_difficulty) request["metadata"] = dict(request["metadata"] or {}) request["metadata"]["PROVIDER"] = provider_to_use # Materialize chat_history generator so it can be consumed by process_request @@ -1222,7 +1233,11 @@ def _cli_error_append(cli_idx, text): return False processed_response = LLM_Response(step_data=processed_step, stats=response.stats) - job.append_chat_log(request, processed_response) + job.append_chat_log( + request, + processed_response, + request_difficulty=request_difficulty, + ) # Step 15: MD_DIALOG/DIRECT — dispatch LLM response text to user via send_message dialog_error = False diff --git a/tests/test_extra_resources_integration.py b/tests/test_extra_resources_integration.py index 1201af6..94d1f0b 100644 --- a/tests/test_extra_resources_integration.py +++ b/tests/test_extra_resources_integration.py @@ -7,13 +7,16 @@ import pytest from statek.agents.dialog_agent import DialogAgent -from statek.executors.job import Job, JobDef, job_def_identity_tag_for_job_def +from statek.executors.chat_log_item import LLM_LogItem +from statek.executors.job import Job, JobDef, JobStatus, job_def_identity_tag_for_job_def from statek.executors.utils import ( AgentLoopDef, find_existing_job_def, run_agentic_fleet, run_agentic_loop, ) from statek.extra_resources import parse_extra_resources +from statek.llm_api import LLM_Response, LLM_StepData, LLM_Stats from statek.prompt_config import PromptDef, parse_prompt_file, update_prompt_config from statek.task import create_new_job +from statek.task_difficulty import TaskDifficulty from tests.conftest import DB0_DIR @@ -147,6 +150,69 @@ def test_extra_resources_persistence_and_ordered_identity(supervised_agent) -> N assert find_existing_job_def(agent, None, extra_resources=(["b", "a"], None, ["c"])) is None +def test_request_difficulty_and_escalation_survive_reopen(supervised_agent) -> None: + """Persisted jobs retain request snapshots and their current escalation level.""" + definition = JobDef( + agent=supervised_agent, + metadata={ + "MODEL": "L:small,M:medium,H:large", + "DEFAULT_DIFFICULTY": "low", + }, + ) + job = Job(job_def=definition, job_status=JobStatus.STARTED) + request_difficulty = job.get_current_difficulty() + request = job.get_next_request() + job.append_chat_log( + request, + LLM_Response( + step_data=LLM_StepData(text="response", call_requests=None), + stats=LLM_Stats(0, 0, None), + ), + request_difficulty=request_difficulty, + ) + job.panic() + identifier = db0.uuid(job) + + db0.close() + db0.init(DB0_DIR, read_write=True) + db0.open("test_prefix", "rw") + restored = db0.fetch(identifier) + + assert restored.get_current_difficulty() == TaskDifficulty.medium + assert restored.chat_log[0].request_difficulty == TaskDifficulty.low + assert restored.get_request_data(0)["model"] == "small" + + +def test_legacy_job_and_log_item_without_difficulty_fields_survive_reopen( + supervised_agent, +) -> None: + """Records written without difficulty fields retain legacy fallback behavior.""" + definition = JobDef( + agent=supervised_agent, + metadata={ + "MODEL": "L:small,M:medium,H:large", + "DEFAULT_DIFFICULTY": "low", + }, + ) + item = LLM_LogItem(console_pos=0, llm_resp="legacy response") + del item.request_difficulty + job = Job( + job_def=definition, + job_status=JobStatus.STARTED, + chat_log=[item], + ) + del job._Job__last_difficulty # pylint: disable=protected-access + identifier = db0.uuid(job) + + db0.close() + db0.init(DB0_DIR, read_write=True) + db0.open("test_prefix", "rw") + restored = db0.fetch(identifier) + + assert restored.get_current_difficulty() == TaskDifficulty.low + assert restored.get_request_data(0)["model"] == "small" + + def test_extra_resources_empty_and_redundant_internal_values(supervised_agent) -> None: empty = supervised_agent.create_job_def(extra_resources=([], None, [])) assert empty.extra_resources is None diff --git a/tests/test_extra_resources_runtime.py b/tests/test_extra_resources_runtime.py index d809aae..86a0517 100644 --- a/tests/test_extra_resources_runtime.py +++ b/tests/test_extra_resources_runtime.py @@ -18,7 +18,7 @@ from statek.executors.job import Job, JobDef from statek.executors.utils import exec_step, exec_tool from statek.extra_resources import parse_extra_resources -from statek.llm_api import select_request_tools +from statek.llm_api import LLM_Response, LLM_StepData, LLM_Stats, select_request_tools from statek.prompt_config import make_system_prompt from statek.system import tool from statek.utils import CallSpec, get_current_agent, get_current_job @@ -313,6 +313,19 @@ async def test_exec_step_calls_inherited_action_in_receiver_context( assert resource_tool_job.py_env.local_state["result"] == "receiver:receiver:saved" +@pytest.mark.asyncio +async def test_exec_step_refreshes_inherited_tools_after_panic_in_same_block( + resource_tool_job: Job, +) -> None: + """A tool activated by panic is callable by the following statement.""" + await exec_step( + "panic()\nresult = schedule_action('expanded')", + resource_tool_job, + ) + + assert resource_tool_job.py_env.local_state["result"] == "schedule:expanded" + + @pytest.mark.asyncio async def test_exec_tool_calls_inherited_action_in_receiver_context( resource_tool_job: Job, @@ -380,6 +393,45 @@ def donor_receiver_action(value: str, **kwargs) -> str: # pylint: disable=unuse assert resource_tool_job.py_env.local_state["collision_result"] == "receiver:ok" +def test_historical_requests_keep_resources_from_their_original_difficulty( + resource_tool_job: Job, +) -> None: + """Later escalation must not add tools to an earlier request preview.""" + low_difficulty = resource_tool_job.get_current_difficulty() + low_request = resource_tool_job.get_next_request() + resource_tool_job.append_chat_log( + low_request, + LLM_Response( + step_data=LLM_StepData(text="low response", call_requests=None), + stats=LLM_Stats(0, 0, None), + ), + request_difficulty=low_difficulty, + ) + resource_tool_job.panic() + medium_difficulty = resource_tool_job.get_current_difficulty() + medium_request = resource_tool_job.get_next_request() + resource_tool_job.append_chat_log( + medium_request, + LLM_Response( + step_data=LLM_StepData(text="medium response", call_requests=None), + stats=LLM_Stats(0, 0, None), + ), + request_difficulty=medium_difficulty, + ) + resource_tool_job.panic() + + low_tools = [ + tool_fn.__name__ for tool_fn in resource_tool_job.get_request_data(0)["available_tools"] + ] + medium_tools = [ + tool_fn.__name__ for tool_fn in resource_tool_job.get_request_data(1)["available_tools"] + ] + + assert "preference_action" in low_tools + assert "schedule_action" not in low_tools + assert "schedule_action" in medium_tools + + def test_resource_wrappers_pass_active_agent_roles(resource_tool_job: Job) -> None: """Example and document wrappers pass receiver-first active roles to their helpers.""" expected_roles = ["receiver", "preferences_assistant"] diff --git a/tests/test_job.py b/tests/test_job.py index c604e32..9b64f9f 100644 --- a/tests/test_job.py +++ b/tests/test_job.py @@ -1070,6 +1070,20 @@ def test_panic_raises_when_already_high(job_def_factory): _run_with_current_job(job, job.panic) +def test_old_job_without_dynamic_difficulty_field_uses_static_default(job_def_factory): + """Jobs persisted before the dynamic field existed remain executable.""" + job_def = job_def_factory(metadata={ + "MODEL": "L:small,M:medium,H:large", + "DEFAULT_DIFFICULTY": "low", + }) + job = Job(job_def=job_def, job_status=JobStatus.READY) + del job._Job__last_difficulty # pylint: disable=protected-access + + assert job.get_current_difficulty() == TaskDifficulty.low + job.panic() + assert job.get_current_difficulty() == TaskDifficulty.medium + + def test_get_current_model_returns_plain_model(job_def_factory): """MODEL without difficulty labels is returned unchanged.""" job_def = job_def_factory(metadata={"MODEL": "test-model"}) @@ -1530,6 +1544,71 @@ def test_get_request_data_reconstructs_first_and_second_turn(self, job_factory): assert history_2[1].content_src == ContentSource.CONSOLE assert historical_2["model"] == "test-model" + def test_get_request_data_keeps_historical_model_and_prompt_difficulty( + self, job_def_factory, + ): + """Historical previews use the difficulty captured for each request.""" + job_def = job_def_factory(metadata={ + "MODEL": "L:small,M:medium,H:large", + "DEFAULT_DIFFICULTY": "low", + }) + job_def.agent.update_system_prompt(parse_system_prompt( + "Intro.\n\n" + "--- low: Scope ---\nLow instructions.\n\n" + "--- medium: Scope ---\nMedium instructions.\n\n" + "--- high: Scope ---\nHigh instructions." + )) + job = Job(job_def=job_def, job_status=JobStatus.STARTED) + + low_difficulty = job.get_current_difficulty() + low_request = job.get_next_request() + job.append_chat_log( + low_request, + LLM_Response( + step_data=LLM_StepData(text="low response", call_requests=None), + stats=LLM_Stats(0, 0, None), + ), + request_difficulty=low_difficulty, + ) + job.panic() + medium_difficulty = job.get_current_difficulty() + medium_request = job.get_next_request() + job.append_chat_log( + medium_request, + LLM_Response( + step_data=LLM_StepData(text="medium response", call_requests=None), + stats=LLM_Stats(0, 0, None), + ), + request_difficulty=medium_difficulty, + ) + job.panic() + + historical_low = job.get_request_data(0) + historical_medium = job.get_request_data(1) + + assert historical_low["model"] == "small" + assert "Low instructions." in historical_low["system_prompt"] + assert "Medium instructions." not in historical_low["system_prompt"] + assert historical_medium["model"] == "medium" + assert "Medium instructions." in historical_medium["system_prompt"] + assert "High instructions." not in historical_medium["system_prompt"] + + def test_get_request_data_supports_legacy_log_item_without_difficulty_snapshot( + self, job_def_factory, + ): + """Old log items fall back to the job's current request behavior.""" + job_def = job_def_factory(metadata={ + "MODEL": "L:small,M:medium,H:large", + "DEFAULT_DIFFICULTY": "low", + }) + job = Job(job_def=job_def, job_status=JobStatus.STARTED) + job.chat_log.append(LLM_LogItem(console_pos=0, llm_resp="legacy response")) + job.panic() + + historical = job.get_request_data(0) + + assert historical["model"] == "medium" + def test_get_request_data_rejects_out_of_range_turn(self, job_factory): """Missing historical turns raise IndexError.""" job = job_factory() diff --git a/tests/test_run_job_step.py b/tests/test_run_job_step.py index a232f39..a8a7cd0 100644 --- a/tests/test_run_job_step.py +++ b/tests/test_run_job_step.py @@ -27,6 +27,7 @@ from statek.llm_api import LLM_Response, LLM_StepData, LLM_Stats, CallParams, OpenRouter_API from statek.settings import LLM_API_Settings from statek.task import SubTaskHandler +from statek.task_difficulty import TaskDifficulty from statek.utils import CodeBlock, CallSpec @@ -299,6 +300,34 @@ async def test_stores_str_when_no_tool_calls( assert isinstance(llm_resp, str) assert llm_resp == "x = 42" + @pytest.mark.asyncio + async def test_stores_request_difficulty_before_waiting_for_response( + self, job_def_factory, db0_fixture # pylint: disable=unused-argument + ): + """A difficulty change while awaiting the provider keeps the request snapshot.""" + job_def = job_def_factory(metadata={ + "MODEL": "L:small,M:medium,H:large", + "DEFAULT_DIFFICULTY": "low", + }) + job = Job(job_def=job_def, job_status=JobStatus.STARTED) + mock_response = _llm_response("x = 42") + + async def process_request(**kwargs): # pylint: disable=unused-argument + job.panic() + return mock_response + + mock_api = MagicMock() + mock_api.process_request = process_request + mock_harness = MagicMock() + + with patch("statek.executors.utils.LLM_API") as mock_llm_api_cls, \ + patch("statek.executors.utils.get_llm_harness", return_value=mock_harness): + mock_llm_api_cls.get.return_value = mock_api + await run_job_step(job, provider="OPENROUTER") + + assert job.get_current_difficulty() == TaskDifficulty.medium + assert job.chat_log[0].request_difficulty == TaskDifficulty.low + class TestRunJobStepHarnessIsolation: """Tests that harness token usage is isolated per job.""" From 86ffa539d8b4eac573dbe24681bd529a3cb58e09 Mon Sep 17 00:00:00 2001 From: Kiki McArron Date: Wed, 16 Sep 2026 15:50:41 +0200 Subject: [PATCH 8/9] Fixes after review --- statek/agents/list_of_documents.py | 2 - statek/document.py | 43 +++++----- statek/executors/utils.py | 29 +++---- statek/system.py | 21 ++++- tests/test_extra_resources_runtime.py | 110 +++++++++++++++++++++++++- tests/test_find_document.py | 40 +++++++++- tests/test_list_of_documents.py | 22 +++++- tests/test_python_sandbox.py | 99 +++++++++++++++++++++++ tests/test_show_document.py | 22 ++++++ 9 files changed, 345 insertions(+), 43 deletions(-) diff --git a/statek/agents/list_of_documents.py b/statek/agents/list_of_documents.py index c50e8f0..6905936 100644 --- a/statek/agents/list_of_documents.py +++ b/statek/agents/list_of_documents.py @@ -28,7 +28,6 @@ find_document, find_topic, load_documents, - validate_unique_document_ids, ) from statek.settings import get_statek_settings from statek.system import tool @@ -91,7 +90,6 @@ def _list_documents(topic_key, agent_name, all_topics, start_index, limit): perm_ctx_set(last_topic_id=matched.ord_no) docs = [d for d in matched.documents if d.match_audience(agent_name)] - validate_unique_document_ids(docs) total = len(docs) print(f"# Document ID: Document name ({total} total)") for d in docs[start_index:start_index + limit]: diff --git a/statek/document.py b/statek/document.py index 9a461c9..bfb6a4b 100644 --- a/statek/document.py +++ b/statek/document.py @@ -63,16 +63,6 @@ def count(self, agent_name: str | Sequence[str]) -> int: return sum(1 for doc in self.documents if doc.match_audience(agent_name)) -def validate_unique_document_ids(documents: Sequence[Document]) -> None: - """Reject an accessible document view containing an ambiguous numeric ID.""" - seen: set[int] = set() - for document in documents: - document_id = document.document_metadata["ord_no"] - if document_id in seen: - raise ValueError(f"Ambiguous document ID '{document_id}'") - seen.add(document_id) - - def parse_document(document: str) -> Document: """Parse a document string into a Document instance. @@ -223,20 +213,37 @@ def find_document(key: Union[int, str], agent_name: str | Sequence[str], document titles and indexes. """ accessible = [d for d in topic.documents if d.match_audience(agent_name)] - validate_unique_document_ids(accessible) if isinstance(key, int): - for doc in accessible: - if doc.document_metadata["ord_no"] == key: - return doc - return None + matches = [ + doc for doc in accessible + if doc.document_metadata["ord_no"] == key + ] + if len(matches) == 1: + return matches[0] + if not matches: + return None + listing = ", ".join( + f"{doc.document_metadata['title']} (#{doc.document_metadata['ord_no']})" + for doc in matches + ) + raise ValueError(f"Ambiguous document ID '{key}', matches: {listing}") key_lower = key.lower() # Exact title match takes priority - for doc in accessible: - if doc.document_metadata["title"].lower() == key_lower: - return doc + exact_matches = [ + doc for doc in accessible + if doc.document_metadata["title"].lower() == key_lower + ] + if len(exact_matches) == 1: + return exact_matches[0] + if len(exact_matches) > 1: + listing = ", ".join( + f"{doc.document_metadata['title']} (#{doc.document_metadata['ord_no']})" + for doc in exact_matches + ) + raise ValueError(f"Ambiguous document '{key}', matches: {listing}") matches = [d for d in accessible if key_lower in d.document_metadata["title"].lower()] diff --git a/statek/executors/utils.py b/statek/executors/utils.py index 2bbe638..7644d3d 100644 --- a/statek/executors/utils.py +++ b/statek/executors/utils.py @@ -393,18 +393,12 @@ def _execution_sandbox_policy(allowed_tools: Optional[Set[str]] = None): return policy -def _is_hidden_tool_call(call_spec: CallSpec, job: Job, policy) -> bool: +def _is_hidden_tool_call(call_spec: CallSpec, tool_fn: Optional[Callable], policy) -> bool: if policy is None or call_spec.func_name in policy.allowed_tools: return False - if call_spec.func_name in policy.blocked_tools: - return True - - for tool_fn in job.get_resource_tools(): - if (getattr(tool_fn, "__name__", None) == call_spec.func_name - and getattr(tool_fn, "tool_hidden", False)): - return True - - return False + if tool_fn is not None: + return bool(getattr(tool_fn, "tool_hidden", False)) + return call_spec.func_name in policy.blocked_tools def _compile_exec_node(node: ast.AST, filename: str): @@ -834,13 +828,6 @@ def _private_print(*args, sep=' ', end='\n', **kwargs): private_console.append(output.rstrip('\n')) policy = _execution_sandbox_policy() - if _is_hidden_tool_call(call_spec, job, policy): - error_msg = f"NameError: tool '{call_spec.func_name}' is not exposed to this job" - full_agent_trace( - "tool.error", - {"job_uuid": job_uuid, "tool_name": call_spec.func_name, "error": error_msg}, - ) - return error_msg, error_msg # Build global and local contexts — mirrors exec_step if job.py_env.global_state is None: @@ -873,6 +860,14 @@ def _private_print(*args, sep=' ', end='\n', **kwargs): else: func = global_context.get(call_spec.func_name) or local_context.get(call_spec.func_name) + if _is_hidden_tool_call(call_spec, func, policy): + error_msg = f"NameError: tool '{call_spec.func_name}' is not exposed to this job" + full_agent_trace( + "tool.error", + {"job_uuid": job_uuid, "tool_name": call_spec.func_name, "error": error_msg}, + ) + return error_msg, error_msg + if func is None: error_msg = f"NameError: tool '{call_spec.func_name}' not found" full_agent_trace( diff --git a/statek/system.py b/statek/system.py index c5022b7..c777a80 100644 --- a/statek/system.py +++ b/statek/system.py @@ -32,7 +32,25 @@ _TOOL_REGISTRY: list[Callable] = [] +def _accepts_local_context(func: Callable) -> bool: + """Return whether a callable accepts Statek's injected local context.""" + try: + parameters = inspect.signature(func).parameters.values() + except (TypeError, ValueError): + return True + return any( + parameter.kind == inspect.Parameter.VAR_KEYWORD + or ( + parameter.name == "_local_context" + and parameter.kind != inspect.Parameter.POSITIONAL_ONLY + ) + for parameter in parameters + ) + + def inject_context(func, __local_context): + accepts_local_context = _accepts_local_context(func) + @functools.wraps(func) def wrapped(*args, **kwargs): if "_local_context" in kwargs: @@ -40,7 +58,8 @@ def wrapped(*args, **kwargs): # defensive copy per invocation local_context = copy(__local_context) - kwargs["_local_context"] = local_context + if accepts_local_context: + kwargs["_local_context"] = local_context statek_ctx = local_context.get("_STATEK_CTX") if isinstance(statek_ctx, dict): diff --git a/tests/test_extra_resources_runtime.py b/tests/test_extra_resources_runtime.py index 86a0517..a455524 100644 --- a/tests/test_extra_resources_runtime.py +++ b/tests/test_extra_resources_runtime.py @@ -15,8 +15,10 @@ show_document as show_document_wrapper, show_example as show_example_wrapper, ) +from statek.agents.dialog_agent import DialogAgent +from statek.agents.researcher import Researcher from statek.executors.job import Job, JobDef -from statek.executors.utils import exec_step, exec_tool +from statek.executors.utils import exec_cli_step, exec_step, exec_tool from statek.extra_resources import parse_extra_resources from statek.llm_api import LLM_Response, LLM_StepData, LLM_Stats, select_request_tools from statek.prompt_config import make_system_prompt @@ -24,6 +26,24 @@ from statek.utils import CallSpec, get_current_agent, get_current_job +DIALOG_MESSAGES: list[tuple[str, object]] = [] + + +def record_dialog_message(body: str, media=None) -> None: + """Record dialog output without accepting framework context keywords.""" + DIALOG_MESSAGES.append((body, media)) + + +def research_message(message: str) -> str: + """Return a researcher response without accepting framework context keywords.""" + return f"response:{message}" + + +def context_observing_message(message: str) -> str: + """Read the active job without accepting framework context keywords.""" + return f"{get_current_job().job_def.agent.role}:{message}" + + @tool def receiver_action(value: str, **kwargs) -> str: # pylint: disable=unused-argument """Return a value from the receiving agent. @@ -62,6 +82,12 @@ def _preference_internal(value: str, **kwargs) -> str: # pylint: disable=unused return f"internal:{value}" +@tool(hidden=True) +def hidden_preference_action(**kwargs) -> str: # pylint: disable=unused-argument + """Return a value that must remain unavailable to formal tool calls.""" + return "hidden preference" + + @pytest.fixture def resource_agents(agent_factory: Callable[..., Agent]) -> dict[str, Agent]: """Configure donors in a different order from their activation order.""" @@ -343,6 +369,88 @@ async def test_exec_tool_calls_inherited_action_in_receiver_context( assert result == "receiver:receiver:saved" +@pytest.mark.asyncio +async def test_exec_tool_rejects_inherited_hidden_action( + resource_tool_job: Job, resource_agents: dict[str, Agent], +) -> None: + """Hidden donor tools remain unavailable after capability expansion.""" + resource_agents["preferences_assistant"]._tools.append(hidden_preference_action) + + result, error = await exec_tool( + CallSpec(id="hidden-donor", func_name="hidden_preference_action", args=[], kwargs={}), + resource_tool_job, + ) + + assert error is not None + assert "hidden_preference_action" in result + + +@pytest.mark.asyncio +async def test_exec_step_calls_dialog_answer_without_context_keyword( + db0_fixture, # pylint: disable=unused-argument +) -> None: + """Name-registered dialog tools retain their original callable signature.""" + DIALOG_MESSAGES.clear() + agent = DialogAgent( + send_message=record_dialog_message, + _metadata={"MODEL": "test-model"}, + ) + job = Job(job_def=agent.create_job_def()) + + running = await exec_step("answer('finished')", job) + + assert running is False + assert DIALOG_MESSAGES == [("finished", None)] + assert job.py_env.exit_status == "Success" + + +@pytest.mark.asyncio +async def test_exec_cli_step_calls_researcher_tool_without_context_keyword( + db0_fixture, # pylint: disable=unused-argument +) -> None: + """Python CLI can call a name-registered tool that has no framework kwargs.""" + agent = Researcher(send_message=research_message) + job = Job(job_def=JobDef(agent=agent, metadata={"MODEL": "test-model"})) + + await exec_cli_step("result = ask('question')", job, lambda _output: None) + + assert job.py_env.local_state["result"] == "response:question" + + +@pytest.mark.asyncio +async def test_exec_tool_calls_researcher_tool_without_context_keyword( + db0_fixture, # pylint: disable=unused-argument +) -> None: + """Formal calls preserve signatures of name-registered context tools.""" + agent = Researcher(send_message=research_message) + job = Job(job_def=JobDef(agent=agent, metadata={"MODEL": "test-model"})) + + result, error = await exec_tool( + CallSpec(id="research", func_name="ask", args=[], kwargs={"question": "question"}), + job, + ) + + assert error is None + assert result == "response:question" + + +@pytest.mark.asyncio +async def test_non_injectable_context_tool_observes_receiver_job( + db0_fixture, # pylint: disable=unused-argument +) -> None: + """Signature compatibility does not remove the receiver's Statek context.""" + agent = Researcher(send_message=context_observing_message, role="receiver") + job = Job(job_def=JobDef(agent=agent, metadata={"MODEL": "test-model"})) + + result, error = await exec_tool( + CallSpec(id="context", func_name="ask", args=[], kwargs={"question": "question"}), + job, + ) + + assert error is None + assert result == "receiver:question" + + @pytest.mark.asyncio async def test_exec_step_calls_inherited_dynamic_and_internal_tools( resource_tool_job: Job, resource_agents: dict[str, Agent], diff --git a/tests/test_find_document.py b/tests/test_find_document.py index dd98f11..652f12d 100644 --- a/tests/test_find_document.py +++ b/tests/test_find_document.py @@ -57,5 +57,43 @@ def test_find_rejects_duplicate_ids_exposed_by_combined_audiences(): _doc(1, "Agent B", audience=["agent_b"]), ]) - with pytest.raises(ValueError, match="Ambiguous document ID '1'"): + with pytest.raises(ValueError, match="Ambiguous document ID '1'") as error: find_document(1, ["agent_a", "agent_b"], topic) + + assert "Agent A (#1)" in str(error.value) + assert "Agent B (#1)" in str(error.value) + + +def test_duplicate_ids_allow_unique_exact_title_lookup(): + """An unrelated numeric collision does not block a unique exact title.""" + topic = Topic(ord_no=0, topic="T", documents=[ + _doc(1, "Receiver Guide", audience=["agent_a"]), + _doc(1, "Donor Guide", audience=["agent_b"]), + ]) + + document = find_document("Donor Guide", ["agent_a", "agent_b"], topic) + + assert document.document_metadata["title"] == "Donor Guide" + + +def test_duplicate_ids_allow_unique_title_fragment_lookup(): + """An unrelated numeric collision does not block a unique title fragment.""" + topic = Topic(ord_no=0, topic="T", documents=[ + _doc(1, "Receiver Guide", audience=["agent_a"]), + _doc(1, "Donor Manual", audience=["agent_b"]), + ]) + + document = find_document("Manual", ["agent_a", "agent_b"], topic) + + assert document.document_metadata["title"] == "Donor Manual" + + +def test_duplicate_exact_titles_are_ambiguous(): + """Exact title priority does not silently select one of multiple exact matches.""" + topic = Topic(ord_no=0, topic="T", documents=[ + _doc(1, "Shared Guide", audience=["agent_a"]), + _doc(2, "Shared Guide", audience=["agent_b"]), + ]) + + with pytest.raises(ValueError, match="Ambiguous document 'Shared Guide'"): + find_document("Shared Guide", ["agent_a", "agent_b"], topic) diff --git a/tests/test_list_of_documents.py b/tests/test_list_of_documents.py index 5d17a4c..7642d02 100644 --- a/tests/test_list_of_documents.py +++ b/tests/test_list_of_documents.py @@ -31,7 +31,7 @@ Body.""" DOC_PREFS_2 = """\ -# ord_no: 2 +# ord_no: 1 # topic: Preferences # audience: agent_b # title: Donor Details @@ -93,8 +93,24 @@ def test_list_documents_unions_resource_audiences_without_duplicate_public_docs( out = capsys.readouterr().out assert out.count("Overview") == 1 - assert "Details" in out - assert "Donor Details" in out + assert "1: Details" in out + assert "1: Donor Details" in out + + +def test_list_documents_paginates_across_duplicate_ids(docs_dir, capsys): + """Pagination uses document order rather than treating IDs as unique keys.""" + job = StatekContextJob() + run_with_statek_job( + job, + lambda: _impl( + ["agent_a", "agent_b"], docs_dir, topic="Preferences", + start_index=2, limit=1, + ), + ) + out = capsys.readouterr().out + + assert "1: Donor Details" in out + assert "1: Details" not in out def test_list_documents_sets_last_topic_id(docs_dir): diff --git a/tests/test_python_sandbox.py b/tests/test_python_sandbox.py index f27eb4e..fca9520 100644 --- a/tests/test_python_sandbox.py +++ b/tests/test_python_sandbox.py @@ -128,6 +128,34 @@ def hidden_tool(**kwargs): # pylint: disable=unused-argument return "hidden" +HIDDEN_ALIAS_CALLS: list[str] = [] + + +@tool(hidden=True) +def hidden_alias_target(**kwargs): # pylint: disable=unused-argument + """Record if a hidden callable is reached through a context alias.""" + HIDDEN_ALIAS_CALLS.append("called") + return "hidden alias" + + +@tool(hidden=True) +def formal_visibility_collision(**kwargs): # pylint: disable=unused-argument + """Hidden registration used to test selected-callable precedence.""" + return "hidden collision" + + +HIDDEN_FORMAL_VISIBILITY_COLLISION = formal_visibility_collision + + +@tool # pylint: disable=function-redefined +def formal_visibility_collision(**kwargs): # pylint: disable=unused-argument,function-redefined + """Visible receiver tool sharing a name with a hidden registration.""" + return "visible collision" + + +VISIBLE_FORMAL_VISIBILITY_COLLISION = formal_visibility_collision + + @tool def _internal_tool(**kwargs): # pylint: disable=unused-argument return "internal" @@ -991,6 +1019,77 @@ async def test_hidden_internal_and_unknown_tools_fail(db0_fixture, tool_name): assert tool_name in result +@pytest.mark.asyncio +async def test_hidden_receiver_context_alias_fails_before_call(db0_fixture): # pylint: disable=unused-argument + """Formal execution rejects the hidden callable selected by alias fallback.""" + HIDDEN_ALIAS_CALLS.clear() + job = _job_with_tools([visible_tool]) + job.job_def.agent.context["hidden_alias"] = hidden_alias_target + + result, error = await exec_tool(_call_spec("hidden_alias"), job) + + assert error is not None + assert "hidden_alias" in result + assert not HIDDEN_ALIAS_CALLS + + +@pytest.mark.asyncio +async def test_visible_receiver_context_alias_executes(db0_fixture): # pylint: disable=unused-argument + """Visible fallback context callables remain executable by alias.""" + job = _job_with_tools() + job.job_def.agent.context["visible_alias"] = visible_tool + + result, error = await exec_tool( + _call_spec("visible_alias", kwargs={"value": "ok"}), job, + ) + + assert error is None + assert result == "visible: ok" + + +@pytest.mark.asyncio +async def test_visible_receiver_tool_wins_over_hidden_registered_name(db0_fixture): # pylint: disable=unused-argument + """Formal resolution validates the selected visible receiver callable.""" + job = _job_with_tools([VISIBLE_FORMAL_VISIBILITY_COLLISION]) + + result, error = await exec_tool(_call_spec("formal_visibility_collision"), job) + + assert error is None + assert result == "visible collision" + + +@pytest.mark.asyncio +async def test_explicit_allow_list_permits_hidden_context_alias(db0_fixture): # pylint: disable=unused-argument + """An explicit sandbox allow-list entry retains precedence.""" + statek.init(StatekSettings( + prompt_defs={}, python_sandbox_allowed_tools="hidden_alias", + )) + HIDDEN_ALIAS_CALLS.clear() + job = _job_with_tools() + job.job_def.agent.context["hidden_alias"] = hidden_alias_target + + result, error = await exec_tool(_call_spec("hidden_alias"), job) + + assert error is None + assert result == "hidden alias" + assert HIDDEN_ALIAS_CALLS == ["called"] + + +@pytest.mark.asyncio +async def test_sandbox_off_permits_hidden_context_alias(db0_fixture): # pylint: disable=unused-argument + """Disabling the sandbox retains unrestricted formal-call behavior.""" + statek.init(StatekSettings(prompt_defs={}), restricted=False) + HIDDEN_ALIAS_CALLS.clear() + job = _job_with_tools() + job.job_def.agent.context["hidden_alias"] = hidden_alias_target + + result, error = await exec_tool(_call_spec("hidden_alias"), job) + + assert error is None + assert result == "hidden alias" + assert HIDDEN_ALIAS_CALLS == ["called"] + + @pytest.mark.asyncio async def test_direct_hidden_tool_call_in_code_fails(db0_fixture): # pylint: disable=unused-argument job = _job_with_tools([hidden_tool]) diff --git a/tests/test_show_document.py b/tests/test_show_document.py index fe0ac8d..459d2f6 100644 --- a/tests/test_show_document.py +++ b/tests/test_show_document.py @@ -24,6 +24,13 @@ # title: Secret Guide Restricted content.""" +DOC_DONOR = """\ +# ord_no: 1 +# topic: Guide +# audience: agent_b +# title: Donor Guide +Donor content.""" + @pytest.fixture(autouse=True) def clear_cache(): @@ -36,6 +43,7 @@ def clear_cache(): def docs_dir(temp_dir): _write(temp_dir, "doc0.txt", DOC_LONG) _write(temp_dir, "doc1.txt", DOC_RESTRICTED) + _write(temp_dir, "doc2.txt", DOC_DONOR) return temp_dir @@ -120,6 +128,20 @@ def test_show_document_accepts_any_active_resource_audience(docs_dir, capsys): assert "Restricted content." in out +def test_show_document_by_unique_title_with_duplicate_resource_ids(docs_dir, capsys): + """A donor document remains retrievable by title when active IDs collide.""" + job = StatekContextJob() + run_with_statek_job( + job, + lambda: show_document( + ["agent_a", "agent_b"], docs_dir, key="Donor Guide", topic="Guide", + ), + ) + out = capsys.readouterr().out + + assert "Donor content." in out + + def _write(directory, filename, content): with open(os.path.join(directory, filename), "w", encoding="utf-8") as f: f.write(content) From 6e761cd1e8abef251c61f61595834b1e91cd9d2d Mon Sep 17 00:00:00 2001 From: Kiki McArron Date: Fri, 18 Sep 2026 14:37:22 +0200 Subject: [PATCH 9/9] Fix after PR review --- statek/executors/chat_log_item.py | 2 +- statek/executors/job.py | 13 +++++++++++-- tests/test_extra_resources_integration.py | 7 ++++--- tests/test_job.py | 23 +++++++++++++++++++++-- tests/test_run_job_step.py | 7 ++++--- 5 files changed, 41 insertions(+), 11 deletions(-) diff --git a/statek/executors/chat_log_item.py b/statek/executors/chat_log_item.py index 327184b..7f4fc36 100644 --- a/statek/executors/chat_log_item.py +++ b/statek/executors/chat_log_item.py @@ -91,7 +91,7 @@ class LLM_LogItem(ChatLogItem): llm_resp: Union[str, CodeBlock] = None # Opaque provider continuation material replayed only by its matching formatter. llm_reasoning_payload: Optional[Any] = None - # Difficulty used to build this request; None identifies legacy log items. + # Non-default difficulty used to build this request; None means the job's static default. request_difficulty: Optional[TaskDifficulty] = None diff --git a/statek/executors/job.py b/statek/executors/job.py index c992263..de6698b 100644 --- a/statek/executors/job.py +++ b/statek/executors/job.py @@ -1491,11 +1491,14 @@ def get_request_data(self, turn_num: int) -> Dict[str, Any]: key: value for key, value in push_log.items() if key <= console_limit } + request_difficulty = getattr(target_item, "request_difficulty", None) + if request_difficulty is None: + request_difficulty = _get_static_task_difficulty(self.job_def.metadata) return self._build_request_data( chat_log=history_chat_log, console=history_console, push_log=history_push_log, - difficulty=getattr(target_item, "request_difficulty", None), + difficulty=request_difficulty, ) def get_resource_agents( @@ -1683,11 +1686,17 @@ def append_chat_log( Args: request: The original request parameters (compatible with LLM_API.process_request) llm_resp: The LLM's response as an LLM_Response object - request_difficulty: Difficulty captured when the request was built. + request_difficulty: Concrete difficulty captured when the request + was built, or None for the static default. The console_pos is set to len(console), marking the position past the end of the current console output. """ + if ( + request_difficulty is not None + and request_difficulty == _get_static_task_difficulty(self.job_def.metadata) + ): + request_difficulty = None chat_style = self.job_def.chat_style is_md_style = chat_style in ( # pylint: disable=no-member ChatStyle.MARKDOWN, ChatStyle.MD_DIALOG) diff --git a/tests/test_extra_resources_integration.py b/tests/test_extra_resources_integration.py index 94d1f0b..c39cefc 100644 --- a/tests/test_extra_resources_integration.py +++ b/tests/test_extra_resources_integration.py @@ -179,14 +179,14 @@ def test_request_difficulty_and_escalation_survive_reopen(supervised_agent) -> N restored = db0.fetch(identifier) assert restored.get_current_difficulty() == TaskDifficulty.medium - assert restored.chat_log[0].request_difficulty == TaskDifficulty.low + assert restored.chat_log[0].request_difficulty is None assert restored.get_request_data(0)["model"] == "small" def test_legacy_job_and_log_item_without_difficulty_fields_survive_reopen( supervised_agent, ) -> None: - """Records written without difficulty fields retain legacy fallback behavior.""" + """Records written without difficulty fields reconstruct at the static default.""" definition = JobDef( agent=supervised_agent, metadata={ @@ -209,7 +209,8 @@ def test_legacy_job_and_log_item_without_difficulty_fields_survive_reopen( db0.open("test_prefix", "rw") restored = db0.fetch(identifier) - assert restored.get_current_difficulty() == TaskDifficulty.low + restored.panic() + assert restored.get_current_difficulty() == TaskDifficulty.medium assert restored.get_request_data(0)["model"] == "small" diff --git a/tests/test_job.py b/tests/test_job.py index 9b64f9f..5ab46dc 100644 --- a/tests/test_job.py +++ b/tests/test_job.py @@ -1592,11 +1592,13 @@ def test_get_request_data_keeps_historical_model_and_prompt_difficulty( assert historical_medium["model"] == "medium" assert "Medium instructions." in historical_medium["system_prompt"] assert "High instructions." not in historical_medium["system_prompt"] + assert job.chat_log[0].request_difficulty is None + assert job.chat_log[1].request_difficulty == medium_difficulty def test_get_request_data_supports_legacy_log_item_without_difficulty_snapshot( self, job_def_factory, ): - """Old log items fall back to the job's current request behavior.""" + """Old log items without a snapshot reconstruct at the static default.""" job_def = job_def_factory(metadata={ "MODEL": "L:small,M:medium,H:large", "DEFAULT_DIFFICULTY": "low", @@ -1607,7 +1609,24 @@ def test_get_request_data_supports_legacy_log_item_without_difficulty_snapshot( historical = job.get_request_data(0) - assert historical["model"] == "medium" + assert historical["model"] == "small" + + def test_get_request_data_resolves_none_to_medium_static_default( + self, job_def_factory, + ): + """A None snapshot resolves to the configured default rather than current difficulty.""" + job_def = job_def_factory(metadata={ + "MODEL": "L:small,M:medium,H:large", + "DEFAULT_DIFFICULTY": "medium", + }) + job = Job(job_def=job_def, job_status=JobStatus.STARTED) + job.chat_log.append(LLM_LogItem( + console_pos=0, llm_resp="default response", request_difficulty=None, + )) + job.panic() + + assert job.get_current_difficulty() == TaskDifficulty.high + assert job.get_request_data(0)["model"] == "medium" def test_get_request_data_rejects_out_of_range_turn(self, job_factory): """Missing historical turns raise IndexError.""" diff --git a/tests/test_run_job_step.py b/tests/test_run_job_step.py index a8a7cd0..cc3333b 100644 --- a/tests/test_run_job_step.py +++ b/tests/test_run_job_step.py @@ -301,10 +301,10 @@ async def test_stores_str_when_no_tool_calls( assert llm_resp == "x = 42" @pytest.mark.asyncio - async def test_stores_request_difficulty_before_waiting_for_response( + async def test_stores_default_request_as_none_before_waiting_for_response( self, job_def_factory, db0_fixture # pylint: disable=unused-argument ): - """A difficulty change while awaiting the provider keeps the request snapshot.""" + """A provider-time escalation cannot rewrite a default request snapshot.""" job_def = job_def_factory(metadata={ "MODEL": "L:small,M:medium,H:large", "DEFAULT_DIFFICULTY": "low", @@ -326,7 +326,8 @@ async def process_request(**kwargs): # pylint: disable=unused-argument await run_job_step(job, provider="OPENROUTER") assert job.get_current_difficulty() == TaskDifficulty.medium - assert job.chat_log[0].request_difficulty == TaskDifficulty.low + assert job.chat_log[0].request_difficulty is None + assert job.get_request_data(0)["model"] == "small" class TestRunJobStepHarnessIsolation: