From 68162bdc6b2638f7af2f99d84f881b5d81a4803d Mon Sep 17 00:00:00 2001
From: Sami Rusani <sr@samirusani>
Date: Tue, 25 Aug 2026 17:51:37 +0200
Subject: [PATCH] fix: inherit stored task context id

---
 .../default_request_handler_v2.py             |  2 ++
 .../test_default_request_handler_v2.py        | 36 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/src/a2a/server/request_handlers/default_request_handler_v2.py b/src/a2a/server/request_handlers/default_request_handler_v2.py
index 1b25e4f2b..a47826adb 100644
--- a/src/a2a/server/request_handlers/default_request_handler_v2.py
+++ b/src/a2a/server/request_handlers/default_request_handler_v2.py
@@ -221,6 +221,8 @@ async def _setup_active_task(
             task = await self.task_store.get(original_task_id, call_context)
             if not task:
                 raise TaskNotFoundError(f'Task {original_task_id} not found')
+            if not original_context_id:
+                original_context_id = task.context_id
 
         # Build context to resolve or generate missing IDs
         request_context = await self._request_context_builder.build(
diff --git a/tests/server/request_handlers/test_default_request_handler_v2.py b/tests/server/request_handlers/test_default_request_handler_v2.py
index 6011174e4..dd8566e96 100644
--- a/tests/server/request_handlers/test_default_request_handler_v2.py
+++ b/tests/server/request_handlers/test_default_request_handler_v2.py
@@ -969,6 +969,7 @@ async def test_on_message_send_task_in_terminal_state(terminal_state):
         task_id=task_id, status_state=terminal_state
     )
     mock_task_store = AsyncMock(spec=TaskStore)
+    mock_task_store.get.return_value = terminal_task
     request_handler = DefaultRequestHandlerV2(
         agent_executor=MockAgentExecutor(),
         task_store=mock_task_store,
@@ -1008,6 +1009,7 @@ async def test_on_message_send_stream_task_in_terminal_state(terminal_state):
         task_id=task_id, status_state=terminal_state
     )
     mock_task_store = AsyncMock(spec=TaskStore)
+    mock_task_store.get.return_value = terminal_task
     request_handler = DefaultRequestHandlerV2(
         agent_executor=MockAgentExecutor(),
         task_store=mock_task_store,
@@ -1038,6 +1040,40 @@ async def test_on_message_send_stream_task_in_terminal_state(terminal_state):
     )
 
 
+@pytest.mark.asyncio
+async def test_setup_active_task_inherits_stored_task_context_id():
+    """A task-only continuation must retain its existing conversation context."""
+    task_id = 'existing_task'
+    stored_task = create_sample_task(
+        task_id=task_id, context_id='stored_context'
+    )
+    task_store = AsyncMock(spec=TaskStore)
+    task_store.get.return_value = stored_task
+    request_handler = DefaultRequestHandlerV2(
+        agent_executor=MockAgentExecutor(),
+        task_store=task_store,
+        agent_card=create_default_agent_card(),
+    )
+    request_handler._active_task_registry.get_or_create = AsyncMock()
+    params = SendMessageRequest(
+        message=Message(
+            role=Role.ROLE_USER,
+            message_id='continue_existing_task',
+            parts=[Part(text='continue')],
+            task_id=task_id,
+        )
+    )
+    call_context = create_server_call_context()
+
+    _, request_context = await request_handler._setup_active_task(
+        params, call_context
+    )
+
+    assert request_context.context_id == stored_task.context_id
+    assert params.message.context_id == stored_task.context_id
+    task_store.get.assert_awaited_once_with(task_id, call_context)
+
+
 @pytest.mark.asyncio
 async def test_on_message_send_task_id_provided_but_task_not_found():
     """Test on_message_send when task_id is provided but task doesn't exist."""