Skip to content

Python: Fix FHA session ID traslation #7601

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Tao Chen (TaoChenOSU) wants to merge 1 commit into from
+62 −71
Closed

Python: Fix FHA session ID traslation #7601

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions python/packages/foundry/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,9 +126,10 @@ session = await agent.create_conversation()
response = await agent.run("Help me plan a trip to Seattle.", session=session)
```

This is separate from hosted-agent `isolation_key` sessions: the created
conversation ID is stored on `AgentSession.service_session_id`, while the local
`session_id` remains available for application/session storage.
This is separate from hosted-agent compute sessions. The conversation ID is
stored on `AgentSession.service_session_id`, while a pre-created Foundry
`agent_session_id` is stored in `AgentSession.state`. Reusing a Foundry session
routes requests to the same hosted compute and does not continue a conversation.

## Publishing an agent as a Foundry prompt agent

Expand Down
72 changes: 29 additions & 43 deletions python/packages/foundry/agent_framework_foundry/_agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,8 @@

logger: logging.Logger = logging.getLogger("agent_framework.foundry")

_FOUNDRY_AGENT_SESSION_ID_STATE_KEY = "agent_framework_foundry.agent_session_id"

AzureTokenProvider = Callable[[], str | Awaitable[str]]
AzureCredentialTypes = TokenCredential | AsyncTokenCredential

Expand Down Expand Up @@ -128,16 +130,6 @@ def _merge_extra_body(extra_body: Any | None, *, additions: Mapping[str, Any] |
return merged


def _uses_foundry_agent_session(conversation_id: Any) -> bool:
"""Return whether a conversation_id should be treated as a Foundry agent session id."""
return (
isinstance(conversation_id, str)
and bool(conversation_id)
and not conversation_id.startswith("resp_")
and not conversation_id.startswith("conv_")
)


def _build_agent_reference(agent_name: str, agent_version: str | None) -> dict[str, str]:
"""Build the Responses API ``agent_reference`` payload for non-preview Foundry agent calls.

Expand Down Expand Up @@ -375,23 +367,14 @@ async def _prepare_options(
run_options["input"] = self._transform_input_for_azure_ai(cast(list[dict[str, Any]], run_options["input"]))

# Merge caller-supplied extra_body with any agent-specific request payload.
conversation_id = options.get("conversation_id")
extra_body = _merge_extra_body(run_options.pop("extra_body", None))
if _uses_foundry_agent_session(conversation_id):
run_options.pop("previous_response_id", None)
run_options.pop("conversation", None)
run_options.pop("model", None)
extra_body["agent_session_id"] = conversation_id
# Non-preview Prompt/Hosted Agent calls need agent_reference in the request body to
# tell the Responses API which Foundry agent (and version) is in use, since ``model``
# is stripped below. The preview path injects the reference via the OpenAI client kwarg
# ``agent_name`` instead, so skip there. See issue #5582.
if not self.allow_preview:
extra_body.setdefault("agent_reference", _build_agent_reference(self.agent_name, self.agent_version))
should_strip_model = _uses_foundry_agent_session(conversation_id) or (
conversation_id is None and not options.get("model")
)
if should_strip_model:
if not options.get("model"):
run_options.pop("model", None)
if extra_body:
run_options["extra_body"] = extra_body
Expand Down Expand Up @@ -425,10 +408,7 @@ def _parse_response_from_openai(
response: Any,
options: dict[str, Any],
) -> Any:
parsed_response = super()._parse_response_from_openai(response, options)
if _uses_foundry_agent_session(options.get("conversation_id")):
parsed_response.conversation_id = None
return parsed_response
return super()._parse_response_from_openai(response, options)

@override
def _parse_chunk_from_openai(
Expand All @@ -438,7 +418,7 @@ def _parse_chunk_from_openai(
function_call_ids: dict[int, tuple[str, str]],
seen_reasoning_delta_item_ids: set[str] | None = None,
) -> ChatResponseUpdate:
"""Parse streaming events while preserving hosted-agent session state."""
"""Parse streaming events, including Foundry OAuth consent events."""
update = try_parse_oauth_consent_event(event, self.model)
if update is None:
update = super()._parse_chunk_from_openai(
Expand All @@ -447,8 +427,6 @@ def _parse_chunk_from_openai(
function_call_ids,
seen_reasoning_delta_item_ids,
)
if _uses_foundry_agent_session(options.get("conversation_id")):
update.conversation_id = None
return update

@override
Expand Down Expand Up @@ -752,7 +730,7 @@ def __init__(
additional_properties=dict(additional_properties) if additional_properties is not None else None,
)

def _resolve_service_session_isolation_key(self, isolation_key: str | None = None) -> str:
def _resolve_agent_session_isolation_key(self, isolation_key: str | None = None) -> str:
"""Resolve the isolation key from an explicit value or default_options."""
resolved_isolation_key = (
isolation_key if isolation_key is not None else self.default_options.get("isolation_key")
Expand All @@ -761,20 +739,20 @@ def _resolve_service_session_isolation_key(self, isolation_key: str | None = Non
raise ValueError("isolation_key is required. Pass it explicitly or set default_options['isolation_key'].")
return resolved_isolation_key

async def _create_service_session_id(
async def _create_agent_session_id(
self,
*,
isolation_key: str | None = None,
) -> str:
"""Create a hosted Foundry service session and return the service session ID."""
"""Create a hosted Foundry compute session and return its agent session ID."""
if not isinstance(self.client, RawFoundryAgentChatClient):
raise TypeError("_create_service_session_id requires a RawFoundryAgentChatClient-based client.")
raise TypeError("_create_agent_session_id requires a RawFoundryAgentChatClient-based client.")
if not self.client.allow_preview:
raise RuntimeError("Hosted Foundry service sessions require allow_preview=True.")

create_session_kwargs: dict[str, Any] = {
"agent_name": self.client.agent_name,
"isolation_key": self._resolve_service_session_isolation_key(isolation_key),
"isolation_key": self._resolve_agent_session_isolation_key(isolation_key),
}
if version := await self.client.get_agent_version():
from azure.ai.projects.models import VersionRefIndicator
Expand Down Expand Up @@ -831,14 +809,20 @@ async def _prepare_run_context(
**{key: value for key, value in runtime_options.items() if value is not None},
}

if (
session is not None
and session.service_session_id is None
and effective_options.get("isolation_key") is not None
):
session.service_session_id = await self._create_service_session_id(
isolation_key=cast(str | None, effective_options.get("isolation_key")),
)
if session is not None:
agent_session_id = session.state.get(_FOUNDRY_AGENT_SESSION_ID_STATE_KEY)
if agent_session_id is None and effective_options.get("isolation_key") is not None:
agent_session_id = await self._create_agent_session_id(
isolation_key=cast(str | None, effective_options.get("isolation_key")),
)
session.state[_FOUNDRY_AGENT_SESSION_ID_STATE_KEY] = agent_session_id
Comment on lines +814 to +818
if agent_session_id is not None:
if not isinstance(agent_session_id, str) or not agent_session_id:
raise ValueError("Foundry agent session state must contain a non-empty string ID.")
runtime_options["extra_body"] = _merge_extra_body(
effective_options.get("extra_body"),
additions={"agent_session_id": agent_session_id},
)

return await super()._prepare_run_context(
messages=messages,
Expand Down Expand Up @@ -999,10 +983,12 @@ def __init__(
already be configured for preview APIs before being passed to
``FoundryAgent``.

To lazily create HostedAgent service sessions inside the agent, pass an
To lazily create a HostedAgent compute session inside the agent, pass an
``isolation_key`` through ``default_options`` (or per-run options). The
agent stores the resulting HostedAgent session ID in
``AgentSession.service_session_id`` and reuses it on subsequent runs.
agent stores the resulting Foundry ``agent_session_id`` in
``AgentSession.state`` and reuses it on subsequent runs. This is separate
from ``AgentSession.service_session_id``, which tracks conversation
continuation.

Keyword Args:
project_endpoint: The Foundry project endpoint URL.
Expand Down
54 changes: 29 additions & 25 deletions python/packages/foundry/tests/foundry/test_foundry_agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -537,8 +537,8 @@ async def test_raw_foundry_agent_chat_client_prepare_options_no_tool_warning_whe
assert not any("cannot be sent when an agent is specified" in record.message for record in caplog.records)


async def test_raw_foundry_agent_chat_client_prepare_options_strips_model_for_hosted_session() -> None:
"""Test that model is stripped when using a hosted agent session (not a PromptAgent)."""
async def test_raw_foundry_agent_chat_client_preserves_separate_compute_and_conversation_ids() -> None:
"""Test that compute-session and conversation IDs use separate request fields."""

mock_project = MagicMock()
mock_openai = MagicMock()
Expand All @@ -555,15 +555,16 @@ async def test_raw_foundry_agent_chat_client_prepare_options_strips_model_for_ho
return_value={
"model": "gpt-4.1",
"previous_response_id": "resp_abc",
"extra_body": {"agent_session_id": "agent-session-123"},
},
):
result = await client._prepare_options(
messages=[Message(role="user", contents="hi")],
options={"conversation_id": "agent-session-123"},
options={"conversation_id": "resp_abc"},
)

assert "model" not in result
assert "previous_response_id" not in result
assert result["previous_response_id"] == "resp_abc"
assert result["extra_body"]["agent_session_id"] == "agent-session-123"
assert result["extra_body"]["agent_reference"] == {"name": "test-agent", "type": "agent_reference"}

Expand Down Expand Up @@ -745,8 +746,8 @@ async def test_raw_foundry_agent_chat_client_prepare_options_preserves_model_for
assert "extra_body" not in result or "agent_session_id" not in result.get("extra_body", {})


async def test_raw_foundry_agent_chat_client_prepare_options_maps_agent_session_id_to_extra_body() -> None:
"""Test that service_session_id is forwarded as agent_session_id for hosted sessions."""
async def test_raw_foundry_agent_chat_client_prepare_options_preserves_conversation_continuation() -> None:
"""Test that conversation IDs are not treated as hosted-agent compute session IDs."""

mock_project = MagicMock()
mock_openai = MagicMock()
Expand All @@ -767,21 +768,19 @@ async def test_raw_foundry_agent_chat_client_prepare_options_maps_agent_session_
):
result = await client._prepare_options(
messages=[Message(role="user", contents="hi")],
options={"conversation_id": "agent-session-123", "isolation_key": "iso-key"},
options={"conversation_id": "conversation-123", "isolation_key": "iso-key"},
)

assert result["extra_body"] == {
"custom": "value",
"agent_session_id": "agent-session-123",
"agent_reference": {"name": "test-agent", "type": "agent_reference"},
}
assert "previous_response_id" not in result
assert "conversation" not in result
assert result["previous_response_id"] == "should-be-removed"
assert "isolation_key" not in result


def test_raw_foundry_agent_chat_client_parse_response_suppresses_conversation_id_for_agent_sessions() -> None:
"""Test that agent-session continuations do not overwrite session.service_session_id."""
def test_raw_foundry_agent_chat_client_parse_response_preserves_conversation_id() -> None:
"""Test that response conversation IDs remain available for framework continuation."""

mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
Expand All @@ -798,14 +797,14 @@ def test_raw_foundry_agent_chat_client_parse_response_suppresses_conversation_id
):
result = client._parse_response_from_openai(
response=MagicMock(),
options={"conversation_id": "agent-session-123"},
options={"conversation_id": "conversation-123"},
)

assert result.conversation_id is None
assert result.conversation_id == "resp_123"


def test_raw_foundry_agent_chat_client_parse_chunk_suppresses_conversation_id_for_agent_sessions() -> None:
"""Test that agent-session stream updates do not overwrite session.service_session_id."""
def test_raw_foundry_agent_chat_client_parse_chunk_preserves_conversation_id() -> None:
"""Test that streaming conversation IDs remain available for framework continuation."""

mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
Expand All @@ -822,11 +821,11 @@ def test_raw_foundry_agent_chat_client_parse_chunk_suppresses_conversation_id_fo
):
result = client._parse_chunk_from_openai(
event=MagicMock(type="response.output_text.delta"),
options={"conversation_id": "agent-session-123"},
options={"conversation_id": "conversation-123"},
function_call_ids={},
)

assert result.conversation_id is None
assert result.conversation_id == "resp_123"


def test_raw_foundry_agent_chat_client_check_model_presence_is_noop() -> None:
Expand Down Expand Up @@ -1051,8 +1050,8 @@ def my_func() -> str:
assert agent.default_options.get("tools") is not None


async def test_raw_foundry_agent_prepare_run_context_creates_service_session_from_isolation_key() -> None:
"""Test that RawFoundryAgent lazily creates a hosted session and stores it on service_session_id."""
async def test_raw_foundry_agent_prepare_run_context_creates_agent_session_from_isolation_key() -> None:
"""Test that a compute session is stored in state without replacing the conversation ID."""

mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
Expand All @@ -1069,7 +1068,7 @@ async def test_raw_foundry_agent_prepare_run_context_creates_service_session_fro
agent_version="1.0",
allow_preview=True,
)
session = AgentSession()
session = AgentSession(service_session_id="conv_123")

with patch(
"agent_framework._agents.RawAgent._prepare_run_context",
Expand All @@ -1087,16 +1086,21 @@ async def test_raw_foundry_agent_prepare_run_context_creates_service_session_fro
)

assert result == {"ok": True}
assert session.service_session_id == "agent-session-123"
assert session.service_session_id == "conv_123"
assert session.state["agent_framework_foundry.agent_session_id"] == "agent-session-123"
mock_project.beta.agents.create_session.assert_awaited_once()
create_session_kwargs = mock_project.beta.agents.create_session.await_args.kwargs
assert create_session_kwargs["agent_name"] == "test-agent"
assert create_session_kwargs["isolation_key"] == "iso-key"
assert "version_indicator" in create_session_kwargs
assert mock_prepare_run_context.await_args.kwargs["options"] == {
"isolation_key": "iso-key",
"extra_body": {"agent_session_id": "agent-session-123"},
}
mock_prepare_run_context.assert_awaited_once()


async def test_raw_foundry_agent_create_service_session_uses_stable_agents_operations() -> None:
async def test_raw_foundry_agent_create_agent_session_uses_stable_agents_operations() -> None:
"""Test that hosted sessions use the stable agents operations when available."""

create_session = AsyncMock(return_value=SimpleNamespace(agent_session_id="agent-session-123"))
Expand All @@ -1111,7 +1115,7 @@ async def test_raw_foundry_agent_create_service_session_uses_stable_agents_opera
allow_preview=True,
)

result = await agent._create_service_session_id(isolation_key="iso-key")
result = await agent._create_agent_session_id(isolation_key="iso-key")

assert result == "agent-session-123"
create_session.assert_awaited_once()
Expand All @@ -1122,7 +1126,7 @@ async def test_raw_foundry_agent_create_service_session_uses_stable_agents_opera
assert "version_indicator" in create_session_kwargs


async def test_raw_foundry_agent_prepare_run_context_requires_preview_for_hosted_sessions() -> None:
async def test_raw_foundry_agent_prepare_run_context_requires_preview_for_agent_sessions() -> None:
"""Test that hosted-agent sessions require allow_preview=True."""

mock_project = MagicMock()
Expand Down
Loading