fix(server): harden task lifecycle and database upgrades

docs/guide.md

@@ -217,6 +217,8 @@ The package declares `sqlalchemy[asyncio]` so standalone installs such as `uv to
 
 Published package metadata also enforces the security floors `click>=8.3.3` and `pyasn1>=0.6.4`. These apply to standalone wheel/sdist installs as well as repository environments; no separate uv constraint configuration is needed.
 
+Unhandled execution errors and shutdown interruptions persist `FAILED` when storage is available, preserving existing terminal states. Shutdown drains buffered output before closing clients and stores. SSE disconnects allow execution to continue; use `CancelTask` to request an upstream abort before stopping the application.
+
 ### SQLite Persistence Hardening
 
 File-backed SQLite databases are hardened at startup and on every new connection on POSIX systems:
@@ -241,7 +243,7 @@ Database-backed task persistence also keeps the existing first-terminal-state-wi
 
 At startup, the runtime logs a concise persistence summary covering the active backend, the redacted database URL when applicable, the shared persistence scope, and whether the SQLite local durability profile is active.
 
-The adapter-owned state tables listed above remain managed by the internal migration runner. The SDK-owned `tasks` table does not use runtime auto-migration here; upgrade existing SDK task schemas explicitly with upstream `a2a-db` before starting the service after an SDK schema change. If `a2a-db` is unavailable in your environment, install the `a2a-sdk[db-cli]` extra first.
+SDK task schemas require an explicit [upgrade](#upgrading-an-existing-database) when incompatible; startup then converts legacy task payloads transactionally. Existing v1 rows and task ownership are preserved.
 
 In-flight asyncio locks, outbound A2A client caches, and stream-local aggregation buffers remain process-local runtime state.
 
@@ -251,6 +253,21 @@ To opt into an ephemeral development profile, set:
 A2A_TASK_STORE_BACKEND=memory
 ```
 
+### Upgrading an Existing Database
+
+If startup reports an incompatible task schema, stop the application and back up SQLite, then use the installed release's migration tool:
+
+```bash
+OPENCODE_A2A_VERSION="$(opencode-a2a --version | awk '{print $2}')"
+uvx --from "opencode-a2a[migrations]==${OPENCODE_A2A_VERSION}" a2a-db upgrade head \
+  --database-url 'sqlite+aiosqlite:////absolute/path/opencode-a2a.db' \
+  --add_columns_owner_last_updated-default-owner automation
+```
+
+For source checkouts, use `uv run --extra migrations a2a-db` with the same arguments. Replace `automation` with the intended bearer principal or Basic username: adding an owner column assigns all old tasks to that identity; existing owners are unchanged. Multiple historical owners require an explicit mapping before migration.
+
+Restart and verify an old task under its owner. Startup converts legacy payloads and fills missing query timestamps; malformed payloads roll back conversion and prevent startup. Restore the backup to roll back the upgrade. Fresh installations and compatible schemas need no migration command.
+
 ## Troubleshooting Provider Auth State
 
 If one deployment works while another fails against the same upstream provider, check the deployed OpenCode user's local state before assuming the difference comes from the `opencode-a2a` package itself.

pyproject.toml

@@ -43,6 +43,9 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+migrations = [
+    "a2a-sdk[db-cli]==1.1.5",
+]
 dev = [
     "mypy>=1.19.1,<3.0",
     "pip-audit>=2.9,<3.0",

src/opencode_a2a/server/application.py

@@ -9,16 +9,11 @@
 import uvicorn
 from a2a.server.agent_execution import RequestContext
 from a2a.server.events import EventConsumer, EventQueueLegacy
-from a2a.server.request_handlers.default_request_handler import (
-    TERMINAL_TASK_STATES,
-    LegacyRequestHandler,
-)
 from a2a.server.request_handlers.response_helpers import agent_card_to_dict
 from a2a.server.routes.common import DefaultServerCallContextBuilder
 from a2a.server.routes.rest_dispatcher import RestDispatcher
 from a2a.server.tasks import ResultAggregator, TaskManager
 from a2a.types import (
-    AgentCard,
     Artifact,
     CancelTaskRequest,
     GetTaskRequest,
@@ -46,7 +41,7 @@
     A2AError,
     ErrorMapping,
 )
-from a2a.utils.task import apply_history_length, validate_history_length
+from a2a.utils.task import apply_history_length
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse, PlainTextResponse
 from google.protobuf.json_format import MessageToDict, ParseDict, ParseError
@@ -88,6 +83,7 @@
 )
 from ..profile.runtime import build_runtime_profile
 from ..protocol_versions import A2A_PROTOCOL_VERSION
+from ..task_states import TERMINAL_TASK_STATES
 from ..trace_context import install_log_record_factory
 from .agent_card import (
     _CHAT_OUTPUT_MODES,
@@ -96,6 +92,7 @@
 from .client_manager import A2AClientManager
 from .context_helpers import AuthenticatedIdentityUser
 from .database import build_database_engine
+from .handler_lifecycle import ManagedLegacyRequestHandler
 from .lifespan import build_lifespan
 from .middleware import (
     build_agent_card_etag,
@@ -234,46 +231,15 @@ def _parse_rest_send_message_request(body: bytes):
 
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncIterator, Awaitable, Callable
+    from collections.abc import AsyncIterator, Callable
 
-    from a2a.server.agent_execution import AgentExecutor, RequestContextBuilder
     from a2a.server.context import ServerCallContext
     from a2a.server.request_handlers.request_handler import RequestHandler
     from a2a.server.routes.common import ServerCallContextBuilder
-    from a2a.server.tasks import (
-        PushNotificationConfigStore,
-        PushNotificationSender,
-        TaskStore,
-    )
-
 
-class OpencodeRequestHandler(LegacyRequestHandler):
-    """Custom request handler to gracefully handle client disconnects and prevent dead loops."""
 
-    def __init__(  # noqa: PLR0913
-        self,
-        agent_executor: AgentExecutor,
-        task_store: TaskStore,
-        agent_card: AgentCard,
-        queue_manager: Any | None = None,
-        push_config_store: PushNotificationConfigStore | None = None,
-        push_sender: PushNotificationSender | None = None,
-        request_context_builder: RequestContextBuilder | None = None,
-        extended_agent_card: AgentCard | None = None,
-        extended_card_modifier: Callable[[AgentCard, ServerCallContext], Awaitable[AgentCard]]
-        | None = None,
-    ) -> None:
-        super().__init__(
-            agent_executor=agent_executor,
-            task_store=task_store,
-            agent_card=agent_card,
-            queue_manager=queue_manager,
-            push_config_store=push_config_store,
-            push_sender=push_sender,
-            request_context_builder=request_context_builder,
-            extended_agent_card=extended_agent_card,
-            extended_card_modifier=extended_card_modifier,
-        )
+class OpencodeRequestHandler(ManagedLegacyRequestHandler):
+    """Preserve adapter contracts around the SDK's legacy execution pipeline."""
 
     @staticmethod
     def _task_store_failure_message(operation: str) -> str:
@@ -522,11 +488,8 @@ async def on_get_task(
         context=None,
     ) -> Task | None:
         try:
-            validate_history_length(params)
-            task = await self.task_store.get(params.id, context)
-            if not task:
-                raise TaskNotFoundError()
-            return self._apply_task_output_negotiation(apply_history_length(task, params), context)
+            task = await super().on_get_task(params, context)
+            return self._apply_task_output_negotiation(task, context) if task is not None else None
         except TaskStoreOperationError as exc:
             raise self._task_store_server_error(exc) from exc
 
@@ -682,6 +645,8 @@ async def on_message_send_stream(self, params, context=None):
                     self._validate_task_id_match(task_id, event.id)
                 await self._send_push_notification_if_needed(task_id, event)
                 yield event
+            if producer_task.done() and not producer_task.cancelled():
+                producer_task.result()
             stream_completed = True
         except TaskStoreOperationError as exc:
             logger.exception(
@@ -708,7 +673,7 @@ async def on_message_send_stream(self, params, context=None):
                 )
             )
             detached_task.set_name(f"continue_stream_after_disconnect:{task_id}")
-            self._track_background_task(detached_task)
+            self._track_consumer_task(detached_task)
             stream_detached = True
             raise
         finally:
@@ -744,22 +709,47 @@ async def on_message_send(self, params, context=None):
         interrupted_or_non_blocking = False
         bg_consume_task: asyncio.Task | None = None
         try:
+            try:
 
-            async def push_notification_callback() -> None:
-                await self._send_push_notification_if_needed(task_id, result_aggregator)
-
-            (
-                result,
-                interrupted_or_non_blocking,
-                bg_consume_task,
-            ) = await result_aggregator.consume_and_break_on_interrupt(
-                consumer,
-                blocking=blocking,
-                event_callback=push_notification_callback,
-            )
-            if bg_consume_task is not None:
-                bg_consume_task.set_name(f"continue_consuming:{task_id}")
-                self._track_background_task(bg_consume_task)
+                async def push_notification_callback() -> None:
+                    await self._send_push_notification_if_needed(task_id, result_aggregator)
+
+                (
+                    result,
+                    interrupted_or_non_blocking,
+                    bg_consume_task,
+                ) = await result_aggregator.consume_and_break_on_interrupt(
+                    consumer,
+                    blocking=blocking,
+                    event_callback=push_notification_callback,
+                )
+                if bg_consume_task is not None:
+                    bg_consume_task.set_name(f"continue_consuming:{task_id}")
+                    self._track_consumer_task(bg_consume_task)
+            finally:
+                if interrupted_or_non_blocking:
+                    cleanup_task = asyncio.create_task(
+                        self._cleanup_producer(producer_task, task_id)
+                    )
+                    cleanup_task.set_name(f"cleanup_producer:{task_id}")
+                    self._track_background_task(cleanup_task)
+                else:
+                    try:
+                        current_task = asyncio.current_task()
+                        if current_task is not None and current_task.cancelled():
+                            logger.debug(
+                                "Client disconnected from message request. Cancelling task %s",
+                                task_id,
+                            )
+                            producer_task.cancel()
+                            await queue.close(immediate=True)
+
+                        await asyncio.shield(self._cleanup_producer(producer_task, task_id))
+                    except asyncio.CancelledError:
+                        pass
+
+            if producer_task.done() and not producer_task.cancelled():
+                producer_task.result()
         except TaskStoreOperationError as exc:
             logger.exception(
                 "Task store operation failed during SendMessage task_id=%s operation=%s",
@@ -774,24 +764,6 @@ async def push_notification_callback() -> None:
         except Exception:
             logger.exception("Agent execution failed")
             raise
-        finally:
-            if interrupted_or_non_blocking:
-                cleanup_task = asyncio.create_task(self._cleanup_producer(producer_task, task_id))
-                cleanup_task.set_name(f"cleanup_producer:{task_id}")
-                self._track_background_task(cleanup_task)
-            else:
-                try:
-                    current_task = asyncio.current_task()
-                    if current_task is not None and current_task.cancelled():
-                        logger.debug(
-                            "Client disconnected from message request. Cancelling task %s", task_id
-                        )
-                        producer_task.cancel()
-                        await queue.close(immediate=True)
-
-                    await asyncio.shield(self._cleanup_producer(producer_task, task_id))
-                except asyncio.CancelledError:
-                    pass
 
         if not result:
             raise InternalError()
@@ -970,6 +942,7 @@ def create_app(settings: Settings) -> FastAPI:
     extended_card_etag = build_agent_card_etag(extended_agent_card)
     persistence_summary = describe_lightweight_persistence_backend(settings)
     lifespan = build_lifespan(
+        request_handler=handler,
         database_engine=database_engine,
         task_store_runtime=task_store_runtime,
         runtime_state_runtime=runtime_state_runtime,