Cleanup agent based on initial PR feedback

Author: wbreza

cli/azd/cmd/init.go

@@ -384,7 +384,6 @@ func (i *initAction) initAppWithAgent(ctx context.Context) error {
 	}
 
 	defer azdAgent.Stop()
-	agentThoughts := azdAgent.Thoughts()
 
 	type initStep struct {
 		Name        string

cli/azd/internal/agent/agent.go

@@ -25,12 +25,18 @@ type agentBase struct {
 	callbacksHandler callbacks.Handler
 	thoughtChan      chan logging.Thought
 	cleanupFunc      AgentCleanup
+	maxIterations    int
 }
 
+// AgentCleanup is a function that performs cleanup tasks for an agent.
 type AgentCleanup func() error
 
+// Agent represents an AI agent that can execute tools and interact with language models.
 type Agent interface {
+	// SendMessage sends a message to the agent and returns the response
 	SendMessage(ctx context.Context, args ...string) (string, error)
+
+	// Stop terminates the agent and performs any necessary cleanup
 	Stop() error
 }
 
@@ -53,6 +59,13 @@ func WithDebug(debug bool) AgentCreateOption {
 	}
 }
 
+// WithMaxIterations returns an option that sets the maximum number of iterations for the agent
+func WithMaxIterations(maxIterations int) AgentCreateOption {
+	return func(agent *agentBase) {
+		agent.maxIterations = maxIterations
+	}
+}
+
 // WithDefaultModel returns an option that sets the default language model for the agent
 func WithDefaultModel(model llms.Model) AgentCreateOption {
 	return func(agent *agentBase) {
@@ -74,12 +87,20 @@ func WithCallbacksHandler(handler callbacks.Handler) AgentCreateOption {
 	}
 }
 
+// WithThoughtChannel returns an option that sets the thought channel for the agent
 func WithThoughtChannel(thoughtChan chan logging.Thought) AgentCreateOption {
 	return func(agent *agentBase) {
 		agent.thoughtChan = thoughtChan
 	}
 }
 
+// WithCleanup returns an option that sets the cleanup function for the agent
+func WithCleanup(cleanupFunc AgentCleanup) AgentCreateOption {
+	return func(agent *agentBase) {
+		agent.cleanupFunc = cleanupFunc
+	}
+}
+
 // toolNames returns a comma-separated string of all tool names in the provided slice
 func toolNames(tools []common.AnnotatedTool) string {
 	var tn strings.Builder

cli/azd/internal/agent/agent_factory.go

@@ -13,12 +13,14 @@ import (
 	"github.com/azure/azure-dev/cli/azd/pkg/llm"
 )
 
+// AgentFactory is responsible for creating agent instances
 type AgentFactory struct {
 	consentManager consent.ConsentManager
 	llmManager     *llm.Manager
 	console        input.Console
 }
 
+// NewAgentFactory creates a new instance of AgentFactory
 func NewAgentFactory(
 	consentManager consent.ConsentManager,
 	console input.Console,
@@ -31,13 +33,16 @@ func NewAgentFactory(
 	}
 }
 
+// CreateAgent creates a new agent instance
 func (f *AgentFactory) Create(opts ...AgentCreateOption) (Agent, error) {
+	// Create a daily log file for all agent activity
 	fileLogger, loggerCleanup, err := logging.NewFileLoggerDefault()
 	if err != nil {
 		defer loggerCleanup()
 		return nil, err
 	}
 
+	// Create a channel for logging thoughts & actions
 	thoughtChan := make(chan logging.Thought)
 	thoughtHandler := logging.NewThoughtLogger(thoughtChan)
 	chainedHandler := logging.NewChainedHandler(fileLogger, thoughtHandler)
@@ -47,13 +52,16 @@ func (f *AgentFactory) Create(opts ...AgentCreateOption) (Agent, error) {
 		return loggerCleanup()
 	}
 
+	// Default model gets the chained handler to expose the UX experience for the agent
 	defaultModelContainer, err := f.llmManager.GetDefaultModel(llm.WithLogger(chainedHandler))
 	if err != nil {
 		defer cleanup()
 		return nil, err
 	}
 
-	samplingModelContainer, err := f.llmManager.GetDefaultModel(llm.WithLogger(chainedHandler))
+	// Sampling model only gets the file logger to output sampling actions
+	// We don't need UX for sampling requests right now
+	samplingModelContainer, err := f.llmManager.GetDefaultModel(llm.WithLogger(fileLogger))
 	if err != nil {
 		defer cleanup()
 		return nil, err
@@ -66,7 +74,8 @@ func (f *AgentFactory) Create(opts ...AgentCreateOption) (Agent, error) {
 		samplingModelContainer,
 	)
 
-	toolLoaders := []localtools.ToolLoader{
+	// Loads build-in tools & any referenced MCP servers
+	toolLoaders := []common.ToolLoader{
 		localtools.NewLocalToolsLoader(),
 		mcptools.NewMcpToolsLoader(samplingHandler),
 	}
@@ -95,22 +104,23 @@ func (f *AgentFactory) Create(opts ...AgentCreateOption) (Agent, error) {
 		}
 	}
 
+	// Wraps all tools in consent workflow
 	protectedTools := f.consentManager.WrapTools(allTools)
 
+	// Finalize agent creation options
 	allOptions := []AgentCreateOption{}
 	allOptions = append(allOptions, opts...)
 	allOptions = append(allOptions,
 		WithCallbacksHandler(chainedHandler),
 		WithThoughtChannel(thoughtChan),
 		WithTools(protectedTools...),
+		WithCleanup(cleanup),
 	)
 
 	azdAgent, err := NewConversationalAzdAiAgent(defaultModelContainer.Model, allOptions...)
 	if err != nil {
 		return nil, err
 	}
 
-	azdAgent.cleanupFunc = cleanup
-
 	return azdAgent, nil
 }

cli/azd/internal/agent/consent/checker.go

@@ -13,6 +13,9 @@ import (
 	"github.com/mark3labs/mcp-go/mcp"
 )
 
+var ErrToolExecutionDenied = fmt.Errorf("tool execution denied by user")
+var ErrSamplingDenied = fmt.Errorf("sampling denied by user")
+
 // ConsentChecker provides shared consent checking logic for different tool types
 type ConsentChecker struct {
 	consentMgr ConsentManager
@@ -66,6 +69,49 @@ func (cc *ConsentChecker) CheckSamplingConsent(
 	return cc.consentMgr.CheckConsent(ctx, consentRequest)
 }
 
+// PromptAndGrantConsent shows consent prompt and grants permission based on user choice
+func (cc *ConsentChecker) PromptAndGrantConsent(
+	ctx context.Context,
+	toolName, toolDesc string,
+	annotations mcp.ToolAnnotation,
+) error {
+	toolId := fmt.Sprintf("%s/%s", cc.serverName, toolName)
+
+	choice, err := cc.promptForToolConsent(ctx, toolName, toolDesc, annotations)
+	if err != nil {
+		return err
+	}
+
+	if choice == "deny" {
+		return ErrToolExecutionDenied
+	}
+
+	// Grant consent based on user choice
+	return cc.grantConsentFromChoice(ctx, toolId, choice, OperationTypeTool)
+}
+
+// PromptAndGrantSamplingConsent shows sampling consent prompt and grants permission based on user choice
+func (cc *ConsentChecker) PromptAndGrantSamplingConsent(
+	ctx context.Context,
+	toolName, toolDesc string,
+) error {
+	toolId := fmt.Sprintf("%s/%s", cc.serverName, toolName)
+
+	choice, err := cc.promptForSamplingConsent(ctx, toolName, toolDesc)
+	if err != nil {
+		return fmt.Errorf("sampling consent prompt failed: %w", err)
+	}
+
+	if choice == "deny" {
+		return ErrSamplingDenied
+	}
+
+	// Grant sampling consent based on user choice
+	return cc.grantConsentFromChoice(ctx, toolId, choice, OperationTypeSampling)
+}
+
+// Private Struct Methods
+
 // formatToolDescriptionWithAnnotations creates a formatted description with tool annotations as bullet points
 func (cc *ConsentChecker) formatToolDescriptionWithAnnotations(
 	toolDesc string,
@@ -128,27 +174,6 @@ func (cc *ConsentChecker) formatToolDescriptionWithAnnotations(
 	return description
 }
 
-// PromptAndGrantConsent shows consent prompt and grants permission based on user choice
-func (cc *ConsentChecker) PromptAndGrantConsent(
-	ctx context.Context,
-	toolName, toolDesc string,
-	annotations mcp.ToolAnnotation,
-) error {
-	toolId := fmt.Sprintf("%s/%s", cc.serverName, toolName)
-
-	choice, err := cc.promptForToolConsent(ctx, toolName, toolDesc, annotations)
-	if err != nil {
-		return fmt.Errorf("consent prompt failed: %w", err)
-	}
-
-	if choice == "deny" {
-		return fmt.Errorf("tool execution denied by user")
-	}
-
-	// Grant consent based on user choice
-	return cc.grantConsentFromChoice(ctx, toolId, choice, OperationTypeTool)
-}
-
 // promptForToolConsent shows an interactive consent prompt and returns the user's choice
 func (cc *ConsentChecker) promptForToolConsent(
 	ctx context.Context,
@@ -361,26 +386,6 @@ func (cc *ConsentChecker) grantConsentFromChoice(
 	return cc.consentMgr.GrantConsent(ctx, rule)
 }
 
-// PromptAndGrantSamplingConsent shows sampling consent prompt and grants permission based on user choice
-func (cc *ConsentChecker) PromptAndGrantSamplingConsent(
-	ctx context.Context,
-	toolName, toolDesc string,
-) error {
-	toolId := fmt.Sprintf("%s/%s", cc.serverName, toolName)
-
-	choice, err := cc.promptForSamplingConsent(ctx, toolName, toolDesc)
-	if err != nil {
-		return fmt.Errorf("sampling consent prompt failed: %w", err)
-	}
-
-	if choice == "deny" {
-		return fmt.Errorf("sampling denied by user")
-	}
-
-	// Grant sampling consent based on user choice
-	return cc.grantConsentFromChoice(ctx, toolId, choice, OperationTypeSampling)
-}
-
 // promptForSamplingConsent shows an interactive sampling consent prompt and returns the user's choice
 func (cc *ConsentChecker) promptForSamplingConsent(
 	ctx context.Context,

cli/azd/internal/agent/conversational_agent.go

@@ -4,11 +4,9 @@
 package agent
 
 import (
-	"bufio"
 	"context"
 	_ "embed"
 	"fmt"
-	"os"
 	"strings"
 	"time"
 
@@ -34,7 +32,7 @@ type ConversationalAzdAiAgent struct {
 // NewConversationalAzdAiAgent creates a new conversational agent with memory, tool loading,
 // and MCP sampling capabilities. It filters out excluded tools and configures the agent
 // for interactive conversations with a high iteration limit for complex tasks.
-func NewConversationalAzdAiAgent(llm llms.Model, opts ...AgentCreateOption) (*ConversationalAzdAiAgent, error) {
+func NewConversationalAzdAiAgent(llm llms.Model, opts ...AgentCreateOption) (Agent, error) {
 	azdAgent := &ConversationalAzdAiAgent{
 		agentBase: &agentBase{
 			defaultModel: llm,
@@ -46,6 +44,11 @@ func NewConversationalAzdAiAgent(llm llms.Model, opts ...AgentCreateOption) (*Co
 		opt(azdAgent.agentBase)
 	}
 
+	// Default max iterations
+	if azdAgent.maxIterations <= 0 {
+		azdAgent.maxIterations = 100
+	}
+
 	smartMemory := memory.NewConversationBuffer(
 		memory.WithInputKey("input"),
 		memory.WithOutputKey("output"),
@@ -74,7 +77,7 @@ func NewConversationalAzdAiAgent(llm llms.Model, opts ...AgentCreateOption) (*Co
 
 	// 5. Create executor without separate memory configuration since agent already has it
 	executor := agents.NewExecutor(conversationAgent,
-		agents.WithMaxIterations(100),
+		agents.WithMaxIterations(azdAgent.maxIterations),
 		agents.WithMemory(smartMemory),
 		agents.WithCallbacksHandler(azdAgent.callbacksHandler),
 		agents.WithReturnIntermediateSteps(),
@@ -86,62 +89,6 @@ func NewConversationalAzdAiAgent(llm llms.Model, opts ...AgentCreateOption) (*Co
 
 // SendMessage processes a single message through the agent and returns the response
 func (aai *ConversationalAzdAiAgent) SendMessage(ctx context.Context, args ...string) (string, error) {
-	return aai.runChain(ctx, strings.Join(args, "\n"))
-}
-
-// StartConversation runs an interactive conversation loop with the agent.
-// It accepts an optional initial query and handles user input/output with proper formatting.
-// The conversation continues until the user types "exit" or "quit".
-func (aai *ConversationalAzdAiAgent) StartConversation(ctx context.Context, args ...string) (string, error) {
-	// Handle initial query if provided
-	var initialQuery string
-	if len(args) > 0 {
-		initialQuery = strings.Join(args, " ")
-	}
-
-	scanner := bufio.NewScanner(os.Stdin)
-
-	for {
-		var userInput string
-
-		if initialQuery != "" {
-			userInput = initialQuery
-			initialQuery = "" // Clear after first use
-			color.Cyan("💬 You: %s\n", userInput)
-		} else {
-			fmt.Print(color.CyanString("\n💬 You: "))
-			color.Set(color.FgCyan) // Set blue color for user input
-			if !scanner.Scan() {
-				color.Unset() // Reset color
-				break         // EOF or error
-			}
-			userInput = strings.TrimSpace(scanner.Text())
-			color.Unset() // Reset color after input
-		}
-
-		// Check for exit commands
-		if userInput == "" {
-			continue
-		}
-
-		if strings.ToLower(userInput) == "exit" || strings.ToLower(userInput) == "quit" {
-			fmt.Println("👋 Goodbye! Thanks for using azd Agent!")
-			break
-		}
-
-		// Process the query with the enhanced agent
-		return aai.runChain(ctx, userInput)
-	}
-
-	if err := scanner.Err(); err != nil {
-		return "", fmt.Errorf("error reading input: %w", err)
-	}
-
-	return "", nil
-}
-
-// runChain executes a user query through the agent's chain with memory and returns the response
-func (aai *ConversationalAzdAiAgent) runChain(ctx context.Context, userInput string) (string, error) {
 	thoughtsCtx, cancelCtx := context.WithCancel(ctx)
 	cleanup, err := aai.renderThoughts(thoughtsCtx)
 	if err != nil {
@@ -154,8 +101,7 @@ func (aai *ConversationalAzdAiAgent) runChain(ctx context.Context, userInput str
 		cancelCtx()
 	}()
 
-	// Execute with enhanced input - agent should automatically handle memory
-	output, err := chains.Run(ctx, aai.executor, userInput)
+	output, err := chains.Run(ctx, aai.executor, strings.Join(args, "\n"))
 	if err != nil {
 		return "", err
 	}