[codex] Add MCP integration

pyproject.toml

@@ -35,6 +35,7 @@ dependencies = [
     "keyring>=25.0",
     "tree-sitter>=0.25,<0.26",
     "tree-sitter-bash>=0.25,<0.26",
+    "mcp>=1.28.0",
 ]
 
 [project.optional-dependencies]

src/iac_code/a2a/events.py

@@ -24,8 +24,10 @@
 )
 from iac_code.a2a.exposure import A2AExposureType, normalize_a2a_exposure_types
 from iac_code.a2a.metadata_redaction import A2AMetadataEchoRedactor
+from iac_code.i18n import _
 from iac_code.types.stream_events import (
     ErrorEvent,
+    MCPProgressEvent,
     MessageEndEvent,
     PermissionRequestEvent,
     TextDeltaEvent,
@@ -79,6 +81,45 @@ def make_text_part(text: str) -> Part:
     return Part(text=text)
 
 
+async def publish_mcp_warnings(
+    event_queue: Any,
+    *,
+    task_id: str,
+    context_id: str,
+    runtime: Any,
+    state: int = TaskState.TASK_STATE_WORKING,
+    iac_code_session_id: str | None = None,
+) -> None:
+    warnings = list(getattr(runtime, "mcp_config_warnings", None) or [])
+    pushed_count = getattr(runtime, "_a2a_mcp_warnings_pushed_count", 0)
+    if pushed_count >= len(warnings):
+        return
+    for warning in warnings[pushed_count:]:
+        message = str(getattr(warning, "message", None) or warning)
+        await _enqueue_status(
+            event_queue,
+            task_id=task_id,
+            context_id=context_id,
+            state=state,
+            message=_agent_text_message(
+                task_id=task_id,
+                context_id=context_id,
+                text=_("MCP warning: {message}").format(message=message),
+            ),
+            metadata={
+                "iac_code": {
+                    "mcpWarning": {
+                        "serverName": str(getattr(warning, "server_name", "")),
+                        "code": str(getattr(warning, "code", "")),
+                        "message": message,
+                    }
+                }
+            },
+            iac_code_session_id=iac_code_session_id,
+        )
+    setattr(runtime, "_a2a_mcp_warnings_pushed_count", len(warnings))
+
+
 def _extract_artifact_metadata(result: Any, artifact_store: Any | None) -> dict[str, Any] | None:
     if artifact_store is None or not isinstance(result, dict):
         return None
@@ -303,6 +344,31 @@ async def publish_stream_event(
         )
         return None
 
+    if isinstance(event, MCPProgressEvent):
+        if A2AExposureType.TOOL_TRACE not in enabled_exposure_types:
+            return None
+        progress_metadata = {
+            "status": "progress",
+            "toolUseId": event.tool_use_id or "",
+            "name": "mcp__{}__{}".format(event.server_name, event.tool_name),
+            "mcp": {
+                "serverName": event.server_name,
+                "toolName": event.tool_name,
+                "progress": event.progress,
+                "total": event.total,
+                "message": _truncate(event.message) if event.message else None,
+            },
+        }
+        await _enqueue_status(
+            event_queue,
+            task_id=task_id,
+            context_id=context_id,
+            state=TaskState.TASK_STATE_WORKING,
+            metadata={"iac_code": {"tool": progress_metadata}},
+            iac_code_session_id=iac_code_session_id,
+        )
+        return None
+
     if isinstance(event, PermissionRequestEvent):
         approved = auto_approve_permissions
         if permission_resolver is not None:

src/iac_code/a2a/executor.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import json
 import logging
 import os
@@ -19,6 +20,7 @@
 from iac_code.a2a.events import (
     iac_code_session_metadata,
     make_text_part,
+    publish_mcp_warnings,
     publish_stream_event,
     with_iac_code_session_metadata,
 )
@@ -44,7 +46,7 @@
     credentials_with_metadata_api_key,
     refresh_runtime_cloud_tools,
 )
-from iac_code.a2a.task_store import A2ATaskStore
+from iac_code.a2a.task_store import A2ATaskStore, _close_runtime
 from iac_code.a2a.types import (
     TASK_STATE_CANCELED,
     TASK_STATE_FAILED,
@@ -934,6 +936,10 @@ async def publish_initial_task_if_missing() -> None:
         if route_pipeline_handoff_to_normal:
             await self._ensure_pipeline_handoff_context_in_session(context_id=context_id, cwd=cwd)
 
+        task.active_task = asyncio.current_task()
+        task.state = TASK_STATE_WORKING
+        self._task_store.mirror_task(task)
+
         def runtime_factory(session_id: str) -> Any:
             session_storage = SessionStorage()
             resume_messages = None
@@ -957,10 +963,28 @@ def runtime_factory(session_id: str) -> Any:
                     runtime_factory=runtime_factory,
                 )
                 if not hasattr(ctx.runtime, "agent_loop"):
+                    old_runtime = ctx.runtime
                     ctx.runtime = runtime_factory(ctx.session_id)
                     self._task_store.mirror_context(ctx)
+                    await _close_runtime(old_runtime)
+        except asyncio.CancelledError:
+            task.active_task = None
+            task.state = TASK_STATE_CANCELED
+            self._task_store.mirror_task(task)
+            with contextlib.suppress(Exception):
+                await self._publish_status(
+                    event_queue,
+                    task_id=task_id,
+                    context_id=context_id,
+                    state=TaskState.TASK_STATE_CANCELED,
+                    text=_("Task canceled."),
+                )
+                await self._notify_terminal_task(task_id=task.task_id, context_id=task.context_id, state=task.state)
+            self._metrics.record_task_canceled()
+            raise
         except Exception as exc:
             self._log_executor_exception("runtime setup", task_id=task_id, context_id=context_id)
+            task.active_task = None
             await self._publish_status(
                 event_queue,
                 task_id=task_id,
@@ -978,6 +1002,7 @@ def runtime_factory(session_id: str) -> Any:
         if ctx.lock is None:
             ctx.lock = asyncio.Lock()
         if ctx.active_task_id is not None:
+            task.active_task = None
             task.state = TASK_STATE_FAILED
             await self._publish_status(
                 event_queue,
@@ -995,7 +1020,24 @@ def runtime_factory(session_id: str) -> Any:
         lock = ctx.lock
         try:
             await asyncio.wait_for(lock.acquire(), timeout=_CONTEXT_LOCK_ACQUIRE_TIMEOUT_SECONDS)
+        except asyncio.CancelledError:
+            task.active_task = None
+            task.state = TASK_STATE_CANCELED
+            self._task_store.mirror_task(task)
+            with contextlib.suppress(Exception):
+                await self._publish_status(
+                    event_queue,
+                    task_id=task_id,
+                    context_id=context_id,
+                    state=TaskState.TASK_STATE_CANCELED,
+                    text=_("Task canceled."),
+                    session_id=ctx.session_id,
+                )
+                await self._notify_terminal_task(task_id=task.task_id, context_id=task.context_id, state=task.state)
+            self._metrics.record_task_canceled()
+            raise
         except TimeoutError:
+            task.active_task = None
             task.state = TASK_STATE_FAILED
             await self._publish_status(
                 event_queue,
@@ -1020,6 +1062,13 @@ def runtime_factory(session_id: str) -> Any:
                 runtime = ctx.runtime
                 if runtime is None:
                     raise RuntimeError("A2A context runtime missing")
+                await publish_mcp_warnings(
+                    event_queue,
+                    task_id=task_id,
+                    context_id=context_id,
+                    runtime=runtime,
+                    iac_code_session_id=ctx.session_id,
+                )
                 await self._publish_status(
                     event_queue,
                     task_id=task_id,
@@ -1070,6 +1119,13 @@ def runtime_factory(session_id: str) -> Any:
                         session_id=ctx.session_id,
                     )
                     async for event in stream:
+                        await publish_mcp_warnings(
+                            event_queue,
+                            task_id=task_id,
+                            context_id=context_id,
+                            runtime=runtime,
+                            iac_code_session_id=ctx.session_id,
+                        )
                         text_chunk = await publish_stream_event(
                             event_queue,
                             task_id=task_id,
@@ -1083,6 +1139,13 @@ def runtime_factory(session_id: str) -> Any:
                         )
                         if text_chunk:
                             task.output_text.append(text_chunk)
+                    await publish_mcp_warnings(
+                        event_queue,
+                        task_id=task_id,
+                        context_id=context_id,
+                        runtime=runtime,
+                        iac_code_session_id=ctx.session_id,
+                    )
                 task.state = TASK_STATE_INPUT_REQUIRED
                 await self._publish_status(
                     event_queue,

src/iac_code/a2a/pipeline_events.py

@@ -15,6 +15,7 @@
     AskUserQuestionEvent,
     CandidateDetailEvent,
     DiagramEvent,
+    MCPProgressEvent,
     PermissionRequestEvent,
     SubPipelineStreamEvent,
     TextDeltaEvent,
@@ -148,6 +149,10 @@ def __init__(self, context: PipelineA2AContext):
     def last_sequence(self) -> int:
         return self._sequence
 
+    @property
+    def context(self) -> PipelineA2AContext:
+        return self._context
+
     def hydrate_from_events(self, events: list[dict[str, Any]]) -> None:
         latest_parent_step_sequence = -1
         for event in events:
@@ -273,6 +278,8 @@ def translate(self, event: Any) -> list[dict[str, Any]]:
             return [self._translate_diagram_event(event)]
         if isinstance(event, ToolResultEvent):
             return self._translate_tool_result_event(event)
+        if isinstance(event, MCPProgressEvent):
+            return [self._translate_parent_scoped_display_event("tool_progress", _mcp_progress_data(event))]
         if isinstance(event, SubPipelineStreamEvent):
             return self._translate_sub_pipeline_stream_event(event)
         return []
@@ -539,6 +546,11 @@ def _translate_sub_pipeline_stream_event(self, event: SubPipelineStreamEvent) ->
             data = _tool_result_data(inner)
             input_data = None
             permission = None
+        elif isinstance(inner, MCPProgressEvent):
+            event_type = "tool_progress"
+            data = _mcp_progress_data(inner)
+            input_data = None
+            permission = None
         else:
             return []
 
@@ -1241,6 +1253,22 @@ def _tool_result_data(event: ToolResultEvent) -> dict[str, Any]:
     }
 
 
+def _mcp_progress_data(event: MCPProgressEvent) -> dict[str, Any]:
+    data: dict[str, Any] = {
+        "toolName": "mcp__{}__{}".format(event.server_name, event.tool_name),
+        "toolUseId": event.tool_use_id or "",
+        "serverName": event.server_name,
+        "mcpToolName": event.tool_name,
+    }
+    if event.progress is not None:
+        data["progress"] = event.progress
+    if event.total is not None:
+        data["total"] = event.total
+    if event.message:
+        data["message"] = _truncate(event.message)
+    return data
+
+
 def _completion_step_id(envelope: dict[str, Any]) -> str | None:
     candidate_step = envelope.get("candidateStep")
     if isinstance(candidate_step, dict):

src/iac_code/a2a/pipeline_executor.py

@@ -13,7 +13,7 @@
 from a2a.types import Message, Role, TaskState, TaskStatus, TaskStatusUpdateEvent
 from a2a.utils.errors import InvalidParamsError
 
-from iac_code.a2a.events import make_text_part
+from iac_code.a2a.events import make_text_part, publish_mcp_warnings
 from iac_code.a2a.pipeline_events import PipelineA2AContext, PipelineEventTranslator
 from iac_code.a2a.pipeline_journal import A2APipelineJournal
 from iac_code.a2a.pipeline_paths import (
@@ -275,6 +275,13 @@ def runtime_factory(session_id: str) -> Any:
                         cwd=cwd,
                     )
                     agent_runtime = pipeline_runtime.agent_runtime
+                    await publish_mcp_warnings(
+                        event_queue,
+                        task_id=task_id,
+                        context_id=context_id,
+                        runtime=agent_runtime,
+                        iac_code_session_id=ctx.session_id,
+                    )
                     configure_runtime_model(
                         agent_runtime,
                         self._model,
@@ -790,11 +797,25 @@ async def _consume_stream_until_restart(
                     event = await next_task
                 except StopAsyncIteration:
                     next_task = None
+                    await publish_mcp_warnings(
+                        publisher.event_queue,
+                        task_id=publisher.translator.context.task_id,
+                        context_id=publisher.translator.context.context_id,
+                        runtime=runtime.agent_runtime,
+                        iac_code_session_id=publisher.translator.context.iac_code_session_id,
+                    )
                     return _StreamConsumeResult(had_events=had_events, restart_requested=False)
                 finally:
                     next_task = None
 
                 had_events = True
+                await publish_mcp_warnings(
+                    publisher.event_queue,
+                    task_id=publisher.translator.context.task_id,
+                    context_id=publisher.translator.context.context_id,
+                    runtime=runtime.agent_runtime,
+                    iac_code_session_id=publisher.translator.context.iac_code_session_id,
+                )
                 text = await publisher.publish(
                     event,
                     permission_resolver=self._permission_resolver,