diff --git a/cli/azd/extensions/azure.ai.agents/go.mod b/cli/azd/extensions/azure.ai.agents/go.mod
index 78453895d53..f236dc2e8b3 100644
--- a/cli/azd/extensions/azure.ai.agents/go.mod
+++ b/cli/azd/extensions/azure.ai.agents/go.mod
@@ -10,13 +10,14 @@ require (
 	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/cognitiveservices/armcognitiveservices v1.8.0
 	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources v1.2.0
 	github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armsubscriptions v1.3.0
-	github.com/azure/azure-dev/cli/azd v0.0.0-20251121010829-d5e0a142e813
+	github.com/azure/azure-dev/cli/azd v0.0.0-20251212003342-848978091314
 	github.com/braydonk/yaml v0.9.0
 	github.com/drone/envsubst v1.0.3
 	github.com/fatih/color v1.18.0
 	github.com/google/uuid v1.6.0
 	github.com/mark3labs/mcp-go v0.41.1
 	github.com/spf13/cobra v1.10.1
+	github.com/spf13/pflag v1.0.10
 	go.yaml.in/yaml/v3 v3.0.4
 	google.golang.org/protobuf v1.36.10
 	gopkg.in/yaml.v3 v3.0.1
@@ -75,7 +76,6 @@ require (
 	github.com/rivo/uniseg v0.4.7 // indirect
 	github.com/sethvargo/go-retry v0.3.0 // indirect
 	github.com/spf13/cast v1.10.0 // indirect
-	github.com/spf13/pflag v1.0.10 // indirect
 	github.com/stretchr/testify v1.11.1 // indirect
 	github.com/theckman/yacspin v0.13.12 // indirect
 	github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
diff --git a/cli/azd/extensions/azure.ai.agents/go.sum b/cli/azd/extensions/azure.ai.agents/go.sum
index deddeec82e7..5045ad2c862 100644
--- a/cli/azd/extensions/azure.ai.agents/go.sum
+++ b/cli/azd/extensions/azure.ai.agents/go.sum
@@ -51,8 +51,8 @@ github.com/aymanbagabas/go-udiff v0.2.0 h1:TK0fH4MteXUDspT88n8CKzvK0X9O2xu9yQjWp
 github.com/aymanbagabas/go-udiff v0.2.0/go.mod h1:RE4Ex0qsGkTAJoQdQQCA0uG+nAzJO/pI/QwceO5fgrA=
 github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk=
 github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4=
-github.com/azure/azure-dev/cli/azd v0.0.0-20251121010829-d5e0a142e813 h1:6RgPxlo9PsEc4q/IDkompYhL7U0+XdW0V4iP+1tpoKc=
-github.com/azure/azure-dev/cli/azd v0.0.0-20251121010829-d5e0a142e813/go.mod h1:k86H7K6vCw8UmimYs0/gDTilxQwXUZDaikRYfDweB/U=
+github.com/azure/azure-dev/cli/azd v0.0.0-20251212003342-848978091314 h1:2COt/tcJlZauO+Vd47SGD//isdVqSj2K1DhMfa3J3Vo=
+github.com/azure/azure-dev/cli/azd v0.0.0-20251212003342-848978091314/go.mod h1:9+M/plQRg5MGyLdTOm8MMxgKohlUdBF04pzZrIugmPs=
 github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=
 github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg=
 github.com/benbjohnson/clock v1.3.5 h1:VvXlSJBzZpA/zum6Sj74hxwYI2DIxRWuNIoXAzHZz5o=
diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/debug.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/debug.go
new file mode 100644
index 00000000000..c207f0f5ded
--- /dev/null
+++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/debug.go
@@ -0,0 +1,44 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package cmd
+
+import (
+	"fmt"
+	"os"
+	"regexp"
+	"strconv"
+	"time"
+
+	azcorelog "github.com/Azure/azure-sdk-for-go/sdk/azcore/log"
+	"github.com/spf13/pflag"
+)
+
+var connectionStringJSONRegex = regexp.MustCompile(`("[\w]*(?:CONNECTION_STRING|ConnectionString)":\s*)"[^"]*"`)
+
+// setupDebugLogging configures the Azure SDK logger if debug mode is enabled.
+func setupDebugLogging(flags *pflag.FlagSet) {
+	if isDebug(flags) {
+		currentDate := time.Now().Format("2006-01-02")
+		logFileName := fmt.Sprintf("azd-ai-agents-%s.log", currentDate)
+
+		logFile, err := os.OpenFile(logFileName, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
+		if err != nil {
+			logFile = os.Stderr
+		}
+		azcorelog.SetListener(func(event azcorelog.Event, msg string) {
+			msg = connectionStringJSONRegex.ReplaceAllString(msg, `${1}"REDACTED"`)
+			fmt.Fprintf(logFile, "[%s] %s: %s\n", time.Now().Format(time.RFC3339), event, msg)
+		})
+	}
+}
+
+// isDebug checks if debug mode is enabled via --debug flag or AZD_EXT_DEBUG environment variable
+func isDebug(flags *pflag.FlagSet) bool {
+	if debugFlag, err := flags.GetBool("debug"); err == nil && debugFlag {
+		return true
+	}
+
+	debug, _ := strconv.ParseBool(os.Getenv("AZD_EXT_DEBUG"))
+	return debug
+}
diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go
index bc371fd39a9..0a2dc9e767e 100644
--- a/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go
+++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go
@@ -91,6 +91,8 @@ func newInitCommand(rootFlags rootFlagsDefinition) *cobra.Command {
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := azdext.WithAccessToken(cmd.Context())
 
+			setupDebugLogging(cmd.Flags())
+
 			azdClient, err := azdext.NewAzdClient()
 			if err != nil {
 				return fmt.Errorf("failed to create azd client: %w", err)
@@ -317,7 +319,7 @@ func ensureEnvironment(ctx context.Context, flags *initFlags, azdClient *azdext.
 		}
 
 		// Create Cognitive Services Projects client
-		projectsClient, err := armcognitiveservices.NewProjectsClient(foundryProject.SubscriptionId, credential, nil)
+		projectsClient, err := armcognitiveservices.NewProjectsClient(foundryProject.SubscriptionId, credential, azure.NewArmClientOptions())
 		if err != nil {
 			return nil, fmt.Errorf("failed to create Cognitive Services Projects client: %w", err)
 		}
@@ -1739,7 +1741,7 @@ func (a *InitAction) getModelDeploymentDetails(ctx context.Context, model agent_
 			accountName = parts[8]   // accounts/{account}
 		}
 
-		deploymentsClient, err := armcognitiveservices.NewDeploymentsClient(subscription, a.credential, nil)
+		deploymentsClient, err := armcognitiveservices.NewDeploymentsClient(subscription, a.credential, azure.NewArmClientOptions())
 		if err != nil {
 			return nil, fmt.Errorf("failed to create deployments client: %w", err)
 		}
diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/listen.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/listen.go
index 9c195b9c932..46d26f73a58 100644
--- a/cli/azd/extensions/azure.ai.agents/internal/cmd/listen.go
+++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/listen.go
@@ -9,19 +9,14 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
-	"strconv"
 	"strings"
-	"time"
 
 	"azureaiagent/internal/pkg/agents/agent_yaml"
-	"azureaiagent/internal/pkg/azure"
 	"azureaiagent/internal/project"
 
-	azcorelog "github.com/Azure/azure-sdk-for-go/sdk/azcore/log"
 	"github.com/azure/azure-dev/cli/azd/pkg/azdext"
 	"github.com/braydonk/yaml"
 	"github.com/spf13/cobra"
-	"github.com/spf13/pflag"
 	"google.golang.org/protobuf/types/known/structpb"
 )
 
@@ -34,19 +29,7 @@ func newListenCommand() *cobra.Command {
 			// Create a new context that includes the AZD access token.
 			ctx := azdext.WithAccessToken(cmd.Context())
 
-			if isDebug(cmd.Flags()) {
-				currentDate := time.Now().Format("2006-01-02")
-				logFileName := fmt.Sprintf("azd-ai-agents-%s.log", currentDate)
-
-				logFile, err := os.OpenFile(logFileName, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
-				if err != nil {
-					logFile = os.Stderr
-				}
-				azcorelog.SetListener(func(event azcorelog.Event, msg string) {
-					msg = azure.ConnectionStringJSONRegex.ReplaceAllString(msg, `${1}"REDACTED"`)
-					fmt.Fprintf(logFile, "[%s] %s: %s\n", time.Now().Format(time.RFC3339), event, msg)
-				})
-			}
+			setupDebugLogging(cmd.Flags())
 
 			// Create a new AZD client.
 			azdClient, err := azdext.NewAzdClient()
@@ -334,13 +317,3 @@ func populateContainerSettings(ctx context.Context, azdClient *azdext.AzdClient,
 
 	return nil
 }
-
-// isDebug checks if debug mode is enabled via --debug flag or AZD_EXT_DEBUG environment variable
-func isDebug(flags *pflag.FlagSet) bool {
-	if debugFlag, err := flags.GetBool("debug"); err == nil && debugFlag {
-		return true
-	}
-
-	debug, _ := strconv.ParseBool(os.Getenv("AZD_EXT_DEBUG"))
-	return debug
-}
diff --git a/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/agent_api/operations.go b/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/agent_api/operations.go
index f5763244a5b..de07dc25ed5 100644
--- a/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/agent_api/operations.go
+++ b/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/agent_api/operations.go
@@ -34,6 +34,9 @@ func NewAgentClient(endpoint string, cred azcore.TokenCredential) *AgentClient {
 
 	clientOptions := &policy.ClientOptions{
 		Logging: policy.LogOptions{
+			AllowedHeaders: []string{"X-Ms-Correlation-Request-Id", "X-Request-Id"},
+			// Include request/response bodies in logs when debug mode is enabled.
+			// Sensitive data is sanitized in internal/cmd/debug.go.
 			IncludeBody: true,
 		},
 		PerCallPolicies: []policy.Policy{
@@ -71,15 +74,15 @@ func (c *AgentClient) GetAgent(ctx context.Context, agentName, apiVersion string
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusOK) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 200 {
-		return nil, fmt.Errorf("failed to get agent. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var agent AgentObject
 	if err := json.Unmarshal(body, &agent); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -112,15 +115,15 @@ func (c *AgentClient) CreateAgent(ctx context.Context, request *CreateAgentReque
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusOK, http.StatusCreated) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 200 && resp.StatusCode != 201 {
-		return nil, fmt.Errorf("failed to create agent. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var agent AgentObject
 	if err := json.Unmarshal(body, &agent); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -153,15 +156,15 @@ func (c *AgentClient) UpdateAgent(ctx context.Context, agentName string, request
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusOK) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 200 {
-		return nil, fmt.Errorf("failed to update agent. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var agent AgentObject
 	if err := json.Unmarshal(body, &agent); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -185,15 +188,15 @@ func (c *AgentClient) DeleteAgent(ctx context.Context, agentName, apiVersion str
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusOK) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 200 {
-		return nil, fmt.Errorf("failed to delete agent. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var deleteResponse DeleteAgentResponse
 	if err := json.Unmarshal(body, &deleteResponse); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -245,15 +248,15 @@ func (c *AgentClient) ListAgents(ctx context.Context, params *ListAgentQueryPara
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusOK) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 200 {
-		return nil, fmt.Errorf("failed to list agents. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var agentList AgentList
 	if err := json.Unmarshal(body, &agentList); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -286,15 +289,15 @@ func (c *AgentClient) CreateAgentVersion(ctx context.Context, agentName string,
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusOK, http.StatusCreated) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 200 && resp.StatusCode != 201 {
-		return nil, fmt.Errorf("failed to create agent version. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var agentVersion AgentVersionObject
 	if err := json.Unmarshal(body, &agentVersion); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -318,15 +321,15 @@ func (c *AgentClient) GetAgentVersion(ctx context.Context, agentName, agentVersi
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusOK) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 200 {
-		return nil, fmt.Errorf("failed to get agent version. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var version AgentVersionObject
 	if err := json.Unmarshal(body, &version); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -350,15 +353,15 @@ func (c *AgentClient) DeleteAgentVersion(ctx context.Context, agentName, agentVe
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusOK) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 200 {
-		return nil, fmt.Errorf("failed to delete agent version. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var deleteResponse DeleteAgentVersionResponse
 	if err := json.Unmarshal(body, &deleteResponse); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -415,15 +418,15 @@ func (c *AgentClient) ListAgentVersions(ctx context.Context, agentName string, p
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusOK) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 200 {
-		return nil, fmt.Errorf("failed to list agent versions. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var versionList AgentVersionList
 	if err := json.Unmarshal(body, &versionList); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -458,15 +461,15 @@ func (c *AgentClient) CreateOrUpdateAgentEventHandler(ctx context.Context, agent
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusOK, http.StatusCreated) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 200 && resp.StatusCode != 201 {
-		return nil, fmt.Errorf("failed to create/update event handler. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var eventHandler AgentEventHandlerObject
 	if err := json.Unmarshal(body, &eventHandler); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -490,15 +493,15 @@ func (c *AgentClient) GetAgentEventHandler(ctx context.Context, agentName, event
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusOK) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 200 {
-		return nil, fmt.Errorf("failed to get event handler. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var eventHandler AgentEventHandlerObject
 	if err := json.Unmarshal(body, &eventHandler); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -522,15 +525,15 @@ func (c *AgentClient) DeleteAgentEventHandler(ctx context.Context, agentName, ev
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusOK) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 200 {
-		return nil, fmt.Errorf("failed to delete event handler. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var deleteResponse DeleteAgentEventHandlerResponse
 	if err := json.Unmarshal(body, &deleteResponse); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -578,15 +581,15 @@ func (c *AgentClient) StartAgentContainer(ctx context.Context, agentName, agentV
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusAccepted) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 202 {
-		return nil, fmt.Errorf("failed to start agent container. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var operation AgentContainerOperationObject
 	if err := json.Unmarshal(body, &operation); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -632,15 +635,15 @@ func (c *AgentClient) UpdateAgentContainer(ctx context.Context, agentName, agent
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusAccepted) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 202 {
-		return nil, fmt.Errorf("failed to update agent container. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var operation AgentContainerOperationObject
 	if err := json.Unmarshal(body, &operation); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -673,15 +676,15 @@ func (c *AgentClient) StopAgentContainer(ctx context.Context, agentName, agentVe
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusAccepted) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 202 {
-		return nil, fmt.Errorf("failed to stop agent container. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var operation AgentContainerOperationObject
 	if err := json.Unmarshal(body, &operation); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -714,15 +717,15 @@ func (c *AgentClient) DeleteAgentContainer(ctx context.Context, agentName, agent
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusAccepted) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 202 {
-		return nil, fmt.Errorf("failed to delete agent container. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var operation AgentContainerOperationObject
 	if err := json.Unmarshal(body, &operation); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -751,15 +754,15 @@ func (c *AgentClient) GetAgentContainer(ctx context.Context, agentName, agentVer
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusOK) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 200 {
-		return nil, fmt.Errorf("failed to get agent container. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var container AgentContainerObject
 	if err := json.Unmarshal(body, &container); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
@@ -783,15 +786,15 @@ func (c *AgentClient) GetAgentContainerOperation(ctx context.Context, agentName,
 	}
 	defer resp.Body.Close()
 
+	if !runtime.HasStatusCode(resp, http.StatusOK) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	if resp.StatusCode != 200 {
-		return nil, fmt.Errorf("failed to get container operation. Status code: %d, Response: %s", resp.StatusCode, string(body))
-	}
-
 	var operation AgentContainerOperationObject
 	if err := json.Unmarshal(body, &operation); err != nil {
 		return nil, fmt.Errorf("failed to parse response: %w", err)
diff --git a/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/registry_api/operations.go b/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/registry_api/operations.go
index 8c7f7e1827c..6792a8034da 100644
--- a/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/registry_api/operations.go
+++ b/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/registry_api/operations.go
@@ -5,6 +5,7 @@ package registry_api
 
 import (
 	"azureaiagent/internal/pkg/agents/agent_api"
+	"azureaiagent/internal/version"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -13,22 +14,43 @@ import (
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
+	"github.com/azure/azure-dev/cli/azd/pkg/azsdk"
 )
 
 // RegistryAgentManifestClient provides methods to interact with Azure ML registry agent manifests
 type RegistryAgentManifestClient struct {
 	baseEndpoint string
-	cred         azcore.TokenCredential
-	client       *http.Client
+	pipeline     runtime.Pipeline
 }
 
 // NewRegistryAgentManifestClient creates a new instance of RegistryAgentManifestClient
 func NewRegistryAgentManifestClient(registryName string, cred azcore.TokenCredential) *RegistryAgentManifestClient {
 	baseEndpoint := fmt.Sprintf("https://int.api.azureml-test.ms/agent-asset/v1.0/registries/%s/agentManifests", registryName)
+
+	userAgent := fmt.Sprintf("azd-ext-azure-ai-agents/%s", version.Version)
+
+	clientOptions := &policy.ClientOptions{
+		Logging: policy.LogOptions{
+			AllowedHeaders: []string{azsdk.MsCorrelationIdHeader},
+		},
+		PerCallPolicies: []policy.Policy{
+			runtime.NewBearerTokenPolicy(cred, []string{"https://ai.azure.com/.default"}, nil),
+			azsdk.NewMsCorrelationPolicy(),
+			azsdk.NewUserAgentPolicy(userAgent),
+		},
+	}
+
+	pipeline := runtime.NewPipeline(
+		"azure-ai-agents",
+		"v1.0.0",
+		runtime.PipelineOptions{},
+		clientOptions,
+	)
+
 	return &RegistryAgentManifestClient{
 		baseEndpoint: baseEndpoint,
-		cred:         cred,
-		client:       &http.Client{},
+		pipeline:     pipeline,
 	}
 }
 
@@ -36,15 +58,25 @@ func NewRegistryAgentManifestClient(registryName string, cred azcore.TokenCreden
 func (c *RegistryAgentManifestClient) GetManifest(ctx context.Context, manifestName string, manifestVersion string) (*Manifest, error) {
 	targetEndpoint := fmt.Sprintf("%s/%s/versions/%s", c.baseEndpoint, manifestName, manifestVersion)
 
-	req, err := http.NewRequestWithContext(ctx, "GET", targetEndpoint, nil)
+	req, err := runtime.NewRequest(ctx, http.MethodGet, targetEndpoint)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create request: %w", err)
 	}
 
 	fmt.Println("Making HTTP request to retrieve manifest...")
-	body, err := c.makeHTTPRequest(ctx, req)
+	resp, err := c.pipeline.Do(req)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("HTTP request failed: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if !runtime.HasStatusCode(resp, http.StatusOK) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
 	var manifest Manifest
@@ -214,111 +246,30 @@ func HandleTools(manifest *Manifest) ([]any, error) {
 
 // GetAllLatest retrieves all latest agent manifests from the specified registry
 func (c *RegistryAgentManifestClient) GetAllLatest(ctx context.Context) ([]Manifest, error) {
-	req, err := http.NewRequestWithContext(ctx, "GET", c.baseEndpoint, nil)
+	req, err := runtime.NewRequest(ctx, http.MethodGet, c.baseEndpoint)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create request: %w", err)
 	}
 
-	body, err := c.makeHTTPRequest(ctx, req)
+	resp, err := c.pipeline.Do(req)
 	if err != nil {
-		return nil, err
-	}
-
-	var manifestList ManifestList
-	if err := json.Unmarshal(body, &manifestList); err != nil {
-		return nil, fmt.Errorf("failed to unmarshal manifest list response: %w", err)
+		return nil, fmt.Errorf("HTTP request failed: %w", err)
 	}
+	defer resp.Body.Close()
 
-	return manifestList.Value, nil
-}
-
-// Helper methods
-
-// makeHTTPRequest makes an HTTP request with proper authentication and error handling
-func (c *RegistryAgentManifestClient) makeHTTPRequest(ctx context.Context, req *http.Request) ([]byte, error) {
-	// Log the request details - uncomment for debugging
-	// c.logRequest(req.Method, req.URL.String(), nil)
-
-	// Add authentication header
-	if err := c.setAuthHeader(ctx, req); err != nil {
-		return nil, fmt.Errorf("failed to set authentication header: %w", err)
-	}
-
-	// Set common headers
-	req.Header.Set("Content-Type", "application/json")
-	req.Header.Set("Accept", "application/json")
-
-	// Make the HTTP request
-	resp, err := c.client.Do(req)
-	if err != nil {
-		return nil, fmt.Errorf("failed to make HTTP request: %w", err)
+	if !runtime.HasStatusCode(resp, http.StatusOK) {
+		return nil, runtime.NewResponseError(resp)
 	}
-	defer resp.Body.Close()
 
-	// Read response body
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	// Log the response details - uncomment for debugging
-	// c.logResponse(body)
-
-	// Check for HTTP errors
-	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
-		return nil, fmt.Errorf("HTTP request failed with status %d: %s", resp.StatusCode, string(body))
-	}
-
-	return body, nil
-}
-
-// setAuthHeader sets the authorization header using the credential
-func (c *RegistryAgentManifestClient) setAuthHeader(ctx context.Context, req *http.Request) error {
-	token, err := c.getAiFoundryAzureToken(ctx, c.cred)
-	if err != nil {
-		return fmt.Errorf("failed to get Azure token: %w", err)
-	}
-
-	req.Header.Set("Authorization", "Bearer "+token)
-	return nil
-}
-
-// getAiFoundryAzureToken gets an Azure access token using the provided credential
-func (c *RegistryAgentManifestClient) getAiFoundryAzureToken(ctx context.Context, cred azcore.TokenCredential) (string, error) {
-	tokenRequestOptions := policy.TokenRequestOptions{
-		Scopes: []string{"https://ai.azure.com/.default"},
-	}
-
-	token, err := cred.GetToken(ctx, tokenRequestOptions)
-	if err != nil {
-		return "", err
-	}
-
-	return token.Token, nil
-}
-
-// logRequest logs the request details to stderr for debugging
-func (c *RegistryAgentManifestClient) logRequest(method, url string, payload []byte) {
-	fmt.Printf("%s %s\n", method, url)
-	if len(payload) > 0 {
-		var prettyPayload interface{}
-		if err := json.Unmarshal(payload, &prettyPayload); err == nil {
-			prettyJSON, _ := json.MarshalIndent(prettyPayload, "", "  ")
-			fmt.Printf("Payload:\n%s\n", string(prettyJSON))
-		} else {
-			fmt.Printf("Payload: %s\n", string(payload))
-		}
+	var manifestList ManifestList
+	if err := json.Unmarshal(body, &manifestList); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal manifest list response: %w", err)
 	}
-}
 
-// logResponse logs the response body to stderr for debugging
-func (c *RegistryAgentManifestClient) logResponse(body []byte) {
-	fmt.Println("Response:")
-	var jsonResponse interface{}
-	if err := json.Unmarshal(body, &jsonResponse); err == nil {
-		prettyJSON, _ := json.MarshalIndent(jsonResponse, "", "  ")
-		fmt.Println(string(prettyJSON))
-	} else {
-		fmt.Println(string(body))
-	}
+	return manifestList.Value, nil
 }
diff --git a/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/ai/model_catalog.go b/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/ai/model_catalog.go
index a36f3a45b94..6009d108bbd 100644
--- a/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/ai/model_catalog.go
+++ b/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/ai/model_catalog.go
@@ -407,7 +407,7 @@ func createModelsClient(
 	subscriptionId string,
 	credential azcore.TokenCredential,
 ) (*armcognitiveservices.ModelsClient, error) {
-	client, err := armcognitiveservices.NewModelsClient(subscriptionId, credential, nil)
+	client, err := armcognitiveservices.NewModelsClient(subscriptionId, credential, azure.NewArmClientOptions())
 	if err != nil {
 		return nil, err
 	}
diff --git a/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/azure_client.go b/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/azure_client.go
index 2cea9b3fb45..006984a4496 100644
--- a/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/azure_client.go
+++ b/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/azure_client.go
@@ -47,7 +47,7 @@ func (c *AzureClient) ListLocations(ctx context.Context, subscriptionId string)
 }
 
 func createSubscriptionsClient(subscriptionId string, credential azcore.TokenCredential) (*armsubscriptions.Client, error) {
-	client, err := armsubscriptions.NewClient(credential, nil)
+	client, err := armsubscriptions.NewClient(credential, NewArmClientOptions())
 	if err != nil {
 		return nil, err
 	}
diff --git a/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/client_options.go b/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/client_options.go
new file mode 100644
index 00000000000..bc7f7e2fe07
--- /dev/null
+++ b/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/client_options.go
@@ -0,0 +1,31 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azure
+
+import (
+	"azureaiagent/internal/version"
+	"fmt"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	"github.com/azure/azure-dev/cli/azd/pkg/azsdk"
+)
+
+// NewArmClientOptions creates a new arm.ClientOptions with standard policies for Azure SDK clients.
+// This includes correlation headers, user agent, and logging configuration.
+func NewArmClientOptions() *arm.ClientOptions {
+	userAgent := fmt.Sprintf("azd-ext-azure-ai-agents/%s", version.Version)
+
+	return &arm.ClientOptions{
+		ClientOptions: policy.ClientOptions{
+			Logging: policy.LogOptions{
+				AllowedHeaders: []string{azsdk.MsCorrelationIdHeader},
+			},
+			PerCallPolicies: []policy.Policy{
+				azsdk.NewMsCorrelationPolicy(),
+				azsdk.NewUserAgentPolicy(userAgent),
+			},
+		},
+	}
+}
diff --git a/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/foundry_projects_client.go b/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/foundry_projects_client.go
index afd48ed6ded..02b32ed4f98 100644
--- a/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/foundry_projects_client.go
+++ b/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/foundry_projects_client.go
@@ -4,6 +4,7 @@
 package azure
 
 import (
+	"azureaiagent/internal/version"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -12,24 +13,45 @@ import (
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
+	"github.com/azure/azure-dev/cli/azd/pkg/azsdk"
 )
 
 // FoundryProjectsClient provides methods to interact with Microsoft Foundry projects
 type FoundryProjectsClient struct {
 	baseEndpoint string
 	apiVersion   string
-	cred         azcore.TokenCredential
-	client       *http.Client
+	pipeline     runtime.Pipeline
 }
 
 // NewFoundryProjectsClient creates a new instance of FoundryProjectsClient
 func NewFoundryProjectsClient(accountName string, projectName string, cred azcore.TokenCredential) *FoundryProjectsClient {
 	baseEndpoint := fmt.Sprintf("https://%s.services.ai.azure.com/api/projects/%s", accountName, projectName)
+
+	userAgent := fmt.Sprintf("azd-ext-azure-ai-agents/%s", version.Version)
+
+	clientOptions := &policy.ClientOptions{
+		Logging: policy.LogOptions{
+			AllowedHeaders: []string{azsdk.MsCorrelationIdHeader},
+		},
+		PerCallPolicies: []policy.Policy{
+			runtime.NewBearerTokenPolicy(cred, []string{"https://ai.azure.com/.default"}, nil),
+			azsdk.NewMsCorrelationPolicy(),
+			azsdk.NewUserAgentPolicy(userAgent),
+		},
+	}
+
+	pipeline := runtime.NewPipeline(
+		"azure-ai-agents",
+		"v1.0.0",
+		runtime.PipelineOptions{},
+		clientOptions,
+	)
+
 	return &FoundryProjectsClient{
 		baseEndpoint: baseEndpoint,
 		apiVersion:   "2025-11-15-preview",
-		cred:         cred,
-		client:       &http.Client{},
+		pipeline:     pipeline,
 	}
 }
 
@@ -90,14 +112,24 @@ type PagedConnection struct {
 func (c *FoundryProjectsClient) GetPagedConnections(ctx context.Context) (*PagedConnection, error) {
 	targetEndpoint := fmt.Sprintf("%s/connections?api-version=%s", c.baseEndpoint, c.apiVersion)
 
-	req, err := http.NewRequestWithContext(ctx, "GET", targetEndpoint, nil)
+	req, err := runtime.NewRequest(ctx, http.MethodGet, targetEndpoint)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create request: %w", err)
 	}
 
-	body, err := c.makeHTTPRequest(ctx, req)
+	resp, err := c.pipeline.Do(req)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("HTTP request failed: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if !runtime.HasStatusCode(resp, http.StatusOK) {
+		return nil, runtime.NewResponseError(resp)
+	}
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
 	var pagedConnections PagedConnection
@@ -125,21 +157,11 @@ func (c *FoundryProjectsClient) GetAllConnections(ctx context.Context) ([]Connec
 
 	// Continue fetching pages while there's a next link
 	for nextLink != nil && *nextLink != "" {
-		req, err := http.NewRequestWithContext(ctx, "GET", *nextLink, nil)
-		if err != nil {
-			return nil, fmt.Errorf("failed to create request for next page: %w", err)
-		}
-
-		body, err := c.makeHTTPRequest(ctx, req)
+		pagedConnections, err := c.getNextPage(ctx, *nextLink)
 		if err != nil {
 			return nil, err
 		}
 
-		var pagedConnections PagedConnection
-		if err := json.Unmarshal(body, &pagedConnections); err != nil {
-			return nil, fmt.Errorf("failed to unmarshal connections response: %w", err)
-		}
-
 		// Add connections from this page
 		allConnections = append(allConnections, pagedConnections.Value...)
 		nextLink = pagedConnections.NextLink
@@ -148,67 +170,32 @@ func (c *FoundryProjectsClient) GetAllConnections(ctx context.Context) ([]Connec
 	return allConnections, nil
 }
 
-// Helper methods
-
-// makeHTTPRequest makes an HTTP request with proper authentication and error handling
-func (c *FoundryProjectsClient) makeHTTPRequest(ctx context.Context, req *http.Request) ([]byte, error) {
-	// Log the request details - uncomment for debugging
-	// c.logRequest(req.Method, req.URL.String(), nil)
-
-	// Add authentication header
-	if err := c.setAuthHeader(ctx, req); err != nil {
-		return nil, fmt.Errorf("failed to set authentication header: %w", err)
-	}
-
-	// Set common headers
-	req.Header.Set("Content-Type", "application/json")
-	req.Header.Set("Accept", "application/json")
-
-	// Make the HTTP request
-	resp, err := c.client.Do(req)
+// getNextPage fetches a single page of connections from the given URL
+func (c *FoundryProjectsClient) getNextPage(ctx context.Context, url string) (*PagedConnection, error) {
+	req, err := runtime.NewRequest(ctx, http.MethodGet, url)
 	if err != nil {
-		return nil, fmt.Errorf("failed to make HTTP request: %w", err)
+		return nil, fmt.Errorf("failed to create request for next page: %w", err)
 	}
-	defer resp.Body.Close()
 
-	// Read response body
-	body, err := io.ReadAll(resp.Body)
+	resp, err := c.pipeline.Do(req)
 	if err != nil {
-		return nil, fmt.Errorf("failed to read response body: %w", err)
+		return nil, fmt.Errorf("HTTP request failed: %w", err)
 	}
+	defer resp.Body.Close()
 
-	// Log the response details - uncomment for debugging
-	// c.logResponse(body)
-
-	// Check for HTTP errors
-	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
-		return nil, fmt.Errorf("HTTP request failed with status %d: %s", resp.StatusCode, string(body))
+	if !runtime.HasStatusCode(resp, http.StatusOK) {
+		return nil, runtime.NewResponseError(resp)
 	}
 
-	return body, nil
-}
-
-// setAuthHeader sets the authorization header using the credential
-func (c *FoundryProjectsClient) setAuthHeader(ctx context.Context, req *http.Request) error {
-	token, err := c.getAiFoundryAzureToken(ctx, c.cred)
+	body, err := io.ReadAll(resp.Body)
 	if err != nil {
-		return fmt.Errorf("failed to get Azure token: %w", err)
-	}
-
-	req.Header.Set("Authorization", "Bearer "+token)
-	return nil
-}
-
-// getAiFoundryAzureToken gets an Azure access token using the provided credential
-func (c *FoundryProjectsClient) getAiFoundryAzureToken(ctx context.Context, cred azcore.TokenCredential) (string, error) {
-	tokenRequestOptions := policy.TokenRequestOptions{
-		Scopes: []string{"https://ai.azure.com/.default"},
+		return nil, fmt.Errorf("failed to read response body: %w", err)
 	}
 
-	token, err := cred.GetToken(ctx, tokenRequestOptions)
-	if err != nil {
-		return "", err
+	var pagedConnections PagedConnection
+	if err := json.Unmarshal(body, &pagedConnections); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal connections response: %w", err)
 	}
 
-	return token.Token, nil
+	return &pagedConnections, nil
 }
diff --git a/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/logging.go b/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/logging.go
deleted file mode 100644
index fdbd74e3cdb..00000000000
--- a/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/logging.go
+++ /dev/null
@@ -1,8 +0,0 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
-// Licensed under the MIT License.
-
-package azure
-
-import "regexp"
-
-var ConnectionStringJSONRegex = regexp.MustCompile(`("[\w]*(?:CONNECTION_STRING|ConnectionString)":\s*)"[^"]*"`)
diff --git a/cli/azd/extensions/azure.ai.agents/internal/project/service_target_agent.go b/cli/azd/extensions/azure.ai.agents/internal/project/service_target_agent.go
index 1bc1288fa1e..05d49738db1 100644
--- a/cli/azd/extensions/azure.ai.agents/internal/project/service_target_agent.go
+++ b/cli/azd/extensions/azure.ai.agents/internal/project/service_target_agent.go
@@ -15,6 +15,7 @@ import (
 
 	"azureaiagent/internal/pkg/agents/agent_api"
 	"azureaiagent/internal/pkg/agents/agent_yaml"
+	"azureaiagent/internal/pkg/azure"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
@@ -210,7 +211,7 @@ func (p *AgentServiceTargetProvider) GetTargetResource(
 	projectName := p.foundryProject.Name
 
 	// Create Cognitive Services Projects client
-	projectsClient, err := armcognitiveservices.NewProjectsClient(p.foundryProject.SubscriptionID, p.credential, nil)
+	projectsClient, err := armcognitiveservices.NewProjectsClient(p.foundryProject.SubscriptionID, p.credential, azure.NewArmClientOptions())
 	if err != nil {
 		return nil, fmt.Errorf("failed to create Cognitive Services Projects client: %w", err)
 	}
diff --git a/cli/azd/extensions/azure.ai.agents/main.go b/cli/azd/extensions/azure.ai.agents/main.go
index e51dfe1f609..7aa35770d38 100644
--- a/cli/azd/extensions/azure.ai.agents/main.go
+++ b/cli/azd/extensions/azure.ai.agents/main.go
@@ -4,11 +4,11 @@
 package main
 
 import (
-	"context"
 	"os"
 
 	"azureaiagent/internal/cmd"
 
+	"github.com/azure/azure-dev/cli/azd/pkg/azdext"
 	"github.com/fatih/color"
 )
 
@@ -21,7 +21,7 @@ func init() {
 
 func main() {
 	// Execute the root command
-	ctx := context.Background()
+	ctx := azdext.NewContext()
 	rootCmd := cmd.NewRootCommand()
 
 	if err := rootCmd.ExecuteContext(ctx); err != nil {