feat: implement single-task execution lifecycle (#21) #144
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -4,6 +4,7 @@ | |||||
| tool invocation, and budget tracking into a single ``run()`` entry point. | ||||||
| """ | ||||||
|
|
||||||
| import asyncio | ||||||
| import time | ||||||
| from datetime import UTC, datetime | ||||||
| from typing import TYPE_CHECKING | ||||||
|
|
@@ -17,6 +18,7 @@ | |||||
| TerminationReason, | ||||||
| TurnRecord, | ||||||
| ) | ||||||
| from ai_company.engine.metrics import TaskCompletionMetrics | ||||||
| from ai_company.engine.prompt import SystemPrompt, build_system_prompt | ||||||
| from ai_company.engine.react_loop import ReactLoop | ||||||
| from ai_company.engine.run_result import AgentRunResult | ||||||
|
|
@@ -31,7 +33,9 @@ | |||||
| EXECUTION_ENGINE_INVALID_INPUT, | ||||||
| EXECUTION_ENGINE_PROMPT_BUILT, | ||||||
| EXECUTION_ENGINE_START, | ||||||
| EXECUTION_ENGINE_TASK_METRICS, | ||||||
| EXECUTION_ENGINE_TASK_TRANSITION, | ||||||
| EXECUTION_ENGINE_TIMEOUT, | ||||||
| ) | ||||||
| from ai_company.providers.enums import MessageRole | ||||||
| from ai_company.providers.models import ChatMessage | ||||||
|
|
@@ -90,14 +94,15 @@ def __init__( | |||||
| has_cost_tracker=self._cost_tracker is not None, | ||||||
| ) | ||||||
|
|
||||||
| async def run( # noqa: PLR0913 | ||||||
| self, | ||||||
| *, | ||||||
| identity: AgentIdentity, | ||||||
| task: Task, | ||||||
| completion_config: CompletionConfig | None = None, | ||||||
| max_turns: int = DEFAULT_MAX_TURNS, | ||||||
| memory_messages: tuple[ChatMessage, ...] = (), | ||||||
| timeout_seconds: float | None = None, | ||||||
| ) -> AgentRunResult: | ||||||
|
Comment on lines
+98
to
107
Contributor
There was a problem hiding this comment.
The |
||||||
| """Execute an agent on a task. | ||||||
|
|
||||||
|
|
@@ -108,6 +113,10 @@ async def run( | |||||
| max_turns: Maximum LLM turns allowed (must be >= 1). | ||||||
| memory_messages: Optional working memory messages to inject | ||||||
| between the system prompt and task instruction. | ||||||
| timeout_seconds: Optional wall-clock timeout in seconds. | ||||||
| When exceeded, the execution loop is cancelled and the | ||||||
| run returns with ``TerminationReason.ERROR``. Cost | ||||||
| recording and post-execution processing still occur. | ||||||
|
|
||||||
| Returns: | ||||||
| ``AgentRunResult`` with execution outcome and metadata. | ||||||
|
|
@@ -118,7 +127,8 @@ async def run( | |||||
| Raises: | ||||||
| ExecutionStateError: If pre-flight validation fails (agent | ||||||
| not ACTIVE or task not ASSIGNED/IN_PROGRESS). | ||||||
| ValueError: If ``max_turns`` is less than 1, or if | ||||||
| ``timeout_seconds`` is not positive. | ||||||
| MemoryError: Re-raised unconditionally (non-recoverable). | ||||||
| RecursionError: Re-raised unconditionally (non-recoverable). | ||||||
| """ | ||||||
|
|
@@ -135,6 +145,16 @@ async def run( | |||||
| ) | ||||||
| raise ValueError(msg) | ||||||
|
|
||||||
| if timeout_seconds is not None and timeout_seconds <= 0: | ||||||
| msg = f"timeout_seconds must be > 0, got {timeout_seconds}" | ||||||
| logger.warning( | ||||||
| EXECUTION_ENGINE_INVALID_INPUT, | ||||||
| agent_id=agent_id, | ||||||
| task_id=task_id, | ||||||
| reason=msg, | ||||||
| ) | ||||||
| raise ValueError(msg) | ||||||
|
|
||||||
| self._validate_agent(identity, agent_id) | ||||||
| self._validate_task(task, agent_id, task_id) | ||||||
|
|
||||||
|
|
@@ -167,6 +187,7 @@ async def run( | |||||
| ctx=ctx, | ||||||
| system_prompt=system_prompt, | ||||||
| start=start, | ||||||
| timeout_seconds=timeout_seconds, | ||||||
| ) | ||||||
| except MemoryError, RecursionError: | ||||||
|
There was a problem hiding this comment. Python 2
Suggested change
Prompt To Fix With AIThis is a comment left during a code review.
Path: src/ai_company/engine/agent_engine.py
Line: 192
Comment:
**Python 2 `except` comma syntax — not valid Python 3**
`except MemoryError, RecursionError:` is Python 2 syntax and is a `SyntaxError` in Python 3. Python 2 parsed this as "catch `MemoryError`, bind it to the name `RecursionError`" — it does NOT catch both exception types. The Python 3 form to catch multiple exceptions requires parentheses. This same pattern appears on lines **584** and **669** as well.
```suggestion
except (MemoryError, RecursionError):
```
How can I resolve this? If you propose a fix, please make it concise.
Contributor
There was a problem hiding this comment. The except (MemoryError, RecursionError): |
||||||
| logger.error( | ||||||
|
|
@@ -200,6 +221,7 @@ async def _execute( # noqa: PLR0913 | |||||
| ctx: AgentContext, | ||||||
| system_prompt: SystemPrompt, | ||||||
| start: float, | ||||||
| timeout_seconds: float | None = None, | ||||||
| ) -> AgentRunResult: | ||||||
| """Run the execution loop, record costs, and build result.""" | ||||||
| budget_checker = _make_budget_checker(task) | ||||||
|
|
@@ -212,16 +234,49 @@ async def _execute( # noqa: PLR0913 | |||||
| estimated_tokens=system_prompt.estimated_tokens, | ||||||
| ) | ||||||
|
|
||||||
| try: | ||||||
| coro = self._loop.execute( | ||||||
| context=ctx, | ||||||
| provider=self._provider, | ||||||
| tool_invoker=tool_invoker, | ||||||
| budget_checker=budget_checker, | ||||||
| completion_config=completion_config, | ||||||
| ) | ||||||
| if timeout_seconds is not None: | ||||||
| execution_result = await asyncio.wait_for( | ||||||
| coro, | ||||||
| timeout=timeout_seconds, | ||||||
| ) | ||||||
| else: | ||||||
| execution_result = await coro | ||||||
| except TimeoutError: | ||||||
|
There was a problem hiding this comment.
Suggested change
Contributor
There was a problem hiding this comment.
except asyncio.TimeoutError: |
||||||
| if timeout_seconds is None: | ||||||
| raise | ||||||
|
greptile-apps[bot] marked this conversation as resolved.
Outdated
|
||||||
| duration = time.monotonic() - start | ||||||
| error_msg = ( | ||||||
| f"Wall-clock timeout after {duration:.1f}s (limit: {timeout_seconds}s)" | ||||||
| ) | ||||||
| logger.warning( | ||||||
| EXECUTION_ENGINE_TIMEOUT, | ||||||
| agent_id=agent_id, | ||||||
| task_id=task_id, | ||||||
| duration_seconds=duration, | ||||||
| timeout_seconds=timeout_seconds, | ||||||
| ) | ||||||
|
coderabbitai[bot] marked this conversation as resolved.
Outdated
|
||||||
| execution_result = ExecutionResult( | ||||||
| context=ctx, | ||||||
| termination_reason=TerminationReason.ERROR, | ||||||
| error_message=error_msg, | ||||||
| ) | ||||||
|
|
||||||
| await self._record_costs(execution_result, identity, agent_id, task_id) | ||||||
| execution_result = self._apply_post_execution_transitions( | ||||||
| execution_result, | ||||||
| agent_id, | ||||||
| task_id, | ||||||
| ) | ||||||
|
|
||||||
| duration = time.monotonic() - start | ||||||
|
|
||||||
| result = AgentRunResult( | ||||||
| execution_result=execution_result, | ||||||
|
|
@@ -230,7 +285,7 @@ async def _execute( # noqa: PLR0913 | |||||
| agent_id=agent_id, | ||||||
| task_id=task_id, | ||||||
| ) | ||||||
| self._log_completion(result, agent_id, task_id, duration) | ||||||
| return result | ||||||
|
|
||||||
| # ── Setup ──────────────────────────────────────────────────── | ||||||
|
|
@@ -341,6 +396,78 @@ def _transition_task_if_needed( | |||||
| ) | ||||||
| return ctx | ||||||
|
|
||||||
| def _apply_post_execution_transitions( | ||||||
| self, | ||||||
| execution_result: ExecutionResult, | ||||||
| agent_id: str, | ||||||
| task_id: str, | ||||||
| ) -> ExecutionResult: | ||||||
| """Apply post-execution task transitions based on termination reason. | ||||||
|
|
||||||
| Only ``TerminationReason.COMPLETED`` triggers transitions: | ||||||
| IN_PROGRESS → IN_REVIEW → COMPLETED (two-hop auto-complete). | ||||||
| All other reasons leave the task in its current state. | ||||||
|
|
||||||
| Note: | ||||||
| The IN_REVIEW → COMPLETED auto-complete is M3 scaffolding | ||||||
| (no reviewers yet). Later milestones will gate COMPLETED | ||||||
| on reviewer approval. | ||||||
|
|
||||||
| Transition failures are logged but do not discard the | ||||||
| successful execution result — a bookkeeping error must never | ||||||
| destroy the agent's work. | ||||||
|
|
||||||
| Args: | ||||||
| execution_result: Result from the execution loop. | ||||||
| agent_id: Agent identifier for logging. | ||||||
| task_id: Task identifier for logging. | ||||||
|
|
||||||
| Returns: | ||||||
| New ``ExecutionResult`` with updated context if transitions | ||||||
| were applied, or the original result unchanged. | ||||||
| """ | ||||||
| ctx = execution_result.context | ||||||
| if ctx.task_execution is None: | ||||||
| return execution_result | ||||||
|
|
||||||
| if execution_result.termination_reason != TerminationReason.COMPLETED: | ||||||
| return execution_result | ||||||
|
|
||||||
| try: | ||||||
| ctx = ctx.with_task_transition( | ||||||
| TaskStatus.IN_REVIEW, | ||||||
| reason="Agent completed execution", | ||||||
| ) | ||||||
| logger.info( | ||||||
| EXECUTION_ENGINE_TASK_TRANSITION, | ||||||
| agent_id=agent_id, | ||||||
| task_id=task_id, | ||||||
| from_status=TaskStatus.IN_PROGRESS.value, | ||||||
| to_status=TaskStatus.IN_REVIEW.value, | ||||||
| ) | ||||||
| # TODO(M4): Replace auto-complete with review gate | ||||||
| ctx = ctx.with_task_transition( | ||||||
| TaskStatus.COMPLETED, | ||||||
| reason="Auto-completed (no reviewers in M3)", | ||||||
| ) | ||||||
| logger.info( | ||||||
| EXECUTION_ENGINE_TASK_TRANSITION, | ||||||
| agent_id=agent_id, | ||||||
| task_id=task_id, | ||||||
| from_status=TaskStatus.IN_REVIEW.value, | ||||||
| to_status=TaskStatus.COMPLETED.value, | ||||||
| ) | ||||||
| except (ValueError, ExecutionStateError) as exc: | ||||||
| logger.exception( | ||||||
| EXECUTION_ENGINE_ERROR, | ||||||
| agent_id=agent_id, | ||||||
| task_id=task_id, | ||||||
| error=f"Post-execution transition failed: {exc}", | ||||||
| ) | ||||||
| return execution_result | ||||||
|
|
||||||
| return execution_result.model_copy(update={"context": ctx}) | ||||||
|
|
||||||
| def _make_tool_invoker(self) -> ToolInvoker | None: | ||||||
| """Create a ToolInvoker from the registry, or None.""" | ||||||
| if self._tool_registry is None: | ||||||
|
|
@@ -350,23 +477,34 @@ def _make_tool_invoker(self) -> ToolInvoker | None: | |||||
| def _log_completion( | ||||||
| self, | ||||||
| result: AgentRunResult, | ||||||
| agent_id: str, | ||||||
| task_id: str, | ||||||
| duration: float, | ||||||
| ) -> None: | ||||||
| """Log structured completion event and proxy overhead metrics.""" | ||||||
| accumulated = result.execution_result.context.accumulated_cost | ||||||
| logger.info( | ||||||
| EXECUTION_ENGINE_COMPLETE, | ||||||
| agent_id=agent_id, | ||||||
| task_id=task_id, | ||||||
| termination_reason=result.termination_reason.value, | ||||||
| total_turns=result.total_turns, | ||||||
| total_tokens=accumulated.total_tokens, | ||||||
| duration_seconds=duration, | ||||||
| cost_usd=result.total_cost_usd, | ||||||
| ) | ||||||
|
|
||||||
| metrics = TaskCompletionMetrics.from_run_result(result) | ||||||
| logger.info( | ||||||
| EXECUTION_ENGINE_TASK_METRICS, | ||||||
| agent_id=agent_id, | ||||||
| task_id=task_id, | ||||||
|
Comment on lines
+570
to
+574
There was a problem hiding this comment.
|
||||||
| turns_per_task=metrics.turns_per_task, | ||||||
| tokens_per_task=metrics.tokens_per_task, | ||||||
| cost_per_task=metrics.cost_per_task, | ||||||
| duration_seconds=metrics.duration_seconds, | ||||||
| ) | ||||||
|
|
||||||
| async def _record_costs( | ||||||
| self, | ||||||
| result: ExecutionResult, | ||||||
|
|
@@ -545,7 +683,7 @@ def _handle_fatal_error( # noqa: PLR0913 | |||||
| error=f"Failed to build error result: {build_exc}", | ||||||
| original_error=error_msg, | ||||||
| ) | ||||||
| raise exc from build_exc | ||||||
|
greptile-apps[bot] marked this conversation as resolved.
Outdated
|
||||||
|
|
||||||
|
|
||||||
| def _format_task_instruction(task: Task) -> str: | ||||||
|
|
||||||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
DESIGN_SPEC references
asyncio.wait_forbut implementation usesasyncio.waitThe spec description (line 928) states:
However, the actual implementation in
_run_loop_with_timeout(lines 278–300) deliberately usesasyncio.waitinstead, with clear rationale: this prevents conflating an internalTimeoutErrorfrom the loop with the engine's wall-clock deadline (see the method's docstring at lines 280–282).The implementation choice is correct and well-reasoned; the spec just needs to match:
Prompt To Fix With AI