diff --git a/README.md b/README.md index 827b00a..fd6a0a7 100644 --- a/README.md +++ b/README.md @@ -174,7 +174,9 @@ def handle_message(msg: StreamMessage) -> None: print(f"Tool call: {msg.tool_name}({msg.tool_input})") elif isinstance(msg, ToolResult): status = "❌" if msg.is_error else "✅" - print(f"{status} {msg.content}") + # tool_use_id correlates the result with its ToolUse; tool_name is + # backfilled from that call (None if the call was never seen). + print(f"{status} [{msg.tool_use_id}] {msg.tool_name}: {msg.content}") elif isinstance(msg, ToolProgress): print(f" ⏳ {msg.tool_name}: {msg.content}") elif isinstance(msg, WorkingStateChanged): @@ -184,7 +186,9 @@ def handle_message(msg: StreamMessage) -> None: elif isinstance(msg, TurnComplete): print("\n--- Turn complete ---") elif isinstance(msg, ErrorEvent): - print(f"Error [{msg.error_type}]: {msg.message}") + # error_type is often the unhelpful "Error"; error_name exposes the + # nested error.name (e.g. "LLMInvalidRequestError") to branch on. + print(f"Error [{msg.error_name or msg.error_type}]: {msg.message}") ``` ## Permission Handler @@ -264,12 +268,52 @@ async def main(): The main client class. Wraps a transport and provides typed async methods for all `droid.*` RPC methods. **Session methods:** -- `initialize_session(...)` — Create a new session +- `initialize_session(...)` — Create a new session (supports `enabled_tool_ids` and `disabled_tool_ids`) - `load_session(session_id=...)` — Load an existing session -- `add_user_message(text=...)` — Send a user message +- `add_user_message(text=..., output_format=...)` — Send a user message, optionally with a structured-output (JSON Schema) contract - `interrupt_session()` — Interrupt the current session - `kill_worker_session(worker_session_id=...)` — Kill a worker session -- `update_session_settings(...)` — Update session settings +- `update_session_settings(...)` — Update session settings (supports `enabled_tool_ids`/`disabled_tool_ids`) +- `close_session(reason=...)` — Close the active session +- `compact_session(custom_instructions=...)` — Compact the conversation to reclaim context +- `fork_session(title=..., tags=...)` — Fork the session into a new one +- `rename_session(title=...)` — Rename the session + +**Discovery methods:** +- `list_tools(...)` — List native CLI tools with `default_allowed`/`currently_allowed` (useful for locking the tool set down) +- `list_commands()` — List custom slash commands + +**Context and rewind methods:** +- `get_context_stats()` — Context-window usage (used/remaining/limit) +- `get_context_breakdown()` — Per-category/skill/MCP/droid token breakdown +- `get_rewind_info(message_id=...)` — Restorable/created/evicted files for a rewind point +- `execute_rewind(...)` — Rewind to a message, forking the session + +**Locking the tool set down:** + +```python +# enabled_tool_ids is additive, so pass an explicit disable list to +# actually restrict native tools. list_tools() lets you verify the result. +catalog = await client.list_tools() +tool_ids = [t.id for t in catalog.tools] +await client.update_session_settings(enabled_tool_ids=[], disabled_tool_ids=tool_ids) +``` + +**Structured output:** + +```python +await client.add_user_message( + text="Return an answer.", + output_format={ + "type": "json_schema", + "schema": { + "type": "object", + "properties": {"answer": {"type": "integer"}}, + "required": ["answer"], + }, + }, +) +``` **MCP methods:** - `toggle_mcp_server(...)` — Enable/disable an MCP server @@ -308,7 +352,16 @@ Protocol (interface) that all transport implementations must satisfy. Use this t uv sync # Run tests -uv run pytest +uv run --group dev python -m pytest + +# Run opt-in tests against the installed, authenticated droid exec CLI. +# These create real sessions and consume model usage. +DROID_LIVE_TESTS=1 uv run --group dev python -m pytest \ + tests/test_live_droid_exec.py -v + +# Override the executable path when droid is not on PATH. +DROID_LIVE_TESTS=1 DROID_EXEC_PATH=/path/to/droid \ + uv run --group dev python -m pytest tests/test_live_droid_exec.py -v # Type check (strict mode) uv run mypy --strict src/ diff --git a/examples/interactive_session.py b/examples/interactive_session.py index d366797..eed251c 100644 --- a/examples/interactive_session.py +++ b/examples/interactive_session.py @@ -337,6 +337,8 @@ def handle_permission(params: dict[str, object]) -> str: ) return ToolConfirmationOutcome.ProceedOnce.value + client.set_permission_handler(handle_permission) + # Initialize session print(f"Executable: {exec_path}") print(f"Working directory: {cwd}") diff --git a/pyproject.toml b/pyproject.toml index eef7028..3980cd4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "droid-sdk" -version = "0.1.2" +version = "0.1.3" description = "Python asyncio SDK for Factory Droid" readme = "README.md" license = {text = "Apache-2.0"} diff --git a/src/droid_sdk/client.py b/src/droid_sdk/client.py index 5c52676..2b95316 100644 --- a/src/droid_sdk/client.py +++ b/src/droid_sdk/client.py @@ -9,10 +9,11 @@ import asyncio import contextlib +import dataclasses import logging from collections.abc import AsyncIterator, Callable from types import TracebackType # noqa: TC003 -from typing import Any +from typing import TYPE_CHECKING, Any, Literal from droid_sdk.errors import ( ConnectionError as DroidConnectionError, @@ -21,6 +22,7 @@ SessionError, ) from droid_sdk.protocol import ( + COMPACTION_TIMEOUT, MCP_AUTH_TIMEOUT, SESSION_INIT_TIMEOUT, ProtocolEngine, @@ -30,22 +32,48 @@ ) from droid_sdk.schemas.client import ( AddMcpServerResult, + AddUserMessageRequestParams, AuthenticateMcpServerResult, Base64ImageSource, CancelMcpAuthResult, ClearMcpAuthResult, + CloseSessionRequestParams, + CloseSessionResult, + CompactSessionRequestParams, + CompactSessionResult, DocumentSource, + ExecuteRewindRequestParams, + ExecuteRewindResult, + ForkSessionRequestParams, + ForkSessionResult, + GetContextBreakdownResult, + GetContextStatsResult, + GetRewindInfoRequestParams, + GetRewindInfoResult, + InitializeSessionRequestParams, InitializeSessionResult, + ListCommandsResult, ListMcpRegistryResult, ListMcpServersResult, ListMcpToolsResult, ListSkillsResult, + ListToolsRequestParams, + ListToolsResult, + LoadSessionRequestParams, LoadSessionResult, + OutputFormat, RemoveMcpServerResult, + RenameSessionRequestParams, + RenameSessionResult, + RewindFileCreation, + RewindFileSnapshot, + SessionSource, + SessionTag, SubmitBugReportResult, SubmitMcpAuthCodeResult, ToggleMcpServerResult, ToggleMcpToolResult, + UpdateSessionSettingsRequestParams, ) from droid_sdk.schemas.enums import ( AutonomyLevel, @@ -62,12 +90,17 @@ from droid_sdk.stream import ( StreamMessage, TokenUsageUpdate, + ToolProgress, + ToolResult, TurnComplete, WorkingStateChanged, _notification_to_stream_message, ) from droid_sdk.types import DroidClientTransport # noqa: TC001 +if TYPE_CHECKING: + from pydantic import BaseModel + logger = logging.getLogger(__name__) # Type alias for notification listener callbacks @@ -269,9 +302,10 @@ async def initialize_session( decomp_mission_id: str | None = None, skip_permissions_unsafe: bool | None = None, enabled_tool_ids: list[str] | None = None, + disabled_tool_ids: list[str] | None = None, session_location: str | None = None, - session_source: dict[str, Any] | None = None, - tags: list[dict[str, Any]] | None = None, + session_source: SessionSource | dict[str, Any] | None = None, + tags: list[SessionTag | dict[str, Any]] | None = None, mcp_oauth_callback_uri: str | None = None, ) -> InitializeSessionResult: """Initialize a new session. @@ -296,7 +330,10 @@ async def initialize_session( decomp_session_type: Session type for mission decomposition. decomp_mission_id: Mission ID for worker sessions. skip_permissions_unsafe: Skip permission checks. - enabled_tool_ids: Additional tool IDs to enable. + enabled_tool_ids: Additional tool IDs to enable beyond defaults. + disabled_tool_ids: Tool IDs to disable (subtractive). Combine with + ``enabled_tool_ids=[]`` and an explicit disable list to lock + the tool set down. session_location: Session metadata location. session_source: Session source information. tags: Optional session tags. @@ -313,33 +350,43 @@ async def initialize_session( self._ensure_not_closed() protocol = self._ensure_protocol() - params: dict[str, Any] = { - "machineId": machine_id, - "cwd": cwd, - } - # Add optional params only if set - _set_if_not_none(params, "sessionId", session_id) - _set_if_not_none(params, "workspaceId", workspace_id) - _set_if_not_none(params, "mcpServers", mcp_servers) - _set_if_not_none(params, "autonomyMode", _enum_value(autonomy_mode)) - _set_if_not_none(params, "interactionMode", _enum_value(interaction_mode)) - _set_if_not_none(params, "autonomyLevel", _enum_value(autonomy_level)) - _set_if_not_none(params, "modelId", model_id) - _set_if_not_none(params, "reasoningEffort", _enum_value(reasoning_effort)) - _set_if_not_none(params, "specModeModelId", spec_mode_model_id) - _set_if_not_none( - params, - "specModeReasoningEffort", - _enum_value(spec_mode_reasoning_effort), + validated_session_source = ( + SessionSource.model_validate(session_source) + if session_source is not None + else None + ) + validated_tags = ( + [SessionTag.model_validate(tag) for tag in tags] + if tags is not None + else None ) - _set_if_not_none(params, "decompSessionType", _enum_value(decomp_session_type)) - _set_if_not_none(params, "decompMissionId", decomp_mission_id) - _set_if_not_none(params, "skipPermissionsUnsafe", skip_permissions_unsafe) - _set_if_not_none(params, "enabledToolIds", enabled_tool_ids) - _set_if_not_none(params, "sessionLocation", session_location) - _set_if_not_none(params, "sessionSource", session_source) - _set_if_not_none(params, "tags", tags) - _set_if_not_none(params, "mcpOAuthCallbackUri", mcp_oauth_callback_uri) + params = _serialize_params( + InitializeSessionRequestParams( + machine_id=machine_id, + cwd=cwd, + session_id=session_id, + workspace_id=workspace_id, + autonomy_mode=autonomy_mode, + interaction_mode=interaction_mode, + autonomy_level=autonomy_level, + model_id=model_id, + reasoning_effort=reasoning_effort, + spec_mode_model_id=spec_mode_model_id, + spec_mode_reasoning_effort=spec_mode_reasoning_effort, + decomp_session_type=decomp_session_type, + decomp_mission_id=decomp_mission_id, + skip_permissions_unsafe=skip_permissions_unsafe, + enabled_tool_ids=enabled_tool_ids, + disabled_tool_ids=disabled_tool_ids, + session_location=session_location, + session_source=validated_session_source, + tags=validated_tags, + mcp_oauth_callback_uri=mcp_oauth_callback_uri, + ) + ) + # Preserve compatibility with legacy unvalidated MCP dictionaries. + if mcp_servers is not None: + params["mcpServers"] = mcp_servers response = await protocol.send_request( method=DroidServerMethod.INITIALIZE_SESSION.value, @@ -380,11 +427,15 @@ async def load_session( self._ensure_not_closed() protocol = self._ensure_protocol() - params: dict[str, Any] = { - "sessionId": session_id, - } - _set_if_not_none(params, "mcpServers", mcp_servers) - _set_if_not_none(params, "mcpOAuthCallbackUri", mcp_oauth_callback_uri) + params = _serialize_params( + LoadSessionRequestParams( + session_id=session_id, + mcp_oauth_callback_uri=mcp_oauth_callback_uri, + ) + ) + # Preserve compatibility with legacy unvalidated MCP dictionaries. + if mcp_servers is not None: + params["mcpServers"] = mcp_servers response = await protocol.send_request( method=DroidServerMethod.LOAD_SESSION.value, @@ -402,6 +453,7 @@ async def add_user_message( text: str, images: list[Base64ImageSource | dict[str, Any]] | None = None, files: list[DocumentSource | dict[str, Any]] | None = None, + output_format: OutputFormat | dict[str, Any] | None = None, request_id: str | None = None, ) -> None: """Add a user message to the session. @@ -412,6 +464,10 @@ async def add_user_message( text: Message text content. images: Optional attached images. files: Optional attached documents. + output_format: Optional structured-output contract. Either an + :class:`~droid_sdk.schemas.client.OutputFormat` or a raw dict + of the form ``{"type": "json_schema", "schema": {...}}`` to + constrain the reply to a JSON value matching the schema. request_id: Optional custom request ID for the JSON-RPC envelope. Raises: @@ -423,9 +479,21 @@ async def add_user_message( self._ensure_session() protocol = self._ensure_protocol() - params: dict[str, Any] = {"text": text} - _set_if_not_none(params, "images", images) - _set_if_not_none(params, "files", files) + validated_output_format = ( + OutputFormat.model_validate(output_format) + if output_format is not None + else None + ) + params = _serialize_params( + AddUserMessageRequestParams( + text=text, + output_format=validated_output_format, + ) + ) + if images is not None: + params["images"] = images + if files is not None: + params["files"] = files # Use custom request_id by monkey-patching the protocol temporarily, # or pass via overridden send_request. The protocol engine generates @@ -496,6 +564,8 @@ async def update_session_settings( autonomy_level: AutonomyLevel | None = None, spec_mode_model_id: str | None = None, spec_mode_reasoning_effort: ReasoningEffort | None = None, + enabled_tool_ids: list[str] | None = None, + disabled_tool_ids: list[str] | None = None, ) -> None: """Update session settings. @@ -510,6 +580,8 @@ async def update_session_settings( autonomy_level: Optional autonomy level. spec_mode_model_id: Optional spec mode model ID. spec_mode_reasoning_effort: Optional spec mode reasoning effort. + enabled_tool_ids: Additional tool IDs to enable beyond defaults. + disabled_tool_ids: Tool IDs to disable (subtractive). Raises: SessionError: If no active session. @@ -520,17 +592,18 @@ async def update_session_settings( self._ensure_session() protocol = self._ensure_protocol() - params: dict[str, Any] = {} - _set_if_not_none(params, "modelId", model_id) - _set_if_not_none(params, "reasoningEffort", _enum_value(reasoning_effort)) - _set_if_not_none(params, "autonomyMode", _enum_value(autonomy_mode)) - _set_if_not_none(params, "interactionMode", _enum_value(interaction_mode)) - _set_if_not_none(params, "autonomyLevel", _enum_value(autonomy_level)) - _set_if_not_none(params, "specModeModelId", spec_mode_model_id) - _set_if_not_none( - params, - "specModeReasoningEffort", - _enum_value(spec_mode_reasoning_effort), + params = _serialize_params( + UpdateSessionSettingsRequestParams( + model_id=model_id, + reasoning_effort=reasoning_effort, + autonomy_mode=autonomy_mode, + interaction_mode=interaction_mode, + autonomy_level=autonomy_level, + spec_mode_model_id=spec_mode_model_id, + spec_mode_reasoning_effort=spec_mode_reasoning_effort, + enabled_tool_ids=enabled_tool_ids, + disabled_tool_ids=disabled_tool_ids, + ) ) await protocol.send_request( @@ -1004,6 +1077,376 @@ async def submit_bug_report( return SubmitBugReportResult.model_validate(response.get("result", {})) + # ---------------------------------------------------------- + # Tool / command discovery + # ---------------------------------------------------------- + + async def list_tools( + self, + *, + enabled_tool_ids: list[str] | None = None, + disabled_tool_ids: list[str] | None = None, + ) -> ListToolsResult: + """List native CLI tools with their allow-state. + + Sends ``droid.list_tools``. This can be called before session + initialization, allowing callers to discover tool IDs for initial + allow/deny settings. The result reports each tool's ``id`` plus + ``default_allowed`` and ``currently_allowed``. + + Args: + enabled_tool_ids: Optional hypothetical enable list to evaluate + ``currently_allowed`` against (does not mutate the session). + disabled_tool_ids: Optional hypothetical disable list. + + Returns: + Typed ``ListToolsResult`` with a ``tools`` list. + + Raises: + ProtocolError: If the server returns an error. + ConnectionError: If the client has been closed. + """ + self._ensure_not_closed() + protocol = self._ensure_protocol() + + params = _serialize_params( + ListToolsRequestParams( + enabled_tool_ids=enabled_tool_ids, + disabled_tool_ids=disabled_tool_ids, + ) + ) + + response = await protocol.send_request( + method=DroidServerMethod.LIST_TOOLS.value, + params=params, + ) + + return ListToolsResult.model_validate(response.get("result", {})) + + async def list_commands(self) -> ListCommandsResult: + """List custom slash commands. + + Sends ``droid.list_commands``. Requires an active session. + + Returns: + Typed ``ListCommandsResult`` with a ``commands`` list. + + Raises: + SessionError: If no active session. + ProtocolError: If the server returns an error. + ConnectionError: If the client has been closed. + """ + self._ensure_not_closed() + self._ensure_session() + protocol = self._ensure_protocol() + + response = await protocol.send_request( + method=DroidServerMethod.LIST_COMMANDS.value, + params={}, + ) + + return ListCommandsResult.model_validate(response.get("result", {})) + + # ---------------------------------------------------------- + # Session lifecycle: close / compact / fork / rename + # ---------------------------------------------------------- + + async def close_session( + self, + *, + reason: Literal["clear", "logout", "prompt_input_exit", "other"] | None = None, + ) -> CloseSessionResult: + """Close the active session. + + Sends ``droid.close_session``. Requires an active session. + + Args: + reason: Optional close reason (``"clear"``, ``"logout"``, + ``"prompt_input_exit"`` or ``"other"``). + + Returns: + Typed ``CloseSessionResult`` (empty payload). + + Raises: + SessionError: If no active session. + ProtocolError: If the server returns an error. + ConnectionError: If the client has been closed. + """ + self._ensure_not_closed() + self._ensure_session() + protocol = self._ensure_protocol() + + params = _serialize_params(CloseSessionRequestParams(reason=reason)) + + response = await protocol.send_request( + method=DroidServerMethod.CLOSE_SESSION.value, + params=params, + ) + + result = CloseSessionResult.model_validate(response.get("result", {})) + self._session_id = None + return result + + async def compact_session( + self, + *, + custom_instructions: str | None = None, + ) -> CompactSessionResult: + """Compact the conversation to reclaim context. + + Sends ``droid.compact_session`` with an extended timeout, since + compaction runs an LLM summarization pass. Requires an active + session. + + Args: + custom_instructions: Optional instructions to steer the + compaction summary. + + Returns: + Typed ``CompactSessionResult`` with ``new_session_id`` and + ``removed_count``. + + Raises: + SessionError: If no active session. + ProtocolError: If the server returns an error. + ConnectionError: If the client has been closed. + """ + self._ensure_not_closed() + self._ensure_session() + protocol = self._ensure_protocol() + + params = _serialize_params( + CompactSessionRequestParams(custom_instructions=custom_instructions) + ) + + response = await protocol.send_request( + method=DroidServerMethod.COMPACT_SESSION.value, + params=params, + timeout=COMPACTION_TIMEOUT, + ) + + return CompactSessionResult.model_validate(response.get("result", {})) + + async def fork_session( + self, + *, + title: str | None = None, + tags: list[SessionTag | dict[str, Any]] | None = None, + ) -> ForkSessionResult: + """Fork the active session into a new one. + + Sends ``droid.fork_session``. Requires an active session. + + Args: + title: Optional title for the fork. + tags: Optional session tags for the fork. + + Returns: + Typed ``ForkSessionResult`` with ``new_session_id``. + + Raises: + SessionError: If no active session. + ProtocolError: If the server returns an error. + ConnectionError: If the client has been closed. + """ + self._ensure_not_closed() + self._ensure_session() + protocol = self._ensure_protocol() + + validated_tags = ( + [SessionTag.model_validate(tag) for tag in tags] + if tags is not None + else None + ) + params = _serialize_params( + ForkSessionRequestParams(title=title, tags=validated_tags) + ) + + response = await protocol.send_request( + method=DroidServerMethod.FORK_SESSION.value, + params=params, + ) + + return ForkSessionResult.model_validate(response.get("result", {})) + + async def rename_session(self, *, title: str) -> RenameSessionResult: + """Rename the active session. + + Sends ``droid.rename_session``. Requires an active session. + + Args: + title: New session title. + + Returns: + Typed ``RenameSessionResult`` with a ``success`` flag. + + Raises: + SessionError: If no active session. + ProtocolError: If the server returns an error. + ConnectionError: If the client has been closed. + """ + self._ensure_not_closed() + self._ensure_session() + protocol = self._ensure_protocol() + + params = _serialize_params(RenameSessionRequestParams(title=title)) + + response = await protocol.send_request( + method=DroidServerMethod.RENAME_SESSION.value, + params=params, + ) + + return RenameSessionResult.model_validate(response.get("result", {})) + + # ---------------------------------------------------------- + # Context introspection + # ---------------------------------------------------------- + + async def get_context_stats(self) -> GetContextStatsResult: + """Get context-window usage statistics. + + Sends ``droid.get_context_stats``. Requires an active session. + + Returns: + Typed ``GetContextStatsResult`` with used/remaining/limit tokens. + + Raises: + SessionError: If no active session. + ProtocolError: If the server returns an error. + ConnectionError: If the client has been closed. + """ + self._ensure_not_closed() + self._ensure_session() + protocol = self._ensure_protocol() + + response = await protocol.send_request( + method=DroidServerMethod.GET_CONTEXT_STATS.value, + params={}, + ) + + return GetContextStatsResult.model_validate(response.get("result", {})) + + async def get_context_breakdown(self) -> GetContextBreakdownResult: + """Get a detailed breakdown of context-window usage. + + Sends ``droid.get_context_breakdown``. Requires an active session. + + Returns: + Typed ``GetContextBreakdownResult`` with per-category, per-skill, + per-MCP-server and per-droid token usage. + + Raises: + SessionError: If no active session. + ProtocolError: If the server returns an error. + ConnectionError: If the client has been closed. + """ + self._ensure_not_closed() + self._ensure_session() + protocol = self._ensure_protocol() + + response = await protocol.send_request( + method=DroidServerMethod.GET_CONTEXT_BREAKDOWN.value, + params={}, + ) + + return GetContextBreakdownResult.model_validate(response.get("result", {})) + + # ---------------------------------------------------------- + # Rewind + # ---------------------------------------------------------- + + async def get_rewind_info(self, *, message_id: str) -> GetRewindInfoResult: + """Get file-restore information for rewinding to a message. + + Sends ``droid.get_rewind_info``. Requires an active session. + + Args: + message_id: The message to rewind to. + + Returns: + Typed ``GetRewindInfoResult`` with restorable, created and + evicted file lists. + + Raises: + SessionError: If no active session. + ProtocolError: If the server returns an error. + ConnectionError: If the client has been closed. + """ + self._ensure_not_closed() + session_id = self._ensure_session() + protocol = self._ensure_protocol() + + params = _serialize_params( + GetRewindInfoRequestParams( + session_id=session_id, + message_id=message_id, + ) + ) + + response = await protocol.send_request( + method=DroidServerMethod.GET_REWIND_INFO.value, + params=params, + ) + + return GetRewindInfoResult.model_validate(response.get("result", {})) + + async def execute_rewind( + self, + *, + message_id: str, + files_to_restore: list[RewindFileSnapshot | dict[str, Any]], + files_to_delete: list[RewindFileCreation | dict[str, Any]], + fork_title: str, + ) -> ExecuteRewindResult: + """Execute a rewind, forking the session at ``message_id``. + + Sends ``droid.execute_rewind`` with an extended timeout. Requires + an active session. + + Args: + message_id: The message to rewind to. + files_to_restore: File snapshots to restore, each + ``{"filePath", "contentHash", "size"}``. + files_to_delete: Files to delete, each ``{"filePath"}``. + fork_title: Title for the new forked session. + + Returns: + Typed ``ExecuteRewindResult`` with ``new_session_id`` and + restore/delete counts. + + Raises: + SessionError: If no active session. + ProtocolError: If the server returns an error. + ConnectionError: If the client has been closed. + """ + self._ensure_not_closed() + session_id = self._ensure_session() + protocol = self._ensure_protocol() + + validated_files_to_restore = [ + RewindFileSnapshot.model_validate(file) for file in files_to_restore + ] + validated_files_to_delete = [ + RewindFileCreation.model_validate(file) for file in files_to_delete + ] + params = _serialize_params( + ExecuteRewindRequestParams( + session_id=session_id, + message_id=message_id, + files_to_restore=validated_files_to_restore, + files_to_delete=validated_files_to_delete, + fork_title=fork_title, + ) + ) + + response = await protocol.send_request( + method=DroidServerMethod.EXECUTE_REWIND.value, + params=params, + timeout=SESSION_INIT_TIMEOUT, + ) + + return ExecuteRewindResult.model_validate(response.get("result", {})) + # ---------------------------------------------------------- # Streaming: receive_response() async iterator # ---------------------------------------------------------- @@ -1042,6 +1485,9 @@ async def receive_response(self) -> AsyncIterator[StreamMessage]: queue: asyncio.Queue[StreamMessage | None] = asyncio.Queue() was_not_idle = False last_token_usage: TokenUsageUpdate | None = None + # Correlate tool_use_id -> tool_name so tool_result / tool_progress + # notifications (which omit the name) can be backfilled. + tool_names: dict[str, str] = {} def _on_notification(notification_dict: dict[str, Any]) -> None: nonlocal was_not_idle, last_token_usage @@ -1065,9 +1511,21 @@ def _on_notification(notification_dict: dict[str, Any]) -> None: # Handle list of ToolUse (from CREATE_MESSAGE) if isinstance(result, list): for item in result: + tool_names[item.tool_use_id] = item.tool_name queue.put_nowait(item) return + # Backfill the tool name on results/progress from the matching + # tool_use, correlating via tool_use_id. + if isinstance(result, (ToolResult, ToolProgress)): + should_backfill_tool_name = ( + isinstance(result, ToolResult) and result.tool_name is None + ) or (isinstance(result, ToolProgress) and not result.tool_name) + if should_backfill_tool_name and result.tool_use_id is not None: + name = tool_names.get(result.tool_use_id) + if name is not None: + result = dataclasses.replace(result, tool_name=name) + # Track token usage for TurnComplete if isinstance(result, TokenUsageUpdate): last_token_usage = result @@ -1248,12 +1706,13 @@ def _ensure_not_closed(self) -> None: "or call connect() to reconnect." ) - def _ensure_session(self) -> None: - """Raise SessionError if no active session.""" + def _ensure_session(self) -> str: + """Return the active session ID, or raise SessionError if absent.""" if self._session_id is None: raise SessionError( "No active session. Call initialize_session or load_session first." ) + return self._session_id def _ensure_protocol(self) -> ProtocolEngine: """Return the protocol engine, raising if not connected.""" @@ -1268,6 +1727,15 @@ def _set_if_not_none(d: dict[str, Any], key: str, value: Any) -> None: d[key] = value +def _serialize_params(model: BaseModel) -> dict[str, Any]: + """Serialize validated request parameters to their wire aliases.""" + return model.model_dump( + mode="json", + by_alias=True, + exclude_none=True, + ) + + def _enum_value(val: Any) -> Any: """Extract .value from an enum, or return None.""" if val is None: diff --git a/src/droid_sdk/protocol.py b/src/droid_sdk/protocol.py index 47e452b..a9e87ef 100644 --- a/src/droid_sdk/protocol.py +++ b/src/droid_sdk/protocol.py @@ -52,6 +52,9 @@ SESSION_INIT_TIMEOUT: Final[float] = 60.0 """Extended timeout for session initialization (60 seconds).""" +COMPACTION_TIMEOUT: Final[float] = 240.0 +"""Extended timeout for session compaction (4 minutes).""" + MCP_AUTH_TIMEOUT: Final[float] = 300.0 """Extended timeout for MCP OAuth authentication (5 minutes).""" @@ -631,6 +634,7 @@ def _reject_all_pending(self, error: Exception) -> None: __all__ = [ + "COMPACTION_TIMEOUT", "DEFAULT_REQUEST_TIMEOUT", "MCP_AUTH_TIMEOUT", "SESSION_INIT_TIMEOUT", diff --git a/src/droid_sdk/schemas/__init__.py b/src/droid_sdk/schemas/__init__.py index dbe9ee1..468ead2 100644 --- a/src/droid_sdk/schemas/__init__.py +++ b/src/droid_sdk/schemas/__init__.py @@ -80,7 +80,41 @@ ClearMcpAuthResponse, ClearMcpAuthResult, ClientRequest, + CloseSessionRequest, + CloseSessionRequestParams, + CloseSessionResponse, + CloseSessionResult, + CompactSessionRequest, + CompactSessionRequestParams, + CompactSessionResponse, + CompactSessionResult, + ContextBreakdownCategory, + ContextBreakdownDroidEntry, + ContextBreakdownMcpServerEntry, + ContextBreakdownSkillEntry, + CustomCommandInfo, DocumentSource, + ExecToolInfo, + ExecuteRewindRequest, + ExecuteRewindRequestParams, + ExecuteRewindResponse, + ExecuteRewindResult, + ForkSessionRequest, + ForkSessionRequestParams, + ForkSessionResponse, + ForkSessionResult, + GetContextBreakdownRequest, + GetContextBreakdownRequestParams, + GetContextBreakdownResponse, + GetContextBreakdownResult, + GetContextStatsRequest, + GetContextStatsRequestParams, + GetContextStatsResponse, + GetContextStatsResult, + GetRewindInfoRequest, + GetRewindInfoRequestParams, + GetRewindInfoResponse, + GetRewindInfoResult, GitRepoInfo, HttpHeader, HttpMcpConfig, @@ -96,6 +130,10 @@ KillWorkerSessionRequestParams, KillWorkerSessionResponse, KillWorkerSessionResult, + ListCommandsRequest, + ListCommandsRequestParams, + ListCommandsResponse, + ListCommandsResult, ListMcpRegistryRequest, ListMcpRegistryRequestParams, ListMcpRegistryResponse, @@ -112,15 +150,27 @@ ListSkillsRequestParams, ListSkillsResponse, ListSkillsResult, + ListToolsRequest, + ListToolsRequestParams, + ListToolsResponse, + ListToolsResult, LoadSessionRequest, LoadSessionRequestParams, LoadSessionResponse, LoadSessionResult, MissionSnapshot, + OutputFormat, RemoveMcpServerRequest, RemoveMcpServerRequestParams, RemoveMcpServerResponse, RemoveMcpServerResult, + RenameSessionRequest, + RenameSessionRequestParams, + RenameSessionResponse, + RenameSessionResult, + RewindEvictedFile, + RewindFileCreation, + RewindFileSnapshot, SessionSettings, SessionSource, SessionTag, @@ -310,10 +360,23 @@ "CliRequestOrNotification", "ClientRequest", "ClientType", + "CloseSessionRequest", + "CloseSessionRequestParams", + "CloseSessionResponse", + "CloseSessionResult", + "CompactSessionRequest", + "CompactSessionRequestParams", + "CompactSessionResponse", + "CompactSessionResult", # --- Messages --- "ContentBlock", + "ContextBreakdownCategory", + "ContextBreakdownDroidEntry", + "ContextBreakdownMcpServerEntry", + "ContextBreakdownSkillEntry", "CreateMessageNotification", "CreateToolConfirmationDetails", + "CustomCommandInfo", "DecompSessionType", # --- Mission --- "DiscoveredIssue", @@ -333,11 +396,32 @@ "EditToolConfirmationDetails", "ErrorDetail", "ErrorNotification", + "ExecToolInfo", + "ExecuteRewindRequest", + "ExecuteRewindRequestParams", + "ExecuteRewindResponse", + "ExecuteRewindResult", "ExecuteToolConfirmationDetails", "ExitSpecModeConfirmationDetails", "FactoryDroidMessage", "FeatureStatus", "FeatureSuccessState", + "ForkSessionRequest", + "ForkSessionRequestParams", + "ForkSessionResponse", + "ForkSessionResult", + "GetContextBreakdownRequest", + "GetContextBreakdownRequestParams", + "GetContextBreakdownResponse", + "GetContextBreakdownResult", + "GetContextStatsRequest", + "GetContextStatsRequestParams", + "GetContextStatsResponse", + "GetContextStatsResult", + "GetRewindInfoRequest", + "GetRewindInfoRequestParams", + "GetRewindInfoResponse", + "GetRewindInfoResult", "GitRepoInfo", "Handoff", "HandoffItemsDismissedEntry", @@ -366,6 +450,10 @@ "KillWorkerSessionRequestParams", "KillWorkerSessionResponse", "KillWorkerSessionResult", + "ListCommandsRequest", + "ListCommandsRequestParams", + "ListCommandsResponse", + "ListCommandsResult", "ListMcpRegistryRequest", "ListMcpRegistryRequestParams", "ListMcpRegistryResponse", @@ -382,6 +470,10 @@ "ListSkillsRequestParams", "ListSkillsResponse", "ListSkillsResult", + "ListToolsRequest", + "ListToolsRequestParams", + "ListToolsResponse", + "ListToolsResult", "LoadSessionRequest", "LoadSessionRequestParams", "LoadSessionResponse", @@ -421,6 +513,7 @@ "MissionWorkerCompletedNotification", "MissionWorkerStartedNotification", "ModelProvider", + "OutputFormat", "PermissionResolvedNotification", "Platform", "ProgressLogEntry", @@ -432,9 +525,16 @@ "RemoveMcpServerRequestParams", "RemoveMcpServerResponse", "RemoveMcpServerResult", + "RenameSessionRequest", + "RenameSessionRequestParams", + "RenameSessionResponse", + "RenameSessionResult", "RequestPermissionRequest", "RequestPermissionRequestParams", "RequestPermissionResult", + "RewindEvictedFile", + "RewindFileCreation", + "RewindFileSnapshot", "SessionNotification", "SessionNotificationParams", "SessionNotificationType", diff --git a/src/droid_sdk/schemas/client.py b/src/droid_sdk/schemas/client.py index 994a934..2ba33ce 100644 --- a/src/droid_sdk/schemas/client.py +++ b/src/droid_sdk/schemas/client.py @@ -1,10 +1,10 @@ """Client→server request/response Pydantic schemas for the Factory Droid protocol. -All 19 client→server RPC method request/response pairs, plus supporting types +All 29 client→server RPC method request/response pairs, plus supporting types and the ClientRequest discriminated union. Ported from TypeScript source: -- packages/common/src/droid/schemas/client.ts +- packages/droid-sdk-core/src/protocol/droid/schemas/client.ts """ from __future__ import annotations @@ -70,7 +70,41 @@ "ClearMcpAuthResponse", "ClearMcpAuthResult", "ClientRequest", + "CloseSessionRequest", + "CloseSessionRequestParams", + "CloseSessionResponse", + "CloseSessionResult", + "CompactSessionRequest", + "CompactSessionRequestParams", + "CompactSessionResponse", + "CompactSessionResult", + "ContextBreakdownCategory", + "ContextBreakdownDroidEntry", + "ContextBreakdownMcpServerEntry", + "ContextBreakdownSkillEntry", + "CustomCommandInfo", "DocumentSource", + "ExecToolInfo", + "ExecuteRewindRequest", + "ExecuteRewindRequestParams", + "ExecuteRewindResponse", + "ExecuteRewindResult", + "ForkSessionRequest", + "ForkSessionRequestParams", + "ForkSessionResponse", + "ForkSessionResult", + "GetContextBreakdownRequest", + "GetContextBreakdownRequestParams", + "GetContextBreakdownResponse", + "GetContextBreakdownResult", + "GetContextStatsRequest", + "GetContextStatsRequestParams", + "GetContextStatsResponse", + "GetContextStatsResult", + "GetRewindInfoRequest", + "GetRewindInfoRequestParams", + "GetRewindInfoResponse", + "GetRewindInfoResult", "GitRepoInfo", "HttpHeader", "HttpMcpConfig", @@ -86,6 +120,10 @@ "KillWorkerSessionRequestParams", "KillWorkerSessionResponse", "KillWorkerSessionResult", + "ListCommandsRequest", + "ListCommandsRequestParams", + "ListCommandsResponse", + "ListCommandsResult", "ListMcpRegistryRequest", "ListMcpRegistryRequestParams", "ListMcpRegistryResponse", @@ -102,15 +140,27 @@ "ListSkillsRequestParams", "ListSkillsResponse", "ListSkillsResult", + "ListToolsRequest", + "ListToolsRequestParams", + "ListToolsResponse", + "ListToolsResult", "LoadSessionRequest", "LoadSessionRequestParams", "LoadSessionResponse", "LoadSessionResult", "MissionSnapshot", + "OutputFormat", "RemoveMcpServerRequest", "RemoveMcpServerRequestParams", "RemoveMcpServerResponse", "RemoveMcpServerResult", + "RenameSessionRequest", + "RenameSessionRequestParams", + "RenameSessionResponse", + "RenameSessionResult", + "RewindEvictedFile", + "RewindFileCreation", + "RewindFileSnapshot", "SessionSettings", "SessionSource", "SessionTag", @@ -189,6 +239,22 @@ class DocumentSource(BaseModel): """Optional additional MIME type info.""" +class OutputFormat(BaseModel): + """Structured-output contract for a user message or session. + + Constrains the assistant's reply to a JSON value matching the given + JSON Schema. + """ + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + type: Literal["json_schema"] + """Output format type, always 'json_schema'.""" + + schema_: dict[str, Any] = Field(alias="schema") + """JSON Schema describing the required output shape.""" + + class SessionTag(BaseModel): """Session tag metadata.""" @@ -201,6 +267,42 @@ class SessionTag(BaseModel): """Optional key-value metadata.""" +class RewindFileSnapshot(BaseModel): + """A file snapshot that can be restored during a rewind.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + file_path: str = Field(alias="filePath") + """Path to the file.""" + + content_hash: str = Field(alias="contentHash") + """Content hash of the snapshot.""" + + size: int + """File size in bytes.""" + + +class RewindFileCreation(BaseModel): + """A file created after the rewind point (deleted on rewind).""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + file_path: str = Field(alias="filePath") + """Path to the file.""" + + +class RewindEvictedFile(BaseModel): + """A file that cannot be restored, with the reason why.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + file_path: str = Field(alias="filePath") + """Path to the file.""" + + reason: str + """Why the file cannot be restored.""" + + class SessionSource(BaseModel): """Session source information. @@ -537,6 +639,9 @@ class InitializeSessionRequestParams(BaseModel): enabled_tool_ids: list[str] | None = Field(default=None, alias="enabledToolIds") """Additional tool IDs to enable beyond defaults.""" + disabled_tool_ids: list[str] | None = Field(default=None, alias="disabledToolIds") + """Tool IDs to disable (subtractive; applied on top of the default set).""" + session_location: str | None = Field(default=None, alias="sessionLocation") """Session metadata location.""" @@ -588,6 +693,9 @@ class AddUserMessageRequestParams(BaseModel): files: list[DocumentSource] | None = None """Optional attached files.""" + output_format: OutputFormat | None = Field(default=None, alias="outputFormat") + """Optional structured-output (JSON Schema) contract for the reply.""" + class InterruptSessionRequestParams(BaseModel): """Parameters for droid.interrupt_session request (empty).""" @@ -636,6 +744,12 @@ class UpdateSessionSettingsRequestParams(BaseModel): ) """Optional spec mode reasoning effort (nullable to clear).""" + enabled_tool_ids: list[str] | None = Field(default=None, alias="enabledToolIds") + """Additional tool IDs to enable beyond defaults.""" + + disabled_tool_ids: list[str] | None = Field(default=None, alias="disabledToolIds") + """Tool IDs to disable (subtractive).""" + class ToggleMcpServerRequestParams(BaseModel): """Parameters for droid.toggle_mcp_server request.""" @@ -790,6 +904,108 @@ class SubmitBugReportRequestParams(BaseModel): """Optional client log data.""" +class ListToolsRequestParams(BaseModel): + """Parameters for droid.list_tools request.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + enabled_tool_ids: list[str] | None = Field(default=None, alias="enabledToolIds") + """Optional hypothetical additional tool IDs.""" + + disabled_tool_ids: list[str] | None = Field(default=None, alias="disabledToolIds") + """Optional hypothetical disabled tool IDs.""" + + +class ListCommandsRequestParams(BaseModel): + """Parameters for droid.list_commands request (empty).""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + +class CloseSessionRequestParams(BaseModel): + """Parameters for droid.close_session request.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + reason: Literal["clear", "logout", "prompt_input_exit", "other"] | None = None + """Optional reason for closing the session.""" + + +class CompactSessionRequestParams(BaseModel): + """Parameters for droid.compact_session request.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + custom_instructions: str | None = Field(default=None, alias="customInstructions") + """Optional instructions for the compaction summary.""" + + +class ForkSessionRequestParams(BaseModel): + """Parameters for droid.fork_session request.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + title: str | None = None + """Optional title for the fork.""" + + tags: list[SessionTag] | None = None + """Optional tags for the fork.""" + + +class RenameSessionRequestParams(BaseModel): + """Parameters for droid.rename_session request.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + title: str + """New session title.""" + + +class GetContextStatsRequestParams(BaseModel): + """Parameters for droid.get_context_stats request (empty).""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + +class GetContextBreakdownRequestParams(BaseModel): + """Parameters for droid.get_context_breakdown request (empty).""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + +class GetRewindInfoRequestParams(BaseModel): + """Parameters for droid.get_rewind_info request.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + session_id: str = Field(alias="sessionId") + """Session containing the rewind point.""" + + message_id: str = Field(alias="messageId") + """Message identifying the rewind point.""" + + +class ExecuteRewindRequestParams(BaseModel): + """Parameters for droid.execute_rewind request.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + + session_id: str = Field(alias="sessionId") + """Session containing the rewind point.""" + + message_id: str = Field(alias="messageId") + """Message identifying the rewind point.""" + + files_to_restore: list[RewindFileSnapshot] = Field(alias="filesToRestore") + """File snapshots to restore.""" + + files_to_delete: list[RewindFileCreation] = Field(alias="filesToDelete") + """Files to delete.""" + + fork_title: str = Field(alias="forkTitle") + """Title for the rewind fork.""" + + # ============================================================ # Request schemas (JsonRpcRequest + method literal + typed params) # ============================================================ @@ -1023,6 +1239,86 @@ class SubmitBugReportRequest(JsonRpcRequest): """Typed request parameters.""" +class ListToolsRequest(JsonRpcRequest): + """Request to list native CLI tools.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + method: Literal[DroidServerMethod.LIST_TOOLS] + params: ListToolsRequestParams # type: ignore[assignment] + + +class ListCommandsRequest(JsonRpcRequest): + """Request to list custom slash commands.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + method: Literal[DroidServerMethod.LIST_COMMANDS] + params: ListCommandsRequestParams # type: ignore[assignment] + + +class CloseSessionRequest(JsonRpcRequest): + """Request to close the active session.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + method: Literal[DroidServerMethod.CLOSE_SESSION] + params: CloseSessionRequestParams # type: ignore[assignment] + + +class CompactSessionRequest(JsonRpcRequest): + """Request to compact the active session.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + method: Literal[DroidServerMethod.COMPACT_SESSION] + params: CompactSessionRequestParams # type: ignore[assignment] + + +class ForkSessionRequest(JsonRpcRequest): + """Request to fork the active session.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + method: Literal[DroidServerMethod.FORK_SESSION] + params: ForkSessionRequestParams # type: ignore[assignment] + + +class RenameSessionRequest(JsonRpcRequest): + """Request to rename the active session.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + method: Literal[DroidServerMethod.RENAME_SESSION] + params: RenameSessionRequestParams # type: ignore[assignment] + + +class GetContextStatsRequest(JsonRpcRequest): + """Request for context-window usage statistics.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + method: Literal[DroidServerMethod.GET_CONTEXT_STATS] + params: GetContextStatsRequestParams # type: ignore[assignment] + + +class GetContextBreakdownRequest(JsonRpcRequest): + """Request for detailed context-window usage.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + method: Literal[DroidServerMethod.GET_CONTEXT_BREAKDOWN] + params: GetContextBreakdownRequestParams # type: ignore[assignment] + + +class GetRewindInfoRequest(JsonRpcRequest): + """Request for file restore information at a rewind point.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + method: Literal[DroidServerMethod.GET_REWIND_INFO] + params: GetRewindInfoRequestParams # type: ignore[assignment] + + +class ExecuteRewindRequest(JsonRpcRequest): + """Request to execute a rewind and fork the session.""" + + model_config = ConfigDict(populate_by_name=True, extra="forbid") + method: Literal[DroidServerMethod.EXECUTE_REWIND] + params: ExecuteRewindRequestParams # type: ignore[assignment] + + # ============================================================ # Result schemas # @@ -1300,6 +1596,280 @@ class SubmitBugReportResult(BaseModel): """Created bug report ID.""" +# ============================================================ +# Tool / command discovery (droid.list_tools, droid.list_commands) +# ============================================================ + + +class ExecToolInfo(BaseModel): + """A native CLI tool entry returned by droid.list_tools.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + id: str + """Tool identifier (the ID used in enabled/disabled tool lists).""" + + llm_id: str | None = Field(default=None, alias="llmId") + """Identifier presented to the model.""" + + display_name: str | None = Field(default=None, alias="displayName") + """Human-readable display name.""" + + description: str | None = None + """Tool description.""" + + category: str | None = None + """Tool catalog category.""" + + default_allowed: bool = Field(alias="defaultAllowed") + """Whether the tool is allowed by default.""" + + currently_allowed: bool = Field(alias="currentlyAllowed") + """Whether the tool is currently allowed given the session config.""" + + +class ListToolsResult(BaseModel): + """Result for droid.list_tools response.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + tools: list[ExecToolInfo] + """Available native CLI tools with their allow-state.""" + + +class CustomCommandInfo(BaseModel): + """A custom slash command entry returned by droid.list_commands.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + name: str + """Command name.""" + + description: str + """Command description.""" + + argument_hint: str | None = Field(default=None, alias="argumentHint") + """Optional argument hint.""" + + is_executable: bool | None = Field(default=None, alias="isExecutable") + """Whether the command is backed by an executable script.""" + + +class ListCommandsResult(BaseModel): + """Result for droid.list_commands response.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + commands: list[CustomCommandInfo] + """Available custom slash commands.""" + + +# ============================================================ +# Session lifecycle (close / compact / fork / rename) +# ============================================================ + + +class CloseSessionResult(BaseModel): + """Result for droid.close_session response (empty).""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + +class CompactSessionResult(BaseModel): + """Result for droid.compact_session response.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + new_session_id: str = Field(alias="newSessionId") + """Session ID created by the compaction.""" + + removed_count: int = Field(alias="removedCount") + """Number of messages removed by the compaction.""" + + +class ForkSessionResult(BaseModel): + """Result for droid.fork_session response.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + new_session_id: str = Field(alias="newSessionId") + """Session ID of the fork.""" + + +class RenameSessionResult(BaseModel): + """Result for droid.rename_session response.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + success: bool + """Whether the rename succeeded.""" + + +# ============================================================ +# Context stats / breakdown +# ============================================================ + + +class GetContextStatsResult(BaseModel): + """Result for droid.get_context_stats response.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + used: int + """Tokens used in the context window.""" + + remaining: int + """Tokens remaining in the context window.""" + + limit: int + """Total context window size.""" + + accuracy: str + """Accuracy of the estimate (server-driven string enum).""" + + updated_at: str = Field(alias="updatedAt") + """ISO 8601 timestamp of the last update.""" + + +class ContextBreakdownCategory(BaseModel): + """A top-level context usage category.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + name: str + """Category name.""" + + tokens: int + """Tokens attributed to the category.""" + + color_key: str = Field(alias="colorKey") + """UI color key for the category.""" + + +class ContextBreakdownSkillEntry(BaseModel): + """A per-skill context usage entry.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + name: str + """Skill name.""" + + location: str + """Skill location.""" + + tokens: int + """Tokens attributed to the skill.""" + + +class ContextBreakdownMcpServerEntry(BaseModel): + """A per-MCP-server context usage entry.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + name: str + """MCP server name.""" + + tool_count: int = Field(alias="toolCount") + """Number of tools contributed by the server.""" + + tokens: int + """Tokens attributed to the server.""" + + +class ContextBreakdownDroidEntry(BaseModel): + """A per-droid context usage entry.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + name: str + """Droid name.""" + + location: str + """Droid location.""" + + tokens: int + """Tokens attributed to the droid.""" + + +class GetContextBreakdownResult(BaseModel): + """Result for droid.get_context_breakdown response.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + model_id: str = Field(alias="modelId") + """Active model identifier.""" + + model_display_name: str = Field(alias="modelDisplayName") + """Human-readable model name.""" + + context_budget: int = Field(alias="contextBudget") + """Total context budget in tokens.""" + + last_call_compaction_tokens: int | None = Field( + default=None, alias="lastCallCompactionTokens" + ) + """Tokens saved by compaction on the last call, if any.""" + + used_tokens: int = Field(alias="usedTokens") + """Total tokens used.""" + + free_tokens: int = Field(alias="freeTokens") + """Total tokens free.""" + + categories: list[ContextBreakdownCategory] + """Top-level usage categories.""" + + skills: list[ContextBreakdownSkillEntry] + """Per-skill usage.""" + + mcp_servers: list[ContextBreakdownMcpServerEntry] = Field(alias="mcpServers") + """Per-MCP-server usage.""" + + droids: list[ContextBreakdownDroidEntry] + """Per-droid usage.""" + + +# ============================================================ +# Rewind (droid.get_rewind_info, droid.execute_rewind) +# ============================================================ + + +class GetRewindInfoResult(BaseModel): + """Result for droid.get_rewind_info response.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + available_files: list[RewindFileSnapshot] = Field(alias="availableFiles") + """Files that can be restored.""" + + created_files: list[RewindFileCreation] = Field(alias="createdFiles") + """Files created after the rewind point (candidates for deletion).""" + + evicted_files: list[RewindEvictedFile] = Field(alias="evictedFiles") + """Files that cannot be restored.""" + + +class ExecuteRewindResult(BaseModel): + """Result for droid.execute_rewind response.""" + + model_config = ConfigDict(populate_by_name=True, extra="allow") + + new_session_id: str = Field(alias="newSessionId") + """Session ID created by the rewind (a fork).""" + + restored_count: int = Field(alias="restoredCount") + """Number of files restored.""" + + deleted_count: int = Field(alias="deletedCount") + """Number of files deleted.""" + + failed_restore_count: int = Field(alias="failedRestoreCount") + """Number of files that failed to restore.""" + + failed_delete_count: int = Field(alias="failedDeleteCount") + """Number of files that failed to delete.""" + + # ============================================================ # Response schemas (union of success + failure) # ============================================================ @@ -1438,6 +2008,56 @@ class _SubmitBugReportResponseSuccess(JsonRpcResponseSuccess): result: SubmitBugReportResult # type: ignore[assignment] +class _ListToolsResponseSuccess(JsonRpcResponseSuccess): + model_config = ConfigDict(populate_by_name=True, extra="allow") + result: ListToolsResult # type: ignore[assignment] + + +class _ListCommandsResponseSuccess(JsonRpcResponseSuccess): + model_config = ConfigDict(populate_by_name=True, extra="allow") + result: ListCommandsResult # type: ignore[assignment] + + +class _CloseSessionResponseSuccess(JsonRpcResponseSuccess): + model_config = ConfigDict(populate_by_name=True, extra="allow") + result: CloseSessionResult # type: ignore[assignment] + + +class _CompactSessionResponseSuccess(JsonRpcResponseSuccess): + model_config = ConfigDict(populate_by_name=True, extra="allow") + result: CompactSessionResult # type: ignore[assignment] + + +class _ForkSessionResponseSuccess(JsonRpcResponseSuccess): + model_config = ConfigDict(populate_by_name=True, extra="allow") + result: ForkSessionResult # type: ignore[assignment] + + +class _RenameSessionResponseSuccess(JsonRpcResponseSuccess): + model_config = ConfigDict(populate_by_name=True, extra="allow") + result: RenameSessionResult # type: ignore[assignment] + + +class _GetContextStatsResponseSuccess(JsonRpcResponseSuccess): + model_config = ConfigDict(populate_by_name=True, extra="allow") + result: GetContextStatsResult # type: ignore[assignment] + + +class _GetContextBreakdownResponseSuccess(JsonRpcResponseSuccess): + model_config = ConfigDict(populate_by_name=True, extra="allow") + result: GetContextBreakdownResult # type: ignore[assignment] + + +class _GetRewindInfoResponseSuccess(JsonRpcResponseSuccess): + model_config = ConfigDict(populate_by_name=True, extra="allow") + result: GetRewindInfoResult # type: ignore[assignment] + + +class _ExecuteRewindResponseSuccess(JsonRpcResponseSuccess): + model_config = ConfigDict(populate_by_name=True, extra="allow") + result: ExecuteRewindResult # type: ignore[assignment] + + # Union response types (success | failure) InitializeSessionResponse = _InitializeSessionResponseSuccess | JsonRpcResponseFailure @@ -1463,10 +2083,22 @@ class _SubmitBugReportResponseSuccess(JsonRpcResponseSuccess): ToggleMcpToolResponse = _ToggleMcpToolResponseSuccess | JsonRpcResponseFailure ListSkillsResponse = _ListSkillsResponseSuccess | JsonRpcResponseFailure SubmitBugReportResponse = _SubmitBugReportResponseSuccess | JsonRpcResponseFailure +ListToolsResponse = _ListToolsResponseSuccess | JsonRpcResponseFailure +ListCommandsResponse = _ListCommandsResponseSuccess | JsonRpcResponseFailure +CloseSessionResponse = _CloseSessionResponseSuccess | JsonRpcResponseFailure +CompactSessionResponse = _CompactSessionResponseSuccess | JsonRpcResponseFailure +ForkSessionResponse = _ForkSessionResponseSuccess | JsonRpcResponseFailure +RenameSessionResponse = _RenameSessionResponseSuccess | JsonRpcResponseFailure +GetContextStatsResponse = _GetContextStatsResponseSuccess | JsonRpcResponseFailure +GetContextBreakdownResponse = ( + _GetContextBreakdownResponseSuccess | JsonRpcResponseFailure +) +GetRewindInfoResponse = _GetRewindInfoResponseSuccess | JsonRpcResponseFailure +ExecuteRewindResponse = _ExecuteRewindResponseSuccess | JsonRpcResponseFailure # ============================================================ -# ClientRequest discriminated union over all 19 request types +# ClientRequest discriminated union over all 29 request types # ============================================================ ClientRequestUnion = Annotated[ @@ -1488,13 +2120,23 @@ class _SubmitBugReportResponseSuccess(JsonRpcResponseSuccess): | ListMcpServersRequest | ToggleMcpToolRequest | ListSkillsRequest - | SubmitBugReportRequest, + | SubmitBugReportRequest + | ListToolsRequest + | ListCommandsRequest + | CloseSessionRequest + | CompactSessionRequest + | ForkSessionRequest + | RenameSessionRequest + | GetContextStatsRequest + | GetContextBreakdownRequest + | GetRewindInfoRequest + | ExecuteRewindRequest, Field(discriminator="method"), ] class ClientRequest(RootModel[ClientRequestUnion]): - """Discriminated union over all 19 client→server request types. + """Discriminated union over all 29 client→server request types. Dispatches on the ``method`` field to the appropriate request model. """ diff --git a/src/droid_sdk/schemas/enums.py b/src/droid_sdk/schemas/enums.py index 8d8917d..f3c9488 100644 --- a/src/droid_sdk/schemas/enums.py +++ b/src/droid_sdk/schemas/enums.py @@ -69,6 +69,16 @@ class DroidServerMethod(str, Enum): SUBMIT_MCP_AUTH_CODE = "droid.submit_mcp_auth_code" LIST_SKILLS = "droid.list_skills" SUBMIT_BUG_REPORT = "droid.submit_bug_report" + LIST_TOOLS = "droid.list_tools" + LIST_COMMANDS = "droid.list_commands" + CLOSE_SESSION = "droid.close_session" + COMPACT_SESSION = "droid.compact_session" + FORK_SESSION = "droid.fork_session" + RENAME_SESSION = "droid.rename_session" + GET_CONTEXT_STATS = "droid.get_context_stats" + GET_CONTEXT_BREAKDOWN = "droid.get_context_breakdown" + GET_REWIND_INFO = "droid.get_rewind_info" + EXECUTE_REWIND = "droid.execute_rewind" class DroidClientMethod(str, Enum): diff --git a/src/droid_sdk/stream.py b/src/droid_sdk/stream.py index 3ea7e8f..42fd04a 100644 --- a/src/droid_sdk/stream.py +++ b/src/droid_sdk/stream.py @@ -86,19 +86,35 @@ class ToolUse: @dataclass(frozen=True, slots=True) class ToolResult: - """The result returned from a tool execution.""" + """The result returned from a tool execution. + + ``tool_name`` is ``None`` when the name is unknown. The + ``tool_result`` notification does not carry the tool name, so + :meth:`~droid_sdk.client.DroidClient.receive_response` backfills it + from the matching ``tool_use`` (correlated via ``tool_use_id``) when + that call was seen earlier in the turn. It stays ``None`` when no + match is available, which is distinct from a tool that reported an + empty name. ``tool_use_id`` is optional for compatibility with + callers that construct stream events directly. + """ - tool_name: str + tool_name: str | None content: str | list[Any] is_error: bool + tool_use_id: str | None = None @dataclass(frozen=True, slots=True) class ToolProgress: - """A streaming progress update from a tool execution.""" + """A streaming progress update from a tool execution. + + ``tool_use_id`` is optional for compatibility with callers that + construct stream events directly. + """ tool_name: str content: str + tool_use_id: str | None = None @dataclass(frozen=True, slots=True) @@ -127,10 +143,20 @@ class TurnComplete: @dataclass(frozen=True, slots=True) class ErrorEvent: - """An error event from the droid process.""" + """An error event from the droid process. + + ``error_type`` mirrors the top-level ``errorType`` field, which is + frequently the unhelpful discriminator ``"Error"``. ``error_name`` + exposes the nested ``error.name`` the protocol carries (for example + ``"LLMInvalidRequestError"``), giving callers a stable value to + branch on. ``error_detail`` holds the full nested error payload when + present. + """ message: str error_type: str + error_name: str | None = field(default=None) + error_detail: dict[str, Any] | None = field(default=None) StreamMessage = ( @@ -187,8 +213,11 @@ def _notification_to_stream_message( content = list(notification.content) else: content = str(notification.content) + # The tool_result notification carries no tool name; leave it None + # so receive_response() can backfill it from the matching tool_use. return ToolResult( - tool_name="", + tool_name=None, + tool_use_id=notification.tool_use_id, content=content, is_error=bool(notification.is_error), ) @@ -199,6 +228,7 @@ def _notification_to_stream_message( text = update.text or update.status or update.details or "" return ToolProgress( tool_name=notification.tool_name, + tool_use_id=notification.tool_use_id, content=text, ) @@ -217,9 +247,16 @@ def _notification_to_stream_message( ) if isinstance(notification, ErrorNotification): + error_name: str | None = None + error_detail: dict[str, Any] | None = None + if notification.error is not None: + error_name = notification.error.name + error_detail = notification.error.model_dump(by_alias=True) return ErrorEvent( message=notification.message, error_type=notification.error_type.value, + error_name=error_name, + error_detail=error_detail, ) if isinstance(notification, CreateMessageNotification): diff --git a/tests/test_client_protocol_surface.py b/tests/test_client_protocol_surface.py new file mode 100644 index 0000000..af2173f --- /dev/null +++ b/tests/test_client_protocol_surface.py @@ -0,0 +1,569 @@ +"""Tests for the extended protocol surface added for issue #3. + +Covers the new typed ``DroidClient`` methods (tool/command discovery, +session lifecycle, context introspection, rewind) plus the new +``disabled_tool_ids`` and ``output_format`` request fields. +""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest +from pydantic import ValidationError + +from droid_sdk.client import DroidClient +from droid_sdk.errors import SessionError +from droid_sdk.protocol import COMPACTION_TIMEOUT +from droid_sdk.schemas.client import ( + OutputFormat, + RewindFileCreation, + RewindFileSnapshot, + SessionTag, +) +from droid_sdk.schemas.enums import DroidServerMethod +from tests.helpers import InMemoryTransport, make_success_response + +_background_tasks: set[asyncio.Task[Any]] = set() + + +def _fire(coro: Any) -> asyncio.Task[Any]: + task: asyncio.Task[Any] = asyncio.create_task(coro) + _background_tasks.add(task) + task.add_done_callback(_background_tasks.discard) + return task + + +async def _setup_client(transport: InMemoryTransport) -> DroidClient: + """Create a connected client with an active session.""" + client = DroidClient(transport=transport) + await client.connect() + + init_task = asyncio.create_task( + client.initialize_session(machine_id="test", cwd="/tmp") + ) + await asyncio.sleep(0) + sent = transport.get_last_sent_parsed() + transport.inject_message( + make_success_response( + sent["id"], + { + "sessionId": "sess-1", + "session": {"id": "sess-1"}, + "settings": {"modelId": "claude-sonnet-4", "reasoningEffort": "medium"}, + }, + ) + ) + await init_task + return client + + +async def _call( + transport: InMemoryTransport, + coro: Any, + result: dict[str, Any], +) -> tuple[dict[str, Any], Any]: + """Run *coro*, capture the sent request, inject *result*, return both.""" + task = _fire(coro) + await asyncio.sleep(0.01) + sent = transport.get_last_sent_parsed() + transport.inject_message(make_success_response(sent["id"], result)) + value = await task + return sent, value + + +# --------------------------------------------------------------------------- +# Tool / command discovery +# --------------------------------------------------------------------------- + + +class TestListTools: + @pytest.mark.asyncio + async def test_works_before_session_initialization(self) -> None: + transport = InMemoryTransport() + client = DroidClient(transport=transport) + await client.connect() + + sent, result = await _call( + transport, + client.list_tools(), + {"tools": []}, + ) + + assert sent["method"] == DroidServerMethod.LIST_TOOLS.value + assert result.tools == [] + + await client.close() + + @pytest.mark.asyncio + async def test_sends_method_and_parses_result(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + + sent, result = await _call( + transport, + client.list_tools(), + { + "tools": [ + { + "id": "read-cli", + "llmId": "Read", + "displayName": "Read", + "description": "Read a file", + "category": "filesystem", + "defaultAllowed": True, + "currentlyAllowed": False, + } + ] + }, + ) + + assert sent["method"] == DroidServerMethod.LIST_TOOLS.value + assert len(result.tools) == 1 + assert result.tools[0].id == "read-cli" + assert result.tools[0].default_allowed is True + assert result.tools[0].currently_allowed is False + + await client.close() + + @pytest.mark.asyncio + async def test_forwards_tool_id_filters(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + + sent, _ = await _call( + transport, + client.list_tools(enabled_tool_ids=[], disabled_tool_ids=["read-cli"]), + {"tools": []}, + ) + + assert sent["params"]["enabledToolIds"] == [] + assert sent["params"]["disabledToolIds"] == ["read-cli"] + + await client.close() + + +class TestListCommands: + @pytest.mark.asyncio + async def test_sends_method_and_parses_result(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + + sent, result = await _call( + transport, + client.list_commands(), + {"commands": [{"name": "deploy", "description": "Deploy the app"}]}, + ) + + assert sent["method"] == DroidServerMethod.LIST_COMMANDS.value + assert result.commands[0].name == "deploy" + + await client.close() + + +# --------------------------------------------------------------------------- +# Session lifecycle +# --------------------------------------------------------------------------- + + +class TestSessionLifecycle: + @pytest.mark.asyncio + async def test_close_session(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + + sent, _ = await _call(transport, client.close_session(reason="clear"), {}) + + assert sent["method"] == DroidServerMethod.CLOSE_SESSION.value + assert sent["params"]["reason"] == "clear" + assert client.session_id is None + + sent_count = len(transport.sent_messages) + with pytest.raises(SessionError): + await client.get_context_stats() + assert len(transport.sent_messages) == sent_count + + await client.close() + + @pytest.mark.asyncio + async def test_close_session_rejects_invalid_reason_locally(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + sent_count = len(transport.sent_messages) + + with pytest.raises(ValidationError): + await client.close_session(reason="invalid") # type: ignore[arg-type] + + assert client.session_id == "sess-1" + assert len(transport.sent_messages) == sent_count + + await client.close() + + @pytest.mark.asyncio + async def test_compact_session(self, monkeypatch: pytest.MonkeyPatch) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + protocol = client._protocol + assert protocol is not None + send_request = protocol.send_request + captured_timeout: float | None = None + + async def capture_timeout(**kwargs: Any) -> dict[str, Any]: + nonlocal captured_timeout + captured_timeout = kwargs.get("timeout") + return await send_request(**kwargs) + + monkeypatch.setattr(protocol, "send_request", capture_timeout) + + sent, result = await _call( + transport, + client.compact_session(custom_instructions="keep decisions"), + {"newSessionId": "sess-2", "removedCount": 12}, + ) + + assert sent["method"] == DroidServerMethod.COMPACT_SESSION.value + assert sent["params"]["customInstructions"] == "keep decisions" + assert result.new_session_id == "sess-2" + assert result.removed_count == 12 + assert captured_timeout == COMPACTION_TIMEOUT == 240.0 + + await client.close() + + @pytest.mark.asyncio + async def test_fork_session(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + + sent, result = await _call( + transport, + client.fork_session(title="experiment"), + {"newSessionId": "sess-3"}, + ) + + assert sent["method"] == DroidServerMethod.FORK_SESSION.value + assert sent["params"]["title"] == "experiment" + assert result.new_session_id == "sess-3" + + await client.close() + + @pytest.mark.asyncio + async def test_rename_session(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + + sent, result = await _call( + transport, + client.rename_session(title="New title"), + {"success": True}, + ) + + assert sent["method"] == DroidServerMethod.RENAME_SESSION.value + assert sent["params"]["title"] == "New title" + assert result.success is True + + await client.close() + + +# --------------------------------------------------------------------------- +# Context introspection +# --------------------------------------------------------------------------- + + +class TestContextIntrospection: + @pytest.mark.asyncio + async def test_get_context_stats(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + + sent, result = await _call( + transport, + client.get_context_stats(), + { + "used": 1000, + "remaining": 199000, + "limit": 200000, + "accuracy": "exact", + "updatedAt": "2026-08-04T00:00:00Z", + }, + ) + + assert sent["method"] == DroidServerMethod.GET_CONTEXT_STATS.value + assert result.used == 1000 + assert result.remaining == 199000 + assert result.limit == 200000 + + await client.close() + + @pytest.mark.asyncio + async def test_get_context_breakdown(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + + sent, result = await _call( + transport, + client.get_context_breakdown(), + { + "modelId": "claude-sonnet-4", + "modelDisplayName": "Claude Sonnet 4", + "contextBudget": 200000, + "usedTokens": 1000, + "freeTokens": 199000, + "categories": [{"name": "System", "tokens": 500, "colorKey": "blue"}], + "skills": [], + "mcpServers": [], + "droids": [], + }, + ) + + assert sent["method"] == DroidServerMethod.GET_CONTEXT_BREAKDOWN.value + assert result.model_id == "claude-sonnet-4" + assert result.categories[0].name == "System" + assert result.categories[0].color_key == "blue" + + await client.close() + + +# --------------------------------------------------------------------------- +# Rewind +# --------------------------------------------------------------------------- + + +class TestRewind: + @pytest.mark.asyncio + async def test_get_rewind_info(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + + sent, result = await _call( + transport, + client.get_rewind_info(message_id="msg-7"), + { + "availableFiles": [ + {"filePath": "a.py", "contentHash": "abc", "size": 10} + ], + "createdFiles": [{"filePath": "b.py"}], + "evictedFiles": [{"filePath": "c.py", "reason": "too large"}], + }, + ) + + assert sent["method"] == DroidServerMethod.GET_REWIND_INFO.value + assert sent["params"]["sessionId"] == "sess-1" + assert sent["params"]["messageId"] == "msg-7" + assert result.available_files[0].file_path == "a.py" + assert result.created_files[0].file_path == "b.py" + assert result.evicted_files[0].reason == "too large" + + await client.close() + + @pytest.mark.asyncio + async def test_execute_rewind(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + + sent, result = await _call( + transport, + client.execute_rewind( + message_id="msg-7", + files_to_restore=[ + {"filePath": "a.py", "contentHash": "abc", "size": 10} + ], + files_to_delete=[{"filePath": "b.py"}], + fork_title="rewound", + ), + { + "newSessionId": "sess-9", + "restoredCount": 1, + "deletedCount": 1, + "failedRestoreCount": 0, + "failedDeleteCount": 0, + }, + ) + + assert sent["method"] == DroidServerMethod.EXECUTE_REWIND.value + assert sent["params"]["messageId"] == "msg-7" + assert sent["params"]["forkTitle"] == "rewound" + assert result.new_session_id == "sess-9" + assert result.restored_count == 1 + + await client.close() + + @pytest.mark.asyncio + async def test_execute_rewind_serializes_typed_files(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + + sent, _ = await _call( + transport, + client.execute_rewind( + message_id="msg-7", + files_to_restore=[ + RewindFileSnapshot( + filePath="a.py", + contentHash="abc", + size=10, + ) + ], + files_to_delete=[RewindFileCreation(filePath="b.py")], + fork_title="rewound", + ), + { + "newSessionId": "sess-9", + "restoredCount": 1, + "deletedCount": 1, + "failedRestoreCount": 0, + "failedDeleteCount": 0, + }, + ) + + assert sent["params"]["filesToRestore"] == [ + {"filePath": "a.py", "contentHash": "abc", "size": 10} + ] + assert sent["params"]["filesToDelete"] == [{"filePath": "b.py"}] + + await client.close() + + @pytest.mark.asyncio + async def test_execute_rewind_rejects_malformed_files_locally(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + sent_count = len(transport.sent_messages) + + with pytest.raises(ValidationError): + await client.execute_rewind( + message_id="msg-7", + files_to_restore=[{"filePath": "a.py", "size": 10}], + files_to_delete=[], + fork_title="rewound", + ) + + assert len(transport.sent_messages) == sent_count + + await client.close() + + +# --------------------------------------------------------------------------- +# New request fields (disabled_tool_ids, output_format) +# --------------------------------------------------------------------------- + + +class TestNewRequestFields: + @pytest.mark.asyncio + async def test_initialize_session_sends_disabled_tool_ids(self) -> None: + transport = InMemoryTransport() + client = DroidClient(transport=transport) + await client.connect() + + sent, _ = await _call( + transport, + client.initialize_session( + machine_id="m", + cwd="/tmp", + enabled_tool_ids=[], + disabled_tool_ids=["read-cli", "execute-cli"], + ), + { + "sessionId": "sess-1", + "session": {"id": "sess-1"}, + "settings": {"modelId": "claude-sonnet-4", "reasoningEffort": "medium"}, + }, + ) + + assert sent["params"]["disabledToolIds"] == ["read-cli", "execute-cli"] + assert sent["params"]["enabledToolIds"] == [] + + await client.close() + + @pytest.mark.asyncio + async def test_update_session_settings_sends_disabled_tool_ids(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + + sent, _ = await _call( + transport, + client.update_session_settings(disabled_tool_ids=["read-cli"]), + {}, + ) + + assert sent["params"]["disabledToolIds"] == ["read-cli"] + + await client.close() + + @pytest.mark.asyncio + async def test_add_user_message_output_format_dict(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + + schema = {"type": "object", "properties": {"answer": {"type": "integer"}}} + sent, _ = await _call( + transport, + client.add_user_message( + text="Return an answer.", + output_format={"type": "json_schema", "schema": schema}, + ), + {}, + ) + + assert sent["params"]["outputFormat"]["type"] == "json_schema" + assert sent["params"]["outputFormat"]["schema"] == schema + + await client.close() + + @pytest.mark.asyncio + async def test_add_user_message_output_format_model(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + + schema = {"type": "object"} + sent, _ = await _call( + transport, + client.add_user_message( + text="Return an answer.", + output_format=OutputFormat(type="json_schema", schema=schema), + ), + {}, + ) + + assert sent["params"]["outputFormat"] == { + "type": "json_schema", + "schema": schema, + } + + await client.close() + + @pytest.mark.asyncio + async def test_add_user_message_rejects_invalid_output_format_locally( + self, + ) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + sent_count = len(transport.sent_messages) + + with pytest.raises(ValidationError): + await client.add_user_message( + text="Return an answer.", + output_format={"type": "text", "schema": {}}, + ) + + assert len(transport.sent_messages) == sent_count + + await client.close() + + @pytest.mark.asyncio + async def test_fork_session_serializes_typed_tags(self) -> None: + transport = InMemoryTransport() + client = await _setup_client(transport) + + sent, _ = await _call( + transport, + client.fork_session( + tags=[SessionTag(name="live-test", metadata={"source": "sdk"})] + ), + {"newSessionId": "sess-3"}, + ) + + assert sent["params"]["tags"] == [ + {"name": "live-test", "metadata": {"source": "sdk"}} + ] + + await client.close() diff --git a/tests/test_enums.py b/tests/test_enums.py index 59c4771..6eb5e27 100644 --- a/tests/test_enums.py +++ b/tests/test_enums.py @@ -77,7 +77,7 @@ def test_factory_client_version(self) -> None: ENUM_MEMBER_COUNTS: list[tuple[type[Enum], int]] = [ - (DroidServerMethod, 19), + (DroidServerMethod, 29), (DroidClientMethod, 3), (SessionNotificationType, 20), (ToolConfirmationOutcome, 8), @@ -217,6 +217,16 @@ def test_json_rpc_error_code_int_mixin() -> None: (DroidServerMethod, "SUBMIT_MCP_AUTH_CODE", "droid.submit_mcp_auth_code"), (DroidServerMethod, "LIST_SKILLS", "droid.list_skills"), (DroidServerMethod, "SUBMIT_BUG_REPORT", "droid.submit_bug_report"), + (DroidServerMethod, "LIST_TOOLS", "droid.list_tools"), + (DroidServerMethod, "LIST_COMMANDS", "droid.list_commands"), + (DroidServerMethod, "CLOSE_SESSION", "droid.close_session"), + (DroidServerMethod, "COMPACT_SESSION", "droid.compact_session"), + (DroidServerMethod, "FORK_SESSION", "droid.fork_session"), + (DroidServerMethod, "RENAME_SESSION", "droid.rename_session"), + (DroidServerMethod, "GET_CONTEXT_STATS", "droid.get_context_stats"), + (DroidServerMethod, "GET_CONTEXT_BREAKDOWN", "droid.get_context_breakdown"), + (DroidServerMethod, "GET_REWIND_INFO", "droid.get_rewind_info"), + (DroidServerMethod, "EXECUTE_REWIND", "droid.execute_rewind"), # DroidClientMethod (DroidClientMethod, "SESSION_NOTIFICATION", "droid.session_notification"), (DroidClientMethod, "REQUEST_PERMISSION", "droid.request_permission"), diff --git a/tests/test_extended_client_schemas.py b/tests/test_extended_client_schemas.py new file mode 100644 index 0000000..05ac4b2 --- /dev/null +++ b/tests/test_extended_client_schemas.py @@ -0,0 +1,139 @@ +"""Schema coverage for the extended client protocol surface.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from pydantic import TypeAdapter + +import droid_sdk.schemas as schemas + +_REQUEST_ENVELOPE = { + "jsonrpc": "2.0", + "factoryApiVersion": "1.0.0", + "type": "request", + "id": "req-extended", +} + +_RESPONSE_ENVELOPE = { + "jsonrpc": "2.0", + "factoryApiVersion": "1.0.0", + "type": "response", + "id": "req-extended", +} + +_REQUEST_CASES: list[tuple[str, dict[str, Any], str]] = [ + ( + "droid.list_tools", + {"enabledToolIds": [], "disabledToolIds": ["read-cli"]}, + "ListToolsRequest", + ), + ("droid.list_commands", {}, "ListCommandsRequest"), + ("droid.close_session", {"reason": "other"}, "CloseSessionRequest"), + ( + "droid.compact_session", + {"customInstructions": "Preserve decisions"}, + "CompactSessionRequest", + ), + ( + "droid.fork_session", + {"title": "fork", "tags": [{"name": "test"}]}, + "ForkSessionRequest", + ), + ("droid.rename_session", {"title": "renamed"}, "RenameSessionRequest"), + ("droid.get_context_stats", {}, "GetContextStatsRequest"), + ("droid.get_context_breakdown", {}, "GetContextBreakdownRequest"), + ( + "droid.get_rewind_info", + {"sessionId": "sess-1", "messageId": "msg-1"}, + "GetRewindInfoRequest", + ), + ( + "droid.execute_rewind", + { + "sessionId": "sess-1", + "messageId": "msg-1", + "filesToRestore": [{"filePath": "a.py", "contentHash": "abc", "size": 10}], + "filesToDelete": [{"filePath": "b.py"}], + "forkTitle": "rewound", + }, + "ExecuteRewindRequest", + ), +] + +_RESPONSE_CASES: list[tuple[str, dict[str, Any]]] = [ + ("ListToolsResponse", {"tools": []}), + ("ListCommandsResponse", {"commands": []}), + ("CloseSessionResponse", {}), + ( + "CompactSessionResponse", + {"newSessionId": "sess-2", "removedCount": 3}, + ), + ("ForkSessionResponse", {"newSessionId": "sess-3"}), + ("RenameSessionResponse", {"success": True}), + ( + "GetContextStatsResponse", + { + "used": 1, + "remaining": 9, + "limit": 10, + "accuracy": "exact", + "updatedAt": "2026-08-04T00:00:00Z", + }, + ), + ( + "GetContextBreakdownResponse", + { + "modelId": "model", + "modelDisplayName": "Model", + "contextBudget": 10, + "usedTokens": 1, + "freeTokens": 9, + "categories": [], + "skills": [], + "mcpServers": [], + "droids": [], + }, + ), + ( + "GetRewindInfoResponse", + {"availableFiles": [], "createdFiles": [], "evictedFiles": []}, + ), + ( + "ExecuteRewindResponse", + { + "newSessionId": "sess-4", + "restoredCount": 0, + "deletedCount": 0, + "failedRestoreCount": 0, + "failedDeleteCount": 0, + }, + ), +] + + +@pytest.mark.parametrize(("method", "params", "expected_type"), _REQUEST_CASES) +def test_client_request_union_includes_extended_methods( + method: str, + params: dict[str, Any], + expected_type: str, +) -> None: + request = schemas.ClientRequest.model_validate( + {**_REQUEST_ENVELOPE, "method": method, "params": params} + ) + + assert type(request.root).__name__ == expected_type + + +@pytest.mark.parametrize(("response_name", "result"), _RESPONSE_CASES) +def test_extended_response_types_parse_success( + response_name: str, + result: dict[str, Any], +) -> None: + response_type = getattr(schemas, response_name) + response = TypeAdapter(response_type).validate_python( + {**_RESPONSE_ENVELOPE, "result": result} + ) + + assert response.result is not None diff --git a/tests/test_live_droid_exec.py b/tests/test_live_droid_exec.py new file mode 100644 index 0000000..6083791 --- /dev/null +++ b/tests/test_live_droid_exec.py @@ -0,0 +1,268 @@ +"""Opt-in integration tests against the installed ``droid exec`` CLI. + +These tests create real sessions and consume model usage. They are excluded +from normal test runs; set ``DROID_LIVE_TESTS=1`` to enable them. +""" + +from __future__ import annotations + +import asyncio +import os +import re +from contextlib import asynccontextmanager +from typing import TYPE_CHECKING, Any + +import pytest + +from droid_sdk import ( + AssistantTextDelta, + DroidClient, + ProcessTransport, + ToolConfirmationOutcome, + TurnComplete, +) + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + from pathlib import Path + +pytestmark = pytest.mark.skipif( + os.environ.get("DROID_LIVE_TESTS") != "1", + reason="set DROID_LIVE_TESTS=1 to run tests against droid exec", +) + +_EXEC_PATH = os.environ.get("DROID_EXEC_PATH", "droid") +_TURN_TIMEOUT = 180.0 +_RPC_TIMEOUT = 180.0 + + +@asynccontextmanager +async def _client(cwd: Path) -> AsyncIterator[DroidClient]: + transport = ProcessTransport(exec_path=_EXEC_PATH, cwd=str(cwd)) + async with DroidClient(transport=transport) as client: + client.set_permission_handler( + lambda _: ToolConfirmationOutcome.ProceedOnce.value + ) + yield client + + +async def _initialize(client: DroidClient, cwd: Path) -> str: + result = await asyncio.wait_for( + client.initialize_session( + machine_id="droid-sdk-python-live-tests", + cwd=str(cwd), + ), + timeout=_RPC_TIMEOUT, + ) + return result.session_id + + +async def _run_turn(client: DroidClient, text: str) -> str: + chunks: list[str] = [] + + async def consume() -> None: + async for message in client.receive_response(): + if isinstance(message, AssistantTextDelta): + chunks.append(message.text) + elif isinstance(message, TurnComplete): + return + + consumer = asyncio.create_task(consume()) + await asyncio.sleep(0) + await client.add_user_message(text=text) + await asyncio.wait_for(consumer, timeout=_TURN_TIMEOUT) + return "".join(chunks) + + +def _inner_notification(raw: dict[str, Any]) -> dict[str, Any] | None: + params = raw.get("params") + if not isinstance(params, dict): + return None + notification = params.get("notification") + return notification if isinstance(notification, dict) else None + + +@pytest.mark.asyncio +async def test_live_load_session(tmp_path: Path) -> None: + async with _client(tmp_path) as client: + session_id = await _initialize(client, tmp_path) + + async with _client(tmp_path) as client: + result = await asyncio.wait_for( + client.load_session(session_id=session_id), + timeout=_RPC_TIMEOUT, + ) + + assert result.session + assert client.session_id == session_id + + +@pytest.mark.asyncio +async def test_live_interrupt_session(tmp_path: Path) -> None: + async with _client(tmp_path) as client: + await _initialize(client, tmp_path) + became_busy = asyncio.Event() + + def capture_state(raw: dict[str, Any]) -> None: + notification = _inner_notification(raw) + if ( + notification is not None + and notification.get("type") == "droid_working_state_changed" + and notification.get("newState") != "idle" + ): + became_busy.set() + + client.on_notification(capture_state) + + consumer = asyncio.create_task(_run_turn(client, "Count slowly to 10000.")) + await asyncio.wait_for(became_busy.wait(), timeout=30) + await asyncio.wait_for(client.interrupt_session(), timeout=_RPC_TIMEOUT) + await asyncio.wait_for(consumer, timeout=_TURN_TIMEOUT) + + +@pytest.mark.asyncio +async def test_live_update_session_tool_settings(tmp_path: Path) -> None: + async with _client(tmp_path) as client: + tools = await asyncio.wait_for(client.list_tools(), timeout=_RPC_TIMEOUT) + assert tools.tools + tool_id = tools.tools[0].id + + await _initialize(client, tmp_path) + await asyncio.wait_for( + client.update_session_settings( + enabled_tool_ids=[], + disabled_tool_ids=[tool_id], + ), + timeout=_RPC_TIMEOUT, + ) + + +@pytest.mark.asyncio +async def test_live_close_session(tmp_path: Path) -> None: + async with _client(tmp_path) as client: + await _initialize(client, tmp_path) + result = await asyncio.wait_for( + client.close_session(reason="other"), + timeout=_RPC_TIMEOUT, + ) + + assert result is not None + + +@pytest.mark.asyncio +async def test_live_compact_session(tmp_path: Path) -> None: + async with _client(tmp_path) as client: + await _initialize(client, tmp_path) + await _run_turn(client, "Reply exactly COMPACTION_READY.") + + result = await asyncio.wait_for( + client.compact_session( + custom_instructions="Preserve the COMPACTION_READY fact." + ), + timeout=_RPC_TIMEOUT, + ) + + assert result.new_session_id + assert result.removed_count >= 0 + + +@pytest.mark.asyncio +async def test_live_fork_session(tmp_path: Path) -> None: + async with _client(tmp_path) as client: + session_id = await _initialize(client, tmp_path) + result = await asyncio.wait_for( + client.fork_session( + title="droid-sdk-python live fork", + tags=[{"name": "live-test", "metadata": {"source": "sdk"}}], + ), + timeout=_RPC_TIMEOUT, + ) + + assert result.new_session_id + assert result.new_session_id != session_id + + +@pytest.mark.asyncio +async def test_live_rewind_info_and_execute(tmp_path: Path) -> None: + async with _client(tmp_path) as client: + await _initialize(client, tmp_path) + user_message_ids: list[str] = [] + + def capture_user_message(raw: dict[str, Any]) -> None: + notification = _inner_notification(raw) + if notification is None or notification.get("type") != "create_message": + return + message = notification.get("message") + if not isinstance(message, dict) or message.get("role") != "user": + return + content = message.get("content") + if isinstance(content, list) and content: + message_id = message.get("id") + if isinstance(message_id, str): + user_message_ids.append(message_id) + + client.on_notification(capture_user_message) + await _run_turn(client, "Reply exactly REWIND_READY.") + assert user_message_ids + message_id = user_message_ids[-1] + + info = await asyncio.wait_for( + client.get_rewind_info(message_id=message_id), + timeout=_RPC_TIMEOUT, + ) + assert isinstance(info.available_files, list) + assert isinstance(info.created_files, list) + assert isinstance(info.evicted_files, list) + + result = await asyncio.wait_for( + client.execute_rewind( + message_id=message_id, + files_to_restore=[], + files_to_delete=[], + fork_title="droid-sdk-python live rewind", + ), + timeout=_RPC_TIMEOUT, + ) + assert result.new_session_id + assert result.restored_count == 0 + assert result.deleted_count == 0 + + +@pytest.mark.asyncio +async def test_live_kill_worker_session(tmp_path: Path) -> None: + async with _client(tmp_path) as client: + await _initialize(client, tmp_path) + worker_session_ids: list[str] = [] + + def capture_worker(raw: dict[str, Any]) -> None: + notification = _inner_notification(raw) + if notification is None or notification.get("type") != "tool_result": + return + content = notification.get("content") + if ( + not isinstance(content, str) + or "Task launched in background" not in content + ): + return + match = re.search(r"session_id: ([0-9a-f-]+)", content) + if match is not None: + worker_session_ids.append(match.group(1)) + + client.on_notification(capture_worker) + response = await _run_turn( + client, + "Use the Task tool exactly once with run_in_background=true. " + "Ask the worker to run `python -c 'import time; time.sleep(60)'` " + "and then finish. After launching it, reply exactly WORKER_STARTED " + "without waiting for it.", + ) + + assert worker_session_ids + worker_session_id = worker_session_ids[-1] + try: + assert "WORKER_STARTED" in response + finally: + await asyncio.wait_for( + client.kill_worker_session(worker_session_id=worker_session_id), + timeout=_RPC_TIMEOUT, + ) diff --git a/tests/test_receive_response.py b/tests/test_receive_response.py index aa55ea6..631e0bb 100644 --- a/tests/test_receive_response.py +++ b/tests/test_receive_response.py @@ -1153,3 +1153,244 @@ async def inject() -> None: assert turn_complete.token_usage.output_tokens == 50 await client.close() + + +# --------------------------------------------------------------------------- +# Tool-name backfill / correlation (issue #4) +# --------------------------------------------------------------------------- + + +class TestToolNameBackfill: + """receive_response() correlates tool_result/progress with tool_use.""" + + @pytest.mark.asyncio + async def test_tool_result_backfills_name_from_tool_use(self) -> None: + """A tool_result with no name is enriched from the matching tool_use.""" + transport = InMemoryTransport() + client = await _setup_client(transport) + + async def inject() -> None: + await asyncio.sleep(0) + transport.inject_message( + _make_session_notification( + SessionNotificationType.DROID_WORKING_STATE_CHANGED, + {"newState": "streaming_assistant_message"}, + ) + ) + await asyncio.sleep(0) + transport.inject_message( + _make_session_notification( + SessionNotificationType.CREATE_MESSAGE, + { + "message": { + "id": "msg1", + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "tu_abc", + "name": "ToolSearch", + "input": {"query": "weather"}, + }, + ], + "createdAt": 1700000000.0, + "updatedAt": 1700000000.0, + }, + }, + ) + ) + await asyncio.sleep(0) + transport.inject_message( + _make_session_notification( + SessionNotificationType.TOOL_RESULT, + { + "messageId": "msg2", + "toolUseId": "tu_abc", + "content": "Error: No tools matched", + "isError": True, + }, + ) + ) + await asyncio.sleep(0) + transport.inject_message( + _make_session_notification( + SessionNotificationType.DROID_WORKING_STATE_CHANGED, + {"newState": "idle"}, + ) + ) + + _fire_task(inject()) + + messages: list[StreamMessage] = [] + async for msg in client.receive_response(): + messages.append(msg) + + results = [m for m in messages if isinstance(m, ToolResult)] + assert len(results) == 1 + assert results[0].tool_use_id == "tu_abc" + assert results[0].tool_name == "ToolSearch" + assert results[0].is_error is True + + await client.close() + + @pytest.mark.asyncio + async def test_tool_result_name_none_when_no_tool_use(self) -> None: + """Without a preceding tool_use, tool_name stays None (not '').""" + transport = InMemoryTransport() + client = await _setup_client(transport) + + async def inject() -> None: + await asyncio.sleep(0) + transport.inject_message( + _make_session_notification( + SessionNotificationType.DROID_WORKING_STATE_CHANGED, + {"newState": "executing_tool"}, + ) + ) + await asyncio.sleep(0) + transport.inject_message( + _make_session_notification( + SessionNotificationType.TOOL_RESULT, + { + "messageId": "msg2", + "toolUseId": "tu_orphan", + "content": "done", + "isError": False, + }, + ) + ) + await asyncio.sleep(0) + transport.inject_message( + _make_session_notification( + SessionNotificationType.DROID_WORKING_STATE_CHANGED, + {"newState": "idle"}, + ) + ) + + _fire_task(inject()) + + messages: list[StreamMessage] = [] + async for msg in client.receive_response(): + messages.append(msg) + + results = [m for m in messages if isinstance(m, ToolResult)] + assert len(results) == 1 + assert results[0].tool_use_id == "tu_orphan" + assert results[0].tool_name is None + + await client.close() + + @pytest.mark.asyncio + async def test_tool_progress_carries_tool_use_id(self) -> None: + """ToolProgress exposes the correlating tool_use_id.""" + transport = InMemoryTransport() + client = await _setup_client(transport) + + async def inject() -> None: + await asyncio.sleep(0) + transport.inject_message( + _make_session_notification( + SessionNotificationType.DROID_WORKING_STATE_CHANGED, + {"newState": "executing_tool"}, + ) + ) + await asyncio.sleep(0) + transport.inject_message( + _make_session_notification( + SessionNotificationType.TOOL_PROGRESS_UPDATE, + { + "toolUseId": "tu_9", + "toolName": "execute", + "update": {"type": "status", "text": "Running..."}, + }, + ) + ) + await asyncio.sleep(0) + transport.inject_message( + _make_session_notification( + SessionNotificationType.DROID_WORKING_STATE_CHANGED, + {"newState": "idle"}, + ) + ) + + _fire_task(inject()) + + messages: list[StreamMessage] = [] + async for msg in client.receive_response(): + messages.append(msg) + + progress = [m for m in messages if isinstance(m, ToolProgress)] + assert len(progress) == 1 + assert progress[0].tool_use_id == "tu_9" + assert progress[0].tool_name == "execute" + + await client.close() + + @pytest.mark.asyncio + async def test_tool_progress_backfills_empty_name(self) -> None: + """An empty toolName on progress is backfilled from the tool_use.""" + transport = InMemoryTransport() + client = await _setup_client(transport) + + async def inject() -> None: + await asyncio.sleep(0) + transport.inject_message( + _make_session_notification( + SessionNotificationType.DROID_WORKING_STATE_CHANGED, + {"newState": "streaming_assistant_message"}, + ) + ) + await asyncio.sleep(0) + transport.inject_message( + _make_session_notification( + SessionNotificationType.CREATE_MESSAGE, + { + "message": { + "id": "msg1", + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "tu_p", + "name": "ToolSearch", + "input": {}, + }, + ], + "createdAt": 1700000000.0, + "updatedAt": 1700000000.0, + }, + }, + ) + ) + await asyncio.sleep(0) + transport.inject_message( + _make_session_notification( + SessionNotificationType.TOOL_PROGRESS_UPDATE, + { + "toolUseId": "tu_p", + "toolName": "", + "update": {"type": "status", "text": "working"}, + }, + ) + ) + await asyncio.sleep(0) + transport.inject_message( + _make_session_notification( + SessionNotificationType.DROID_WORKING_STATE_CHANGED, + {"newState": "idle"}, + ) + ) + + _fire_task(inject()) + + messages: list[StreamMessage] = [] + async for msg in client.receive_response(): + messages.append(msg) + + progress = [m for m in messages if isinstance(m, ToolProgress)] + assert len(progress) == 1 + assert progress[0].tool_use_id == "tu_p" + # Empty name on the wire is backfilled from the matching tool_use. + assert progress[0].tool_name == "ToolSearch" + + await client.close() diff --git a/tests/test_stream.py b/tests/test_stream.py index e5a7965..a1f55be 100644 --- a/tests/test_stream.py +++ b/tests/test_stream.py @@ -74,11 +74,27 @@ def test_field_names(self) -> None: class TestToolResult: """Tests for ToolResult dataclass.""" + def test_legacy_constructor_remains_compatible(self) -> None: + from droid_sdk.stream import ToolResult + + msg = ToolResult("read_file", "file contents", False) + + assert msg.tool_name == "read_file" + assert msg.content == "file contents" + assert msg.is_error is False + assert msg.tool_use_id is None + def test_construction_string_content(self) -> None: from droid_sdk.stream import ToolResult - msg = ToolResult(tool_name="read_file", content="file contents", is_error=False) + msg = ToolResult( + tool_name="read_file", + tool_use_id="tu_1", + content="file contents", + is_error=False, + ) assert msg.tool_name == "read_file" + assert msg.tool_use_id == "tu_1" assert msg.content == "file contents" assert msg.is_error is False @@ -86,32 +102,62 @@ def test_construction_list_content(self) -> None: from droid_sdk.stream import ToolResult content: list[Any] = [{"type": "text", "text": "result"}] - msg = ToolResult(tool_name="execute", content=content, is_error=True) + msg = ToolResult( + tool_name="execute", + tool_use_id="tu_2", + content=content, + is_error=True, + ) assert msg.content == content assert msg.is_error is True + def test_tool_name_optional(self) -> None: + from droid_sdk.stream import ToolResult + + msg = ToolResult( + tool_name=None, + tool_use_id="tu_3", + content="x", + is_error=False, + ) + assert msg.tool_name is None + def test_field_names(self) -> None: from droid_sdk.stream import ToolResult names = {f.name for f in fields(ToolResult)} - assert names == {"tool_name", "content", "is_error"} + assert names == {"tool_name", "tool_use_id", "content", "is_error"} class TestToolProgress: """Tests for ToolProgress dataclass.""" + def test_legacy_constructor_remains_compatible(self) -> None: + from droid_sdk.stream import ToolProgress + + msg = ToolProgress("execute", "running step 2...") + + assert msg.tool_name == "execute" + assert msg.content == "running step 2..." + assert msg.tool_use_id is None + def test_construction(self) -> None: from droid_sdk.stream import ToolProgress - msg = ToolProgress(tool_name="execute", content="running step 2...") + msg = ToolProgress( + tool_name="execute", + tool_use_id="tu_1", + content="running step 2...", + ) assert msg.tool_name == "execute" + assert msg.tool_use_id == "tu_1" assert msg.content == "running step 2..." def test_field_names(self) -> None: from droid_sdk.stream import ToolProgress names = {f.name for f in fields(ToolProgress)} - assert names == {"tool_name", "content"} + assert names == {"tool_name", "tool_use_id", "content"} class TestWorkingStateChanged: @@ -204,12 +250,29 @@ def test_construction(self) -> None: msg = ErrorEvent(message="Something went wrong", error_type="ConnectionError") assert msg.message == "Something went wrong" assert msg.error_type == "ConnectionError" + assert msg.error_name is None + assert msg.error_detail is None + + def test_construction_with_nested_error(self) -> None: + from droid_sdk.stream import ErrorEvent + + msg = ErrorEvent( + message="Requested model was not found on the API provider", + error_type="Error", + error_name="LLMInvalidRequestError", + error_detail={"name": "LLMInvalidRequestError", "message": "..."}, + ) + assert msg.error_name == "LLMInvalidRequestError" + assert msg.error_detail == { + "name": "LLMInvalidRequestError", + "message": "...", + } def test_field_names(self) -> None: from droid_sdk.stream import ErrorEvent names = {f.name for f in fields(ErrorEvent)} - assert names == {"message", "error_type"} + assert names == {"message", "error_type", "error_name", "error_detail"} # --------------------------------------------------------------------------- @@ -346,7 +409,7 @@ def test_tool_result(self) -> None: assert result.is_error is False def test_tool_result_with_missing_tool_name(self) -> None: - """ToolResultNotification doesn't have toolName; we should get empty string.""" + """ToolResultNotification has no toolName; converter yields None.""" from droid_sdk.schemas.cli import SessionNotification from droid_sdk.stream import ( ToolResult, @@ -365,8 +428,10 @@ def test_tool_result_with_missing_tool_name(self) -> None: notif = SessionNotification.model_validate(raw) result = _notification_to_stream_message(notif.params.notification) assert isinstance(result, ToolResult) - # ToolResultNotification doesn't carry tool_name, so it should be "" - assert result.tool_name == "" + # ToolResultNotification doesn't carry tool_name; None means "unknown" + # (distinct from a tool that reported an empty name). + assert result.tool_name is None + assert result.tool_use_id == "tu_1" assert result.is_error is True def test_tool_progress_update(self) -> None: @@ -391,6 +456,7 @@ def test_tool_progress_update(self) -> None: result = _notification_to_stream_message(notif.params.notification) assert isinstance(result, ToolProgress) assert result.tool_name == "execute" + assert result.tool_use_id == "tu_2" assert result.content == "Compiling..." def test_tool_progress_update_fallback_content(self) -> None: @@ -482,6 +548,36 @@ def test_error_notification(self) -> None: assert isinstance(result, ErrorEvent) assert result.message == "Something went wrong" assert result.error_type == "ConnectionError" + assert result.error_name is None + assert result.error_detail is None + + def test_error_notification_with_nested_error(self) -> None: + """The nested error.name is surfaced as error_name plus raw detail.""" + from droid_sdk.schemas.cli import SessionNotification + from droid_sdk.stream import ( + ErrorEvent, + _notification_to_stream_message, + ) + + raw = _make_session_notification( + SessionNotificationType.ERROR, + { + "message": "Requested model was not found on the API provider", + "errorType": "Error", + "timestamp": "2026-08-02T11:00:13.647Z", + "error": { + "name": "LLMInvalidRequestError", + "message": "Requested model was not found on the API provider", + }, + }, + ) + notif = SessionNotification.model_validate(raw) + result = _notification_to_stream_message(notif.params.notification) + assert isinstance(result, ErrorEvent) + assert result.error_type == "Error" + assert result.error_name == "LLMInvalidRequestError" + assert result.error_detail is not None + assert result.error_detail["name"] == "LLMInvalidRequestError" def test_create_message_with_tool_use_blocks(self) -> None: from droid_sdk.schemas.cli import SessionNotification diff --git a/uv.lock b/uv.lock index 0529c02..2bf360f 100644 --- a/uv.lock +++ b/uv.lock @@ -149,7 +149,7 @@ toml = [ [[package]] name = "droid-sdk" -version = "0.1.0" +version = "0.1.3" source = { editable = "." } dependencies = [ { name = "pydantic" },