Python: feat: cross-session origin attribution on context messages#7041
Python: feat: cross-session origin attribution on context messages#7041moonbox3 wants to merge 7 commits into
Conversation
Add an optional origin_session_id parameter to SessionContext.extend_messages that propagates into the existing _attribution payload on Message.additional_properties. Downstream context observers can use it to detect when a provider injects content stored under a different session than the requesting one. Populate the field from the harness memory consolidation pipeline (_harness/_memory.py) when injected topics include contributions from sessions other than the current one. Add a self-contained sample observer under samples/02-agents/context_providers/cross_session_observer.py demonstrating how to subscribe to the signal. Backward-compatible: omitting the parameter preserves the existing attribution shape exactly. Tests added in test_sessions.py and test_harness_memory.py cover the new parameter, the harness cross-session case, and the same-session case. Motivated by Dai et al., Stateful Agent Backdoor (arXiv:2605.06158, May 2026), which specifically surveys MAF in section 6.1 / Table 10. See microsoft#5914 for design discussion. Surfaced during independent audit conducted by @finnoybu (Ken Tannenbaum, AEGIS Initiative); [MEDIUM, python/packages/core].
giles17
added
documentation
There was a problem hiding this comment.
Pull request overview
This PR extends the Python SessionContext attribution mechanism so context messages can carry ordered, deduplicated cross-session provenance (origin_session_ids), enabling downstream governance/audit observers to detect when injected context originated from prior sessions (e.g., consolidated memory).
Changes:
- Add optional
origin_session_idssupport toSessionContext.extend_messages(...), merging with existing_attributionwithout mutating caller-owned message objects. - Populate
origin_session_idsin the harnessMemoryContextProviderwhen injected memory topics include contributors from sessions other than the current one. - Add a self-contained sample observer and new tests covering default behavior, merge behavior, and harness cross-session attribution.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| python/packages/core/agent_framework/_sessions.py | Adds origin_session_ids parameter and attribution merge/copy behavior in extend_messages. |
| python/packages/core/agent_framework/_harness/_memory.py | Computes and passes cross-session contributor session IDs when injecting memory context. |
| python/packages/core/tests/core/test_sessions.py | Adds unit tests for extend_messages default/merge behavior and origin recording. |
| python/packages/core/tests/core/test_harness_memory.py | Adds tests validating harness memory injection sets/omits origin_session_ids correctly. |
| python/samples/02-agents/context_providers/cross_session_observer.py | New sample observer demonstrating how to detect cross-session injected messages. |
| python/samples/02-agents/context_providers/README.md | Documents the new sample and prerequisites. |
| python/packages/core/AGENTS.md | Updates session docs to mention the new origin_session_ids attribution capability. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Automated Code Review
Reviewers: 5 | Confidence: 84%
✓ Correctness
The PR adds cross-session origin attribution to context messages. The implementation is correct: deduplication via dict.fromkeys() preserves order, each message gets its own list copy (line 280), original messages are never mutated (verified via shallow copy + new dict for additional_properties + new merged dict), the falsy check on empty list correctly omits the field, and all imports (Mapping, cast, Any) are already present. The merge behavior using setdefault intentionally preserves existing attribution keys, which matches the stated design. No bugs, race conditions, or incorrect API usage found.
✓ Security Reliability
The PR adds cross-session origin attribution to context messages. The implementation is sound: messages are properly copied before mutation, session IDs are deduplicated, and the observer sample includes defensive type checks. One medium concern exists around the merge behavior in
extend_messages: when a message already carries_attribution["origin_session_ids"], the framework-computed value is silently discarded due tosetdefaultsemantics. While not currently exploitable (the memory provider creates fresh messages without pre-existing attribution), this is a fragility in a feature whose stated purpose is governance and audit integrity. No blocking issues found.
✓ Test Coverage
Test coverage for the new cross-session origin attribution feature is solid. The unit tests in test_sessions.py thoroughly cover the SessionContext.extend_messages changes (deduplication, merge behavior, provider objects, default omission, no-mutation of originals). The integration tests in test_harness_memory.py cover the MemoryContextProvider path for both cross-session and current-session-only scenarios. One modest gap: there is no test exercising multiple MemoryTopicRecords each contributing distinct session_ids, which would verify the flattening/dedup logic across multiple topics in the production code path. This is a low-severity nice-to-have since the logic is straightforward.
✓ Failure Modes
The cross-session origin attribution feature is well-implemented with no significant failure modes. The deduplication uses order-preserving
dict.fromkeys, shallow copies correctly isolate per-message attribution dicts and lists, the falsy check on emptyorigin_session_idscorrectly omits the field when no cross-session origins exist, andMemoryTopicRecord.session_idsis guaranteed to be a list (never None) via_dedupe_strings(session_ids or []). The merge path for messages with existing_attributionMappings correctly addsorigin_session_idsviasetdefaultwithout clobbering existing keys. The memory provider's generator guardif contributor and contributor != current_session_idproperly filters both empty strings and the current session.
✗ Design Approach
I found two design issues in the new provenance surface. The framework-level merge path in
SessionContext.extend_messagesdoes not actually mergeorigin_session_idswhen a message already carries its own_atribution, some cross-session contributors can still be silently dropped. Separately, the new sample observer reports thecontext_messagesbucket key as the source instead of the message’s preserved attribution source, which can misidentify the injector on inputs the core tests explicitly support.
Suggestions
python/samples/02-agents/context_providers/cross_session_observer.py:72-87should report the message attribution'ssource_idwhen present rather than thecontext_messagesbucket key, becausetest_sessions.py:147-163showsextend_messages(...)intentionally preserves an existing_attribution["source_id"]even when the message is stored under a different bucket.
Automated review by moonbox3's agents
Python Test Coverage Report •
Python Unit Test Overview
|
||||||||||||||||||||||||||||||||||||||||
There was a problem hiding this comment.
Automated Code Review
Reviewers: 5 | Confidence: 89%
✓ Correctness
The implementation is correct. The cross-session origin attribution logic properly deduplicates session IDs, handles union merging with pre-existing attribution, maintains copy isolation (original messages are never mutated), and correctly omits the field when no cross-session contributors exist. Tests comprehensively cover all paths including the merge, no-origin, and single-session scenarios. No correctness issues found.
✓ Security Reliability
The cross-session origin attribution implementation is well-structured from a security/reliability standpoint. It properly uses shallow copies to avoid mutating caller objects, performs defensive type checks on existing attribution data, and correctly deduplicates session IDs using insertion-order-preserving dict.fromkeys. The merge logic for existing attribution is sound — union-dedup for origin_session_ids and setdefault for other keys preserves audit integrity. No injection risks, resource leaks, or unhandled failure modes identified in the production code paths.
✓ Test Coverage
Test coverage for this PR is thorough. The unit tests cover: default behavior (no origin), deduplication, provider objects, merge/union with existing attribution, and immutability of original messages. Integration tests exercise the memory provider in both cross-session and same-session scenarios. Two minor gaps exist: (1) no test verifies list independence when multiple messages are passed in one
extend_messagescall withorigin_session_ids, and (2) the memory provider integration test uses only one topic record, so cross-record deduplication isn't explicitly exercised (thoughdict.fromkeysguarantees correctness regardless).
✓ Failure Modes
The implementation is well-structured with proper defensive copying and deduplication. The merge logic correctly handles union semantics for existing origin_session_ids, and shallow copies are managed to avoid mutation of originals. The memory provider correctly computes cross-session origins and passes an empty list (falsy) when only the current session contributes, which the framework correctly interprets as "omit the field." No concrete silent failure modes or operational bugs were identified in the changed code paths.
✗ Design Approach
The provenance feature is close, but the current memory-provider implementation can silently under-report cross-session contributors for the injected
MEMORY.mdblock. That message includes both the full memory index and the auto-loaded topic files, while the neworigin_session_idslist is derived only from the auto-loaded subset.
Flagged Issues
- python/packages/core/agent_framework/_harness/_memory.py:1317-1324 computes
origin_session_idsfromselected_topicsonly, but the injected message also includesindex_textbuilt from all indexed topics at :1173-1185 and rendered into the same message at :133-1337. If an indexed topic from another session is not selected for auto-loading, its provenance is silently omitted from the governance/audit signal.
Automated review by moonbox3's agents
|
Flagged issue python/packages/core/agent_framework/_harness/_memory.py:1317-1324 computes Source: automated DevFlow PR review |
Motivation & Context
Context providers can inject memory consolidated from multiple sessions, but downstream governance and audit providers currently cannot identify which sessions contributed that context. This change exposes accurate cross-session provenance instead of collapsing multiple contributors into an arbitrary single session.
Description & Review Guide
SessionContext.extend_messagesaccepts optional ordered, deduplicatedorigin_session_ids; the harness memory provider supplies every distinct non-current contributor; existing attribution keys and origins are preserved; tests and a self-contained observer sample demonstrate the behavior._attribution["origin_session_ids"].Related Issue
Fixes #5914.
Supersedes #5916 because its head is an organization-owned fork that cannot grant upstream maintainers push access. This branch includes the outstanding review fixes and current
main.Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.