Dmytro Struk (dmytrostruk) · Jul 24, 2025 · Jul 22, 2025 · Jul 22, 2025 · Jul 22, 2025 · Jul 22, 2025
diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py
@@ -2,6 +2,7 @@
 
 import contextlib
 import json
+import sys
 from collections.abc import AsyncIterable, MutableMapping, MutableSequence
 from typing import Any, ClassVar
 
@@ -19,6 +20,7 @@
     DataContent,
     FunctionCallContent,
     FunctionResultContent,
+    HostedCodeInterpreterTool,
     TextContent,
     UriContent,
     UsageContent,
@@ -34,6 +36,7 @@
     AgentStreamEvent,
     AsyncAgentEventHandler,
     AsyncAgentRunStream,
+    CodeInterpreterToolDefinition,
     FunctionName,
     ListSortOrder,
     MessageDeltaChunk,
@@ -56,6 +59,11 @@
 from azure.core.credentials_async import AsyncTokenCredential
 from pydantic import Field, PrivateAttr, ValidationError
 
+if sys.version_info >= (3, 11):
+    from typing import Self  # pragma: no cover
+else:
+    from typing_extensions import Self  # pragma: no cover
+
 
 class FoundrySettings(AFBaseSettings):
     """Foundry model settings.
@@ -85,9 +93,12 @@ class FoundrySettings(AFBaseSettings):
 @use_tool_calling
 class FoundryChatClient(ChatClientBase):
     client: AIProjectClient = Field(...)
+    credential: AsyncTokenCredential | None = Field(...)
     agent_id: str | None = Field(default=None)
     thread_id: str | None = Field(default=None)
     _should_delete_agent: bool = PrivateAttr(default=False)  # Track whether we should delete the agent
+    _should_close_client: bool = PrivateAttr(default=False)  # Track whether we should close client connection
+    _should_close_credential: bool = PrivateAttr(default=False)  # Track whether we should close credential
     _foundry_settings: FoundrySettings = PrivateAttr()
 
     def __init__(
@@ -133,6 +144,8 @@ def __init__(
             raise ServiceInitializationError("Failed to create Foundry settings.", ex) from ex
 
         # If no client is provided, create one
+        should_close_client = False
+        should_close_credential = False
         if client is None:
             if not foundry_settings.project_endpoint:
                 raise ServiceInitializationError("Project endpoint is required when client is not provided.")
@@ -145,20 +158,25 @@ def __init__(
                 from azure.identity.aio import DefaultAzureCredential
 
                 credential = DefaultAzureCredential()
+                should_close_credential = True
 
             client = AIProjectClient(endpoint=foundry_settings.project_endpoint, credential=credential)
+            should_close_client = True
 
         super().__init__(
             client=client,  # type: ignore[reportCallIssue]
+            credential=credential,  # type: ignore[reportCallIssue]
             agent_id=agent_id,  # type: ignore[reportCallIssue]
             thread_id=thread_id,  # type: ignore[reportCallIssue]
             **kwargs,
         )
 
         self._should_delete_agent = False
+        self._should_close_client = should_close_client
+        self._should_close_credential = should_close_credential
         self._foundry_settings = foundry_settings
 
-    async def __aenter__(self) -> "FoundryChatClient":
+    async def __aenter__(self) -> "Self":
         """Async context manager entry."""
         return self
 
@@ -169,6 +187,8 @@ async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseExc
     async def close(self) -> None:
         """Close the client and clean up any agents we created."""
         await self._cleanup_agent_if_needed()
+        await self._close_client_if_needed()
+        await self._close_credential_if_needed()
 
     @classmethod
     def from_dict(cls, settings: dict[str, Any]) -> "FoundryChatClient":
@@ -432,6 +452,16 @@ def _create_function_call_contents(self, event_data: ThreadRun, response_id: str
 
         return contents
 
+    async def _close_credential_if_needed(self) -> None:
+        """Close credential if we created it."""
+        if self._should_close_credential and self.credential is not None:
+            await self.credential.close()
+
+    async def _close_client_if_needed(self) -> None:
+        """Close client session if we created it."""
+        if self._should_close_client:
+            await self.client.close()
+
     async def _cleanup_agent_if_needed(self) -> None:
         """Clean up the agent if we created it."""
         if self._should_delete_agent and self.agent_id is not None:
@@ -455,14 +485,15 @@ def _create_run_options(
             run_options["parallel_tool_calls"] = chat_options.allow_multiple_tool_calls
 
             if chat_options.tools is not None:
-                # TODO (eavanvalkenburg): replace with _prepare_tools_and_tool_choice overload
                 tool_definitions: list[MutableMapping[str, Any]] = []
 
                 for tool in chat_options.tools:
                     if isinstance(tool, AIFunction):
                         tool_definitions.append(ai_function_to_json_schema_spec(tool))
-                    else:
-                        tool_definitions.append(tool)  # type: ignore
+                    elif isinstance(tool, HostedCodeInterpreterTool):
+                        tool_definitions.append(CodeInterpreterToolDefinition())
+                    elif isinstance(tool, MutableMapping):
+                        tool_definitions.append(tool)
 
                 if len(tool_definitions) > 0:
                     run_options["tools"] = tool_definitions

@@ -441,9 +441,11 @@ async def run(
             kwargs: Additional keyword arguments for the agent.
                 will only be passed to functions that are called.
         """
+        input_messages = self._normalize_messages(messages)
+
         thread, thread_messages = await self._prepare_thread_and_messages(
             thread=thread,
-            input_messages=messages,
+            input_messages=input_messages,
             construct_thread=lambda: ChatClientAgentThread(),
             expected_type=ChatClientAgentThread,
         )
@@ -453,6 +455,7 @@ async def run(
             chat_options=self.chat_options
             & ChatOptions(
                 ai_model_id=model,
+                conversation_id=thread.id,
                 frequency_penalty=frequency_penalty,
                 logit_bias=logit_bias,
                 max_tokens=max_tokens,
@@ -476,7 +479,7 @@ async def run(
 
         # Only notify the thread of new messages if the chatResponse was successful
         # to avoid inconsistent messages state in the thread.
-        await self._notify_thread_of_new_messages(thread, thread_messages)
+        await self._notify_thread_of_new_messages(thread, input_messages)
         await self._notify_thread_of_new_messages(thread, response.messages)
 
         return AgentRunResponse(
@@ -548,9 +551,11 @@ async def run_stream(
                 will only be passed to functions that are called.
 
         """
+        input_messages = self._normalize_messages(messages)
+
         thread, thread_messages = await self._prepare_thread_and_messages(
             thread=thread,
-            input_messages=messages,
+            input_messages=input_messages,
             construct_thread=lambda: ChatClientAgentThread(),
             expected_type=ChatClientAgentThread,
         )
@@ -561,6 +566,7 @@ async def run_stream(
             messages=thread_messages,
             chat_options=self.chat_options
             & ChatOptions(
+                conversation_id=thread.id,
                 frequency_penalty=frequency_penalty,
                 logit_bias=logit_bias,
                 max_tokens=max_tokens,
@@ -598,7 +604,7 @@ async def run_stream(
 
         # Only notify the thread of new messages if the chatResponse was successful
         # to avoid inconsistent messages state in the thread.
-        await self._notify_thread_of_new_messages(thread, thread_messages)
+        await self._notify_thread_of_new_messages(thread, input_messages)
         await self._notify_thread_of_new_messages(thread, response.messages)
 
     def get_new_thread(self) -> AgentThread:
@@ -637,15 +643,15 @@ async def _prepare_thread_and_messages(
         self,
         *,
         thread: AgentThread | None,
-        input_messages: str | ChatMessage | Sequence[str] | Sequence[ChatMessage] | None = None,
+        input_messages: list[ChatMessage] | None = None,
         construct_thread: Callable[[], TThreadType],
         expected_type: type[TThreadType],
     ) -> tuple[TThreadType, list[ChatMessage]]:
         """Prepare thread and messages for agent execution.
 
         Args:
             thread: The conversation thread, or None to create a new one.
-            input_messages: Messages to process, can be string, ChatMessage, or sequence.
+            input_messages: Messages to process.
             construct_thread: Factory function to create a new thread.
             expected_type: Expected thread type for validation.
 
@@ -671,15 +677,24 @@ async def _prepare_thread_and_messages(
         if isinstance(thread, MessagesRetrievableThread):
             async for message in thread.get_messages():
                 messages.append(message)
+
         if input_messages is None:
             return thread, messages
-        if isinstance(input_messages, str):
-            messages.append(ChatMessage(role=ChatRole.USER, text=input_messages))
-            return thread, messages
-        if isinstance(input_messages, ChatMessage):
-            messages.append(input_messages)
-            return thread, messages
-        messages.extend([
-            ChatMessage(role=ChatRole.USER, text=msg) if isinstance(msg, str) else msg for msg in input_messages
-        ])
+
+        messages.extend(input_messages)
         return thread, messages
+
+    def _normalize_messages(
+        self,
+        messages: str | ChatMessage | Sequence[str] | Sequence[ChatMessage] | None = None,
+    ) -> list[ChatMessage]:
+        if messages is None:
+            return []
+
+        if isinstance(messages, str):
+            return [ChatMessage(role=ChatRole.USER, text=messages)]
+
+        if isinstance(messages, ChatMessage):
+            return [messages]
+
+        return [ChatMessage(role=ChatRole.USER, text=msg) if isinstance(msg, str) else msg for msg in messages]
diff --git a/python/packages/main/agent_framework/_tools.py b/python/packages/main/agent_framework/_tools.py
@@ -1,13 +1,13 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import inspect
-from collections.abc import Awaitable, Callable, Mapping
+from collections.abc import Awaitable, Callable
 from functools import wraps
 from typing import Any, Generic, Protocol, TypeVar, runtime_checkable
 
 from pydantic import BaseModel, create_model
 
-__all__ = ["AIFunction", "AITool", "ai_function"]
+__all__ = ["AIFunction", "AITool", "HostedCodeInterpreterTool", "ai_function"]
 
 
 @runtime_checkable
@@ -34,10 +34,6 @@ def __str__(self) -> str:
         """Return a string representation of the tool."""
         ...
 
-    def parameters(self) -> Mapping[str, Any]:
-        """Return the parameters of the tool as a JSON schema."""
-        ...
-
 
 ArgsT = TypeVar("ArgsT", bound=BaseModel)
 ReturnT = TypeVar("ReturnT")
@@ -159,3 +155,32 @@ def wrapper(f: Callable[..., ReturnT | Awaitable[ReturnT]]) -> AIFunction[Any, R
         return wrapper(func)
 
     return decorator(func) if func else decorator  # type: ignore[reportReturnType, return-value]
+
+
+class HostedCodeInterpreterTool(AITool):
+    """Represents a hosted tool that can be specified to an AI service to enable it to execute generated code.
+
+    This tool does not implement code interpretation itself. It serves as a marker to inform a service
+    that it is allowed to execute generated code if the service is capable of doing so.
+    """
+
+    def __init__(
+        self,
+        name: str = "code_interpreter",
+        description: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+    ):
+        """Initialize a HostedCodeInterpreterTool.
+
+        Args:
+            name: The name of the tool. Defaults to "code_interpreter".
+            description: A description of the tool.
+            additional_properties: Additional properties associated with the tool.
+        """
+        self.name = name
+        self.description = description
+        self.additional_properties = additional_properties
+
+    def __str__(self) -> str:
+        """Return a string representation of the tool."""
+        return f"HostedCodeInterpreterTool(name={self.name})"
@@ -229,12 +229,12 @@ async def test_chat_client_agent_get_new_thread(chat_client: ChatClient) -> None
 
 async def test_chat_client_agent_prepare_thread_and_messages(chat_client: ChatClient) -> None:
     agent = ChatClientAgent(chat_client=chat_client)
-    message = ChatMessage(role=ChatRole.USER, contents=[TextContent("Hello")])
+    message = ChatMessage(role=ChatRole.USER, text="Hello")
     thread = ChatClientAgentThread(messages=[message])
 
     result_thread, result_messages = await agent._prepare_thread_and_messages(  # type: ignore[reportPrivateUsage]
         thread=thread,
-        input_messages="Test",
+        input_messages=[ChatMessage(role=ChatRole.USER, text="Test")],
         construct_thread=lambda: ChatClientAgentThread(),
         expected_type=ChatClientAgentThread,
     )