diff --git a/cli/azd/extensions/azure.ai.agents/internal/connections/cmd/connection.go b/cli/azd/extensions/azure.ai.agents/internal/connections/cmd/connection.go index a6d8639ce26..5a55daa7a6c 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/connections/cmd/connection.go +++ b/cli/azd/extensions/azure.ai.agents/internal/connections/cmd/connection.go @@ -199,6 +199,9 @@ type connectionCreateFlags struct { metadata []string force bool projectEndpoint string + clientID string // OAuth2 client ID + clientSecret string // OAuth2 client secret + audience string // Token audience for user-entra-token / agentic-identity } // ConnectionCreateAction implements connection creation. @@ -236,6 +239,28 @@ func (a *ConnectionCreateAction) Run(ctx context.Context) error { "Specify at least one custom key (e.g., --custom-key x-api-key=value).", ) } + if a.flags.authType == "oauth2" && (a.flags.clientID == "" || a.flags.clientSecret == "") { + return exterrors.Validation( + exterrors.CodeMissingConnectionField, + "Missing required flags --client-id and --client-secret for oauth2 auth.", + "Specify both OAuth2 client credentials.", + ) + } + if a.flags.authType != "oauth2" && (a.flags.clientID != "" || a.flags.clientSecret != "") { + return exterrors.Validation( + exterrors.CodeConflictingArguments, + "--client-id and --client-secret are only valid with --auth-type oauth2.", + "", + ) + } + if a.flags.audience != "" && a.flags.authType != "user-entra-token" && + a.flags.authType != "agentic-identity" { + return exterrors.Validation( + exterrors.CodeConflictingArguments, + "--audience is only valid with --auth-type user-entra-token or agentic-identity.", + "", + ) + } connCtx, err := resolveConnectionContext(ctx, a.flags.projectEndpoint) if err != nil { @@ -256,21 +281,37 @@ func (a *ConnectionCreateAction) Run(ctx context.Context) error { } } - body, err := buildConnectionBody( - a.flags.kind, a.flags.target, a.flags.authType, - a.flags.key, a.flags.customKeys, a.flags.metadata, - ) - if err != nil { - return err + // Route to raw REST or typed SDK based on auth type + switch a.flags.authType { + case "user-entra-token", "project-managed-identity", "agentic-identity": + err = rawCreateConnection( + ctx, connCtx, + a.flags.name, + rawConnectionProperties{ + AuthType: normalizeAuthTypeToARM(a.flags.authType), + Category: normalizeKind(a.flags.kind), + Target: a.flags.target, + Audience: a.flags.audience, + Metadata: parseKVMap(a.flags.metadata), + }, + ) + default: + body, buildErr := buildConnectionBody( + a.flags.kind, a.flags.target, a.flags.authType, + a.flags.key, a.flags.customKeys, a.flags.metadata, + a.flags.clientID, a.flags.clientSecret, + ) + if buildErr != nil { + return buildErr + } + _, err = connCtx.armClient.Create( + ctx, connCtx.rg, connCtx.account, connCtx.project, + a.flags.name, + &armcognitiveservices.ProjectConnectionsClientCreateOptions{ + Connection: body, + }, + ) } - - _, err = connCtx.armClient.Create( - ctx, connCtx.rg, connCtx.account, connCtx.project, - a.flags.name, - &armcognitiveservices.ProjectConnectionsClientCreateOptions{ - Connection: body, - }, - ) if err != nil { return exterrors.ServiceFromAzure(err, exterrors.OpCreateConnection) } @@ -305,11 +346,12 @@ func newConnectionCreateCommand(extCtx *azdext.ExtensionContext) *cobra.Command } cmd.Flags().StringVar(&flags.kind, "kind", "", - "Connection kind (e.g., remote-tool, cognitive-search)") + "Connection kind (e.g., remote-tool, remote-a2a, cognitive-search)") cmd.Flags().StringVar(&flags.target, "target", "", "Target URL or ARM resource ID") cmd.Flags().StringVar(&flags.authType, "auth-type", "none", - "Auth type: api-key, custom-keys, none") + "Auth type: api-key, custom-keys, none, oauth2, user-entra-token, "+ + "project-managed-identity, agentic-identity") cmd.Flags().StringVar(&flags.key, "key", "", "API key (for api-key auth)") cmd.Flags().StringArrayVar(&flags.customKeys, "custom-key", nil, @@ -318,6 +360,12 @@ func newConnectionCreateCommand(extCtx *azdext.ExtensionContext) *cobra.Command "Metadata key=value (repeatable)") cmd.Flags().BoolVar(&flags.force, "force", false, "Replace existing connection (upsert)") + cmd.Flags().StringVar(&flags.clientID, "client-id", "", + "OAuth2 client ID (required for oauth2 auth)") + cmd.Flags().StringVar(&flags.clientSecret, "client-secret", "", + "OAuth2 client secret (required for oauth2 auth)") + cmd.Flags().StringVar(&flags.audience, "audience", "", + "Token audience for user-entra-token/agentic-identity auth") return cmd } @@ -421,21 +469,37 @@ func (a *ConnectionUpdateAction) Run(ctx context.Context) error { credCustomKeys = append(credCustomKeys, k+"="+v) } - body, err := buildConnectionBody( - kindStr, newTarget, normalizedAuth, - credKey, credCustomKeys, metaPairs, - ) - if err != nil { - return err + // Route to raw REST or typed SDK based on auth type + switch normalizedAuth { + case "user-entra-token", "project-managed-identity", "agentic-identity": + // Identity auth types lack ARM SDK structs — update via raw REST + err = rawCreateConnection( + ctx, connCtx, + a.flags.name, + rawConnectionProperties{ + AuthType: normalizeAuthTypeToARM(normalizedAuth), + Category: kindStr, + Target: newTarget, + Metadata: parseKVMap(metaPairs), + }, + ) + default: + body, buildErr := buildConnectionBody( + kindStr, newTarget, normalizedAuth, + credKey, credCustomKeys, metaPairs, + "", "", + ) + if buildErr != nil { + return buildErr + } + _, err = connCtx.armClient.Create( + ctx, connCtx.rg, connCtx.account, connCtx.project, + a.flags.name, + &armcognitiveservices.ProjectConnectionsClientCreateOptions{ + Connection: body, + }, + ) } - - _, err = connCtx.armClient.Create( - ctx, connCtx.rg, connCtx.account, connCtx.project, - a.flags.name, - &armcognitiveservices.ProjectConnectionsClientCreateOptions{ - Connection: body, - }, - ) if err != nil { return exterrors.ServiceFromAzure(err, exterrors.OpUpdateConnection) } @@ -629,6 +693,7 @@ func buildCredentialReferences( func buildConnectionBody( kind, target, authType, key string, customKeys, metadata []string, + clientID, clientSecret string, ) (*armcognitiveservices.ConnectionPropertiesV2BasicResource, error) { metaMap := parseKVPtrMap(metadata) cat := armcognitiveservices.ConnectionCategory(normalizeKind(kind)) @@ -671,11 +736,32 @@ func buildConnectionBody( }, }, nil + case "oauth2": + at := armcognitiveservices.ConnectionAuthTypeOAuth2 + creds := &armcognitiveservices.ConnectionOAuth2{} + if clientID != "" { + creds.ClientID = &clientID + } + if clientSecret != "" { + creds.ClientSecret = &clientSecret + } + return &armcognitiveservices.ConnectionPropertiesV2BasicResource{ + Properties: &armcognitiveservices.OAuth2AuthTypeConnectionProperties{ + AuthType: &at, + Category: &cat, + Target: &target, + Credentials: creds, + Metadata: metaMap, + }, + }, nil + default: return nil, exterrors.Validation( exterrors.CodeInvalidAuthType, fmt.Sprintf("Unsupported auth type %q.", authType), - "Supported: api-key, custom-keys, none", + "Supported: api-key, custom-keys, none, oauth2. "+ + "For identity-based auth types (user-entra-token, project-managed-identity, "+ + "agentic-identity), use 'connection create' directly.", ) } } @@ -772,6 +858,7 @@ func authTypeStr(a *armcognitiveservices.ConnectionAuthType) string { func normalizeKind(cliKind string) string { mapping := map[string]string{ "remote-tool": "RemoteTool", + "remote-a2a": "RemoteA2A", "cognitive-search": "CognitiveSearch", "api-key": "ApiKey", "app-insights": "AppInsights", @@ -795,7 +882,30 @@ func normalizeAuthType(armAuthType string) string { return "custom-keys" case "None": return "none" + case "OAuth2": + return "oauth2" + case "UserEntraToken": + return "user-entra-token" + case "ProjectManagedIdentity": + return "project-managed-identity" + case "AgenticIdentityToken": + return "agentic-identity" default: return armAuthType } } + +// normalizeAuthTypeToARM converts CLI kebab-case auth type to the ARM wire format. +// Used for auth types that lack ARM SDK structs and require raw REST. +func normalizeAuthTypeToARM(cliAuthType string) string { + switch cliAuthType { + case "user-entra-token": + return "UserEntraToken" + case "project-managed-identity": + return "ProjectManagedIdentity" + case "agentic-identity": + return "AgenticIdentityToken" // ARM expects "Token" suffix + default: + return cliAuthType + } +} diff --git a/cli/azd/extensions/azure.ai.agents/internal/connections/cmd/connection_test.go b/cli/azd/extensions/azure.ai.agents/internal/connections/cmd/connection_test.go index a840d164389..e6bd829407f 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/connections/cmd/connection_test.go +++ b/cli/azd/extensions/azure.ai.agents/internal/connections/cmd/connection_test.go @@ -4,10 +4,12 @@ package cmd import ( + "encoding/json" "testing" "azureaiagent/internal/connections/pkg/connections" + "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/cognitiveservices/armcognitiveservices/v2" "github.com/stretchr/testify/require" ) @@ -116,6 +118,7 @@ func TestNormalizeKind(t *testing.T) { want string }{ {"remote-tool", "RemoteTool"}, + {"remote-a2a", "RemoteA2A"}, {"cognitive-search", "CognitiveSearch"}, {"api-key", "ApiKey"}, {"app-insights", "AppInsights"}, @@ -145,6 +148,10 @@ func TestNormalizeAuthType(t *testing.T) { {"ApiKey", "api-key"}, {"CustomKeys", "custom-keys"}, {"None", "none"}, + {"OAuth2", "oauth2"}, + {"UserEntraToken", "user-entra-token"}, + {"ProjectManagedIdentity", "project-managed-identity"}, + {"AgenticIdentityToken", "agentic-identity"}, // Unknown — pass through {"AAD", "AAD"}, {"", ""}, @@ -157,6 +164,110 @@ func TestNormalizeAuthType(t *testing.T) { } } +func TestNormalizeAuthTypeToARM(t *testing.T) { + tests := []struct { + input string + want string + }{ + {"user-entra-token", "UserEntraToken"}, + {"project-managed-identity", "ProjectManagedIdentity"}, + {"agentic-identity", "AgenticIdentityToken"}, + {"unknown", "unknown"}, + {"", ""}, + } + + for _, tt := range tests { + t.Run(tt.input, func(t *testing.T) { + require.Equal(t, tt.want, normalizeAuthTypeToARM(tt.input)) + }) + } +} + +func TestBuildConnectionBody_OAuth2(t *testing.T) { + body, err := buildConnectionBody( + "RemoteTool", "https://example.com", "oauth2", + "", nil, nil, + "test-client-id", "test-client-secret", + ) + require.NoError(t, err) + + props, ok := body.Properties.(*armcognitiveservices.OAuth2AuthTypeConnectionProperties) + require.True(t, ok, "expected OAuth2AuthTypeConnectionProperties") + require.Equal(t, armcognitiveservices.ConnectionAuthTypeOAuth2, *props.AuthType) + require.Equal(t, "https://example.com", *props.Target) + require.Equal(t, "test-client-id", *props.Credentials.ClientID) + require.Equal(t, "test-client-secret", *props.Credentials.ClientSecret) +} + +func TestBuildConnectionBody_UnsupportedAuthType(t *testing.T) { + _, err := buildConnectionBody( + "RemoteTool", "https://example.com", "invalid-type", + "", nil, nil, "", "", + ) + require.Error(t, err) + require.Contains(t, err.Error(), "Unsupported auth type") +} + +func TestRawConnectionBody_MarshalJSON(t *testing.T) { + props := rawConnectionProperties{ + AuthType: "UserEntraToken", + Category: "RemoteTool", + Target: "https://example.com", + Audience: "https://mcp.ai.azure.com", + } + body := rawConnectionBody{Properties: props} + data, err := json.Marshal(body) + require.NoError(t, err) + + var parsed map[string]any + require.NoError(t, json.Unmarshal(data, &parsed)) + + p := parsed["properties"].(map[string]any) + require.Equal(t, "UserEntraToken", p["authType"]) + require.Equal(t, "RemoteTool", p["category"]) + require.Equal(t, "https://example.com", p["target"]) + require.Equal(t, "https://mcp.ai.azure.com", p["audience"]) +} + +func TestRawConnectionBody_OmitsEmptyAudience(t *testing.T) { + props := rawConnectionProperties{ + AuthType: "ProjectManagedIdentity", + Category: "RemoteA2A", + Target: "https://example.com", + } + body := rawConnectionBody{Properties: props} + data, err := json.Marshal(body) + require.NoError(t, err) + + var parsed map[string]any + require.NoError(t, json.Unmarshal(data, &parsed)) + + p := parsed["properties"].(map[string]any) + _, hasAudience := p["audience"] + require.False(t, hasAudience, "audience should be omitted when empty") +} + +func TestParseKVMap(t *testing.T) { + tests := []struct { + name string + pairs []string + want map[string]string + }{ + {"nil", nil, nil}, + {"empty", []string{}, nil}, + {"single", []string{"k=v"}, map[string]string{"k": "v"}}, + {"value-with-equals", []string{"k=v=1"}, map[string]string{"k": "v=1"}}, + {"multiple", []string{"a=1", "b=2"}, map[string]string{"a": "1", "b": "2"}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := parseKVMap(tt.pairs) + require.Equal(t, tt.want, got) + }) + } +} + func TestParseKVPtrMap(t *testing.T) { tests := []struct { name string diff --git a/cli/azd/extensions/azure.ai.agents/internal/connections/cmd/context.go b/cli/azd/extensions/azure.ai.agents/internal/connections/cmd/context.go index 47da15e6321..5bf382802cb 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/connections/cmd/context.go +++ b/cli/azd/extensions/azure.ai.agents/internal/connections/cmd/context.go @@ -25,6 +25,8 @@ type connectionContext struct { rg string account string project string + sub string // subscription ID for raw REST calls + cred azcore.TokenCredential // credential for raw REST calls } // resolveConnectionContext resolves the project endpoint, discovers ARM context, @@ -71,6 +73,8 @@ func resolveConnectionContext( rg: armCtx.ResourceGroup, account: account, project: project, + sub: armCtx.SubscriptionID, + cred: cred, }, nil } diff --git a/cli/azd/extensions/azure.ai.agents/internal/connections/cmd/raw_connection.go b/cli/azd/extensions/azure.ai.agents/internal/connections/cmd/raw_connection.go new file mode 100644 index 00000000000..669742c0295 --- /dev/null +++ b/cli/azd/extensions/azure.ai.agents/internal/connections/cmd/raw_connection.go @@ -0,0 +1,110 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package cmd + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log" + "net/http" + "net/url" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy" + "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime" +) + +// rawConnectionProperties represents the JSON body for connection PUT requests +// that use auth types not covered by the ARM Go SDK (e.g., UserEntraToken, +// ProjectManagedIdentity, AgenticIdentityToken). +type rawConnectionProperties struct { + AuthType string `json:"authType"` + Category string `json:"category"` + Target string `json:"target"` + Metadata map[string]string `json:"metadata,omitempty"` + Audience string `json:"audience,omitempty"` +} + +type rawConnectionBody struct { + Properties rawConnectionProperties `json:"properties"` +} + +// rawCreateConnection performs a PUT to the ARM connections endpoint using raw REST, +// bypassing the typed ARM SDK. Used for auth types like UserEntraToken, +// ProjectManagedIdentity, and AgenticIdentityToken that lack SDK structs. +func rawCreateConnection( + ctx context.Context, + connCtx *connectionContext, + name string, + props rawConnectionProperties, +) error { + apiVersion := "2025-04-01-preview" + armURL := fmt.Sprintf( + "https://management.azure.com/subscriptions/%s/resourceGroups/%s/"+ + "providers/Microsoft.CognitiveServices/accounts/%s/projects/%s/"+ + "connections/%s?api-version=%s", + connCtx.sub, connCtx.rg, connCtx.account, connCtx.project, + url.PathEscape(name), apiVersion, + ) + + body := rawConnectionBody{Properties: props} + bodyBytes, err := json.Marshal(body) + if err != nil { + return fmt.Errorf("failed to marshal connection body: %w", err) + } + + pipeline := runtime.NewPipeline("azd-connection-raw", "1.0.0", + runtime.PipelineOptions{ + PerCall: []policy.Policy{ + runtime.NewBearerTokenPolicy(connCtx.cred, + []string{"https://management.azure.com/.default"}, nil), + }, + }, + &policy.ClientOptions{}, + ) + + req, err := runtime.NewRequest(ctx, http.MethodPut, armURL) + if err != nil { + return fmt.Errorf("failed to create request: %w", err) + } + req.Raw().Header.Set("Content-Type", "application/json") + req.Raw().Body = io.NopCloser(bytes.NewReader(bodyBytes)) + req.Raw().ContentLength = int64(len(bodyBytes)) + + resp, err := pipeline.Do(req) + if err != nil { + return fmt.Errorf("ARM request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return runtime.NewResponseError(resp) + } + + return nil +} + +// parseKVMap parses "key=value" pairs into a map[string]string. +func parseKVMap(pairs []string) map[string]string { + if len(pairs) == 0 { + return nil + } + result := make(map[string]string, len(pairs)) + for _, pair := range pairs { + found := false + for i := range len(pair) { + if pair[i] == '=' { + result[pair[:i]] = pair[i+1:] + found = true + break + } + } + if !found { + log.Printf("warning: ignoring malformed key=value pair: %q", pair) + } + } + return result +}