Add --client-header flag to azd ai agent invoke (#8478) (#9043)

* Add --client-header flag to azd ai agent invoke (#8478)

Add a --client-header flag to send custom x-client-* request headers with
azd ai agent invoke. The responses and invocations protocol proxies forward
the x-client-* header family to the agent container, so the flag accepts only
x-client-* names (case-insensitive) and rejects any other name with guidance
toward --user-identity/--call-id for identity headers.

The a2a protocol uses its own fixed internal-routing header set and does not
propagate x-client-*, so combining --client-header with --protocol a2a is a
hard error. Custom headers are applied before managed headers (Authorization,
Content-Type, user identity) so managed headers always win.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

* fix: address PR review comments on invoke.go and invoke_headers_test.go

- Fix comment 'translating head' -> 'translating proxy' in invoke.go
- Add a2a + client-header re-validation in Run() after protocol
  auto-detection to prevent silent header drops when --protocol is omitted
- Remove misleading/dead lines from TestApplyCustomHeaders that set
  Content-Type before calling applyCustomHeaders (reverse order from
  what the comment claimed), leaving only the correct builder-order test

Co-authored-by: trangevi <26490000+trangevi@users.noreply.github.com>

---------

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: trangevi <26490000+trangevi@users.noreply.github.com>
Co-authored-by: Travis Angevine <trangevi@microsoft.com>

Author: 4 people

cli/azd/extensions/azure.ai.agents/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Features Added
 
+- [[#8478]](https://github.com/Azure/azure-dev/issues/8478) Add `--client-header` to `azd ai agent invoke` for sending custom `x-client-*` request headers in `"Name: Value"` format (repeatable). The responses and invocations protocols forward the `x-client-*` header family to the agent; other header names are rejected, and the flag is not supported with the `a2a` protocol (which does not propagate `x-client-*` headers). Managed headers (`Authorization`, `Content-Type`, user identity) always take precedence.
 - [[#8939]](https://github.com/Azure/azure-dev/pull/8939) Add native support for the Activity protocol to `azd ai agent`. `azd ai agent init` can now scaffold an Activity-protocol agent (defaulting to the service-recommended version `2.0.0`), `azd deploy` provisions a companion Azure Bot Service registration authenticated with `BotServiceRbac` and prints Microsoft Teams setup guidance, and `azd down` tears the bot down. Both init-from-code and init-from-manifest flows are supported.
 - [[#8989]](https://github.com/Azure/azure-dev/pull/8989) Add `a2a` protocol support to `azd ai agent invoke`. A plain message is wrapped in a JSON-RPC 2.0 `message/send` request, `--input-file` sends a complete JSON-RPC request, and `--output raw` dumps the response verbatim. A2A is remote-only (not available with `--local`).
 

cli/azd/extensions/azure.ai.agents/internal/cmd/invoke.go

@@ -45,6 +45,7 @@ type invokeFlags struct {
 	version         string
 	outputFmt       string
 	callID          string
+	clientHeaders   []string
 }
 
 // outputRaw is the sentinel value of the inherited --output flag that selects
@@ -63,9 +64,10 @@ const maxInvokeVersionLength = 128
 var createInvokeVersionSession = createInvokeVersionSessionImpl
 
 type InvokeAction struct {
-	flags    *invokeFlags
-	noPrompt bool
-	endpoint *parsedAgentEndpoint
+	flags         *invokeFlags
+	noPrompt      bool
+	endpoint      *parsedAgentEndpoint
+	clientHeaders http.Header
 }
 
 func newInvokeCommand(extCtx *azdext.ExtensionContext) *cobra.Command {
@@ -105,6 +107,12 @@ Use --call-id to send a call ID with a local invoke. It is sent as the
 x-agent-foundry-call-id header and applies only to local invocations; it is
 ignored for remote requests.
 
+Use --client-header to send custom x-client-* request headers in "Name: Value"
+format (repeatable). The responses and invocations protocols forward the
+x-client-* header family to the agent; other header names are rejected, and
+the flag is not supported with the a2a protocol (which does not propagate
+x-client-* headers). For identity headers use --user-identity or --call-id.
+
 Use --output raw (or -o raw) to dump the unmodified server response (status
 line, headers, and body verbatim) to stdout. Useful for debugging server
 behavior and inspecting response headers (for example, the agent version
@@ -146,6 +154,9 @@ suppressed in raw mode.`,
   # Dump the raw server response (status line, headers, body) for debugging
   azd ai agent invoke --output raw "Hello!"
 
+  # Send custom x-client-* headers (repeatable)
+  azd ai agent invoke --client-header "x-client-request-id: abc123" --client-header "x-client-tenant: contoso" "Hello!"
+
   # Invoke a deployed agent from any directory using the endpoint URL shown by 'azd ai agent show'
   azd ai agent invoke \
 	  --agent-endpoint https://<acct>.services.ai.azure.com/api/projects/<proj>/agents/<name>/endpoint/protocols/openai/responses?api-version=v1 \
@@ -230,6 +241,23 @@ suppressed in raw mode.`,
 				}
 			}
 
+			clientHeaders, err := parseCustomHeaders(flags.clientHeaders)
+			if err != nil {
+				return err
+			}
+			// A2A routes through a translating proxy with its own fixed header
+			// set and does not propagate x-client-* to the agent, so accepting
+			// the flag there would silently drop the headers.
+			if len(clientHeaders) > 0 && agent_api.AgentProtocol(flags.protocol) == agent_api.AgentProtocolA2A {
+				return exterrors.Validation(
+					exterrors.CodeInvalidParameter,
+					"--client-header is not supported with the a2a protocol",
+					"the a2a protocol does not forward x-client-* headers to the agent; "+
+						"use --protocol responses or invocations to send client headers",
+				)
+			}
+			action.clientHeaders = clientHeaders
+
 			return action.Run(ctx)
 		},
 	}
@@ -258,6 +286,14 @@ suppressed in raw mode.`,
 		"Call ID header value (sent as "+agent_api.AgentFoundryCallIDHeader+" for local invocations only; "+
 			"ignored for remote requests)",
 	)
+	cmd.Flags().StringArrayVar(
+		&flags.clientHeaders,
+		"client-header",
+		nil,
+		`Custom x-client-* request header in "Name: Value" format (repeatable). `+
+			"The responses and invocations protocols forward the x-client-* header family to the agent; "+
+			"other header names are rejected and the flag is not supported with a2a.",
+	)
 	cmd.Flags().StringVar(
 		&flags.agentEndpoint,
 		"agent-endpoint",
@@ -386,6 +422,21 @@ func (a *InvokeAction) Run(ctx context.Context) error {
 		return err
 	}
 
+	// Re-validate after protocol resolution: when --protocol was omitted the
+	// protocol may have been auto-detected as a2a (e.g. from agent.yaml). In
+	// that case the flag-parse guard above was skipped and clientHeaders was
+	// populated, but a2aRemote never calls applyCustomHeaders — the headers
+	// would be silently dropped, which is the exact silent no-op the guard
+	// intends to prevent.
+	if len(a.clientHeaders) > 0 && protocol == agent_api.AgentProtocolA2A {
+		return exterrors.Validation(
+			exterrors.CodeInvalidParameter,
+			"--client-header is not supported with the a2a protocol",
+			"the a2a protocol does not forward x-client-* headers to the agent; "+
+				"use --protocol responses or invocations to send client headers",
+		)
+	}
+
 	if a.flags.local {
 		switch protocol {
 		case agent_api.AgentProtocolInvocations:
@@ -619,6 +670,7 @@ func (a *InvokeAction) responsesLocal(ctx context.Context) error {
 	if err != nil {
 		return fmt.Errorf("failed to create request: %w", err)
 	}
+	applyCustomHeaders(req, a.clientHeaders)
 	req.Header.Set("Content-Type", "application/json")
 	applyLocalUserIdentityHeader(req, &a.flags.userIdentityFlags)
 	applyLocalCallIDHeader(req, a.flags.callID)
@@ -1018,6 +1070,7 @@ func (a *InvokeAction) responsesRemote(ctx context.Context) error {
 	if err != nil {
 		return fmt.Errorf("failed to create request: %w", err)
 	}
+	applyCustomHeaders(req, a.clientHeaders)
 	req.Header.Set("Content-Type", "application/json")
 	req.Header.Set("Authorization", "Bearer "+rc.bearerToken)
 	applyRemoteUserIdentityHeader(req, &a.flags.userIdentityFlags)
@@ -1145,6 +1198,7 @@ func (a *InvokeAction) invocationsLocal(ctx context.Context) error {
 	if err != nil {
 		return fmt.Errorf("failed to create request: %w", err)
 	}
+	applyCustomHeaders(req, a.clientHeaders)
 	req.Header.Set("Content-Type", contentTypeForBody(body))
 	applyLocalUserIdentityHeader(req, &a.flags.userIdentityFlags)
 	applyLocalCallIDHeader(req, a.flags.callID)
@@ -1255,6 +1309,7 @@ func (a *InvokeAction) invocationsRemote(ctx context.Context) error {
 	if err != nil {
 		return fmt.Errorf("failed to create request: %w", err)
 	}
+	applyCustomHeaders(req, a.clientHeaders)
 	req.Header.Set("Content-Type", contentTypeForBody(body))
 	req.Header.Set("Authorization", "Bearer "+rc.bearerToken)
 	applyRemoteUserIdentityHeader(req, &a.flags.userIdentityFlags)

cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_headers.go

@@ -0,0 +1,102 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package cmd
+
+import (
+	"fmt"
+	"net/http"
+	"strings"
+
+	"azureaiagent/internal/exterrors"
+)
+
+// clientHeaderPrefix is the required prefix (case-insensitive) for headers
+// passed via --client-header. The Responses protocol proxy forwards the x-client-*
+// family to the agent container; other header names are rejected so users are
+// not misled into thinking arbitrary headers reach the agent.
+const clientHeaderPrefix = "x-client-"
+
+// parseCustomHeaders converts raw "Name: Value" flag entries (curl -H style)
+// into an http.Header with canonicalized keys. Repeating a name adds multiple
+// values. Only x-client-* names are accepted; any other name (or a malformed
+// entry) is rejected so the mistake surfaces before any request is sent.
+func parseCustomHeaders(entries []string) (http.Header, error) {
+	if len(entries) == 0 {
+		return nil, nil
+	}
+
+	headers := http.Header{}
+	for _, entry := range entries {
+		name, value, ok := strings.Cut(entry, ":")
+		if !ok {
+			return nil, exterrors.Validation(
+				exterrors.CodeInvalidParameter,
+				fmt.Sprintf("invalid --client-header value %q", entry),
+				`use the "Name: Value" format, for example --client-header "x-client-request-id: abc123"`,
+			)
+		}
+
+		name = strings.TrimSpace(name)
+		if name == "" {
+			return nil, exterrors.Validation(
+				exterrors.CodeInvalidParameter,
+				fmt.Sprintf("invalid --client-header value %q: header name is empty", entry),
+				`use the "Name: Value" format, for example --client-header "x-client-request-id: abc123"`,
+			)
+		}
+		if !isValidHeaderName(name) {
+			return nil, exterrors.Validation(
+				exterrors.CodeInvalidParameter,
+				fmt.Sprintf("invalid --client-header name %q", name),
+				"header names may only contain letters, digits, and the characters !#$%&'*+-.^_`|~",
+			)
+		}
+		if !strings.HasPrefix(strings.ToLower(name), clientHeaderPrefix) {
+			return nil, exterrors.Validation(
+				exterrors.CodeInvalidParameter,
+				fmt.Sprintf("unsupported --client-header name %q", name),
+				fmt.Sprintf(
+					"only %s* headers can be sent with --client-header; "+
+						"use --user-identity or --call-id for identity headers",
+					clientHeaderPrefix,
+				),
+			)
+		}
+
+		headers.Add(name, strings.TrimSpace(value))
+	}
+
+	return headers, nil
+}
+
+// isValidHeaderName reports whether name is a valid HTTP field name per the
+// RFC 7230 token grammar. This guards against sending malformed headers (for
+// example, names containing spaces) that would otherwise fail deeper in the
+// stack with a less actionable error.
+func isValidHeaderName(name string) bool {
+	for _, r := range name {
+		switch {
+		case r >= 'a' && r <= 'z',
+			r >= 'A' && r <= 'Z',
+			r >= '0' && r <= '9':
+			continue
+		}
+		if !strings.ContainsRune("!#$%&'*+-.^_`|~", r) {
+			return false
+		}
+	}
+	return true
+}
+
+// applyCustomHeaders adds the parsed custom headers to req. It is called before
+// the request builders set their managed headers (Content-Type, Authorization,
+// user identity, and so on) so those always take precedence and a custom header
+// can never accidentally clobber authentication.
+func applyCustomHeaders(req *http.Request, headers http.Header) {
+	for name, values := range headers {
+		for _, value := range values {
+			req.Header.Add(name, value)
+		}
+	}
+}

cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_headers_test.go

@@ -0,0 +1,142 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package cmd
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestParseCustomHeaders(t *testing.T) {
+	t.Parallel()
+
+	t.Run("nil for no entries", func(t *testing.T) {
+		t.Parallel()
+		got, err := parseCustomHeaders(nil)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if got != nil {
+			t.Fatalf("expected nil header, got %v", got)
+		}
+	})
+
+	t.Run("parses and canonicalizes name, trims value", func(t *testing.T) {
+		t.Parallel()
+		got, err := parseCustomHeaders([]string{"x-client-request-id:  abc123  "})
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if v := got.Get("X-Client-Request-Id"); v != "abc123" {
+			t.Fatalf("expected trimmed canonical value, got %q", v)
+		}
+	})
+
+	t.Run("repeated name adds multiple values", func(t *testing.T) {
+		t.Parallel()
+		got, err := parseCustomHeaders([]string{"X-Client-Tag: a", "X-Client-Tag: b"})
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		values := got.Values("X-Client-Tag")
+		if len(values) != 2 || values[0] != "a" || values[1] != "b" {
+			t.Fatalf("expected [a b], got %v", values)
+		}
+	})
+
+	t.Run("allows empty value", func(t *testing.T) {
+		t.Parallel()
+		got, err := parseCustomHeaders([]string{"X-Client-Empty:"})
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if _, ok := got["X-Client-Empty"]; !ok {
+			t.Fatalf("expected header to be present with empty value, got %v", got)
+		}
+	})
+
+	t.Run("accepts prefix case-insensitively", func(t *testing.T) {
+		t.Parallel()
+		got, err := parseCustomHeaders([]string{"X-CLIENT-Foo: bar"})
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		if v := got.Get("X-Client-Foo"); v != "bar" {
+			t.Fatalf("expected header to be accepted, got %q", v)
+		}
+	})
+
+	t.Run("rejects missing colon", func(t *testing.T) {
+		t.Parallel()
+		if _, err := parseCustomHeaders([]string{"NoColonHere"}); err == nil {
+			t.Fatal("expected error for entry without a colon")
+		}
+	})
+
+	t.Run("rejects empty name", func(t *testing.T) {
+		t.Parallel()
+		if _, err := parseCustomHeaders([]string{":value"}); err == nil {
+			t.Fatal("expected error for empty header name")
+		}
+	})
+
+	t.Run("rejects invalid name", func(t *testing.T) {
+		t.Parallel()
+		if _, err := parseCustomHeaders([]string{"x-client-bad name: value"}); err == nil {
+			t.Fatal("expected error for header name with a space")
+		}
+	})
+
+	t.Run("rejects non x-client name", func(t *testing.T) {
+		t.Parallel()
+		if _, err := parseCustomHeaders([]string{"X-Custom-Header: value"}); err == nil {
+			t.Fatal("expected error for header name outside the x-client- family")
+		}
+	})
+}
+
+func TestIsValidHeaderName(t *testing.T) {
+	t.Parallel()
+
+	valid := []string{"X-Client-Request-Id", "Authorization", "a", "X-Foo_Bar", "x.y", "A1!#$%&'*+-.^_`|~"}
+	for _, name := range valid {
+		if !isValidHeaderName(name) {
+			t.Errorf("expected %q to be valid", name)
+		}
+	}
+
+	invalid := []string{"has space", "with:colon", "tab\ttab", "new\nline"}
+	for _, name := range invalid {
+		if isValidHeaderName(name) {
+			t.Errorf("expected %q to be invalid", name)
+		}
+	}
+}
+
+func TestApplyCustomHeaders(t *testing.T) {
+	t.Parallel()
+
+	req := httptest.NewRequest(http.MethodPost, "http://localhost/responses", nil)
+	headers := http.Header{
+		"X-Client-Request-Id": {"abc"},
+		"X-Client-Tag":        {"a", "b"},
+	}
+	applyCustomHeaders(req, headers)
+
+	if v := req.Header.Get("X-Client-Request-Id"); v != "abc" {
+		t.Fatalf("expected header to be applied, got %q", v)
+	}
+	if values := req.Header.Values("X-Client-Tag"); len(values) != 2 {
+		t.Fatalf("expected 2 values for X-Client-Tag, got %v", values)
+	}
+
+	// Simulate builder order: custom applied first, managed overrides.
+	req2 := httptest.NewRequest(http.MethodPost, "http://localhost/responses", nil)
+	applyCustomHeaders(req2, http.Header{"Content-Type": {"text/plain"}})
+	req2.Header.Set("Content-Type", "application/json")
+	if v := req2.Header.Get("Content-Type"); v != "application/json" {
+		t.Fatalf("expected managed Content-Type to win, got %q", v)
+	}
+}