fix: preserve referenced agent service fields

Author: huimiu

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

@@ -734,6 +734,9 @@ func populateContainerSettings(
 	if err != nil {
 		return err
 	}
+	if persisted == nil {
+		return nil
+	}
 
 	if err := project.AddServiceSerialized(
 		ctx,
@@ -752,14 +755,18 @@ func prepareContainerSettings(
 ) (*azdext.ServiceConfig, error) {
 	rawAdditional := svc.GetAdditionalProperties()
 	rawConfig := svc.GetConfig()
-	if err := project.ResolveServiceConfigInPlace(
-		svc,
-		projectRoot,
-	); err != nil {
-		return nil, fmt.Errorf(
-			"failed to resolve agent config: %w",
-			err,
-		)
+	hasFileRef := project.ConfigContainsFileRef(rawAdditional) ||
+		project.ConfigContainsFileRef(rawConfig)
+	if hasFileRef {
+		if err := project.ResolveServiceConfigInPlace(
+			svc,
+			projectRoot,
+		); err != nil {
+			return nil, fmt.Errorf(
+				"failed to resolve agent config: %w",
+				err,
+			)
+		}
 	}
 	foundryAgentConfig, err := project.LoadServiceTargetAgentConfig(svc)
 	if err != nil {
@@ -785,19 +792,19 @@ func prepareContainerSettings(
 		result.Cpu = project.DefaultCpu
 	}
 
-	persisted := *svc
-	persisted.AdditionalProperties = rawAdditional
-	persisted.Config = rawConfig
 	if err := project.SetAgentContainerSettings(
-		&persisted,
+		svc,
 		&project.ContainerSettings{Resources: result},
 	); err != nil {
 		return nil, fmt.Errorf(
 			"failed to update agent container settings: %w",
 			err,
 		)
 	}
-	return &persisted, nil
+	if hasFileRef {
+		return nil, nil
+	}
+	return svc, nil
 }
 
 // resolveToolboxEnvVars resolves ${VAR} references in toolbox name, description,

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

@@ -131,7 +131,9 @@ func TestIsHostedAgentServiceRejectsTraversal(t *testing.T) {
 	}
 }
 
-func TestPrepareContainerSettings_ResolvesFileRef(t *testing.T) {
+func TestPopulateContainerSettings_DoesNotPersistResolvedFileRef(
+	t *testing.T,
+) {
 	t.Parallel()
 
 	root := t.TempDir()
@@ -158,17 +160,11 @@ func TestPrepareContainerSettings_ResolvesFileRef(t *testing.T) {
 		AdditionalProperties: props,
 	}
 
-	persisted, err := prepareContainerSettings(svc, root)
+	err = populateContainerSettings(t.Context(), nil, svc, root)
 
 	require.NoError(t, err)
 	require.Equal(t, "src/echo", svc.GetRelativePath())
-	require.Equal(t, "src/echo", persisted.GetRelativePath())
-	require.Contains(
-		t,
-		persisted.GetAdditionalProperties().GetFields(),
-		"$ref",
-	)
-	cfg, err := project.LoadServiceTargetAgentConfig(persisted)
+	cfg, err := project.LoadServiceTargetAgentConfig(svc)
 	require.NoError(t, err)
 	require.NotNil(t, cfg.Container)
 	require.NotNil(t, cfg.Container.Resources)

cli/azd/extensions/azure.ai.agents/internal/project/service_target_agent.go

@@ -41,6 +41,7 @@ import (
 	"github.com/azure/azure-dev/cli/azd/pkg/output"
 	"github.com/fatih/color"
 	"github.com/google/uuid"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/structpb"
 )
 
@@ -146,8 +147,7 @@ var _ azdext.ServiceTargetProvider = &AgentServiceTargetProvider{}
 type AgentServiceTargetProvider struct {
 	azdClient           *azdext.AzdClient
 	serviceConfig       *azdext.ServiceConfig
-	sourceProperties    *structpb.Struct
-	sourceConfig        *structpb.Struct
+	sourceServiceConfig *azdext.ServiceConfig
 	agentDefinitionPath string
 	projectPath         string
 	servicePath         string
@@ -185,8 +185,7 @@ func NewAgentServiceTargetProvider(azdClient *azdext.AzdClient) azdext.ServiceTa
 // only when a deploy-time entrypoint needs it.
 func (p *AgentServiceTargetProvider) Initialize(ctx context.Context, serviceConfig *azdext.ServiceConfig) error {
 	p.serviceConfig = serviceConfig
-	p.sourceProperties = serviceConfig.GetAdditionalProperties()
-	p.sourceConfig = serviceConfig.GetConfig()
+	p.sourceServiceConfig = proto.Clone(serviceConfig).(*azdext.ServiceConfig)
 	return nil
 }
 
@@ -545,7 +544,7 @@ func (p *AgentServiceTargetProvider) Package(
 	serviceConfig *azdext.ServiceConfig,
 	serviceContext *azdext.ServiceContext,
 	progress azdext.ProgressReporter,
-) (*azdext.ServicePackageResult, error) {
+) (result *azdext.ServicePackageResult, err error) {
 	if err := p.ensureDeployContext(ctx); err != nil {
 		return nil, err
 	}
@@ -590,12 +589,26 @@ func (p *AgentServiceTargetProvider) Package(
 			Artifacts: []*azdext.Artifact{preBuiltImageArtifact(agentDef.Image)},
 		}, nil
 	}
-	if err := p.persistBuildService(ctx); err != nil {
+	restore, err := p.persistBuildService(ctx)
+	if err != nil {
 		return nil, exterrors.Internal(
 			exterrors.OpContainerBuild,
 			fmt.Sprintf("prepare referenced service: %s", err),
 		)
 	}
+	if restore != nil {
+		defer func() {
+			err = combineBuildServiceRestore(
+				ctx,
+				restore,
+				err,
+				exterrors.OpContainerBuild,
+			)
+			if err != nil {
+				result = nil
+			}
+		}()
+	}
 
 	var packageArtifact *azdext.Artifact
 	var newArtifacts []*azdext.Artifact
@@ -651,25 +664,61 @@ func (p *AgentServiceTargetProvider) Package(
 
 func (p *AgentServiceTargetProvider) persistBuildService(
 	ctx context.Context,
-) error {
-	if !configContainsFileRef(p.sourceProperties) &&
-		!configContainsFileRef(p.sourceConfig) {
-		return nil
-	}
-	persisted := *p.serviceConfig
-	persisted.AdditionalProperties = p.sourceProperties
-	persisted.Config = p.sourceConfig
+) (func(context.Context) error, error) {
+	source := p.sourceServiceConfig
+	if source == nil ||
+		(!ConfigContainsFileRef(source.GetAdditionalProperties()) &&
+			!ConfigContainsFileRef(source.GetConfig())) {
+		return nil, nil
+	}
+	resolvedPath := p.serviceConfig.GetRelativePath()
+	if resolvedPath == source.GetRelativePath() {
+		return nil, nil
+	}
+	persisted := proto.Clone(source).(*azdext.ServiceConfig)
+	persisted.RelativePath = resolvedPath
 	if err := AddServiceSerialized(
 		ctx,
 		p.azdClient,
-		&persisted,
+		persisted,
 	); err != nil {
-		return fmt.Errorf("persist resolved service fields: %w", err)
+		return nil, fmt.Errorf("persist resolved service fields: %w", err)
 	}
-	return nil
+	return func(restoreCtx context.Context) error {
+		original := proto.Clone(source).(*azdext.ServiceConfig)
+		if err := AddServiceSerialized(
+			restoreCtx,
+			p.azdClient,
+			original,
+		); err != nil {
+			return fmt.Errorf("restore referenced service: %w", err)
+		}
+		return nil
+	}, nil
 }
 
-func configContainsFileRef(config *structpb.Struct) bool {
+func combineBuildServiceRestore(
+	ctx context.Context,
+	restore func(context.Context) error,
+	operationErr error,
+	code string,
+) error {
+	if restore == nil {
+		return operationErr
+	}
+	restoreErr := restore(context.WithoutCancel(ctx))
+	if restoreErr == nil {
+		return operationErr
+	}
+	classified := exterrors.Internal(code, restoreErr.Error())
+	if operationErr == nil {
+		return classified
+	}
+	return errors.Join(operationErr, classified)
+}
+
+// ConfigContainsFileRef reports whether config contains a local $ref.
+func ConfigContainsFileRef(config *structpb.Struct) bool {
 	if config == nil {
 		return false
 	}
@@ -703,7 +752,7 @@ func (p *AgentServiceTargetProvider) Publish(
 	targetResource *azdext.TargetResource,
 	publishOptions *azdext.PublishOptions,
 	progress azdext.ProgressReporter,
-) (*azdext.ServicePublishResult, error) {
+) (result *azdext.ServicePublishResult, err error) {
 	// Pre-built image: nothing to package or push. Skip deploy-context
 	// resolution so this path stays cheap and doesn't require agent.yaml.
 	if preBuiltArtifact := findPreBuiltImageArtifact(serviceContext.Package); preBuiltArtifact != nil {
@@ -729,6 +778,27 @@ func (p *AgentServiceTargetProvider) Publish(
 		return &azdext.ServicePublishResult{}, nil
 	}
 
+	restore, err := p.persistBuildService(ctx)
+	if err != nil {
+		return nil, exterrors.Internal(
+			exterrors.OpContainerPublish,
+			fmt.Sprintf("prepare referenced service: %s", err),
+		)
+	}
+	if restore != nil {
+		defer func() {
+			err = combineBuildServiceRestore(
+				ctx,
+				restore,
+				err,
+				exterrors.OpContainerPublish,
+			)
+			if err != nil {
+				result = nil
+			}
+		}()
+	}
+
 	progress("Publishing container")
 	publishResponse, err := p.azdClient.
 		Container().