codex-a2a 1.3.3__tar.gz → 1.3.4__tar.gz

{codex_a2a-1.3.3 → codex_a2a-1.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: codex-a2a
-Version: 1.3.3
+Version: 1.3.4
 Summary: Codex A2A runtime adapter
 Author: liujuanjuan1984@Intelligent-Internet
 License-Expression: Apache-2.0

{codex_a2a-1.3.3 → codex_a2a-1.3.4}/docs/maintainer-architecture.md

@@ -74,6 +74,16 @@ flowchart TD
 -   **Session State**: Stores the binding of `context_id` to `session_id`.
 -   **Interrupt Store**: Stores pending interrupt requests to survive service restarts.
 
+## Background Streams and Shutdown
+
+Background streams use SDK Source/Sink queues: producers persist once before
+broadcasting, and slow or disconnected subscribers cannot block other consumers.
+Shutdown drains SDK and adapter tasks before closing clients and stores.
+
+Use the [standard upgrade procedure](./guide.md#released-cli-self-start) and restart;
+the task-store format is unchanged. SDK upgrade regressions run in normal CI and
+can be checked with `uv run pytest --no-cov tests/server/test_sdk_upgrade.py`.
+
 ## Configuration Layering
 
 Configuration is handled in `src/codex_a2a/config.py` using `pydantic-settings`. It is categorized by prefix:

{codex_a2a-1.3.3 → codex_a2a-1.3.4}/src/codex_a2a/execution/discovery_runtime.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import asyncio
 import logging
 import uuid
+from contextlib import aclosing
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
@@ -94,22 +95,27 @@ class CodexDiscoveryRuntime:
 
     async def _run_watch(self, *, handle: DiscoveryWatchHandle, event_queue) -> None:  # noqa: ANN001
         append = False
-        async for event in self._client.stream_events(stop_event=handle.stop_event):
-            payload = self._payload_from_event(event, events=handle.events)
-            if payload is None:
-                continue
-            await enqueue_artifact_update(
-                event_queue=event_queue,
-                task_id=handle.task_id,
-                context_id=handle.context_id,
-                artifact_id=f"{handle.task_id}:discovery",
-                part=new_data_part(payload),
-                append=append,
-                last_chunk=None,
-                artifact_metadata={"codex": {"discovery_watch": {"events": sorted(handle.events)}}},
-                event_metadata={"codex": {"discovery_watch": {"events": sorted(handle.events)}}},
-            )
-            append = True
+        async with aclosing(self._client.stream_events(stop_event=handle.stop_event)) as events:
+            async for event in events:
+                payload = self._payload_from_event(event, events=handle.events)
+                if payload is None:
+                    continue
+                await enqueue_artifact_update(
+                    event_queue=event_queue,
+                    task_id=handle.task_id,
+                    context_id=handle.context_id,
+                    artifact_id=f"{handle.task_id}:discovery",
+                    part=new_data_part(payload),
+                    append=append,
+                    last_chunk=None,
+                    artifact_metadata={
+                        "codex": {"discovery_watch": {"events": sorted(handle.events)}}
+                    },
+                    event_metadata={
+                        "codex": {"discovery_watch": {"events": sorted(handle.events)}}
+                    },
+                )
+                append = True
 
     def _payload_from_event(
         self,

{codex_a2a-1.3.3 → codex_a2a-1.3.4}/src/codex_a2a/execution/exec_runtime.py

@@ -278,10 +278,16 @@ class CodexExecRuntime:
             )
         finally:
             handle.stop_event.set()
+            child_tasks: list[asyncio.Future] = [exec_task]
             if pending_event_task is not None:
-                pending_event_task.cancel()
-                with contextlib.suppress(asyncio.CancelledError):
-                    await pending_event_task
+                child_tasks.append(pending_event_task)
+            for child_task in child_tasks:
+                if not child_task.done() and not (
+                    isinstance(child_task, asyncio.Task) and child_task.cancelling()
+                ):
+                    child_task.cancel()
+            await asyncio.gather(*child_tasks, return_exceptions=True)
+            await stream_iter.aclose()
 
     async def _emit_output_delta(
         self,

{codex_a2a-1.3.3 → codex_a2a-1.3.4}/src/codex_a2a/execution/review_runtime.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import asyncio
 import logging
 import uuid
+from contextlib import aclosing
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
@@ -132,39 +133,40 @@ class CodexReviewRuntime:
             )
             append = True
 
-        async for event in self._client.stream_events(stop_event=handle.stop_event):
-            payload = self._payload_from_event(event, handle=handle)
-            if payload is None:
-                continue
-            is_terminal = payload["event"] in {"review.completed", "review.failed"}
-            await enqueue_artifact_update(
-                event_queue=event_queue,
-                task_id=handle.task_id,
-                context_id=handle.context_id,
-                artifact_id=f"{handle.task_id}:review-watch",
-                part=new_data_part(payload),
-                append=append,
-                last_chunk=is_terminal,
-                artifact_metadata=metadata,
-                event_metadata=metadata,
-            )
-            append = True
-            if is_terminal:
-                await event_queue.enqueue_event(
-                    TaskStatusUpdateEvent(
-                        task_id=handle.task_id,
-                        context_id=handle.context_id,
-                        status=TaskStatus(
-                            state=(
-                                TaskState.TASK_STATE_COMPLETED
-                                if payload["event"] == "review.completed"
-                                else TaskState.TASK_STATE_FAILED
+        async with aclosing(self._client.stream_events(stop_event=handle.stop_event)) as events:
+            async for event in events:
+                payload = self._payload_from_event(event, handle=handle)
+                if payload is None:
+                    continue
+                is_terminal = payload["event"] in {"review.completed", "review.failed"}
+                await enqueue_artifact_update(
+                    event_queue=event_queue,
+                    task_id=handle.task_id,
+                    context_id=handle.context_id,
+                    artifact_id=f"{handle.task_id}:review-watch",
+                    part=new_data_part(payload),
+                    append=append,
+                    last_chunk=is_terminal,
+                    artifact_metadata=metadata,
+                    event_metadata=metadata,
+                )
+                append = True
+                if is_terminal:
+                    await event_queue.enqueue_event(
+                        TaskStatusUpdateEvent(
+                            task_id=handle.task_id,
+                            context_id=handle.context_id,
+                            status=TaskStatus(
+                                state=(
+                                    TaskState.TASK_STATE_COMPLETED
+                                    if payload["event"] == "review.completed"
+                                    else TaskState.TASK_STATE_FAILED
+                                ),
                             ),
-                        ),
-                        metadata=metadata,
+                            metadata=metadata,
+                        )
                     )
-                )
-                break
+                    break
 
     def _started_payload(self, handle: ReviewWatchHandle) -> dict[str, Any] | None:
         if "review.started" not in handle.events:

{codex_a2a-1.3.3 → codex_a2a-1.3.4}/src/codex_a2a/execution/thread_lifecycle_runtime.py

@@ -5,6 +5,7 @@ import hashlib
 import json
 import logging
 import uuid
+from contextlib import aclosing
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, TypeGuard
 
@@ -212,22 +213,23 @@ class CodexThreadLifecycleRuntime:
             }
         }
         try:
-            async for event in self._client.stream_events(stop_event=handle.stop_event):
-                payload = self._payload_from_event(event, handle=handle)
-                if payload is None:
-                    continue
-                await enqueue_artifact_update(
-                    event_queue=event_queue,
-                    task_id=handle.task_id,
-                    context_id=handle.context_id,
-                    artifact_id=f"{handle.task_id}:thread-lifecycle",
-                    part=new_data_part(payload),
-                    append=append,
-                    last_chunk=None,
-                    artifact_metadata=metadata,
-                    event_metadata=metadata,
-                )
-                append = True
+            async with aclosing(self._client.stream_events(stop_event=handle.stop_event)) as events:
+                async for event in events:
+                    payload = self._payload_from_event(event, handle=handle)
+                    if payload is None:
+                        continue
+                    await enqueue_artifact_update(
+                        event_queue=event_queue,
+                        task_id=handle.task_id,
+                        context_id=handle.context_id,
+                        artifact_id=f"{handle.task_id}:thread-lifecycle",
+                        part=new_data_part(payload),
+                        append=append,
+                        last_chunk=None,
+                        artifact_metadata=metadata,
+                        event_metadata=metadata,
+                    )
+                    append = True
         except asyncio.CancelledError:
             release_reason = "task_cancel"
             raise

{codex_a2a-1.3.3 → codex_a2a-1.3.4}/src/codex_a2a/server/application.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from contextlib import asynccontextmanager
+from contextlib import AsyncExitStack, asynccontextmanager
 from typing import Any
 
 import uvicorn
@@ -288,22 +288,24 @@ def create_app(settings: Settings) -> FastAPI:
             persistence_summary["database_url"],
             persistence_summary["sqlite_tuning"],
         )
-        await task_store_runtime.startup()
-        await push_config_store_runtime.startup()
-        await runtime_state_runtime.startup()
-        try:
+        async with AsyncExitStack() as cleanup:
+            if shared_database_engine is not None:
+                cleanup.push_async_callback(shared_database_engine.dispose)
+            await task_store_runtime.startup()
+            cleanup.push_async_callback(task_store_runtime.shutdown)
+            await push_config_store_runtime.startup()
+            cleanup.push_async_callback(push_config_store_runtime.shutdown)
+            await runtime_state_runtime.startup()
+            cleanup.push_async_callback(runtime_state_runtime.shutdown)
+            cleanup.push_async_callback(client.close)
+            cleanup.push_async_callback(a2a_client_manager.close_all)
+            # Callbacks run in reverse order. Producers may still need clients
+            # and stores while unwinding, even if another cleanup callback fails.
+            cleanup.push_async_callback(handler.aclose)
             await client.restore_persisted_interrupt_requests()
             await client.startup_preflight()
             await thread_lifecycle_runtime.reconcile_persisted_watches()
             yield
-        finally:
-            await a2a_client_manager.close_all()
-            await client.close()
-            await runtime_state_runtime.shutdown()
-            await push_config_store_runtime.shutdown()
-            await task_store_runtime.shutdown()
-            if shared_database_engine is not None:
-                await shared_database_engine.dispose()
 
     context_builder = IdentityAwareCallContextBuilder()
     jsonrpc_methods = {

{codex_a2a-1.3.3 → codex_a2a-1.3.4}/src/codex_a2a/server/output_negotiation.py

@@ -1,19 +1,15 @@
 from __future__ import annotations
 
-import asyncio
 import json
 from collections.abc import Iterable
 from typing import Any, cast
 
-from a2a.server.events import EventConsumer
-from a2a.server.tasks import ResultAggregator, TaskManager
 from a2a.types import (
     Artifact,
     Message,
     Part,
     Task,
     TaskArtifactUpdateEvent,
-    TaskState,
     TaskStatusUpdateEvent,
 )
 from google.protobuf.message import Message as ProtoMessage  # type: ignore[import-untyped]
@@ -168,111 +164,6 @@ def apply_accepted_output_modes(
     return payload
 
 
-class NegotiatingResultAggregator(ResultAggregator):
-    def __init__(
-        self,
-        task_manager: TaskManager,
-        accepted_output_modes: Iterable[str] | None,
-    ) -> None:
-        super().__init__(task_manager)
-        self._accepted_output_modes = normalize_accepted_output_modes(accepted_output_modes)
-
-    def _transform_event(self, event: Any) -> Any | None:
-        negotiated_event = apply_accepted_output_modes(event, self._accepted_output_modes)
-        if negotiated_event is None:
-            return None
-        return annotate_output_negotiation_metadata(negotiated_event, self._accepted_output_modes)
-
-    async def _persist_output_negotiation_metadata(self, event: Any) -> None:
-        if not isinstance(event, TaskArtifactUpdateEvent):
-            return
-
-        accepted_output_modes = extract_accepted_output_modes_from_metadata(event.metadata)
-        if accepted_output_modes is None:
-            return
-
-        task = await self.task_manager.ensure_task(event)
-        merged_metadata = merge_output_negotiation_metadata(task.metadata, accepted_output_modes)
-        if merged_metadata == task.metadata:
-            return
-        task.ClearField("metadata")
-        if merged_metadata is not None:
-            task.metadata.CopyFrom(merged_metadata)
-        await self.task_manager._save_task(task)
-
-    async def consume_and_emit(self, consumer: EventConsumer):  # noqa: ANN201
-        async for event in consumer.consume_all():
-            transformed_event = self._transform_event(event)
-            if transformed_event is None:
-                continue
-            await self._persist_output_negotiation_metadata(transformed_event)
-            await self.task_manager.process(transformed_event)
-            yield transformed_event
-
-    async def consume_all(self, consumer: EventConsumer) -> Task | Message | None:
-        async for event in consumer.consume_all():
-            transformed_event = self._transform_event(event)
-            if transformed_event is None:
-                continue
-            if isinstance(transformed_event, Message):
-                self._message = transformed_event
-                return transformed_event
-            await self._persist_output_negotiation_metadata(transformed_event)
-            await self.task_manager.process(transformed_event)
-        return await self.task_manager.get_task()
-
-    async def consume_and_break_on_interrupt(
-        self,
-        consumer: EventConsumer,
-        blocking: bool = True,
-        event_callback=None,  # noqa: ANN001
-    ) -> tuple[Task | Message | None, bool, asyncio.Task | None]:
-        event_stream = consumer.consume_all()
-        interrupted = False
-        bg_task: asyncio.Task | None = None
-        async for event in event_stream:
-            transformed_event = self._transform_event(event)
-            if transformed_event is None:
-                continue
-            if isinstance(transformed_event, Message):
-                self._message = transformed_event
-                return transformed_event, False, None
-            await self._persist_output_negotiation_metadata(transformed_event)
-            await self.task_manager.process(transformed_event)
-
-            should_interrupt = False
-            is_auth_required = (
-                isinstance(transformed_event, Task | TaskStatusUpdateEvent)
-                and transformed_event.status.state == TaskState.TASK_STATE_AUTH_REQUIRED
-            )
-            if is_auth_required:
-                should_interrupt = True
-            elif not blocking:
-                should_interrupt = True
-
-            if should_interrupt:
-                bg_task = asyncio.create_task(
-                    self._continue_consuming(event_stream, event_callback)
-                )
-                interrupted = True
-                break
-        return await self.task_manager.get_task(), interrupted, bg_task
-
-    async def _continue_consuming(
-        self,
-        event_stream,
-        event_callback=None,  # noqa: ANN001
-    ) -> None:
-        async for event in event_stream:
-            transformed_event = self._transform_event(event)
-            if transformed_event is None:
-                continue
-            await self._persist_output_negotiation_metadata(transformed_event)
-            await self.task_manager.process(transformed_event)
-            if event_callback:
-                await event_callback()
-
-
 def _filter_task(task: Task, accepted_modes: frozenset[str]) -> Task:
     filtered_history: list[Message] = []
     for message in task.history or []: