diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index c7a65e80d3..b820187562 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -48,6 +48,17 @@ } +def _deduplicate_origin_session_ids(origin_session_ids: Iterable[str]) -> list[str]: + """Return origin session IDs in first-seen order without duplicates.""" + unique_origin_session_ids: list[str] = [] + seen_origin_session_ids: set[str] = set() + for origin_session_id in origin_session_ids: + if origin_session_id not in seen_origin_session_ids: + seen_origin_session_ids.add(origin_session_id) + unique_origin_session_ids.append(origin_session_id) + return unique_origin_session_ids + + @runtime_checkable class TokenizerProtocol(Protocol): """Protocol for token counters used by token-aware compaction strategies.""" @@ -1706,6 +1717,16 @@ async def before_run( if not all_messages: return + # Track each original message's source before compaction + source_by_id: dict[int, str] = { + id(message): sid for sid, msgs in context.context_messages.items() for message in msgs + } + + # Track original messages by message_id for attribution preservation + message_by_message_id: dict[str, Message] = { + message.message_id: message for message in all_messages if message.message_id + } + await _run_compaction_strategy( all_messages, strategy=self.before_strategy, @@ -1714,9 +1735,64 @@ async def before_run( ) projected = project_included_messages(all_messages) - projected_set = {id(m) for m in projected} - for sid in list(context.context_messages): - context.context_messages[sid] = [m for m in context.context_messages[sid] if id(m) in projected_set] + + # Rebuild provider message lists from the projected list, preserving source attribution + # and including new synthetic messages created by compaction strategies + rebuilt: dict[str, list[Message]] = {sid: [] for sid in context.context_messages} + fallback_sid = next(iter(rebuilt), self.source_id) + last_sid = fallback_sid + for message in projected: + # For new synthetic messages, use the last known source; for original messages, use their tracked source + sid = source_by_id.get(id(message), last_sid) + if sid not in rebuilt: + # If the source was somehow removed during compaction, fall back to the last known source + sid = last_sid + rebuilt[sid].append(message) + last_sid = sid + + context.context_messages.clear() + context.context_messages.update(rebuilt) + + # Preserve attribution metadata on synthetic summary messages + # This ensures cross-session governance signals are not lost when content is summarized + for message in projected: + # Check if this is a synthetic summary message + annotation = _read_group_annotation_raw(message) + if annotation is None: + continue + + summarized_message_ids: Any = annotation.get(SUMMARY_OF_MESSAGE_IDS_KEY) + if not isinstance(summarized_message_ids, list) or not summarized_message_ids: + continue + + # Collect origin_session_ids from all summarized messages + origin_session_ids: list[str] = [] + for msg_id in cast("list[Any]", summarized_message_ids): + if not isinstance(msg_id, str): + continue + original_message = message_by_message_id.get(msg_id) + if original_message is None: + continue + original_attribution = original_message.additional_properties.get("_attribution") + if isinstance(original_attribution, Mapping): + original_origins = original_attribution.get("origin_session_ids") + if isinstance(original_origins, Sequence) and not isinstance(original_origins, str): + for origin in cast("Sequence[Any]", original_origins): + if isinstance(origin, str): + origin_session_ids.append(origin) + + if origin_session_ids: + # Deduplicate and attach to the synthetic summary + deduplicated_ids = _deduplicate_origin_session_ids(origin_session_ids) + summary_attribution = message.additional_properties.get("_attribution") + if isinstance(summary_attribution, Mapping): + # Merge with existing attribution if present + merged_attribution = dict(cast("Mapping[str, Any]", summary_attribution)) + merged_attribution["origin_session_ids"] = deduplicated_ids + message.additional_properties["_attribution"] = merged_attribution + else: + # Create new attribution dict + message.additional_properties["_attribution"] = {"origin_session_ids": deduplicated_ids} async def after_run( self, diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index b0e82ae94f..501036626d 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -102,10 +102,11 @@ def _assistant_reasoning_and_function_calls(*call_ids: str) -> Message: return Message(role="assistant", contents=contents) -def _tool_result(call_id: str, result: str) -> Message: +def _tool_result(call_id: str, result: str, message_id: str | None = None) -> Message: return Message( role="tool", contents=[Content.from_function_result(call_id=call_id, result=result)], + message_id=message_id, ) @@ -2010,3 +2011,351 @@ def test_serialize_message_preserves_non_ascii_for_token_count() -> None: assert text in serialized assert "\\u3053" not in serialized assert tokenizer.count_tokens(serialized) < tokenizer.count_tokens(escaped) + + +async def test_compaction_provider_before_run_preserves_synthetic_summary_messages() -> None: + """Test that before_run preserves synthetic summary messages created by compaction strategies. + + This is a regression test for a bug where ToolResultCompactionStrategy and SummarizationStrategy + create new synthetic Message objects during compaction, but CompactionProvider.before_run only + filtered existing messages by id(), causing the new summary messages to be silently dropped. + """ + from agent_framework._sessions import SessionContext + + # Create a session context with some messages from a history provider + messages = [ + Message(role="user", contents=["hello"]), + _assistant_function_call("c1"), + _tool_result("c1", "sunny, 18C"), + Message(role="assistant", contents=["final response"]), + ] + + ctx = SessionContext(input_messages=[]) + ctx.extend_messages("history", messages) + + # Verify initial state - 4 messages from history provider + assert len(ctx.context_messages["history"]) == 4 + assert len(ctx.get_messages()) == 4 + + # Apply ToolResultCompactionStrategy via CompactionProvider.before_run + provider = CompactionProvider( + before_strategy=ToolResultCompactionStrategy(keep_last_tool_call_groups=0) + ) + + await provider.before_run(agent=None, session=None, context=ctx, state={}) + + # After compaction, we should have 3 messages: + # - user message + # - synthetic summary message "[Tool results: tool: sunny, 18C]" + # - final assistant response + final_messages = ctx.get_messages() + assert len(final_messages) == 3, f"Expected 3 messages after compaction, got {len(final_messages)}" + + # Verify the summary message is present - it should be an assistant message with tool results summary + summary_messages = [ + m for m in final_messages + if m.role == "assistant" and any(hasattr(c, 'text') and "Tool results" in c.text for c in m.contents) + ] + assert len(summary_messages) == 1, "Summary message should be present in final messages" + + # Verify the context_messages dict was also updated correctly + assert len(ctx.context_messages["history"]) == 3, "History provider should have 3 messages after compaction" + + # Verify the summary is in the history provider's list + history_summary_messages = [ + m for m in ctx.context_messages["history"] + if m.role == "assistant" and any(hasattr(c, 'text') and "Tool results" in c.text for c in m.contents) + ] + assert len(history_summary_messages) == 1, "Summary message should be in history provider's message list" + + +async def test_compaction_provider_before_run_preserves_summarization_strategy_messages() -> None: + """Test that before_run preserves synthetic summary messages created by SummarizationStrategy. + + This is a regression test for the same bug affecting SummarizationStrategy, which creates + new synthetic Message objects during compaction. + """ + from agent_framework._sessions import SessionContext + from unittest.mock import AsyncMock, MagicMock + + # Create a mock chat client for summarization + mock_client = AsyncMock() + mock_response = MagicMock() + mock_response.text = "Summary of conversation" + mock_client.get_response.return_value = mock_response + + # Create messages that will trigger summarization + messages = [ + Message(role="user", contents=["first message"]), + Message(role="assistant", contents=["first response"]), + Message(role="user", contents=["second message"]), + Message(role="assistant", contents=["second response"]), + Message(role="user", contents=["third message"]), + Message(role="assistant", contents=["third response"]), + ] + + ctx = SessionContext(input_messages=[]) + ctx.extend_messages("history", messages) + + # Verify initial state - 6 messages from history provider + assert len(ctx.context_messages["history"]) == 6 + assert len(ctx.get_messages()) == 6 + + # Apply SummarizationStrategy via CompactionProvider.before_run + # Set target_count=2 to trigger summarization since we have 6 non-system messages + provider = CompactionProvider( + before_strategy=SummarizationStrategy( + client=mock_client, + target_count=2, + threshold=0, + ) + ) + + await provider.before_run(agent=None, session=None, context=ctx, state={}) + + # After compaction, we should have messages including the synthetic summary + final_messages = ctx.get_messages() + assert len(final_messages) >= 3, f"Expected at least 3 messages after compaction (summary + retained), got {len(final_messages)}" + + # Verify the summary message is present + summary_messages = [ + m for m in final_messages + if m.role == "assistant" and any(hasattr(c, 'text') and "Summary of conversation" in c.text for c in m.contents) + ] + assert len(summary_messages) == 1, "Summary message should be present in final messages" + + # Verify the context_messages dict was also updated correctly + assert len(ctx.context_messages["history"]) >= 3, "History provider should have at least 3 messages after compaction" + + # Verify the summary is in the history provider's list + history_summary_messages = [ + m for m in ctx.context_messages["history"] + if m.role == "assistant" and any(hasattr(c, 'text') and "Summary of conversation" in c.text for c in m.contents) + ] + assert len(history_summary_messages) == 1, "Summary message should be in history provider's message list" + + +async def test_compaction_provider_preserves_attribution_on_synthetic_summaries() -> None: + """Test that synthetic summary messages preserve origin_session_ids from summarized messages. + + This is a regression test for a security/governance issue where cross-session attribution + was lost when messages were summarized, potentially bypassing CrossSessionObserver governance. + """ + from agent_framework._sessions import SessionContext + from unittest.mock import AsyncMock, MagicMock + + # Create a mock chat client for summarization + mock_client = AsyncMock() + mock_response = MagicMock() + mock_response.text = "Summary of conversation" + mock_client.get_response.return_value = mock_response + + # Create messages with cross-session attribution + messages = [ + Message(role="user", contents=["first message"], message_id="msg1"), + Message(role="assistant", contents=["first response"], message_id="msg2"), + Message(role="user", contents=["second message"], message_id="msg3"), + Message(role="assistant", contents=["second response"], message_id="msg4"), + ] + + # Add attribution to some messages to simulate cross-session content + messages[0].additional_properties["_attribution"] = { + "source_id": "history", + "source_type": "HistoryProvider", + "origin_session_ids": ["session-prior-1"], + } + messages[1].additional_properties["_attribution"] = { + "source_id": "history", + "source_type": "HistoryProvider", + "origin_session_ids": ["session-prior-2"], + } + + ctx = SessionContext(input_messages=[]) + ctx.extend_messages("history", messages) + + # Apply SummarizationStrategy via CompactionProvider.before_run + provider = CompactionProvider( + before_strategy=SummarizationStrategy( + client=mock_client, + target_count=1, + threshold=0, + ) + ) + + await provider.before_run(agent=None, session=None, context=ctx, state={}) + + # Find the synthetic summary message + final_messages = ctx.get_messages() + summary_messages = [ + m for m in final_messages + if m.role == "assistant" and any(hasattr(c, 'text') and "Summary of conversation" in c.text for c in m.contents) + ] + assert len(summary_messages) == 1, "Summary message should be present" + + summary = summary_messages[0] + + # Verify the summary preserves the aggregated origin_session_ids + summary_attribution = summary.additional_properties.get("_attribution") + assert summary_attribution is not None, "Summary should have attribution" + assert isinstance(summary_attribution, dict), "Attribution should be a dict" + + origin_session_ids = summary_attribution.get("origin_session_ids") + assert origin_session_ids is not None, "Summary should have origin_session_ids" + assert isinstance(origin_session_ids, list), "origin_session_ids should be a list" + + # Should contain both session IDs, deduplicated + assert set(origin_session_ids) == {"session-prior-1", "session-prior-2"}, \ + f"Expected both session IDs, got {origin_session_ids}" + + +async def test_compaction_provider_preserves_attribution_on_tool_result_summaries() -> None: + """Test that ToolResultCompactionStrategy summaries preserve origin_session_ids.""" + from agent_framework._sessions import SessionContext + + # Create messages with cross-session attribution + messages = [ + Message(role="user", contents=["hello"], message_id="msg1"), + _assistant_function_call("c1"), + _tool_result("c1", "sunny, 18C", message_id="msg2"), + Message(role="assistant", contents=["final response"], message_id="msg4"), + ] + + # Add attribution to the tool result to simulate cross-session content + messages[2].additional_properties["_attribution"] = { + "source_id": "history", + "source_type": "HistoryProvider", + "origin_session_ids": ["session-prior"], + } + + ctx = SessionContext(input_messages=[]) + ctx.extend_messages("history", messages) + + # Apply ToolResultCompactionStrategy via CompactionProvider.before_run + provider = CompactionProvider( + before_strategy=ToolResultCompactionStrategy(keep_last_tool_call_groups=0) + ) + + await provider.before_run(agent=None, session=None, context=ctx, state={}) + + # Find the synthetic summary message + final_messages = ctx.get_messages() + summary_messages = [ + m for m in final_messages + if m.role == "assistant" and any(hasattr(c, 'text') and "Tool results" in c.text for c in m.contents) + ] + assert len(summary_messages) == 1, "Tool result summary should be present" + + summary = summary_messages[0] + + # Verify the summary preserves the origin_session_ids + summary_attribution = summary.additional_properties.get("_attribution") + assert summary_attribution is not None, "Summary should have attribution" + assert isinstance(summary_attribution, dict), "Attribution should be a dict" + + origin_session_ids = summary_attribution.get("origin_session_ids") + assert origin_session_ids is not None, "Summary should have origin_session_ids" + assert isinstance(origin_session_ids, list), "origin_session_ids should be a list" + assert origin_session_ids == ["session-prior"], \ + f"Expected session-prior, got {origin_session_ids}" + + +async def test_compaction_provider_deduplicates_origin_session_ids() -> None: + """Test that duplicate origin_session_ids are deduplicated in synthetic summaries.""" + from agent_framework._sessions import SessionContext + from unittest.mock import AsyncMock, MagicMock + + # Create a mock chat client for summarization + mock_client = AsyncMock() + mock_response = MagicMock() + mock_response.text = "Summary" + mock_client.get_response.return_value = mock_response + + # Create messages where multiple messages have the same origin + messages = [ + Message(role="user", contents=["msg1"], message_id="msg1"), + Message(role="assistant", contents=["resp1"], message_id="msg2"), + Message(role="user", contents=["msg2"], message_id="msg3"), + Message(role="assistant", contents=["resp2"], message_id="msg4"), + ] + + # All messages have the same origin_session_id + for msg in messages: + msg.additional_properties["_attribution"] = { + "source_id": "history", + "source_type": "HistoryProvider", + "origin_session_ids": ["session-prior"], + } + + ctx = SessionContext(input_messages=[]) + ctx.extend_messages("history", messages) + + provider = CompactionProvider( + before_strategy=SummarizationStrategy( + client=mock_client, + target_count=1, + threshold=0, + ) + ) + + await provider.before_run(agent=None, session=None, context=ctx, state={}) + + # Find the summary + final_messages = ctx.get_messages() + summary_messages = [ + m for m in final_messages + if m.role == "assistant" and any(hasattr(c, 'text') and "Summary" in c.text for c in m.contents) + ] + assert len(summary_messages) == 1 + + summary = summary_messages[0] + origin_session_ids = summary.additional_properties.get("_attribution", {}).get("origin_session_ids") + + # Should be deduplicated - only one instance of session-prior + assert origin_session_ids == ["session-prior"], \ + f"Expected single deduplicated session ID, got {origin_session_ids}" + + +async def test_compaction_provider_no_attribution_when_sources_have_none() -> None: + """Test that summaries of messages without attribution don't get attribution added.""" + from agent_framework._sessions import SessionContext + from unittest.mock import AsyncMock, MagicMock + + # Create a mock chat client for summarization + mock_client = AsyncMock() + mock_response = MagicMock() + mock_response.text = "Summary" + mock_client.get_response.return_value = mock_response + + # Create messages without any attribution + messages = [ + Message(role="user", contents=["msg1"], message_id="msg1"), + Message(role="assistant", contents=["resp1"], message_id="msg2"), + ] + + ctx = SessionContext(input_messages=[]) + ctx.extend_messages("history", messages) + + provider = CompactionProvider( + before_strategy=SummarizationStrategy( + client=mock_client, + target_count=1, + threshold=0, + ) + ) + + await provider.before_run(agent=None, session=None, context=ctx, state={}) + + # Find the summary + final_messages = ctx.get_messages() + summary_messages = [ + m for m in final_messages + if m.role == "assistant" and any(hasattr(c, 'text') and "Summary" in c.text for c in m.contents) + ] + assert len(summary_messages) == 1 + + summary = summary_messages[0] + summary_attribution = summary.additional_properties.get("_attribution") + + # Should not have attribution since sources had none + assert summary_attribution is None or "origin_session_ids" not in summary_attribution, \ + "Summary should not have origin_session_ids when sources have none"