Skip to content

fix: inherit stored task context id #1208

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
samrusani wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

fix: inherit stored task context id #1208

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/a2a/server/request_handlers/default_request_handler_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
36 changes: 36 additions & 0 deletions tests/server/request_handlers/test_default_request_handler_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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."""
Expand Down
Loading