fix: Ensure agent wakes up on A2A task completion (#384)

Enabled A2A monitoring in config and fixed agent event channel
integration. Agent now properly receives wake-up events when A2A tasks
complete. Fixed task state handling with case-insensitive comparisons
and prevented premature agent exit.

## Changes Made:

1. **Enabled A2A monitoring** in `.infer/config.yaml`
2. **Fixed import ordering** in multiple files for consistency
3. **Added agent event channel integration** to wake up the agent when
A2A tasks complete
4. **Fixed A2A task state handling** with case-insensitive state
comparisons
5. **Prevented agent from exiting** while A2A tasks are pending in idle
state

## Key Benefits:
- A2A monitoring is now enabled by default
- Agent stays alive while A2A tasks are pending
- Proper event notification when tasks complete
- Robust state handling for various task status formats
- Improved reliability of A2A task execution

Author: edenreich

.infer/config.yaml

@@ -651,7 +651,7 @@ chat:
       cost: true
       git_branch: true
 a2a:
-  enabled: false
+  enabled: true
   cache:
     enabled: true
     ttl: 300

cmd/agent_test.go

@@ -5,10 +5,12 @@ import (
 	"testing"
 	"time"
 
+	sdk "github.com/inference-gateway/sdk"
+
+	domainmocks "github.com/inference-gateway/cli/tests/mocks/domain"
+
 	config "github.com/inference-gateway/cli/config"
 	domain "github.com/inference-gateway/cli/internal/domain"
-	domainmocks "github.com/inference-gateway/cli/tests/mocks/domain"
-	sdk "github.com/inference-gateway/sdk"
 )
 
 func TestIsModelAvailable(t *testing.T) {

internal/agent/agent.go

@@ -7,11 +7,12 @@ import (
 	"sync"
 	"time"
 
+	sdk "github.com/inference-gateway/sdk"
+
 	constants "github.com/inference-gateway/cli/internal/constants"
 	domain "github.com/inference-gateway/cli/internal/domain"
 	logger "github.com/inference-gateway/cli/internal/logger"
 	services "github.com/inference-gateway/cli/internal/services"
-	sdk "github.com/inference-gateway/sdk"
 )
 
 // AgentServiceImpl implements the AgentService interface with direct chat functionality
@@ -457,6 +458,10 @@ func (s *AgentServiceImpl) RunWithStream(ctx context.Context, req *domain.AgentR
 			taskTracker,
 		)
 
+		if monitor != nil {
+			monitor.SetAgentEventChannel(agent.GetEventChannel())
+		}
+
 		agent.Start()
 		agent.Wait()
 	}()

internal/agent/agent_event_driven.go

@@ -175,16 +175,11 @@ func (a *EventDrivenAgent) registerHandler(handler domain.StateHandler) {
 
 // Start begins the event-driven agent execution
 func (a *EventDrivenAgent) Start() {
-	logger.Debug("starting event-driven agent",
-		"request_id", a.req.RequestID,
-		"max_turns", a.agentCtx.MaxTurns)
-
 	_ = a.stateMachine.Transition(a.agentCtx, domain.StateIdle)
 
 	a.wg.Add(1)
 	go a.processEvents()
 
-	logger.Debug("triggering initial message received event")
 	a.events <- domain.MessageReceivedEvent{}
 }
 
@@ -194,42 +189,48 @@ func (a *EventDrivenAgent) Wait() {
 	close(a.events)
 }
 
+// GetEventChannel returns the agent's internal event channel for external components
+// to send wake-up events (e.g., when A2A tasks complete)
+func (a *EventDrivenAgent) GetEventChannel() chan<- domain.AgentEvent {
+	return a.events
+}
+
 // processEvents is the main event processing loop
 func (a *EventDrivenAgent) processEvents() {
 	defer a.wg.Done()
-	defer logger.Debug("agent event processing stopped", "request_id", a.req.RequestID)
 
 	for {
 		select {
 		case <-a.cancelChan:
-			logger.Debug("agent cancelled", "request_id", a.req.RequestID)
 			_ = a.stateMachine.Transition(a.agentCtx, domain.StateCancelled)
 			a.eventPublisher.publishChatComplete("", []sdk.ChatCompletionMessageToolCall{}, a.service.GetMetrics(a.req.RequestID))
 			return
 
 		case event, ok := <-a.events:
 			if !ok {
-				logger.Debug("event channel closed")
 				return
 			}
 
-			logger.Debug("received event",
-				"event_type", event.EventType(),
-				"current_state", a.stateMachine.GetCurrentState(),
-				"turn", a.agentCtx.Turns)
-
 			a.handleEvent(event)
 
 			currentState := a.stateMachine.GetCurrentState()
-			if currentState == domain.StateIdle ||
-				currentState == domain.StateStopped ||
+			if currentState == domain.StateStopped ||
 				currentState == domain.StateCancelled ||
 				currentState == domain.StateError {
-				logger.Debug("agent reached terminal state",
-					"state", currentState,
-					"total_turns", a.agentCtx.Turns)
 				return
 			}
+
+			if currentState == domain.StateIdle {
+				hasPendingTasks := a.taskTracker != nil && len(a.taskTracker.GetAllPollingTasks()) > 0
+				if hasPendingTasks {
+					logger.Debug("agent in Idle state but has pending A2A tasks, staying alive",
+						"pending_tasks", len(a.taskTracker.GetAllPollingTasks()))
+				} else {
+					logger.Debug("agent reached Idle state with no pending tasks",
+						"total_turns", a.agentCtx.Turns)
+					return
+				}
+			}
 		}
 	}
 }
@@ -240,11 +241,6 @@ func (a *EventDrivenAgent) handleEvent(event domain.AgentEvent) {
 	defer a.mu.Unlock()
 
 	currentState := a.stateMachine.GetCurrentState()
-	logger.Debug("handling event in state",
-		"event", event.EventType(),
-		"state", currentState,
-		"turn", a.agentCtx.Turns,
-		"queue_empty", a.agentCtx.MessageQueue.IsEmpty())
 
 	handler, exists := a.stateHandlers[currentState]
 	if !exists {
@@ -257,8 +253,4 @@ func (a *EventDrivenAgent) handleEvent(event domain.AgentEvent) {
 			"state", currentState.String(),
 			"error", err)
 	}
-
-	logger.Debug("event handled",
-		"event", event.EventType(),
-		"new_state", a.stateMachine.GetCurrentState())
 }

internal/agent/tools/a2a_task.go

@@ -7,11 +7,12 @@ import (
 	"strings"
 	"time"
 
-	client "github.com/inference-gateway/adk/client"
 	adk "github.com/inference-gateway/adk/types"
+	sdk "github.com/inference-gateway/sdk"
+
+	client "github.com/inference-gateway/adk/client"
 	config "github.com/inference-gateway/cli/config"
 	domain "github.com/inference-gateway/cli/internal/domain"
-	sdk "github.com/inference-gateway/sdk"
 )
 
 // A2ASubmitTaskTool handles A2A task submission and management
@@ -286,7 +287,6 @@ func (t *A2ASubmitTaskTool) pollTaskInBackground(
 	adkClient := t.getOrCreateClient(agentURL)
 
 	strategy := t.config.A2A.Task.PollingStrategy
-
 	currentInterval := t.initializePollingStrategy(agentURL, taskID, strategy)
 	state.CurrentInterval = currentInterval
 	state.NextPollTime = time.Now().Add(currentInterval)
@@ -427,8 +427,10 @@ func (t *A2ASubmitTaskTool) publishStatusUpdate(state *domain.TaskPollingState,
 }
 
 func (t *A2ASubmitTaskTool) handleTaskState(agentURL, _ /* taskID */ string, _ /* pollAttempt */ int, state *domain.TaskPollingState, currentTask adk.Task, _ /* pollingDetails */ string) (bool, *domain.ToolExecutionResult) {
-	switch currentTask.Status.State {
-	case adk.TaskStateCompleted:
+	normalizedState := strings.ToLower(string(currentTask.Status.State))
+
+	switch {
+	case normalizedState == strings.ToLower(string(adk.TaskStateCompleted)) || normalizedState == "completed":
 		finalResult := ""
 		if currentTask.Status.Message != nil {
 			finalResult = t.extractTextFromParts(currentTask.Status.Message.Parts)
@@ -451,7 +453,7 @@ func (t *A2ASubmitTaskTool) handleTaskState(agentURL, _ /* taskID */ string, _ /
 		}
 		return true, result
 
-	case adk.TaskStateFailed:
+	case normalizedState == strings.ToLower(string(adk.TaskStateFailed)) || normalizedState == "failed":
 		finalResult := ""
 		if currentTask.Status.Message != nil {
 			finalResult = t.extractTextFromParts(currentTask.Status.Message.Parts)
@@ -473,7 +475,7 @@ func (t *A2ASubmitTaskTool) handleTaskState(agentURL, _ /* taskID */ string, _ /
 		}
 		return true, result
 
-	case adk.TaskStateInputRequired:
+	case normalizedState == strings.ToLower(string(adk.TaskStateInputRequired)) || normalizedState == "input-required" || normalizedState == "input_required":
 		inputMessage := ""
 		if currentTask.Status.Message != nil {
 			inputMessage = t.extractTextFromParts(currentTask.Status.Message.Parts)
@@ -495,7 +497,7 @@ func (t *A2ASubmitTaskTool) handleTaskState(agentURL, _ /* taskID */ string, _ /
 		}
 		return true, result
 
-	case adk.TaskStateCancelled:
+	case normalizedState == strings.ToLower(string(adk.TaskStateCancelled)) || normalizedState == "cancelled" || normalizedState == "canceled":
 		cancelMessage := ""
 		if currentTask.Status.Message != nil {
 			cancelMessage = t.extractTextFromParts(currentTask.Status.Message.Parts)

internal/handlers/chat_handler.go

@@ -13,14 +13,16 @@ import (
 	"time"
 
 	tea "github.com/charmbracelet/bubbletea"
+
+	sdk "github.com/inference-gateway/sdk"
+
 	config "github.com/inference-gateway/cli/config"
 	tools "github.com/inference-gateway/cli/internal/agent/tools"
 	domain "github.com/inference-gateway/cli/internal/domain"
 	logger "github.com/inference-gateway/cli/internal/logger"
 	services "github.com/inference-gateway/cli/internal/services"
 	shortcuts "github.com/inference-gateway/cli/internal/shortcuts"
 	utils "github.com/inference-gateway/cli/internal/utils"
-	sdk "github.com/inference-gateway/sdk"
 )
 
 type ChatHandler struct {
@@ -1583,6 +1585,8 @@ func (h *ChatHandler) handleA2ATaskStatusUpdate(
 func (h *ChatHandler) handleMessageQueued(
 	_ domain.MessageQueuedEvent,
 ) (tea.Model, tea.Cmd) {
+	chatSession := h.stateManager.GetChatSession()
+
 	var cmds []tea.Cmd
 
 	cmds = append(cmds, func() tea.Msg {
@@ -1600,7 +1604,6 @@ func (h *ChatHandler) handleMessageQueued(
 		}
 	})
 
-	chatSession := h.stateManager.GetChatSession()
 	if chatSession != nil && chatSession.EventChannel != nil {
 		cmds = append(cmds, h.ListenForChatEvents(chatSession.EventChannel))
 	}

internal/services/a2a_polling_monitor.go

@@ -7,9 +7,10 @@ import (
 	"time"
 
 	adk "github.com/inference-gateway/adk/types"
+	sdk "github.com/inference-gateway/sdk"
+
 	domain "github.com/inference-gateway/cli/internal/domain"
 	logger "github.com/inference-gateway/cli/internal/logger"
-	sdk "github.com/inference-gateway/sdk"
 )
 
 type A2APollingMonitor struct {
@@ -22,6 +23,7 @@ type A2APollingMonitor struct {
 	activeMonitors   map[string]context.CancelFunc
 	stopChan         chan struct{}
 	stopped          bool
+	agentEventChan   chan<- domain.AgentEvent
 }
 
 func NewA2APollingMonitor(
@@ -43,6 +45,14 @@ func NewA2APollingMonitor(
 	}
 }
 
+// SetAgentEventChannel sets the agent's internal event channel for waking up the agent
+// when an A2A task completes. This should be called after the agent is created.
+func (m *A2APollingMonitor) SetAgentEventChannel(eventChan chan<- domain.AgentEvent) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.agentEventChan = eventChan
+}
+
 func (m *A2APollingMonitor) Start(ctx context.Context) {
 	ticker := time.NewTicker(1 * time.Second)
 	defer ticker.Stop()
@@ -305,6 +315,21 @@ func (m *A2APollingMonitor) addResultToMessageQueue(taskID string, result *domai
 
 	m.messageQueue.Enqueue(message, m.requestID)
 
+	m.mu.RLock()
+	agentChan := m.agentEventChan
+	m.mu.RUnlock()
+
+	if agentChan != nil {
+		select {
+		case agentChan <- domain.MessageReceivedEvent{}:
+			logger.Debug("Sent wake-up event to agent for A2A task completion",
+				"task_id", taskID)
+		default:
+			logger.Debug("Failed to send wake-up event to agent - channel full",
+				"task_id", taskID)
+		}
+	}
+
 	if m.eventChan != nil {
 		event := domain.MessageQueuedEvent{
 			RequestID: m.requestID,