From 5b2b8e9668d1e4f63fc26fa61579f8816d3fbd5a Mon Sep 17 00:00:00 2001 From: John Miller Date: Thu, 9 Apr 2026 11:01:17 -0400 Subject: [PATCH 1/3] Add container resource management to agent initialization and parsing. FIxes #7599. Fixes #7598. --- .../azure.ai.agents/internal/cmd/init.go | 44 ++++++-- .../pkg/agents/agent_yaml/parse_test.go | 103 ++++++++++++++++++ .../internal/pkg/agents/agent_yaml/yaml.go | 7 ++ 3 files changed, 145 insertions(+), 9 deletions(-) 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 ffbf0ffb163..e972cd13adf 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go @@ -66,8 +66,9 @@ type InitAction struct { flags *initFlags models *modelSelector - deploymentDetails []project.Deployment - httpClient *http.Client + deploymentDetails []project.Deployment + containerSettings *project.ContainerSettings + httpClient *http.Client } // modelSelector encapsulates the dependencies needed for model selection and @@ -475,6 +476,23 @@ func (a *InitAction) Run(ctx context.Context) error { return fmt.Errorf("configuring model choice: %w", err) } + // For hosted agents, prompt for container resources before writing agent.yaml + // so the selected values are persisted into the definition file. + if hostedAgent, ok := agentManifest.Template.(agent_yaml.ContainerAgent); ok { + containerSettings, err := a.populateContainerSettings(ctx, hostedAgent.Resources) + if err != nil { + return fmt.Errorf("failed to populate container settings: %w", err) + } + a.containerSettings = containerSettings + + // Update the agent definition with the selected resources + hostedAgent.Resources = &agent_yaml.ContainerResources{ + Cpu: containerSettings.Resources.Cpu, + Memory: containerSettings.Resources.Memory, + } + agentManifest.Template = hostedAgent + } + // Write the final agent.yaml to disk (after deployment names have been injected) if err := writeAgentDefinitionFile(targetDir, agentManifest); err != nil { return fmt.Errorf("writing agent definition: %w", err) @@ -1226,12 +1244,8 @@ func (a *InitAction) addToProject(ctx context.Context, targetDir string, agentMa } } - // Prompt user for container settings - containerSettings, err := a.populateContainerSettings(ctx) - if err != nil { - return fmt.Errorf("failed to populate container settings: %w", err) - } - agentConfig.Container = containerSettings + // Use container settings that were already populated before writing agent.yaml + agentConfig.Container = a.containerSettings } agentConfig.Deployments = a.deploymentDetails @@ -1282,7 +1296,10 @@ func (a *InitAction) addToProject(ctx context.Context, targetDir string, agentMa return nil } -func (a *InitAction) populateContainerSettings(ctx context.Context) (*project.ContainerSettings, error) { +func (a *InitAction) populateContainerSettings( + ctx context.Context, + manifestResources *agent_yaml.ContainerResources, +) (*project.ContainerSettings, error) { choices := make([]*azdext.SelectChoice, len(project.ResourceTiers)) for i, t := range project.ResourceTiers { choices[i] = &azdext.SelectChoice{ @@ -1292,6 +1309,15 @@ func (a *InitAction) populateContainerSettings(ctx context.Context) (*project.Co } defaultIndex := int32(0) + if manifestResources != nil { + for i, t := range project.ResourceTiers { + if t.Cpu == manifestResources.Cpu && t.Memory == manifestResources.Memory { + defaultIndex = int32(i) + break + } + } + } + resp, err := a.azdClient.Prompt().Select(ctx, &azdext.SelectRequest{ Options: &azdext.SelectOptions{ Message: "Select container resource allocation (CPU and Memory) for your agent. You can adjust these settings later in the azure.yaml file if needed.", diff --git a/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/agent_yaml/parse_test.go b/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/agent_yaml/parse_test.go index 3679d528713..ce891d49f25 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/agent_yaml/parse_test.go +++ b/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/agent_yaml/parse_test.go @@ -6,6 +6,8 @@ package agent_yaml import ( "strings" "testing" + + "go.yaml.in/yaml/v3" ) // TestExtractAgentDefinition_WithTemplateField tests parsing YAML with a template field (manifest format) @@ -40,6 +42,107 @@ template: } } +// TestExtractAgentDefinition_WithResources tests that resources (cpu/memory) are parsed and round-tripped +func TestExtractAgentDefinition_WithResources(t *testing.T) { + yamlContent := []byte(` +name: echo-agent +template: + kind: hosted + name: echo-agent + description: A simple echo agent + protocols: + - protocol: invocations + version: 1.0.0 + resources: + cpu: "0.25" + memory: 0.5Gi +`) + + agent, err := ExtractAgentDefinition(yamlContent) + if err != nil { + t.Fatalf("ExtractAgentDefinition failed: %v", err) + } + + containerAgent, ok := agent.(ContainerAgent) + if !ok { + t.Fatalf("Expected ContainerAgent, got %T", agent) + } + + if containerAgent.Resources == nil { + t.Fatal("Expected Resources to be set, got nil") + } + if containerAgent.Resources.Cpu != "0.25" { + t.Errorf("Expected cpu '0.25', got '%s'", containerAgent.Resources.Cpu) + } + if containerAgent.Resources.Memory != "0.5Gi" { + t.Errorf("Expected memory '0.5Gi', got '%s'", containerAgent.Resources.Memory) + } + + // Verify YAML round-trip: marshal the agent and check resources are preserved + marshaled, err := yaml.Marshal(containerAgent) + if err != nil { + t.Fatalf("Failed to marshal ContainerAgent: %v", err) + } + + marshaledStr := string(marshaled) + if !strings.Contains(marshaledStr, "cpu:") || !strings.Contains(marshaledStr, "memory:") { + t.Errorf("Marshaled YAML should contain cpu and memory under resources, got:\n%s", marshaledStr) + } + + // Unmarshal back and verify + var roundTripped ContainerAgent + if err := yaml.Unmarshal(marshaled, &roundTripped); err != nil { + t.Fatalf("Failed to unmarshal ContainerAgent: %v", err) + } + + if roundTripped.Resources == nil { + t.Fatal("Expected Resources after round-trip, got nil") + } + if roundTripped.Resources.Cpu != "0.25" { + t.Errorf("Round-trip cpu: expected '0.25', got '%s'", roundTripped.Resources.Cpu) + } + if roundTripped.Resources.Memory != "0.5Gi" { + t.Errorf("Round-trip memory: expected '0.5Gi', got '%s'", roundTripped.Resources.Memory) + } +} + +// TestExtractAgentDefinition_WithoutResources tests that ContainerAgent without resources still parses +func TestExtractAgentDefinition_WithoutResources(t *testing.T) { + yamlContent := []byte(` +name: test-manifest +template: + kind: hosted + name: test-agent + protocols: + - protocol: responses + version: v1 +`) + + agent, err := ExtractAgentDefinition(yamlContent) + if err != nil { + t.Fatalf("ExtractAgentDefinition failed: %v", err) + } + + containerAgent, ok := agent.(ContainerAgent) + if !ok { + t.Fatalf("Expected ContainerAgent, got %T", agent) + } + + if containerAgent.Resources != nil { + t.Errorf("Expected Resources to be nil when not specified, got %+v", containerAgent.Resources) + } + + // Verify marshaling omits resources when nil + marshaled, err := yaml.Marshal(containerAgent) + if err != nil { + t.Fatalf("Failed to marshal ContainerAgent: %v", err) + } + + if strings.Contains(string(marshaled), "resources:") { + t.Errorf("Marshaled YAML should not contain resources when nil, got:\n%s", string(marshaled)) + } +} + // TestExtractAgentDefinition_EmptyTemplateField tests that an empty or null template field returns an error func TestExtractAgentDefinition_EmptyTemplateField(t *testing.T) { testCases := []struct { diff --git a/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/agent_yaml/yaml.go b/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/agent_yaml/yaml.go index bb9231c7651..52aa1220308 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/agent_yaml/yaml.go +++ b/cli/azd/extensions/azure.ai.agents/internal/pkg/agents/agent_yaml/yaml.go @@ -98,12 +98,19 @@ type Workflow struct { Trigger *map[string]any `json:"trigger,omitempty" yaml:"trigger,omitempty"` } +// ContainerResources represents the resource allocation for a containerized agent. +type ContainerResources struct { + Cpu string `json:"cpu" yaml:"cpu"` + Memory string `json:"memory" yaml:"memory"` +} + // ContainerAgent This represents a container based agent hosted by the provider/publisher. // The intent is to represent a container application that the user wants to run // in a hosted environment that the provider manages. type ContainerAgent struct { AgentDefinition `json:",inline" yaml:",inline"` Protocols []ProtocolVersionRecord `json:"protocols" yaml:"protocols"` + Resources *ContainerResources `json:"resources,omitempty" yaml:"resources,omitempty"` EnvironmentVariables *[]EnvironmentVariable `json:"environmentVariables,omitempty" yaml:"environment_variables,omitempty"` } From 81467860e107a462dc727f5b8d7e3dd7f5629699 Mon Sep 17 00:00:00 2001 From: John Miller Date: Thu, 9 Apr 2026 11:02:21 -0400 Subject: [PATCH 2/3] fix: remove replica mention from the suggestion. Fixes #7605. --- .../extensions/azure.ai.agents/internal/cmd/init_from_code.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/init_from_code.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/init_from_code.go index 166f02b7601..1de04f11bbd 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/init_from_code.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/init_from_code.go @@ -128,7 +128,7 @@ func (a *InitFromCodeAction) Run(ctx context.Context) error { fmt.Printf(" %s %s\n", color.GreenString("+"), color.GreenString("%s/agent.yaml", srcDir)) } - fmt.Println("\nYou can customize environment variables, cpu, memory, and replica settings in the agent.yaml.") + fmt.Println("\nYou can customize environment variables and other settings in the agent.yaml.") if projectID, _ := a.azdClient.Environment().GetValue(ctx, &azdext.GetEnvRequest{ EnvName: a.environment.Name, Key: "AZURE_AI_PROJECT_ID", From 22fe3d1d49cac1729ca587ac814dab6ea34b78c3 Mon Sep 17 00:00:00 2001 From: John Miller Date: Thu, 9 Apr 2026 11:44:55 -0400 Subject: [PATCH 3/3] style: format initialization struct fields for better readability --- cli/azd/extensions/azure.ai.agents/internal/cmd/init.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 e972cd13adf..45c9c19d18f 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go @@ -66,9 +66,9 @@ type InitAction struct { flags *initFlags models *modelSelector - deploymentDetails []project.Deployment - containerSettings *project.ContainerSettings - httpClient *http.Client + deploymentDetails []project.Deployment + containerSettings *project.ContainerSettings + httpClient *http.Client } // modelSelector encapsulates the dependencies needed for model selection and