Fix agent init: container resource management and stale UX text

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

@@ -67,6 +67,7 @@ type InitAction struct {
 	models        *modelSelector
 
 	deploymentDetails []project.Deployment
+	containerSettings *project.ContainerSettings
 	httpClient        *http.Client
 }
 
@@ -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.",

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",

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 {

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"`
 }