bma-executor-sdk 0.4.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- bma_executor_sdk-0.4.0/.gitignore +4 -0
- bma_executor_sdk-0.4.0/PKG-INFO +28 -0
- bma_executor_sdk-0.4.0/README.md +17 -0
- bma_executor_sdk-0.4.0/examples/agent_minimal.py +79 -0
- bma_executor_sdk-0.4.0/examples/executor_minimal.py +64 -0
- bma_executor_sdk-0.4.0/examples/harness_loopback.py +50 -0
- bma_executor_sdk-0.4.0/pyproject.toml +26 -0
- bma_executor_sdk-0.4.0/src/bma_executor_sdk/__init__.py +77 -0
- bma_executor_sdk-0.4.0/src/bma_executor_sdk/_rpc.py +192 -0
- bma_executor_sdk-0.4.0/src/bma_executor_sdk/agent.py +558 -0
- bma_executor_sdk-0.4.0/src/bma_executor_sdk/client.py +675 -0
- bma_executor_sdk-0.4.0/src/bma_executor_sdk/generated.py +1155 -0
- bma_executor_sdk-0.4.0/src/bma_executor_sdk/harness.py +769 -0
- bma_executor_sdk-0.4.0/src/bma_executor_sdk/transport.py +189 -0
- bma_executor_sdk-0.4.0/tests/executor_helpers.py +57 -0
- bma_executor_sdk-0.4.0/tests/test_executor_conformance.py +151 -0
- bma_executor_sdk-0.4.0/tests/test_executor_contract.py +33 -0
- bma_executor_sdk-0.4.0/tests/test_executor_interop.py +128 -0
- bma_executor_sdk-0.4.0/tests/test_executor_runtime.py +409 -0
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: bma-executor-sdk
|
|
3
|
+
Version: 0.4.0
|
|
4
|
+
Summary: BMA executor protocol (bma.executor.v1.1): executor, harness and agent ends on asyncio with WebSocket, Unix socket and in-process transports.
|
|
5
|
+
Project-URL: Repository, https://github.com/beeos-ai/bma-sdks
|
|
6
|
+
License-Expression: Apache-2.0
|
|
7
|
+
Requires-Python: >=3.10
|
|
8
|
+
Requires-Dist: httpx<1,>=0.27
|
|
9
|
+
Requires-Dist: websockets<17,>=14
|
|
10
|
+
Description-Content-Type: text/markdown
|
|
11
|
+
|
|
12
|
+
# bma-executor-sdk (Python)
|
|
13
|
+
|
|
14
|
+
asyncio executor (`ExecutorClient`), harness (`Harness`) and agent ends of the
|
|
15
|
+
BMA executor protocol `bma.executor.v1.1`, with TypedDict protocol types
|
|
16
|
+
generated from the BMA contract (`contracts/bma/executor`). Transports:
|
|
17
|
+
WebSocket, WebSocket over a Unix domain socket (`unix://`), in-process.
|
|
18
|
+
|
|
19
|
+
```sh
|
|
20
|
+
pip install bma-executor-sdk
|
|
21
|
+
```
|
|
22
|
+
|
|
23
|
+
```python
|
|
24
|
+
import bma_executor_sdk as bx
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
See the [repository README](../../README.md#python-executor) and
|
|
28
|
+
[examples](examples). The REST/façade client is the separate `bma-sdk` package.
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
# bma-executor-sdk (Python)
|
|
2
|
+
|
|
3
|
+
asyncio executor (`ExecutorClient`), harness (`Harness`) and agent ends of the
|
|
4
|
+
BMA executor protocol `bma.executor.v1.1`, with TypedDict protocol types
|
|
5
|
+
generated from the BMA contract (`contracts/bma/executor`). Transports:
|
|
6
|
+
WebSocket, WebSocket over a Unix domain socket (`unix://`), in-process.
|
|
7
|
+
|
|
8
|
+
```sh
|
|
9
|
+
pip install bma-executor-sdk
|
|
10
|
+
```
|
|
11
|
+
|
|
12
|
+
```python
|
|
13
|
+
import bma_executor_sdk as bx
|
|
14
|
+
```
|
|
15
|
+
|
|
16
|
+
See the [repository README](../../README.md#python-executor) and
|
|
17
|
+
[examples](examples). The REST/façade client is the separate `bma-sdk` package.
|
|
@@ -0,0 +1,79 @@
|
|
|
1
|
+
"""Agent executor: a scripted local agent served over the executor protocol,
|
|
2
|
+
driven by an in-process harness with an in-memory platform.
|
|
3
|
+
|
|
4
|
+
python agent_minimal.py
|
|
5
|
+
|
|
6
|
+
Against BMA, keep EchoAgent and serve_agent() and build the client as in
|
|
7
|
+
executor_minimal.py; BMA is then the platform.
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
import asyncio
|
|
13
|
+
from typing import Any
|
|
14
|
+
|
|
15
|
+
from bma_executor_sdk import AgentEventRecord, AgentTurn, ExecutorClient, Harness, static_token
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class EchoAgent:
|
|
19
|
+
async def open_turn(self, params: dict[str, Any]) -> dict[str, Any]:
|
|
20
|
+
return {"accepted": True, "turn_id": params["turn_id"], "executor_session_ref": "echo_" + params["session_id"], "resumed": False, "effective": {"model": "echo"}}
|
|
21
|
+
|
|
22
|
+
async def run_turn(self, turn: AgentTurn) -> None:
|
|
23
|
+
text = turn.params["input"][0]["text"]
|
|
24
|
+
message_id = turn.new_event_id()
|
|
25
|
+
await turn.delta(message_id, "agent.message", 0, "Echo: ") # streaming preview
|
|
26
|
+
turn.emit_with_id(message_id, {"type": "agent.message", "content": [{"type": "text", "text": "Echo: " + text}]})
|
|
27
|
+
tool_use_id = turn.emit({"type": "agent.tool_use", "name": "bash", "input": {"command": "date"}, "evaluated_permission": {"type": "ask"}})
|
|
28
|
+
answer = await turn.request_interaction({"kind": "permission", "tool_use_id": tool_use_id, "tool_name": "bash", "tool_input": {"command": "date"}, "scopes": ["once"]}) # type: ignore[typeddict-item]
|
|
29
|
+
allowed = answer["outcome"] == "allow"
|
|
30
|
+
turn.emit({"type": "agent.tool_result", "tool_use_id": tool_use_id, "is_error": not allowed, "content": [{"type": "text", "text": "ran" if allowed else "denied"}]})
|
|
31
|
+
|
|
32
|
+
async def close_session(self, session_id: str, reason: str) -> None:
|
|
33
|
+
print("closed", session_id, reason)
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
class PrintingPlatform:
|
|
37
|
+
"""AgentPlatform that prints events and allows every interaction."""
|
|
38
|
+
|
|
39
|
+
def __init__(self) -> None:
|
|
40
|
+
self.harness: Harness | None = None
|
|
41
|
+
self.idle = asyncio.Event()
|
|
42
|
+
|
|
43
|
+
async def acked_sequence(self, session_id: str) -> int:
|
|
44
|
+
return 0
|
|
45
|
+
|
|
46
|
+
async def agent_event(self, executor: Any, record: AgentEventRecord) -> None:
|
|
47
|
+
print(f"#{record.sequence} {record.type}")
|
|
48
|
+
if record.type == "session.status_idle":
|
|
49
|
+
self.idle.set()
|
|
50
|
+
|
|
51
|
+
def agent_delta(self, executor: Any, delta: dict[str, Any]) -> None:
|
|
52
|
+
print(" delta:", delta["text"])
|
|
53
|
+
|
|
54
|
+
async def interaction_requested(self, executor: Any, request: dict[str, Any]) -> str:
|
|
55
|
+
print(" permission asked for", request["tool_name"])
|
|
56
|
+
asyncio.ensure_future(executor.respond_interaction({"session_id": request["session_id"], "interaction_id": request["interaction_id"], "outcome": "allow"}))
|
|
57
|
+
return "pint_" + request["interaction_id"]
|
|
58
|
+
|
|
59
|
+
async def interaction_canceled(self, executor: Any, cancel: dict[str, Any]) -> None:
|
|
60
|
+
return None
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
async def main() -> None:
|
|
64
|
+
platform = PrintingPlatform()
|
|
65
|
+
harness = Harness(agent=platform)
|
|
66
|
+
client = ExecutorClient(environment_id="env_local", executor={"name": "echo-agent", "os": "local", "arch": "local", "version": "0.1.0"}, token_source=static_token("inprocess://"), dialer=harness.inprocess_dialer()) # type: ignore[typeddict-item]
|
|
67
|
+
client.serve_agent(EchoAgent(), {"models": [{"id": "echo"}], "permission_modes": ["default"], "interaction_kinds": ["permission"]}) # type: ignore[arg-type,typeddict-item]
|
|
68
|
+
await client.start()
|
|
69
|
+
try:
|
|
70
|
+
executor = await harness.wait_executor()
|
|
71
|
+
await executor.start_turn({"session_id": "ses_demo", "turn_id": "turn_1", "input_id": "in_1", "input": [{"type": "text", "text": "hello"}]}) # type: ignore[typeddict-item]
|
|
72
|
+
await asyncio.wait_for(platform.idle.wait(), 10)
|
|
73
|
+
await executor.close_agent_session({"session_id": "ses_demo", "reason": "done"}) # type: ignore[typeddict-item]
|
|
74
|
+
finally:
|
|
75
|
+
await client.stop("done")
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
if __name__ == "__main__":
|
|
79
|
+
asyncio.run(main())
|
|
@@ -0,0 +1,64 @@
|
|
|
1
|
+
"""Minimal executor: registers with BMA and answers bash tool calls.
|
|
2
|
+
|
|
3
|
+
BMA_ENVIRONMENT_KEY=... BMA_ENVIRONMENT_ID=... python executor_minimal.py
|
|
4
|
+
|
|
5
|
+
Set BMA_HARNESS_URL (ws://..., unix:///path.sock) to connect to a local
|
|
6
|
+
harness instead of BMA, for example the one harness_loopback.py starts.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
import asyncio
|
|
12
|
+
import os
|
|
13
|
+
import platform
|
|
14
|
+
|
|
15
|
+
from bma_executor_sdk import ExecutorClient, static_token
|
|
16
|
+
from bma_executor_sdk.generated import ToolCallParams, ToolResultParams
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
async def run_bash(call: ToolCallParams) -> ToolResultParams:
|
|
20
|
+
proc = await asyncio.create_subprocess_shell(
|
|
21
|
+
call["input"].get("command", ""), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT,
|
|
22
|
+
)
|
|
23
|
+
try:
|
|
24
|
+
output, _ = await proc.communicate()
|
|
25
|
+
except asyncio.CancelledError: # tool.cancel or the deadline: stop the process
|
|
26
|
+
proc.kill()
|
|
27
|
+
raise
|
|
28
|
+
return {
|
|
29
|
+
"session_id": call["session_id"],
|
|
30
|
+
"tool_call_id": call["tool_call_id"],
|
|
31
|
+
"outcome": "succeeded" if proc.returncode == 0 else "failed",
|
|
32
|
+
"is_error": proc.returncode != 0,
|
|
33
|
+
"content": [{"type": "text", "text": output.decode(errors="replace")}],
|
|
34
|
+
} # type: ignore[typeddict-item]
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def make_client() -> ExecutorClient:
|
|
38
|
+
info = {"name": "py-executor", "os": platform.system().lower(), "arch": platform.machine(), "version": "0.1.0"}
|
|
39
|
+
capabilities = {"shell": {"enabled": True}}
|
|
40
|
+
harness_url = os.environ.get("BMA_HARNESS_URL")
|
|
41
|
+
if harness_url:
|
|
42
|
+
return ExecutorClient(environment_id="env_local", executor=info, capabilities=capabilities, token_source=static_token(harness_url)) # type: ignore[arg-type]
|
|
43
|
+
return ExecutorClient(
|
|
44
|
+
environment_id=os.environ["BMA_ENVIRONMENT_ID"],
|
|
45
|
+
environment_key=os.environ["BMA_ENVIRONMENT_KEY"],
|
|
46
|
+
base_url=os.environ.get("BMA_BASE_URL", "https://api.bma.beeos.ai"),
|
|
47
|
+
executor=info, # type: ignore[arg-type]
|
|
48
|
+
capabilities=capabilities, # type: ignore[arg-type]
|
|
49
|
+
)
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
async def main() -> None:
|
|
53
|
+
client = make_client()
|
|
54
|
+
client.on_tool_call(run_bash)
|
|
55
|
+
registered = await client.start()
|
|
56
|
+
print("registered as", registered["executor_id"])
|
|
57
|
+
try:
|
|
58
|
+
await asyncio.Event().wait()
|
|
59
|
+
finally:
|
|
60
|
+
await client.stop("shutdown")
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
if __name__ == "__main__":
|
|
64
|
+
asyncio.run(main())
|
|
@@ -0,0 +1,50 @@
|
|
|
1
|
+
"""Local brain and hand in one process: a harness and an executor connected
|
|
2
|
+
in-process (or over a Unix socket with --unix PATH), one bash tool call.
|
|
3
|
+
|
|
4
|
+
python harness_loopback.py [--unix /tmp/bma-harness.sock]
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import argparse
|
|
10
|
+
import asyncio
|
|
11
|
+
|
|
12
|
+
from bma_executor_sdk import ExecutorClient, Harness, static_token
|
|
13
|
+
from executor_minimal import run_bash
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
async def main(unix_path: str | None) -> None:
|
|
17
|
+
harness = Harness() # no authenticate hook: in-process or a private socket only
|
|
18
|
+
server = None
|
|
19
|
+
if unix_path:
|
|
20
|
+
server = await harness.serve_unix(unix_path)
|
|
21
|
+
client = ExecutorClient(environment_id="env_local", executor={"name": "loopback", "os": "local", "arch": "local", "version": "0.1.0"}, token_source=static_token(f"unix://{unix_path}")) # type: ignore[typeddict-item]
|
|
22
|
+
else:
|
|
23
|
+
client = ExecutorClient(environment_id="env_local", executor={"name": "loopback", "os": "local", "arch": "local", "version": "0.1.0"}, token_source=static_token("inprocess://"), dialer=harness.inprocess_dialer()) # type: ignore[typeddict-item]
|
|
24
|
+
|
|
25
|
+
async def on_assign(params): # type: ignore[no-untyped-def]
|
|
26
|
+
# Report ready once the workspace is prepared (here: nothing to do).
|
|
27
|
+
asyncio.ensure_future(client.session_ready({"session_id": params["session_id"], "work_id": params["work_id"], "lease_token": params["lease"]["token"]}))
|
|
28
|
+
return {"accepted": True, "session_id": params["session_id"]}
|
|
29
|
+
|
|
30
|
+
client.on_session_assign(on_assign)
|
|
31
|
+
client.on_tool_call(run_bash)
|
|
32
|
+
await client.start()
|
|
33
|
+
try:
|
|
34
|
+
executor = await harness.wait_executor()
|
|
35
|
+
session = await executor.assign_session({"session_id": "ses_local"}) # type: ignore[typeddict-item]
|
|
36
|
+
await session.ready()
|
|
37
|
+
result = await session.call_tool({"tool": "bash", "input": {"command": "echo hello from the hand"}}) # type: ignore[typeddict-item]
|
|
38
|
+
print(result["outcome"], result["content"][0]["text"].strip())
|
|
39
|
+
await session.stop("done")
|
|
40
|
+
finally:
|
|
41
|
+
await client.stop("done")
|
|
42
|
+
if server is not None:
|
|
43
|
+
server.close()
|
|
44
|
+
await server.wait_closed() # also removes the socket file
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
if __name__ == "__main__":
|
|
48
|
+
parser = argparse.ArgumentParser()
|
|
49
|
+
parser.add_argument("--unix")
|
|
50
|
+
asyncio.run(main(parser.parse_args().unix))
|
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "bma-executor-sdk"
|
|
3
|
+
version = "0.4.0"
|
|
4
|
+
description = "BMA executor protocol (bma.executor.v1.1): executor, harness and agent ends on asyncio with WebSocket, Unix socket and in-process transports."
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
license = "Apache-2.0"
|
|
7
|
+
requires-python = ">=3.10"
|
|
8
|
+
dependencies = ["httpx>=0.27,<1", "websockets>=14,<17"]
|
|
9
|
+
|
|
10
|
+
[project.urls]
|
|
11
|
+
Repository = "https://github.com/beeos-ai/bma-sdks"
|
|
12
|
+
|
|
13
|
+
[dependency-groups]
|
|
14
|
+
dev = ["pytest>=8", "pytest-asyncio>=0.24"]
|
|
15
|
+
|
|
16
|
+
[build-system]
|
|
17
|
+
requires = ["hatchling>=1.26"]
|
|
18
|
+
build-backend = "hatchling.build"
|
|
19
|
+
|
|
20
|
+
[tool.hatch.build.targets.wheel]
|
|
21
|
+
packages = ["src/bma_executor_sdk"]
|
|
22
|
+
|
|
23
|
+
[tool.pytest.ini_options]
|
|
24
|
+
testpaths = ["tests"]
|
|
25
|
+
asyncio_mode = "auto"
|
|
26
|
+
asyncio_default_fixture_loop_scope = "function"
|
|
@@ -0,0 +1,77 @@
|
|
|
1
|
+
"""BMA executor protocol (bma.executor.v1.1): executor, harness and agent ends.
|
|
2
|
+
|
|
3
|
+
Install with ``pip install bma-executor-sdk``.
|
|
4
|
+
Protocol types are generated from the BMA contract into
|
|
5
|
+
:mod:`bma_executor_sdk.generated` (TypedDicts: plain JSON dicts on the wire).
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from . import generated
|
|
9
|
+
from ._rpc import ClientStoppedError, ExecutorError, NotConnectedError, RPCError
|
|
10
|
+
from .agent import AgentEventRecord, AgentHandler, AgentPlatform, AgentSessionClosedError, AgentTurn
|
|
11
|
+
from .client import (
|
|
12
|
+
DEFAULT_BASE_URL,
|
|
13
|
+
PROTOCOL_V1,
|
|
14
|
+
PROTOCOL_V11,
|
|
15
|
+
ExecutorClient,
|
|
16
|
+
ExecutorTokenError,
|
|
17
|
+
ExecutorTokenScopeError,
|
|
18
|
+
LogEvent,
|
|
19
|
+
MCPChannel,
|
|
20
|
+
static_token,
|
|
21
|
+
)
|
|
22
|
+
from .harness import (
|
|
23
|
+
AssignRejectedError,
|
|
24
|
+
Executor,
|
|
25
|
+
Harness,
|
|
26
|
+
HarnessMCPChannel,
|
|
27
|
+
Principal,
|
|
28
|
+
Session,
|
|
29
|
+
SessionEndedError,
|
|
30
|
+
ToolOutcomeUnknownError,
|
|
31
|
+
Unauthorized,
|
|
32
|
+
)
|
|
33
|
+
from .transport import (
|
|
34
|
+
TERMINAL_CLOSE_CODES,
|
|
35
|
+
ChannelClosed,
|
|
36
|
+
ConnectRequest,
|
|
37
|
+
Connection,
|
|
38
|
+
Dialer,
|
|
39
|
+
dial,
|
|
40
|
+
)
|
|
41
|
+
|
|
42
|
+
__all__ = [
|
|
43
|
+
"DEFAULT_BASE_URL",
|
|
44
|
+
"PROTOCOL_V1",
|
|
45
|
+
"PROTOCOL_V11",
|
|
46
|
+
"TERMINAL_CLOSE_CODES",
|
|
47
|
+
"AgentEventRecord",
|
|
48
|
+
"AgentHandler",
|
|
49
|
+
"AgentPlatform",
|
|
50
|
+
"AgentSessionClosedError",
|
|
51
|
+
"AgentTurn",
|
|
52
|
+
"AssignRejectedError",
|
|
53
|
+
"ChannelClosed",
|
|
54
|
+
"ClientStoppedError",
|
|
55
|
+
"ConnectRequest",
|
|
56
|
+
"Connection",
|
|
57
|
+
"Dialer",
|
|
58
|
+
"Executor",
|
|
59
|
+
"ExecutorClient",
|
|
60
|
+
"ExecutorError",
|
|
61
|
+
"ExecutorTokenError",
|
|
62
|
+
"ExecutorTokenScopeError",
|
|
63
|
+
"Harness",
|
|
64
|
+
"HarnessMCPChannel",
|
|
65
|
+
"LogEvent",
|
|
66
|
+
"MCPChannel",
|
|
67
|
+
"NotConnectedError",
|
|
68
|
+
"Principal",
|
|
69
|
+
"RPCError",
|
|
70
|
+
"Session",
|
|
71
|
+
"SessionEndedError",
|
|
72
|
+
"ToolOutcomeUnknownError",
|
|
73
|
+
"Unauthorized",
|
|
74
|
+
"dial",
|
|
75
|
+
"generated",
|
|
76
|
+
"static_token",
|
|
77
|
+
]
|
|
@@ -0,0 +1,192 @@
|
|
|
1
|
+
"""JSON-RPC 2.0 over a Connection, shared by the executor and harness ends."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import asyncio
|
|
6
|
+
import json
|
|
7
|
+
from collections.abc import Awaitable, Callable
|
|
8
|
+
from typing import Any
|
|
9
|
+
|
|
10
|
+
from .generated import JsonRpcErrorCode
|
|
11
|
+
from .transport import ChannelClosed, Connection
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class ExecutorError(Exception):
|
|
15
|
+
"""Base class of bma_executor_sdk errors."""
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class NotConnectedError(ExecutorError):
|
|
19
|
+
"""The channel is not connected (yet, or any more)."""
|
|
20
|
+
|
|
21
|
+
def __init__(self, message: str = "executor channel is not connected") -> None:
|
|
22
|
+
super().__init__(message)
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class ClientStoppedError(ExecutorError):
|
|
26
|
+
"""The client was stopped, or closed with a terminal close code."""
|
|
27
|
+
|
|
28
|
+
def __init__(self, message: str = "client stopped") -> None:
|
|
29
|
+
super().__init__(message)
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
class RPCError(ExecutorError):
|
|
33
|
+
"""A JSON-RPC error. A handler that raises it keeps its code on the wire."""
|
|
34
|
+
|
|
35
|
+
def __init__(self, code: int, message: str, data: dict[str, Any] | None = None) -> None:
|
|
36
|
+
super().__init__(f"json-rpc error {code}: {message}")
|
|
37
|
+
self.code = code
|
|
38
|
+
self.message = message
|
|
39
|
+
self.data = data
|
|
40
|
+
|
|
41
|
+
@property
|
|
42
|
+
def retryable(self) -> bool | None:
|
|
43
|
+
return None if self.data is None else self.data.get("retryable")
|
|
44
|
+
|
|
45
|
+
def to_wire(self) -> dict[str, Any]:
|
|
46
|
+
error: dict[str, Any] = {"code": int(self.code), "message": self.message}
|
|
47
|
+
if self.data is not None:
|
|
48
|
+
error["data"] = self.data
|
|
49
|
+
return error
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def rpc_error(code: JsonRpcErrorCode, message: str, retryable: bool) -> RPCError:
|
|
53
|
+
return RPCError(code, message, {"retryable": retryable})
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
RequestHandler = Callable[[str, str, Any], Awaitable[Any]]
|
|
57
|
+
NotificationHandler = Callable[[str, Any], None]
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
class Peer:
|
|
61
|
+
"""One JSON-RPC endpoint on a Connection.
|
|
62
|
+
|
|
63
|
+
Inbound requests run as tasks (except methods in ``inline``, which are
|
|
64
|
+
handled in read order); notifications are handled in read order.
|
|
65
|
+
"""
|
|
66
|
+
|
|
67
|
+
def __init__(
|
|
68
|
+
self,
|
|
69
|
+
conn: Connection,
|
|
70
|
+
*,
|
|
71
|
+
id_prefix: str,
|
|
72
|
+
on_request: RequestHandler,
|
|
73
|
+
on_notification: NotificationHandler,
|
|
74
|
+
on_log: Callable[[str, str, BaseException | None], None],
|
|
75
|
+
inline: frozenset[str] = frozenset(),
|
|
76
|
+
) -> None:
|
|
77
|
+
self.conn = conn
|
|
78
|
+
self._id_prefix = id_prefix
|
|
79
|
+
self._on_request = on_request
|
|
80
|
+
self._on_notification = on_notification
|
|
81
|
+
self._log = on_log
|
|
82
|
+
self._inline = inline
|
|
83
|
+
self._pending: dict[str, asyncio.Future[Any]] = {}
|
|
84
|
+
self._next_id = 0
|
|
85
|
+
self._tasks: set[asyncio.Task[Any]] = set()
|
|
86
|
+
self.closed: ChannelClosed | None = None
|
|
87
|
+
|
|
88
|
+
async def call(self, method: str, params: Any, timeout: float) -> Any:
|
|
89
|
+
if self.closed is not None:
|
|
90
|
+
raise NotConnectedError()
|
|
91
|
+
self._next_id += 1
|
|
92
|
+
request_id = f"{self._id_prefix}_{self._next_id}"
|
|
93
|
+
future: asyncio.Future[Any] = asyncio.get_running_loop().create_future()
|
|
94
|
+
self._pending[request_id] = future
|
|
95
|
+
try:
|
|
96
|
+
await self._send({"jsonrpc": "2.0", "id": request_id, "method": method, "params": params})
|
|
97
|
+
return await asyncio.wait_for(future, timeout)
|
|
98
|
+
except asyncio.TimeoutError as error:
|
|
99
|
+
raise TimeoutError(f"{method}: no response within {timeout}s") from error
|
|
100
|
+
finally:
|
|
101
|
+
self._pending.pop(request_id, None)
|
|
102
|
+
|
|
103
|
+
async def notify(self, method: str, params: Any) -> None:
|
|
104
|
+
if self.closed is not None:
|
|
105
|
+
raise NotConnectedError()
|
|
106
|
+
await self._send({"jsonrpc": "2.0", "method": method, "params": params})
|
|
107
|
+
|
|
108
|
+
async def _send(self, message: dict[str, Any]) -> None:
|
|
109
|
+
try:
|
|
110
|
+
await self.conn.send(json.dumps(message, separators=(",", ":")))
|
|
111
|
+
except ChannelClosed as error:
|
|
112
|
+
raise NotConnectedError(f"write {message.get('method', 'response')}: {error}") from error
|
|
113
|
+
|
|
114
|
+
async def run(self) -> ChannelClosed:
|
|
115
|
+
"""Read until the channel closes; then fail pending calls."""
|
|
116
|
+
try:
|
|
117
|
+
while True:
|
|
118
|
+
raw = await self.conn.recv()
|
|
119
|
+
try:
|
|
120
|
+
message = json.loads(raw)
|
|
121
|
+
except ValueError as error:
|
|
122
|
+
self._log("warn", "invalid JSON-RPC frame", error)
|
|
123
|
+
continue
|
|
124
|
+
if isinstance(message, dict) and message.get("method") in self._inline and message.get("id") is not None:
|
|
125
|
+
await self._handle_request(str(message["id"]), message["method"], message.get("params"))
|
|
126
|
+
continue
|
|
127
|
+
self._dispatch(message)
|
|
128
|
+
except ChannelClosed as closed:
|
|
129
|
+
self.closed = closed
|
|
130
|
+
except asyncio.CancelledError:
|
|
131
|
+
self.closed = ChannelClosed(1006, "reader canceled")
|
|
132
|
+
raise
|
|
133
|
+
finally:
|
|
134
|
+
if self.closed is None:
|
|
135
|
+
self.closed = ChannelClosed(1006, "reader failed")
|
|
136
|
+
self.fail_pending(NotConnectedError())
|
|
137
|
+
return self.closed
|
|
138
|
+
|
|
139
|
+
def fail_pending(self, error: Exception) -> None:
|
|
140
|
+
pending, self._pending = self._pending, {}
|
|
141
|
+
for future in pending.values():
|
|
142
|
+
if not future.done():
|
|
143
|
+
future.set_exception(error)
|
|
144
|
+
|
|
145
|
+
def _dispatch(self, message: Any) -> None:
|
|
146
|
+
if not isinstance(message, dict):
|
|
147
|
+
self._log("warn", "JSON-RPC frame is not an object", None)
|
|
148
|
+
return
|
|
149
|
+
method, request_id = message.get("method"), message.get("id")
|
|
150
|
+
if method and request_id is not None:
|
|
151
|
+
self._spawn(self._handle_request(str(request_id), method, message.get("params")))
|
|
152
|
+
elif method:
|
|
153
|
+
try:
|
|
154
|
+
self._on_notification(method, message.get("params"))
|
|
155
|
+
except Exception as error: # a handler bug must not stop the channel
|
|
156
|
+
self._log("error", f"{method} notification handler failed", error)
|
|
157
|
+
elif request_id is not None:
|
|
158
|
+
future = self._pending.pop(str(request_id), None)
|
|
159
|
+
if future is None or future.done():
|
|
160
|
+
return
|
|
161
|
+
error = message.get("error")
|
|
162
|
+
if error is not None:
|
|
163
|
+
future.set_exception(RPCError(error.get("code", 0), error.get("message", ""), error.get("data")))
|
|
164
|
+
else:
|
|
165
|
+
future.set_result(message.get("result"))
|
|
166
|
+
else:
|
|
167
|
+
self._log("warn", "JSON-RPC frame is neither request nor response", None)
|
|
168
|
+
|
|
169
|
+
def _spawn(self, coro: Awaitable[None]) -> None:
|
|
170
|
+
task = asyncio.ensure_future(coro)
|
|
171
|
+
self._tasks.add(task)
|
|
172
|
+
task.add_done_callback(self._tasks.discard)
|
|
173
|
+
|
|
174
|
+
async def _handle_request(self, request_id: str, method: str, params: Any) -> None:
|
|
175
|
+
try:
|
|
176
|
+
result = await self._on_request(request_id, method, params)
|
|
177
|
+
response: dict[str, Any] = {"jsonrpc": "2.0", "id": request_id, "result": result}
|
|
178
|
+
except RPCError as error:
|
|
179
|
+
response = {"jsonrpc": "2.0", "id": request_id, "error": error.to_wire()}
|
|
180
|
+
except asyncio.CancelledError:
|
|
181
|
+
raise
|
|
182
|
+
except Exception as error:
|
|
183
|
+
self._log("warn", f"{method} handler failed", error)
|
|
184
|
+
response = {"jsonrpc": "2.0", "id": request_id, "error": {"code": int(JsonRpcErrorCode.InternalError), "message": str(error)}}
|
|
185
|
+
try:
|
|
186
|
+
await self._send(response)
|
|
187
|
+
except NotConnectedError as error:
|
|
188
|
+
self._log("debug", f"write {method} response", error)
|
|
189
|
+
|
|
190
|
+
def cancel_tasks(self) -> None:
|
|
191
|
+
for task in list(self._tasks):
|
|
192
|
+
task.cancel()
|