agent-client-protocol 0.12.0__tar.gz → 0.12.1__tar.gz

{agent_client_protocol-0.12.0 → agent_client_protocol-0.12.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: agent-client-protocol
-Version: 0.12.0
+Version: 0.12.1
 Summary: A Python implement of Agent Client Protocol (ACP, by Zed Industries)
 Keywords: python
 Author-Email: Chojan Shang <psiace@apache.org>, Frost Ming <me@frostming.com>

{agent_client_protocol-0.12.0 → agent_client_protocol-0.12.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "agent-client-protocol"
-version = "0.12.0"
+version = "0.12.1"
 description = "A Python implement of Agent Client Protocol (ACP, by Zed Industries)"
 authors = [
     { name = "Chojan Shang", email = "psiace@apache.org" },
@@ -43,7 +43,7 @@ http = [
 
 [dependency-groups]
 dev = [
-    "datamodel-code-generator>=0.25",
+    "datamodel-code-generator>=0.71.0",
     "pytest>=7.2.0",
     "pytest-asyncio>=0.21.0",
     "tox-uv>=1.11.3",

{agent_client_protocol-0.12.0 → agent_client_protocol-0.12.1}/scripts/gen_schema.py

@@ -20,13 +20,6 @@ SCHEMA_JSON = SCHEMA_DIR / "schema.json"
 VERSION_FILE = SCHEMA_DIR / "VERSION"
 SCHEMA_OUT = ROOT / "src" / "acp" / "schema.py"
 
-# Pattern caches used when post-processing generated schema.
-FIELD_DECLARATION_PATTERN = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\s*:")
-DESCRIPTION_PATTERN = re.compile(
-    r"description\s*=\s*(?P<prefix>[rRbBuU]*)?(?P<quote>'''|\"\"\"|'|\")(?P<value>.*?)(?P=quote)",
-    re.DOTALL,
-)
-
 STDIO_TYPE_LITERAL = 'Literal["2#-datamodel-code-generator-#-object-#-special-#"]'
 MODELS_TO_REMOVE = [
     "AgentClientProtocol",
@@ -137,6 +130,31 @@ ENUM_LITERAL_MAP: dict[str, tuple[str, ...]] = {
     "ToolKind": ("read", "edit", "delete", "move", "search", "execute", "think", "fetch", "switch_mode", "other"),
 }
 
+# datamodel-code-generator 0.64 promotes referenced string enums to Enum classes.
+# Keep the existing Python API, where these schema types are plain strings; the
+# selected public fields below are narrowed back to the named Literal aliases.
+STRING_ENUM_TYPES = (
+    *ENUM_LITERAL_MAP,
+    "ElicitationSchemaType",
+    "NesDiagnosticSeverity",
+    "NesRejectReason",
+    "NesTriggerKind",
+    "PositionEncodingKind",
+    "Role",
+    "StringFormat",
+    "TextDocumentSyncKind",
+)
+
+# Preserve RootModel classes that existed in the generated public surface before
+# 0.64; other unreferenced RootModels are intermediates left after collapsing.
+PUBLIC_ROOT_MODELS = {
+    "AgentResponse",
+    "ClientResponse",
+    "ElicitationContentValue",
+    "ElicitationFormMode",
+    "ElicitationUrlMode",
+}
+
 FIELD_TYPE_OVERRIDES: tuple[tuple[str, str, str, bool], ...] = (
     ("PermissionOption", "kind", "PermissionOptionKind", False),
     ("PlanEntry", "priority", "PlanEntryPriority", False),
@@ -259,7 +277,16 @@ def generate_schema() -> None:
             "--collapse-root-models",
             "--output-model-type",
             "pydantic_v2.BaseModel",
+            "--no-use-specialized-enum",
+            "--no-use-standard-collections",
+            "--no-use-union-operator",
+            "--type-overrides",
+            json.dumps(dict.fromkeys(STRING_ENUM_TYPES, "builtins.str")),
+            "--formatters",
+            "black",
+            "isort",
             "--use-annotated",
+            "--use-field-description",
             "--snake-case-field",
         ]
 
@@ -474,6 +501,9 @@ def postprocess_generated_schema(output_path: Path) -> list[str]:
     header_block = _build_header_block()
 
     content = _strip_existing_header(raw_content)
+    # Type overrides for builtins are rendered as imports in 0.64, but the
+    # annotations should continue to use Python's builtin `str` directly.
+    content = content.replace("from builtins import str\n", "")
     content = _remove_unused_models(content)
     content, leftover_classes = _rename_numbered_models(content)
 
@@ -481,8 +511,8 @@ def postprocess_generated_schema(output_path: Path) -> list[str]:
         _ProcessingStep("apply field overrides", _apply_field_overrides),
         _ProcessingStep("apply default overrides", _apply_default_overrides),
         _ProcessingStep("restore required nullable fields", _restore_required_nullable_fields),
-        _ProcessingStep("attach description comments", _add_description_comments),
         _ProcessingStep("ensure custom BaseModel", _ensure_custom_base_model),
+        _ProcessingStep("enable RootModel attribute docstrings", _enable_root_model_attribute_docstrings),
         _ProcessingStep("inject field validators", _inject_field_validators),
         _ProcessingStep("inject deserialize defaults", _inject_deserialize_defaults),
         _ProcessingStep("inject schema aliases", _inject_schema_aliases),
@@ -494,6 +524,7 @@ def postprocess_generated_schema(output_path: Path) -> list[str]:
     missing_targets = _find_missing_targets(content)
 
     content = _inject_enum_aliases(content)
+    content = _remove_unreferenced_root_models(content)
     final_content = header_block + content.rstrip() + "\n"
     if not final_content.endswith("\n"):
         final_content += "\n"
@@ -646,7 +677,7 @@ def _ensure_custom_base_model(content: str) -> str:
         lines[idx] = "from pydantic import " + ", ".join(new_imports)
         to_insert = textwrap.dedent("""\
             class BaseModel(_BaseModel):
-                model_config = ConfigDict(populate_by_name=True)
+                model_config = ConfigDict(populate_by_name=True, use_attribute_docstrings=True)
 
                 def __getattr__(self, item: str) -> Any:
                     if item.lower() != item:
@@ -662,6 +693,24 @@ def _ensure_custom_base_model(content: str) -> str:
     return "\n".join(lines) + "\n"
 
 
+def _enable_root_model_attribute_docstrings(content: str) -> str:
+    lines = content.splitlines(keepends=True)
+    tree = ast.parse(content)
+    insertion_points = [
+        node.body[0].lineno - 1
+        for node in tree.body
+        if isinstance(node, ast.ClassDef)
+        and node.body
+        and any(
+            isinstance(base, ast.Subscript) and isinstance(base.value, ast.Name) and base.value.id == "RootModel"
+            for base in node.bases
+        )
+    ]
+    for line_index in reversed(insertion_points):
+        lines.insert(line_index, "    model_config = ConfigDict(use_attribute_docstrings=True)\n\n")
+    return "".join(lines)
+
+
 def _ensure_pydantic_import(content: str, name: str) -> str:
     """Add *name* to the ``from pydantic import ...`` line if not already present."""
     lines = content.splitlines()
@@ -916,8 +965,14 @@ def _restore_required_nullable_fields(content: str, schema: dict[str, Any] | Non
         def restore_block(match: re.Match[str], _field_names: list[str] = field_names) -> str:
             header, block = match.group(1), match.group(2)
             for field_name in _field_names:
-                field_pattern = re.compile(rf"(\n\s+{re.escape(field_name)}:\s+Annotated\[[\s\S]*?\n\s+\]\s*)=\s*None")
-                block = field_pattern.sub(r"\1", block, count=1)
+                field_patterns = (
+                    re.compile(rf"(\n\s+{re.escape(field_name)}:[^\n]*?)\s*=\s*None(?=\n)"),
+                    re.compile(rf"(\n\s+{re.escape(field_name)}:[^\n]*\[\s*\n[\s\S]*?\n\s+\]\s*)=\s*None"),
+                )
+                for field_pattern in field_patterns:
+                    block, count = field_pattern.subn(r"\1", block, count=1)
+                    if count:
+                        break
             return header + block
 
         content = class_pattern.sub(restore_block, content, count=1)
@@ -926,18 +981,14 @@ def _restore_required_nullable_fields(content: str, schema: dict[str, Any] | Non
 
 def _apply_field_overrides(content: str) -> str:
     for class_name, field_name, new_type, optional in FIELD_TYPE_OVERRIDES:
-        if optional:
-            pattern = re.compile(
-                rf"(class {class_name}\(BaseModel\):.*?\n\s+{field_name}:\s+Annotated\[\s*)Optional\[str],",
-                re.DOTALL,
-            )
-            content, count = pattern.subn(rf"\1Optional[{new_type}],", content)
-        else:
-            pattern = re.compile(
-                rf"(class {class_name}\(BaseModel\):.*?\n\s+{field_name}:\s+Annotated\[\s*)str,",
-                re.DOTALL,
-            )
-            content, count = pattern.subn(rf"\1{new_type},", content)
+        old_type = "Optional[str]" if optional else "str"
+        replacement_type = f"Optional[{new_type}]" if optional else new_type
+        pattern = re.compile(
+            rf"(class {re.escape(class_name)}\(BaseModel\):.*?\n\s+{re.escape(field_name)}:\s+"
+            rf"(?:Annotated\[\s*)?){re.escape(old_type)}(?=\s*(?:,|=|\n))",
+            re.DOTALL,
+        )
+        content, count = pattern.subn(rf"\g<1>{replacement_type}", content, count=1)
         if count == 0:
             print(
                 f"Warning: failed to apply type override for {class_name}.{field_name} -> {new_type}",
@@ -963,7 +1014,8 @@ def _apply_default_overrides(content: str) -> str:
             field_patterns: tuple[tuple[re.Pattern[str], Callable[[re.Match[str]], str]], ...] = (
                 (
                     re.compile(
-                        rf"(\n\s+{_field_name}:.*?\]\s*=\s*)([\s\S]*?)(?=\n\s{{4}}[A-Za-z_]|$)",
+                        rf"(\n\s+{_field_name}:.*?\]\s*=\s*)([\s\S]*?)"
+                        rf"(?=\n\s{{4}}(?:[A-Za-z_][A-Za-z0-9_]*\s*:|[rRuUbBfF]*(?:'''|\"\"\"))|$)",
                         re.DOTALL,
                     ),
                     lambda m, _rep=_replacement: m.group(1) + _rep,
@@ -995,77 +1047,6 @@ def _apply_default_overrides(content: str) -> str:
     return content
 
 
-def _add_description_comments(content: str) -> str:
-    lines = content.splitlines()
-    new_lines: list[str] = []
-    index = 0
-
-    while index < len(lines):
-        line = lines[index]
-        stripped = line.lstrip()
-        indent = len(line) - len(stripped)
-
-        if indent == 4 and FIELD_DECLARATION_PATTERN.match(stripped or ""):
-            block_lines, next_index = _collect_field_block(lines, index, indent)
-            block_text = "\n".join(block_lines)
-            description = _extract_description(block_text)
-
-            if description:
-                indent_str = " " * indent
-                comment_lines = [
-                    f"{indent_str}# {comment_line}" if comment_line else f"{indent_str}#"
-                    for comment_line in description.splitlines()
-                ]
-                if comment_lines:
-                    new_lines.extend(comment_lines)
-
-            new_lines.extend(block_lines)
-            index = next_index
-            continue
-
-        new_lines.append(line)
-        index += 1
-
-    return "\n".join(new_lines)
-
-
-def _collect_field_block(lines: list[str], start: int, indent: int) -> tuple[list[str], int]:
-    block: list[str] = []
-    index = start
-
-    while index < len(lines):
-        current_line = lines[index]
-        current_indent = len(current_line) - len(current_line.lstrip())
-        if index != start and current_line.strip() and current_indent <= indent:
-            break
-
-        block.append(current_line)
-        index += 1
-
-    return block, index
-
-
-def _extract_description(block_text: str) -> str | None:
-    match = DESCRIPTION_PATTERN.search(block_text)
-    if not match:
-        return None
-
-    prefix = match.group("prefix") or ""
-    quote = match.group("quote")
-    value = match.group("value")
-    literal = f"{prefix}{quote}{value}{quote}"
-
-    # datamodel-code-generator emits standard string literals, but fall back to raw text on parse errors.
-    try:
-        parsed = ast.literal_eval(literal)
-    except (SyntaxError, ValueError):
-        return value.replace("\\n", "\n")
-
-    if isinstance(parsed, str):
-        return parsed
-    return str(parsed)
-
-
 def _inject_enum_aliases(content: str) -> str:
     enum_lines = [
         f"{name} = Literal[{', '.join(repr(value) for value in values)}]" for name, values in ENUM_LITERAL_MAP.items()
@@ -1080,6 +1061,45 @@ def _inject_enum_aliases(content: str) -> str:
     return content[:insertion_point] + block + content[insertion_point:]
 
 
+def _remove_unreferenced_root_models(content: str) -> str:
+    tree = ast.parse(content)
+    root_models = {
+        node.name: node
+        for node in tree.body
+        if isinstance(node, ast.ClassDef)
+        and any(
+            isinstance(base, ast.Subscript) and isinstance(base.value, ast.Name) and base.value.id == "RootModel"
+            for base in node.bases
+        )
+    }
+
+    referenced_roots = set(PUBLIC_ROOT_MODELS)
+    root_dependencies: dict[str, set[str]] = {}
+    for statement in tree.body:
+        loaded_names = {
+            node.id
+            for node in ast.walk(statement)
+            if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load) and node.id in root_models
+        }
+        if isinstance(statement, ast.ClassDef) and statement.name in root_models:
+            root_dependencies[statement.name] = loaded_names
+        else:
+            referenced_roots.update(loaded_names)
+
+    pending = list(referenced_roots)
+    while pending:
+        root_name = pending.pop()
+        for dependency in root_dependencies.get(root_name, set()) - referenced_roots:
+            referenced_roots.add(dependency)
+            pending.append(dependency)
+
+    unused_models = [model for name, model in root_models.items() if name not in referenced_roots]
+    lines = content.splitlines(keepends=True)
+    for model in sorted(unused_models, key=lambda item: item.lineno, reverse=True):
+        del lines[model.lineno - 1 : model.end_lineno]
+    return re.sub(r"\n{4,}", "\n\n\n", "".join(lines))
+
+
 def _remove_unused_models(content: str) -> str:
     for model_name in MODELS_TO_REMOVE:
         pattern = re.compile(

{agent_client_protocol-0.12.0 → agent_client_protocol-0.12.1}/scripts/gen_signature.py

@@ -127,7 +127,10 @@ class NodeTransformer(ast.NodeTransformer):
         return arg, default
 
     def _format_annotation(self, annotation: t.Any) -> ast.expr:
-        if t.get_origin(annotation) is t.Literal and annotation in self._literals.values():
+        origin = t.get_origin(annotation)
+        if origin is t.Annotated:
+            return self._format_annotation(t.get_args(annotation)[0])
+        if origin is t.Literal and annotation in self._literals.values():
             name = next(name for name, value in self._literals.items() if value is annotation)
             self._add_schema_import(name)
             return ast.Name(id=name)
@@ -139,7 +142,6 @@ class NodeTransformer(ast.NodeTransformer):
             self._add_schema_import(annotation.__name__)
             return ast.Name(id=annotation.__name__)
         elif args := t.get_args(annotation):
-            origin = t.get_origin(annotation)
             return ast.Subscript(
                 value=self._format_annotation(origin),
                 slice=ast.Tuple(elts=[self._format_annotation(arg) for arg in args], ctx=ast.Load())

{agent_client_protocol-0.12.0 → agent_client_protocol-0.12.1}/src/acp/client/connection.py

@@ -2,12 +2,15 @@ from __future__ import annotations
 
 import asyncio
 from collections.abc import Callable
+from contextvars import ContextVar
 from typing import Any, cast, final
 
 from .._transport import Transport
 from ..connection import Connection
+from ..exceptions import RequestError
 from ..interfaces import Agent, Client
-from ..meta import AGENT_METHODS
+from ..meta import AGENT_METHODS, CLIENT_METHODS
+from ..router import _resolve_handler, _warn_legacy_handler
 from ..schema import (
     AcpMcpServer,
     AudioContentBlock,
@@ -37,6 +40,7 @@ from ..schema import (
     ResourceContentBlock,
     ResumeSessionRequest,
     ResumeSessionResponse,
+    SessionNotification,
     SetSessionConfigOptionBooleanRequest,
     SetSessionConfigOptionResponse,
     SetSessionConfigOptionSelectRequest,
@@ -52,6 +56,56 @@ __all__ = ["ClientSideConnection"]
 _CLIENT_CONNECTION_ERROR = "ClientSideConnection requires asyncio StreamWriter/StreamReader"
 
 
+class _SessionUpdateTracker:
+    """Client proxy that tracks in-flight session updates."""
+
+    def __init__(self, client: Client) -> None:
+        self._client = client
+        self._session_update, self._session_update_attr, self._legacy_session_update = _resolve_handler(
+            client, "session_update"
+        )
+        self._pending: dict[str, set[asyncio.Future[None]]] = {}
+        self._current_update: ContextVar[asyncio.Future[None] | None] = ContextVar(
+            "acp_current_session_update", default=None
+        )
+
+    async def session_update(self, session_id: str, update: Any, **kwargs: Any) -> None:
+        completed: asyncio.Future[None] = asyncio.get_running_loop().create_future()
+        pending = self._pending.setdefault(session_id, set())
+        pending.add(completed)
+        token = self._current_update.set(completed)
+
+        try:
+            if self._session_update is None:
+                raise RequestError.method_not_found(CLIENT_METHODS["session_update"])
+            if self._legacy_session_update:
+                _warn_legacy_handler(self._client, self._session_update_attr)
+                notification = SessionNotification(session_id=session_id, update=update, field_meta=kwargs or None)
+                await self._session_update(notification)
+            else:
+                await self._session_update(session_id=session_id, update=update, **kwargs)
+        finally:
+            self._current_update.reset(token)
+            if not completed.done():
+                completed.set_result(None)
+            pending.discard(completed)
+            if not pending:
+                self._pending.pop(session_id, None)
+
+    async def wait(self, session_id: str) -> None:
+        # Snapshot before yielding so updates received after the response are
+        # not associated with this prompt.
+        current = self._current_update.get()
+        notifications = tuple(
+            completed for completed in self._pending.get(session_id, set()) if completed is not current
+        )
+        if notifications:
+            await asyncio.gather(*(asyncio.shield(completed) for completed in notifications))
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._client, name)
+
+
 @final
 @compatible_class
 class ClientSideConnection:
@@ -69,7 +123,9 @@ class ClientSideConnection:
         **connection_kwargs: Any,
     ) -> None:
         client = to_client(self) if callable(to_client) else to_client
-        handler = build_client_router(cast(Client, client), use_unstable_protocol=use_unstable_protocol)
+        self._session_updates = _SessionUpdateTracker(cast(Client, client))
+        handler = build_client_router(cast(Client, self._session_updates), use_unstable_protocol=use_unstable_protocol)
+
         if isinstance(input_stream, Transport):
             if output_stream is not None:
                 raise TypeError(_CLIENT_CONNECTION_ERROR)
@@ -206,12 +262,18 @@ class ClientSideConnection:
         ],
         **kwargs: Any,
     ) -> PromptResponse:
-        return await request_model(
-            self._conn,
-            AGENT_METHODS["session_prompt"],
-            PromptRequest(prompt=prompt, session_id=session_id, field_meta=kwargs or None),
-            PromptResponse,
-        )
+        try:
+            response = await request_model(
+                self._conn,
+                AGENT_METHODS["session_prompt"],
+                PromptRequest(prompt=prompt, session_id=session_id, field_meta=kwargs or None),
+                PromptResponse,
+            )
+        except Exception:
+            await self._session_updates.wait(session_id)
+            raise
+        await self._session_updates.wait(session_id)
+        return response
 
     @param_model(ForkSessionRequest)
     async def fork_session(