diff --git a/.azure-pipelines/hosted-agent-tests/python/agent-framework/responses/25-custom-store/test-spec.yml b/.azure-pipelines/hosted-agent-tests/python/agent-framework/responses/25-custom-store/test-spec.yml new file mode 100644 index 000000000..c2095da1c --- /dev/null +++ b/.azure-pipelines/hosted-agent-tests/python/agent-framework/responses/25-custom-store/test-spec.yml @@ -0,0 +1,65 @@ +version: 1 + +sample: + owner: sapsing + experiences: [azd, vscode] + +tests: + - name: validates-response-and-session-stores + turns: + - input: Remember the exact marker INC-2301 for my next message. Reply briefly. + - input: | + What exact marker did I ask you to remember? Then call + inspect_custom_store exactly once and include its result verbatim. + assertions: + - source: assistant_text + type: contains + value: INC-2301 + + - source: assistant_text + type: regex + value: 'agent_session_store=[1-9][0-9]*' + + - name: validates-checkpoint-store + turns: + - input: | + Call validate_checkpoint_store exactly once with marker CHK-2301. + Include the tool result verbatim. + assertions: + - source: assistant_text + type: contains + value: CHECKPOINT_STORE_OK CHK-2301 + + - name: validates-function-approval-store + turns: + - input: | + Call approve_incident_change exactly once with change_id CHG-2301. + Include the tool result verbatim after approval. + approvals: + mcp: + - server_label: agent_framework + name: approve_incident_change + assertions: + - source: assistant_text + type: contains + value: APPROVAL_STORE_OK CHG-2301 + + - name: confirms-all-four-stores-contain-data + turns: + - input: Call inspect_custom_store exactly once and include its result verbatim. + assertions: + - source: assistant_text + type: regex + value: 'response_store=[1-9][0-9]*' + + - source: assistant_text + type: regex + value: 'agent_session_store=[1-9][0-9]*' + + - source: assistant_text + type: regex + value: 'checkpoint_store=[1-9][0-9]*' + + - source: assistant_text + type: regex + value: 'function_approval_store=[1-9][0-9]*' diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/AGENTS.md b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/AGENTS.md new file mode 100644 index 000000000..360362376 --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/AGENTS.md @@ -0,0 +1,60 @@ +--- +description: Coding agent guidance for the custom store hosted-agent sample +--- + +# Coding Agent Instructions + +This sample is a Microsoft Foundry hosted Agent Framework agent served through +the Responses protocol. It replaces all four hosting storage extension points +with custom Redis-backed implementations: the Responses store, the agent +session store, the workflow checkpoint store, and the function approval store. + +## Key files + +* `azure.yaml` defines the Foundry project, model, agent, runtime configuration, + deployed environment variables, and the layered infrastructure (Foundry layer + plus a Bicep layer that provisions Azure Managed Redis). +* `docker-compose.yml` runs a local Redis for `python main.py` development. +* `infra/redis/main.bicep` provisions Azure Managed Redis and exports its `REDIS_HOST` and `REDIS_PORT` outputs consumed by `azure.yaml`; the `hooks/` postdeploy scripts grant the deployed agent identity Redis data access. +* `src/agent-framework-agent-custom-store-responses/main.py` creates the agent + and supplies the response store plus the three store providers to + `ResponsesHostServer`. +* `src/agent-framework-agent-custom-store-responses/redis_database.py` owns the + shared async Redis client, key namespacing, and per-user partitioning helper. +* `src/agent-framework-agent-custom-store-responses/redis_response_store.py` + implements `ResponseProviderProtocol`. +* `src/agent-framework-agent-custom-store-responses/redis_state_stores.py` + implements the session, checkpoint, and function approval stores and their + providers. +* `README.md` documents local and hosted execution paths. + +Preserve per-user partitioning, namespaced keys, and native key TTLs when +changing the stores. All four stores share one Redis instance and each owns its +own key prefix. + +Do not commit `.vscode` launch or task files generated by Foundry Toolkit. Keep +the service virtual environment beside `requirements.txt` under the service +directory. + +## Development workflow + +Use `azd` from an initialized agent project: + +``` + +azd provision +azd ai agent run +azd ai agent invoke --local "your message" +azd deploy +azd ai agent invoke "your message" + +``` + +This project was built with the `microsoft-foundry` skill. Before working on or +answering questions about Foundry agents, read the `vscode-microsoft-foundry` +skill first when using VS Code, then read the `microsoft-foundry` skill. + +## References + +* [Hosted agents overview](https://learn.microsoft.com/azure/foundry/agents/concepts/hosted-agents) + diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/CLAUDE.md b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/CLAUDE.md new file mode 100644 index 000000000..89066b045 --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/CLAUDE.md @@ -0,0 +1,10 @@ +--- +description: Claude Code entry point for the sample coding agent instructions +--- + +# Claude Code Instructions + +This project uses [AGENTS.md](./AGENTS.md) as the source of truth for coding +agent instructions. + +@AGENTS.md diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/README.md b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/README.md new file mode 100644 index 000000000..a585ac9ee --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/README.md @@ -0,0 +1,341 @@ +--- +description: Replace all four Foundry hosting storage extension points with custom Redis implementations +ms.date: 2026-09-11 +ms.topic: sample +--- + +# Bring Your Own Store (Responses Protocol) + +`ResponsesHostServer` persists agent state through four storage extension +points, each backed by a Foundry native store by default: + +| Extension point | Purpose | +|-----------------------------------|------------------------------------------------| +| `store` | Responses transcript and history | +| `agent_session_store_provider` | Agent Framework agent session snapshots | +| `checkpoint_store_provider` | Workflow checkpoints, scoped to one context | +| `function_approval_store_provider`| Human-in-the-loop function approval requests | + +This sample replaces all four defaults with stores that you own and manage. It +uses [Redis](https://redis.io/) as a single, networked backend so the extension +points and their contracts stay in focus while the state survives restarts and +can be shared across multiple agent instances. You can use the same extension +points to integrate another storage backend that meets your data ownership, +retention, compliance, or application integration requirements. + +Local development runs against a [Docker](https://www.docker.com/) Redis +container. Deployment provisions [Azure Managed Redis](https://learn.microsoft.com/azure/redis/) +through a layered Bicep stage. An `azd` post-deployment hook adds the deployed +agent identity to the Redis access policy, enabling passwordless authentication. + +## What this sample demonstrates + +The sample implements each hosting storage contract against Redis: + +* `RedisResponseStore` implements `ResponseProviderProtocol` and persists + response envelopes, transcript items, previous-response history chains, and + conversation indexes with a configurable retention TTL. +* `RedisAgentSessionStore` implements `SessionStore` for agent session + snapshots. +* `RedisCheckpointStore` implements `CheckpointStorage` for workflow + checkpoints, scoped to one workflow context. +* `RedisFunctionApprovalStore` implements `FunctionApprovalStore` for function + approval requests. + +All four stores share one Redis instance, each owns its own key prefix, and +partitions every key by a hash of the Foundry-provided user ID so one caller +cannot read another caller's state. + +## How it works + +[main.py](src/agent-framework-agent-custom-store-responses/main.py) creates one +shared Redis client and passes the response store plus the three store providers +to `ResponsesHostServer`: + +```python +database = create_database() + +server = ResponsesHostServer( + agent, + store=create_response_store(database), + agent_session_store_provider=RedisAgentSessionStoreProvider(database), + checkpoint_store_provider=RedisCheckpointStoreProvider(database), + function_approval_store_provider=RedisFunctionApprovalStoreProvider(database), +) +``` + +The `store` argument takes a store instance directly. The three providers return +a per-request store so each request is partitioned by its own user identity. The +checkpoint provider is context-scoped: it also receives the workflow context ID +so checkpoints for one workflow run can be listed and pruned as a collection. + +The storage implementation uses +[redis_database.py](src/agent-framework-agent-custom-store-responses/redis_database.py) +to manage one shared asynchronous Redis client, apply a common prefix to every +key, and serialize conflicting read-modify-write operations across agent +instances with resource-scoped Redis locks. During local development, the client +connects to the Docker Redis instance without authentication. When deployed, it +connects to Azure Managed Redis over TLS and uses `DefaultAzureCredential` to +authenticate with the agent's managed identity. + +[redis_response_store.py](src/agent-framework-agent-custom-store-responses/redis_response_store.py) +implements the Responses storage contract, while +[redis_state_stores.py](src/agent-framework-agent-custom-store-responses/redis_state_stores.py) +implements the agent session, workflow checkpoint, and function approval +contracts. Response envelopes and transcript items use native Redis key TTLs, +allowing Redis to remove expired state without a separate cleanup process. + +The agent keeps `default_options={"store": False}` because the Responses host +supplies the complete transcript to the model for every turn. This model-client +option is separate from the custom stores configured on the host. + +Three local tools provide deterministic validation scenarios: + +* `inspect_custom_store` reports per-user key counts for all four stores. +* `validate_checkpoint_store` runs a small Agent Framework workflow, persists + its checkpoints through `RedisCheckpointStoreProvider`, and loads the last + checkpoint before returning success. +* `approve_incident_change` requires human approval. The host persists the + approval request before resuming the tool call from the approval response. + +> [!NOTE] +> Redis stores state on a networked service that is shared across every agent +> instance, so conversation continuity and approvals survive restarts and scale +> horizontally. The local Docker Redis is ephemeral by design; the deployed +> Azure Managed Redis is durable. + +## Prerequisites + +1. Python 3.11 or later. The `azd ai agent run` path uses Python 3.13. +2. A Foundry project with a deployed chat model, or permission to create both + through `azd provision`. +3. The Foundry User role on the Foundry project for the local or hosted identity. +4. Permission to assign an Azure Managed Redis data access policy during the + `azd deploy` post-deployment hook, so the deployed agent identity can reach + Redis without a password or access key. +5. [Docker](https://www.docker.com/) for the local Redis container. Deployment + provisions Azure Managed Redis instead and requires no local Redis. + +The `redis` and `redis-entraid` clients are installed from `requirements.txt`. + +## Option 1: Azure Developer CLI + +### Install prerequisites + +1. Install the [Azure Developer CLI](https://learn.microsoft.com/azure/developer/azure-developer-cli/install-azd). +2. Install the Foundry extension: + + ```bash + azd ext install microsoft.foundry + ``` + +3. Authenticate: + + ```bash + azd auth login + ``` + +### Initialize the agent project + +No cloning required. Create a new folder and initialize from the manifest: + +```bash +mkdir my-custom-store-agent && cd my-custom-store-agent +azd ai agent init . -m https://github.com/microsoft-foundry/foundry-samples/blob/main/samples/python/hosted-agents/agent-framework/responses/25-custom-store/azure.yaml +``` + +Follow the prompts to select or create a Foundry project and model deployment. + +### Provision Foundry resources + +Provision the resources declared in [azure.yaml](azure.yaml) before the first +deployment, even if the selected project and model deployment already exist. +This runs the `redis` Bicep layer, which provisions Azure Managed Redis: + +```bash +azd provision +``` + +### Run locally + +Start a local Redis with Docker from the sample root: + +```bash +docker compose up -d +``` + +Change to the generated service directory, which contains `main.py`. Copy +`.env.example` to `.env`, set the Foundry endpoint and model deployment, and +point `REDIS_URL` at the local Redis for direct `python main.py` runs: + +```dotenv +FOUNDRY_PROJECT_ENDPOINT="https://.services.ai.azure.com/api/projects/" +AZURE_AI_MODEL_DEPLOYMENT_NAME="" + +REDIS_URL=redis://localhost:6379/0 +LOCAL_STORE_USER_ID=local-developer +REDIS_RESPONSE_TTL_SECONDS=86400 +``` + +To run through `azd`, run the agent from the initialized project root: + +```bash +azd ai agent run +``` + +The Responses endpoint listens on `http://localhost:8088/responses`. Keys are +written to Redis under the `custom-store:` prefix on the first request. + +### Verify conversation continuity + +Send an initial request, then continue from its response ID: + +```powershell +$first = Invoke-RestMethod -Uri http://localhost:8088/responses -Method Post -ContentType "application/json" -Body '{"input":"Remember that incident INC-2048 affects the dispatch API."}' +$body = @{ input = "Which API is affected?"; previous_response_id = $first.id } | ConvertTo-Json +Invoke-RestMethod -Uri http://localhost:8088/responses -Method Post -ContentType "application/json" -Body $body +``` + +The second response identifies the dispatch API. Stop and restart the agent host, +then send the follow-up with the same response ID. The context remains available +in Redis until its TTL expires. + +### Validate all four custom stores + +Use these prompts in Agent Inspector or with `azd ai agent invoke`. Run them in +order so the final prompt can confirm that every store contains data: + +1. `Remember marker INC-2301 for my next message.` +2. `Recall my marker, then call inspect_custom_store and include its result verbatim.` +3. `Call validate_checkpoint_store with marker CHK-2301 and include its result verbatim.` +4. `Call approve_incident_change with change_id CHG-2301 and include its result verbatim after approval.` +5. `Call inspect_custom_store and include its result verbatim.` + +The first two prompts exercise agent-session persistence and confirm that the +Responses store contains response records. The third runs a real checkpointed +workflow against the checkpoint provider. The fourth creates and reloads an +approval request through the Responses continuation flow. Approve that request +when Agent Inspector asks. The final result has this shape, with positive counts +for every provider: + +```text +STORE_COUNTS response_store=4 agent_session_store=4 checkpoint_store=2 function_approval_store=1 +``` + +Counts can differ because a workflow may create more than one checkpoint and +the host writes response and session records for each turn. + +### Inspect the Redis data + +Connect to the local Redis with `redis-cli` to list the stored keys and read a +value: + +```bash +docker exec -it sh #or /bin/bash +redis-cli --scan --pattern "custom-store:*" # Run this command from inside the container +``` + +Each store owns a distinct key prefix: `custom-store:resp:` for responses, +`custom-store:session:` for agent sessions, `custom-store:ckpt:` for workflow +checkpoints, and `custom-store:approval:` for function approvals. + +### Deploy + +Provision Azure Managed Redis, then deploy the agent with its store +configuration: + +```bash +azd provision +azd deploy +``` + +`azd deploy` deploys the agent but does not run the infrastructure layers. After +the initial provisioning, use `azd deploy` by itself for code-only updates. The +`redis` Bicep layer enables public network access, disables access-key +authentication, and exports the Azure Managed Redis host and port as the +`REDIS_HOST` and `REDIS_PORT` outputs, which `azure.yaml` injects into the +deployed agent. Direct local runs continue to use the Docker Redis from `.env`. +After deployment, the `postdeploy` hook assigns the built-in Redis `default` +access policy to the deployed agent identity. +Invoke the deployed agent twice to verify continuation: + +```bash +azd ai agent invoke "Remember that incident INC-2048 affects the dispatch API." +azd ai agent invoke "Which API is affected?" +``` + +> [!NOTE] +> The deployed agent authenticates to Azure Managed Redis without a password or +> access key. After deployment, the `postdeploy` hook assigns the built-in +> `default` Redis access policy to the deployed agent identity. At runtime, +> `DefaultAzureCredential` and `redis-entraid` use that identity to obtain a +> Redis data-plane token. The agent configuration contains only the Redis host +> and port, not a connection secret. This sample deploys Azure Managed Redis +> with public network access enabled for demonstration purposes only. For +> production workloads, configure virtual network integration and private +> endpoints to restrict network access. + +## Option 2: VS Code + +### Install VS Code prerequisites + +1. Install VS Code and the + [Foundry Toolkit](https://marketplace.visualstudio.com/items?itemName=ms-windows-ai-studio.windows-ai-studio) + extension. +2. Install the + [Python extension](https://marketplace.visualstudio.com/items?itemName=ms-python.python). + +### Set up the Python environment + +Open the service directory in VS Code. Create a `.venv`, then install the +dependencies: + +```bash +pip install uv +uv pip install -r requirements.txt +``` + +Copy `.env.example` to `.env` and configure the Foundry endpoint, model +deployment, and store settings. Start a local Redis before running the agent: + +```bash +docker compose up -d +``` + +### Run and debug the agent + +Press **F5** to start the agent and open Agent Inspector. You can also run +`python main.py`, open **Foundry Toolkit: Open Agent Inspector** from the Command +Palette, and connect to port `8088`. + +To deploy, run **Foundry Toolkit: Deploy Hosted Agent** from the Command Palette +and select the project described by [azure.yaml](azure.yaml). + +## Troubleshooting + +### A previous response is not found + +Confirm that the follow-up uses the exact `id` returned by the first response, +that both requests use the same hosted user identity, and that the TTL has not +expired. Read operations do not extend retention. + +### The agent cannot connect to Redis + +For local runs, confirm `REDIS_URL` is set, start the Docker Redis with +`docker compose up -d`, and verify it is reachable with +`docker compose exec redis redis-cli ping`. For deployed runs, confirm +`azd provision` completed the `redis` layer and that `REDIS_HOST` is present in +the azd environment. Authentication failures usually mean the deployed agent +identity is missing its Redis access policy assignment; re-run `azd deploy` to +rerun the `postdeploy` hook. + +### State is missing after the local container restarts + +The local Docker Redis in `docker-compose.yml` keeps data in memory only, so +stopping the container clears all state. The deployed Azure Managed Redis is +durable and shared across agent instances. + +## Next steps + +* [Quickstart: Create a hosted agent](https://learn.microsoft.com/azure/foundry/agents/quickstarts/quickstart-hosted-agent) +* [Deploy a hosted agent](https://learn.microsoft.com/azure/foundry/agents/how-to/deploy-hosted-agent) diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/azure.yaml b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/azure.yaml new file mode 100644 index 000000000..c63bce3d8 --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/azure.yaml @@ -0,0 +1,74 @@ +# yaml-language-server: $schema=https://raw.githubusercontent.com/Azure/azure-dev/main/schemas/v1.0/azure.yaml.json + +requiredVersions: + azd: '>=1.27.1' + extensions: + azure.ai.agents: '>=1.0.0-beta.9' +name: agent-framework-agent-custom-store-responses +services: + ai-project: + host: azure.ai.project + deployments: + - name: gpt-5.4-mini + model: + format: OpenAI + name: gpt-5.4-mini + version: '2026-03-17' + sku: + name: GlobalStandard + capacity: 10 + agent-framework-agent-custom-store-responses: + host: azure.ai.agent + metadata: + tags: + - Agent Framework + - AI Agent Hosting + - Azure AI AgentServer + - Responses Protocol + - Custom Store + - Redis + - Azure Managed Redis + project: src/agent-framework-agent-custom-store-responses + language: python + codeConfiguration: + runtime: python_3_13 + entryPoint: main.py + uses: + - ai-project + kind: hosted + name: agent-framework-agent-custom-store-responses + displayName: Custom Store Agent + description: An Agent Framework Responses agent that replaces all four hosting storage extension points (responses, agent sessions, workflow checkpoints, and function approvals) with Redis-backed implementations that isolate data per user. Local development uses a Docker Redis; deployment provisions Azure Managed Redis through a layered Bicep stage. + protocols: + - protocol: responses + version: 2.0.0 + env: + AZURE_AI_MODEL_DEPLOYMENT_NAME: ${AZURE_AI_MODEL_DEPLOYMENT_NAME} + # The redis layer provisions Azure Managed Redis and exports its host and + # port. The postdeploy hook grants the hosted agent identity data access + # for passwordless authentication through DefaultAzureCredential. + REDIS_HOST: ${REDIS_HOST} + REDIS_PORT: ${REDIS_PORT} + REDIS_RESPONSE_TTL_SECONDS: ${REDIS_RESPONSE_TTL_SECONDS=86400} + container: + resources: + cpu: '0.5' + memory: 1Gi +hooks: + postdeploy: + windows: + shell: pwsh + run: hooks/Set-RedisAgentAccess.ps1 + posix: + shell: sh + run: hooks/postdeploy.sh +infra: + layers: + - provider: microsoft.foundry + path: infra/foundry + name: foundry + - provider: bicep + path: infra/redis + name: redis + dependsOn: + - foundry diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/docker-compose.yml b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/docker-compose.yml new file mode 100644 index 000000000..dfc20a089 --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/docker-compose.yml @@ -0,0 +1,11 @@ +# Local Redis for developing the custom-store agent with `python main.py`. +# Start with: docker compose up -d +# The agent reads REDIS_URL=redis://localhost:6379/0 from .env by default. +name: agent-framework-agent-custom-store-responses +services: + redis: + image: redis:7-alpine + ports: + - "6379:6379" + # Keep data in memory only; this is a throwaway local development store. + command: ["redis-server", "--save", "", "--appendonly", "no"] diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/hooks/Set-RedisAgentAccess.ps1 b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/hooks/Set-RedisAgentAccess.ps1 new file mode 100644 index 000000000..af6912254 --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/hooks/Set-RedisAgentAccess.ps1 @@ -0,0 +1,111 @@ +#!/usr/bin/env pwsh +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: MIT +#Requires -Version 7.0 + +<# +.SYNOPSIS + Grants the deployed hosted agent identity access to Azure Managed Redis. +.DESCRIPTION + Gets the deployed agent identity, then adds that identity to Redis + authentication with the default access policy. +.EXAMPLE + ./hooks/Set-RedisAgentAccess.ps1 +.NOTES + Runs automatically as the azd postdeploy hook. +#> + +[CmdletBinding()] +param() + +$ErrorActionPreference = 'Stop' + +function Set-RedisAgentAccess { + [CmdletBinding()] + [OutputType([void])] + param() + + $AzdValues = @{} + azd env get-values | ForEach-Object { + if ($_ -match '^([^=]+)="(.*)"$') { + $AzdValues[$Matches[1]] = $Matches[2] + } + } + if ($LASTEXITCODE -ne 0) { + throw 'Unable to read the selected azd environment.' + } + + $ProjectEndpoint = $AzdValues['AZURE_AI_PROJECT_ENDPOINT'] ?? + $AzdValues['FOUNDRY_PROJECT_ENDPOINT'] + $AgentName = $AzdValues.Keys | + Where-Object { $_ -match '^AGENT_.*_NAME$' } | + ForEach-Object { $AzdValues[$_] } | + Select-Object -First 1 + $RedisDatabaseId = $AzdValues['REDIS_DATABASE_RESOURCE_ID'] + + if ([string]::IsNullOrWhiteSpace($ProjectEndpoint) -or + [string]::IsNullOrWhiteSpace($AgentName) -or + [string]::IsNullOrWhiteSpace($RedisDatabaseId)) { + throw 'Required project, agent, or Redis deployment values are missing.' + } + + # 1. Get the deployed agent identity. + $FoundryToken = az account get-access-token ` + --resource 'https://ai.azure.com' ` + --query accessToken ` + --output tsv + if ($LASTEXITCODE -ne 0) { + throw 'Unable to acquire a Foundry data-plane token.' + } + + $AgentUri = "$($ProjectEndpoint.TrimEnd('/'))/agents/$AgentName`?api-version=v1" + $Agent = Invoke-RestMethod ` + -Method Get ` + -Uri $AgentUri ` + -Headers @{ Authorization = "Bearer $FoundryToken" } + $PrincipalId = $Agent.instance_identity.principal_id + if ([string]::IsNullOrWhiteSpace($PrincipalId)) { + throw 'The deployed agent identity has no principal ID.' + } + + # 2. Add the identity to Redis authentication. + $ManagementToken = az account get-access-token ` + --resource 'https://management.azure.com/' ` + --query accessToken ` + --output tsv + if ($LASTEXITCODE -ne 0) { + throw 'Unable to acquire an Azure management token.' + } + + $AssignmentName = "agent-$PrincipalId" + $AssignmentUri = 'https://management.azure.com{0}/accessPolicyAssignments/{1}?api-version=2025-07-01' -f $RedisDatabaseId, $AssignmentName + $Body = @{ + properties = @{ + accessPolicyName = 'default' + user = @{ objectId = $PrincipalId } + } + } | ConvertTo-Json -Depth 4 -Compress + + Invoke-RestMethod ` + -Method Put ` + -Uri $AssignmentUri ` + -MaximumRetryCount 5 ` + -RetryIntervalSec 5 ` + -Headers @{ Authorization = "Bearer $ManagementToken" } ` + -ContentType 'application/json' ` + -Body $Body | Out-Null + + Write-Host "Granted Azure Managed Redis access to agent identity $PrincipalId." ` + -ForegroundColor Green +} + +if ($MyInvocation.InvocationName -ne '.') { + try { + Set-RedisAgentAccess + exit 0 + } + catch { + Write-Error -ErrorAction Continue "Redis agent access setup failed: $($_.Exception.Message)" + exit 1 + } +} diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/hooks/postdeploy.sh b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/hooks/postdeploy.sh new file mode 100755 index 000000000..2fe0db0a0 --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/hooks/postdeploy.sh @@ -0,0 +1,58 @@ +#!/bin/sh +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: MIT +# +# Grants the deployed hosted agent identity access to Azure Managed Redis. + +set -eu + +main() { + azd_values="$(azd env get-values)" + project_endpoint="$(printf '%s\n' "${azd_values}" | + sed -n 's/^\(AZURE_AI_PROJECT_ENDPOINT\|FOUNDRY_PROJECT_ENDPOINT\)="\(.*\)"$/\2/p' | + head -n 1)" + agent_name="$(printf '%s\n' "${azd_values}" | + sed -n 's/^AGENT_.*_NAME="\(.*\)"$/\1/p' | + head -n 1)" + redis_database_id="$(printf '%s\n' "${azd_values}" | + sed -n 's/^REDIS_DATABASE_RESOURCE_ID="\(.*\)"$/\1/p')" + + if [ -z "${project_endpoint}" ] || [ -z "${agent_name}" ] || + [ -z "${redis_database_id}" ]; then + echo "ERROR: Required project, agent, or Redis deployment values are missing." >&2 + exit 1 + fi + + # 1. Get the deployed agent identity. + foundry_token="$(az account get-access-token \ + --resource 'https://ai.azure.com' \ + --query accessToken \ + --output tsv)" + principal_id="$(curl --fail --silent --show-error \ + --header "Authorization: Bearer ${foundry_token}" \ + "${project_endpoint%/}/agents/${agent_name}?api-version=v1" | + python3 -c \ + "import json, sys; print(json.load(sys.stdin)['instance_identity']['principal_id'])")" + + # 2. Add the identity to Redis authentication. + management_token="$(az account get-access-token \ + --resource 'https://management.azure.com/' \ + --query accessToken \ + --output tsv)" + assignment_name="agent-${principal_id}" + assignment_uri="https://management.azure.com${redis_database_id}/accessPolicyAssignments/${assignment_name}?api-version=2025-07-01" + + curl --fail --silent --show-error \ + --retry 5 \ + --retry-delay 5 \ + --retry-max-time 120 \ + --request PUT \ + --header "Authorization: Bearer ${management_token}" \ + --header 'Content-Type: application/json' \ + --data "{\"properties\":{\"accessPolicyName\":\"default\",\"user\":{\"objectId\":\"${principal_id}\"}}}" \ + "${assignment_uri}" >/dev/null + + echo "Granted Azure Managed Redis access to agent identity ${principal_id}." +} + +main "$@" diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/infra/foundry/README.md b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/infra/foundry/README.md new file mode 100644 index 000000000..d37179160 --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/infra/foundry/README.md @@ -0,0 +1,3 @@ +# Foundry infrastructure layer + +The `microsoft.foundry` provider provisions the Foundry project and agent for this layer. No Bicep template is required. diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/infra/redis/main.bicep b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/infra/redis/main.bicep new file mode 100644 index 000000000..dfb385199 --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/infra/redis/main.bicep @@ -0,0 +1,44 @@ +@description('Location for the Azure Managed Redis cluster. Defaults to the resource group location.') +param location string = resourceGroup().location + +@description('Name of the Azure Managed Redis cluster. Defaults to a name derived from the resource group.') +param clusterName string = 'amr-${uniqueString(resourceGroup().id)}' + +@description('Azure Managed Redis SKU. Balanced_B0 is the smallest tier and is sufficient for this sample.') +param skuName string = 'Balanced_B0' + +// Azure Managed Redis (Microsoft.Cache/redisEnterprise). The database uses the +// EnterpriseCluster policy so the four custom stores can run multi-key commands +// (pipelines, sorted sets) against a single logical endpoint without OSS +// cluster hash-slot restrictions. +resource cluster 'Microsoft.Cache/redisEnterprise@2025-07-01' = { + name: clusterName + location: location + sku: { + name: skuName + } + properties: { + publicNetworkAccess: 'Enabled' + } +} + +resource database 'Microsoft.Cache/redisEnterprise/databases@2025-07-01' = { + parent: cluster + name: 'default' + properties: { + accessKeysAuthentication: 'Disabled' + clientProtocol: 'Encrypted' + port: 10000 + clusteringPolicy: 'EnterpriseCluster' + evictionPolicy: 'NoEviction' + } +} + +@description('Azure Managed Redis hostname. Consumed by azure.yaml as REDIS_HOST.') +output REDIS_HOST string = cluster.properties.hostName + +@description('Azure Managed Redis TLS port. Consumed by azure.yaml as REDIS_PORT.') +output REDIS_PORT string = string(database.properties.port) + +@description('Azure Managed Redis database resource ID. Consumed by the postdeploy hook.') +output REDIS_DATABASE_RESOURCE_ID string = database.id diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/infra/redis/main.parameters.json b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/infra/redis/main.parameters.json new file mode 100644 index 000000000..10195bd53 --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/infra/redis/main.parameters.json @@ -0,0 +1,5 @@ +{ + "$schema": "https://schema.management.azure.com/schemas/2019-04-01/deploymentParameters.json#", + "contentVersion": "1.0.0.0", + "parameters": {} +} diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/.agentignore b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/.agentignore new file mode 100644 index 000000000..8b78c6bfd --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/.agentignore @@ -0,0 +1,21 @@ +agent.yaml +agent.manifest.yaml +azure.yaml +.agentignore +.env +.env.* +.azure/ +.git/ +.vscode/ +__pycache__/ +.venv/ +venv/ +*.py[cod] +.mypy_cache/ +.pytest_cache/ +.ruff_cache/ +build/ +dist/ +*.egg-info/ +.foundry/ +*.log \ No newline at end of file diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/.dockerignore b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/.dockerignore new file mode 100644 index 000000000..04196e09f --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/.dockerignore @@ -0,0 +1,8 @@ +.git/ +.venv/ +__pycache__/ +*.py[cod] +.env +.foundry/ +.vscode/ +*.log \ No newline at end of file diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/.env.example b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/.env.example new file mode 100644 index 000000000..307d2fde3 --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/.env.example @@ -0,0 +1,13 @@ +FOUNDRY_PROJECT_ENDPOINT="..." +AZURE_AI_MODEL_DEPLOYMENT_NAME="..." + +# All four custom stores share one Redis instance. For direct `python main.py` +# runs, REDIS_URL selects the local Docker Redis. When deployed, this file is +# excluded and `azd` injects REDIS_HOST / REDIS_PORT for Azure Managed Redis; +# the agent then connects passwordless with its managed identity +# (DefaultAzureCredential). +REDIS_URL=redis://localhost:6379/0 +# Identity used to partition local, unhosted requests. +LOCAL_STORE_USER_ID=local-developer +# Retention window applied to Responses store writes. +REDIS_RESPONSE_TTL_SECONDS=86400 \ No newline at end of file diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/Dockerfile b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/Dockerfile new file mode 100644 index 000000000..c9c8298e1 --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/Dockerfile @@ -0,0 +1,14 @@ +FROM python:3.13-slim + +WORKDIR /app + +COPY ./ . + +RUN pip install --upgrade pip && \ + pip install --no-cache-dir -r requirements.txt + +EXPOSE 8088 + +RUN PYTHONDONTWRITEBYTECODE= python -m compileall -q . + +CMD ["python", "main.py"] \ No newline at end of file diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/main.py b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/main.py new file mode 100644 index 000000000..fdbe8439b --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/main.py @@ -0,0 +1,67 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Host an Agent Framework agent with four custom Redis-backed stores.""" + +import os + +from agent_framework import Agent +from agent_framework.foundry import FoundryChatClient +from agent_framework_foundry_hosting import ResponsesHostServer +from azure.identity import DefaultAzureCredential +from dotenv import load_dotenv +from redis_database import create_database +from redis_response_store import create_response_store +from redis_state_stores import ( + RedisAgentSessionStoreProvider, + RedisCheckpointStoreProvider, + RedisFunctionApprovalStoreProvider, +) +from store_validation_tools import create_store_validation_tools + +load_dotenv() + +INSTRUCTIONS = ( + "You are the Contoso Service Operations assistant. Help operators triage and " + "resolve incidents. Keep answers brief and actionable. When asked to validate " + "a custom store, call the named validation tool exactly once and include its " + "result verbatim in your response." +) + + +def main() -> None: + """Create and host the operations agent with fully custom Redis storage.""" + credential = DefaultAzureCredential() + client = FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], + credential=credential, + ) + host_kwargs: dict[str, object] = {} + if "APPLICATIONINSIGHTS_CONNECTION_STRING" not in os.environ: + host_kwargs["configure_observability"] = None + + # One shared Redis instance backs all four hosting storage extension points. + database = create_database() + agent = Agent( + client=client, + instructions=INSTRUCTIONS, + tools=create_store_validation_tools(database), + # The Responses host supplies the complete transcript on each turn. + # This setting controls downstream model-provider storage, not the + # custom stores configured below. + default_options={"store": False}, + ) + + server = ResponsesHostServer( + agent, + store=create_response_store(database), + agent_session_store_provider=RedisAgentSessionStoreProvider(database), + checkpoint_store_provider=RedisCheckpointStoreProvider(database), + function_approval_store_provider=RedisFunctionApprovalStoreProvider(database), + **host_kwargs, + ) + server.run() + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/redis_database.py b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/redis_database.py new file mode 100644 index 000000000..efca6ea7c --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/redis_database.py @@ -0,0 +1,139 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Shared Redis client used by every custom store in this sample. + +All four custom stores (responses, agent sessions, workflow checkpoints, and +function approvals) persist to a single Redis instance. One shared async client +is opened lazily and reused. Every store namespaces its keys by a common prefix +and a per-user partition, so one caller cannot read another caller's state. A +shared write lock serializes the few read-modify-write flows so concurrent +requests cannot interleave. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import os +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager + +import redis.asyncio as redis +from redis_entraid.cred_provider import create_from_default_azure_credential + +DEFAULT_REDIS_URL = "redis://localhost:6379/0" +DEFAULT_LOCAL_USER_ID = "local-developer" +DEFAULT_REDIS_PORT = 10000 +DEFAULT_LOCK_TIMEOUT_SECONDS = 30 +KEY_PREFIX = "custom-store" +# Microsoft Entra ID scope for the Azure Cache for Redis / Azure Managed Redis +# data plane. DefaultAzureCredential exchanges the agent's managed identity for +# a token in this scope. +REDIS_ENTRA_SCOPE = "https://redis.azure.com/.default" + + +def partition_for(user_id: str | None, local_user_id: str) -> str: + """Return a non-identifying tenant partition for the resolved user. + + Hashing keeps caller identifiers out of stored keys while still isolating + every user's data. Local development requests fall back to ``local_user_id``. + """ + resolved = user_id if user_id else local_user_id + return hashlib.sha256(resolved.encode("utf-8")).hexdigest() + + +class RedisDatabase: + """Own the shared async Redis client and build namespaced keys. + + When ``host`` is set the client connects to Azure Managed Redis over TLS and + authenticates passwordless with the agent's managed identity through + ``DefaultAzureCredential``. Otherwise it connects to the local Redis at + ``url`` without authentication for development. + """ + + def __init__( + self, + *, + url: str | None = None, + host: str | None = None, + port: int = DEFAULT_REDIS_PORT, + ) -> None: + self._url = url + self._host = host + self._port = port + self._client: redis.Redis | None = None + self._init_lock = asyncio.Lock() + + @asynccontextmanager + async def write_lock(self, resource_key: str) -> AsyncGenerator[None]: + """Serialize mutations of one resource across agent instances.""" + client = await self.connect() + lock = client.lock( + f"{resource_key}:lock", + timeout=DEFAULT_LOCK_TIMEOUT_SECONDS, + blocking_timeout=DEFAULT_LOCK_TIMEOUT_SECONDS, + ) + async with lock: + yield + + def key(self, *parts: str) -> str: + """Build a namespaced key from the shared prefix and the given parts.""" + return ":".join((KEY_PREFIX, *parts)) + + def partition_key(self, namespace: str, partition: str, *parts: str) -> str: + """Build a partitioned key that is safe for Redis Cluster transactions. + + Redis Cluster hashes only the value inside braces when a key contains a + hash tag. Using the per-user partition as that tag keeps every key that + can participate in one transaction or multi-key command in the same + slot while preserving isolation between users. + """ + return self.key(namespace, f"{{{partition}}}", *parts) + + async def connect(self) -> redis.Redis: + """Open the shared client on first use and reuse it thereafter.""" + if self._client is not None: + return self._client + async with self._init_lock: + if self._client is None: + # decode_responses keeps every stored value as ``str`` so the + # stores can serialize and deserialize JSON without byte handling. + if self._host: + # Passwordless: the credential provider refreshes the Entra + # ID token in the background and supplies it to every + # (re)connection. + credential_provider = create_from_default_azure_credential( + (REDIS_ENTRA_SCOPE,) + ) + self._client = redis.Redis( + host=self._host, + port=self._port, + ssl=True, + decode_responses=True, + credential_provider=credential_provider, + ) + else: + self._client = redis.from_url(self._url, decode_responses=True) + return self._client + + +def create_database() -> RedisDatabase: + """Create the shared database from Redis configuration. + + An explicitly configured ``REDIS_URL`` selects local Redis. Otherwise, + ``REDIS_HOST`` selects passwordless Azure Managed Redis (injected by the azd + ``redis`` layer). With neither setting, the local Redis default is used. + """ + url = os.getenv("REDIS_URL", "").strip() + if url: + return RedisDatabase(url=url) + host = os.getenv("REDIS_HOST", "").strip() + if host: + port = int(os.getenv("REDIS_PORT", str(DEFAULT_REDIS_PORT)).strip() or DEFAULT_REDIS_PORT) + return RedisDatabase(host=host, port=port) + return RedisDatabase(url=DEFAULT_REDIS_URL) + + +def get_local_user_id() -> str: + """Read the developer identity used to partition local, unhosted requests.""" + return os.getenv("LOCAL_STORE_USER_ID", DEFAULT_LOCAL_USER_ID).strip() or DEFAULT_LOCAL_USER_ID diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/redis_response_store.py b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/redis_response_store.py new file mode 100644 index 000000000..37e9e355f --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/redis_response_store.py @@ -0,0 +1,364 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Persist Responses protocol conversation history in Redis.""" + +from __future__ import annotations + +import json +import os +import time +import uuid +from typing import Any, Iterable, cast + +from azure.ai.agentserver.responses import ( + PlatformContext, + ResponseObject, + ResponseProviderProtocol, +) +from azure.ai.agentserver.responses.models import OutputItem, get_conversation_id +from azure.ai.agentserver.responses.store import ResponseAlreadyExistsError + +from redis_database import RedisDatabase, get_local_user_id, partition_for + +DEFAULT_RESPONSE_TTL_SECONDS = 86_400 + + +def _get_positive_int(name: str, default: int) -> int: + """Read a positive integer environment setting.""" + configured_value = os.getenv(name, str(default)) + try: + value = int(configured_value) + except ValueError as error: + raise ValueError(f"{name} must be a positive integer.") from error + if value <= 0: + raise ValueError(f"{name} must be a positive integer.") + return value + + +def _serialize(value: Any) -> str: + """Serialize a Responses model or mapping to compact JSON.""" + if hasattr(value, "as_dict") and callable(value.as_dict): + value = value.as_dict() + return json.dumps(value, separators=(",", ":"), default=str) + + +def _deserialize_mapping(value: str, *, kind: str) -> dict[str, Any]: + """Deserialize one stored value and reject corrupt non-object payloads.""" + decoded = json.loads(value) + if not isinstance(decoded, dict): + raise RuntimeError(f"Stored {kind} is not a JSON object.") + return decoded + + +def _item_id(item: Any) -> str | None: + """Extract an item identifier from a mapping or model.""" + value = item.get("id") if isinstance(item, dict) else getattr(item, "id", None) + return str(value) if value is not None else None + + +def _collect_items(items: Iterable[OutputItem]) -> list[tuple[str, str]]: + """Serialize identified items while preserving input order.""" + item_values: list[tuple[str, str]] = [] + for item in items: + identifier = _item_id(item) + if identifier is not None: + item_values.append((identifier, _serialize(item))) + return item_values + + +class RedisResponseStore(ResponseProviderProtocol): + """Store response envelopes, items, and history indexes in Redis. + + Every response envelope is a Redis hash whose retention is enforced by a + native key TTL, so expired responses disappear without a sweep. Transcript + items are individual TTL keys, and a per-conversation sorted set indexes the + responses in creation order for continuation lookups. + """ + + def __init__( + self, + *, + database: RedisDatabase, + ttl_seconds: int, + local_user_id: str, + ) -> None: + if ttl_seconds <= 0: + raise ValueError("ttl_seconds must be a positive integer.") + if not local_user_id: + raise ValueError("local_user_id must not be empty.") + self._database = database + self._ttl_seconds = ttl_seconds + self._local_user_id = local_user_id + + def _partition(self, context: PlatformContext | None) -> str: + """Return a non-identifying tenant partition for stored keys.""" + user_id = context.user_id_key if context is not None else None + return partition_for(user_id, self._local_user_id) + + def _response_key(self, partition: str, response_id: str) -> str: + return self._database.partition_key("resp", partition, response_id) + + def _item_key(self, partition: str, item_id: str) -> str: + return self._database.partition_key("item", partition, item_id) + + def _conversation_key(self, partition: str, conversation_id: str) -> str: + return self._database.partition_key("conv", partition, conversation_id) + + def _sequence_key(self, partition: str) -> str: + return self._database.partition_key("seq", partition) + + async def create_response( + self, + response: ResponseObject, + input_items: Iterable[OutputItem] | None, + history_item_ids: Iterable[str] | None, + *, + context: PlatformContext | None = None, + ) -> None: + """Persist a response and its input, output, and history references.""" + partition = self._partition(context) + response_id = str(response.get("id")) + history_ids = list(history_item_ids or []) + input_values = _collect_items(input_items or []) + output_values = _collect_items(response.get("output") or []) + item_values = [*input_values, *output_values] + conversation_id = get_conversation_id(response) + + client = await self._database.connect() + response_key = self._response_key(partition, response_id) + async with self._database.write_lock(response_key): + existing = await client.hgetall(response_key) + if existing and existing.get("deleted") != "1": + raise ResponseAlreadyExistsError(response_id) + + sequence = await client.incr(self._sequence_key(partition)) if conversation_id else 0 + mapping = { + "generation": uuid.uuid4().hex, + "response": _serialize(response), + "input_item_ids": json.dumps([identifier for identifier, _ in input_values]), + "output_item_ids": json.dumps([identifier for identifier, _ in output_values]), + "history_item_ids": json.dumps(history_ids), + "conversation_id": conversation_id or "", + "deleted": "0", + } + await self._persist_items(client, partition, item_values, history_ids) + pipe = client.pipeline(transaction=True) + # Drop any tombstone fields before rewriting the envelope. + pipe.delete(response_key) + pipe.hset(response_key, mapping=mapping) + pipe.expire(response_key, self._ttl_seconds) + if conversation_id: + conversation_key = self._conversation_key(partition, conversation_id) + # NX preserves the original ordering score across re-creation. + pipe.zadd(conversation_key, {response_id: sequence}, nx=True) + pipe.expire(conversation_key, self._ttl_seconds) + await pipe.execute() + + async def get_response( + self, + response_id: str, + *, + context: PlatformContext | None = None, + ) -> ResponseObject: + """Load a response envelope by identifier.""" + partition = self._partition(context) + client = await self._database.connect() + envelope = await client.hgetall(self._response_key(partition, response_id)) + if not envelope or envelope.get("deleted") == "1": + raise KeyError(f"response '{response_id}' not found") + response = json.loads(envelope["response"]) + if not isinstance(response, dict): + raise RuntimeError(f"Stored response '{response_id}' is corrupt.") + return cast(ResponseObject, response) + + async def update_response( + self, + response: ResponseObject, + *, + context: PlatformContext | None = None, + ) -> None: + """Replace a response envelope and its output-item index.""" + partition = self._partition(context) + response_id = str(response.get("id")) + output_values = _collect_items(response.get("output") or []) + conversation_id = get_conversation_id(response) + + client = await self._database.connect() + response_key = self._response_key(partition, response_id) + async with self._database.write_lock(response_key): + envelope = await client.hgetall(response_key) + if not envelope or envelope.get("deleted") == "1": + raise KeyError(f"response '{response_id}' not found") + + retained_item_ids = [ + *json.loads(envelope["history_item_ids"]), + *json.loads(envelope["input_item_ids"]), + ] + sequence = await client.incr(self._sequence_key(partition)) if conversation_id else 0 + await self._persist_items( + client, partition, output_values, retained_item_ids + ) + pipe = client.pipeline(transaction=True) + pipe.hset( + response_key, + mapping={ + "response": _serialize(response), + "output_item_ids": json.dumps([identifier for identifier, _ in output_values]), + "conversation_id": conversation_id or "", + }, + ) + pipe.expire(response_key, self._ttl_seconds) + if conversation_id: + conversation_key = self._conversation_key(partition, conversation_id) + pipe.zadd(conversation_key, {response_id: sequence}, nx=True) + pipe.expire(conversation_key, self._ttl_seconds) + await pipe.execute() + + async def delete_response( + self, + response_id: str, + *, + context: PlatformContext | None = None, + ) -> None: + """Soft-delete a response while retaining a temporary tombstone.""" + partition = self._partition(context) + client = await self._database.connect() + response_key = self._response_key(partition, response_id) + async with self._database.write_lock(response_key): + envelope = await client.hgetall(response_key) + if not envelope or envelope.get("deleted") == "1": + raise KeyError(f"response '{response_id}' not found") + pipe = client.pipeline(transaction=True) + pipe.hset(response_key, "deleted", "1") + pipe.expire(response_key, self._ttl_seconds) + await pipe.execute() + + async def get_input_items( + self, + response_id: str, + limit: int = 20, + ascending: bool = False, + after: str | None = None, + before: str | None = None, + *, + context: PlatformContext | None = None, + ) -> list[OutputItem]: + """Load paginated history and input items for one response.""" + partition = self._partition(context) + client = await self._database.connect() + envelope = await client.hgetall(self._response_key(partition, response_id)) + if not envelope: + raise KeyError(f"response '{response_id}' not found") + if envelope.get("deleted") == "1": + raise ValueError(f"response '{response_id}' has been deleted") + + item_ids = [*json.loads(envelope["history_item_ids"]), *json.loads(envelope["input_item_ids"])] + ordered_ids = item_ids if ascending else list(reversed(item_ids)) + if after is not None and after in ordered_ids: + ordered_ids = ordered_ids[ordered_ids.index(after) + 1 :] + if before is not None and before in ordered_ids: + ordered_ids = ordered_ids[: ordered_ids.index(before)] + safe_limit = max(1, min(100, int(limit))) + items = await self.get_items(ordered_ids[:safe_limit], context=context) + return [item for item in items if item is not None] + + async def get_items( + self, + item_ids: Iterable[str], + *, + context: PlatformContext | None = None, + ) -> list[OutputItem | None]: + """Load items by identifier while preserving order and misses.""" + identifiers = list(item_ids) + if not identifiers: + return [] + partition = self._partition(context) + client = await self._database.connect() + values = await client.mget([self._item_key(partition, item_id) for item_id in identifiers]) + return [ + cast(OutputItem, _deserialize_mapping(value, kind="response item")) + if value is not None + else None + for value in values + ] + + async def get_history_item_ids( + self, + previous_response_id: str | None, + conversation_id: str | None, + limit: int, + *, + context: PlatformContext | None = None, + ) -> list[str]: + """Resolve chronological item identifiers for response continuation.""" + if limit == 0 or limit < -1: + return [] + partition = self._partition(context) + client = await self._database.connect() + + response_ids: list[str] = [] + if previous_response_id is not None: + response_ids.append(previous_response_id) + if conversation_id is not None: + # The sorted set orders responses by their creation sequence. + members = await client.zrange(self._conversation_key(partition, conversation_id), 0, -1) + response_ids.extend(members) + if not response_ids: + return [] + + resolved: list[str] = [] + for response_id in response_ids: + envelope = await client.hgetall(self._response_key(partition, response_id)) + if not envelope or envelope.get("deleted") == "1": + continue + if conversation_id is not None and envelope.get("conversation_id") != conversation_id: + continue + resolved.extend(json.loads(envelope["history_item_ids"])) + resolved.extend(json.loads(envelope["input_item_ids"])) + resolved.extend(json.loads(envelope["output_item_ids"])) + unique_resolved = list(dict.fromkeys(resolved)) + return unique_resolved if limit == -1 else unique_resolved[-limit:] + + async def _persist_items( + self, + client: Any, + partition: str, + item_values: list[tuple[str, str]], + retained_item_ids: list[str], + ) -> None: + """Persist referenced items before an envelope makes them readable.""" + if not item_values and not retained_item_ids: + return + pipe = client.pipeline(transaction=True) + self._write_items(pipe, partition, item_values) + self._refresh_items(pipe, partition, retained_item_ids) + await pipe.execute() + + def _write_items( + self, + pipe: Any, + partition: str, + item_values: list[tuple[str, str]], + ) -> None: + """Queue item writes with a refreshed retention window.""" + for item_id, value in item_values: + pipe.set(self._item_key(partition, item_id), value, ex=self._ttl_seconds) + + def _refresh_items( + self, + pipe: Any, + partition: str, + item_ids: list[str], + ) -> None: + """Queue retention-window extensions on referenced items already stored.""" + for item_id in item_ids: + pipe.expire(self._item_key(partition, item_id), self._ttl_seconds) + + +def create_response_store(database: RedisDatabase) -> RedisResponseStore: + """Create the Redis-backed Responses store.""" + return RedisResponseStore( + database=database, + ttl_seconds=_get_positive_int("REDIS_RESPONSE_TTL_SECONDS", DEFAULT_RESPONSE_TTL_SECONDS), + local_user_id=get_local_user_id(), + ) diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/redis_state_stores.py b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/redis_state_stores.py new file mode 100644 index 000000000..4859a765f --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/redis_state_stores.py @@ -0,0 +1,258 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Redis-backed agent session, workflow checkpoint, and function approval stores. + +These three stores complement the Redis Responses store. Each one implements the +Agent Framework hosting extension point it plugs into and persists to the shared +Redis instance: + +* ``RedisAgentSessionStore`` implements ``SessionStore`` for MAF agent sessions. +* ``RedisCheckpointStore`` implements ``CheckpointStorage`` for workflow + checkpoints, scoped to one workflow context. +* ``RedisFunctionApprovalStore`` implements ``FunctionApprovalStore`` for + human-in-the-loop function approval requests. + +Every key is namespaced by the requesting user so one caller cannot read another +caller's state. +""" + +from __future__ import annotations + +import json +from datetime import datetime + +from agent_framework import ( + AgentSession, + CheckpointID, + CheckpointStorage, + Content, + SessionStore, + WorkflowCheckpoint, + WorkflowCheckpointException, +) +from agent_framework._workflows._checkpoint_encoding import ( + decode_checkpoint_value, + encode_checkpoint_value, +) +from agent_framework_foundry_hosting import ( + ContextScopedStoreProvider, + FunctionApprovalStore, + StoreProvider, +) +from azure.ai.agentserver.core import AgentConfig, FoundryAgentRequestContext + +from redis_database import RedisDatabase, get_local_user_id, partition_for + +# region Agent session persistence + + +class RedisAgentSessionStore(SessionStore): + """Agent session store backed by the shared Redis instance.""" + + def __init__(self, database: RedisDatabase, partition: str) -> None: + super().__init__() + self._database = database + self._partition = partition + + def _key(self, session_id: str) -> str: + return self._database.partition_key("session", self._partition, session_id) + + async def get(self, session_id: str) -> AgentSession | None: + SessionStore.validate_session_id(session_id) + client = await self._database.connect() + value = await client.get(self._key(session_id)) + if value is None: + return None + return AgentSession.from_dict(json.loads(value)) + + async def set(self, session_id: str, session: AgentSession) -> None: + SessionStore.validate_session_id(session_id) + client = await self._database.connect() + await client.set(self._key(session_id), json.dumps(session.to_dict())) + + async def delete(self, session_id: str) -> None: + SessionStore.validate_session_id(session_id) + client = await self._database.connect() + await client.delete(self._key(session_id)) + + +class RedisAgentSessionStoreProvider(StoreProvider[SessionStore]): + """Provide the Redis agent session store for each request.""" + + def __init__(self, database: RedisDatabase) -> None: + self._database = database + self._local_user_id = get_local_user_id() + + def get_store( + self, + *, + config: AgentConfig, + platform_context: FoundryAgentRequestContext, + ) -> SessionStore: + partition = partition_for(platform_context.user_id, self._local_user_id) + return RedisAgentSessionStore(self._database, partition) + + +# endregion Agent session persistence + +# region Checkpoint persistence + + +class RedisCheckpointStore(CheckpointStorage): + """Workflow checkpoint store scoped to a single workflow context. + + Each checkpoint is a Redis hash, and a per-workflow set indexes the + checkpoint IDs so a workflow run can be listed and pruned as a collection. + """ + + def __init__(self, database: RedisDatabase, partition: str, context_id: str) -> None: + if not context_id: + raise ValueError("context_id must be provided to initialize a checkpoint store.") + self._database = database + self._partition = partition + self._context_id = context_id + + def _checkpoint_key(self, checkpoint_id: str) -> str: + return self._database.partition_key( + "ckpt", self._partition, self._context_id, checkpoint_id + ) + + def _index_key(self, workflow_name: str) -> str: + return self._database.partition_key( + "ckptidx", self._partition, self._context_id, workflow_name + ) + + async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: + client = await self._database.connect() + checkpoint_key = self._checkpoint_key(checkpoint.checkpoint_id) + async with self._database.write_lock(checkpoint_key): + pipe = client.pipeline(transaction=True) + pipe.hset( + checkpoint_key, + mapping={ + "workflow_name": checkpoint.workflow_name, + "timestamp": checkpoint.timestamp, + "value": json.dumps(encode_checkpoint_value(checkpoint.to_dict())), + }, + ) + pipe.sadd(self._index_key(checkpoint.workflow_name), checkpoint.checkpoint_id) + await pipe.execute() + return checkpoint.checkpoint_id + + async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint: + client = await self._database.connect() + value = await client.hget(self._checkpoint_key(checkpoint_id), "value") + if value is None: + raise WorkflowCheckpointException(f"No checkpoint found with ID {checkpoint_id}") + return WorkflowCheckpoint.from_dict( + decode_checkpoint_value(json.loads(value), allowed_types=frozenset()) + ) + + async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoint]: + client = await self._database.connect() + checkpoint_ids = await client.smembers(self._index_key(workflow_name)) + checkpoints: list[WorkflowCheckpoint] = [] + for checkpoint_id in checkpoint_ids: + value = await client.hget(self._checkpoint_key(checkpoint_id), "value") + if value is None: + continue + checkpoints.append( + WorkflowCheckpoint.from_dict( + decode_checkpoint_value(json.loads(value), allowed_types=frozenset()) + ) + ) + return checkpoints + + async def delete(self, checkpoint_id: CheckpointID) -> bool: + client = await self._database.connect() + checkpoint_key = self._checkpoint_key(checkpoint_id) + async with self._database.write_lock(checkpoint_key): + workflow_name = await client.hget(checkpoint_key, "workflow_name") + pipe = client.pipeline(transaction=True) + pipe.delete(checkpoint_key) + if workflow_name is not None: + pipe.srem(self._index_key(workflow_name), checkpoint_id) + results = await pipe.execute() + return bool(results[0]) + + async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: + checkpoints = await self.list_checkpoints(workflow_name=workflow_name) + if not checkpoints: + return None + return max(checkpoints, key=lambda checkpoint: datetime.fromisoformat(checkpoint.timestamp)) + + async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID]: + checkpoints = await self.list_checkpoints(workflow_name=workflow_name) + return [checkpoint.checkpoint_id for checkpoint in checkpoints] + + +class RedisCheckpointStoreProvider(ContextScopedStoreProvider[CheckpointStorage]): + """Provide the Redis checkpoint store scoped to a workflow context.""" + + def __init__(self, database: RedisDatabase) -> None: + self._database = database + self._local_user_id = get_local_user_id() + + def get_store( + self, + *, + config: AgentConfig, + context_id: str, + platform_context: FoundryAgentRequestContext, + ) -> CheckpointStorage: + if not context_id: + raise ValueError("context_id must be provided to get a checkpoint store.") + partition = partition_for(platform_context.user_id, self._local_user_id) + return RedisCheckpointStore(self._database, partition, context_id) + + +# endregion Checkpoint persistence + +# region Function approval persistence + + +class RedisFunctionApprovalStore(FunctionApprovalStore): + """Function approval store backed by the shared Redis instance.""" + + def __init__(self, database: RedisDatabase, partition: str) -> None: + self._database = database + self._partition = partition + + def _key(self, approval_request_id: str) -> str: + return self._database.partition_key("approval", self._partition, approval_request_id) + + async def save_approval_request(self, approval_request_id: str, request: Content) -> None: + client = await self._database.connect() + # SET NX makes the create atomic and rejects duplicate approval IDs. + created = await client.set( + self._key(approval_request_id), json.dumps(request.to_dict()), nx=True + ) + if not created: + raise ValueError(f"Approval request with ID '{approval_request_id}' already exists.") + + async def load_approval_request(self, approval_request_id: str) -> Content: + client = await self._database.connect() + value = await client.get(self._key(approval_request_id)) + if value is None: + raise KeyError(f"Approval request with ID '{approval_request_id}' does not exist.") + return Content.from_dict(json.loads(value)) + + +class RedisFunctionApprovalStoreProvider(StoreProvider[FunctionApprovalStore]): + """Provide the Redis function approval store for each request.""" + + def __init__(self, database: RedisDatabase) -> None: + self._database = database + self._local_user_id = get_local_user_id() + + def get_store( + self, + *, + config: AgentConfig, + platform_context: FoundryAgentRequestContext, + ) -> FunctionApprovalStore: + partition = partition_for(platform_context.user_id, self._local_user_id) + return RedisFunctionApprovalStore(self._database, partition) + + +# endregion Function approval persistence diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/requirements.in b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/requirements.in new file mode 100644 index 000000000..b18a7b810 --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/requirements.in @@ -0,0 +1,9 @@ +agent-framework-foundry +agent-framework-foundry-hosting>=1.0.0b260813 +python-dotenv +redis>=5.0 +# Passwordless Microsoft Entra ID auth for Azure Managed Redis via DefaultAzureCredential. +redis-entraid>=1.2,<2 + +# debugpy enables local debugging of this agent with the Foundry Toolkit VS Code extension. +debugpy \ No newline at end of file diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/requirements.txt b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/requirements.txt new file mode 100644 index 000000000..7ae3f9b31 --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/requirements.txt @@ -0,0 +1,108 @@ +# This file was autogenerated by uv via the following command: +# uv pip compile requirements.in --python 3.13 --python-platform x86_64-unknown-linux-gnu --no-annotate --output-file requirements.txt +agent-framework-core==1.18.0 +agent-framework-foundry==1.13.0 +agent-framework-foundry-hosting==1.0.0b260821 +agent-framework-openai==1.14.3 +aiohappyeyeballs==2.7.1 +aiohttp==3.14.3 +aiosignal==1.4.0 +annotated-types==0.8.0 +anyio==4.15.1 +asgiref==3.12.1 +attrs==26.1.0 +azure-ai-agentserver-core==2.1.0 +azure-ai-agentserver-invocations==1.1.0 +azure-ai-agentserver-responses==2.1.0 +azure-ai-inference==1.0.0b9 +azure-ai-projects==2.6.1 +azure-core==1.41.0 +azure-core-tracing-opentelemetry==1.0.0b13 +azure-identity==1.25.3 +azure-monitor-opentelemetry-exporter==1.0.0b57 +azure-storage-blob==12.30.2 +certifi==2026.7.22 +cffi==2.1.1 +charset-normalizer==3.5.1 +click==8.5.0 +cryptography==50.0.1 +debugpy==1.8.22 +frozenlist==1.8.0 +googleapis-common-protos==1.75.3 +h11==0.16.0 +h2==4.4.1 +hpack==4.2.0 +httpcore==1.0.9 +httpcore2==2.12.0 +httpx==0.28.1 +httpx-sse==0.4.3 +httpx2==2.12.0 +hypercorn==0.18.0 +hyperframe==6.1.0 +idna==3.19 +isodate==0.7.2 +jiter==0.16.0 +jsonschema==4.26.0 +jsonschema-specifications==2025.9.1 +mcp==1.30.0 +microsoft-opentelemetry==1.3.9 +msal==1.39.0 +msal-extensions==1.3.1 +msgspec==0.21.1 +msrest==0.7.1 +multidict==6.8.0 +oauthlib==3.3.1 +openai==3.13.0 +opentelemetry-api==1.44.0 +opentelemetry-exporter-otlp-proto-common==1.44.0 +opentelemetry-exporter-otlp-proto-http==1.44.0 +opentelemetry-instrumentation==0.65b0 +opentelemetry-instrumentation-asgi==0.65b0 +opentelemetry-instrumentation-dbapi==0.65b0 +opentelemetry-instrumentation-django==0.65b0 +opentelemetry-instrumentation-fastapi==0.65b0 +opentelemetry-instrumentation-flask==0.65b0 +opentelemetry-instrumentation-httpx==0.65b0 +opentelemetry-instrumentation-logging==0.65b0 +opentelemetry-instrumentation-openai-agents-v2==0.1.0 +opentelemetry-instrumentation-openai-v2==2.3b0 +opentelemetry-instrumentation-psycopg2==0.65b0 +opentelemetry-instrumentation-requests==0.65b0 +opentelemetry-instrumentation-urllib==0.65b0 +opentelemetry-instrumentation-urllib3==0.65b0 +opentelemetry-instrumentation-wsgi==0.65b0 +opentelemetry-proto==1.44.0 +opentelemetry-resource-detector-azure==0.2.0 +opentelemetry-sdk==1.44.0 +opentelemetry-semantic-conventions==0.65b0 +opentelemetry-util-genai==0.3b0 +opentelemetry-util-http==0.65b0 +packaging==26.3 +priority==2.0.0 +propcache==0.5.2 +protobuf==7.36.1 +psutil==7.2.2 +pycparser==3.0 +pydantic==2.13.5 +pydantic-core==2.46.5 +pydantic-settings==2.15.0 +pyjwt==2.13.0 +python-dotenv==1.2.3 +python-multipart==0.0.32 +redis==8.1.0 +redis-entraid==1.2.1 +referencing==0.37.0 +requests==2.34.2 +requests-oauthlib==2.0.0 +rpds-py==2026.6.3 +sniffio==1.3.1 +sse-starlette==3.4.11 +starlette==1.6.0 +truststore==0.10.4 +typing-extensions==4.16.0 +typing-inspection==0.4.4 +urllib3==2.7.0 +uvicorn==0.52.4 +wrapt==1.17.3 +wsproto==1.3.2 +yarl==1.25.1 diff --git a/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/store_validation_tools.py b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/store_validation_tools.py new file mode 100644 index 000000000..f1af109bb --- /dev/null +++ b/samples/python/hosted-agents/agent-framework/responses/25-custom-store/src/agent-framework-agent-custom-store-responses/store_validation_tools.py @@ -0,0 +1,113 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tools that exercise and inspect the four custom Redis stores.""" + +from __future__ import annotations + +from typing import Annotated, Any, Never +from uuid import uuid4 + +from agent_framework import ( + Executor, + FunctionTool, + WorkflowBuilder, + WorkflowContext, + handler, + tool, +) +from azure.ai.agentserver.core import AgentConfig, get_request_context +from pydantic import Field +from redis_database import RedisDatabase, get_local_user_id, partition_for +from redis_state_stores import RedisCheckpointStoreProvider + +CHECKPOINT_WORKFLOW_NAME = "custom-store-validation" + + +class _CheckpointMarkerExecutor(Executor): + """Persist and yield a marker so checkpoint validation is deterministic.""" + + _marker: str | None = None + + @handler + async def record_marker(self, marker: str, ctx: WorkflowContext[Never, str]) -> None: + self._marker = marker + await ctx.yield_output(marker) + + async def on_checkpoint_save(self) -> dict[str, Any]: + return {"marker": self._marker} + + +def create_store_validation_tools(database: RedisDatabase) -> list[FunctionTool]: + """Create tools that validate checkpoint, approval, and stored-record data.""" + checkpoint_provider = RedisCheckpointStoreProvider(database) + + @tool(approval_mode="never_require") + async def validate_checkpoint_store( + marker: Annotated[str, Field(description="Exact marker to persist in a workflow checkpoint.")], + ) -> str: + """Run a checkpointed workflow, then load and verify its saved checkpoint.""" + request_context = get_request_context() + context_id = request_context.call_id or uuid4().hex + checkpoint_store = checkpoint_provider.get_store( + config=AgentConfig.from_env(), + context_id=f"validation-{context_id}", + platform_context=request_context, + ) + workflow = WorkflowBuilder( + start_executor=_CheckpointMarkerExecutor(id="record-marker"), + name=CHECKPOINT_WORKFLOW_NAME, + ).build() + + events = await workflow.run(marker, checkpoint_storage=checkpoint_store) + if events.get_outputs() != [marker]: + raise RuntimeError("The checkpoint validation workflow returned an unexpected output.") + + checkpoint_ids = await checkpoint_store.list_checkpoint_ids(workflow_name=CHECKPOINT_WORKFLOW_NAME) + if not checkpoint_ids: + raise RuntimeError("The checkpoint validation workflow did not persist a checkpoint.") + checkpoints = [await checkpoint_store.load(checkpoint_id) for checkpoint_id in checkpoint_ids] + if any(checkpoint.workflow_name != CHECKPOINT_WORKFLOW_NAME for checkpoint in checkpoints): + raise RuntimeError("The checkpoint validation workflow loaded unexpected data.") + stored_markers = [ + executor_state.get("marker") + for checkpoint in checkpoints + for executor_states in [checkpoint.state.get("_executor_state")] + if isinstance(executor_states, dict) + for executor_state in [executor_states.get("record-marker")] + if isinstance(executor_state, dict) + ] + if marker not in stored_markers: + raise RuntimeError("The checkpoint validation workflow did not reload the persisted marker.") + + return f"CHECKPOINT_STORE_OK {marker}" + + @tool(approval_mode="always_require") + def approve_incident_change( + change_id: Annotated[str, Field(description="Change request identifier to approve.")], + ) -> str: + """Approve a simulated incident change after host-managed human approval.""" + return f"APPROVAL_STORE_OK {change_id}" + + @tool(approval_mode="never_require") + async def inspect_custom_store() -> str: + """Return record counts for the current user across all four custom stores.""" + request_context = get_request_context() + partition = partition_for(request_context.user_id, get_local_user_id()) + client = await database.connect() + # Each store owns a distinct key prefix, so a per-store SCAN counts only + # that store's keys for this partition. Expired keys are excluded natively. + stores = ( + ("response_store", database.partition_key("resp", partition, "*")), + ("agent_session_store", database.partition_key("session", partition, "*")), + ("checkpoint_store", database.partition_key("ckpt", partition, "*")), + ("function_approval_store", database.partition_key("approval", partition, "*")), + ) + counts: list[str] = [] + for label, pattern in stores: + count = 0 + async for _ in client.scan_iter(match=pattern, count=100): + count += 1 + counts.append(f"{label}={count}") + return f"STORE_COUNTS {' '.join(counts)}" + + return [validate_checkpoint_store, approve_incident_change, inspect_custom_store] \ No newline at end of file