feat: add azd ai agent connection commands + credential resolution in run #8174
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,184 @@ | ||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
| // Licensed under the MIT License. | ||
|
|
||
| package cmd | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
| "log" | ||
| "os" | ||
| "path/filepath" | ||
| "regexp" | ||
| "strings" | ||
|
|
||
| "azureaiagent/internal/connections/pkg/connections" | ||
| "azureaiagent/internal/pkg/agents/agent_yaml" | ||
|
|
||
| "github.com/Azure/azure-sdk-for-go/sdk/azidentity" | ||
| "gopkg.in/yaml.v3" | ||
| ) | ||
|
|
||
| // connectionRefPattern matches ${{connections.<name>.credentials.<key>}} references | ||
| // in agent manifest environment variable values. | ||
| var connectionRefPattern = regexp.MustCompile(`\$\{\{connections\.([^.]+)\.credentials\.([^}]+)\}\}`) | ||
|
|
||
| // resolveConnectionCredentials reads the agent manifest from projectDir, | ||
| // scans environment_variables for ${{connections.<name>.credentials.<key>}} patterns, | ||
| // fetches credential values from the Foundry data plane, and returns them as | ||
| // KEY=VALUE strings ready to inject into the agent process environment. | ||
| // | ||
| // This is additive to existing env var handling in run.go: | ||
| // - ${VAR} references are already resolved via loadAzdEnvironment | ||
| // - ${{connections...}} references are resolved here via data-plane API | ||
| // - Literal values pass through unchanged | ||
| // | ||
| // Returns nil (no error) if no manifest is found, no env vars are declared, | ||
| // or no connection references are present — the agent still starts normally. | ||
| func resolveConnectionCredentials( | ||
| ctx context.Context, | ||
| projectDir string, | ||
| endpoint string, | ||
| ) ([]string, error) { | ||
| if endpoint == "" { | ||
| return nil, nil | ||
| } | ||
|
|
||
| // Find and parse the agent manifest | ||
| manifestPath := findManifestInDir(projectDir) | ||
| if manifestPath == "" { | ||
| return nil, nil | ||
| } | ||
|
|
||
| manifestBytes, err := os.ReadFile(manifestPath) | ||
| if err != nil { | ||
| log.Printf("run: could not read manifest %s: %v", manifestPath, err) | ||
| return nil, nil | ||
| } | ||
|
|
||
| // Try parsing as AgentManifest (agent.manifest.yaml — has "template:" wrapper) | ||
| var envVars []agent_yaml.EnvironmentVariable | ||
|
|
||
| manifest, err := agent_yaml.LoadAndValidateAgentManifest(manifestBytes) | ||
| if err == nil { | ||
| if containerAgent, ok := manifest.Template.(agent_yaml.ContainerAgent); ok && | ||
| containerAgent.EnvironmentVariables != nil { | ||
| envVars = *containerAgent.EnvironmentVariables | ||
| } | ||
| } | ||
|
|
||
| // Fall back to parsing as ContainerAgent directly (agent.yaml — no wrapper) | ||
| if len(envVars) == 0 { | ||
| var agentDef agent_yaml.ContainerAgent | ||
| if yamlErr := yaml.Unmarshal(manifestBytes, &agentDef); yamlErr == nil && | ||
| agentDef.EnvironmentVariables != nil { | ||
| envVars = *agentDef.EnvironmentVariables | ||
| } | ||
|
Nathandrake229 marked this conversation as resolved.
|
||
| } | ||
|
|
||
| if len(envVars) == 0 { | ||
| return nil, nil | ||
| } | ||
|
|
||
| // Scan for connection references | ||
| type connRef struct { | ||
| envName string // the env var name (e.g., TAVILY_API_KEY) | ||
| connName string // connection name (e.g., my-test-conn) | ||
| credKey string // credential key (e.g., x-api-key) | ||
| } | ||
|
|
||
| var refs []connRef | ||
| for _, ev := range envVars { | ||
| matches := connectionRefPattern.FindStringSubmatch(ev.Value) | ||
| if matches != nil { | ||
| refs = append(refs, connRef{ | ||
| envName: ev.Name, | ||
| connName: matches[1], | ||
| credKey: matches[2], | ||
| }) | ||
| } | ||
| } | ||
|
|
||
| if len(refs) == 0 { | ||
| return nil, nil | ||
| } | ||
|
|
||
| // Create data-plane credential and client | ||
| cred, err := azidentity.NewAzureDeveloperCLICredential( | ||
| &azidentity.AzureDeveloperCLICredentialOptions{}, | ||
| ) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to create credential for connection resolution: %w", err) | ||
| } | ||
|
|
||
| dpClient := connections.NewDataClient(endpoint, cred) | ||
|
|
||
| // Resolve each reference, caching per connection name | ||
| connCache := map[string]*connections.Connection{} | ||
| var result []string | ||
|
|
||
| for _, ref := range refs { | ||
| conn, cached := connCache[ref.connName] | ||
| if !cached { | ||
| conn, err = dpClient.GetConnectionWithCredentials(ctx, ref.connName) | ||
| if err != nil { | ||
| return nil, fmt.Errorf( | ||
| "failed to resolve credential for %s (connection %q): %w", | ||
| ref.envName, ref.connName, err, | ||
| ) | ||
| } | ||
| connCache[ref.connName] = conn | ||
| } | ||
|
|
||
| // Look up the credential key | ||
| var credValue string | ||
| if ref.credKey == "key" && conn.Credentials != nil && conn.Credentials.Key != "" { | ||
| credValue = conn.Credentials.Key | ||
| } else if conn.Credentials != nil { | ||
| if v, ok := conn.Credentials.CustomKeys[ref.credKey]; ok { | ||
| credValue = v | ||
| } | ||
| } | ||
|
|
||
| if credValue == "" { | ||
| return nil, fmt.Errorf( | ||
| "credential key %q not found on connection %q (for env var %s)", | ||
| ref.credKey, ref.connName, ref.envName, | ||
| ) | ||
| } | ||
|
|
||
| result = append(result, fmt.Sprintf("%s=%s", ref.envName, credValue)) | ||
| // Log the key name only — NEVER log the value | ||
| log.Printf("run: resolved connection credential: %s (connection: %s, key: %s)", | ||
| ref.envName, ref.connName, ref.credKey) | ||
| } | ||
|
|
||
| if len(result) > 0 { | ||
| fmt.Fprintf(os.Stderr, " %d connection credential(s) resolved\n", len(result)) | ||
| } | ||
|
|
||
| return result, nil | ||
| } | ||
|
|
||
| // findManifestInDir looks for an agent manifest or definition file in the given directory. | ||
| // Checks agent.yaml first (the definition the agent app uses), then agent.manifest.yaml. | ||
| // Returns the first file that exists and contains environment_variables with connection references. | ||
| func findManifestInDir(dir string) string { | ||
| // Check agent.yaml first — this is the file the agent app code references | ||
| candidates := []string{ | ||
|
trangevi marked this conversation as resolved.
|
||
| "agent.yaml", | ||
| "agent.manifest.yaml", | ||
| "agent.yml", | ||
| "agent.manifest.yml", | ||
| } | ||
| for _, name := range candidates { | ||
| path := filepath.Join(dir, name) | ||
| if _, err := os.Stat(path); err == nil { | ||
| data, err := os.ReadFile(path) | ||
| if err == nil && strings.Contains(string(data), "${{connections.") { | ||
| return path | ||
| } | ||
| } | ||
| } | ||
| return "" | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.