[Feature]: Allow passing threadId as the ADK session_id

[umax-imagination-media]

Pre-flight Checklist

• I have searched existing issues and this hasn't been requested yet.

Problem or Motivation

Summary

SessionManager currently ignores the CopilotKit threadId when creating the underlying ADK session, generating a random UUID instead. This makes it impossible to do an O(1) session lookup by threadId and forces an expensive list_sessions + linear scan on every reconnect.

I'd like to propose a use_thread_id_as_session_id flag that lets integrators use the threadId directly as the ADK session_id.

Problem

When CopilotKit connects, it sends a threadId that serves as the canonical conversation identifier on the frontend. The current SessionManager flow:

1. Creates a session via DatabaseSessionService.create_session without passing session_id, so the database generates a random UUID.
2. Stores the threadId inside session.state["_ag_ui_thread_id"].
3. On reconnect, calls list_sessions and does a linear scan over every session's state to find the one whose _ag_ui_thread_id matches — an O(n) operation that degrades as sessions accumulate.

This causes several issues:

Issue	Impact
threadId ≠ session.id	Cannot build a shareable URL like /chat/{threadId} that maps directly to a DB row
Reconnect requires full table scan	Latency grows linearly with the number of sessions
/agents/state resolution	Must repeat the same list-and-scan to resolve a threadId
Race conditions on reconnect	Two concurrent requests can each create a new session for the same threadId because the scan-then-create is not atomic

Proposed Solution

Proposed Solution

Add a use_thread_id_as_session_id boolean flag (default False) to ADKAgent / SessionManager:

ADKAgent(
    ...,
    use_thread_id_as_session_id=True,
)

When the flag is True:

1. get_or_create_session passes session_id=thread_id to create_session instead of letting the database generate a random UUID.
2. _find_session_by_thread_id attempts a direct get_session(session_id=thread_id) first (O(1) primary-key lookup), falling back to the existing list_sessions scan only for backward compatibility with sessions created before the flag was enabled.

When the flag is False (default), behaviour is identical to today — fully backward compatible.

Handling concurrent creates

When session_id=thread_id is passed, two concurrent requests for the same threadId can race past the existence check. The create path should catch the resulting AlreadyExistsError / IntegrityError and fall back to get_session:

try:
    session = await self._session_service.create_session(
        ..., session_id=thread_id,
    )
except (AlreadyExistsError, IntegrityError):
    session = await self._session_service.get_session(
        session_id=thread_id, app_name=app_name, user_id=user_id,
    )

Alternatives Considered

No response

Additional Context

Current Workaround

We are currently monkey-patching SessionManager.get_or_create_session and SessionManager._find_session_by_thread_id at startup to achieve this. The patch works reliably in production with ag-ui-adk ~=0.5.0 and google-adk ~=1.25.1, but it's fragile — any internal refactor of SessionManager could break it.

Monkey-patch implementation (click to expand)

"""Patch ag_ui_adk SessionManager to use threadId as the ADK session_id."""

from __future__ import annotations

import logging
from typing import Any, Dict, Optional, Tuple

from google.adk.errors.already_exists_error import AlreadyExistsError
from sqlalchemy.exc import IntegrityError

from ag_ui_adk.session_manager import (
    APP_NAME_STATE_KEY,
    THREAD_ID_STATE_KEY,
    USER_ID_STATE_KEY,
    SessionManager,
)

logger = logging.getLogger(__name__)


async def _get_or_create_session_using_thread_id(
    self: SessionManager,
    thread_id: str,
    app_name: str,
    user_id: str,
    initial_state: Optional[Dict[str, Any]] = None,
) -> Tuple[Any, str]:
    if self._max_per_user:
        user_count = len(self._user_sessions.get(user_id, set()))
        if user_count >= self._max_per_user:
            await self._remove_oldest_user_session(user_id)

    session = await self._find_session_by_thread_id(app_name, user_id, thread_id)
    if session:
        session_key = self._make_session_key(app_name, session.id)
        self._track_session(session_key, user_id)
        return session, session.id

    state = {
        **(initial_state or {}),
        THREAD_ID_STATE_KEY: thread_id,
        APP_NAME_STATE_KEY: app_name,
        USER_ID_STATE_KEY: user_id,
    }

    try:
        session = await self._session_service.create_session(
            user_id=user_id,
            app_name=app_name,
            state=state,
            session_id=thread_id,  # <-- the key change
        )
    except (AlreadyExistsError, IntegrityError):
        session = await self._session_service.get_session(
            session_id=thread_id, app_name=app_name, user_id=user_id,
        )
        if session is None:
            raise RuntimeError(
                f"Session {thread_id} reported as duplicate but could not be read back"
            )

    session_key = self._make_session_key(app_name, session.id)
    self._track_session(session_key, user_id)

    if not self._cleanup_task:
        self._start_cleanup_task()

    return session, session.id


async def _find_session_by_thread_id_direct(
    self: SessionManager,
    app_name: str,
    user_id: str,
    thread_id: str,
) -> Optional[Any]:
    # O(1) lookup — session_id IS the thread_id
    try:
        session = await self._session_service.get_session(
            session_id=thread_id, app_name=app_name, user_id=user_id,
        )
        if session:
            return session
    except Exception:
        pass

    # Fallback for legacy sessions (thread_id stored in state)
    if hasattr(self._session_service, "list_sessions"):
        response = await self._session_service.list_sessions(
            app_name=app_name, user_id=user_id,
        )
        for session in response.sessions:
            if session.state and session.state.get(THREAD_ID_STATE_KEY) == thread_id:
                return session

    return None


def apply_session_id_patch() -> None:
    SessionManager.get_or_create_session = _get_or_create_session_using_thread_id
    SessionManager._find_session_by_thread_id = _find_session_by_thread_id_direct

Expected Benefits

Before	After
threadId ≠ session.id (random UUID)	threadId == session.id
Reconnect = list_sessions + O(n) scan	Reconnect = get_session(id) O(1)
Cannot derive DB row from frontend ID	/chat/{threadId} maps directly to a DB row
Concurrent creates can produce duplicates	Race condition handled gracefully

Environment

• ag-ui-adk: ~=0.5.0
• google-adk: ~=1.25.1
• Session backend: DatabaseSessionService (PostgreSQL)
• Frontend: CopilotKit

[umax-imagination-media]

@contextablemark fyi: I have confirmed that the same issue still exists in ag-ui-adk 0.5.1 https://pypi.org/project/ag-ui-adk/0.5.1/

[contextablemark]

@umax-imagination-media Apologies... I thought that I had already followed up on this, but I don't see my comment here.
When you indicate that the same issue still exists, do you mean that when the "use_thread_id_as_session_id" flag is set that it is not functioning correctly?

[umax-imagination-media]

Hello @contextablemark. Sorry for the delay in response and misunderstanding. There are no use_thread_id_as_session_id flag yet in v 0.5.1 :)

And because of that, in my project, my custom patch is still needed. The installed ag-ui-adk 0.5.1 (which is also the latest on PyPI) still has both issues:
get_or_create_session explicitly omits session_id (line 164 comment: # Note: session_id intentionally omitted - let backend generate it), so threadId ≠ session.id.
_find_session_by_thread_id still does a full list_sessions + O(n) linear scan to find a session by thread_id in state — no O(1) get_session attempt.

[contextablemark]

@umax-imagination-media I see... so you were just verifying that the open issue is still not addressed - I was thinking that I had fixed it and you were indicating that the fix was not working 😆 I'll take a look.

[contextablemark]

@umax-imagination-media It is fixed, just in the "Unreleased" section because its... er... unreleased. But on giving it a once over I did find one more gap that needs to be addressed.

[contextablemark]

More changes in #1383

[umax-imagination-media]

Hello @contextablemark . Thank you for helping with this issue!

just adding link to the related PR #1271 since PR #1383 above does not have all the changes