github-featured/microsoft--agent-framework
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

Python: feat: cross-session origin attribution on context messages (#7041)

Browse Source 
* Python: feat: cross-session origin attribution on context messages

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 #5914 for design discussion.

Surfaced during independent audit conducted by @finnoybu (Ken Tannenbaum, AEGIS Initiative); [MEDIUM, python/packages/core].

* Address cross-session attribution review feedback

* Address follow-up review feedback

* Address follow-up review comments

* Address origin attribution review feedback

---------

Co-authored-by: finnoybu <21694570+finnoybu@users.noreply.github.com>
This commit is contained in:
Evan Mattson
2026-07-15 02:34:52 +09:00
committed by GitHub
parent 47cd0a508d
commit c35a63ed8d
8 changed files with 409 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/core/AGENTS.md
+2 -1
View File
@@ -75,7 +75,8 @@ agent_framework/
- **`AgentSession`** - Manages conversation state and session metadata
- **`ServiceSessionId`** - Mapping alias for structured service-owned continuation handles used in `AgentSession.service_session_id`
- **`SessionContext`** - Context object for session-scoped data during agent runs
- **`SessionContext`** - Context object for session-scoped data during agent runs. `extend_messages(...)` can attach
ordered, deduplicated `origin_session_ids` attribution when a provider injects content from other sessions.
- **`ContextProvider`** - Base class for context providers (RAG, memory systems)
- **`HistoryProvider`** - Base class for conversation history storage
- **`InMemoryHistoryProvider`** - Built-in session-state history provider for local runs
python/packages/core/agent_framework/_harness/_memory.py
+14
View File
@@ -1306,6 +1306,19 @@ class MemoryContextProvider(HistoryProvider):
)
if recent_history_messages:
context.extend_messages(self.source_id, recent_history_messages)
# Surface every cross-session origin so downstream context observers
# can distinguish injected memory from content native to the current
# session. Loaded topic files may carry contributions from multiple
# earlier sessions, tracked in ``MemoryTopicRecord.session_ids``.
current_session_id = context.session_id
cross_session_origins = [
contributor
for record in selected_topics
for contributor in record.session_ids
if contributor and contributor != current_session_id
]
context.extend_messages(
self.source_id,
[
@@ -1322,6 +1335,7 @@ class MemoryContextProvider(HistoryProvider):
],
)
],
origin_session_ids=cross_session_origins,
)
async def after_run(
python/packages/core/agent_framework/_sessions.py
+58 -4
View File
@@ -22,7 +22,7 @@ import uuid
import weakref
from abc import abstractmethod
from base64 import urlsafe_b64encode
from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence
from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, Sequence
from contextlib import suppress
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias, TypeGuard, cast
@@ -66,6 +66,17 @@ def _default_json_loads(value: str | bytes) -> Any:
return json.loads(value)
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
def _is_middleware_sequence(
middleware: MiddlewareTypes | Sequence[MiddlewareTypes],
) -> TypeGuard[Sequence[MiddlewareTypes]]:
@@ -230,7 +241,13 @@ class SessionContext:
"""The agent's response. Set by the framework after invocation, read-only for providers."""
return self._response
def extend_messages(self, source: str | object, messages: Sequence[Message]) -> None:
def extend_messages(
self,
source: str | object,
messages: Sequence[Message],
*,
origin_session_ids: Sequence[str] | None = None,
) -> None:
"""Add context messages from a specific source.
Messages are copied before attribution is added, so the caller's
@@ -245,19 +262,56 @@ class SessionContext:
object is passed, its class name is recorded as
``source_type`` in the attribution.
messages: The messages to add.
Keyword Args:
origin_session_ids: Optional session IDs that originally produced
these messages, when different from the current session. Set
by providers that inject content stored under other sessions
(cross-session memory). The IDs describe the contributing
sessions for every message supplied in this call; they are not
positionally paired with messages, and a composed message can
have multiple origins. The values are exposed under
``additional_properties["_attribution"]["origin_session_ids"]``
so downstream context observers can detect cross-session
content for governance, audit, or behavioral-analysis
purposes. Omit (default) when content originates in the
current session; absence of the field means that no origin
information was supplied.
"""
if isinstance(source, str):
source_id = source
attribution: dict[str, str] = {"source_id": source_id}
attribution: dict[str, Any] = {"source_id": source_id}
else:
source_id = source.source_id # type: ignore[attr-defined]
attribution = {"source_id": source_id, "source_type": type(source).__name__}
if origin_session_ids:
attribution["origin_session_ids"] = _deduplicate_origin_session_ids(origin_session_ids)
copied: list[Message] = []
for message in messages:
msg_copy = copy.copy(message)
msg_copy.additional_properties = dict(message.additional_properties)
msg_copy.additional_properties.setdefault("_attribution", attribution)
message_attribution = dict(attribution)
if "origin_session_ids" in message_attribution:
message_attribution["origin_session_ids"] = list(message_attribution["origin_session_ids"])
existing_attribution = msg_copy.additional_properties.get("_attribution")
if isinstance(existing_attribution, Mapping):
merged_attribution = dict(cast(Mapping[str, Any], existing_attribution))
for key, value in message_attribution.items():
if key == "origin_session_ids":
existing_origins = merged_attribution.get(key)
if isinstance(existing_origins, Sequence) and not isinstance(existing_origins, str):
existing_origin_values = cast(Sequence[Any], existing_origins)
value = _deduplicate_origin_session_ids(
[origin for origin in existing_origin_values if isinstance(origin, str)]
+ cast(list[str], value)
)
merged_attribution[key] = value
else:
merged_attribution.setdefault(key, value)
msg_copy.additional_properties["_attribution"] = merged_attribution
else:
msg_copy.additional_properties.setdefault("_attribution", message_attribution)
copied.append(msg_copy)
if source_id not in self.context_messages:
self.context_messages[source_id] = []
python/packages/core/tests/core/test_harness_memory.py
+88
View File
@@ -505,6 +505,94 @@ async def test_memory_context_provider_recent_turns_can_skip_tool_call_groups(tm
assert with_tools_messages[4].text == "Second final answer"
async def test_memory_context_provider_marks_cross_session_origins(tmp_path) -> None:
"""Injected memory should carry all prior session origins without duplicates.
Exercises the cross-session attribution surface added to support downstream observers
detecting attacks of the class documented in Dai et al. (arXiv:2605.06158).
"""
session = AgentSession(session_id="session-current")
session.state["owner_id"] = "alice"
store = MemoryFileStore(
tmp_path,
owner_state_key="owner_id",
dumps=lambda value: json.dumps(value, separators=(",", ":"), sort_keys=True),
loads=json.loads,
)
updated_at = datetime(2026, 4, 21, tzinfo=timezone.utc).replace(microsecond=0).isoformat()
store.write_topic(
session,
MemoryTopicRecord(
topic="travel preferences",
summary="Loves Oslo trips.",
memories=["Prefers Oslo in summer."],
updated_at=updated_at,
session_ids=["session-current", "session-prior-1", "session-prior-2", "session-prior-1"],
),
source_id=DEFAULT_MEMORY_SOURCE_ID,
)
agent = Agent(
client=_MemoryHarnessClient(), # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type]
context_providers=[MemoryContextProvider(store=store)],
default_options=_no_store_options(),
)
session_context, _ = await agent._prepare_session_and_messages( # type: ignore[reportPrivateUsage]
session=session,
input_messages=[Message(role="user", contents=["Tell me about my travel preferences."])],
)
memory_messages = [
m for m in session_context.context_messages.get(DEFAULT_MEMORY_SOURCE_ID, []) if "### MEMORY.md" in m.text
]
assert memory_messages, "expected an injected memory block under the memory source"
attribution: dict[str, Any] = memory_messages[0].additional_properties.get("_attribution") or {}
assert attribution.get("origin_session_ids") == ["session-prior-1", "session-prior-2"]
async def test_memory_context_provider_omits_origin_when_only_current_session(tmp_path) -> None:
"""When all contributing topics are from the current session, attribution must NOT advertise an origin."""
session = AgentSession(session_id="session-current")
session.state["owner_id"] = "alice"
store = MemoryFileStore(
tmp_path,
owner_state_key="owner_id",
dumps=lambda value: json.dumps(value, separators=(",", ":"), sort_keys=True),
loads=json.loads,
)
updated_at = datetime(2026, 4, 21, tzinfo=timezone.utc).replace(microsecond=0).isoformat()
store.write_topic(
session,
MemoryTopicRecord(
topic="travel preferences",
summary="Loves Oslo trips.",
memories=["Prefers Oslo in summer."],
updated_at=updated_at,
session_ids=["session-current"],
),
source_id=DEFAULT_MEMORY_SOURCE_ID,
)
agent = Agent(
client=_MemoryHarnessClient(), # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type]
context_providers=[MemoryContextProvider(store=store)],
default_options=_no_store_options(),
)
session_context, _ = await agent._prepare_session_and_messages( # type: ignore[reportPrivateUsage]
session=session,
input_messages=[Message(role="user", contents=["Tell me about my travel preferences."])],
)
memory_messages = [
m for m in session_context.context_messages.get(DEFAULT_MEMORY_SOURCE_ID, []) if "### MEMORY.md" in m.text
]
assert memory_messages
attribution: dict[str, Any] = memory_messages[0].additional_properties.get("_attribution") or {}
assert "origin_session_ids" not in attribution
async def test_memory_context_provider_uses_explicit_consolidation_client(tmp_path) -> None:
"""The memory provider should use the explicit consolidation client when one is configured."""
session = AgentSession(session_id="session-1")
python/packages/core/tests/core/test_sessions.py
+86
View File
@@ -107,6 +107,92 @@ class TestSessionContext:
stored = ctx.context_messages["rag"][0]
assert stored.additional_properties["_attribution"] == {"source_id": "rag", "source_type": "MyProvider"}
def test_extend_messages_origin_session_ids_default_omits_field(self) -> None:
ctx = SessionContext(input_messages=[])
msg = Message(role="system", contents=["ctx"])
ctx.extend_messages("rag", [msg])
stored = ctx.context_messages["rag"][0]
# Default (no origin_session_ids passed) preserves the historical attribution shape
# so observers can distinguish "no origin info" from "explicit cross-session marker."
assert "origin_session_ids" not in stored.additional_properties["_attribution"]
def test_extend_messages_origin_session_ids_recorded_on_attribution(self) -> None:
ctx = SessionContext(session_id="current", input_messages=[])
msg = Message(role="system", contents=["loaded from a prior session"])
ctx.extend_messages(
"memory_provider",
[msg],
origin_session_ids=["prior-session-id", "another-session", "prior-session-id"],
)
stored = ctx.context_messages["memory_provider"][0]
assert stored.additional_properties["_attribution"] == {
"source_id": "memory_provider",
"origin_session_ids": ["prior-session-id", "another-session"],
}
def test_extend_messages_origin_session_ids_with_provider_object(self) -> None:
class MyMemoryProvider:
source_id = "memory"
ctx = SessionContext(session_id="current", input_messages=[])
msg = Message(role="assistant", contents=["consolidated memory content"])
ctx.extend_messages(MyMemoryProvider(), [msg], origin_session_ids=["prior"])
stored = ctx.context_messages["memory"][0]
assert stored.additional_properties["_attribution"] == {
"source_id": "memory",
"source_type": "MyMemoryProvider",
"origin_session_ids": ["prior"],
}
def test_extend_messages_applies_all_origins_to_each_message(self) -> None:
ctx = SessionContext(session_id="current", input_messages=[])
messages = [
Message(role="assistant", contents=["first composed memory"]),
Message(role="assistant", contents=["second composed memory"]),
]
ctx.extend_messages("memory_provider", messages, origin_session_ids=["session-a", "session-b"])
stored_messages = ctx.context_messages["memory_provider"]
assert [message.additional_properties["_attribution"] for message in stored_messages] == [
{
"source_id": "memory_provider",
"origin_session_ids": ["session-a", "session-b"],
},
{
"source_id": "memory_provider",
"origin_session_ids": ["session-a", "session-b"],
},
]
def test_extend_messages_adds_origin_to_existing_attribution(self) -> None:
ctx = SessionContext(session_id="current", input_messages=[])
msg = Message(
role="system",
contents=["loaded from a prior session"],
additional_properties={
"_attribution": {
"source_id": "custom",
"custom_key": "value",
"origin_session_ids": ["existing", "prior"],
}
},
)
ctx.extend_messages("memory_provider", [msg], origin_session_ids=["prior", "new"])
stored = ctx.context_messages["memory_provider"][0]
assert stored.additional_properties["_attribution"] == {
"source_id": "custom",
"custom_key": "value",
"origin_session_ids": ["existing", "prior", "new"],
}
assert msg.additional_properties["_attribution"] == {
"source_id": "custom",
"custom_key": "value",
"origin_session_ids": ["existing", "prior"],
}
def test_extend_instructions_string(self) -> None:
ctx = SessionContext(input_messages=[])
ctx.extend_instructions("sys", "Be helpful")
python/samples/02-agents/context_providers/README.md
+4
View File
@@ -7,6 +7,7 @@ These samples demonstrate how to use context providers to enrich agent conversat
| File / Folder | Description |
|---------------|-------------|
| [`simple_context_provider.py`](simple_context_provider.py) | Implement a custom context provider by extending `ContextProvider` to extract and inject structured user information across turns. |
| [`cross_session_observer.py`](cross_session_observer.py) | Detect injected context messages whose origins differ from the current session, via the `Message.additional_properties["_attribution"]["origin_session_ids"]` field. Self-contained — no LLM credentials required. |
| [`azure_ai_foundry_memory.py`](azure_ai_foundry_memory.py) | Use `FoundryMemoryProvider` to add semantic memory — automatically retrieves, searches, and stores memories via Microsoft Foundry. |
| [`file_access_data_processing/`](file_access_data_processing/) | Use `FileAccessProvider` with `FileSystemAgentFileStore` to give an agent read/write/search access to a folder of CSV data files. See its own [README](file_access_data_processing/README.md). |
| [`azure_ai_search/`](azure_ai_search/) | Retrieval Augmented Generation (RAG) with Azure AI Search in semantic and agentic modes. See its own [README](azure_ai_search/README.md). |
@@ -15,6 +16,9 @@ These samples demonstrate how to use context providers to enrich agent conversat
## Prerequisites
**For `cross_session_observer.py`:**
- No external dependencies; runs against in-memory `SessionContext`.
**For `simple_context_provider.py`:**
- `FOUNDRY_PROJECT_ENDPOINT`: Your Microsoft Foundry project endpoint
- `FOUNDRY_MODEL`: Model deployment name
python/samples/02-agents/context_providers/cross_session_observer.py
+156
View File
@@ -0,0 +1,156 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
from collections.abc import Callable, Mapping, Sequence
from typing import Any, cast
from agent_framework import AgentSession, ContextProvider, Message, SessionContext
"""This sample demonstrates how to detect cross-session memory injection.
When a context provider injects messages from a different ``session_id`` than
the requesting one — the legitimate cross-session memory use case (consolidated
memories, Mem0 with default scope, shared knowledge bases) — the framework
records the originating sessions under
``message.additional_properties["_attribution"]["origin_session_ids"]``.
Downstream context observers can subscribe to this signal for governance,
audit, and behavioral analysis purposes. This is useful for defending against
the stateful-agent-backdoor attack class documented in Dai et al.,
arXiv:2605.06158, in which an adversary chains sub-backdoors across sessions
under permission isolation via persisted memory state.
The sample is self-contained: it constructs ``SessionContext`` directly and
invokes provider lifecycle methods manually, so no LLM credentials are
required to run it.
"""
class CrossSessionObserver(ContextProvider):
"""Detect injected context messages whose origin differs from the current session.
Subscribes via the standard ``ContextProvider`` pipeline. In ``before_run``,
walks the accumulated context messages and invokes a user-supplied
callback for each message whose ``_attribution["origin_session_ids"]``
contains one or more sessions other than the current ``session_id``.
The callback receives the source_id that injected the content, the
originating session IDs, the current session_id, and the message itself.
Use it to log, alert, increment metrics, or enforce policy — the observer
itself only surfaces the signal, leaving the response policy to the caller.
"""
DEFAULT_SOURCE_ID = "cross_session_observer"
def __init__(
self,
on_cross_session_access: Callable[[str, Sequence[str], str | None, Message], None],
*,
source_id: str = DEFAULT_SOURCE_ID,
) -> None:
"""Initialize the observer.
Args:
on_cross_session_access: Callback invoked for each detected
cross-session message. Signature is
``(source_id, origin_session_ids, current_session_id, message)``.
source_id: Unique identifier for this observer instance.
"""
super().__init__(source_id)
self._on_cross_session_access = on_cross_session_access
async def before_run(
self,
*,
agent: Any,
session: AgentSession | None,
context: SessionContext,
state: dict[str, Any],
) -> None:
"""Inspect accumulated context messages for cross-session origin."""
current_session_id = context.session_id
for source_id, messages in context.context_messages.items():
if source_id == self.source_id:
continue
for message in messages:
attribution_raw = message.additional_properties.get("_attribution")
if not isinstance(attribution_raw, Mapping):
continue
attribution = cast(Mapping[str, Any], attribution_raw)
origins = attribution.get("origin_session_ids")
if not isinstance(origins, Sequence) or isinstance(origins, str):
continue
cross_session_origins = [
origin for origin in origins if isinstance(origin, str) and origin != current_session_id
]
if cross_session_origins:
self._on_cross_session_access(source_id, cross_session_origins, current_session_id, message)
def _on_detected(source_id: str, origins: Sequence[str], current: str | None, message: Message) -> None:
"""Sample callback that logs cross-session detections to stdout."""
preview = " ".join(message.text.split())[:80]
print(
f"[cross-session detected] source={source_id!r} "
f"origin_sessions={list(origins)!r} current_session={current!r} "
f"preview={preview!r}"
)
async def main() -> None:
"""Demonstrate the observer firing on cross-session injection."""
observer = CrossSessionObserver(_on_detected)
# --- Case 1: same-session injection (observer should be silent) ---
same_session_context = SessionContext(
session_id="session-A",
input_messages=[Message("user", ["What did we discuss last time?"])],
)
# Simulate a same-session provider injecting same-session history. Omitting
# origin_session_ids means "no origin info"; observers treat it as equivalent
# to same-session for backward compatibility.
same_session_context.extend_messages(
"history_provider",
[Message("assistant", ["We talked about Q3 revenue projections."])],
)
await observer.before_run(
agent=None,
session=None,
context=same_session_context,
state={},
)
print("--- Same-session case complete (no detections expected above) ---\n")
# --- Case 2: cross-session injection (observer should fire) ---
cross_session_context = SessionContext(
session_id="session-B",
input_messages=[Message("user", ["Continue from where we left off."])],
)
# Simulate a cross-session memory provider injecting content originally
# written in sessions A and C while we're now running in session B.
cross_session_context.extend_messages(
"memory_provider",
[Message("assistant", ["Remember: API key for prod is <REDACTED> (from prior sessions)."])],
origin_session_ids=["session-A", "session-C"],
)
await observer.before_run(
agent=None,
session=None,
context=cross_session_context,
state={},
)
print("--- Cross-session case complete (one detection expected above) ---")
if __name__ == "__main__":
asyncio.run(main())
"""
Sample output:
--- Same-session case complete (no detections expected above) ---
[cross-session detected] source='memory_provider' origin_sessions=['session-A', 'session-C'] \
current_session='session-B' preview='Remember: API key for prod is <REDACTED> (from prior sessions).'
--- Cross-session case complete (one detection expected above) ---
"""
python/samples/02-agents/observability/microsoft_opentelemetry_distro.py
+1 -3
View File
@@ -54,9 +54,7 @@ async def main():
"Why is the sky blue?",
]
with get_tracer().start_as_current_span(
"Scenario: Agent Chat", kind=SpanKind.CLIENT
) as current_span:
with get_tracer().start_as_current_span("Scenario: Agent Chat", kind=SpanKind.CLIENT) as current_span:
print(f"Trace ID: {format_trace_id(current_span.get_span_context().trace_id)}")
agent = Agent(