[P1] Implement LeadAgent.assign_task() logic

[frankbria]

Overview

The Lead Agent's task assignment method is stubbed, preventing dynamic task allocation to worker agents.

Current State

File: codeframe/agents/lead_agent.py lines 590-591

def assign_task(self, task_id: int, agent_id: str) -> None:
    """Assign task to worker agent."""
    # TODO: Implement task assignment logic
    pass

Expected Behavior

• Validate that task exists and is not already assigned
• Validate that agent exists and is available
• Update task assignment in database
• Notify agent via WebSocket or internal message queue
• Track assignment timestamp for bottleneck detection
• Log assignment for audit trail

Implementation Requirements

1. Query database to validate task and agent
2. Update tasks table with assigned_to field
3. Broadcast assignment event via WebSocket
4. Handle errors (task not found, agent unavailable)
5. Support reassignment (when agent fails)

Acceptance Criteria

• Tasks successfully assigned to specific agents
• Database updated with assignment information
• WebSocket event broadcast to notify agents
• Error handling for invalid task/agent IDs
• Unit tests verify assignment logic
• Integration tests verify end-to-end assignment flow

Priority

P1 - Critical: Required for multi-agent coordination.

References

• User Documentation: CODEFRAME_USER_DOCUMENTATION.md lines 81-82
• Issues Analysis: CODEFRAME_ISSUES_ANALYSIS.md lines 22-29

[traycerai]

Plan

Observations

The assign_task() method in file:codeframe/agents/lead_agent.py (lines 589-592) is currently stubbed. The codebase has a robust infrastructure for task management, agent coordination, WebSocket broadcasting, and database operations. The Task model includes an assigned_to field (line 148 in file:codeframe/core/models.py), the database has an update_task() method for field updates, and file:codeframe/ui/websocket_broadcasts.py provides a broadcast_task_assigned() function (lines 391-419). The multi-agent execution loop in _assign_and_execute_task() (lines 1287-1330) already handles task assignment dynamically but doesn't use the assign_task() method.

Approach

The implementation will follow the existing patterns in the codebase: validate inputs using database queries, update the database atomically, broadcast events via WebSocket for real-time UI updates, log operations for audit trails, and handle errors gracefully with descriptive messages. The method will support both initial assignment and reassignment scenarios (when agents fail), integrate with the existing agent pool manager, and track assignment timestamps for bottleneck detection. This approach aligns with the codebase's emphasis on observability, error handling, and real-time updates.

Implementation Steps

1. Input Validation

In file:codeframe/agents/lead_agent.py, implement the assign_task() method starting at line 589:

• Validate task_id: Query the database using self.db.get_task(task_id) to verify the task exists

  • If task is None, raise ValueError with message: "Task {task_id} not found"
  • Extract project_id from task dict and verify it matches self.project_id
  • If mismatch, raise ValueError with message: "Task {task_id} does not belong to project {self.project_id}"

• Validate agent_id: Check if agent exists in the agent pool

  • Use self.agent_pool_manager.get_agent_status() to get all agents
  • If agent_id not in the returned status dict, raise ValueError with message: "Agent {agent_id} not found in agent pool"
  • Check agent status - if status is "blocked", raise ValueError with message: "Agent {agent_id} is blocked and cannot accept new tasks"

2. Check Assignment Status

• Verify task is assignable: Check task status from the retrieved task dict
  • If status is "completed", raise ValueError with message: "Cannot assign completed task {task_id}"
  • If status is "in_progress" and assigned_to field is not None and not equal to agent_id, log a warning: "Task {task_id} is already assigned to {current_agent}, reassigning to {agent_id}"
  • For reassignment scenarios, this allows recovery from agent failures

3. Update Database

• Update task assignment: Call self.db.update_task(task_id, updates) with the following fields:

  • "assigned_to": agent_id
  • "status": "assigned" (use TaskStatus.ASSIGNED.value for enum compatibility)
  • Store current timestamp for bottleneck detection (the database schema should be checked for a timestamp field like assigned_at or use the existing updated_at field)

• Handle database errors: Wrap the update in a try-except block

  • Catch Exception and log error: "Failed to update task assignment in database: {error}"
  • Re-raise the exception to propagate to caller

4. Broadcast WebSocket Event

• Notify connected clients: Use the existing broadcast_task_assigned() function from file:codeframe/ui/websocket_broadcasts.py
  • Import at top of file: from codeframe.ui.websocket_broadcasts import broadcast_task_assigned
  • Check if self.ws_manager exists (it's optional)
  • If WebSocket manager is available, call the broadcast function asynchronously:

    if self.ws_manager:
        try:
            loop = asyncio.get_running_loop()
            loop.create_task(
                broadcast_task_assigned(
                    self.ws_manager,
                    self.project_id,
                    task_id,
                    agent_id,
                    task_title=task.get("title")
                )
            )
        except RuntimeError:
            # No event loop running (sync context)
            logger.debug(f"Skipped WebSocket broadcast for task {task_id} assignment")

  • Handle broadcast failures gracefully - log but don't fail the assignment

5. Logging and Audit Trail

• Log successful assignment: Use the existing logger

  • At INFO level: "✅ Task {task_id} ({task_title}) assigned to agent {agent_id}"
  • Include task title for better observability

• Log reassignment: If this is a reassignment (task was already assigned)

  • At WARNING level: "⚠️ Task {task_id} reassigned from {old_agent} to {agent_id}"

6. Error Handling Summary

Implement comprehensive error handling for all failure scenarios:

Error Condition	Exception Type	Error Message
Task not found	ValueError	"Task {task_id} not found"
Task belongs to different project	ValueError	"Task {task_id} does not belong to project {self.project_id}"
Agent not found	ValueError	"Agent {agent_id} not found in agent pool"
Agent is blocked	ValueError	"Agent {agent_id} is blocked and cannot accept new tasks"
Task already completed	ValueError	"Cannot assign completed task {task_id}"
Database update fails	Re-raise original exception	Log: "Failed to update task assignment in database: {error}"

7. Integration Points

The implementation integrates with existing components:

• Database Layer: Uses self.db.get_task() and self.db.update_task() from file:codeframe/persistence/database.py
• Agent Pool Manager: Validates agent existence via self.agent_pool_manager.get_agent_status() from file:codeframe/agents/agent_pool_manager.py
• WebSocket Broadcasting: Uses broadcast_task_assigned() from file:codeframe/ui/websocket_broadcasts.py for real-time UI updates
• Task Model: References TaskStatus.ASSIGNED enum from file:codeframe/core/models.py

8. Method Signature and Return Type

The method signature should remain as defined:

def assign_task(self, task_id: int, agent_id: str) -> None:
    """Assign task to worker agent.
    
    Args:
        task_id: ID of the task to assign
        agent_id: ID of the agent to assign the task to
        
    Raises:
        ValueError: If task or agent not found, or assignment is invalid
        Exception: If database update fails
    """

9. Testing Considerations

The implementation should be verified with unit tests covering:

• Happy path: Valid task and agent assignment succeeds
• Task validation: Non-existent task raises ValueError
• Agent validation: Non-existent agent raises ValueError
• Status checks: Completed tasks cannot be assigned
• Reassignment: Task already assigned to different agent logs warning
• Database errors: Database failures are propagated
• WebSocket broadcasting: Events are broadcast when manager is available
• Logging: All operations are logged appropriately

Test file location: file:tests/agents/test_lead_agent.py (add new test class TestLeadAgentTaskAssignment)

10. Sequence Diagram

sequenceDiagram
    participant Caller
    participant LeadAgent
    participant Database
    participant AgentPool
    participant WebSocket
    
    Caller->>LeadAgent: assign_task(task_id, agent_id)
    
    LeadAgent->>Database: get_task(task_id)
    Database-->>LeadAgent: task_dict or None
    
    alt Task not found
        LeadAgent-->>Caller: raise ValueError("Task not found")
    end
    
    LeadAgent->>LeadAgent: Validate project_id matches
    
    alt Project mismatch
        LeadAgent-->>Caller: raise ValueError("Task belongs to different project")
    end
    
    LeadAgent->>AgentPool: get_agent_status()
    AgentPool-->>LeadAgent: agent_status_dict
    
    alt Agent not found
        LeadAgent-->>Caller: raise ValueError("Agent not found")
    end
    
    alt Agent blocked
        LeadAgent-->>Caller: raise ValueError("Agent is blocked")
    end
    
    LeadAgent->>LeadAgent: Check task status
    
    alt Task completed
        LeadAgent-->>Caller: raise ValueError("Cannot assign completed task")
    end
    
    LeadAgent->>Database: update_task(task_id, {assigned_to, status})
    Database-->>LeadAgent: success
    
    LeadAgent->>WebSocket: broadcast_task_assigned(...)
    WebSocket-->>LeadAgent: (async, fire-and-forget)
    
    LeadAgent->>LeadAgent: Log assignment
    LeadAgent-->>Caller: return (success)

Loading

Import In IDE

🤖 Prompt for AI Agents

## Observations

The `assign_task()` method in `file:codeframe/agents/lead_agent.py` (lines 589-592) is currently stubbed. The codebase has a robust infrastructure for task management, agent coordination, WebSocket broadcasting, and database operations. The `Task` model includes an `assigned_to` field (line 148 in `file:codeframe/core/models.py`), the database has an `update_task()` method for field updates, and `file:codeframe/ui/websocket_broadcasts.py` provides a `broadcast_task_assigned()` function (lines 391-419). The multi-agent execution loop in `_assign_and_execute_task()` (lines 1287-1330) already handles task assignment dynamically but doesn't use the `assign_task()` method.

## Approach

The implementation will follow the existing patterns in the codebase: validate inputs using database queries, update the database atomically, broadcast events via WebSocket for real-time UI updates, log operations for audit trails, and handle errors gracefully with descriptive messages. The method will support both initial assignment and reassignment scenarios (when agents fail), integrate with the existing agent pool manager, and track assignment timestamps for bottleneck detection. This approach aligns with the codebase's emphasis on observability, error handling, and real-time updates.

## Implementation Steps

### 1. Input Validation

In `file:codeframe/agents/lead_agent.py`, implement the `assign_task()` method starting at line 589:

- **Validate task_id**: Query the database using `self.db.get_task(task_id)` to verify the task exists
  - If task is `None`, raise `ValueError` with message: `"Task {task_id} not found"`
  - Extract `project_id` from task dict and verify it matches `self.project_id`
  - If mismatch, raise `ValueError` with message: `"Task {task_id} does not belong to project {self.project_id}"`

- **Validate agent_id**: Check if agent exists in the agent pool
  - Use `self.agent_pool_manager.get_agent_status()` to get all agents
  - If `agent_id` not in the returned status dict, raise `ValueError` with message: `"Agent {agent_id} not found in agent pool"`
  - Check agent status - if status is `"blocked"`, raise `ValueError` with message: `"Agent {agent_id} is blocked and cannot accept new tasks"`

### 2. Check Assignment Status

- **Verify task is assignable**: Check task status from the retrieved task dict
  - If status is `"completed"`, raise `ValueError` with message: `"Cannot assign completed task {task_id}"`
  - If status is `"in_progress"` and `assigned_to` field is not None and not equal to `agent_id`, log a warning: `"Task {task_id} is already assigned to {current_agent}, reassigning to {agent_id}"`
  - For reassignment scenarios, this allows recovery from agent failures

### 3. Update Database

- **Update task assignment**: Call `self.db.update_task(task_id, updates)` with the following fields:
  - `"assigned_to": agent_id`
  - `"status": "assigned"` (use `TaskStatus.ASSIGNED.value` for enum compatibility)
  - Store current timestamp for bottleneck detection (the database schema should be checked for a timestamp field like `assigned_at` or use the existing `updated_at` field)

- **Handle database errors**: Wrap the update in a try-except block
  - Catch `Exception` and log error: `"Failed to update task assignment in database: {error}"`
  - Re-raise the exception to propagate to caller

### 4. Broadcast WebSocket Event

- **Notify connected clients**: Use the existing `broadcast_task_assigned()` function from `file:codeframe/ui/websocket_broadcasts.py`
  - Import at top of file: `from codeframe.ui.websocket_broadcasts import broadcast_task_assigned`
  - Check if `self.ws_manager` exists (it's optional)
  - If WebSocket manager is available, call the broadcast function asynchronously:
    ```python
    if self.ws_manager:
        try:
            loop = asyncio.get_running_loop()
            loop.create_task(
                broadcast_task_assigned(
                    self.ws_manager,
                    self.project_id,
                    task_id,
                    agent_id,
                    task_title=task.get("title")
                )
            )
        except RuntimeError:
            # No event loop running (sync context)
            logger.debug(f"Skipped WebSocket broadcast for task {task_id} assignment")
    ```
  - Handle broadcast failures gracefully - log but don't fail the assignment

### 5. Logging and Audit Trail

- **Log successful assignment**: Use the existing logger
  - At INFO level: `"✅ Task {task_id} ({task_title}) assigned to agent {agent_id}"`
  - Include task title for better observability

- **Log reassignment**: If this is a reassignment (task was already assigned)
  - At WARNING level: `"⚠️  Task {task_id} reassigned from {old_agent} to {agent_id}"`

### 6. Error Handling Summary

Implement comprehensive error handling for all failure scenarios:

| Error Condition | Exception Type | Error Message |
|----------------|---------------|---------------|
| Task not found | `ValueError` | `"Task {task_id} not found"` |
| Task belongs to different project | `ValueError` | `"Task {task_id} does not belong to project {self.project_id}"` |
| Agent not found | `ValueError` | `"Agent {agent_id} not found in agent pool"` |
| Agent is blocked | `ValueError` | `"Agent {agent_id} is blocked and cannot accept new tasks"` |
| Task already completed | `ValueError` | `"Cannot assign completed task {task_id}"` |
| Database update fails | Re-raise original exception | Log: `"Failed to update task assignment in database: {error}"` |

### 7. Integration Points

The implementation integrates with existing components:

- **Database Layer**: Uses `self.db.get_task()` and `self.db.update_task()` from `file:codeframe/persistence/database.py`
- **Agent Pool Manager**: Validates agent existence via `self.agent_pool_manager.get_agent_status()` from `file:codeframe/agents/agent_pool_manager.py`
- **WebSocket Broadcasting**: Uses `broadcast_task_assigned()` from `file:codeframe/ui/websocket_broadcasts.py` for real-time UI updates
- **Task Model**: References `TaskStatus.ASSIGNED` enum from `file:codeframe/core/models.py`

### 8. Method Signature and Return Type

The method signature should remain as defined:
```python
def assign_task(self, task_id: int, agent_id: str) -> None:
    """Assign task to worker agent.
    
    Args:
        task_id: ID of the task to assign
        agent_id: ID of the agent to assign the task to
        
    Raises:
        ValueError: If task or agent not found, or assignment is invalid
        Exception: If database update fails
    """
```

### 9. Testing Considerations

The implementation should be verified with unit tests covering:

- **Happy path**: Valid task and agent assignment succeeds
- **Task validation**: Non-existent task raises `ValueError`
- **Agent validation**: Non-existent agent raises `ValueError`
- **Status checks**: Completed tasks cannot be assigned
- **Reassignment**: Task already assigned to different agent logs warning
- **Database errors**: Database failures are propagated
- **WebSocket broadcasting**: Events are broadcast when manager is available
- **Logging**: All operations are logged appropriately

Test file location: `file:tests/agents/test_lead_agent.py` (add new test class `TestLeadAgentTaskAssignment`)

### 10. Sequence Diagram

```mermaid
sequenceDiagram
    participant Caller
    participant LeadAgent
    participant Database
    participant AgentPool
    participant WebSocket
    
    Caller->>LeadAgent: assign_task(task_id, agent_id)
    
    LeadAgent->>Database: get_task(task_id)
    Database-->>LeadAgent: task_dict or None
    
    alt Task not found
        LeadAgent-->>Caller: raise ValueError("Task not found")
    end
    
    LeadAgent->>LeadAgent: Validate project_id matches
    
    alt Project mismatch
        LeadAgent-->>Caller: raise ValueError("Task belongs to different project")
    end
    
    LeadAgent->>AgentPool: get_agent_status()
    AgentPool-->>LeadAgent: agent_status_dict
    
    alt Agent not found
        LeadAgent-->>Caller: raise ValueError("Agent not found")
    end
    
    alt Agent blocked
        LeadAgent-->>Caller: raise ValueError("Agent is blocked")
    end
    
    LeadAgent->>LeadAgent: Check task status
    
    alt Task completed
        LeadAgent-->>Caller: raise ValueError("Cannot assign completed task")
    end
    
    LeadAgent->>Database: update_task(task_id, {assigned_to, status})
    Database-->>LeadAgent: success
    
    LeadAgent->>WebSocket: broadcast_task_assigned(...)
    WebSocket-->>LeadAgent: (async, fire-and-forget)
    
    LeadAgent->>LeadAgent: Log assignment
    LeadAgent-->>Caller: return (success)
```

---

Execution Information

Branch: main
Commit: 1a87e4e

---

💡 Tips

Supported Commands (Inside Comments)

• Use @traycerai generate to iterate on the previous version of the implementation plan.

Supported Commands (Inside Description)

• Add @traycerai ignore anywhere in the ticket description to prevent this ticket from being processed.
• Add @traycerai branch:<branch-name> anywhere in the ticket description to specify the target branch for the implementation plan.

Community

• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

📝 CodeRabbit Plan Mode

Generate an implementation plan and prompts that you can use with your favorite coding agent.

• Create Plan

Examples

• Example 1
• Example 2

_{You can configure auto-planning by selecting labels in the issue_enrichment configuration.}

🔗 Similar Issues

Related Issues

• [P1] Implement WorkerAgent.execute_task() with Anthropic API integration #98
• [Future] Evaluate Subagent Context Isolation per Task #68

🔗 Related PRs

#3 - feat(sprint-4): Multi-Agent Coordination System - P0 & P1 Complete [merged]

👤 Suggested Assignees

• frankbria

---

🧪 Issue enrichment is currently in open beta.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord or schedule a call!

[frankbria]

Closed with PR #127