diff --git a/src/stirrup/core/models.py b/src/stirrup/core/models.py index c911f2c..8397ab4 100644 --- a/src/stirrup/core/models.py +++ b/src/stirrup/core/models.py @@ -304,6 +304,13 @@ def _aggregate_list[T: Addable](metadata_list: list[T]) -> T | None: for m in metadata_list[1:]: if isinstance(aggregated, dict) and isinstance(m, dict): aggregated = _merge_dicts(aggregated, m) # ty: ignore[invalid-assignment] + elif isinstance(aggregated, dict) and isinstance(m, BaseModel): + # Metadata restored from the JSON cache is plain dicts, while the + # current run appends live models. Rehydrate the dict into the live + # item's type so the pair can be added instead of raising TypeError. + aggregated = type(m).model_validate(aggregated) + m + elif isinstance(aggregated, BaseModel) and isinstance(m, dict): + aggregated = aggregated + type(aggregated).model_validate(m) else: aggregated = aggregated + m return aggregated diff --git a/tests/test_agent.py b/tests/test_agent.py index 89c8b9b..49f336e 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -34,8 +34,10 @@ ToolCall, ToolMessage, ToolResult, + ToolUseCountMetadata, TurnWarningMessage, UserMessage, + aggregate_metadata, joined_text, ) from stirrup.prompts import MESSAGE_SUMMARIZER @@ -648,6 +650,82 @@ async def test_v01_cache_state_preserves_aggregate_run_metadata() -> None: assert restored.run_metadata_by_turn == {"__legacy__": {"search": [{"queries": 1}]}} +class SearchLikeMetadata(BaseModel): + """Stand-in for tool-specific metadata models (e.g. SearchMetadata).""" + + queries: int = 0 + + def __add__(self, other: "SearchLikeMetadata") -> "SearchLikeMetadata": + return SearchLikeMetadata(queries=self.queries + other.queries) + + +async def test_resumed_run_metadata_aggregates_cached_dicts_with_live_models() -> None: + """Regression test: metadata restored from the cache is plain dicts, while a + resumed run appends live models. Aggregating the merged list at run end must + combine them instead of raising TypeError (dict + ToolUseCountMetadata).""" + state = CacheState( + msgs=[UserMessage(content="Test task")], + full_msg_history=[], + run_metadata_by_turn={"turn-1": {"my_tool": [ToolUseCountMetadata(num_uses=2)]}}, + task_hash="task", + ) + restored = CacheState.from_dict(json.loads(json.dumps(state.to_dict()))) + assert isinstance(restored.run_metadata_by_turn["turn-1"]["my_tool"][0], dict) + + # The resumed run appends a live metadata object for the same tool, keyed by + # the new turn's assistant message id; run end merges the per-turn lists. + run_metadata_by_turn = dict(restored.run_metadata_by_turn) + run_metadata_by_turn["turn-2"] = {"my_tool": [ToolUseCountMetadata(num_uses=3)]} + run_metadata: dict = {} + for turn_metadata in run_metadata_by_turn.values(): + for name, metadata_list in turn_metadata.items(): + run_metadata.setdefault(name, []).extend(metadata_list) + + assert aggregate_metadata(run_metadata) == {"my_tool": {"num_uses": 5}} + + +def test_aggregate_metadata_combines_dicts_and_models_in_either_order() -> None: + cached = ToolUseCountMetadata(num_uses=2).model_dump(mode="json") + + dict_first = aggregate_metadata( + {"tool": [cached, ToolUseCountMetadata(num_uses=3)]}, return_json_serializable=False + ) + assert dict_first["tool"] == ToolUseCountMetadata(num_uses=5) + + model_first = aggregate_metadata( + {"tool": [ToolUseCountMetadata(num_uses=3), cached]}, return_json_serializable=False + ) + assert model_first["tool"] == ToolUseCountMetadata(num_uses=5) + + two_dicts = aggregate_metadata( + {"tool": [{"num_uses": 2}, {"num_uses": 1}, ToolUseCountMetadata(num_uses=3)]}, + return_json_serializable=False, + ) + assert two_dicts["tool"] == ToolUseCountMetadata(num_uses=6) + + +def test_aggregate_metadata_rehydrates_dicts_into_the_live_model_type() -> None: + """Dict entries take the type of the live model in the same list, so custom + tool metadata types aggregate after a resume too, not just ToolUseCountMetadata.""" + result = aggregate_metadata( + {"search": [{"queries": 2}, SearchLikeMetadata(queries=3)]}, + return_json_serializable=False, + ) + assert result["search"] == SearchLikeMetadata(queries=5) + + +def test_aggregate_metadata_combines_cached_sub_agent_dicts_with_live_models() -> None: + """A cache-restored SubAgentMetadata dict aggregates with a live one, and the + nested per-tool lists (themselves a mix of dicts and models) aggregate too.""" + cached = SubAgentMetadata(message_history=[], run_metadata={"tool": [ToolUseCountMetadata(num_uses=2)]}).model_dump( + mode="json" + ) + live = SubAgentMetadata(message_history=[], run_metadata={"tool": [ToolUseCountMetadata(num_uses=1)]}) + + result = aggregate_metadata({"sub": [cached, live]}, return_json_serializable=False) + assert result["sub"]["tool"] == ToolUseCountMetadata(num_uses=3) + + def _v01_idless_tool_history() -> list[dict[str, object]]: return [ {