Support header isolation keys for agent sessions by glharper · Pull Request #8206 · Azure/azure-dev

glharper · 2026-05-15T16:15:00Z

Summary

add --user-isolation-key and --chat-isolation-key to session-level Azure AI Agents commands
propagate the Foundry isolation headers through invoke, monitor, files, and sessions APIs
preserve existing --isolation-key behavior for session ownership headers

Tests

go test .\... -count=1

Add user and chat isolation key flags to session-level agent commands and propagate them through invoke, monitor, files, and sessions requests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

github-actions · 2026-05-18T15:48:03Z

🔗 Linked Issue Required

Thanks for the contribution! Please link a GitHub issue to this PR by adding Fixes #123 to the description or using the sidebar.
No issue yet? Feel free to create one!

Copilot

Pull request overview

Adds support for propagating Foundry header-based isolation keys across session-scoped Azure AI Agents operations in the azure.ai.agents extension, while preserving the existing session-ownership isolation key behavior.

Changes:

Introduces SessionRequestOptions + well-known header constants in agent_api and threads these options through session lifecycle, files, and log streaming APIs.
Adds --user-isolation-key and --chat-isolation-key flags to session-level CLI commands (sessions, invoke, monitor, files) and applies them to outbound requests (including LRO polling and conversation creation).
Expands unit tests to validate header propagation across the affected operations and command flag registration.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
cli/azd/extensions/azure.ai.agents/internal/pkg/agents/agent_api/operations.go	Adds isolation header constants + `SessionRequestOptions`, and applies headers across session/files/logstream operations.
cli/azd/extensions/azure.ai.agents/internal/pkg/agents/agent_api/operations_test.go	Adds tests ensuring isolation headers are applied for session lifecycle, file ops, and log streaming.
cli/azd/extensions/azure.ai.agents/internal/cmd/isolation_headers.go	New shared flag plumbing for `--user-isolation-key` / `--chat-isolation-key` and conversion to request options.
cli/azd/extensions/azure.ai.agents/internal/cmd/session.go	Registers isolation header flags on session commands and propagates options to agent API calls.
cli/azd/extensions/azure.ai.agents/internal/cmd/session_test.go	Extends tests to assert new isolation flags exist with correct defaults.
cli/azd/extensions/azure.ai.agents/internal/cmd/monitor.go	Adds isolation header flags and passes options to session log streaming.
cli/azd/extensions/azure.ai.agents/internal/cmd/monitor_test.go	Extends tests to assert new isolation flags exist with correct defaults.
cli/azd/extensions/azure.ai.agents/internal/cmd/invoke.go	Adds isolation header flags and propagates options through invoke flows (requests, LRO polling, conversation creation).
cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_test.go	Adds tests asserting flag registration and propagation through LRO polling and conversation creation.
cli/azd/extensions/azure.ai.agents/internal/cmd/helpers.go	Threads session request options into conversation ID resolution and conversation creation.
cli/azd/extensions/azure.ai.agents/internal/cmd/files.go	Adds isolation header flags for file operations and propagates options to agent API calls.
cli/azd/extensions/azure.ai.agents/internal/cmd/files_test.go	Extends tests to assert new isolation flags exist across file subcommands.

Comments suppressed due to low confidence (1)

cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_test.go:1162

This test stores the request pointer in an outer variable from the server handler goroutine and reads it from the test goroutine without synchronization; with t.Parallel() this will be flagged as a race by the race detector. Use a channel to hand the request back to the test (or otherwise synchronize) before making assertions.

func TestCreateConversation_PropagatesIsolationHeaders(t *testing.T) {
	t.Parallel()

	var request *http.Request
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		request = r
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(`{"id":"conv-headers"}`))
	}))
	defer srv.Close()

+	"azureaiagent/internal/pkg/agents/agent_api"
+
+	"github.com/spf13/cobra"


+func TestGetAgentSessionLogStream_ApplyIsolationHeaders(t *testing.T) {
+	var request *http.Request
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		request = r
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("event: log\n\n"))
+	}))


+func TestHandleInvocationLRO_PropagatesIsolationHeaders(t *testing.T) {
+	origInterval := defaultLROPollInterval
+	defaultLROPollInterval = time.Millisecond
+	t.Cleanup(func() { defaultLROPollInterval = origInterval })
+
+	var pollRequest *http.Request
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		pollRequest = r
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte(`{"status":"completed"}`))
+	}))


glharper · 2026-05-18T19:15:49Z

@copilot apply changes based on the comments in this thread

Agent-Logs-Url: https://github.com/Azure/azure-dev/sessions/4eb6559e-3bc1-409c-942a-12ccc66e54e4 Co-authored-by: glharper <64209257+glharper@users.noreply.github.com>

Copilot · 2026-05-18T19:22:01Z

@copilot apply changes based on the comments in this thread

Applied all three review comments in af6d390:

Import grouping in isolation_headers.go: reordered to stdlib → external (github.com/spf13/cobra) → local (azureaiagent/...).
Data race in operations_test.go: replaced the shared var request *http.Request assignment in the httptest handler with a buffered channel (reqCh := make(chan *http.Request, 1)).
Data race in invoke_test.go: same fix applied to both TestHandleInvocationLRO_PropagatesIsolationHeaders and TestCreateConversation_PropagatesIsolationHeaders.

All three tests pass under go test -race.

wbreza

Supplemental review — focused on what the prior reviewer didn't already cover. Prior race-condition findings stand; please address those.

Test Coverage

LRO polling test only verifies the last poll request — TestHandleInvocationLRO_PropagatesIsolationHeaders (invoke_test.go ~L446–497) overwrites pollRequest on each iteration, so headers are only asserted on the final poll. Recommend pairing this with the race fix already requested: switch to a buffered channel (or mutex-guarded slice) and assert headers on every poll request. That single change resolves the race and tightens the contract.

No partial-set coverage for the new flags — every new test sets both UserIsolationKey and ChatIsolationKey (or neither). Please add cases for user-only and chat-only to lock in the independence of the two headers. ApplyHeaders is implemented correctly today, but a future change that gates one header on the other being set would currently slip through.

Minor

Duplicated test logic — TestApplyIsolationHeaders (invoke_test.go ~L396–419) duplicates TestSessionRequestOptions_ApplyHeaders (operations_test.go ~L1170–1185). Consider removing it, or repurposing it to test the CLI-flag-struct → SessionRequestOptions conversion (sessionRequestOptions / sessionRequestOptionsWithSessionKey), which is currently only exercised indirectly through end-to-end tests.

Flag help text disambiguation — isolation_headers.go adds --user-isolation-key and --chat-isolation-key alongside the pre-existing --isolation-key. The three names look very similar in --help. Consider noting in the descriptions that the new flags are independent of the session-ownership --isolation-key to reduce CLI confusion.

Looks good

Headers are applied consistently across sessions, files, invoke (incl. LRO polling), monitor/log-stream, and conversation creation.
SessionRequestOptions.ApplyHeaders is nil-safe and filters empty values.
Existing --isolation-key semantics preserved via sessionRequestOptionsWithSessionKey.
Header constants are exported with doc comments.
No isolation-key values are logged in debug output.

Support header isolation keys for agent sessions

b787124

Add user and chat isolation key flags to session-level agent commands and propagate them through invoke, monitor, files, and sessions requests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

microsoft-github-policy-service Bot assigned glharper May 15, 2026

glharper marked this pull request as ready for review May 18, 2026 15:47

Copilot AI review requested due to automatic review settings May 18, 2026 15:47

glharper requested review from JeffreyCA, hemarina, jongio, rajeshkamal5050, tg-msft, therealjohn, trangevi, trrwilson, vhvb1989, wbreza and weikanglim as code owners May 18, 2026 15:47

Copilot started reviewing on behalf of glharper May 18, 2026 15:48 View session

Copilot AI reviewed May 18, 2026

View reviewed changes

Copilot started work on behalf of glharper May 18, 2026 19:16 View session

fix: import grouping and data races in isolation header tests

af6d390

Agent-Logs-Url: https://github.com/Azure/azure-dev/sessions/4eb6559e-3bc1-409c-942a-12ccc66e54e4 Co-authored-by: glharper <64209257+glharper@users.noreply.github.com>

Copilot finished work on behalf of glharper May 18, 2026 19:22

wbreza reviewed May 18, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support header isolation keys for agent sessions#8206

Support header isolation keys for agent sessions#8206
glharper wants to merge 2 commits into
mainfrom
glharper/ai-agent-isolation-headers

glharper commented May 15, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented May 18, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

glharper commented May 18, 2026

Uh oh!

Copilot AI commented May 18, 2026

Uh oh!

wbreza left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		"azureaiagent/internal/pkg/agents/agent_api"

		"github.com/spf13/cobra"

Conversation

glharper commented May 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Tests

Uh oh!

github-actions Bot commented May 18, 2026

🔗 Linked Issue Required

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

glharper commented May 18, 2026

Uh oh!

Copilot AI commented May 18, 2026

Uh oh!

wbreza left a comment

Choose a reason for hiding this comment

Test Coverage

Minor

Looks good

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

glharper commented May 15, 2026 •

edited

Loading