diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 90fb9754e2a..a2e846a207f 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -23,6 +23,8 @@ updates: labels: - ".NET" - "dependencies" + cooldown: + default-days: 7 # Maintain dependencies for python. # TODO: Remove these Python Dependabot entries after we have confidence in the @@ -35,6 +37,9 @@ updates: labels: - "python" - "dependencies" + cooldown: + default-days: 7 + - package-ecosystem: "uv" directory: "python/" schedule: @@ -43,6 +48,8 @@ updates: labels: - "python" - "dependencies" + cooldown: + default-days: 7 # Maintain dependencies for github-actions - package-ecosystem: "github-actions" diff --git a/.gitignore b/.gitignore index 9a528c309bb..34bf943e24a 100644 --- a/.gitignore +++ b/.gitignore @@ -249,6 +249,9 @@ python/dotnet-ref dotnet/filtered-*.slnx **/*.lscache +# Local .NET build overrides +/dotnet/Directory.Build.rsp + # Local tool state .omc/ .omx/ diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 1b0d32bb357..0b0b4bec73e 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -167,6 +167,18 @@ Each language has its own dev setup guide, coding standards, and build scripts: - Integration tests: `dotnet test --filter-query "/*IntegrationTests*/*/*/*"` (requires API keys/endpoints) - Linting (auto-fix): `dotnet format` +#### Microsoft Internal Feed Proxy for GitHub Copilot SDK (.NET) + +Microsoft contributors can route GitHub Copilot SDK npm downloads through the internal proxy without passing extra `dotnet` arguments: + +1. Create `dotnet/Directory.Build.rsp` with: + + ```text + -p:CopilotNpmRegistryUrl=https://packagefeedproxy.microsoft.io/npm/ + ``` + +When running `dotnet build` (or other `dotnet` commands) from the `./dotnet` directory, this property is applied automatically. + ### PR - CI Process The continuous integration (CI) system will automatically perform the required diff --git a/dotnet/.gitignore b/dotnet/.gitignore index 572680831e2..60f723c18b8 100644 --- a/dotnet/.gitignore +++ b/dotnet/.gitignore @@ -83,8 +83,6 @@ StyleCopReport.xml *.pgc *.pgd *.rsp -# but not Directory.Build.rsp, as it configures directory-level build defaults -!Directory.Build.rsp *.sbr *.tlb *.tli diff --git a/dotnet/Directory.Build.props b/dotnet/Directory.Build.props index 6b73159828e..8b1b6330f05 100644 --- a/dotnet/Directory.Build.props +++ b/dotnet/Directory.Build.props @@ -19,6 +19,25 @@ false false + + + + <_AspireUseSdkPickBestRid>false + $(NoWarn);ASPIRE010 + diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props index 54f8ef550aa..2e6e1bd413c 100644 --- a/dotnet/Directory.Packages.props +++ b/dotnet/Directory.Packages.props @@ -7,7 +7,7 @@ - 13.1.0 + 13.5.2 @@ -77,8 +77,7 @@ - - + @@ -89,24 +88,24 @@ - - - - - - + + + + + + - - + + - + - + - + - - + + diff --git a/dotnet/agent-framework-dotnet.slnx b/dotnet/agent-framework-dotnet.slnx index 403c61f54f5..0ba3890a2d6 100644 --- a/dotnet/agent-framework-dotnet.slnx +++ b/dotnet/agent-framework-dotnet.slnx @@ -402,7 +402,8 @@ - + + diff --git a/dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTools/A2AAgent_AsFunctionTools.csproj b/dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTool/A2AAgent_AsFunctionTool.csproj similarity index 77% rename from dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTools/A2AAgent_AsFunctionTools.csproj rename to dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTool/A2AAgent_AsFunctionTool.csproj index d91b20e34bc..b23848c138d 100644 --- a/dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTools/A2AAgent_AsFunctionTools.csproj +++ b/dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTool/A2AAgent_AsFunctionTool.csproj @@ -10,14 +10,12 @@ - - - + diff --git a/dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTool/Program.cs b/dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTool/Program.cs new file mode 100644 index 00000000000..9216c6c19c7 --- /dev/null +++ b/dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTool/Program.cs @@ -0,0 +1,35 @@ +// Copyright (c) Microsoft. All rights reserved. + +// This sample shows how to expose an A2A agent as a function tool that another agent can call. + +using A2A; +using Azure.AI.Projects; +using Azure.Identity; +using Microsoft.Agents.AI; + +var endpoint = Environment.GetEnvironmentVariable("FOUNDRY_PROJECT_ENDPOINT") ?? throw new InvalidOperationException("FOUNDRY_PROJECT_ENDPOINT is not set."); +var model = Environment.GetEnvironmentVariable("FOUNDRY_MODEL") ?? "gpt-5.4-mini"; +var a2aAgentHost = Environment.GetEnvironmentVariable("A2A_AGENT_HOST") ?? throw new InvalidOperationException("A2A_AGENT_HOST is not set."); + +// Initialize an A2ACardResolver to get an A2A agent card. +A2ACardResolver agentCardResolver = new(new Uri(a2aAgentHost)); + +// Get the agent card +AgentCard agentCard = await agentCardResolver.GetAgentCardAsync(); + +// Create an instance of the AIAgent for an existing A2A agent specified by the agent card. +AIAgent a2aAgent = agentCard.AsAIAgent(); + +// Create the main agent, and provide the A2A agent as a function tool. +// WARNING: DefaultAzureCredential is convenient for development but requires careful consideration in production. +// In production, consider using a specific credential (e.g., ManagedIdentityCredential) to avoid +// latency issues, unintended credential probing, and potential security risks from fallback mechanisms. +AIAgent agent = new AIProjectClient(new Uri(endpoint), new DefaultAzureCredential()) + .AsAIAgent( + model: model, + instructions: "You are a helpful assistant. Use the available tools to answer the user's questions.", + tools: [a2aAgent.AsAIFunction()] + ); + +// Invoke the agent and output the text result. +Console.WriteLine(await agent.RunAsync("What can you help me with?")); diff --git a/dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTool/README.md b/dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTool/README.md new file mode 100644 index 00000000000..be055754606 --- /dev/null +++ b/dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTool/README.md @@ -0,0 +1,30 @@ +# A2A Agent as a Function Tool + +This sample demonstrates how to expose an A2A agent as a single function tool that another AI agent can call. + +The sample: + +- Discovers a remote A2A agent from its agent card +- Converts the agent card to an `AIAgent` +- Converts that agent to one function tool with `AsAIFunction()` +- Registers the function tool with a main AI agent + +## Prerequisites + +Before you begin, ensure you have the following prerequisites: + +- .NET 10 SDK or later +- Access to the A2A agent host service +- A Microsoft Foundry project with a deployed model + +**Note**: These samples need to be run against a valid A2A server. If no A2A server is available, they can be run against the echo-agent that can be spun up locally by following the guidelines at: https://github.com/a2aproject/a2a-dotnet/blob/main/samples/AgentServer/README.md + +Set the following environment variables: + +Authenticate with Azure CLI by running `az login`, then set: + +```powershell +$env:A2A_AGENT_HOST="https://your-a2a-agent-host" # Replace with your A2A agent host endpoint +$env:FOUNDRY_PROJECT_ENDPOINT="https://your-project.services.ai.azure.com/api/projects/your-project" # Replace with your Foundry project endpoint +$env:FOUNDRY_MODEL="gpt-5.4-mini" # Optional, defaults to gpt-5.4-mini +``` diff --git a/dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTools/README.md b/dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTools/README.md deleted file mode 100644 index 33b5f692f2e..00000000000 --- a/dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTools/README.md +++ /dev/null @@ -1,22 +0,0 @@ -# A2A Agent as Function Tools - -This sample demonstrates how to represent an A2A agent as a set of function tools, where each function tool corresponds to a skill of the A2A agent, -and register these function tools with another AI agent so it can leverage the A2A agent's skills. - -# Prerequisites - -Before you begin, ensure you have the following prerequisites: - -- .NET 10 SDK or later -- Access to the A2A agent host service - -**Note**: These samples need to be run against a valid A2A server. If no A2A server is available, they can be run against the echo-agent that can be -spun up locally by following the guidelines at: https://github.com/a2aproject/a2a-dotnet/blob/main/samples/AgentServer/README.md - -Set the following environment variables: - -```powershell -$env:A2A_AGENT_HOST="https://your-a2a-agent-host" # Replace with your A2A agent host endpoint -$env:AZURE_OPENAI_ENDPOINT="https://your-resource.openai.azure.com/" # Replace with your Azure OpenAI resource endpoint -$env:AZURE_OPENAI_DEPLOYMENT_NAME="gpt-5.4-mini" # Optional, defaults to gpt-5.4-mini -``` \ No newline at end of file diff --git a/dotnet/samples/02-agents/A2A/A2AAgent_Skills/A2AAgent_Skills.csproj b/dotnet/samples/02-agents/A2A/A2AAgent_Skills/A2AAgent_Skills.csproj new file mode 100644 index 00000000000..b23848c138d --- /dev/null +++ b/dotnet/samples/02-agents/A2A/A2AAgent_Skills/A2AAgent_Skills.csproj @@ -0,0 +1,21 @@ + + + + Exe + net10.0 + + enable + enable + + + + + + + + + + + + + diff --git a/dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTools/Program.cs b/dotnet/samples/02-agents/A2A/A2AAgent_Skills/Program.cs similarity index 77% rename from dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTools/Program.cs rename to dotnet/samples/02-agents/A2A/A2AAgent_Skills/Program.cs index d0b71ef7855..cf46ab51805 100644 --- a/dotnet/samples/02-agents/A2A/A2AAgent_AsFunctionTools/Program.cs +++ b/dotnet/samples/02-agents/A2A/A2AAgent_Skills/Program.cs @@ -1,39 +1,30 @@ // Copyright (c) Microsoft. All rights reserved. -// This sample shows how to represent an A2A agent as a set of function tools, where each function tool -// corresponds to a skill of the A2A agent, and register these function tools with another AI agent so -// it can leverage the A2A agent's skills. +// This sample shows how to expose each skill advertised by an A2A agent as a separate function tool. using System.Text.RegularExpressions; using A2A; -using Azure.AI.OpenAI; +using Azure.AI.Projects; using Azure.Identity; using Microsoft.Agents.AI; using Microsoft.Extensions.AI; -using OpenAI.Chat; -var endpoint = Environment.GetEnvironmentVariable("AZURE_OPENAI_ENDPOINT") ?? throw new InvalidOperationException("AZURE_OPENAI_ENDPOINT is not set."); -var deploymentName = Environment.GetEnvironmentVariable("AZURE_OPENAI_DEPLOYMENT_NAME") ?? "gpt-5.4-mini"; +var endpoint = Environment.GetEnvironmentVariable("FOUNDRY_PROJECT_ENDPOINT") ?? throw new InvalidOperationException("FOUNDRY_PROJECT_ENDPOINT is not set."); +var model = Environment.GetEnvironmentVariable("FOUNDRY_MODEL") ?? "gpt-5.4-mini"; var a2aAgentHost = Environment.GetEnvironmentVariable("A2A_AGENT_HOST") ?? throw new InvalidOperationException("A2A_AGENT_HOST is not set."); -// Initialize an A2ACardResolver to get an A2A agent card. +// Resolve the remote A2A agent and its advertised skills. A2ACardResolver agentCardResolver = new(new Uri(a2aAgentHost)); - -// Get the agent card AgentCard agentCard = await agentCardResolver.GetAgentCardAsync(); - -// Create an instance of the AIAgent for an existing A2A agent specified by the agent card. AIAgent a2aAgent = agentCard.AsAIAgent(); -// Create the main agent, and provide the a2a agent skills as a function tools. +// Create the main agent, and provide each A2A skill as a separate function tool. // WARNING: DefaultAzureCredential is convenient for development but requires careful consideration in production. // In production, consider using a specific credential (e.g., ManagedIdentityCredential) to avoid // latency issues, unintended credential probing, and potential security risks from fallback mechanisms. -AIAgent agent = new AzureOpenAIClient( - new Uri(endpoint), - new DefaultAzureCredential()) - .GetChatClient(deploymentName) +AIAgent agent = new AIProjectClient(new Uri(endpoint), new DefaultAzureCredential()) .AsAIAgent( + model: model, instructions: "You are a helpful assistant that helps people with travel planning.", tools: [.. CreateFunctionTools(a2aAgent, agentCard)] ); diff --git a/dotnet/samples/02-agents/A2A/A2AAgent_Skills/README.md b/dotnet/samples/02-agents/A2A/A2AAgent_Skills/README.md new file mode 100644 index 00000000000..06238052f04 --- /dev/null +++ b/dotnet/samples/02-agents/A2A/A2AAgent_Skills/README.md @@ -0,0 +1,34 @@ +# A2A Agent Skills + +This sample demonstrates how to expose each skill advertised by an A2A agent as a separate function tool for another AI agent. + +Unlike exposing the whole A2A agent as one tool, this approach gives the model a distinct tool for each skill. The sample uses the skill metadata from the A2A agent card to name and describe those tools. + +All generated tools invoke the same A2A agent. Their names and descriptions help the main agent choose an advertised capability; they do not invoke separate skill-specific endpoints. + +The sample: + +- Discovers a remote A2A agent and its skills from the agent card +- Creates one function tool for each advertised skill +- Adds the skill's description, tags, examples, and supported modes to the function description +- Registers all generated function tools with a main AI agent + +## Prerequisites + +Before you begin, ensure you have the following prerequisites: + +- .NET 10 SDK or later +- Access to the A2A agent host service +- A Microsoft Foundry project with a deployed model + +**Note**: These samples need to be run against a valid A2A server. If no A2A server is available, they can be run against the echo-agent that can be spun up locally by following the guidelines at: https://github.com/a2aproject/a2a-dotnet/blob/main/samples/AgentServer/README.md + +Set the following environment variables: + +Authenticate with Azure CLI by running `az login`, then set: + +```powershell +$env:A2A_AGENT_HOST="https://your-a2a-agent-host" # Replace with your A2A agent host endpoint +$env:FOUNDRY_PROJECT_ENDPOINT="https://your-project.services.ai.azure.com/api/projects/your-project" # Replace with your Foundry project endpoint +$env:FOUNDRY_MODEL="gpt-5.4-mini" # Optional, defaults to gpt-5.4-mini +``` diff --git a/dotnet/samples/02-agents/A2A/README.md b/dotnet/samples/02-agents/A2A/README.md index 28f2c0a9102..8f6cecc7e43 100644 --- a/dotnet/samples/02-agents/A2A/README.md +++ b/dotnet/samples/02-agents/A2A/README.md @@ -13,7 +13,8 @@ See the README.md for each sample for the prerequisites for that sample. |Sample|Description| |---|---| -|[A2A Agent As Function Tools](./A2AAgent_AsFunctionTools/)|This sample demonstrates how to represent an A2A agent as a set of function tools, where each function tool corresponds to a skill of the A2A agent, and register these function tools with another AI agent so it can leverage the A2A agent's skills.| +|[A2A Agent as a Function Tool](./A2AAgent_AsFunctionTool/)|This sample demonstrates how to expose an A2A agent as a single function tool that another AI agent can call.| +|[A2A Agent Skills](./A2AAgent_Skills/)|This sample demonstrates how to expose each skill advertised by an A2A agent as a separate function tool.| |[A2A Agent Polling For Task Completion](./A2AAgent_PollingForTaskCompletion/)|This sample demonstrates how to poll for long-running task completion using continuation tokens with an A2A agent.| |[A2A Agent Stream Reconnection](./A2AAgent_StreamReconnection/)|This sample demonstrates how to reconnect to an A2A agent's streaming response using continuation tokens, allowing recovery from stream interruptions.| |[A2A Agent Protocol Selection](./A2AAgent_ProtocolSelection/)|This sample demonstrates how to select the A2A protocol binding (HTTP+JSON vs JSON-RPC) when creating an AIAgent from an A2A agent card using A2AClientOptions.| @@ -23,7 +24,7 @@ See the README.md for each sample for the prerequisites for that sample. To run the samples, navigate to the desired sample directory, e.g. ```powershell -cd A2AAgent_AsFunctionTools +cd A2AAgent_AsFunctionTool ``` Set the required environment variables as documented in the sample readme. diff --git a/dotnet/samples/02-agents/AgentWithMemory/AgentWithMemory_Step08_MemoryUsingCosmosNoSql/AgentWithMemory_Step08_MemoryUsingCosmosNoSql.csproj b/dotnet/samples/02-agents/AgentWithMemory/AgentWithMemory_Step08_MemoryUsingCosmosNoSql/AgentWithMemory_Step08_MemoryUsingCosmosNoSql.csproj index d8950c82053..f4b063d84e6 100644 --- a/dotnet/samples/02-agents/AgentWithMemory/AgentWithMemory_Step08_MemoryUsingCosmosNoSql/AgentWithMemory_Step08_MemoryUsingCosmosNoSql.csproj +++ b/dotnet/samples/02-agents/AgentWithMemory/AgentWithMemory_Step08_MemoryUsingCosmosNoSql/AgentWithMemory_Step08_MemoryUsingCosmosNoSql.csproj @@ -10,7 +10,7 @@ - + diff --git a/dotnet/samples/02-agents/AgentWithMemory/AgentWithMemory_Step08_MemoryUsingCosmosNoSql/Program.cs b/dotnet/samples/02-agents/AgentWithMemory/AgentWithMemory_Step08_MemoryUsingCosmosNoSql/Program.cs index 4552d04653f..a6707d1d749 100644 --- a/dotnet/samples/02-agents/AgentWithMemory/AgentWithMemory_Step08_MemoryUsingCosmosNoSql/Program.cs +++ b/dotnet/samples/02-agents/AgentWithMemory/AgentWithMemory_Step08_MemoryUsingCosmosNoSql/Program.cs @@ -6,7 +6,7 @@ using System.Text.Json; using Azure.AI.Projects; using Azure.Identity; -using CommunityToolkit.VectorData.CosmosNoSql; +using CommunityToolkit.VectorData.AzureCosmosDB; using Microsoft.Agents.AI; using Microsoft.Azure.Cosmos; using Microsoft.Extensions.AI; @@ -40,9 +40,9 @@ DatabaseResponse databaseResponse = await cosmosClient.CreateDatabaseIfNotExistsAsync(cosmosDatabaseName); -VectorStore vectorStore = new CosmosNoSqlVectorStore( +VectorStore vectorStore = new CosmosVectorStore( databaseResponse.Database, - new CosmosNoSqlVectorStoreOptions + new CosmosVectorStoreOptions { JsonSerializerOptions = JsonSerializerOptions.Default, EmbeddingGenerator = aiProjectClient diff --git a/dotnet/samples/02-agents/AgentWithMemory/AgentWithMemory_Step08_MemoryUsingCosmosNoSql/README.md b/dotnet/samples/02-agents/AgentWithMemory/AgentWithMemory_Step08_MemoryUsingCosmosNoSql/README.md index 16775abcd8d..af667e86eff 100644 --- a/dotnet/samples/02-agents/AgentWithMemory/AgentWithMemory_Step08_MemoryUsingCosmosNoSql/README.md +++ b/dotnet/samples/02-agents/AgentWithMemory/AgentWithMemory_Step08_MemoryUsingCosmosNoSql/README.md @@ -1,6 +1,6 @@ # Agent with Memory Using Azure Cosmos DB for NoSQL -This sample uses `ChatHistoryMemoryProvider` with `CosmosNoSqlVectorStore` to persist chat history in Azure Cosmos DB for NoSQL and recall relevant messages in a new agent session. +This sample uses `ChatHistoryMemoryProvider` with `CosmosVectorStore` to persist chat history in Azure Cosmos DB for NoSQL and recall relevant messages in a new agent session. ## Features Demonstrated diff --git a/dotnet/samples/02-agents/AgentWithMemory/AgentWithMemory_Step09_MemoryUsingMem0Sharp/README.md b/dotnet/samples/02-agents/AgentWithMemory/AgentWithMemory_Step09_MemoryUsingMem0Sharp/README.md new file mode 100644 index 00000000000..edde3e1496b --- /dev/null +++ b/dotnet/samples/02-agents/AgentWithMemory/AgentWithMemory_Step09_MemoryUsingMem0Sharp/README.md @@ -0,0 +1,7 @@ +# Agent with Memory Using Mem0Sharp + +Mem0Sharp is an independent long-term memory library for .NET applications and agents. Its repository contains a Microsoft Agent Framework sample that demonstrates cross-session memory with an in-memory store. + +In addition to in-memory storage, Mem0Sharp supports PostgreSQL with pgvector, SQLite, Qdrant, and custom storage providers. + +See the [Microsoft Agent Framework memory sample in the Mem0Sharp repository](https://github.com/jihadkhawaja/mem0sharp/blob/main/samples/AgentFrameworkMemory/README.md) for prerequisites, configuration, and run instructions. \ No newline at end of file diff --git a/dotnet/samples/02-agents/AgentWithMemory/README.md b/dotnet/samples/02-agents/AgentWithMemory/README.md index 16096f44fae..25cc85bcba0 100644 --- a/dotnet/samples/02-agents/AgentWithMemory/README.md +++ b/dotnet/samples/02-agents/AgentWithMemory/README.md @@ -12,6 +12,7 @@ These samples show how to create an agent with the Agent Framework that uses Mem |[Memory Using AgentMemory](./AgentWithMemory_Step06_MemoryUsingAgentMemory/)|This sample demonstrates a retail shopping assistant built with [`AgentMemory`](https://www.nuget.org/packages/AgentMemory), an unofficial .NET port of the Neo4j Labs graph-memory provider, to learn customer preferences and recommend products via graph traversal.| |[File Based Memory](./AgentWithMemory_Step07_FileMemoryProvider/)|This sample demonstrates how to use the `FileMemoryProvider` to give an agent tools for storing and recalling memories as files, and how to configure the folder that those memory files are written to.| |[Memory with Azure Cosmos DB for NoSQL](./AgentWithMemory_Step08_MemoryUsingCosmosNoSql/)|This sample demonstrates how to persist and retrieve chat history across sessions with Azure Cosmos DB for NoSQL.| +|[Memory with Mem0Sharp](./AgentWithMemory_Step09_MemoryUsingMem0Sharp/)|This sample links to the Mem0Sharp repository's Microsoft Agent Framework memory integration example.| > **See also**: [Memory Search with Foundry Agents](../AgentProviders/foundry/Agent_Step22_MemorySearch/) - demonstrates using the built-in Memory Search tool with Microsoft Foundry agents. diff --git a/dotnet/samples/04-hosting/FoundryHostedAgents/invocations/Hosted-Invocations-EchoAgent/Hosted-Invocations-EchoAgent.csproj b/dotnet/samples/04-hosting/FoundryHostedAgents/invocations/Hosted-Invocations-EchoAgent/Hosted-Invocations-EchoAgent.csproj index 32b0ba2065a..edd74f5a74b 100644 --- a/dotnet/samples/04-hosting/FoundryHostedAgents/invocations/Hosted-Invocations-EchoAgent/Hosted-Invocations-EchoAgent.csproj +++ b/dotnet/samples/04-hosting/FoundryHostedAgents/invocations/Hosted-Invocations-EchoAgent/Hosted-Invocations-EchoAgent.csproj @@ -32,7 +32,7 @@ - + diff --git a/dotnet/samples/05-end-to-end/AGUIClientServer/AGUIServer/AGUIServer.csproj b/dotnet/samples/05-end-to-end/AGUIClientServer/AGUIServer/AGUIServer.csproj index 8aa91a3406b..788f6a3f3ff 100644 --- a/dotnet/samples/05-end-to-end/AGUIClientServer/AGUIServer/AGUIServer.csproj +++ b/dotnet/samples/05-end-to-end/AGUIClientServer/AGUIServer/AGUIServer.csproj @@ -9,8 +9,8 @@ - + diff --git a/dotnet/samples/05-end-to-end/AGUIClientServer/AGUIServer/Program.cs b/dotnet/samples/05-end-to-end/AGUIClientServer/AGUIServer/Program.cs index e3de957c1e7..6f1ed5ba3ef 100644 --- a/dotnet/samples/05-end-to-end/AGUIClientServer/AGUIServer/Program.cs +++ b/dotnet/samples/05-end-to-end/AGUIClientServer/AGUIServer/Program.cs @@ -1,13 +1,14 @@ // Copyright (c) Microsoft. All rights reserved. +using System.ClientModel.Primitives; using System.ComponentModel; using AGUIServer; -using Azure.AI.OpenAI; using Azure.Identity; using Microsoft.Agents.AI.Hosting; using Microsoft.Agents.AI.Hosting.AGUI.AspNetCore; using Microsoft.Extensions.AI; -using OpenAI.Chat; +using OpenAI; +using OpenAI.Responses; WebApplicationBuilder builder = WebApplication.CreateBuilder(args); builder.Services.AddHttpClient().AddLogging(); @@ -19,35 +20,15 @@ const string AgentName = "AGUIAssistant"; -// Create the AI agent with tools +// Create a Responses-backed OpenAI client that sends a bearer token to the Azure endpoint. // WARNING: DefaultAzureCredential is convenient for development but requires careful consideration in production. // In production, consider using a specific credential (e.g., ManagedIdentityCredential) to avoid // latency issues, unintended credential probing, and potential security risks from fallback mechanisms. -var agent = new AzureOpenAIClient( - new Uri(endpoint), - new DefaultAzureCredential()) - .GetChatClient(deploymentName) - .AsAIAgent( - name: AgentName, - tools: [ - AIFunctionFactory.Create( - () => DateTimeOffset.UtcNow, - name: "get_current_time", - description: "Get the current UTC time." - ), - AIFunctionFactory.Create( - ([Description("The weather forecast request")]ServerWeatherForecastRequest request) => { - return new ServerWeatherForecastResponse() - { - Summary = "Sunny", - TemperatureC = 25, - Date = request.Date - }; - }, - name: "get_server_weather_forecast", - description: "Gets the forecast for a specific location and date", - AGUIServerSerializerContext.Default.Options) - ]); +IChatClient chatClient = new OpenAIClient( + new BearerTokenPolicy(new DefaultAzureCredential(), "https://ai.azure.com/.default"), + new OpenAIClientOptions { Endpoint = new Uri(endpoint) }) + .GetResponsesClient() + .AsIChatClientWithStoredOutputDisabled(model: deploymentName); // WARNING: When adding session persistence (e.g., WithInMemorySessionStore), or running in production, // make sure to also register an AgentIsolationKeyProvider to scope sessions by principal in multi-user @@ -57,7 +38,26 @@ // Register the agent with the host and configure it to use an in-memory session store // so that conversation state is maintained across requests. In production, you may want to use a persistent session store. builder - .AddAIAgent(AgentName, (_, _) => agent) + .AddAIAgent(AgentName, "You are a helpful assistant.", chatClient) + .WithAITools( + new HostedWebSearchTool(), + AIFunctionFactory.Create( + () => DateTimeOffset.UtcNow, + name: "get_current_time", + description: "Get the current UTC time."), + AIFunctionFactory.Create( + ([Description("The weather forecast request")] ServerWeatherForecastRequest request) => + { + return new ServerWeatherForecastResponse() + { + Summary = "Sunny", + TemperatureC = 25, + Date = request.Date + }; + }, + name: "get_server_weather_forecast", + description: "Gets the forecast for a specific location and date", + AGUIServerSerializerContext.Default.Options)) .WithInMemorySessionStore(); WebApplication app = builder.Build(); diff --git a/dotnet/samples/05-end-to-end/AGUIClientServer/README.md b/dotnet/samples/05-end-to-end/AGUIClientServer/README.md index 0501ec54694..30adeb1eef2 100644 --- a/dotnet/samples/05-end-to-end/AGUIClientServer/README.md +++ b/dotnet/samples/05-end-to-end/AGUIClientServer/README.md @@ -18,11 +18,18 @@ The demonstration has two components: Configure the required Azure OpenAI environment variables: ```powershell -$env:AZURE_OPENAI_ENDPOINT="<>" +$env:AZURE_OPENAI_ENDPOINT="https://.openai.azure.com/openai/v1/" $env:AZURE_OPENAI_DEPLOYMENT_NAME="gpt-5.4-mini" ``` -> **Note:** This sample uses `DefaultAzureCredential` for authentication. Make sure you're authenticated with Azure (e.g., via `az login`, Visual Studio, or environment variables). +> [!NOTE] +> Include `/openai/v1/` in the endpoint. The OpenAI SDK uses `DefaultAzureCredential` to obtain a bearer token. Make sure you're authenticated with Azure, for example through `az login`, Visual Studio, or environment variables. + +> [!NOTE] +> This sample calls Azure OpenAI inference directly through the resource endpoint. It does not require a Microsoft Foundry project. A project-scoped application would instead use a Foundry project endpoint with `Azure.AI.Projects` and the Agent Framework Foundry provider. + +> [!NOTE] +> The server uses the Azure OpenAI Responses API because hosted web search is a Responses API tool. It sets `store` to `false` so Agent Framework persists chat history in the configured session store instead of depending on service-retained responses. Web search uses Grounding with Bing and may incur additional charges; review the [web search documentation and data usage terms](https://learn.microsoft.com/azure/foundry/openai/how-to/web-search) before using it. ## Running the Sample @@ -117,13 +124,18 @@ User (:q or quit to exit): :q The `AGUIServer` uses the `MapAGUIServer` extension method to expose an agent through the AG-UI protocol: ```csharp -AIAgent agent = new OpenAIClient(apiKey) - .GetChatClient(model) - .AsAIAgent( - instructions: "You are a helpful assistant.", - name: "AGUIAssistant"); - -app.MapAGUIServer("/", agent); +IChatClient chatClient = new OpenAIClient( + new BearerTokenPolicy(new DefaultAzureCredential(), "https://ai.azure.com/.default"), + new OpenAIClientOptions { Endpoint = new Uri(endpoint) }) + .GetResponsesClient() + .AsIChatClientWithStoredOutputDisabled(model: deploymentName); + +builder + .AddAIAgent("AGUIAssistant", "You are a helpful assistant.", chatClient) + .WithAITool(new HostedWebSearchTool()) + .WithInMemorySessionStore(); + +app.MapAGUIServer("AGUIAssistant", "/"); ``` This automatically handles: diff --git a/dotnet/src/Aspire.Hosting.AgentFramework.DevUI/DevUIAggregatorHostedService.cs b/dotnet/src/Aspire.Hosting.AgentFramework.DevUI/DevUIAggregatorHostedService.cs index 04ef517892b..d87e6ea28d1 100644 --- a/dotnet/src/Aspire.Hosting.AgentFramework.DevUI/DevUIAggregatorHostedService.cs +++ b/dotnet/src/Aspire.Hosting.AgentFramework.DevUI/DevUIAggregatorHostedService.cs @@ -167,6 +167,7 @@ private async Task ServeDevUIFrontendAsync(HttpContext context, string? path) } context.Response.StatusCode = StatusCodes.Status301MovedPermanently; + // CodeQL [SM04598] False positive: The Location is always /devui/?. Since it is relative, the client re-requests the same host; the query cannot change the destination. context.Response.Headers.Location = redirect; return; } @@ -725,8 +726,8 @@ private static async Task ProxyRequestAsync( ? HttpCompletionOption.ResponseHeadersRead : HttpCompletionOption.ResponseContentRead; - using var response = await client.SendAsync( // CodeQL [SM03781] False positive: ValidateProxyTarget confirms the target host, scheme, and port match the configured backend, so the user-supplied path and query cannot change the destination. - request, completionOption, context.RequestAborted).ConfigureAwait(false); + // CodeQL [SM03781] False positive: ValidateProxyTarget confirms the target host, scheme, and port match the configured backend, so the user-supplied path and query cannot change the destination. + using var response = await client.SendAsync(request, completionOption, context.RequestAborted).ConfigureAwait(false); if (streaming && response.Content.Headers.ContentType?.MediaType == "text/event-stream") { diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs index 725fd28f03b..32248cf30b1 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/AgentFrameworkResponseHandler.cs @@ -536,6 +536,11 @@ bool CheckNotAllowedStoreUsage() => } evt = enumerator.Current; + if (evt is ResponseCompletedEvent) + { + consentCts.Token.ThrowIfCancellationRequested(); + } + shutdownDetected = context.IsShutdownRequested && !emittedTerminal; } diff --git a/dotnet/src/Microsoft.Agents.AI.Mcp/Skills/Loaders/AgentMcpSkillArchiveExtractor.cs b/dotnet/src/Microsoft.Agents.AI.Mcp/Skills/Loaders/AgentMcpSkillArchiveExtractor.cs index 15e37926a77..3957aaf0c8b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Mcp/Skills/Loaders/AgentMcpSkillArchiveExtractor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Mcp/Skills/Loaders/AgentMcpSkillArchiveExtractor.cs @@ -159,7 +159,7 @@ private static void ExtractZip(Stream source, string fullTarget, int maxFileCoun throw new InvalidDataException($"Skill archive exceeds the maximum allowed file count ({maxFileCount})."); } - string? destination = ResolveDestination(fullTarget, entry.FullName); + string? destination = ResolveDestination(fullTarget, entry.FullName); // CodeQL [SM02729] ResolveDestination rejects paths outside fullTarget. if (destination is null) { continue; diff --git a/dotnet/tests/Microsoft.Agents.AI.FeatureRegistry.UnitTests/FeatureRegistryTests.cs b/dotnet/tests/Microsoft.Agents.AI.FeatureRegistry.UnitTests/FeatureRegistryTests.cs index 1b671a0508b..386b6b833f6 100644 --- a/dotnet/tests/Microsoft.Agents.AI.FeatureRegistry.UnitTests/FeatureRegistryTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.FeatureRegistry.UnitTests/FeatureRegistryTests.cs @@ -5,9 +5,11 @@ using System.Globalization; using System.IO; using System.Linq; +#pragma warning disable IDE0005 // Required by Windows builds for Roslyn APIs. using Microsoft.CodeAnalysis; using Microsoft.CodeAnalysis.CSharp; using Microsoft.CodeAnalysis.CSharp.Syntax; +#pragma warning restore IDE0005 namespace Microsoft.Agents.AI.FeatureRegistry.UnitTests; @@ -431,5 +433,5 @@ private static void AssertNoErrors(IReadOnlyCollection errors) private sealed record RegistryEntry(int Index, string Id); - private sealed record FeatureDeclaration(string Owner, string MemberName, int Index, string FilePath); + private readonly record struct FeatureDeclaration(string Owner, string MemberName, int Index, string FilePath); } diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerTests.cs index 070a6ca0f7b..2b7ae5b0061 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerTests.cs @@ -666,15 +666,26 @@ public async Task CreateAsync_CancellationDuringExecution_PropagatesOperationCan .ReturnsAsync(Array.Empty()); using var cts = new CancellationTokenSource(); - cts.Cancel(); + var events = new List(); - // Act & Assert - await Assert.ThrowsAsync(async () => + async Task ExecuteAsync() { - await foreach (var _ in handler.CreateAsync(request, mockContext.Object, cts.Token)) + await foreach (var evt in handler.CreateAsync(request, mockContext.Object, cts.Token)) { + events.Add(evt); } - }); + } + + // Act + var execution = ExecuteAsync(); + await agent.Started.Task.WaitAsync(TimeSpan.FromSeconds(5)); + cts.Cancel(); + + // Assert + await Assert.ThrowsAnyAsync( + async () => await execution.WaitAsync(TimeSpan.FromSeconds(5))); + await agent.CancellationObserved.Task.WaitAsync(TimeSpan.FromSeconds(5)); + Assert.DoesNotContain(events, evt => evt is ResponseCompletedEvent); } [Fact] @@ -1643,15 +1654,29 @@ protected override ValueTask DeserializeSessionCoreAsync( private sealed class CancellationCheckingAgent : AIAgent { + public TaskCompletionSource Started { get; } = + new(TaskCreationOptions.RunContinuationsAsynchronously); + + public TaskCompletionSource CancellationObserved { get; } = + new(TaskCreationOptions.RunContinuationsAsynchronously); + protected override async IAsyncEnumerable RunCoreStreamingAsync( IEnumerable messages, AgentSession? session, AgentRunOptions? options, [EnumeratorCancellation] CancellationToken cancellationToken = default) { - cancellationToken.ThrowIfCancellationRequested(); - yield return new AgentResponseUpdate { Contents = [new MeaiTextContent("test")] }; - await Task.CompletedTask; + this.Started.TrySetResult(); + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + this.CancellationObserved.TrySetResult(); + } + + yield break; } protected override Task RunCoreAsync( diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerWorkflowTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerWorkflowTests.cs index f5194f5f6d1..e3d255d700f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerWorkflowTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/AgentFrameworkResponseHandlerWorkflowTests.cs @@ -161,6 +161,40 @@ public async Task WorkflowAgent_RegisteredWithKey_ResolvesCorrectlyAsync() Assert.True(events.Count >= 3, $"Expected at least 3 events, got {events.Count}"); } + [Fact] + public async Task CreateAsync_WorkflowCancellationDuringExecution_CancelsInnerAgentWithoutCompletionAsync() + { + // Arrange + var innerAgent = new CancellationCheckingWorkflowAgent("blocking-agent"); + var workflow = AgentWorkflowBuilder.BuildSequential("cancellation-workflow", innerAgent); + var workflowAgent = workflow.AsAIAgent( + id: "workflow-agent", + name: "Cancellation Workflow", + executionEnvironment: InProcessExecution.OffThread); + var (handler, request, context) = CreateHandlerWithAgent(workflowAgent, "Hello"); + using var cts = new CancellationTokenSource(); + var events = new List(); + + async Task ExecuteAsync() + { + await foreach (var evt in handler.CreateAsync(request, context, cts.Token)) + { + events.Add(evt); + } + } + + // Act + var execution = ExecuteAsync(); + await innerAgent.Started.Task.WaitAsync(TimeSpan.FromSeconds(5)); + cts.Cancel(); + + // Assert + await Assert.ThrowsAnyAsync( + async () => await execution.WaitAsync(TimeSpan.FromSeconds(5))); + await innerAgent.CancellationObserved.Task.WaitAsync(TimeSpan.FromSeconds(5)); + Assert.DoesNotContain(events, evt => evt is ResponseCompletedEvent); + } + private static (AgentFrameworkResponseHandler handler, CreateResponse request, ResponseContext context) CreateHandlerWithAgent( AIAgent agent, diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/WorkflowTestAgents.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/WorkflowTestAgents.cs index 7e10abc9ebd..0ca010d34fe 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/WorkflowTestAgents.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.Hosting.UnitTests/WorkflowTestAgents.cs @@ -94,3 +94,62 @@ protected override ValueTask DeserializeSessionCoreAsync( CancellationToken cancellationToken = default) => throw new NotImplementedException(); } + +/// +/// A test agent that blocks until its run is cancelled. +/// +internal sealed class CancellationCheckingWorkflowAgent(string id) : AIAgent +{ + public new string Id => id; + + public TaskCompletionSource Started { get; } = + new(TaskCreationOptions.RunContinuationsAsynchronously); + + public TaskCompletionSource CancellationObserved { get; } = + new(TaskCreationOptions.RunContinuationsAsynchronously); + + protected override async IAsyncEnumerable RunCoreStreamingAsync( + IEnumerable messages, + AgentSession? session, + AgentRunOptions? options, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + this.Started.TrySetResult(); + try + { + await Task.Delay(Timeout.InfiniteTimeSpan, cancellationToken); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + this.CancellationObserved.TrySetResult(); + throw; + } + + yield break; + } + + protected override Task RunCoreAsync( + IEnumerable messages, + AgentSession? session, + AgentRunOptions? options, + CancellationToken cancellationToken = default) => + throw new NotImplementedException(); + + protected override ValueTask CreateSessionCoreAsync( + CancellationToken cancellationToken = default) => + new(new TestAgentSession()); + + protected override ValueTask SerializeSessionCoreAsync( + AgentSession session, + JsonSerializerOptions? jsonSerializerOptions, + CancellationToken cancellationToken = default) => + new(JsonSerializer.SerializeToElement(new { })); + + protected override ValueTask DeserializeSessionCoreAsync( + JsonElement serializedState, + JsonSerializerOptions? jsonSerializerOptions, + CancellationToken cancellationToken = default) => + new(new TestAgentSession()); + + private sealed class TestAgentSession : AgentSession; +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/AgentHostingServiceCollectionExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/AgentHostingServiceCollectionExtensionsTests.cs new file mode 100644 index 00000000000..587af5ff339 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/AgentHostingServiceCollectionExtensionsTests.cs @@ -0,0 +1,95 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.ClientModel; +using System.ClientModel.Primitives; +using System.Net; +using System.Net.Http; +using System.Text; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DependencyInjection; +using OpenAI; +using OpenAI.Responses; + +namespace Microsoft.Agents.AI.Hosting.OpenAI.UnitTests; + +public sealed class AgentHostingServiceCollectionExtensionsTests +{ + [Fact] + public async Task HostedWebSearchTool_WithResponsesClient_UsesResponsesWireFormatAsync() + { + // Arrange + using var handler = new RecordingHandler(); +#pragma warning disable CA5399 + using var httpClient = new HttpClient(handler); +#pragma warning restore CA5399 + using IChatClient chatClient = new ResponsesClient( + new ApiKeyCredential("test-key"), + new OpenAIClientOptions + { + Endpoint = new Uri("https://example.test/v1"), + Transport = new HttpClientPipelineTransport(httpClient) + }) + .AsIChatClientWithStoredOutputDisabled(model: "test-model"); + + var services = new ServiceCollection(); + services + .AddAIAgent("test-agent", "You are a helpful assistant.", chatClient) + .WithAITool(new HostedWebSearchTool()); + using ServiceProvider serviceProvider = services.BuildServiceProvider(); + AIAgent agent = serviceProvider.GetRequiredKeyedService("test-agent"); + + // Act + await agent.RunAsync("What happened in the news today?"); + + // Assert + using JsonDocument request = JsonDocument.Parse(Assert.IsType(handler.RequestBody)); + Assert.False(request.RootElement.GetProperty("store").GetBoolean()); + Assert.Contains( + request.RootElement.GetProperty("include").EnumerateArray(), + property => property.GetString() == "reasoning.encrypted_content"); + Assert.False(request.RootElement.TryGetProperty("web_search_options", out _)); + JsonElement webSearchTool = Assert.Single(request.RootElement.GetProperty("tools").EnumerateArray()); + Assert.Equal("web_search", webSearchTool.GetProperty("type").GetString()); + Assert.Equal("/v1/responses", handler.RequestUri?.AbsolutePath); + } + + private sealed class RecordingHandler : HttpMessageHandler + { + public Uri? RequestUri { get; private set; } + + public string? RequestBody { get; private set; } + + protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + this.RequestUri = request.RequestUri; + this.RequestBody = await request.Content!.ReadAsStringAsync(cancellationToken); + + return new HttpResponseMessage(HttpStatusCode.OK) + { + Content = new StringContent( + """ + { + "id": "resp_1", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "test-model", + "output": [], + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2 + } + } + """, + Encoding.UTF8, + "application/json"), + RequestMessage = request + }; + } + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests.csproj b/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests.csproj index 17d9742436c..355e271d767 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests.csproj +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests.csproj @@ -19,6 +19,7 @@ + diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index 389f5038fce..6d4b86b40f2 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -43,6 +43,9 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. - Default stateless conversation history is client-controlled, including historical tool calls and results. Never document conversational tool results as authorization or policy evidence; use deterministic server-side checks, server-validated approvals, or scoped authoritative snapshots. +- AG-UI Thread and Run ids are client-owned protocol correlation ids. Service-session mode stores provider conversation + or response ids privately in the thread snapshot. Set `service_session_id_from_thread_id=True` only for compatibility + when the application intentionally uses a provider continuation id as its AG-UI Thread id. - `confirm_changes` snapshot cleanup resolves the synthetic confirmation back to its original `function_call_id`; it must never concatenate unrelated tool results or record accepted changes without a matching real result. - SSE keepalive is endpoint-owned transport behavior configured through diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent.py index 9ae73c3790e..34dd4b9ac3c 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent.py @@ -26,6 +26,7 @@ def __init__( require_confirmation: bool = True, snapshot_store: AGUIThreadSnapshotStore | None = None, a2ui_config: dict[str, Any] | None = None, + service_session_id_from_thread_id: bool = False, ): """Initialize agent configuration. @@ -33,6 +34,8 @@ def __init__( state_schema: Optional state schema for state management; accepts dict or Pydantic model/class predict_state_config: Configuration for predictive state updates use_service_session: Whether the agent session is service-managed + service_session_id_from_thread_id: Compatibility option that treats the + client-owned AG-UI Thread id as the provider service-session id. require_confirmation: When True (default), emit a ``confirm_changes`` tool call for approval-gated tools so a human-in-the-loop frontend can prompt for approval, and require confirmation for predictive state updates. When False, no ``confirm_changes`` call is @@ -46,6 +49,7 @@ def __init__( self.state_schema = self._normalize_state_schema(state_schema) self.predict_state_config = predict_state_config or {} self.use_service_session = use_service_session + self.service_session_id_from_thread_id = service_session_id_from_thread_id self.require_confirmation = require_confirmation self.snapshot_store = snapshot_store self.a2ui_config = a2ui_config @@ -96,6 +100,7 @@ def __init__( use_service_session: bool = False, snapshot_store: AGUIThreadSnapshotStore | None = None, a2ui_config: dict[str, Any] | None = None, + service_session_id_from_thread_id: bool = False, ): """Initialize the AG-UI compatible agent wrapper. @@ -110,6 +115,8 @@ def __init__( confirmation for predictive state updates. When False, no ``confirm_changes`` call is emitted; tools remain gated server-side. use_service_session: Whether the agent session is service-managed + service_session_id_from_thread_id: Compatibility option that treats the + client-owned AG-UI Thread id as the provider service-session id. snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless endpoint setup also provides an explicit Snapshot Scope resolver. a2ui_config: Optional backend A2UI config consumed by auto-injection. @@ -122,6 +129,7 @@ def __init__( state_schema=state_schema, predict_state_config=predict_state_config, use_service_session=use_service_session, + service_session_id_from_thread_id=service_session_id_from_thread_id, require_confirmation=require_confirmation, snapshot_store=snapshot_store, a2ui_config=a2ui_config, diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 8b99e106dc4..ea5a00211fb 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -120,6 +120,7 @@ # Keys that are internal to AG-UI orchestration and should not be passed to chat clients AG_UI_INTERNAL_METADATA_KEYS = {"ag_ui_thread_id", "ag_ui_run_id", "current_state", "forwarded_props"} _COLLECTED_APPROVAL_RESPONSES_KEY = "collected_approval_responses" +_PROVIDER_SERVICE_SESSION_ID_STATE_KEY = "__ag_ui_provider_service_session_id" @dataclass @@ -2108,16 +2109,24 @@ def _text_events_to_snapshot_messages(events: list[BaseEvent]) -> list[dict[str, return [message for message in messages if message.get("content")] -def _restore_session_continuation_state(session: AgentSession, snapshot: AGUIThreadSnapshot | None) -> None: +def _restore_session_continuation_state( + session: AgentSession, + snapshot: AGUIThreadSnapshot | None, + *, + restore_service_session_id: bool, +) -> None: """Restore typed private state from trusted snapshot storage.""" if snapshot is None or snapshot.session_state is None: return + serialized_state = copy.deepcopy(snapshot.session_state) + service_session_id = serialized_state.pop(_PROVIDER_SERVICE_SESSION_ID_STATE_KEY, None) try: restored = AgentSession.from_dict( { "type": "session", "session_id": session.session_id, - "state": snapshot.session_state, + "service_session_id": service_session_id, + "state": serialized_state, } ) except Exception: @@ -2126,6 +2135,8 @@ def _restore_session_continuation_state(session: AgentSession, snapshot: AGUIThr session.session_id, ) return + if restore_service_session_id and service_session_id is not None: + session.service_session_id = restored.service_session_id session.state.update(restored.state) @@ -2178,21 +2189,31 @@ def _serialize_session_continuation_state( agent: SupportsAgentRun, *, shared_state_keys: set[str], + include_service_session_id: bool, ) -> dict[str, Any] | None: """Serialize server-owned state while preserving each AG-UI State Authority.""" context_providers = cast(list[Any], getattr(agent, "context_providers", [])) excluded_keys = { *shared_state_keys, _TOOL_APPROVAL_STATE_KEY, + _PROVIDER_SERVICE_SESSION_ID_STATE_KEY, *(provider.source_id for provider in context_providers if isinstance(provider, HistoryProvider)), } continuation_state = {key: value for key, value in session.state.items() if key not in excluded_keys} - if not continuation_state: + service_session_id = session.service_session_id if include_service_session_id else None + if not continuation_state and service_session_id is None: return None - serialized_session = AgentSession(session_id=session.session_id) + serialized_session = AgentSession( + session_id=session.session_id, + service_session_id=service_session_id, + ) serialized_session.state.update(continuation_state) - return cast(dict[str, Any], serialized_session.to_dict()["state"]) + serialized_payload = serialized_session.to_dict() + serialized_state = cast(dict[str, Any], serialized_payload["state"]) + if serialized_service_session_id := serialized_payload.get("service_session_id"): + serialized_state[_PROVIDER_SERVICE_SESSION_ID_STATE_KEY] = serialized_service_session_id + return serialized_state def _safe_serialize_session_continuation_state( @@ -2200,6 +2221,7 @@ def _safe_serialize_session_continuation_state( agent: SupportsAgentRun, *, shared_state_keys: set[str], + include_service_session_id: bool, ) -> dict[str, Any] | None: """Return JSON-safe continuation state without failing a completed run.""" try: @@ -2207,6 +2229,7 @@ def _safe_serialize_session_continuation_state( session, agent, shared_state_keys=shared_state_keys, + include_service_session_id=include_service_session_id, ) if serialized_state is None: return None @@ -2225,6 +2248,23 @@ def _safe_serialize_session_continuation_state( return None +def _split_service_session_input( + stored_snapshot_messages: list[dict[str, Any]], + current_turn_messages: list[dict[str, Any]], + stored_interrupt: list[dict[str, Any]] | None = None, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """Return the validated new suffix and backend-authoritative snapshot history.""" + if not current_turn_messages: + snapshot_messages = copy.deepcopy(stored_snapshot_messages) + else: + snapshot_messages = _reconstruct_messages_from_thread_snapshot( + stored_messages=stored_snapshot_messages, + incoming_messages=current_turn_messages, + stored_interrupt=stored_interrupt, + ) + return snapshot_messages[len(stored_snapshot_messages) :], snapshot_messages + + async def run_agent_stream( input_data: dict[str, Any], agent: SupportsAgentRun, @@ -2296,16 +2336,34 @@ async def run_agent_stream( yield event return + snapshot_seed_messages: list[dict[str, Any]] | None = None + if stored_snapshot is not None: if resume_payload is not None and stored_pending_approval_interrupt_ids: - raw_messages = snapshot_session.resume_seeded_messages(raw_messages) seeded_resume_from_snapshot = True - else: + + if not config.use_service_session: + raw_messages = snapshot_session.resume_seeded_messages(raw_messages) + else: + provider_suffix, snapshot_seed_messages = _split_service_session_input( + stored_snapshot_messages=stored_snapshot.messages, + current_turn_messages=raw_messages, + stored_interrupt=stored_snapshot.interrupt, + ) + raw_messages = provider_suffix + elif not config.use_service_session: raw_messages = _reconstruct_messages_from_thread_snapshot( stored_messages=stored_snapshot.messages, incoming_messages=raw_messages, stored_interrupt=stored_snapshot.interrupt, ) + else: + provider_suffix, snapshot_seed_messages = _split_service_session_input( + stored_snapshot_messages=stored_snapshot.messages, + current_turn_messages=raw_messages, + stored_interrupt=stored_snapshot.interrupt, + ) + raw_messages = provider_suffix # Initialize flow state with stored state plus request-provided overrides; # endpoint-deferred defaults apply only to keys missing from both. @@ -2393,9 +2451,13 @@ async def run_agent_stream( if approval_resume_messages: logger.info(f"Appending {len(approval_resume_messages)} synthesized approval resume message(s).") raw_messages.extend(approval_resume_messages) + if snapshot_seed_messages is not None: + snapshot_seed_messages.extend(copy.deepcopy(approval_resume_messages)) if resume_messages: logger.info(f"Appending {len(resume_messages)} synthesized resume message(s) to AG-UI input.") raw_messages.extend(resume_messages) + if snapshot_seed_messages is not None: + snapshot_seed_messages.extend(copy.deepcopy(resume_messages)) if retained_approval_results and not raw_messages: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) for event in _make_approval_tool_result_events(retained_approval_results): @@ -2405,9 +2467,15 @@ async def run_agent_stream( protected_tool_call_ids = _approval_state_tool_call_ids(approval_state_store, approval_thread_id) messages, snapshot_messages = normalize_agui_input_messages( raw_messages, + sanitize_tool_history=not (config.use_service_session and stored_snapshot is not None), protected_tool_call_ids=protected_tool_call_ids, ) + if snapshot_seed_messages is not None: + _, snapshot_messages = normalize_agui_input_messages( + snapshot_seed_messages, + protected_tool_call_ids=protected_tool_call_ids, + ) # Check for structured output mode (skip text content) skip_text = False response_format: type[Any] | None = None @@ -2473,10 +2541,41 @@ async def run_agent_stream( # Create session (with service session support) if config.use_service_session: - session = AgentSession(session_id=thread_id, service_session_id=supplied_thread_id) + if isinstance(default_options, dict) and default_options.get("store") is False: + raise ValueError( + "use_service_session=True requires provider storage. Set agent default_options['store']=True " + "or disable use_service_session." + ) + if not config.service_session_id_from_thread_id and not snapshot_session.enabled: + raise ValueError( + "use_service_session=True requires snapshot persistence unless service_session_id_from_thread_id=True." + ) + service_session_id = supplied_thread_id if config.service_session_id_from_thread_id else None + session = AgentSession(session_id=thread_id, service_session_id=service_session_id) + stored_service_session_id = ( + stored_snapshot.session_state.get(_PROVIDER_SERVICE_SESSION_ID_STATE_KEY) + if stored_snapshot is not None and stored_snapshot.session_state is not None + else None + ) + create_conversation = getattr(agent, "create_conversation", None) + if ( + not config.service_session_id_from_thread_id + and stored_service_session_id is None + and callable(create_conversation) + ): + created_session = create_conversation(session_id=thread_id) + if isinstance(created_session, Awaitable): + created_session = await created_session + if not isinstance(created_session, AgentSession): + raise TypeError("agent.create_conversation() must return AgentSession") + session = created_session else: session = AgentSession(session_id=thread_id) - _restore_session_continuation_state(session, stored_snapshot) + _restore_session_continuation_state( + session, + stored_snapshot, + restore_service_session_id=config.use_service_session, + ) protected_session_state_keys = _request_state_protected_keys(agent) session.state.update( { @@ -2600,7 +2699,7 @@ async def run_agent_stream( # Persist the completed confirmation turn with interrupt=None so hydration # does not replay the stale pending interrupt after the user responded. persisted_messages = snapshot_messages + _text_events_to_snapshot_messages(confirmation_events) - if resume_payload is not None and not seeded_resume_from_snapshot: + if resume_payload is not None and not seeded_resume_from_snapshot and snapshot_seed_messages is None: # Generic resume requests carry only the synthesized response, so prepend # stored history unless this run already seeded raw messages from it. persisted_messages = snapshot_session.resume_seeded_messages(persisted_messages) @@ -2612,6 +2711,7 @@ async def run_agent_stream( session, agent, shared_state_keys=set(flow.current_state).difference(protected_session_state_keys), + include_service_session_id=config.use_service_session, ), ) _save_tool_approval_state(session, approval_state_store, approval_thread_id) @@ -2958,7 +3058,7 @@ async def run_agent_stream( yield snapshot_event persisted_messages = latest_messages_snapshot - if resume_payload is not None and not seeded_resume_from_snapshot: + if resume_payload is not None and not seeded_resume_from_snapshot and snapshot_seed_messages is None: # Generic resume requests carry only the synthesized response, so prepend # stored history unless this run already seeded raw messages from it. persisted_messages = snapshot_session.resume_seeded_messages(persisted_messages) @@ -2970,6 +3070,7 @@ async def run_agent_stream( session, agent, shared_state_keys=set(flow.current_state).difference(protected_session_state_keys), + include_service_session_id=config.use_service_session, ), ) _save_tool_approval_state(session, approval_state_store, approval_thread_id) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index 7cf949341a7..ef57168f739 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -1207,6 +1207,12 @@ def _canonical_snapshot_message(message: dict[str, Any]) -> dict[str, Any]: normalized_message = agui_messages_to_snapshot_format([copy.deepcopy(message)])[0] normalized_message.pop("id", None) + if "toolCalls" in normalized_message: + normalized_message["tool_calls"] = normalized_message.pop("toolCalls") + if "toolCallId" in normalized_message: + normalized_message["tool_call_id"] = normalized_message.pop("toolCallId") + if "encryptedValue" in normalized_message: + normalized_message["encrypted_value"] = normalized_message.pop("encryptedValue") return cast(dict[str, Any], make_json_safe(normalized_message)) @@ -1215,10 +1221,31 @@ def _snapshot_messages_match(stored_message: dict[str, Any], incoming_message: d stored_id = stored_message.get("id") incoming_id = incoming_message.get("id") if stored_id and incoming_id: - return str(stored_id) == str(incoming_id) + if str(stored_id) == str(incoming_id): + return True + if normalize_agui_role(stored_message.get("role", "user")) != "assistant": + return False return _canonical_snapshot_message(stored_message) == _canonical_snapshot_message(incoming_message) +def _snapshot_overlap_length( + stored_messages: list[dict[str, Any]], + incoming_messages: list[dict[str, Any]], + *, + start: int, +) -> int: + """Return the contiguous incoming prefix matching stored history at ``start``.""" + compared_count = min(len(stored_messages) - start, len(incoming_messages)) + return next( + ( + offset + for offset, (stored_message, incoming_message) in enumerate(zip(stored_messages[start:], incoming_messages)) + if not _snapshot_messages_match(stored_message, incoming_message) + ), + compared_count, + ) + + def _latest_user_message_index(messages: list[dict[str, Any]]) -> int | None: """Find the newest incoming user message index.""" for index in range(len(messages) - 1, -1, -1): @@ -1231,17 +1258,24 @@ def _known_tool_call_ids( stored_messages: list[dict[str, Any]], stored_interrupt: list[dict[str, Any]] | None, ) -> set[str]: - """Collect tool call ids the backend previously issued for this thread.""" + """Collect unanswered tool call ids the backend previously issued for this thread.""" known_ids: set[str] = set() + completed_ids: set[str] = set() for message in stored_messages: tool_calls = message.get("tool_calls") or message.get("toolCalls") or [] if not isinstance(tool_calls, list): - continue + tool_calls = [] for tool_call in cast(list[Any], tool_calls): - if isinstance(tool_call, dict): - tool_call_id = cast(dict[str, Any], tool_call).get("id") - if tool_call_id: - known_ids.add(str(tool_call_id)) + if not isinstance(tool_call, dict): + continue + tool_call_id = cast(dict[str, Any], tool_call).get("id") + if tool_call_id: + known_ids.add(str(tool_call_id)) + if normalize_agui_role(message.get("role", "user")) == "tool": + tool_call_id = message.get("toolCallId") or message.get("tool_call_id") or message.get("actionExecutionId") + if tool_call_id: + completed_ids.add(str(tool_call_id)) + known_ids.difference_update(completed_ids) for interrupt in stored_interrupt or []: interrupt_id = interrupt.get("id") if interrupt_id: @@ -1295,17 +1329,36 @@ def _reconstruct_messages_from_thread_snapshot( if not stored_messages or not incoming_messages: return incoming_messages - incoming_suffix: list[dict[str, Any]] - if len(incoming_messages) >= len(stored_messages) and all( - _snapshot_messages_match(stored_message, incoming_message) - for stored_message, incoming_message in zip(stored_messages, incoming_messages) + overlap_candidates = [ + (start, _snapshot_overlap_length(stored_messages, incoming_messages, start=start)) + for start, stored_message in enumerate(stored_messages) + if _snapshot_messages_match(stored_message, incoming_messages[0]) + ] + matched_count = max( + ( + overlap + for start, overlap in overlap_candidates + if overlap >= 2 + or ( + stored_messages[start].get("id") + and incoming_messages[0].get("id") + and str(stored_messages[start]["id"]) == str(incoming_messages[0]["id"]) + ) + ), + default=0, + ) + if ( + matched_count == len(incoming_messages) + and matched_count + and normalize_agui_role(incoming_messages[-1].get("role", "user")) == "user" + and not incoming_messages[-1].get("id") ): - incoming_suffix = incoming_messages[len(stored_messages) :] + matched_count -= 1 + if matched_count: + incoming_suffix = incoming_messages[matched_count:] else: latest_user_index = _latest_user_message_index(incoming_messages) - if latest_user_index is None: - return incoming_messages - incoming_suffix = incoming_messages[latest_user_index:] + incoming_suffix = incoming_messages if latest_user_index is None else incoming_messages[latest_user_index:] incoming_suffix = _filter_untrusted_suffix( incoming_suffix, diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_snapshots.py b/python/packages/ag-ui/agent_framework_ag_ui/_snapshots.py index e862d3b52c7..7729ddb8380 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_snapshots.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_snapshots.py @@ -44,7 +44,8 @@ class AGUIThreadSnapshot: messages: Replayable AG-UI message snapshots. state: Optional AG-UI Shared State snapshot. interrupt: Optional interruption state from ``RUN_FINISHED.outcome.interrupts``. - session_state: Optional private serialized ``AgentSession.state`` payload. + session_state: Optional private serialized ``AgentSession.state`` payload, + including provider continuation metadata when service-session mode is active. """ messages: list[dict[str, Any]] = field(default_factory=list) diff --git a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py index 692b2809634..d2b7e9d6470 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py +++ b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py @@ -713,7 +713,11 @@ async def stream_fn( ) agent = Agent(client=streaming_chat_client_stub(stream_fn)) - wrapper = AgentFrameworkAgent(agent=agent, use_service_session=True) + wrapper = AgentFrameworkAgent( + agent=agent, + use_service_session=True, + service_session_id_from_thread_id=True, + ) input_data = {"messages": [{"role": "user", "content": "Hi"}], "thread_id": "conv_123456"} diff --git a/python/packages/ag-ui/tests/ag_ui/test_multi_turn.py b/python/packages/ag-ui/tests/ag_ui/test_multi_turn.py index 8d912c741ff..ae2d7f2ecea 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_multi_turn.py +++ b/python/packages/ag-ui/tests/ag_ui/test_multi_turn.py @@ -21,7 +21,12 @@ parse_sse_to_event_stream, ) -from agent_framework_ag_ui import AgentFrameworkAgent, add_agent_framework_fastapi_endpoint +from agent_framework_ag_ui import ( + AgentFrameworkAgent, + AGUIThreadSnapshot, + InMemoryAGUIThreadSnapshotStore, + add_agent_framework_fastapi_endpoint, +) def _build_app_with_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> FastAPI: @@ -330,3 +335,412 @@ async def test_workflow_interrupt_resume_round_trip() -> None: interrupt2 = stream2.run_finished_interrupts() assert interrupt2, "Expected hotel interrupt" assert stream2.interrupt_metadata_value(interrupt2[0])["agent"] == "hotels" + + +# ── Service-session snapshot authority ── + + +async def _run_service_turn( + runner: AgentFrameworkAgent, + *, + thread_id: str, + scope: str, + messages: list[dict[str, Any]], + resume: list[dict[str, Any]] | None = None, +) -> list[Any]: + body: dict[str, Any] = { + "threadId": thread_id, + "runId": f"run-{thread_id}", + "__ag_ui_snapshot_scope": scope, + "messages": messages, + } + if resume is not None: + body["resume"] = resume + return [event async for event in runner.run(body)] + + +def _service_session_runner( + updates: list[AgentResponseUpdate] | None = None, +) -> tuple[AgentFrameworkAgent, StubAgent, InMemoryAGUIThreadSnapshotStore]: + stub = StubAgent(updates=updates) + store = InMemoryAGUIThreadSnapshotStore() + return AgentFrameworkAgent(agent=stub, use_service_session=True, snapshot_store=store), stub, store + + +async def test_service_session_full_idless_snapshot_sends_only_new_turn() -> None: + runner, stub, store = _service_session_runner() + await _run_service_turn( + runner, thread_id="service-idless", scope="service", messages=[{"role": "user", "content": "first"}] + ) + snapshot = await store.get(scope="service", thread_id="service-idless") + assert snapshot is not None + idless_history = [{key: value for key, value in message.items() if key != "id"} for message in snapshot.messages] + + await _run_service_turn( + runner, + thread_id="service-idless", + scope="service", + messages=[*idless_history, {"role": "user", "content": "second"}], + ) + + assert [(message.role, message.text) for message in stub.messages_received] == [("user", "second")] + stored = await store.get(scope="service", thread_id="service-idless") + assert stored is not None + assert [message["role"] for message in stored.messages] == ["user", "assistant", "user", "assistant"] + + +async def test_service_session_matches_relayed_assistant_with_regenerated_id() -> None: + runner, stub, store = _service_session_runner() + await store.save( + scope="service", + thread_id="service-assistant-id", + snapshot=AGUIThreadSnapshot( + messages=[ + {"id": "u1", "role": "user", "content": "first"}, + {"id": "backend-a1", "role": "assistant", "content": "response"}, + ] + ), + ) + + await _run_service_turn( + runner, + thread_id="service-assistant-id", + scope="service", + messages=[ + {"id": "u1", "role": "user", "content": "first"}, + {"id": "client-a1", "role": "assistant", "content": "response"}, + {"id": "u2", "role": "user", "content": "second"}, + ], + ) + + assert [(message.role, message.text) for message in stub.messages_received] == [("user", "second")] + + +async def test_service_session_matches_relayed_tool_call_aliases() -> None: + runner, stub, store = _service_session_runner() + await store.save( + scope="service", + thread_id="service-tool-alias", + snapshot=AGUIThreadSnapshot( + messages=[ + {"role": "user", "content": "first"}, + { + "id": "backend-a1", + "role": "assistant", + "tool_calls": [ + {"id": "call-1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} + ], + }, + ] + ), + ) + + await _run_service_turn( + runner, + thread_id="service-tool-alias", + scope="service", + messages=[ + {"role": "user", "content": "first"}, + { + "id": "client-a1", + "role": "assistant", + "toolCalls": [{"id": "call-1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}], + }, + {"role": "tool", "toolCallId": "call-1", "content": "result"}, + ], + ) + + assert len(stub.messages_received) == 1 + assert stub.messages_received[0].role == "tool" + assert stub.messages_received[0].contents[0].call_id == "call-1" + + +async def test_service_session_incremental_only_turn_preserves_snapshot_history() -> None: + runner, stub, store = _service_session_runner() + await _run_service_turn( + runner, + thread_id="service-incremental", + scope="service", + messages=[{"id": "u1", "role": "user", "content": "first"}], + ) + + await _run_service_turn( + runner, + thread_id="service-incremental", + scope="service", + messages=[{"id": "u2", "role": "user", "content": "second"}], + ) + + assert [(message.role, message.text) for message in stub.messages_received] == [("user", "second")] + stored = await store.get(scope="service", thread_id="service-incremental") + assert stored is not None + assert [message["role"] for message in stored.messages] == ["user", "assistant", "user", "assistant"] + + +async def test_service_session_stale_history_preserves_backend_assistant_message() -> None: + runner, stub, store = _service_session_runner() + await _run_service_turn( + runner, + thread_id="service-stale", + scope="service", + messages=[{"id": "u1", "role": "user", "content": "first"}], + ) + + await _run_service_turn( + runner, + thread_id="service-stale", + scope="service", + messages=[ + {"id": "u1", "role": "user", "content": "first"}, + {"id": "u2", "role": "user", "content": "second"}, + ], + ) + + assert [(message.role, message.text) for message in stub.messages_received] == [("user", "second")] + stored = await store.get(scope="service", thread_id="service-stale") + assert stored is not None + assert [message["role"] for message in stored.messages] == ["user", "assistant", "user", "assistant"] + assert stored.messages[1]["content"] == "response" + + +async def test_service_session_matches_truncated_history_overlap() -> None: + runner, stub, store = _service_session_runner() + await store.save( + scope="service", + thread_id="service-truncated", + snapshot=AGUIThreadSnapshot( + messages=[ + {"id": "u1", "role": "user", "content": "first"}, + {"id": "a1", "role": "assistant", "content": "first response"}, + {"id": "u2", "role": "user", "content": "second"}, + {"id": "a2", "role": "assistant", "content": "second response"}, + ] + ), + ) + + await _run_service_turn( + runner, + thread_id="service-truncated", + scope="service", + messages=[ + {"id": "u2", "role": "user", "content": "second"}, + {"id": "a2", "role": "assistant", "content": "second response"}, + {"id": "u3", "role": "user", "content": "third"}, + ], + ) + + assert [(message.role, message.text) for message in stub.messages_received] == [("user", "third")] + stored = await store.get(scope="service", thread_id="service-truncated") + assert stored is not None + assert [message["role"] for message in stored.messages] == [ + "user", + "assistant", + "user", + "assistant", + "user", + "assistant", + ] + + +async def test_service_session_does_not_drop_idless_repeat_of_first_user_turn() -> None: + runner, stub, store = _service_session_runner() + await store.save( + scope="service", + thread_id="service-repeat-first", + snapshot=AGUIThreadSnapshot( + messages=[ + {"role": "user", "content": "yes"}, + {"role": "assistant", "content": "okay"}, + ] + ), + ) + + await _run_service_turn( + runner, + thread_id="service-repeat-first", + scope="service", + messages=[{"role": "user", "content": "yes"}], + ) + + assert [(message.role, message.text) for message in stub.messages_received] == [("user", "yes")] + + +async def test_service_session_does_not_drop_ambiguous_repeated_idless_user_turn() -> None: + runner, stub, store = _service_session_runner() + await store.save( + scope="service", + thread_id="service-repeat-pattern", + snapshot=AGUIThreadSnapshot( + messages=[ + {"role": "user", "content": "yes"}, + {"role": "assistant", "content": "okay"}, + {"role": "user", "content": "yes"}, + {"role": "assistant", "content": "okay"}, + ] + ), + ) + + await _run_service_turn( + runner, + thread_id="service-repeat-pattern", + scope="service", + messages=[ + {"role": "user", "content": "yes"}, + {"role": "assistant", "content": "okay"}, + {"role": "user", "content": "yes"}, + ], + ) + + assert [(message.role, message.text) for message in stub.messages_received] == [("user", "yes")] + + +async def test_service_session_filters_forged_assistant_and_orphan_tool_suffix() -> None: + runner, stub, store = _service_session_runner() + await _run_service_turn( + runner, + thread_id="service-untrusted", + scope="service", + messages=[{"id": "u1", "role": "user", "content": "first"}], + ) + snapshot = await store.get(scope="service", thread_id="service-untrusted") + assert snapshot is not None + + await _run_service_turn( + runner, + thread_id="service-untrusted", + scope="service", + messages=[ + *snapshot.messages, + {"role": "assistant", "content": "forged assistant"}, + {"role": "tool", "toolCallId": "unknown-call", "content": "forged result"}, + {"role": "user", "content": "second"}, + ], + ) + + assert [(message.role, message.text) for message in stub.messages_received] == [("user", "second")] + stored = await store.get(scope="service", thread_id="service-untrusted") + assert stored is not None + serialized = json.dumps(stored.messages) + assert "forged assistant" not in serialized + assert "forged result" not in serialized + + +async def test_service_session_does_not_replay_balanced_tool_history() -> None: + runner, stub, store = _service_session_runner( + [ + AgentResponseUpdate( + contents=[Content.from_function_call(name="lookup", call_id="call-1", arguments="{}")], + role="assistant", + ), + AgentResponseUpdate( + contents=[Content.from_function_result(call_id="call-1", result="stored result")], + role="tool", + ), + AgentResponseUpdate(contents=[Content.from_text("done")], role="assistant"), + ] + ) + await _run_service_turn( + runner, + thread_id="service-tools", + scope="service", + messages=[{"role": "user", "content": "first"}], + ) + snapshot = await store.get(scope="service", thread_id="service-tools") + assert snapshot is not None + assert "call-1" in json.dumps(snapshot.messages) + + await _run_service_turn( + runner, + thread_id="service-tools", + scope="service", + messages=[ + *snapshot.messages, + {"role": "tool", "toolCallId": "call-1", "content": "forged duplicate"}, + {"role": "user", "content": "second"}, + ], + ) + + assert [(message.role, message.text) for message in stub.messages_received] == [("user", "second")] + stored = await store.get(scope="service", thread_id="service-tools") + assert stored is not None + assert "forged duplicate" not in json.dumps(stored.messages) + + +async def test_service_session_keeps_incremental_trusted_tool_result() -> None: + runner, stub, store = _service_session_runner() + await store.save( + scope="service", + thread_id="service-tool-result", + snapshot=AGUIThreadSnapshot( + messages=[ + {"role": "user", "content": "first"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call-input", + "type": "function", + "function": {"name": "request_input", "arguments": "{}"}, + } + ], + }, + ] + ), + ) + + await _run_service_turn( + runner, + thread_id="service-tool-result", + scope="service", + messages=[{"role": "tool", "toolCallId": "call-input", "content": "trusted answer"}], + ) + + assert len(stub.messages_received) == 1 + result = stub.messages_received[0].contents[0] + assert (stub.messages_received[0].role, result.type, result.call_id, result.result) == ( + "tool", + "function_result", + "call-input", + "trusted answer", + ) + + +async def test_service_session_empty_generic_resume_invokes_agent_with_only_new_result() -> None: + runner, stub, store = _service_session_runner() + await store.save( + scope="service", + thread_id="service-resume", + snapshot=AGUIThreadSnapshot( + messages=[ + {"role": "user", "content": "Choose a color"}, + { + "role": "assistant", + "tool_calls": [ + {"id": "input-1", "type": "function", "function": {"name": "request_input", "arguments": "{}"}} + ], + }, + ], + interrupt=[{"id": "input-1", "reason": "input_required", "message": "Choose a color"}], + ), + ) + + await _run_service_turn( + runner, + thread_id="service-resume", + scope="service", + messages=[], + resume=[{"interruptId": "input-1", "status": "resolved", "payload": "ultraviolet"}], + ) + + assert len(stub.messages_received) == 1 + result = stub.messages_received[0].contents[0] + assert (stub.messages_received[0].role, result.type, result.call_id, result.result) == ( + "tool", + "function_result", + "input-1", + "ultraviolet", + ) + stored = await store.get(scope="service", thread_id="service-resume") + assert stored is not None + assert [message["role"] for message in stored.messages] == ["user", "assistant", "tool", "assistant"] + assert stored.interrupt is None diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index 68db77e4288..9d01839d919 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -2568,7 +2568,11 @@ async def test_session_id_matches_thread_id_with_service_session(): from agent_framework_ag_ui import AgentFrameworkAgent stub = StubAgent() - agent = AgentFrameworkAgent(agent=stub, use_service_session=True) + agent = AgentFrameworkAgent( + agent=stub, + use_service_session=True, + service_session_id_from_thread_id=True, + ) payload = { "thread_id": "service-thread-789", @@ -2615,7 +2619,11 @@ async def test_service_session_no_thread_id_generates_uuid(): from agent_framework_ag_ui import AgentFrameworkAgent stub = StubAgent() - agent = AgentFrameworkAgent(agent=stub, use_service_session=True) + agent = AgentFrameworkAgent( + agent=stub, + use_service_session=True, + service_session_id_from_thread_id=True, + ) payload = { "run_id": "run-5", @@ -2629,3 +2637,50 @@ async def test_service_session_no_thread_id_generates_uuid(): uuid.UUID(stub.last_session.session_id) # service_session_id should be None since no thread_id was supplied assert stub.last_session.service_session_id is None + + +async def test_provider_owned_service_session_requires_snapshot_persistence(): + """Provider-owned continuation needs private snapshot storage across AG-UI requests.""" + from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] + + from agent_framework_ag_ui import AgentFrameworkAgent + + agent = AgentFrameworkAgent(agent=StubAgent(), use_service_session=True) + + with pytest.raises(ValueError, match="requires snapshot persistence"): + _ = [ + event + async for event in agent.run( + { + "thread_id": "frontend-thread", + "run_id": "run-6", + "messages": [{"role": "user", "content": "Hello"}], + } + ) + ] + + +async def test_service_session_rejects_disabled_provider_storage(): + """Service-session continuation cannot work when provider storage is disabled.""" + from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] + + from agent_framework_ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore + + agent = AgentFrameworkAgent( + agent=StubAgent(default_options={"store": False}), + use_service_session=True, + snapshot_store=InMemoryAGUIThreadSnapshotStore(), + ) + + with pytest.raises(ValueError, match="requires provider storage"): + _ = [ + event + async for event in agent.run( + { + "thread_id": "frontend-thread", + "run_id": "run-store-false", + "__ag_ui_snapshot_scope": "test", + "messages": [{"role": "user", "content": "Hello"}], + } + ) + ] diff --git a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py index 3ed4d1b664c..f3fab4ca26f 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py @@ -16,6 +16,7 @@ ChatResponse, ChatResponseUpdate, Content, + FinishReason, FinishReasonLiteral, FunctionInvocationConfiguration, FunctionInvocationLayer, @@ -215,6 +216,18 @@ def _apply_option_translations(options: dict[str, Any]) -> None: } +def _map_finish_reason(stop_reason: str | None) -> FinishReason | None: + """Map an Anthropic stop reason onto a framework finish reason. + + Stop reasons that are not in ``FINISH_REASON_MAP`` are passed through unchanged so that + values added by the Anthropic API, such as ``model_context_window_exceeded``, still reach + the caller instead of being reported as no finish reason at all. + """ + if not stop_reason: + return None + return FinishReason(FINISH_REASON_MAP.get(stop_reason, stop_reason)) + + class AnthropicSettings(TypedDict, total=False): """Anthropic Project settings. @@ -1095,7 +1108,7 @@ def _process_message(self, message: BetaMessage, options: Mapping[str, Any]) -> ], usage_details=self._parse_usage_from_anthropic(message.usage), model=message.model, - finish_reason=FINISH_REASON_MAP.get(message.stop_reason) if message.stop_reason else None, + finish_reason=_map_finish_reason(message.stop_reason), response_format=options.get("response_format"), raw_representation=message, ) @@ -1131,9 +1144,7 @@ def _process_stream_event( *usage_details, ], model=event.message.model, - finish_reason=FINISH_REASON_MAP.get(event.message.stop_reason) - if event.message.stop_reason - else None, + finish_reason=_map_finish_reason(event.message.stop_reason), raw_representation=event, ) case "message_delta": @@ -1147,7 +1158,7 @@ def _process_stream_event( ] if usage else [], - finish_reason=FINISH_REASON_MAP.get(event.delta.stop_reason) if event.delta.stop_reason else None, + finish_reason=_map_finish_reason(event.delta.stop_reason), raw_representation=event, ) case "message_stop": diff --git a/python/packages/anthropic/tests/test_anthropic_client.py b/python/packages/anthropic/tests/test_anthropic_client.py index e6f8cc36a21..d0b9a4a8f1a 100644 --- a/python/packages/anthropic/tests/test_anthropic_client.py +++ b/python/packages/anthropic/tests/test_anthropic_client.py @@ -1530,6 +1530,74 @@ def test_process_message_with_tool_use(mock_anthropic_client: MagicMock) -> None assert response.finish_reason == "tool_calls" +@pytest.mark.parametrize( + ("stop_reason", "expected"), + [ + ("end_turn", "stop"), + ("stop_sequence", "stop"), + ("pause_turn", "stop"), + ("max_tokens", "length"), + ("tool_use", "tool_calls"), + ("refusal", "content_filter"), + # Not in FINISH_REASON_MAP: passed through instead of dropped. + ("model_context_window_exceeded", "model_context_window_exceeded"), + (None, None), + ], +) +def test_process_message_finish_reason( + mock_anthropic_client: MagicMock, stop_reason: str | None, expected: str | None +) -> None: + """Known stop reasons map, unmapped ones are preserved, and an absent one stays absent.""" + client = create_test_anthropic_client(mock_anthropic_client) + + mock_message = MagicMock(spec=BetaMessage) + mock_message.id = "msg_123" + mock_message.model = "claude-3-5-sonnet-20241022" + mock_message.content = [BetaTextBlock(type="text", text="Hello there!")] + mock_message.usage = BetaUsage(input_tokens=10, output_tokens=5) + mock_message.stop_reason = stop_reason + + response = client._process_message(mock_message, {}) + + assert response.finish_reason == expected + + +def test_process_stream_event_preserves_unmapped_stop_reason(mock_anthropic_client: MagicMock) -> None: + """Streaming message_delta events pass an unmapped stop reason through unchanged.""" + client = create_test_anthropic_client(mock_anthropic_client) + + mock_event = MagicMock() + mock_event.type = "message_delta" + mock_event.usage = None + mock_event.delta.stop_reason = "model_context_window_exceeded" + + result = client._process_stream_event(mock_event) + + assert result is not None + assert result.finish_reason == "model_context_window_exceeded" + + +def test_process_stream_event_message_start_preserves_unmapped_stop_reason( + mock_anthropic_client: MagicMock, +) -> None: + """Streaming message_start events pass an unmapped stop reason through unchanged.""" + client = create_test_anthropic_client(mock_anthropic_client) + + mock_event = MagicMock() + mock_event.type = "message_start" + mock_event.message.id = "msg_abc" + mock_event.message.role = "assistant" + mock_event.message.model = "claude-3-5-sonnet-20241022" + mock_event.message.content = [] + mock_event.message.usage = None + mock_event.message.stop_reason = "model_context_window_exceeded" + + result = client._process_stream_event(mock_event) + + assert result is not None + assert result.finish_reason == "model_context_window_exceeded" + + def test_parse_usage_from_anthropic_basic(mock_anthropic_client: MagicMock) -> None: """Test _parse_usage_from_anthropic with basic usage.""" client = create_test_anthropic_client(mock_anthropic_client) diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 288d8a29e33..f6224421343 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -1512,6 +1512,9 @@ async def _prepare_run_context( # _merge_options strips unset (None) options, so e.g. an unset `store` is not forwarded # and the service decides its own default. co = _merge_options(chat_options, run_opts) + # The loop marker must remain on SessionContext.options for after_run provider + # scoping, but it is framework-private metadata and must not reach the client. + co.pop(_LOOP_ITERATION_TOKEN_KEY, None) # Build session_messages from session context: context messages + input messages session_messages: list[Message] = session_context.get_messages(include_input=True) diff --git a/python/packages/core/tests/core/test_harness_loop.py b/python/packages/core/tests/core/test_harness_loop.py index cbf20c64dfc..5332f65cd21 100644 --- a/python/packages/core/tests/core/test_harness_loop.py +++ b/python/packages/core/tests/core/test_harness_loop.py @@ -32,9 +32,9 @@ background_tasks_running, background_tasks_running_message, set_agent_mode, - tool, todos_remaining, todos_remaining_message, + tool, ) from agent_framework._harness._loop import ( DEFAULT_JUDGE_MAX_ITERATIONS, @@ -1424,6 +1424,55 @@ async def after_run(self, *, agent: Any, session: Any, context: Any, state: dict self.after_calls += 1 +class _StrictTransportChatClient(RecordingChatClient): + def __init__(self) -> None: + super().__init__() + self.received_provider_options: list[str] = [] + + def _inner_get_response( + self, + *, + messages: Sequence[Message], + stream: bool = False, + options: Mapping[str, Any], + **kwargs: Any, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + self._strict_transport(**options) + return super()._inner_get_response(messages=messages, stream=stream, options=options, **kwargs) + + def _strict_transport(self, *, tool_choice: Any, _provider_option: str) -> None: + self.received_provider_options.append(_provider_option) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) +async def test_loop_marker_is_not_forwarded_to_provider_transport(stream: bool) -> None: + client = _StrictTransportChatClient() + turn_scoped = _TurnScopedRecordingProvider() + agent = Agent( + client=client, + middleware=[AgentLoopMiddleware(always_continue, max_iterations=1)], + context_providers=[turn_scoped], + ) + options = {"_provider_option": "kept"} + + if stream: + response = agent.run( # type: ignore[call-overload] # pyrefly: ignore[no-matching-overload] # ty: ignore[no-matching-overload] + "start", + stream=True, + options=options, # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + ) + _ = [update async for update in response] + await response.get_final_response() + else: + await agent.run( # type: ignore[call-overload] # pyrefly: ignore[no-matching-overload] # ty: ignore[no-matching-overload] + "start", + options=options, # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + ) + + assert client.received_provider_options == ["kept"] + assert turn_scoped.after_calls == 1 + + async def test_turn_scoped_after_run_fires_once_per_loop() -> None: client = RecordingChatClient() turn_scoped = _TurnScopedRecordingProvider() diff --git a/python/packages/foundry/tests/foundry/test_agui_provider_matrix.py b/python/packages/foundry/tests/foundry/test_agui_provider_matrix.py new file mode 100644 index 00000000000..2a23d51236b --- /dev/null +++ b/python/packages/foundry/tests/foundry/test_agui_provider_matrix.py @@ -0,0 +1,244 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Live AG-UI multi-turn coverage for Foundry provider continuation modes.""" + +from __future__ import annotations + +import json +import os +from typing import Any, Literal, cast +from uuid import uuid4 + +import pytest +from agent_framework import Agent, AgentSession, Message +from agent_framework_ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore +from azure.ai.projects import models as projects_models +from azure.ai.projects.aio import AIProjectClient +from azure.identity import AzureCliCredential + +from agent_framework_foundry import FoundryAgent, FoundryChatClient + +skip_if_foundry_integration_tests_disabled = pytest.mark.skipif( + os.getenv("FOUNDRY_PROJECT_ENDPOINT", "") in ("", "https://test-project.services.ai.azure.com/") + or os.getenv("FOUNDRY_MODEL", "") == "", + reason="No real FOUNDRY_PROJECT_ENDPOINT or FOUNDRY_MODEL provided; skipping integration tests.", +) +skip_if_foundry_hosted_agent_integration_tests_disabled = pytest.mark.skipif( + os.getenv("FOUNDRY_PROJECT_ENDPOINT", "") in ("", "https://test-project.services.ai.azure.com/") + or os.getenv("FOUNDRY_AGENT_NAME", "") == "", + reason="No real FOUNDRY_PROJECT_ENDPOINT or FOUNDRY_AGENT_NAME provided; skipping integration tests.", +) + +_Mode = Literal["stateless", "conversation", "previous_response"] +_PROVIDER_SESSION_KEY = "__ag_ui_provider_service_session_id" +_MODES = [ + pytest.param("stateless", False, id="stateless-snapshot-replay"), + pytest.param("conversation", True, id="conversation"), + pytest.param("previous_response", True, id="previous-response-id"), +] +_HOSTED_AGENT_MODES = [ + # TODO: Remove this strict xfail when Foundry Hosted Agents accept standard + # output_text replay without requiring fabricated logprobs. + pytest.param( + "stateless", + False, + id="stateless-snapshot-replay", + marks=pytest.mark.xfail( + strict=True, + reason="Foundry Hosted Agent rejects output_text replay when logprobs are absent", + ), + ), + pytest.param("conversation", True, id="conversation"), + pytest.param("previous_response", True, id="previous-response-id"), +] + + +class _CapturingAgent: + """Delegate to an Agent without exposing its conversation factory.""" + + def __init__(self, agent: Any) -> None: + self._agent = agent + self.inputs: list[list[Message]] = [] + self.service_session_ids: list[Any] = [] + self.created_conversation_ids: list[str] = [] + + def run(self, messages: Any = None, **kwargs: Any) -> Any: + self.inputs.append(list(messages) if isinstance(messages, list) else [messages]) + session = kwargs.get("session") + self.service_session_ids.append(session.service_session_id if session is not None else None) + return self._agent.run(messages, **kwargs) + + def __getattr__(self, name: str) -> Any: + if name == "create_conversation": + raise AttributeError(name) + return getattr(self._agent, name) + + +class _ConversationCapturingAgent(_CapturingAgent): + """Expose backend conversation creation to AgentFrameworkAgent.""" + + async def create_conversation(self, *, session_id: str | None = None) -> AgentSession: + conversation = await self._agent.client.client.conversations.create() + self.created_conversation_ids.append(conversation.id) + return self._agent.get_session(conversation.id, session_id=session_id) + + +async def _exercise_agui_case(agent: Any, mode: _Mode) -> None: + marker = f"AF-AGUI-{uuid4().hex}" + follow_up = "Return only the exact marker from my previous message." + thread_id = f"agui-thread-{uuid4().hex}" + scope = f"agui-scope-{uuid4().hex}" + snapshot_store = InMemoryAGUIThreadSnapshotStore() + capturing = _ConversationCapturingAgent(agent) if mode == "conversation" else _CapturingAgent(agent) + + runner = AgentFrameworkAgent( + agent=cast(Any, capturing), + use_service_session=mode != "stateless", + service_session_id_from_thread_id=False, + snapshot_store=snapshot_store, + ) + + try: + first_events = [ + event + async for event in runner.run({ + "threadId": thread_id, + "runId": f"run-1-{uuid4().hex}", + "__ag_ui_snapshot_scope": scope, + "messages": [ + { + "role": "user", + "content": f"Remember this exact marker for my next message: {marker}", + } + ], + }) + ] + first_snapshot = next( + event for event in reversed(first_events) if getattr(event, "type", None) == "MESSAGES_SNAPSHOT" + ) + replay_messages = cast(list[dict[str, Any]], first_snapshot.model_dump(by_alias=True)["messages"]) + first_stored = await snapshot_store.get(scope=scope, thread_id=thread_id) + assert first_stored is not None + + second_events = [ + event + async for event in runner.run({ + "threadId": thread_id, + "runId": f"run-2-{uuid4().hex}", + "__ag_ui_snapshot_scope": scope, + "messages": [*replay_messages, {"role": "user", "content": follow_up}], + }) + ] + + assert not [event for event in second_events if getattr(event, "type", None) == "RUN_ERROR"] + assert [[message.role for message in messages] for messages in capturing.inputs[:1]] == [["user"]] + if mode == "stateless": + assert [message.role for message in capturing.inputs[1]] == ["user", "assistant", "user"] + assert capturing.service_session_ids == [None, None] + assert first_stored.session_state is None + else: + assert [(message.role, message.text) for message in capturing.inputs[1]] == [("user", follow_up)] + assert first_stored.session_state is not None + provider_id = first_stored.session_state[_PROVIDER_SESSION_KEY] + assert isinstance(provider_id, str) + assert provider_id != thread_id + assert capturing.service_session_ids[1] == provider_id + assert provider_id not in json.dumps(first_stored.messages) + if mode == "conversation": + assert provider_id.startswith("conv_") + assert capturing.service_session_ids == [provider_id, provider_id] + assert capturing.created_conversation_ids == [provider_id] + else: + assert provider_id.startswith(("resp_", "caresp_")) + assert capturing.service_session_ids[0] is None + assert not capturing.created_conversation_ids + + assert marker in capturing.inputs[0][0].text + assert capturing.inputs[1][-1].text == follow_up + response_text = "".join( + str(getattr(event, "delta", "")) + for event in second_events + if getattr(event, "type", None) == "TEXT_MESSAGE_CONTENT" + ) + assert marker in response_text + + final_stored = await snapshot_store.get(scope=scope, thread_id=thread_id) + assert final_stored is not None + assert [message.get("role") for message in final_stored.messages].count("user") == 2 + assert [message.get("role") for message in final_stored.messages].count("assistant") >= 2 + finally: + for conversation_id in capturing.created_conversation_ids: + await agent.client.client.conversations.delete(conversation_id) + + +@pytest.mark.flaky +@pytest.mark.integration +@skip_if_foundry_integration_tests_disabled +@pytest.mark.parametrize(("mode", "store"), _MODES) +async def test_foundry_chat_client_agui_provider_matrix(mode: _Mode, store: bool) -> None: + credential = AzureCliCredential() + client = FoundryChatClient(credential=cast(Any, credential)) + agent = Agent(client=cast(Any, client), default_options=cast(Any, {"store": store})) + try: + await _exercise_agui_case(agent, mode) + finally: + await client.client.close() + await client.project_client.close() + credential.close() + + +@pytest.mark.flaky +@pytest.mark.integration +@skip_if_foundry_integration_tests_disabled +@pytest.mark.parametrize(("mode", "store"), _MODES) +async def test_foundry_prompt_agent_agui_provider_matrix(mode: _Mode, store: bool) -> None: + credential = AzureCliCredential() + project_client = AIProjectClient( + endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + credential=cast(Any, credential), + allow_preview=True, + ) + created_agent: Any | None = None + prompt_agent: FoundryAgent | None = None + try: + created_agent = await project_client.agents.create_version( + agent_name=f"af-agui-{uuid4().hex[:12]}", + definition=projects_models.PromptAgentDefinition( + model=os.environ["FOUNDRY_MODEL"], + instructions="Follow the user instructions exactly and answer concisely.", + ), + ) + prompt_agent = FoundryAgent( + project_client=project_client, + agent_name=created_agent.name, + agent_version=created_agent.version, + allow_preview=False, + default_options=cast(Any, {"store": store}), + ) + await _exercise_agui_case(prompt_agent, mode) + finally: + if created_agent is not None: + await project_client.agents.delete(agent_name=created_agent.name, force=True) + if prompt_agent is not None: + await cast(Any, prompt_agent.client).client.close() + await project_client.close() + credential.close() + + +@pytest.mark.flaky +@pytest.mark.integration +@skip_if_foundry_hosted_agent_integration_tests_disabled +@pytest.mark.parametrize(("mode", "store"), _HOSTED_AGENT_MODES) +async def test_foundry_hosted_agent_agui_provider_matrix(mode: _Mode, store: bool) -> None: + credential = AzureCliCredential() + hosted_agent = FoundryAgent( + credential=cast(Any, credential), + allow_preview=True, + default_options=cast(Any, {"store": store}), + ) + try: + await _exercise_agui_case(hosted_agent, mode) + finally: + await cast(Any, hosted_agent.client).client.close() + await cast(Any, hosted_agent.client).close() + credential.close() diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index 9ae2938ed62..a4c3b8ca25a 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/python/packages/foundry/tests/foundry/test_foundry_agent.py @@ -32,6 +32,7 @@ tool, ) from agent_framework.foundry import FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY +from agent_framework_ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore from agent_framework_openai._chat_client import RawOpenAIChatClient from agent_framework_openai._feature_usage import FeatureIndex as OpenAIFeatureIndex from azure.ai.projects import models as projects_models @@ -1525,6 +1526,98 @@ async def test_foundry_agent_basic_run() -> None: assert "response test" in response.text.lower() +@pytest.mark.flaky +@pytest.mark.integration +@skip_if_foundry_agent_integration_tests_disabled +@pytest.mark.parametrize("continuation_mode", ["conversation", "native_agui"]) +async def test_foundry_agent_ag_ui_service_session_continues_without_history_replay( + monkeypatch: pytest.MonkeyPatch, + continuation_mode: str, +) -> None: + """AG-UI supports Foundry conversation and response continuation without history replay.""" + marker = f"AF-AGUI-{uuid4().hex}" + scope = f"foundry-{continuation_mode}-session-{uuid4().hex}" + store = InMemoryAGUIThreadSnapshotStore() + + async with FoundryAgent(credential=cast(Any, AzureCliCredential()), allow_preview=True) as foundry_agent: + if continuation_mode == "conversation": + conversation = await foundry_agent.create_conversation() + thread_id = cast(str, conversation.service_session_id) + else: + thread_id = str(uuid4()) + provider_inputs: list[list[Message]] = [] + provider_session_ids: list[Any] = [] + original_run = foundry_agent.run + + def capture_provider_input(messages: Any = None, **kwargs: Any) -> Any: + provider_inputs.append(list(messages) if isinstance(messages, list) else [messages]) + session = kwargs.get("session") + provider_session_ids.append(session.service_session_id if session is not None else None) + return original_run(messages, **kwargs) + + monkeypatch.setattr(foundry_agent, "run", capture_provider_input) + runner = AgentFrameworkAgent( + agent=foundry_agent, + use_service_session=True, + service_session_id_from_thread_id=continuation_mode == "conversation", + snapshot_store=store, + ) + first_events = [ + event + async for event in runner.run({ + "threadId": thread_id, + "runId": f"foundry-ag-ui-{continuation_mode}-first", + "__ag_ui_snapshot_scope": scope, + "messages": [ + { + "role": "user", + "content": f"Remember this exact marker for my next message: {marker}", + } + ], + }) + ] + first_snapshot = next( + event for event in reversed(first_events) if getattr(event, "type", None) == "MESSAGES_SNAPSHOT" + ) + prior_messages = cast(list[dict[str, Any]], first_snapshot.model_dump(by_alias=True)["messages"]) + follow_up = "Return only the exact marker I asked you to remember." + + second_events = [ + event + async for event in runner.run({ + "threadId": thread_id, + "runId": f"foundry-ag-ui-{continuation_mode}-second", + "__ag_ui_snapshot_scope": scope, + "messages": [ + *prior_messages, + {"role": "user", "content": follow_up}, + ], + }) + ] + + assert not [event for event in second_events if getattr(event, "type", None) == "RUN_ERROR"] + assert len(provider_inputs) == 2 + assert [(message.role, message.text) for message in provider_inputs[1]] == [("user", follow_up)] + if continuation_mode == "conversation": + assert provider_session_ids == [thread_id, thread_id] + else: + assert all(isinstance(service_session_id, str) for service_session_id in provider_session_ids) + assert provider_session_ids[0] == provider_session_ids[1] + assert provider_session_ids[0].startswith("conv_") + assert provider_session_ids[0] != thread_id + response_text = "".join( + str(getattr(event, "delta", "")) + for event in second_events + if getattr(event, "type", None) == "TEXT_MESSAGE_CONTENT" + ) + assert marker in response_text + stored = await store.get(scope=scope, thread_id=thread_id) + assert stored is not None + assert sum(message.get("role") == "user" for message in stored.messages) == 2 + assert sum(message.get("role") == "assistant" for message in stored.messages) >= 2 + assert marker in json.dumps(stored.messages) + + @pytest.mark.flaky @pytest.mark.integration @skip_if_foundry_agent_integration_tests_disabled diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index f82ffb3e541..11222be8214 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -47,6 +47,9 @@ executor, tool, ) +from agent_framework.ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore +from agent_framework.exceptions import ChatClientException +from agent_framework.openai import OpenAIChatClient from azure.ai.agentserver.core import get_request_context from azure.ai.agentserver.responses import ( FileResponseStore, @@ -59,6 +62,7 @@ from azure.ai.agentserver.responses.streaming._checkpoint import ResponseCheckpointEvent from mcp import McpError from mcp.types import ErrorData +from openai import AsyncOpenAI from typing_extensions import Any from agent_framework_foundry_hosting import ResponsesHostServer @@ -285,6 +289,19 @@ def _make_server(agent: Any, **kwargs: Any) -> ResponsesHostServer: return server +class _CapturingASGITransport(httpx.AsyncBaseTransport): + def __init__(self, app: Any) -> None: + self._transport = httpx.ASGITransport(app=app) + self.payloads: list[dict[str, Any]] = [] + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + self.payloads.append(json.loads(await request.aread())) + return await self._transport.handle_async_request(request) + + async def aclose(self) -> None: + await self._transport.aclose() + + async def _post( server: ResponsesHostServer, *, @@ -342,6 +359,258 @@ def _parse_sse_events(body: str) -> list[dict[str, Any]]: return events +async def test_agui_service_storage_conversation_mode_sends_only_incremental_provider_input() -> None: + """A provider conversation stays authoritative while AG-UI snapshots retain full history.""" + hosted_agent_backend = _make_agent( + response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("ACK")])]) + ) + local_hosted_agent_api = _make_server(hosted_agent_backend) + transport = _CapturingASGITransport(local_hosted_agent_api) + responses_client = AsyncOpenAI( + api_key="test-key", + base_url="http://test", + http_client=httpx.AsyncClient(transport=transport), + max_retries=0, + ) + store = InMemoryAGUIThreadSnapshotStore() + hosted_agent_client = Agent(client=OpenAIChatClient(model="test-model", async_client=responses_client)) + runner = AgentFrameworkAgent( + agent=hosted_agent_client, + use_service_session=True, + service_session_id_from_thread_id=True, + snapshot_store=store, + ) + thread_id = "conv_agui_service_storage" + + try: + first_events = [ + event + async for event in runner.run({ + "thread_id": thread_id, + "__ag_ui_snapshot_scope": "test", + "messages": [{"role": "user", "content": "first"}], + }) + ] + first_snapshot = next( + event.model_dump(by_alias=True)["messages"] + for event in reversed(first_events) + if getattr(event, "type", None) == "MESSAGES_SNAPSHOT" + ) + second_events = [ + event + async for event in runner.run({ + "thread_id": thread_id, + "__ag_ui_snapshot_scope": "test", + "messages": [*first_snapshot, {"role": "user", "content": "second"}], + }) + ] + finally: + await responses_client.close() + + provider_inputs = [ + [item for item in payload["input"] if item.get("type") == "message"] for payload in transport.payloads + ] + assert [[item["role"] for item in items] for items in provider_inputs] == [["user"], ["user"]] + assert [items[0]["content"][0]["text"] for items in provider_inputs] == ["first", "second"] + assert [payload["conversation"] for payload in transport.payloads] == [thread_id, thread_id] + assert all("previous_response_id" not in payload for payload in transport.payloads) + assert hosted_agent_backend.run.call_count == 2 + hosted_turns = [call.kwargs["messages"] for call in hosted_agent_backend.run.call_args_list] + assert [turn[-1].text for turn in hosted_turns] == ["first", "second"] + + second_snapshot = next( + event.model_dump(by_alias=True)["messages"] + for event in reversed(second_events) + if getattr(event, "type", None) == "MESSAGES_SNAPSHOT" + ) + assert [message["role"] for message in second_snapshot] == ["user", "assistant", "user", "assistant"] + + +async def test_agui_service_storage_native_uuid_uses_backend_created_conversation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A backend conversation factory maps a native AG-UI UUID to a provider conversation.""" + hosted_agent_backend = _make_agent( + response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("ACK")])]) + ) + transport = _CapturingASGITransport(_make_server(hosted_agent_backend)) + responses_client = AsyncOpenAI( + api_key="test-key", + base_url="http://test", + http_client=httpx.AsyncClient(transport=transport), + max_retries=0, + ) + hosted_agent_client = Agent(client=OpenAIChatClient(model="test-model", async_client=responses_client)) + + async def create_conversation(*, session_id: str) -> AgentSession: + return AgentSession(session_id=session_id, service_session_id="conv_backend_created") + + monkeypatch.setattr(hosted_agent_client, "create_conversation", create_conversation, raising=False) + runner = AgentFrameworkAgent( + agent=hosted_agent_client, + use_service_session=True, + snapshot_store=InMemoryAGUIThreadSnapshotStore(), + ) + thread_id = "86052504-791b-47d8-a405-60ce167ac93a" + + try: + first_events = [ + event + async for event in runner.run({ + "thread_id": thread_id, + "__ag_ui_snapshot_scope": "test", + "messages": [{"role": "user", "content": "first"}], + }) + ] + first_snapshot = next( + event.model_dump(by_alias=True)["messages"] + for event in reversed(first_events) + if getattr(event, "type", None) == "MESSAGES_SNAPSHOT" + ) + _ = [ + event + async for event in runner.run({ + "thread_id": thread_id, + "__ag_ui_snapshot_scope": "test", + "messages": [*first_snapshot, {"role": "user", "content": "second"}], + }) + ] + finally: + await responses_client.close() + + assert [payload["conversation"] for payload in transport.payloads] == [ + "conv_backend_created", + "conv_backend_created", + ] + assert all("previous_response_id" not in payload for payload in transport.payloads) + provider_inputs = [ + [item for item in payload["input"] if item.get("type") == "message"] for payload in transport.payloads + ] + assert [items[0]["content"][0]["text"] for items in provider_inputs] == ["first", "second"] + assert hosted_agent_backend.run.call_count == 2 + + +async def test_agui_service_storage_response_mode_persists_provider_continuation_for_uuid_thread() -> None: + """A native AG-UI UUID stays separate from the provider-issued Responses continuation.""" + hosted_agent_backend = _make_agent( + response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("ACK")])]) + ) + transport = _CapturingASGITransport(_make_server(hosted_agent_backend)) + responses_client = AsyncOpenAI( + api_key="test-key", + base_url="http://test", + http_client=httpx.AsyncClient(transport=transport), + max_retries=0, + ) + store = InMemoryAGUIThreadSnapshotStore() + runner = AgentFrameworkAgent( + agent=Agent(client=OpenAIChatClient(model="test-model", async_client=responses_client)), + use_service_session=True, + snapshot_store=store, + ) + thread_id = "86052504-791b-47d8-a405-60ce167ac93a" + + try: + first_events = [ + event + async for event in runner.run({ + "thread_id": thread_id, + "__ag_ui_snapshot_scope": "test", + "messages": [{"role": "user", "content": "first"}], + }) + ] + first_snapshot = next( + event.model_dump(by_alias=True)["messages"] + for event in reversed(first_events) + if getattr(event, "type", None) == "MESSAGES_SNAPSHOT" + ) + second_events = [ + event + async for event in runner.run({ + "thread_id": thread_id, + "__ag_ui_snapshot_scope": "test", + "messages": [*first_snapshot, {"role": "user", "content": "second"}], + }) + ] + finally: + await responses_client.close() + + assert "previous_response_id" not in transport.payloads[0] + assert "conversation" not in transport.payloads[0] + assert transport.payloads[1]["previous_response_id"] != thread_id + assert transport.payloads[1]["previous_response_id"].startswith(("resp_", "caresp_", "response_")) + assert hosted_agent_backend.run.call_count == 2 + + stored = await store.get(scope="test", thread_id=thread_id) + assert stored is not None + assert stored.session_state is not None + second_snapshot = next( + event.model_dump(by_alias=True)["messages"] + for event in reversed(second_events) + if getattr(event, "type", None) == "MESSAGES_SNAPSHOT" + ) + assert [message["role"] for message in second_snapshot] == ["user", "assistant", "user", "assistant"] + + +async def test_agui_stateless_store_true_does_not_restore_provider_continuation() -> None: + """Stateless snapshot replay must not combine full history with a stored response id.""" + hosted_agent_backend = _make_agent( + response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("ACK")])]) + ) + transport = _CapturingASGITransport(_make_server(hosted_agent_backend)) + responses_client = AsyncOpenAI( + api_key="test-key", + base_url="http://test", + http_client=httpx.AsyncClient(transport=transport), + max_retries=0, + ) + store = InMemoryAGUIThreadSnapshotStore() + runner = AgentFrameworkAgent( + agent=Agent( + client=OpenAIChatClient( # ty: ignore[invalid-argument-type] + model="test-model", + async_client=responses_client, + ), + default_options={"store": True}, + ), + snapshot_store=store, + ) + thread_id = "stateless-thread" + + try: + first_events = [ + event + async for event in runner.run({ + "thread_id": thread_id, + "__ag_ui_snapshot_scope": "test", + "messages": [{"role": "user", "content": "first"}], + }) + ] + first_snapshot = next( + event.model_dump(by_alias=True)["messages"] + for event in reversed(first_events) + if getattr(event, "type", None) == "MESSAGES_SNAPSHOT" + ) + with pytest.raises(ChatClientException, match="schema validation"): + _ = [ + event + async for event in runner.run({ + "thread_id": thread_id, + "__ag_ui_snapshot_scope": "test", + "messages": [*first_snapshot, {"role": "user", "content": "second"}], + }) + ] + finally: + await responses_client.close() + + assert all("conversation" not in payload for payload in transport.payloads) + assert all("previous_response_id" not in payload for payload in transport.payloads) + assert [item["role"] for item in transport.payloads[1]["input"]] == ["user", "assistant", "user"] + stored = await store.get(scope="test", thread_id=thread_id) + assert stored is not None + assert stored.session_state is None + + def _sse_event_types(events: list[dict[str, Any]]) -> list[str]: """Extract event type strings from parsed SSE events.""" return [e["event"] for e in events] diff --git a/python/packages/foundry_hosting/tests/test_responses_int.py b/python/packages/foundry_hosting/tests/test_responses_int.py index da9fb5c356b..b157ef978b6 100644 --- a/python/packages/foundry_hosting/tests/test_responses_int.py +++ b/python/packages/foundry_hosting/tests/test_responses_int.py @@ -399,6 +399,49 @@ async def test_input_pdf_file_data(self, server: ResponsesHostServer) -> None: class TestMultiTurn: """Multi-round conversations using previous_response_id.""" + @pytest.mark.flaky + @pytest.mark.integration + @skip_if_foundry_hosting_integration_tests_disabled + async def test_explicit_agent_framework_message_conversation(self) -> None: + """Agent Framework Message and Content primitives preserve explicit conversation history.""" + agent = Agent( + client=FoundryChatClient(credential=AzureCliCredential()), # ty: ignore[invalid-argument-type] + instructions="You are a concise assistant. Keep answers very short.", + ) + + response = await agent.run([ + Message(role="user", contents=[Content.from_text("My favorite color is blue.")]), + Message(role="assistant", contents=[Content.from_text("Your favorite color is blue.")]), + Message(role="user", contents=[Content.from_text("What is my favorite color?")]), + ]) + + assert response.text is not None + assert "blue" in response.text.lower() + + @pytest.mark.flaky + @pytest.mark.integration + @skip_if_foundry_hosting_integration_tests_disabled + async def test_explicit_user_assistant_user_conversation(self, server: ResponsesHostServer) -> None: + """Explicit user, assistant, and follow-up user messages form one conversation.""" + resp = await _post_json( + server, + { + "input": [ + {"type": "message", "role": "user", "content": "My favorite color is blue."}, + {"type": "message", "role": "assistant", "content": "Your favorite color is blue."}, + {"type": "message", "role": "user", "content": "What is my favorite color?"}, + ], + "stream": False, + }, + ) + + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "completed" + output_messages = [item for item in body["output"] if item["type"] == "message"] + assert len(output_messages) == 1 + assert "blue" in output_messages[0]["content"][0]["text"].lower() + @pytest.mark.flaky @pytest.mark.integration @skip_if_foundry_hosting_integration_tests_disabled diff --git a/python/packages/mistral/agent_framework_mistral/_chat_client.py b/python/packages/mistral/agent_framework_mistral/_chat_client.py index bc55704392b..67c13cdbc83 100644 --- a/python/packages/mistral/agent_framework_mistral/_chat_client.py +++ b/python/packages/mistral/agent_framework_mistral/_chat_client.py @@ -20,6 +20,7 @@ ChatResponse, ChatResponseUpdate, Content, + FinishReason, FinishReasonLiteral, FunctionInvocationConfiguration, FunctionInvocationLayer, @@ -187,6 +188,20 @@ class MistralSettings(TypedDict, total=False): "tool_calls": "tool_calls", } + +def _map_finish_reason(reason: Any) -> FinishReason | None: + """Map a Mistral finish reason onto a framework finish reason. + + Finish reasons that are not in ``_FINISH_REASON_MAP`` are passed through unchanged so that + values such as ``error`` still reach the caller instead of being reported as no finish + reason at all. + """ + if not reason: + return None + raw = str(reason) + return FinishReason(_FINISH_REASON_MAP.get(raw, raw)) + + # La Plateforme requires tool call IDs to be exactly 9 alphanumeric characters. _MISTRAL_TOOL_CALL_ID_PATTERN = re.compile(r"^[a-zA-Z0-9]{9}$") @@ -754,9 +769,7 @@ def _parse_response( choice: Mapping[str, Any] = choices[0] if choices else {} message: Mapping[str, Any] = choice.get("message") or {} contents = self._parse_message_contents(message) - finish_reason: FinishReasonLiteral | None = None - if reason := choice.get("finish_reason"): - finish_reason = _FINISH_REASON_MAP.get(str(reason)) + finish_reason = _map_finish_reason(choice.get("finish_reason")) return ChatResponse( response_id=response.get("id"), messages=[Message(role="assistant", contents=contents, raw_representation=choice or None)], @@ -775,7 +788,7 @@ def _parse_chunk(self, chunk: Mapping[str, Any], tool_calls: _StreamedToolCalls) emitted as complete calls when their choice finishes. """ contents: list[Content] = [] - finish_reason: FinishReasonLiteral | None = None + finish_reason: FinishReason | None = None choices = cast("Sequence[Mapping[str, Any]]", chunk.get("choices") or ()) for choice in choices: choice_index = index if isinstance(index := choice.get("index"), int) else 0 @@ -786,7 +799,7 @@ def _parse_chunk(self, chunk: Mapping[str, Any], tool_calls: _StreamedToolCalls) if reason := choice.get("finish_reason"): contents.extend(tool_calls.flush_choice(choice_index)) if finish_reason is None: - finish_reason = _FINISH_REASON_MAP.get(str(reason)) + finish_reason = _map_finish_reason(reason) if usage := self._parse_usage(chunk.get("usage")): contents.append(Content.from_usage(usage_details=usage, raw_representation=chunk)) return ChatResponseUpdate( diff --git a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py index cb4a8c95010..2cda9397acd 100644 --- a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py +++ b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py @@ -692,6 +692,30 @@ async def test_parse_finish_reason_model_length() -> None: assert response.finish_reason == "length" +async def test_parse_finish_reason_unmapped_is_preserved() -> None: + """A Mistral finish reason with no framework equivalent is passed through unchanged.""" + client, _ = make_client(json_response(make_response_payload(content="x", finish_reason="error"))) + + response = await client.get_response([Message("user", ["hi"])]) + + assert response.finish_reason == "error" + + +async def test_parse_finish_reason_absent() -> None: + """A response without a finish reason still reports no finish reason.""" + client, _ = make_client( + json_response( + make_response_payload( + choices=[{"index": 0, "message": {"role": "assistant", "content": "x"}}], + ) + ) + ) + + response = await client.get_response([Message("user", ["hi"])]) + + assert response.finish_reason is None + + async def test_function_invocation_loop() -> None: client, server = make_client( json_response( @@ -755,6 +779,17 @@ async def test_streaming_response() -> None: assert server.last_request["stream"] is True +async def test_streaming_finish_reason_unmapped_is_preserved() -> None: + """A streamed Mistral finish reason with no framework equivalent is passed through unchanged.""" + client, _ = make_client(stream_response(make_chunk_payload(content="partial", finish_reason="error"))) + + stream = client.get_response([Message("user", ["hi"])], stream=True) + updates = [update async for update in stream] + + assert [u.finish_reason for u in updates] == ["error"] + assert (await stream.get_final_response()).finish_reason == "error" + + async def test_streaming_tool_calls() -> None: client, _ = make_client( stream_response( diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index 55ff20950ad..9607d4de7e1 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -1866,11 +1866,16 @@ def _prepare_content_for_openai( if role == "assistant": # Assistant history is represented as output text items; Azure validation # requires `annotations` to be present for this type. - return { + output_text = { "type": "output_text", "text": content.text, "annotations": _annotations_to_output_text(getattr(content, "annotations", None)), } + if "logprobs" in content.additional_properties: + output_text["logprobs"] = self._serialize_provider_payload( + content.additional_properties["logprobs"] + ) + return output_text return _attach_prompt_cache_breakpoint( { "type": "input_text", @@ -2614,8 +2619,14 @@ def _parse_response_from_openai( for message_content in item.content: # type: ignore[reportMissingTypeArgument] match message_content.type: case "output_text": + logprobs = getattr(message_content, "logprobs", None) text_content = Content.from_text( text=message_content.text, + additional_properties=( + {"logprobs": self._serialize_provider_payload(logprobs)} + if logprobs is not None + else None + ), raw_representation=message_content, ) metadata.update(self._get_metadata_from_response(message_content)) diff --git a/python/packages/openai/tests/openai/test_agui_provider_matrix.py b/python/packages/openai/tests/openai/test_agui_provider_matrix.py new file mode 100644 index 00000000000..f2df7c8f14f --- /dev/null +++ b/python/packages/openai/tests/openai/test_agui_provider_matrix.py @@ -0,0 +1,215 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Live AG-UI multi-turn coverage for OpenAI provider continuation modes.""" + +from __future__ import annotations + +import json +import os +from typing import Any, Literal, cast +from uuid import uuid4 + +import pytest +from agent_framework import Agent, AgentSession, Message +from agent_framework_ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore + +from agent_framework_openai import OpenAIChatClient, OpenAIChatCompletionClient + +skip_if_openai_integration_tests_disabled = pytest.mark.skipif( + os.getenv("OPENAI_API_KEY", "") in ("", "test-dummy-key"), + reason="No real OPENAI_API_KEY provided; skipping integration tests.", +) + +_Mode = Literal["stateless", "conversation", "previous_response"] +_PROVIDER_SESSION_KEY = "__ag_ui_provider_service_session_id" + + +class _CapturingAgent: + """Delegate to an Agent without exposing its conversation factory.""" + + def __init__(self, agent: Any) -> None: + self._agent = agent + self.inputs: list[list[Message]] = [] + self.service_session_ids: list[Any] = [] + self.created_conversation_ids: list[str] = [] + + def run(self, messages: Any = None, **kwargs: Any) -> Any: + self.inputs.append(list(messages) if isinstance(messages, list) else [messages]) + session = kwargs.get("session") + self.service_session_ids.append(session.service_session_id if session is not None else None) + return self._agent.run(messages, **kwargs) + + def __getattr__(self, name: str) -> Any: + if name == "create_conversation": + raise AttributeError(name) + return getattr(self._agent, name) + + +class _ConversationCapturingAgent(_CapturingAgent): + """Expose backend conversation creation to AgentFrameworkAgent.""" + + async def create_conversation(self, *, session_id: str | None = None) -> AgentSession: + conversation = await self._agent.client.client.conversations.create() + self.created_conversation_ids.append(conversation.id) + return self._agent.get_session(conversation.id, session_id=session_id) + + +async def _exercise_agui_case(agent: Any, mode: _Mode) -> None: + marker = f"AF-AGUI-{uuid4().hex}" + follow_up = "Return only the exact marker from my previous message." + thread_id = f"agui-thread-{uuid4().hex}" + scope = f"agui-scope-{uuid4().hex}" + snapshot_store = InMemoryAGUIThreadSnapshotStore() + capturing = _ConversationCapturingAgent(agent) if mode == "conversation" else _CapturingAgent(agent) + + runner = AgentFrameworkAgent( + agent=cast(Any, capturing), + use_service_session=mode != "stateless", + service_session_id_from_thread_id=False, + snapshot_store=snapshot_store, + ) + + try: + first_events = [ + event + async for event in runner.run({ + "threadId": thread_id, + "runId": f"run-1-{uuid4().hex}", + "__ag_ui_snapshot_scope": scope, + "messages": [ + { + "role": "user", + "content": f"Remember this exact marker for my next message: {marker}", + } + ], + }) + ] + first_snapshot = next( + event for event in reversed(first_events) if getattr(event, "type", None) == "MESSAGES_SNAPSHOT" + ) + replay_messages = cast(list[dict[str, Any]], first_snapshot.model_dump(by_alias=True)["messages"]) + first_stored = await snapshot_store.get(scope=scope, thread_id=thread_id) + assert first_stored is not None + + second_events = [ + event + async for event in runner.run({ + "threadId": thread_id, + "runId": f"run-2-{uuid4().hex}", + "__ag_ui_snapshot_scope": scope, + "messages": [*replay_messages, {"role": "user", "content": follow_up}], + }) + ] + + assert not [event for event in second_events if getattr(event, "type", None) == "RUN_ERROR"] + assert [[message.role for message in messages] for messages in capturing.inputs[:1]] == [["user"]] + if mode == "stateless": + assert [message.role for message in capturing.inputs[1]] == ["user", "assistant", "user"] + assert capturing.service_session_ids == [None, None] + assert first_stored.session_state is None + else: + assert [(message.role, message.text) for message in capturing.inputs[1]] == [("user", follow_up)] + assert first_stored.session_state is not None + provider_id = first_stored.session_state[_PROVIDER_SESSION_KEY] + assert isinstance(provider_id, str) + assert provider_id != thread_id + assert capturing.service_session_ids[1] == provider_id + assert provider_id not in json.dumps(first_stored.messages) + if mode == "conversation": + assert provider_id.startswith("conv_") + assert capturing.service_session_ids == [provider_id, provider_id] + assert capturing.created_conversation_ids == [provider_id] + else: + assert provider_id.startswith("resp_") + assert capturing.service_session_ids[0] is None + assert not capturing.created_conversation_ids + + assert marker in capturing.inputs[0][0].text + assert capturing.inputs[1][-1].text == follow_up + response_text = "".join( + str(getattr(event, "delta", "")) + for event in second_events + if getattr(event, "type", None) == "TEXT_MESSAGE_CONTENT" + ) + assert marker in response_text + + final_stored = await snapshot_store.get(scope=scope, thread_id=thread_id) + assert final_stored is not None + assert [message.get("role") for message in final_stored.messages].count("user") == 2 + assert [message.get("role") for message in final_stored.messages].count("assistant") >= 2 + finally: + for conversation_id in capturing.created_conversation_ids: + await agent.client.client.conversations.delete(conversation_id) + + +@pytest.mark.flaky +@pytest.mark.integration +@skip_if_openai_integration_tests_disabled +@pytest.mark.parametrize( + ("mode", "store"), + [ + pytest.param("stateless", False, id="stateless-snapshot-replay"), + pytest.param("conversation", True, id="conversation"), + pytest.param("previous_response", True, id="previous-response-id"), + ], +) +async def test_openai_responses_agui_provider_matrix(mode: _Mode, store: bool) -> None: + client = OpenAIChatClient() + agent = Agent(client=cast(Any, client), default_options=cast(Any, {"store": store})) + try: + await _exercise_agui_case(agent, mode) + finally: + await client.client.close() + + +@pytest.mark.flaky +@pytest.mark.integration +@skip_if_openai_integration_tests_disabled +@pytest.mark.parametrize("store", [pytest.param(False, id="store-false"), pytest.param(True, id="store-true")]) +async def test_openai_chat_completions_agui_provider_matrix(store: bool) -> None: + client = OpenAIChatCompletionClient() + agent = Agent(client=cast(Any, client), default_options=cast(Any, {"store": store})) + try: + await _exercise_agui_case(agent, "stateless") + finally: + await client.client.close() + + +@pytest.mark.flaky +@pytest.mark.integration +@skip_if_openai_integration_tests_disabled +async def test_openai_responses_replays_real_assistant_logprobs() -> None: + """Real provider logprobs survive direct assistant-message replay without fabrication.""" + client = OpenAIChatClient() + follow_up = Message(role="user", contents=["Reply with exactly: done"]) + try: + first = await client.get_response( + [Message(role="user", contents=["Reply with exactly: hello"])], + options={ + "store": False, + "include": ["message.output_text.logprobs"], + "top_logprobs": 2, + }, + ) + assistant_content = next( + content + for message in first.messages + for content in message.contents + if message.role == "assistant" and content.type == "text" + ) + real_logprobs = assistant_content.additional_properties.get("logprobs") + assert isinstance(real_logprobs, list) + assert real_logprobs + + _, run_options, _ = await client._prepare_request( + [*first.messages, follow_up], + {"store": False}, + ) + assistant_input = next(item for item in run_options["input"] if item.get("role") == "assistant") + replayed_logprobs = assistant_input["content"][0]["logprobs"] + + assert replayed_logprobs == real_logprobs + second = await client.get_response([*first.messages, follow_up], options={"store": False}) + assert "done" in second.text.lower() + finally: + await client.client.close() diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index aab0fc1fe24..da52a9b6db0 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -2881,10 +2881,60 @@ def test_prepare_content_for_openai_text_uses_role_specific_type() -> None: assert user_result["type"] == "input_text" assert assistant_result["type"] == "output_text" assert assistant_result["annotations"] == [] + assert "logprobs" not in assistant_result assert user_result["text"] == "hello" assert assistant_result["text"] == "hello" +def test_prepare_content_for_openai_replays_real_assistant_logprobs() -> None: + """Assistant history replays provider logprobs only when the response supplied them.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + logprobs = [ + { + "token": "hello", + "bytes": [104, 101, 108, 108, 111], + "logprob": -0.1, + "top_logprobs": [], + } + ] + text_content = Content.from_text(text="hello", additional_properties={"logprobs": logprobs}) + + result = client._prepare_content_for_openai("assistant", text_content) + + assert result["logprobs"] == logprobs + + +def test_parse_and_replay_preserves_real_assistant_logprobs() -> None: + """Parsed Responses logprobs remain attached to assistant content for direct replay.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + logprobs = [ + { + "token": "hello", + "bytes": [104, 101, 108, 108, 111], + "logprob": -0.1, + "top_logprobs": [], + } + ] + output_text = MagicMock(type="output_text", text="hello", annotations=[], logprobs=logprobs) + output_message = MagicMock(type="message", content=[output_text]) + response = MagicMock( + output_parsed=None, + output=[output_message], + metadata={}, + usage=None, + id="resp-test", + created_at=1_000_000_000, + model="test-model", + ) + + parsed = client._parse_response_from_openai(response, options={"store": False}) + text_content = parsed.messages[0].contents[0] + replayed = client._prepare_content_for_openai("assistant", text_content) + + assert text_content.additional_properties["logprobs"] == logprobs + assert replayed["logprobs"] == logprobs + + def test_prepare_messages_for_openai_assistant_history_uses_output_text_with_annotations() -> None: """Assistant history should be output_text and include required annotations.""" client = OpenAIChatClient(model="test-model", api_key="test-key") diff --git a/python/samples/05-end-to-end/chatkit-integration/app.py b/python/samples/05-end-to-end/chatkit-integration/app.py index 612d9f9c443..97097cc3945 100644 --- a/python/samples/05-end-to-end/chatkit-integration/app.py +++ b/python/samples/05-end-to-end/chatkit-integration/app.py @@ -43,7 +43,7 @@ # ChatKit imports from chatkit.actions import Action from chatkit.server import ChatKitServer -from chatkit.store import StoreItemType, default_generate_id +from chatkit.store import NotFoundError, StoreItemType, default_generate_id from chatkit.types import ( ThreadItem, ThreadItemDoneEvent, @@ -595,20 +595,25 @@ async def upload_file(attachment_id: str, file: UploadFile = File(...)): # noqa logger.warning(f"Rejected invalid attachment ID: {attachment_id!r}") return JSONResponse(status_code=400, content={"error": "Invalid attachment ID."}) + try: + attachment = await data_store.load_attachment(attachment_id, {"user_id": DEFAULT_USER_ID}) + except NotFoundError: + return JSONResponse(status_code=404, content={"error": "Attachment not found."}) + + if attachment.upload_descriptor is None: + return JSONResponse(status_code=409, content={"error": "Attachment upload is already complete."}) + try: # Read file contents contents = await file.read() # Save to disk - file_path.write_bytes(contents) + file_path.write_bytes(contents) # CodeQL [SM01305] Path is constrained by get_file_path. logger.info(f"Saved {len(contents)} bytes to {file_path}") - # Load the attachment metadata from the data store - attachment = await data_store.load_attachment(attachment_id, {"user_id": DEFAULT_USER_ID}) - - # Clear the upload_url since upload is complete - attachment.upload_url = None + # Clear the upload descriptor since upload is complete + attachment.upload_descriptor = None # Save the updated attachment back to the store await data_store.save_attachment(attachment, {"user_id": DEFAULT_USER_ID}) @@ -637,19 +642,15 @@ async def preview_image(attachment_id: str): return JSONResponse(status_code=400, content={"error": "Invalid attachment ID."}) try: - if not file_path.exists(): - return JSONResponse(status_code=404, content={"error": "File not found"}) + attachment = await data_store.load_attachment(attachment_id, {"user_id": DEFAULT_USER_ID}) + except NotFoundError: + return JSONResponse(status_code=404, content={"error": "Attachment not found."}) - # Determine media type from file extension or attachment metadata - # For simplicity, we'll try to load from the store - try: - attachment = await data_store.load_attachment(attachment_id, {"user_id": DEFAULT_USER_ID}) - media_type = attachment.mime_type - except Exception: - # Default to binary if we can't determine - media_type = "application/octet-stream" + try: + if not file_path.exists(): # CodeQL [SM01305] Path is constrained by get_file_path. + return JSONResponse(status_code=404, content={"error": "File not found"}) - return FileResponse(file_path, media_type=media_type) + return FileResponse(file_path, media_type=attachment.mime_type) # CodeQL [SM01305] Path is constrained by get_file_path. # fmt: skip except Exception as e: logger.error(f"Error serving preview for attachment {attachment_id}: {e}", exc_info=True) diff --git a/python/samples/05-end-to-end/chatkit-integration/attachment_store.py b/python/samples/05-end-to-end/chatkit-integration/attachment_store.py index b08ae9c43a4..c4de028b29a 100644 --- a/python/samples/05-end-to-end/chatkit-integration/attachment_store.py +++ b/python/samples/05-end-to-end/chatkit-integration/attachment_store.py @@ -11,7 +11,13 @@ from typing import TYPE_CHECKING, Any from chatkit.store import AttachmentStore -from chatkit.types import Attachment, AttachmentCreateParams, FileAttachment, ImageAttachment +from chatkit.types import ( + Attachment, + AttachmentCreateParams, + AttachmentUploadDescriptor, + FileAttachment, + ImageAttachment, +) from pydantic import AnyUrl if TYPE_CHECKING: @@ -70,7 +76,7 @@ def get_file_path(self, attachment_id: str) -> Path: if not attachment_id or attachment_id in {".", ".."} or "/" in attachment_id or "\\" in attachment_id: raise ValueError(f"Invalid attachment ID: {attachment_id!r}") - file_path = (self.uploads_dir / attachment_id).resolve() + file_path = (self.uploads_dir / attachment_id).resolve() # CodeQL [SM01305] Path containment is validated below. # fmt: skip if not file_path.is_relative_to(self.uploads_dir) or file_path.parent != self.uploads_dir: raise ValueError(f"Invalid attachment ID: {attachment_id!r}") return file_path @@ -90,8 +96,11 @@ async def create_attachment(self, input: AttachmentCreateParams, context: dict[s # Generate unique ID for this attachment attachment_id = self.generate_attachment_id(input.mime_type, context) - # Generate upload URL that points to our FastAPI upload endpoint - upload_url = f"{self.base_url}/upload/{attachment_id}" + # Generate upload instructions that point to our FastAPI upload endpoint + upload_descriptor = AttachmentUploadDescriptor( + url=AnyUrl(f"{self.base_url}/upload/{attachment_id}"), + method="POST", + ) # Create appropriate attachment type based on MIME type if input.mime_type.startswith("image/"): @@ -103,7 +112,7 @@ async def create_attachment(self, input: AttachmentCreateParams, context: dict[s type="image", mime_type=input.mime_type, name=input.name, - upload_url=AnyUrl(upload_url), + upload_descriptor=upload_descriptor, preview_url=AnyUrl(preview_url), ) else: @@ -113,7 +122,7 @@ async def create_attachment(self, input: AttachmentCreateParams, context: dict[s type="file", mime_type=input.mime_type, name=input.name, - upload_url=AnyUrl(upload_url), + upload_descriptor=upload_descriptor, ) # Save attachment metadata to data store so it's available during upload