get leaf id as input

Author: riturajFi

src/client/acontext-py/src/acontext/resources/async_sessions.py

@@ -228,6 +228,7 @@ async def store_message(
         *,
         blob: MessageBlob,
         format: Literal["acontext", "openai", "anthropic", "gemini"] = "openai",
+        parent_id: str | None = None,
         meta: dict[str, Any] | None = None,
         file_field: str | None = None,
         file: (
@@ -243,6 +244,7 @@ async def store_message(
             session_id: The UUID of the session.
             blob: The message blob in Acontext, OpenAI, Anthropic, or Gemini format.
             format: The format of the message blob. Defaults to "openai".
+            parent_id: Optional parent message UUID for branching. Defaults to None.
             meta: Optional user-provided metadata for the message. This metadata is stored
                 separately from the message content and can be retrieved via get_messages().metas
                 or updated via patch_message_meta(). Works with all formats.
@@ -272,6 +274,8 @@ async def store_message(
         payload: dict[str, Any] = {
             "format": format,
         }
+        if parent_id is not None:
+            payload["parent_id"] = parent_id
         if meta is not None:
             payload["meta"] = meta
 
@@ -369,6 +373,7 @@ async def get_messages(
         *,
         limit: int | None = None,
         cursor: str | None = None,
+        leaf_id: str | None = None,
         with_asset_public_url: bool | None = None,
         with_events: bool | None = None,
         format: Literal["acontext", "openai", "anthropic", "gemini"] = "openai",
@@ -382,6 +387,7 @@ async def get_messages(
             session_id: The UUID of the session.
             limit: Maximum number of messages to return. Defaults to None.
             cursor: Cursor for pagination. Defaults to None.
+            leaf_id: Optional leaf message UUID to read one root-to-leaf branch path. Defaults to None.
             with_asset_public_url: Whether to include presigned URLs for assets. Defaults to None.
             format: The format of the messages. Defaults to "openai". Supports "acontext", "openai", "anthropic", or "gemini".
             time_desc: Order by created_at descending if True, ascending if False. Defaults to None.
@@ -404,9 +410,19 @@ async def get_messages(
         Returns:
             GetMessagesOutput containing the list of messages and pagination information.
         """
+        if leaf_id is not None:
+            if limit is not None:
+                raise ValueError("leaf_id cannot be combined with limit")
+            if cursor is not None:
+                raise ValueError("leaf_id cannot be combined with cursor")
+            if time_desc is not None:
+                raise ValueError("leaf_id cannot be combined with time_desc")
+
         params: dict[str, Any] = {}
         if format is not None:
             params["format"] = format
+        if leaf_id is not None:
+            params["leaf_id"] = leaf_id
         params.update(
             build_params(
                 limit=limit,

src/client/acontext-py/src/acontext/resources/sessions.py

@@ -228,6 +228,7 @@ def store_message(
         *,
         blob: MessageBlob,
         format: Literal["acontext", "openai", "anthropic", "gemini"] = "openai",
+        parent_id: str | None = None,
         meta: dict[str, Any] | None = None,
         file_field: str | None = None,
         file: (
@@ -243,6 +244,7 @@ def store_message(
             session_id: The UUID of the session.
             blob: The message blob in Acontext, OpenAI, Anthropic, or Gemini format.
             format: The format of the message blob. Defaults to "openai".
+            parent_id: Optional parent message UUID for branching. Defaults to None.
             meta: Optional user-provided metadata for the message. This metadata is stored
                 separately from the message content and can be retrieved via get_messages().metas
                 or updated via patch_message_meta(). Works with all formats.
@@ -273,6 +275,8 @@ def store_message(
         payload: dict[str, Any] = {
             "format": format,
         }
+        if parent_id is not None:
+            payload["parent_id"] = parent_id
         if meta is not None:
             payload["meta"] = meta
 
@@ -369,6 +373,7 @@ def get_messages(
         *,
         limit: int | None = None,
         cursor: str | None = None,
+        leaf_id: str | None = None,
         with_asset_public_url: bool | None = None,
         with_events: bool | None = None,
         format: Literal["acontext", "openai", "anthropic", "gemini"] = "openai",
@@ -382,6 +387,7 @@ def get_messages(
             session_id: The UUID of the session.
             limit: Maximum number of messages to return. Defaults to None.
             cursor: Cursor for pagination. Defaults to None.
+            leaf_id: Optional leaf message UUID to read one root-to-leaf branch path. Defaults to None.
             with_asset_public_url: Whether to include presigned URLs for assets. Defaults to None.
             format: The format of the messages. Defaults to "openai". Supports "acontext", "openai", "anthropic", or "gemini".
             time_desc: Order by created_at descending if True, ascending if False. Defaults to None.
@@ -404,9 +410,19 @@ def get_messages(
         Returns:
             GetMessagesOutput containing the list of messages and pagination information.
         """
+        if leaf_id is not None:
+            if limit is not None:
+                raise ValueError("leaf_id cannot be combined with limit")
+            if cursor is not None:
+                raise ValueError("leaf_id cannot be combined with cursor")
+            if time_desc is not None:
+                raise ValueError("leaf_id cannot be combined with time_desc")
+
         params: dict[str, Any] = {}
         if format is not None:
             params["format"] = format
+        if leaf_id is not None:
+            params["leaf_id"] = leaf_id
         params.update(
             build_params(
                 limit=limit,

src/client/acontext-py/tests/test_async_client.py

@@ -210,6 +210,38 @@ def read(self) -> bytes:
     assert attachment[2] == "image/png"
 
 
+@patch("acontext.async_client.AcontextAsyncClient.request", new_callable=AsyncMock)
+@pytest.mark.asyncio
+async def test_async_store_message_forwards_parent_id(
+    mock_request, async_client: AcontextAsyncClient
+) -> None:
+    mock_request.return_value = {
+        "id": "msg-id",
+        "session_id": "session-id",
+        "parent_id": "parent-msg-id",
+        "role": "user",
+        "meta": {},
+        "parts": [],
+        "session_task_process_status": "pending",
+        "created_at": "2024-01-01T00:00:00Z",
+        "updated_at": "2024-01-01T00:00:00Z",
+    }
+
+    await async_client.sessions.store_message(
+        "session-id",
+        blob={"role": "user", "content": "hi"},  # type: ignore[arg-type]
+        format="openai",
+        parent_id="parent-msg-id",
+    )
+
+    mock_request.assert_called_once()
+    args, kwargs = mock_request.call_args
+    method, path = args
+    assert method == "POST"
+    assert path == "/session/session-id/messages"
+    assert kwargs["json_data"]["parent_id"] == "parent-msg-id"
+
+
 @patch("acontext.async_client.AcontextAsyncClient.request", new_callable=AsyncMock)
 @pytest.mark.asyncio
 async def test_async_store_message_allows_nullable_blob_for_other_formats(
@@ -480,6 +512,45 @@ async def test_async_sessions_get_messages_forwards_format(
     assert hasattr(result, "has_more")
 
 
+@patch("acontext.async_client.AcontextAsyncClient.request", new_callable=AsyncMock)
+@pytest.mark.asyncio
+async def test_async_sessions_get_messages_forwards_leaf_id(
+    mock_request, async_client: AcontextAsyncClient
+) -> None:
+    mock_request.return_value = {
+        "items": [],
+        "ids": ["leaf-msg-id"],
+        "has_more": False,
+        "this_time_tokens": 0,
+    }
+
+    result = await async_client.sessions.get_messages(
+        "session-id", format="acontext", leaf_id="leaf-msg-id"
+    )
+
+    mock_request.assert_called_once()
+    args, kwargs = mock_request.call_args
+    method, path = args
+    assert method == "GET"
+    assert path == "/session/session-id/messages"
+    assert kwargs["params"] == {"format": "acontext", "leaf_id": "leaf-msg-id"}
+    assert hasattr(result, "items")
+    assert hasattr(result, "has_more")
+
+
+@patch("acontext.async_client.AcontextAsyncClient.request", new_callable=AsyncMock)
+@pytest.mark.asyncio
+async def test_async_sessions_get_messages_rejects_leaf_id_with_limit(
+    mock_request, async_client: AcontextAsyncClient
+) -> None:
+    with pytest.raises(ValueError, match="leaf_id cannot be combined with limit"):
+        await async_client.sessions.get_messages(
+            "session-id", leaf_id="leaf-msg-id", limit=10
+        )
+
+    mock_request.assert_not_called()
+
+
 @patch("acontext.async_client.AcontextAsyncClient.request", new_callable=AsyncMock)
 @pytest.mark.asyncio
 async def test_async_sessions_get_messages_rejects_non_positive_gt_token(

src/client/acontext-py/tests/test_client.py

@@ -143,6 +143,35 @@ def read(self) -> bytes:
     assert attachment[2] == "image/png"
 
 
+@patch("acontext.client.AcontextClient.request")
+def test_store_message_forwards_parent_id(mock_request, client: AcontextClient) -> None:
+    mock_request.return_value = {
+        "id": "msg-id",
+        "session_id": "session-id",
+        "parent_id": "parent-msg-id",
+        "role": "user",
+        "meta": {},
+        "parts": [],
+        "session_task_process_status": "pending",
+        "created_at": "2024-01-01T00:00:00Z",
+        "updated_at": "2024-01-01T00:00:00Z",
+    }
+
+    client.sessions.store_message(
+        "session-id",
+        blob={"role": "user", "content": "hi"},  # type: ignore[arg-type]
+        format="openai",
+        parent_id="parent-msg-id",
+    )
+
+    mock_request.assert_called_once()
+    args, kwargs = mock_request.call_args
+    method, path = args
+    assert method == "POST"
+    assert path == "/session/session-id/messages"
+    assert kwargs["json_data"]["parent_id"] == "parent-msg-id"
+
+
 @patch("acontext.client.AcontextClient.request")
 def test_store_message_allows_nullable_blob_for_other_formats(
     mock_request, client: AcontextClient
@@ -553,6 +582,41 @@ def test_sessions_get_messages_forwards_format(
     assert hasattr(result, "has_more")
 
 
+@patch("acontext.client.AcontextClient.request")
+def test_sessions_get_messages_forwards_leaf_id(
+    mock_request, client: AcontextClient
+) -> None:
+    mock_request.return_value = {
+        "items": [],
+        "ids": ["leaf-msg-id"],
+        "has_more": False,
+        "this_time_tokens": 0,
+    }
+
+    result = client.sessions.get_messages(
+        "session-id", format="acontext", leaf_id="leaf-msg-id"
+    )
+
+    mock_request.assert_called_once()
+    args, kwargs = mock_request.call_args
+    method, path = args
+    assert method == "GET"
+    assert path == "/session/session-id/messages"
+    assert kwargs["params"] == {"format": "acontext", "leaf_id": "leaf-msg-id"}
+    assert hasattr(result, "items")
+    assert hasattr(result, "has_more")
+
+
+@patch("acontext.client.AcontextClient.request")
+def test_sessions_get_messages_rejects_leaf_id_with_limit(
+    mock_request, client: AcontextClient
+) -> None:
+    with pytest.raises(ValueError, match="leaf_id cannot be combined with limit"):
+        client.sessions.get_messages("session-id", leaf_id="leaf-msg-id", limit=10)
+
+    mock_request.assert_not_called()
+
+
 @patch("acontext.client.AcontextClient.request")
 def test_sessions_get_messages_with_edit_strategies(
     mock_request, client: AcontextClient

src/client/acontext-ts/src/resources/sessions.ts

@@ -229,6 +229,7 @@ export class SessionsAPI {
     blob: MessageBlob,
     options?: {
       format?: 'acontext' | 'openai' | 'anthropic' | 'gemini';
+      parentId?: string | null;
       meta?: Record<string, unknown> | null;
       fileField?: string | null;
       file?: FileUpload | null;
@@ -247,6 +248,10 @@ export class SessionsAPI {
       format,
     };
 
+    if (options?.parentId !== undefined && options?.parentId !== null) {
+      payload.parent_id = options.parentId;
+    }
+
     if (options?.meta !== undefined && options?.meta !== null) {
       payload.meta = options.meta;
     }
@@ -335,6 +340,7 @@ export class SessionsAPI {
    * @param options - Options for retrieving messages.
    * @param options.limit - Maximum number of messages to return.
    * @param options.cursor - Cursor for pagination.
+   * @param options.leafId - Return only the root-to-leaf path for this leaf message ID.
    * @param options.withAssetPublicUrl - Whether to include presigned URLs for assets.
    * @param options.withEvents - Whether to include session events in the response.
    * @param options.format - The format of the messages ('acontext', 'openai', 'anthropic', or 'gemini').
@@ -359,6 +365,7 @@ export class SessionsAPI {
     options?: {
       limit?: number | null;
       cursor?: string | null;
+      leafId?: string | null;
       withAssetPublicUrl?: boolean | null;
       withEvents?: boolean | null;
       format?: 'acontext' | 'openai' | 'anthropic' | 'gemini';
@@ -367,10 +374,25 @@ export class SessionsAPI {
       pinEditingStrategiesAtMessage?: string | null;
     }
   ): Promise<GetMessagesOutput> {
+    if (options?.leafId !== undefined && options?.leafId !== null) {
+      if (options?.limit !== undefined && options?.limit !== null) {
+        throw new Error('leafId cannot be combined with limit');
+      }
+      if (options?.cursor !== undefined && options?.cursor !== null) {
+        throw new Error('leafId cannot be combined with cursor');
+      }
+      if (options?.timeDesc !== undefined && options?.timeDesc !== null) {
+        throw new Error('leafId cannot be combined with timeDesc');
+      }
+    }
+
     const params: Record<string, string | number> = {};
     if (options?.format !== undefined) {
       params.format = options.format;
     }
+    if (options?.leafId !== undefined && options?.leafId !== null) {
+      params.leaf_id = options.leafId;
+    }
     Object.assign(
       params,
       buildParams({