okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

dawnkernel: a job_stop frame lets /jobs stop remote background jobs57a82be

main
branchesmainall refs →
9.7 KiB362 linesraw
1"""Checked JSON messages at the kernel and remote connection boundaries."""
2
3from typing import Literal, NotRequired, TypedDict, cast
4
5
6class Execute(TypedDict):
7    type: Literal["execute"]
8    id: str
9    code: str
10    durable: NotRequired[bool]
11    max_edge: NotRequired[int]
12    timeout_ms: NotRequired[int]
13
14
15class HostSuccess(TypedDict):
16    ok: Literal[True]
17    value: object
18
19
20class HostFailure(TypedDict):
21    ok: Literal[False]
22    code: str
23    message: str
24
25
26HostReply = HostSuccess | HostFailure
27
28
29class Reply(TypedDict):
30    type: Literal["reply"]
31    id: str
32    value: HostReply
33
34
35class Interrupt(TypedDict):
36    type: Literal["interrupt"]
37    id: str
38    # `memory` is the kernel's own, from its memory guard; the owner sends the others.
39    reason: NotRequired[Literal["deadline", "cancelled", "memory"]]
40
41
42class Shutdown(TypedDict):
43    type: Literal["shutdown"]
44
45
46Invoke = TypedDict(
47    "Invoke",
48    {
49        "type": Literal["invoke"],
50        "id": str,
51        "name": NotRequired[str],
52        "target": NotRequired[dict[str, str]],
53        "args": NotRequired[list[object]],
54        "kwargs": NotRequired[dict[str, object]],
55        "await": NotRequired[bool],
56    },
57)
58
59
60class Introspect(TypedDict):
61    type: Literal["introspect"]
62    id: str
63
64
65class Release(TypedDict):
66    type: Literal["release"]
67    handle: str
68
69
70class State(TypedDict):
71    """Save the namespace to disk, or revive it from an earlier save."""
72
73    type: Literal["snapshot", "restore"]
74    id: str
75    path: str
76
77
78class JobStop(TypedDict):
79    """Stop a background job."""
80
81    type: Literal["job_stop"]
82    id: str
83
84
85Incoming = (
86    Execute
87    | Reply
88    | Interrupt
89    | Shutdown
90    | Invoke
91    | Introspect
92    | Release
93    | State
94    | JobStop
95)
96
97
98class SavedCell(TypedDict):
99    id: str
100    source: str
101    started: bool
102    parent: str | None
103    finished: bool
104
105
106class Ready(TypedDict):
107    type: Literal["ready"]
108    pid: int
109    pgid: int | None
110    leader: str | None
111
112
113class StartupError(TypedDict):
114    type: Literal["startup_error"]
115    message: str
116
117
118class RemoteCall(TypedDict):
119    type: Literal["call"]
120    id: str
121    method: str
122    args: dict[str, object]
123
124
125class InvocationError(TypedDict):
126    ename: str
127    evalue: str
128    traceback: list[str]
129
130
131class MirrorState(TypedDict, total=False):
132    job: str | None
133    seen: int
134    tail: str
135    tail_bytes: int
136    exit_code: int | None
137    timed_out: bool
138    duration: float | None
139
140
141class Invoked(TypedDict):
142    type: Literal["invoked"]
143    id: str
144    ok: bool
145    value: NotRequired[object]
146    handle: NotRequired[str]
147    state: NotRequired[MirrorState]
148    error: NotRequired[InvocationError]
149    cancelled: NotRequired[bool]
150
151
152class RemoteHandle(TypedDict):
153    id: str
154    repr: str
155
156
157class Introspected(TypedDict):
158    type: Literal["introspected"]
159    id: str
160    names: list[str]
161    handles: list[RemoteHandle]
162
163
164class Mirror(MirrorState):
165    type: Literal["mirror"]
166    handle: str
167
168
169class RemoteEvent(TypedDict):
170    type: Literal["job_start", "job", "trace", "done"]
171    id: str
172    service: NotRequired[bool]  # job_start: nothing waits on the job's end
173
174
175class RemoteCleanup(TypedDict):
176    type: Literal["cleanup"]
177    module: str
178
179
180class JobStopReply(TypedDict):
181    type: Literal["job_stop"]
182    id: str
183    ok: bool
184    error: NotRequired[str]
185
186
187RemoteMessage = (
188    Ready
189    | StartupError
190    | RemoteCall
191    | Invoked
192    | Introspected
193    | Mirror
194    | RemoteEvent
195    | RemoteCleanup
196    | JobStopReply
197)
198
199
200def _object(value: object, context: str) -> dict[str, object]:
201    if not isinstance(value, dict) or not all(isinstance(key, str) for key in value):
202        raise ValueError(f"{context}: expected an object")
203    return cast(dict[str, object], value)
204
205
206def _field(
207    message: dict[str, object],
208    field: str,
209    expected: type | tuple[type, ...],
210    context: str,
211    *,
212    optional: bool = False,
213) -> None:
214    if optional and field not in message:
215        return
216    if field not in message:
217        raise ValueError(f"{context}.{field}: required field")
218    value = message[field]
219    if not isinstance(value, expected) or (
220        isinstance(value, bool)
221        and expected is not bool
222        and (expected is int or isinstance(expected, tuple) and int in expected)
223    ):
224        raise ValueError(f"{context}.{field}: invalid field type")
225
226
227def _envelope(value: object) -> tuple[dict[str, object], str]:
228    message = _object(value, "message")
229    _field(message, "type", str, "message")
230    return message, cast(str, message["type"])
231
232
233def parse_incoming(value: object) -> Incoming:
234    """Validate a kernel request before it reaches the event loop."""
235    message, kind = _envelope(value)
236    if kind not in {
237        "execute",
238        "reply",
239        "interrupt",
240        "shutdown",
241        "invoke",
242        "introspect",
243        "release",
244        "snapshot",
245        "restore",
246        "job_stop",
247    }:
248        raise ValueError(f"{kind}.type: unknown message type")
249    if kind not in {"shutdown", "release"}:
250        _field(message, "id", str, kind)
251    if kind == "execute":
252        _field(message, "code", str, kind)
253        _field(message, "durable", bool, kind, optional=True)
254        _field(message, "max_edge", int, kind, optional=True)
255        _field(message, "timeout_ms", int, kind, optional=True)
256    elif kind == "reply":
257        reply = _object(message.get("value"), "reply.value")
258        _field(reply, "ok", bool, "reply.value")
259        if reply["ok"]:
260            _field(reply, "value", object, "reply.value")
261        else:
262            _field(reply, "code", str, "reply.value")
263            _field(reply, "message", str, "reply.value")
264    elif kind == "interrupt":
265        _field(message, "reason", str, kind, optional=True)
266        if "reason" in message and message["reason"] not in (
267            "deadline",
268            "cancelled",
269            "memory",
270        ):
271            raise ValueError("interrupt.reason: expected deadline, cancelled or memory")
272    elif kind == "invoke":
273        for field, expected in (
274            ("name", str),
275            ("args", list),
276            ("kwargs", dict),
277            ("await", bool),
278        ):
279            _field(message, field, expected, kind, optional=True)
280        if "kwargs" in message:
281            _object(message["kwargs"], "invoke.kwargs")
282        if "target" in message:
283            target = _object(message["target"], "invoke.target")
284            if set(target) not in ({"handle"}, {"pending"}):
285                raise ValueError(
286                    "invoke.target: expected exactly one handle or pending"
287                )
288            _field(target, next(iter(target)), str, "invoke.target")
289    elif kind == "release":
290        _field(message, "handle", str, kind)
291    elif kind in {"snapshot", "restore"}:
292        _field(message, "path", str, kind)
293    return cast(Incoming, message)
294
295
296def _mirror(value: object, context: str) -> None:
297    state = _object(value, context)
298    for field, expected in (
299        ("job", (str, type(None))),
300        ("seen", int),
301        ("tail", str),
302        ("tail_bytes", int),
303        ("exit_code", (int, type(None))),
304        ("timed_out", bool),
305        ("duration", (int, float, type(None))),
306    ):
307        _field(state, field, expected, context, optional=True)
308
309
310def parse_remote(value: object) -> RemoteMessage:
311    """Validate remote replies and the event fields the owner consumes."""
312    message, kind = _envelope(value)
313    if kind in {"invoked", "introspected", "call", "job_start", "job", "trace", "done"}:
314        _field(message, "id", str, kind)
315    if kind == "ready":
316        _field(message, "pid", int, kind)
317        _field(message, "pgid", (int, type(None)), kind)
318        _field(message, "leader", (str, type(None)), kind)
319    elif kind == "startup_error":
320        _field(message, "message", str, kind)
321    elif kind == "call":
322        _field(message, "method", str, kind)
323        _object(message.get("args"), "call.args")
324    elif kind == "invoked":
325        _field(message, "ok", bool, kind)
326        _field(message, "handle", str, kind, optional=True)
327        _field(message, "cancelled", bool, kind, optional=True)
328        if "state" in message:
329            _mirror(message["state"], "invoked.state")
330        if "error" in message or not message["ok"]:
331            error = _object(message.get("error"), "invoked.error")
332            _field(error, "ename", str, "invoked.error")
333            _field(error, "evalue", str, "invoked.error")
334            _field(error, "traceback", list, "invoked.error")
335            if not all(
336                isinstance(line, str) for line in cast(list[object], error["traceback"])
337            ):
338                raise ValueError("invoked.error.traceback: expected strings")
339        if message["ok"] and "handle" not in message and "value" not in message:
340            raise ValueError("invoked.value: required without handle")
341    elif kind == "introspected":
342        _field(message, "names", list, kind)
343        if not all(
344            isinstance(name, str) for name in cast(list[object], message["names"])
345        ):
346            raise ValueError("introspected.names: expected strings")
347        _field(message, "handles", list, kind)
348        for item in cast(list[object], message["handles"]):
349            handle = _object(item, "introspected.handles")
350            _field(handle, "id", str, "introspected.handles")
351            _field(handle, "repr", str, "introspected.handles")
352    elif kind == "mirror":
353        _field(message, "handle", str, kind)
354        _mirror(message, kind)
355    elif kind == "cleanup":
356        _field(message, "module", str, kind)
357    elif kind == "job_stop":
358        _field(message, "ok", bool, kind)
359        _field(message, "error", str, kind, optional=True)
360    elif kind not in {"job_start", "job", "trace", "done"}:
361        raise ValueError(f"{kind}.type: unknown message type")
362    return cast(RemoteMessage, message)