From 1d2aa8b316efc6f45b1afa5496f030e0acdf4810 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 17:34:10 +0100 Subject: [PATCH 01/45] Enable shared-core's sessions storage (feature sql) restate-sdk-shared-core 8.3.0 with its sql feature: the SQLite-compatible database kept in an object's state, with its page cache (foyer), for restate.sessions. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- Cargo.lock | 1795 +++++++++++++++++++++++++++++++++++++++++++++------- Cargo.toml | 10 +- 2 files changed, 1575 insertions(+), 230 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f630bc59..e5d96bbf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,20 +2,137 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aead" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d122413f284cf2d62fb1b7db97e02edb8cda96d769b16e443a4f6195e35662b0" +dependencies = [ + "crypto-common 0.1.7", + "generic-array", +] + +[[package]] +name = "aegis" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41585f248ed9406567ae6d61ca2a55ed3dc7348e65203a01229881253acb95e3" +dependencies = [ + "cc", + "softaes", +] + +[[package]] +name = "aes" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures 0.2.17", +] + +[[package]] +name = "aes-gcm" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "831010a0f742e1209b3bcea8fab6a8e149051ba6099432c8cb2cc117dec3ead1" +dependencies = [ + "aead", + "aes", + "cipher", + "ctr", + "ghash", + "subtle", +] + [[package]] name = "aho-corasick" -version = "1.1.3" +version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" dependencies = [ "memchr", ] +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + +[[package]] +name = "allocator-api2" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c880a97d28a3681c0267bd29cff89621202715b065127cd445fa0f0fe0aa2880" + +[[package]] +name = "android_system_properties" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae221649c9976a6f6c56ae1facf410f3ddb33cc661c4b7b61020a912d4237fbc" +dependencies = [ + "libc", +] + +[[package]] +name = "antithesis_sdk" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8ed6751597568a284314fd9e2db420665898600109dedc7d8d9b588f00ae9c4" +dependencies = [ + "libc", + "libloading", + "linkme", + "once_cell", + "rand 0.8.8", + "rustc_version_runtime", + "serde", + "serde_json", +] + [[package]] name = "anyhow" -version = "1.0.95" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "arc-swap" +version = "1.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34ac096ce696dc2fcabef30516bb13c0a68a11d30131d3df6f04711467681b04" +checksum = "380962b3d93b99642e50b9febf72c071e97e8b32d86eae9cf4ef582bb960ae6e" +dependencies = [ + "rustversion", +] + +[[package]] +name = "aristo" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fcdeefa800110050103e2459a2f8dbdbd560ad046e30f1f87e24ce19c2cb8bde" +dependencies = [ + "aristo-macros", +] + +[[package]] +name = "aristo-macros" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d64a66d21a80182b35b1741997a6d2456911f54b7eb1918aa4e2382fc205268d" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "assoc" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfdc70193dadb9d7287fa4b633f15f90c876915b31f6af17da307fc59c9859a8" [[package]] name = "async-trait" @@ -36,9 +153,9 @@ checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" [[package]] name = "autocfg" -version = "1.4.0" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" [[package]] name = "base16ct" @@ -58,12 +175,37 @@ version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" +[[package]] +name = "bigdecimal" +version = "0.4.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fa3f3d8cbf4dffcfe4991de61d012bef509a409ecbe9dd41049bfe32b4d4653" +dependencies = [ + "autocfg", + "libm", + "num-bigint", + "num-integer", + "num-traits", +] + [[package]] name = "bitflags" version = "2.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" +[[package]] +name = "bitvec" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddcec3d12c579d40898fe0a9a358a803c23e9c52ca3c425707f81c9436211837" +dependencies = [ + "funty", + "radium", + "tap", + "wyz", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -75,13 +217,22 @@ dependencies = [ [[package]] name = "block-buffer" -version = "0.12.0" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cdd35008169921d80bc60d3d0ab416eecb028c4cd653352907921d95084790be" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" dependencies = [ "hybrid-array", ] +[[package]] +name = "branches" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7436b21fd6195415058eb41c9090e156b556bc7e0377bb57c5c026a421521525" +dependencies = [ + "rustc_version", +] + [[package]] name = "bs58" version = "0.5.1" @@ -93,15 +244,41 @@ dependencies = [ [[package]] name = "bumpalo" -version = "3.16.0" +version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79296716171880943b8470b5f8d03aa55eb2e645a4874bdbb28adb49162e012c" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "bytemuck" +version = "1.25.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a1f896587b6f2c069c73d2f0913e2d590c3990285cd2f0b6aa02b786b4c679c" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.10.1" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" [[package]] name = "bytes-utils" @@ -120,14 +297,28 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "375989decec1a0fa424442bed520fa12ede43a48475ce8c7417d2c69b32b5f6b" dependencies = [ "find-msvc-tools", + "jobserver", + "libc", "shlex", ] [[package]] name = "cfg-if" -version = "1.0.0" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "cfg_aliases" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527" + +[[package]] +name = "cfg_block" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18758054972164c3264f7c8386f5fc6da6114cb46b619fd365d4e3b2dc3ae487" [[package]] name = "chacha20" @@ -136,10 +327,39 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" dependencies = [ "cfg-if", - "cpufeatures 0.3.0", + "cpufeatures 0.3.1", "rand_core 0.10.1", ] +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "js-sys", + "num-traits", + "wasm-bindgen", + "windows-link", +] + +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common 0.1.7", + "inout", +] + +[[package]] +name = "cmsketch" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7ee2cfacbd29706479902b06d75ad8f1362900836aa32799eabc7e004bfd854" + [[package]] name = "combine" version = "4.6.9" @@ -150,6 +370,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "concurrent-queue" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "const-oid" version = "0.9.6" @@ -188,6 +417,17 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "core_affinity" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a034b3a7b624016c6e13f5df875747cc25f884156aad2abd12b6c46797971342" +dependencies = [ + "libc", + "num_cpus", + "winapi", +] + [[package]] name = "cpufeatures" version = "0.2.17" @@ -199,13 +439,22 @@ dependencies = [ [[package]] name = "cpufeatures" -version = "0.3.0" +version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +checksum = "5ca28b0ae3115b884660db4118d803791fd6756b6e88f39c0f3f7859060d7566" dependencies = [ "libc", ] +[[package]] +name = "crc32c" +version = "0.6.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f132a176423ee852e8f0a35d43088d1850966b50745be818a88b76f64b1c6cb" +dependencies = [ + "rustc_version", +] + [[package]] name = "critical-section" version = "1.2.0" @@ -250,23 +499,33 @@ dependencies = [ [[package]] name = "crypto-common" -version = "0.1.6" +version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" dependencies = [ "generic-array", + "rand_core 0.6.4", "typenum", ] [[package]] name = "crypto-common" -version = "0.2.1" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77727bb15fa921304124b128af125e7e3b968275d1b108b379190264f4423710" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" dependencies = [ "hybrid-array", ] +[[package]] +name = "ctr" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835" +dependencies = [ + "cipher", +] + [[package]] name = "curve25519-dalek" version = "4.1.3" @@ -291,7 +550,7 @@ checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3" dependencies = [ "proc-macro2", "quote", - "syn 2.0.96", + "syn 2.0.119", ] [[package]] @@ -319,19 +578,19 @@ checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ "block-buffer 0.10.4", "const-oid 0.9.6", - "crypto-common 0.1.6", + "crypto-common 0.1.7", "subtle", ] [[package]] name = "digest" -version = "0.11.2" +version = "0.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4850db49bf08e663084f7fb5c87d202ef91a3907271aff24a94eb97ff039153c" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" dependencies = [ - "block-buffer 0.12.0", + "block-buffer 0.12.1", "const-oid 0.10.2", - "crypto-common 0.2.1", + "crypto-common 0.2.2", ] [[package]] @@ -385,9 +644,9 @@ dependencies = [ [[package]] name = "either" -version = "1.13.0" +version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" +checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be" [[package]] name = "elliptic-curve" @@ -416,6 +675,49 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "00d174d5400e5e8fd687ad1049e2f578285fa914201b1af7e8b112a4546bd826" +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "fallible-iterator" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" + +[[package]] +name = "fastant" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e825441bfb2d831c47c97d05821552db8832479f44c571b97fededbf0099c07" +dependencies = [ + "small_ctor", + "web-time", +] + +[[package]] +name = "fastbloom" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7f34442dbe69c60fe8eaf58a8cafff81a1f278816d8ab4db255b3bef4ac3c4" +dependencies = [ + "getrandom 0.3.4", + "libm", + "siphasher", +] + +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + [[package]] name = "ff" version = "0.13.1" @@ -444,6 +746,12 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -453,6 +761,148 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "foyer" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0abc0b87814989efa711f9becd9f26969820e2d3905db27d10969c4bd45890" +dependencies = [ + "anyhow", + "equivalent", + "foyer-common", + "foyer-memory", + "foyer-storage", + "foyer-tokio", + "futures-util", + "mea", + "mixtrics", + "pin-project", + "serde", + "tracing", +] + +[[package]] +name = "foyer-common" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3db80d5dece93adb7ad709c84578794724a9cba342a7e566c3551c7ec626789" +dependencies = [ + "anyhow", + "bytes", + "cfg-if", + "foyer-tokio", + "mixtrics", + "parking_lot", + "pin-project", + "twox-hash", +] + +[[package]] +name = "foyer-intrusive-collections" +version = "0.10.0-dev" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e4fee46bea69e0596130e3210e65d3424e0ac1e6df3bde6636304bdf1ca4a3b" +dependencies = [ + "memoffset", +] + +[[package]] +name = "foyer-memory" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db907f40a527ca2aa2f40a5f68b32ea58aa70f050cd233518e9ffd402cfba6ce" +dependencies = [ + "anyhow", + "bitflags", + "cmsketch", + "equivalent", + "foyer-common", + "foyer-intrusive-collections", + "foyer-tokio", + "futures-util", + "hashbrown 0.16.1", + "itertools 0.14.0", + "mea", + "mixtrics", + "parking_lot", + "paste", + "pin-project", + "serde", + "tracing", +] + +[[package]] +name = "foyer-storage" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1983f1db3d0710e9c9d5fc116d9202dccd41a2d1e032572224f1aff5520aa958" +dependencies = [ + "allocator-api2 0.2.21", + "anyhow", + "bytes", + "core_affinity", + "equivalent", + "fastant", + "foyer-common", + "foyer-memory", + "foyer-tokio", + "fs4", + "futures-core", + "futures-util", + "hashbrown 0.16.1", + "io-uring", + "itertools 0.14.0", + "libc", + "lz4", + "mea", + "parking_lot", + "pin-project", + "rand 0.9.5", + "tracing", + "twox-hash", + "zstd", +] + +[[package]] +name = "foyer-tokio" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6577b05a7ffad0db555aedf00bfe52af818220fc4c1c3a7a12520896fc38627" +dependencies = [ + "tokio", +] + +[[package]] +name = "fs4" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8640e34b88f7652208ce9e88b1a37a2ae95227d84abec377ccd3c5cfeb141ed4" +dependencies = [ + "rustix", + "windows-sys 0.59.0", +] + +[[package]] +name = "funty" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c" + +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.34" @@ -460,6 +910,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -468,6 +919,17 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + [[package]] name = "futures-io" version = "0.3.34" @@ -503,18 +965,37 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" dependencies = [ + "futures-channel", "futures-core", + "futures-io", "futures-macro", + "futures-sink", "futures-task", + "memchr", "pin-project-lite", "slab", ] +[[package]] +name = "generator" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1be81b31c9ce4caf7f3178eeb90662bb1864d3ad363ec6bdada3201adff19c37" +dependencies = [ + "cc", + "cfg-if", + "libc", + "log", + "rustversion", + "windows-link", + "windows-result", +] + [[package]] name = "generic-array" -version = "0.14.9" +version = "0.14.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4bb6743198531e02858aeaea5398fcc883e71851fcbcb5a2f773e2fb6cb1edf2" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ "typenum", "version_check", @@ -523,13 +1004,29 @@ dependencies = [ [[package]] name = "getrandom" -version = "0.2.15" +version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" dependencies = [ "cfg-if", + "js-sys", "libc", "wasi", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "r-efi 5.3.0", + "wasip2", + "wasm-bindgen", ] [[package]] @@ -539,9 +1036,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" dependencies = [ "cfg-if", + "js-sys", "libc", - "r-efi", + "r-efi 6.0.0", "rand_core 0.10.1", + "wasm-bindgen", +] + +[[package]] +name = "ghash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0d8a4362ccb29cb0b265253fb0a2728f592895ee6854fd9bc13f2ffda266ff1" +dependencies = [ + "opaque-debug", + "polyval", ] [[package]] @@ -574,6 +1083,17 @@ dependencies = [ "tracing", ] +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2 0.2.21", + "equivalent", + "foldhash", +] + [[package]] name = "hashbrown" version = "0.17.1" @@ -586,6 +1106,18 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17592d60ebacc7d5e169f4663c5f84f9161cc90328abcfe8456f41e4dfcb284" + +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + [[package]] name = "hickory-net" version = "0.26.3" @@ -726,9 +1258,9 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" [[package]] name = "hybrid-array" -version = "0.4.8" +version = "0.4.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8655f91cd07f2b9d0c24137bd650fe69617773435ee5ec83022377777ce65ef1" +checksum = "27f864f10dfb56725ce5ce5472bc52252c8f93a4ab86327122cebf62c5f59a17" dependencies = [ "typenum", ] @@ -766,6 +1298,55 @@ dependencies = [ "tokio", ] +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "icu_collator" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24cb0205c419b2036da3f9082964905d6849b3f2130dd4a517940c45899a544b" +dependencies = [ + "icu_collator_data", + "icu_collections", + "icu_locale_core", + "icu_locale_fallback", + "icu_normalizer", + "icu_properties", + "icu_provider", + "smallvec", + "utf16_iter", + "utf8_iter", + "zerovec", +] + +[[package]] +name = "icu_collator_data" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "183f47e821a327661c1ebe152f61d5b6c2000010d228cadbbf226e0e2b8c2634" + [[package]] name = "icu_collections" version = "2.4.0" @@ -780,6 +1361,21 @@ dependencies = [ "zerovec", ] +[[package]] +name = "icu_locale" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe8fa57bcddc15cb921582ee5259c5b77da39f946f188ca5213fad24d9241236" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_locale_data", + "icu_locale_fallback", + "icu_provider", + "tinystr", + "zerovec", +] + [[package]] name = "icu_locale_core" version = "2.4.0" @@ -788,11 +1384,38 @@ checksum = "25d29b57c17c6528d0dc4c2141cfa2d08200465acb537a3909c8b9b8a4a984c0" dependencies = [ "displaydoc", "litemap", + "serde", "tinystr", "writeable", "zerovec", ] +[[package]] +name = "icu_locale_data" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a64014dac577a11d8d2cee48517c73c1984bd72ab11031e934269bd3cefb2ad9" + +[[package]] +name = "icu_locale_fallback" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "406e9ae5b95d8aa37db6293405931b766a6172d1848daee22d70eb44aa78a43c" +dependencies = [ + "icu_locale_core", + "icu_locale_fallback_data", + "icu_provider", + "potential_utf", + "tinystr", + "zerovec", +] + +[[package]] +name = "icu_locale_fallback_data" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3632ebbd0c5c3637ba1a40bdf8288f2cb9d2268e73441a69d5fc73d7cfbaa87c" + [[package]] name = "icu_normalizer" version = "2.4.0" @@ -804,6 +1427,9 @@ dependencies = [ "icu_properties", "icu_provider", "smallvec", + "utf16_iter", + "utf8_iter", + "write16", "zerovec", ] @@ -842,6 +1468,8 @@ checksum = "413099315636c588ee8f011190b329b971a09a8fae9d157e01dcae8379468514" dependencies = [ "displaydoc", "icu_locale_core", + "serde", + "stable_deref_trait", "writeable", "yoke", "zerofrom", @@ -855,29 +1483,58 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" dependencies = [ - "idna_adapter", - "smallvec", - "utf8_iter", + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "generic-array", ] [[package]] -name = "idna_adapter" -version = "1.2.2" +name = "intrusive-collections" +version = "0.9.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +checksum = "189d0897e4cbe8c75efedf3502c18c887b05046e59d28404d4d8e46cbc4d1e86" dependencies = [ - "icu_normalizer", - "icu_properties", + "memoffset", ] [[package]] -name = "indexmap" -version = "2.14.2" +name = "io-uring" +version = "0.7.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +checksum = "ed3bd0ecfbb87805f538bb7b32e5239ca0763890c623e349860ecba69469f2bb" dependencies = [ - "equivalent", - "hashbrown", + "bitflags", + "cfg-if", + "libc", ] [[package]] @@ -904,18 +1561,27 @@ dependencies = [ [[package]] name = "itertools" -version = "0.13.0" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itertools" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +checksum = "8b4baf93f58d4425749ca49a51c50ebab072c5df6994d08fed93541c331481dc" dependencies = [ "either", ] [[package]] name = "itoa" -version = "1.0.14" +version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d75a2a4b1b190afb6f5425f10f6a8f959d2ea0b9c2b1d79553551850539e4674" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "jni" @@ -944,7 +1610,7 @@ dependencies = [ "quote", "rustc_version", "simd_cesu8", - "syn 2.0.96", + "syn 2.0.119", ] [[package]] @@ -963,45 +1629,57 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" dependencies = [ "quote", - "syn 2.0.96", + "syn 2.0.119", +] + +[[package]] +name = "jobserver" +version = "0.1.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" +dependencies = [ + "getrandom 0.4.3", + "libc", ] [[package]] name = "js-sys" -version = "0.3.77" +version = "0.3.107" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1cfaf33c695fc6e08064efbc1f72ec937429614f25eef83af942d0e227c3a28f" +checksum = "8dc09d2154098bbf4ef54eab6b117a3c24cbbffe94046b16ee6d6d4510584d3a" dependencies = [ - "once_cell", + "cfg-if", + "futures-util", "wasm-bindgen", ] [[package]] name = "jsonwebtoken" -version = "10.3.0" +version = "10.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0529410abe238729a60b108898784df8984c87f6054c9c4fcacc47e4803c1ce1" +checksum = "eba32bfb4ffdeaca3e34431072faf01745c9b26d25504aa7a6cf5684334fc4fc" dependencies = [ "base64", "ed25519-dalek", - "getrandom 0.2.15", + "getrandom 0.2.17", "hmac", "js-sys", "p256", "p384", - "rand 0.8.6", + "rand 0.8.8", "rsa", "serde", "serde_json", "sha2 0.10.9", "signature", + "zeroize", ] [[package]] name = "lazy_static" -version = "1.5.0" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" dependencies = [ "spin", ] @@ -1012,12 +1690,48 @@ version = "0.2.191" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a" +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + [[package]] name = "libm" version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" +[[package]] +name = "linkme" +version = "0.3.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3045e122bd98aef8ec3ad58ce84f0791f64e70163d1a02710af4aa11a4d54cc5" +dependencies = [ + "linkme-impl", +] + +[[package]] +name = "linkme-impl" +version = "0.3.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77060ebe535362c3da75682cd17b0431017b6e7c5661e714fc69a7ad017d1301" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "litemap" version = "0.8.3" @@ -1035,24 +1749,97 @@ dependencies = [ [[package]] name = "log" -version = "0.4.25" +version = "0.4.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d15f86d12da44cd8d4ec3c73984537eb7f4cadcb9776a575a0a8a05ca11005c" + +[[package]] +name = "loom" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "419e0dc8046cb947daa77eb95ae174acfbddb7673b4151f56d1eed8e93fbfaca" +dependencies = [ + "cfg-if", + "generator", + "scoped-tls", + "tracing", + "tracing-subscriber", +] + +[[package]] +name = "lz4" +version = "1.28.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a20b523e860d03443e98350ceaac5e71c6ba89aea7d960769ec3ce37f4de5af4" +dependencies = [ + "lz4-sys", +] + +[[package]] +name = "lz4-sys" +version = "1.11.1+lz4-1.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04cbf5b083de1c7e0222a7a51dbfdba1cbe1c6ab0b15e29fff3f6c077fd9cd9f" +checksum = "6bd8c0d6c6ed0cd30b3652886bb8711dc4bb01d637a68105a3d5158039b418e6" +dependencies = [ + "cc", + "libc", +] [[package]] name = "matchers" -version = "0.1.0" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + +[[package]] +name = "mea" +version = "0.6.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8263075bb86c5a1b1427b5ae862e8889656f126e9f77c484496e8b47cf5c5558" +checksum = "c709842c4ce65cb91e2666ad5319dfc1efc3af0d34f02075eddca9000d9f8afb" dependencies = [ - "regex-automata 0.1.10", + "hashbrown 0.17.1", + "slab", ] [[package]] name = "memchr" -version = "2.7.4" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + +[[package]] +name = "miette" +version = "7.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f98efec8807c63c752b5bd61f862c165c115b0a35685bdcfd9238c7aeb592b7" +dependencies = [ + "cfg-if", + "miette-derive", + "unicode-width", +] + +[[package]] +name = "miette-derive" +version = "7.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" +checksum = "db5b29714e950dbb20d5e6f74f9dcec4edbcc1067bb7f8ed198c097b8c1a818b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] [[package]] name = "mio" @@ -1065,6 +1852,16 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "mixtrics" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1dd94ebec15e61a36cd63b1e8816c83434e42926962cd21ee6e4db898d4d4f05" +dependencies = [ + "itertools 0.15.0", + "parking_lot", +] + [[package]] name = "moka" version = "0.12.16" @@ -1090,12 +1887,21 @@ checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b" [[package]] name = "nu-ansi-term" -version = "0.46.0" +version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "overload", - "winapi", + "windows-sys 0.61.2", +] + +[[package]] +name = "num-bigint" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" +dependencies = [ + "num-integer", + "num-traits", ] [[package]] @@ -1109,27 +1915,26 @@ dependencies = [ "num-integer", "num-iter", "num-traits", - "rand 0.8.6", + "rand 0.8.8", "smallvec", "zeroize", ] [[package]] name = "num-integer" -version = "0.1.46" +version = "0.1.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +checksum = "7ce2d95d4b3734dc35aa2f45e1aa22cd416814592a4f9d9205e11affd5b8e10b" dependencies = [ "num-traits", ] [[package]] name = "num-iter" -version = "0.1.45" +version = "0.1.46" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1429034a0490724d0075ebb2bc9e875d6503c3cf69e235a8941aa757d83ef5bf" +checksum = "c92800bd69a1eac91786bcfe9da64a897eb72911b8dc3095decbd07429e8048b" dependencies = [ - "autocfg", "num-integer", "num-traits", ] @@ -1144,6 +1949,16 @@ dependencies = [ "libm", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -1154,6 +1969,12 @@ dependencies = [ "portable-atomic", ] +[[package]] +name = "opaque-debug" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" + [[package]] name = "openssl-probe" version = "0.2.2" @@ -1161,10 +1982,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e036b411754132fd6b74e94b74e9a13c5cf606fa93d1885c35860d3060820b1" [[package]] -name = "overload" -version = "0.1.1" +name = "owo-colors" +version = "3.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" +checksum = "c1b04fb49957986fdce4d6ee7a65027d55d4b6d2265e5848bbb507b58ccfdb6f" [[package]] name = "p256" @@ -1190,6 +2011,15 @@ dependencies = [ "sha2 0.10.9", ] +[[package]] +name = "pack1" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3b7bb0ecf2e447b1f20ee94ee79ef6eed1e9d4b3c36ce1903b9dea3bf205523" +dependencies = [ + "bytemuck", +] + [[package]] name = "parking_lot" version = "0.12.5" @@ -1213,11 +2043,17 @@ dependencies = [ "windows-link", ] +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "pastey" -version = "0.2.1" +version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b867cad97c0791bbd3aaa6472142568c6c9e8f71937e98379f584cfb0cf35bec" +checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4" [[package]] name = "pem-rfc7468" @@ -1234,11 +2070,31 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "pin-project-lite" -version = "0.2.16" +version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b3cff922bd51709b605d9ead9aa71031d81447142d828eb4a6eba76fe619f9b" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" [[package]] name = "pkcs1" @@ -1261,11 +2117,43 @@ dependencies = [ "spki", ] +[[package]] +name = "pkg-config" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548" + +[[package]] +name = "polling" +version = "3.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d0e4f59085d47d8241c88ead0f274e8a0cb551f3625263c05eb8dd897c34218" +dependencies = [ + "cfg-if", + "concurrent-queue", + "hermit-abi", + "pin-project-lite", + "rustix", + "windows-sys 0.61.2", +] + +[[package]] +name = "polyval" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1fe60d06143b2430aa532c94cfe9e29783047f06c0d7fd359a9a51b729fa25" +dependencies = [ + "cfg-if", + "cpufeatures 0.2.17", + "opaque-debug", + "universal-hash", +] + [[package]] name = "portable-atomic" -version = "1.10.0" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "280dc24453071f1b63954171985a0b0d30058d287960968b9b2aca264c8d4ee6" +checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85" [[package]] name = "potential_utf" @@ -1273,6 +2161,8 @@ version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d83eb9bc6d8e5cf568e7a1101d60ee05e81ed50ea106026f3d18deeb046d7661" dependencies = [ + "serde_core", + "writeable", "zerovec", ] @@ -1307,18 +2197,18 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.93" +version = "1.0.107" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60946a68e5f9d28b0dc1c21bb8a97ee7d018a8b322fa57838ba31cc878e22d99" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" dependencies = [ "unicode-ident", ] [[package]] name = "prost" -version = "0.14.3" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2ea70524a2f82d518bce41317d0fae74151505651af45faf1ffbd6fd33f0568" +checksum = "528ac67416ff8646872a3c02cad9cc4ee5dc9f9540c9b10771855c95cb2e5ae1" dependencies = [ "bytes", "prost-derive", @@ -1326,15 +2216,15 @@ dependencies = [ [[package]] name = "prost-derive" -version = "0.14.3" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" +checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", - "itertools", + "itertools 0.14.0", "proc-macro2", "quote", - "syn 2.0.96", + "syn 2.0.119", ] [[package]] @@ -1417,23 +2307,45 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + [[package]] name = "r-efi" version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "radium" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" + [[package]] name = "rand" -version = "0.8.6" +version = "0.8.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +checksum = "e058c7de0b26af77780c769414d6257830bb240f3c38477dbc2c16e5f54d6d4c" dependencies = [ "libc", - "rand_chacha", + "rand_chacha 0.3.1", "rand_core 0.6.4", ] +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + [[package]] name = "rand" version = "0.10.3" @@ -1455,13 +2367,32 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + [[package]] name = "rand_core" version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" dependencies = [ - "getrandom 0.2.15", + "getrandom 0.2.17", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", ] [[package]] @@ -1471,57 +2402,60 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" [[package]] -name = "redox_syscall" -version = "0.5.18" +name = "rand_pcg" +version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +checksum = "59cad018caf63deb318e5a4586d99a24424a364f40f1e5778c29aca23f4fc73e" dependencies = [ - "bitflags", + "rand_core 0.6.4", ] [[package]] -name = "regex" -version = "1.11.1" +name = "rapidhash" +version = "4.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b544ef1b4eac5dc2db33ea63606ae9ffcfac26c1416a2806ae0bf5f56b201191" +checksum = "5da7e78a036ce858e8d55b7e7dc8ba3a88b78350fd2155d3591bbd966b58589e" dependencies = [ - "aho-corasick", - "memchr", - "regex-automata 0.4.9", - "regex-syntax 0.8.5", + "rustversion", ] [[package]] -name = "regex-automata" -version = "0.1.10" +name = "redox_syscall" +version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c230d73fb8d8c1b9c0b3135c5142a8acee3a0558fb8db5cf1cb65f8d7862132" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "regex-syntax 0.6.29", + "bitflags", ] [[package]] -name = "regex-automata" -version = "0.4.9" +name = "regex" +version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "809e8dc61f6de73b46c85f4c96486310fe304c434cfa43669d7b40f711150908" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" dependencies = [ "aho-corasick", "memchr", - "regex-syntax 0.8.5", + "regex-automata", + "regex-syntax", ] [[package]] -name = "regex-syntax" -version = "0.6.29" +name = "regex-automata" +version = "0.4.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f162c6dd7b008981e4d40210aca20b4bd0f9b60ca9271061b07f78537722f2e1" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] [[package]] name = "regex-syntax" -version = "0.8.5" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "resolv-conf" @@ -1559,6 +2493,12 @@ dependencies = [ "bs58", "bytes", "bytes-utils", + "foyer", + "foyer-common", + "foyer-memory", + "foyer-storage", + "foyer-tokio", + "futures", "futures-util", "hickory-resolver", "hostname", @@ -1568,21 +2508,24 @@ dependencies = [ "hyper", "hyper-util", "jsonwebtoken", + "parking_lot", "pastey", "prost", - "rand 0.8.6", + "rand 0.8.8", "rustls", "rustls-native-certs", "serde", - "sha2 0.11.0", + "sha2 0.11.1", "socket2", - "strum", + "strum 0.27.2", "thiserror", "tokio", "tokio-rustls", "tokio-util", "tracing", + "turso_core", "uuid", + "zstd-safe", ] [[package]] @@ -1603,12 +2546,22 @@ checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" dependencies = [ "cc", "cfg-if", - "getrandom 0.2.15", + "getrandom 0.2.17", "libc", "untrusted", "windows-sys 0.52.0", ] +[[package]] +name = "roaring" +version = "0.11.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18bd8a37d17a58532776dcdf6041ce64929adca78e8489d5cacbafe99229d3e1" +dependencies = [ + "bytemuck", + "byteorder", +] + [[package]] name = "rsa" version = "0.9.10" @@ -1630,14 +2583,43 @@ dependencies = [ ] [[package]] -name = "rustc_version" -version = "0.4.1" +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + +[[package]] +name = "rustc_version_runtime" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +checksum = "2dd18cd2bae1820af0b6ad5e54f4a51d0f3fcc53b05f845675074efcc7af071d" dependencies = [ + "rustc_version", "semver", ] +[[package]] +name = "rustix" +version = "1.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5e76742e5294153d4dde1207dbec5bdd81eef6fabdc0d2118babadabfd858aa" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + [[package]] name = "rustls" version = "0.23.46" @@ -1686,15 +2668,15 @@ dependencies = [ [[package]] name = "rustversion" -version = "1.0.19" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7c45b9784283f1b2e7fb61b42047c2fd678ef0960d4f6f1eba131594cc369d4" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" [[package]] name = "ryu" -version = "1.0.18" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" [[package]] name = "same-file" @@ -1714,6 +2696,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "scoped-tls" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1cf6437eb19a8f4a6cc0f7dca544973b0b78843adbfeb3683d1a94a0024a294" + [[package]] name = "scopeguard" version = "1.2.0" @@ -1765,9 +2753,9 @@ checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" [[package]] name = "serde" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" dependencies = [ "serde_core", "serde_derive", @@ -1775,36 +2763,43 @@ dependencies = [ [[package]] name = "serde_core" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn 2.0.96", + "syn 3.0.9", ] [[package]] name = "serde_json" -version = "1.0.135" +version = "1.0.154" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b0d7ba2887406110130a978386c4e1befb98c674b4fba677954e4db976630d9" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" dependencies = [ "itoa", "memchr", - "ryu", "serde", + "serde_core", + "zmij", ] +[[package]] +name = "sha1_smol" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d" + [[package]] name = "sha2" version = "0.10.9" @@ -1818,13 +2813,13 @@ dependencies = [ [[package]] name = "sha2" -version = "0.11.0" +version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" +checksum = "47d7069beb7d6ac7b9acd1039986e73443f24234f41074da099d6f994ac9ad19" dependencies = [ "cfg-if", - "cpufeatures 0.3.0", - "digest 0.11.2", + "cpufeatures 0.3.1", + "digest 0.11.3", ] [[package]] @@ -1842,6 +2837,26 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" +[[package]] +name = "shuttle" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ab17edba38d63047f46780cf7360acf7467fec2c048928689a5c1dd1c2b4e31" +dependencies = [ + "assoc", + "bitvec", + "cfg-if", + "generator", + "hex", + "owo-colors", + "rand 0.8.8", + "rand_core 0.6.4", + "rand_pcg", + "scoped-tls", + "smallvec", + "tracing", +] + [[package]] name = "signature" version = "2.2.0" @@ -1868,17 +2883,29 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" +[[package]] +name = "siphasher" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33f4fe9184a62d842c9ef383018f3306d8ba224fd9d836f56d7288308847c256" + [[package]] name = "slab" version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" +[[package]] +name = "small_ctor" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88414a5ca1f85d82cc34471e975f0f74f6aa54c40f062efa42c0080e7f763f81" + [[package]] name = "smallvec" -version = "1.13.2" +version = "1.16.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" +checksum = "d0d5d4b1981b8bbbb7aed204cae5726d683d5d6ac7e74ec54122ea28f315e7c8" [[package]] name = "socket2" @@ -1890,11 +2917,17 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "softaes" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4fdfe13e1c42e45f497416afb8dc4e783bf0e25b7be9077c1c77c5fb14cb998" + [[package]] name = "spin" -version = "0.9.8" +version = "0.9.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" +checksum = "3763264f6b73151db08c50ff20d7d8a0b8796e021cdea7ceedad07b80155fa0e" [[package]] name = "spki" @@ -1914,24 +2947,45 @@ checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" [[package]] name = "strum" -version = "0.27.1" +version = "0.26.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fec0f0aef304996cf250b31b5a10dee7980c85da9d759361292b8bca5a18f06" +dependencies = [ + "strum_macros 0.26.4", +] + +[[package]] +name = "strum" +version = "0.27.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f64def088c51c9510a8579e3c5d67c65349dcf755e5479ad3d010aa6454e2c32" +checksum = "af23d6f6c1a224baef9d3f61e287d2761385a5b88fdab4eb4c6f11aeb54c4bcf" dependencies = [ - "strum_macros", + "strum_macros 0.27.2", ] [[package]] name = "strum_macros" -version = "0.27.1" +version = "0.26.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c77a8c5abcaf0f9ce05d62342b7d298c346515365c36b673df4ebe3ced01fde8" +checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" dependencies = [ "heck", "proc-macro2", "quote", "rustversion", - "syn 2.0.96", + "syn 2.0.119", +] + +[[package]] +name = "strum_macros" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7695ce3845ea4b33927c055a39dc438a45b059f7c1b3d91d38d10355fb8cbca7" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.119", ] [[package]] @@ -1942,9 +2996,9 @@ checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" [[package]] name = "syn" -version = "2.0.96" +version = "2.0.119" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d5d0adab1ae378d7f53bdebc67a39f1f151407ef230f0ce2883572f5d8985c80" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" dependencies = [ "proc-macro2", "quote", @@ -2000,40 +3054,58 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" +[[package]] +name = "tap" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" + [[package]] name = "target-lexicon" version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + [[package]] name = "thiserror" -version = "2.0.11" +version = "2.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d452f284b73e6d76dd36758a0c8684b1d5be31f92b89d07fd5822175732206fc" +checksum = "09e52cb86a36cede5cb101bf8908837b3e4c6e5e59fe7fd85c23fb56200d189e" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "2.0.11" +version = "2.0.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26afc1baea8a989337eeb52b6e72a039780ce45c3edfcc9c5b9d112feeb173c2" +checksum = "fe5197923287db20a58125f0bc85c062f7f2c892de97b18c356f9efb14b28524" dependencies = [ "proc-macro2", "quote", - "syn 2.0.96", + "syn 3.0.9", ] [[package]] name = "thread_local" -version = "1.1.8" +version = "1.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b9ef9bad013ada3808854ceac7b46812a6465ba368859a37e2100283d2d719c" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" dependencies = [ "cfg-if", - "once_cell", ] [[package]] @@ -2043,23 +3115,15 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "abc840c238154d58f9e86b2792261edfde565a17696e1cda1a449de27979801a" dependencies = [ "displaydoc", + "serde_core", "zerovec", ] [[package]] name = "tinyvec" -version = "1.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "022db8904dfa342efe721985167e9fcd16c29b226db4397ed752a761cfce81e8" -dependencies = [ - "tinyvec_macros", -] - -[[package]] -name = "tinyvec_macros" -version = "0.1.1" +version = "1.13.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +checksum = "fd3ca314f692efd6c868f8408f53fe444634a845f96c028b97d35f6a1f79f0ee" [[package]] name = "tokio" @@ -2125,9 +3189,9 @@ dependencies = [ [[package]] name = "tracing" -version = "0.1.41" +version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "784e0ac535deb450455cbfa28a6f0df145ea1bb7ae51b821cf5e7927fdcfbdd0" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ "pin-project-lite", "tracing-attributes", @@ -2136,20 +3200,20 @@ dependencies = [ [[package]] name = "tracing-attributes" -version = "0.1.28" +version = "0.1.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "395ae124c09f9e6918a2310af6038fba074bcf474ac352496d5910dd59a2226d" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn 2.0.96", + "syn 2.0.119", ] [[package]] name = "tracing-core" -version = "0.1.33" +version = "0.1.36" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e672c95779cf947c5311f83787af4fa8fffd12fb27e4993211a84bdfd9610f9c" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", "valuable", @@ -2168,14 +3232,14 @@ dependencies = [ [[package]] name = "tracing-subscriber" -version = "0.3.19" +version = "0.3.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e8189decb5ac0fa7bc8b96b7cb9b2701d60d48805aca84a238004d665fcc4008" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" dependencies = [ "matchers", "nu-ansi-term", "once_cell", - "regex", + "regex-automata", "sharded-slab", "smallvec", "thread_local", @@ -2184,17 +3248,159 @@ dependencies = [ "tracing-log", ] +[[package]] +name = "turso_core" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fa832d4ab648bf139f1b18b29bef6d653b18fbd4a77b0e4846c57aaf28c215c" +dependencies = [ + "aegis", + "aes", + "aes-gcm", + "allocator-api2 0.4.0", + "antithesis_sdk", + "arc-swap", + "aristo", + "bigdecimal", + "bitflags", + "branches", + "bumpalo", + "bytemuck", + "cfg_aliases", + "cfg_block", + "chrono", + "crc32c", + "crossbeam-epoch", + "crossbeam-utils", + "either", + "fallible-iterator", + "fastbloom", + "getrandom 0.2.17", + "getrandom 0.3.4", + "hex", + "icu_collator", + "icu_locale", + "intrusive-collections", + "libc", + "libloading", + "libm", + "loom", + "miette", + "num-bigint", + "num-traits", + "pack1", + "parking_lot", + "pastey", + "polling", + "rand 0.9.5", + "rapidhash", + "regex", + "regex-syntax", + "roaring", + "rustc-hash", + "rustix", + "ryu", + "serde_json", + "shuttle", + "simdutf8", + "smallvec", + "strum 0.26.3", + "strum_macros 0.26.4", + "tempfile", + "thiserror", + "tracing", + "tracing-subscriber", + "turso_ext", + "turso_macros", + "turso_parser", + "twox-hash", + "uncased", + "uuid", + "windows-sys 0.61.2", +] + +[[package]] +name = "turso_ext" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5003ed7e98f082fe4114627c8d7c7ab569c27766e9bd350a0c09af06ad3fd46" +dependencies = [ + "chrono", + "getrandom 0.4.3", + "turso_macros", +] + +[[package]] +name = "turso_macros" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74770c0d0c8b272b479b5343f19dbd2f8f2677d85a4d6222fc214dcac2882beb" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "turso_parser" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a69bb856ec8306ad7ed1c13fc61de6a79452bf646ca53d99bb4cc646d5bfc4ff" +dependencies = [ + "bitflags", + "memchr", + "miette", + "strum 0.26.3", + "strum_macros 0.26.4", + "thiserror", + "turso_macros", +] + +[[package]] +name = "twox-hash" +version = "2.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "86a801b3cea342a06d468c8710662aa29e5e05e4f5c0d62f00bbb7f2ad7941c2" +dependencies = [ + "rand 0.10.3", +] + [[package]] name = "typenum" -version = "1.17.0" +version = "1.20.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "uncased" +version = "0.9.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1b88fcfe09e89d3866a5c11019378088af2d24c3fbd4f0543f96b479ec90697" +dependencies = [ + "version_check", +] [[package]] name = "unicode-ident" -version = "1.0.14" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563" + +[[package]] +name = "unicode-width" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" + +[[package]] +name = "universal-hash" +version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83" +checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea" +dependencies = [ + "crypto-common 0.1.7", + "subtle", +] [[package]] name = "untrusted" @@ -2214,6 +3420,15 @@ dependencies = [ "serde", ] +[[package]] +name = "utf16_iter" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c7de35d9d8757dbc1f5125dc8b51438c68b4430546a735c3a9b91aaaaf9a18a" +dependencies = [ + "utf_types", +] + [[package]] name = "utf8_iter" version = "1.1.0" @@ -2237,14 +3452,15 @@ checksum = "4a0e2d11620c50695fe12bbbfde96b3477fe606e7284c22d505c12ce51214514" dependencies = [ "getrandom 0.4.3", "js-sys", + "sha1_smol", "wasm-bindgen", ] [[package]] name = "valuable" -version = "0.1.0" +version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" [[package]] name = "version_check" @@ -2264,40 +3480,36 @@ dependencies = [ [[package]] name = "wasi" -version = "0.11.0+wasi-snapshot-preview1" +version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" [[package]] -name = "wasm-bindgen" -version = "0.2.100" +name = "wasip2" +version = "1.0.4+wasi-0.2.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1edc8929d7499fc4e8f0be2262a241556cfc54a0bea223790e71446f2aab1ef5" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" dependencies = [ - "cfg-if", - "once_cell", - "wasm-bindgen-macro", + "wit-bindgen", ] [[package]] -name = "wasm-bindgen-backend" -version = "0.2.100" +name = "wasm-bindgen" +version = "0.2.130" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f0a0651a5c2bc21487bde11ee802ccaf4c51935d0d3d42a6101f98161700bc6" +checksum = "a0181c14ddd6c7b8b0b0132a5f1ea2f32be0d893cd3ea97b6e3b28886b8a3d1c" dependencies = [ - "bumpalo", - "log", - "proc-macro2", - "quote", - "syn 2.0.96", + "cfg-if", + "rustversion", + "wasm-bindgen-macro", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-macro" -version = "0.2.100" +version = "0.2.130" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7fe63fc6d09ed3792bd0897b314f53de8e16568c2b3f7982f468c0bf9bd0b407" +checksum = "6e931a02a39abdbf66966fc35902b5ed58344750a70885bed4ebb3d8eaf28ed9" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -2305,26 +3517,36 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.100" +version = "0.2.130" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ae87ea40c9f689fc23f209965b6fb8a99ad69aeeb0231408be24920604395de" +checksum = "2ca4b54280e455bd831e33a6b8e734baf432fe63fbc35c42ac14ce67a043a53a" dependencies = [ + "bumpalo", "proc-macro2", "quote", - "syn 2.0.96", - "wasm-bindgen-backend", + "syn 3.0.9", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-shared" -version = "0.2.100" +version = "0.2.130" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a05d73b933a847d6cccdda8f838a22ff101ad9bf93e33684f39c1f5f0eece3d" +checksum = "3b1f9c3798725016c6ee5f2f57869ec0831e8bbc0ab0310a86e787e97f41389e" dependencies = [ "unicode-ident", ] +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "widestring" version = "1.2.1" @@ -2362,6 +3584,41 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "windows-link" version = "0.2.1" @@ -2406,6 +3663,15 @@ dependencies = [ "windows-targets", ] +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + [[package]] name = "windows-sys" version = "0.61.2" @@ -2479,12 +3745,33 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "write16" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1890f4022759daae28ed4fe62859b1236caebfc61ede2f63ed4e695f3f6d936" + [[package]] name = "writeable" version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3ad82d2a33cdc9674dc7465672f271e096168fcdbe0f799d9e6db8c5892679dc" +[[package]] +name = "wyz" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05f360fc0b24296329c78fda852a1e9ae82de9cf7b27dae4b7f62f118f77b9ed" +dependencies = [ + "tap", +] + [[package]] name = "yoke" version = "0.8.4" @@ -2510,22 +3797,22 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.27" +version = "0.8.63" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0894878a5fa3edfd6da3f88c4805f4c8558e2b996227a3d864f47fe11e38282c" +checksum = "e5fe1f8f1b06191a00962174c61aa5005e0bb391a6d80d07e24d115c01a92ed8" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.27" +version = "0.8.63" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88d2b8d9c68ad2b9e4340d7832716a4d21a22a1154777ad56ea55c51a9cf3831" +checksum = "863ad3ac83293fb4d740aedbfdc9240dd8d1a50c1099acd76ce80ce7c7230c7f" dependencies = [ "proc-macro2", "quote", - "syn 2.0.96", + "syn 2.0.119", ] [[package]] @@ -2551,9 +3838,23 @@ dependencies = [ [[package]] name = "zeroize" -version = "1.9.0" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13084392c5e4bc371903e2935a5eaeed24905a7511356b883835e18a78f6879" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" +checksum = "3c50655cbb0fe3fc43170059e702f1ce5e19b84cec58dc87b037a09935c2f328" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] [[package]] name = "zerotrie" @@ -2564,6 +3865,7 @@ dependencies = [ "displaydoc", "yoke", "zerofrom", + "zerovec", ] [[package]] @@ -2572,6 +3874,7 @@ version = "0.11.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f58a2799cfe9a1670c7ea8edbd52ef8edb93cb37a7f83bf31ba9e35d9dae6a40" dependencies = [ + "serde", "yoke", "zerofrom", "zerovec-derive", @@ -2587,3 +3890,37 @@ dependencies = [ "quote", "syn 3.0.9", ] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" + +[[package]] +name = "zstd" +version = "0.13.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdb02c3132d251ab1606e9fc913a7da6c6ed08cc263df15edf082e1dc497d157" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c8ea04668f5de8d8d51051714ba560c9536519cadcdb8829bba4fe2ad159b03" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.1.2+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "481ada9c543b2323cbb7bbdbb5249177fbf0188ed959ed2c15e4b22421c84ed3" +dependencies = [ + "cc", + "pkg-config", +] diff --git a/Cargo.toml b/Cargo.toml index 482c487d..6051ed6b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,15 @@ pyo3 = { version = "0.30.1", features = ["extension-module"] } pyo3-async-runtimes = { version = "0.30", features = ["tokio-runtime"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] } -restate-sdk-shared-core = { version = "8.3.0", features = ["request_identity", "sha2_random_seed", "rust_crypto", "tunnel"] } +# The feature `sql` is the sessions' storage: the SQLite-compatible database kept in an +# object's state, with its page cache (foyer); `tunnel` is Restate Cloud's. +restate-sdk-shared-core = { version = "8.3.0", features = [ + "request_identity", + "sha2_random_seed", + "rust_crypto", + "tunnel", + "sql", +] } # ASGI bridge for the tunnel bytes = "1.10" From a69a15200fa3a12fea84d4b78248769221db11c2 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 14:22:10 +0100 Subject: [PATCH 02/45] Add the restate.sessions and restate.agents packages, empty The Python variants of @restatedev/sessions and @restatedev/agents will live here; for now, each is a package with its docstring. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/agents/__init__.py | 14 ++++++++++++++ python/restate/sessions/__init__.py | 14 ++++++++++++++ 2 files changed, 28 insertions(+) create mode 100644 python/restate/agents/__init__.py create mode 100644 python/restate/sessions/__init__.py diff --git a/python/restate/agents/__init__.py b/python/restate/agents/__init__.py new file mode 100644 index 00000000..ef1b198d --- /dev/null +++ b/python/restate/agents/__init__.py @@ -0,0 +1,14 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +AI agents in session objects, each conversation in its session's database. +The Python variant of ``@restatedev/agents``. +""" diff --git a/python/restate/sessions/__init__.py b/python/restate/sessions/__init__.py new file mode 100644 index 00000000..7a6a8b62 --- /dev/null +++ b/python/restate/sessions/__init__.py @@ -0,0 +1,14 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +Session objects: virtual objects with a SQL database per key, kept in the +object's state (``ctx.db``). The Python variant of ``@restatedev/sessions``. +""" From 3dfdd4c0449cfc4562c59d7ace097718fb98cf02 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 14:37:05 +0100 Subject: [PATCH 03/45] Bind shared-core's sql::Database to Python, for restate.sessions The native side of sessions, in restate._internal (src/sessions.rs): an object's database over a key/value store implemented in Python, as sdk-sessions' node.rs binds it to JavaScript. Every request to the store is scheduled on the asyncio loop with call_soon_threadsafe, in issue order; a call that needs nothing from the store answers at once, one that waits goes on in tokio (pyo3-async-runtimes). The process-wide page cache reads the same RESTATE_SQL_CACHE* variables as the Node addon. restate.sessions._database orders the calls and runs transactions, as database.ts does: records committed with a transaction, with_commit (what runs between a commit's writes and its sync), SqlError for what a retry would repeat, and the store's own error for a session it ended. Ported from sqlexperiment's restate_sql, onto shared-core's API (records in place of effects). Nothing uses it yet. uv now rebuilds the extension when the Rust changes. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- Cargo.lock | 1 + Cargo.toml | 3 + pyproject.toml | 9 + python/restate/sessions/_database.py | 470 +++++++++++++++++++++++ src/lib.rs | 4 + src/sessions.rs | 493 ++++++++++++++++++++++++ tests/sessions_database.py | 551 +++++++++++++++++++++++++++ 7 files changed, 1531 insertions(+) create mode 100644 python/restate/sessions/_database.py create mode 100644 src/sessions.rs create mode 100644 tests/sessions_database.py diff --git a/Cargo.lock b/Cargo.lock index e5d96bbf..5d247e90 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2468,6 +2468,7 @@ name = "restate-sdk-python-core" version = "1.3.0" dependencies = [ "bytes", + "futures", "http", "http-body", "http-body-util", diff --git a/Cargo.toml b/Cargo.toml index 6051ed6b..e7c62566 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,3 +36,6 @@ percent-encoding = "2.3" tokio = { version = "1.49", default-features = false, features = ["sync", "rt-multi-thread"] } tokio-stream = { version = "0.1", default-features = false } tokio-util = { version = "0.7", default-features = false } + +# Sessions (`src/sessions.rs`): the database's futures, polled on Python's event loop. +futures = "0.3" diff --git a/pyproject.toml b/pyproject.toml index 55be4303..1970cec1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,15 @@ features = ["pyo3/extension-module"] module-name = "restate._internal" python-source = "python" +# Rebuilt by `uv sync` when the Rust changes, not only the Python. +[tool.uv] +cache-keys = [ + { file = "pyproject.toml" }, + { file = "Cargo.toml" }, + { file = "Cargo.lock" }, + { file = "src/**/*.rs" }, +] + [tool.ruff] line-length = 120 target-version = "py310" diff --git a/python/restate/sessions/_database.py b/python/restate/sessions/_database.py new file mode 100644 index 00000000..047e2f8c --- /dev/null +++ b/python/restate/sessions/_database.py @@ -0,0 +1,470 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +A session's database, through the native extension (``src/sessions.rs``): +SQLite-compatible SQL, kept in the object's state. The extension reads and +writes the state through a store, in the order it makes its requests; this +module orders the calls and runs transactions:: + + db = await open("Ledger/a", store) + async with db.transaction() as tx: + await tx.execute("INSERT INTO transfers VALUES(?, ?, ?, ?)", [id, src, dst, amount]) + await db.save() # ends the session, leaving its engine to the next one +""" + +from __future__ import annotations + +import asyncio +import contextvars +import inspect +import typing +from collections.abc import AsyncIterator, Awaitable, Callable, Sequence +from contextlib import asynccontextmanager +from typing import Any, Protocol + +from restate._internal import ( # pylint: disable=import-error,no-name-in-module + SQL_CLEAR, + SQL_GET, + SQL_KEYS, + SQL_SET, + SQL_SYNC, + SqlRefusedException, + sql_open, +) + +SqlValue = typing.Union[None, int, float, str, bytes] +"""An SQL value: ``None``, ``int`` (64 bits), ``float``, ``str``, or ``bytes`` (BLOB).""" + +SqlParam = typing.Union[SqlValue, bool, bytearray, memoryview] +"""An SQL parameter: an ``SqlValue``, a ``bool`` (stored as 0 or 1), or a BLOB as any bytes-like object.""" + +Row = dict[str, SqlValue] +"""A row, by column name.""" + +Rows = list[list[SqlValue]] +"""Rows as lists, in column order.""" + + +class SqlError(Exception): + """ + A failure a retry would repeat: SQL the engine rejected (a syntax error, a + constraint, a missing table), which undid just that statement and leaves + the session usable; a parameter it cannot bind; or a call the session + refused, which may have ended it (a ``COMMIT`` in a transaction's body, a + transaction left open). A storage error, which ends the session, raises + as it is. + """ + + +class Store(Protocol): + """ + A key/value store, the object's state. + + ``set`` and ``clear`` return at once, and ``sync`` completes once every + write before it is stored. Writes are stored in the order they were made: + a crash keeps a prefix of the writes since the last sync and loses the + rest. ``get`` issues its read when it is called, and returns the stored + value or this request's own latest write (``None`` when absent), or an + awaitable of it; so a plain ``async def get`` does not do, as its read + would be issued only when it first runs, after writes made later. + + A store may also have ``keys()``: every key, as ``ctx.state_keys()``, or an + awaitable of them. Then a commit that checkpoints deletes the packs no + checkpoint needs any more; a store without it keeps them. + """ + + def get(self, key: str) -> bytes | None | Awaitable[bytes | None]: ... + def set(self, key: str, value: bytes) -> None: ... + def clear(self, key: str) -> None: ... + def sync(self) -> None | Awaitable[None]: ... + + +async def open(namespace: str, store: Store, *, reader: bool = False) -> Connection: # pylint: disable=redefined-builtin + """ + Open the database of ``namespace`` over ``store``. One process-wide page + cache serves every namespace. A ``reader`` reads beside the namespace's + owner (a shared handler): the state as of open, whatever the owner commits + meanwhile, until the owner's second checkpoint after it; it writes nothing. + """ + boundary = _Boundary(store, asyncio.get_running_loop()) + try: + native = await sql_open(namespace, boundary.dispatch, boundary.loop, boundary.lists_keys, reader) + except Exception as error: + if boundary.failed is not None: + raise boundary.failed from error + raise + return Connection(native, boundary) + + +class _Boundary: + """ + The store as the extension calls it, through ``dispatch`` on ``loop``. An + error of the store's crosses the extension as text, which ends the + session: the first is kept here, and the calls that fail on the session it + ended raise it as it was raised (the SDK's suspension, say). A failed write + cannot fail its caller, which has moved on: the next ``sync`` fails. And a + commit's sync first runs what its transaction gave + ``ConnectionTransaction.with_commit``: after the commit's writes, before + anything is awaited. + """ + + def __init__(self, store: Store, loop: asyncio.AbstractEventLoop) -> None: + self.store = store + self.loop = loop + self.lists_keys = callable(getattr(store, "keys", None)) + # The first error of the store's. + self.failed: BaseException | None = None + # What the next sync runs first: the committing transaction's. + self.with_sync: Callable[[], None] | None = None + # The first write that failed: every sync after it fails. + self._write_failed: BaseException | None = None + + def _keep(self, error: BaseException) -> None: + if self.failed is None: + self.failed = error + + def dispatch(self, op: int, key: str, value: bytes | None, reply: Any) -> None: + """One request of the extension's, on the event loop, in the order it made them.""" + if op in (SQL_SET, SQL_CLEAR): + try: + if op == SQL_SET: + self.store.set(key, typing.cast(bytes, value)) + else: + self.store.clear(key) + except BaseException as error: # pylint: disable=broad-exception-caught + self._keep(error) + if self._write_failed is None: + self._write_failed = error + return + try: + if op == SQL_GET: + result: Any = self.store.get(key) + elif op == SQL_KEYS: + result = self.store.keys() # type: ignore[attr-defined] + elif self._write_failed is not None: + reply(None, str(self._write_failed)) + return + else: + run, self.with_sync = self.with_sync, None + if run is not None: + run() + result = self.store.sync() + except BaseException as error: # pylint: disable=broad-exception-caught + self._keep(error) + reply(None, str(error)) + return + if inspect.isawaitable(result): + self._answer(op, reply, result) + else: + reply(_reply_value(op, result)) + + def _answer(self, op: int, reply: Any, awaitable: Awaitable[Any]) -> None: + def done(future: asyncio.Future) -> None: + if future.cancelled(): + self._keep(asyncio.CancelledError()) + reply(None, "the store's request was cancelled") + elif future.exception() is not None: + error = typing.cast(BaseException, future.exception()) + self._keep(error) + reply(None, str(error)) + else: + reply(_reply_value(op, future.result())) + + asyncio.ensure_future(awaitable, loop=self.loop).add_done_callback(done) + + +def _reply_value(op: int, result: Any) -> Any: + """A request's answer as the extension reads it: a value, the keys as a list, or nothing for a sync.""" + if op == SQL_SYNC: + return None + if op == SQL_KEYS: + return list(result) + return result + + +def _objects(columns: list[str], rows: Rows) -> list[Row]: + """The rows as dicts by column name; a column name twice is an ``SqlError``.""" + if len(set(columns)) != len(columns): + twice = next(name for i, name in enumerate(columns) if name in columns[:i]) + raise SqlError(f"the column {twice} is named twice: alias one (AS), or read the rows with values()") + return [dict(zip(columns, row)) for row in rows] + + +class _Sql: + """SQL on a session of the extension: ``query``, ``get``, ``values`` and ``execute`` by way of ``query_columns``.""" + + def __init__(self, native: Any, boundary: _Boundary) -> None: + self._native = native + self._boundary = boundary + + async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: + """The rows of one statement, with its column names: ``(columns, rows)``.""" + raise NotImplementedError + + async def query(self, sql: str, params: Sequence[SqlParam] | None = None) -> list[Row]: + """ + The rows of one statement, each a dict by column name; parameters bind + to a single statement. Columns are named as SQL names them (an + expression needs ``AS``), and two of the same name are an + ``SqlError``: alias them, or use ``values``. + """ + return _objects(*await self.query_columns(sql, params)) + + async def get(self, sql: str, params: Sequence[SqlParam] | None = None) -> Row | None: + """The first row of ``query``, or ``None`` when there is none.""" + rows = await self.query(sql, params) + return rows[0] if rows else None + + async def values(self, sql: str, params: Sequence[SqlParam] | None = None) -> Rows: + """The rows of one statement as lists, in column order: for adapters, and for columns that share a name.""" + return (await self.query_columns(sql, params))[1] + + async def execute(self, sql: str, params: Sequence[SqlParam] | None = None) -> None: + """Run SQL (one or more statements), discarding rows; parameters only with a single statement.""" + await self.query_columns(sql, params) + + async def _request(self, call: Callable[[], Any]) -> Any: + """ + ``call``'s result, or its ``_failure``: ``call`` makes the request, as + the extension raises as it is called (arguments it cannot convert, and + the error of a call that answers at once), and returns a value or an + awaitable of it. + """ + try: + result = call() + if inspect.isawaitable(result): + result = await result + return result + except Exception as error: + raise self._failure(error) from None + + def _failure(self, error: Exception) -> BaseException: + """ + A request's ``error``: an ``SqlError`` if a retry would repeat it; the + store's own error if that ended the session. + """ + if isinstance(error, SqlError): + return error + failed = self._boundary.failed + if failed is not None and not self._native.is_usable(): + return failed + refused = isinstance(error, (SqlRefusedException, TypeError, ValueError)) + if refused or self._native.is_usable(): + sql_error = SqlError(str(error)) + sql_error.__cause__ = error + return sql_error + return error + + +# The transaction whose body is running, in that body's context (and the tasks +# it starts): the innermost, which knows the ones it runs inside. +_running: contextvars.ContextVar[ConnectionTransaction | None] = contextvars.ContextVar( + "restate_sessions_running_transaction", default=None +) + + +class Connection(_Sql): + """ + A session of a database: calls run one at a time, in call order, and a + transaction holds the session until it commits or rolls back. Calls made + meanwhile wait, but a call from inside the transaction's body is refused, + as it would wait for that body: the body uses its ``tx``. + """ + + def __init__(self, native: Any, boundary: _Boundary) -> None: + super().__init__(native, boundary) + # FIFO: calls run in the order they were made. + self._lock = asyncio.Lock() + + @asynccontextmanager + async def _serial(self) -> AsyncIterator[None]: + """Run the body once the calls before it have, unless it comes from the body of a transaction it would wait for.""" + running = _running.get() + if running is not None and running.holds(self): + raise SqlError("inside a transaction, use the transaction handle its body gets") + async with self._lock: + yield + + async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: + async with self._serial(): + return await self._request(lambda: self._native.query_columns(sql, params)) + + @asynccontextmanager + async def transaction(self) -> AsyncIterator[ConnectionTransaction]: + """ + One SQL transaction, as the body of an ``async with``:: + + async with db.transaction() as tx: + await tx.execute("UPDATE stock SET n = n - 1 WHERE id = ?", [7]) + + It commits when the body ends, with the records it set, and rolls back + when the body raises, with the body's own error, once the calls it made + have completed. A statement in it that ends the transaction + (``COMMIT``, ``ROLLBACK``) ends the session. A failing COMMIT rolls + back. + """ + async with self._serial(): + await self._request(self._native.begin) + tx = ConnectionTransaction(self._native, self._boundary, self, _running.get()) + token = _running.set(tx) + try: + yield tx + except BaseException: + await tx._end() # pylint: disable=protected-access + await self._rollback() + raise + finally: + _running.reset(token) + records = await tx._end() # pylint: disable=protected-access + # Its commit's sync runs what the body gave `with_commit`; one that + # wrote nothing synced nothing, and it runs after. + self._boundary.with_sync = tx._commit_runs() # pylint: disable=protected-access + try: + result = self._native.commit(records) + if inspect.isawaitable(result): + await result + except Exception as error: + # A COMMIT that fails (a deferred constraint) leaves the + # transaction open: without the rollback, the calls after it + # would run inside it. + await self._rollback() + raise self._failure(error) from None + finally: + unsynced, self._boundary.with_sync = self._boundary.with_sync, None + if unsynced is not None: + unsynced() + + async def record(self, name: str) -> bytes: + """The record ``name``, as the last commit that set it left it (see ``ConnectionTransaction.set_record``); empty if none did.""" + async with self._serial(): + return await self._request(lambda: self._native.record(name)) + + def is_usable(self) -> bool: + """Whether SQL can still run: an SQL error leaves the session usable, a storage error ends it.""" + return self._native.is_usable() + + async def save(self) -> None: + """ + End the session after the calls before it, leaving its engine to the + next one. It fails if they left a transaction open, or a failure ended + the session. + """ + async with self._serial(): + await self._request(self._native.save) + + async def close(self) -> None: + """End the session now, whatever its state: queued calls fail.""" + await self._native.close() + + async def _rollback(self) -> None: + """Roll back, unless a storage error ended the session: the caller raises its own error, which says why.""" + try: + result = self._native.rollback() + if inspect.isawaitable(result): + await result + except Exception: # pylint: disable=broad-exception-caught + pass + + +class ConnectionTransaction(_Sql): + """ + SQL in one transaction, while its body runs, and the records it commits + with. Its calls run one at a time, in call order, also when the body + awaits several together, or none: each call made while it is open runs + before its commit or rollback, so none outlives it. + """ + + def __init__(self, native: Any, boundary: _Boundary, connection: Connection, outer: ConnectionTransaction | None): + super().__init__(native, boundary) + self._connection = connection + # The transaction whose body this one's runs in, if any: of another + # connection, since a connection's own are refused. + self._outer = outer + self._records: dict[str, bytes] = {} + self._with_commits: list[Callable[[], None]] = [] + self._lock = asyncio.Lock() + self._ended = False + + @property + def open(self) -> bool: + """Whether it is open: its body runs, and holds its connection.""" + return not self._ended + + def holds(self, connection: Connection) -> bool: + """Whether this transaction, or one it runs inside, is open and holds ``connection``.""" + if not self._ended and self._connection is connection: + return True + return self._outer is not None and self._outer.holds(connection) + + async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: + async with self._serial(): + return await self._request(lambda: self._native.query_columns(sql, params)) + + def set_record(self, name: str, value: bytes) -> None: + """ + Set the record ``name`` to ``value`` with this transaction, in the same + write: a whole value kept beside the SQL, in no table, stored exactly + when the transaction is, even if it changes no row. It stays until a + later transaction sets it again; an empty value reads as never set. A + name is at most 255 bytes, and not ``db``; a value at most 16 MiB. + """ + self._refuse_ended() + self._records[name] = bytes(value) + + async def record(self, name: str) -> bytes: + """The record ``name`` as this transaction would commit it: as it set it, or as committed.""" + # As set when it is called, in call order with `set_record`. + value = self._records.get(name) + async with self._serial(): + if value is not None: + return value + return await self._request(lambda: self._native.record(name)) + + def with_commit(self, run: Callable[[], None]) -> None: + """ + Run ``run`` with the commit, right after its writes and before it waits + for them to be stored: the store's host gets what ``run`` makes after + the writes, and before anything that fails the commit's wait (Restate's + sends, which it stores in order with the writes). Never if the + transaction rolls back. + """ + self._refuse_ended() + self._with_commits.append(run) + + async def _end(self) -> dict[str, bytes]: + """End it, refusing calls from now on, once the calls made before have completed; then the records it set.""" + self._ended = True + async with self._lock: + pass + return dict(self._records) + + def _commit_runs(self) -> Callable[[], None] | None: + """What runs with its commit, once it ended: see ``with_commit``.""" + runs = self._with_commits + if not runs: + return None + + def run_all() -> None: + for run in runs: + run() + + return run_all + + @asynccontextmanager + async def _serial(self) -> AsyncIterator[None]: + """Run the body, a call made while the transaction is open, once the calls before it have.""" + self._refuse_ended() + async with self._lock: + yield + + def _refuse_ended(self) -> None: + if self._ended: + raise SqlError("the transaction has ended") diff --git a/src/lib.rs b/src/lib.rs index 5f18c254..abd67b0c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,6 +14,8 @@ use restate_sdk_shared_core::{ use std::fmt; use std::time::{Duration, SystemTime}; +mod sessions; + // Current crate version const CURRENT_VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -1090,6 +1092,8 @@ fn _internal(m: &Bound<'_, PyModule>) -> PyResult<()> { PyNotificationHandle::from(CANCEL_NOTIFICATION_HANDLE), )?; + sessions::register(m)?; + // Set customized error formatter set_error_formatter(PythonErrorFormatter); diff --git a/src/sessions.rs b/src/sessions.rs new file mode 100644 index 00000000..927b9569 --- /dev/null +++ b/src/sessions.rs @@ -0,0 +1,493 @@ +//! Sessions: an object's SQLite-compatible database, kept in its Restate +//! state, over a key/value store implemented in Python. +//! +//! The storage and the SQL on it are shared-core's (`sql::Database`, with its +//! feature `sql`); this module binds them to Python, as `native/src/node.rs` +//! of sdk-sessions binds them to JavaScript. `restate/sessions/_database.py` +//! opens a database over a store, and orders the calls. +//! +//! `sql_open` takes the store as one dispatcher function, +//! `dispatch(op, key, value, reply)`, and the asyncio event loop it runs on. +//! Every request is scheduled on that loop with `call_soon_threadsafe`, which +//! runs callbacks in the order they were scheduled, so a read sees the writes +//! before it. +//! +//! `set` and `clear` queue a write; `sync` completes once every write before +//! it is stored. The object's state keeps writes in issue order, so a crash +//! loses only a suffix of them (see [`KvStore`]): a statement that commits +//! queues its writes, then completes at one `sync`. With `keys`, a checkpoint +//! deletes the packs no checkpoint needs; a store without keeps its garbage. +//! +//! Every database of the process is kept warm in one page cache +//! ([`page_cache`]). +//! +//! A call runs on the calling thread as far as it can ([`answer`]): one that +//! needs nothing from the store, such as SQL on cached pages or a statement +//! inside a transaction, returns its value, or raises its error; one that +//! waits for the store goes on in the tokio runtime and returns an awaitable. +//! The caller awaits the latter. +use bytes::Bytes; +use futures::channel::oneshot; +use futures::future::BoxFuture; +use pyo3::create_exception; +use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyBool, PyByteArray, PyBytes, PyFloat, PyInt, PyList, PyString, PyTuple}; +use pyo3::IntoPyObjectExt; +use restate_sdk_shared_core::sql::{self, Numeric, Value}; +use restate_sdk_shared_core::state::{self, KvStore, PageCache, PageCacheConfig}; +use std::borrow::Cow; +use std::collections::HashMap; +use std::future::Future; +use std::io; +use std::sync::{Arc, Mutex}; + +create_exception!( + restate_sdk_python_core, + SqlRefusedException, + pyo3::exceptions::PyException, + "A failure of the database that a retry would repeat (shared-core's `sql::Refused`): it comes of what the host gave, its SQL, its records, a transaction it left open. It may have ended the session." +); + +/// What a request to the store does: the store's method it calls. +pub(crate) const GET: u32 = 0; +pub(crate) const SET: u32 = 1; +pub(crate) const CLEAR: u32 = 2; +pub(crate) const SYNC: u32 = 3; +pub(crate) const KEYS: u32 = 4; + +/// What a request answers: a `get`'s value, the store's keys, or nothing. +enum Answer { + Nothing, + Value(Bytes), + Keys(Vec), +} + +/// The callable that `dispatch` calls once with a request's outcome: +/// `reply(value)` or `reply(None, error)`. The value is a `get`'s (`bytes`, +/// or `None` when absent), the store's keys (a `list[str]`), or `None`. +#[pyclass(frozen)] +struct SqlReply { + sender: Mutex>>>, +} + +#[pymethods] +impl SqlReply { + #[pyo3(signature = (value, error=None))] + fn __call__(&self, value: Option>, error: Option) -> PyResult<()> { + let answer = match (error, value) { + (Some(error), _) => Err(error), + (None, None) => Ok(Answer::Nothing), + (None, Some(value)) => Ok(if value.is_instance_of::() { + Answer::Keys(value.extract()?) + } else { + Answer::Value(Bytes::from(bytes_of(&value)?)) + }), + }; + let sender = self.sender.lock().ok().and_then(|mut sender| sender.take()); + if let Some(sender) = sender { + let _ = sender.send(answer); + } + Ok(()) + } +} + +/// The bytes of `bytes`, `bytearray` or `memoryview`. +fn bytes_of(value: &Bound<'_, PyAny>) -> PyResult> { + if let Ok(bytes) = value.cast::() { + Ok(bytes.as_bytes().to_vec()) + } else if value.is_instance_of::() { + value.extract() + } else if let Ok(bytes) = value.call_method0("tobytes") { + Ok(bytes.cast::()?.as_bytes().to_vec()) + } else { + Err(PyTypeError::new_err(format!( + "expected bytes, got {}", + value.get_type().name()? + ))) + } +} + +/// The Python store, behind its dispatcher on its event loop. +struct PyStore { + event_loop: Py, + dispatch: Py, + /// Whether the store has `keys`. + lists_keys: bool, +} + +fn io_error(error: impl std::fmt::Display) -> io::Error { + io::Error::other(error.to_string()) +} + +impl PyStore { + /// Schedule `dispatch(op, key, value, reply)` on the event loop; `reply` + /// is `None` for a write, which answers nothing. + fn schedule( + &self, + op: u32, + key: &str, + value: Option, + reply: Option, + ) -> io::Result<()> { + Python::attach(|py| { + let value = value.map(|value| PyBytes::new(py, &value).into_any().unbind()); + let reply = match reply { + Some(reply) => Py::new(py, reply)?.into_any(), + None => py.None(), + }; + self.event_loop.call_method1( + py, + "call_soon_threadsafe", + (self.dispatch.clone_ref(py), op, key, value, reply), + )?; + Ok::<_, PyErr>(()) + }) + .map_err(io_error) + } + + /// A request that answers: a read, a sync, or the keys. + async fn request(&self, op: u32, key: &str) -> io::Result { + let (sender, receiver) = oneshot::channel(); + let reply = SqlReply { + sender: Mutex::new(Some(sender)), + }; + self.schedule(op, key, None, Some(reply))?; + match receiver.await { + Ok(Ok(answer)) => Ok(answer), + Ok(Err(error)) => Err(io_error(error)), + Err(_) => Err(io_error("the store dropped the request")), + } + } +} + +fn unexpected(method: &str) -> io::Error { + io_error(format!("the store's {method} answered with the wrong type")) +} + +impl KvStore for PyStore { + fn get<'a>(&'a self, key: &'a str) -> BoxFuture<'a, io::Result>> { + Box::pin(async move { + match self.request(GET, key).await? { + Answer::Nothing => Ok(None), + Answer::Value(value) => Ok(Some(value)), + Answer::Keys(_) => Err(unexpected("get")), + } + }) + } + + fn put(&self, key: &str, value: Bytes) -> io::Result<()> { + self.schedule(SET, key, Some(value), None) + } + + fn delete(&self, key: &str) -> io::Result<()> { + self.schedule(CLEAR, key, None, None) + } + + fn sync(&self) -> BoxFuture<'_, io::Result<()>> { + Box::pin(async { self.request(SYNC, "").await.map(drop) }) + } + + fn keys(&self) -> BoxFuture<'_, io::Result>> { + Box::pin(async move { + if !self.lists_keys { + return Err(io::Error::new( + io::ErrorKind::Unsupported, + "the store has no keys", + )); + } + match self.request(KEYS, "").await? { + Answer::Keys(keys) => Ok(keys), + _ => Err(unexpected("keys")), + } + }) + } +} + +/// A failure of the database as Python gets it: one a retry would repeat +/// ([`sql::Refused`]) is an `SqlRefusedException`, which +/// `restate/sessions/_database.py` tells apart, as it may have ended the +/// session; any other a `RuntimeError`. +fn error(error: Box) -> PyErr { + match error.is::() { + true => SqlRefusedException::new_err(error.to_string()), + false => PyRuntimeError::new_err(error.to_string()), + } +} + +/// `call`'s value if it is ready at once, polled here on the calling thread; +/// else an awaitable of it, as it goes on in the tokio runtime. Ready, a call +/// spares two thread hops: what a call costs otherwise, several times the SQL +/// it runs. +fn answer<'py, T, F>(py: Python<'py>, call: F) -> PyResult> +where + F: Future> + Send + 'static, + T: for<'a> IntoPyObject<'a> + Send + 'static, +{ + use std::task::{Context, Poll, Waker}; + let mut call = Box::pin(call); + let poll = { + // What the call reaches of tokio (the page cache's disk tier) finds + // the runtime it goes on in when pending. + let _runtime = pyo3_async_runtimes::tokio::get_runtime().enter(); + call.as_mut().poll(&mut Context::from_waker(Waker::noop())) + }; + match poll { + Poll::Ready(result) => result?.into_bound_py_any(py), + Poll::Pending => pyo3_async_runtimes::tokio::future_into_py(py, call), + } +} + +/// The page cache shared by the process's databases: 256 MiB of memory and +/// 1 GiB of disk, in a file under the temporary directory, or the file +/// `RESTATE_SQL_CACHE_FILE` names. `RESTATE_SQL_CACHE_MEMORY` and +/// `RESTATE_SQL_CACHE_DISK` set its sizes (bytes, or with KiB, MiB or GiB); +/// the memory one also bounds the namespaces' states and the parked engines +/// it keeps, so a small cache is a cold one. `RESTATE_SQL_CACHE=off` (or `0`, +/// `false`, `no`) opens every database without a cache (`None`): each session +/// reads what it needs from the store. A size that is not one fails every +/// open, with its message: the process serves, and every invocation retries. +async fn page_cache() -> io::Result> { + static CACHE: futures::lock::Mutex>> = futures::lock::Mutex::new(None); + let mut cache = CACHE.lock().await; + if let Some(cache) = &*cache { + return Ok(cache.clone()); + } + let opened = match cache_config()? { + Some(config) => Some(PageCache::open(config).await?), + None => None, + }; + Ok(cache.insert(opened).clone()) +} + +/// The page cache's configuration, from the environment: `None` when it is off. +fn cache_config() -> io::Result> { + let var = |name: &str| std::env::var(name).ok().filter(|v| !v.is_empty()); + let off = var("RESTATE_SQL_CACHE").map(|v| v.to_ascii_lowercase()); + if matches!(off.as_deref(), Some("off" | "0" | "false" | "no")) { + return Ok(None); + } + let size = |name: &str, default: usize| match var(name) { + Some(value) => state::parse_size(&value).ok_or_else(|| { + io::Error::other(format!( + "{name}={value}: not a size (bytes, or with KiB, MiB or GiB)" + )) + }), + None => Ok(default), + }; + Ok(Some(PageCacheConfig { + file: std::env::var_os("RESTATE_SQL_CACHE_FILE").map(std::path::PathBuf::from), + ..PageCacheConfig::sized( + size("RESTATE_SQL_CACHE_MEMORY", 256 << 20)?, + size("RESTATE_SQL_CACHE_DISK", 1 << 30)?, + ) + })) +} + +/// Open the database of `namespace` over the store behind `dispatch`, which +/// runs on `event_loop` (see the module's doc), kept warm in the process's +/// page cache; `lists_keys` if the store has `keys`. See +/// [`sql::Open::reader`] for a `reader`. +#[pyfunction] +#[pyo3(signature = (namespace, dispatch, event_loop, lists_keys = false, reader = false))] +pub(crate) fn sql_open<'py>( + py: Python<'py>, + namespace: String, + dispatch: Py, + event_loop: Py, + lists_keys: bool, + reader: bool, +) -> PyResult> { + let store = Arc::new(PyStore { + event_loop, + dispatch, + lists_keys, + }); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let mut open = sql::Database::open(store); + let cache = page_cache() + .await + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; + if let Some(cache) = &cache { + open = open.cached(cache, &namespace); + } + if reader { + open = open.reader(); + } + Ok(PySqlDatabase(Arc::new(open.await.map_err(error)?))) + }) +} + +/// One session of a namespace's database (see [`sql::Database`]): every SQL +/// commit is durable when its call completes. Each call returns its value, or +/// an awaitable of it (see [`answer`]). +#[pyclass(frozen)] +pub(crate) struct PySqlDatabase(Arc); + +#[pymethods] +impl PySqlDatabase { + /// Run SQL (one or more statements; parameters bind to a single one) and + /// return `(columns, rows)`: the rows of its last statement, with its + /// column names. A commit is durable when this completes. + #[pyo3(signature = (sql, params=None))] + fn query_columns<'py>( + &self, + py: Python<'py>, + sql: String, + params: Option>, + ) -> PyResult> { + let params = params.map(|params| sql_params(¶ms)).transpose()?; + let db = self.0.clone(); + answer(py, async move { + let (columns, rows) = db.query_columns(&sql, params).await.map_err(error)?; + Ok(QueryResult { columns, rows }) + }) + } + + /// See [`sql::Database::begin`]. + fn begin<'py>(&self, py: Python<'py>) -> PyResult> { + let db = self.0.clone(); + answer(py, async move { db.begin().await.map_err(error) }) + } + + /// See [`sql::Database::commit`]: `records` by name. + #[pyo3(signature = (records=None))] + fn commit<'py>( + &self, + py: Python<'py>, + records: Option>>, + ) -> PyResult> { + let records: Vec<_> = (records.unwrap_or_default().into_iter()) + .map(|(name, value)| (name, Bytes::from(value))) + .collect(); + let db = self.0.clone(); + answer(py, async move { db.commit(records).await.map_err(error) }) + } + + /// See [`sql::Database::rollback`]. + fn rollback<'py>(&self, py: Python<'py>) -> PyResult> { + let db = self.0.clone(); + answer(py, async move { db.rollback().await.map_err(error) }) + } + + /// See [`sql::Database::is_usable`]. + fn is_usable(&self) -> bool { + self.0.is_usable() + } + + /// See [`sql::Database::record`]: `bytes`. + fn record<'py>(&self, py: Python<'py>, name: String) -> PyResult> { + let db = self.0.clone(); + answer(py, async move { + let value = db.record(&name).await.map_err(error)?; + Ok(Cow::<'static, [u8]>::Owned(value.to_vec())) + }) + } + + /// See [`sql::Database::save`]. + fn save<'py>(&self, py: Python<'py>) -> PyResult> { + let db = self.0.clone(); + answer(py, async move { db.save().await.map_err(error) }) + } + + /// See [`sql::Database::close`]: an awaitable. + fn close<'py>(&self, py: Python<'py>) -> PyResult> { + let db = self.0.clone(); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + db.close().await; + Ok(()) + }) + } +} + +/// Rows with the column names of their statement: `(columns, rows)` in Python. +struct QueryResult { + columns: Vec, + rows: sql::Rows, +} + +impl<'py> IntoPyObject<'py> for QueryResult { + type Target = PyTuple; + type Output = Bound<'py, PyTuple>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult> { + let rows = PyList::empty(py); + for row in self.rows { + let values = row + .into_iter() + .map(|value| python_value(py, value)) + .collect::>>()?; + rows.append(PyList::new(py, values)?)?; + } + PyTuple::new( + py, + [PyList::new(py, self.columns)?.into_any(), rows.into_any()], + ) + } +} + +/// SQL values from Python: `None`, `bool` (0 or 1), `int` (64 bits), +/// `float`, `str`, and `bytes`, `bytearray` or `memoryview` (BLOB). +fn sql_params(params: &Bound<'_, PyAny>) -> PyResult> { + if params.is_instance_of::() || params.is_instance_of::() { + return Err(PyTypeError::new_err( + "SQL parameters are a sequence of values", + )); + } + params.try_iter()?.map(|param| sql_value(¶m?)).collect() +} + +fn sql_value(value: &Bound<'_, PyAny>) -> PyResult { + if value.is_none() { + Ok(Value::Null) + } else if value.is_instance_of::() { + Ok(Value::from_i64(i64::from(value.extract::()?))) + } else if value.is_instance_of::() { + let number = value + .extract::() + .map_err(|_| PyValueError::new_err("an integer parameter exceeds 64 bits"))?; + Ok(Value::from_i64(number)) + } else if value.is_instance_of::() { + Ok(Value::from_f64(value.extract::()?)) + } else if value.is_instance_of::() { + Ok(Value::build_text(value.extract::()?)) + } else { + let bytes = bytes_of(value).map_err(|_| { + let name = value.get_type().name().map(|n| n.to_string()); + PyTypeError::new_err(format!( + "unsupported SQL value of type {}", + name.unwrap_or_default() + )) + })?; + Value::from_slice(&bytes).map_err(|e| PyValueError::new_err(e.to_string())) + } +} + +/// A value of a row: `None`, `int`, `float`, `str` or `bytes`. +fn python_value(py: Python<'_>, value: Value) -> PyResult> { + match value { + Value::Null => Ok(py.None().into_bound(py)), + Value::Numeric(Numeric::Integer(i)) => i.into_bound_py_any(py), + Value::Numeric(Numeric::Float(f)) => f64::from(f).into_bound_py_any(py), + Value::Text(text) => Ok(PyString::new(py, text.as_str()).into_any()), + Value::Blob(blob) => Ok(PyBytes::new(py, &blob).into_any()), + } +} + +/// Add the sessions' classes, functions and constants to the module `m`. +pub(crate) fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + m.add_function(wrap_pyfunction!(sql_open, m)?)?; + m.add( + "SqlRefusedException", + m.py().get_type::(), + )?; + m.add("SQL_GET", GET)?; + m.add("SQL_SET", SET)?; + m.add("SQL_CLEAR", CLEAR)?; + m.add("SQL_SYNC", SYNC)?; + m.add("SQL_KEYS", KEYS)?; + Ok(()) +} diff --git a/tests/sessions_database.py b/tests/sessions_database.py new file mode 100644 index 00000000..8a8ac626 --- /dev/null +++ b/tests/sessions_database.py @@ -0,0 +1,551 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""A session's database over an in-memory store: SQL, values, records, transactions, and lost writes.""" + +# pylint: disable=missing-function-docstring,redefined-builtin + +import asyncio +import itertools +import random +import typing + +import pytest + +from restate.sessions._database import SqlError, open + + +@pytest.fixture +def anyio_backend(): + return "asyncio" + + +class Backend: + """Stored values, and writes issued but not yet synced.""" + + def __init__(self, seed=None): + self.stored = {} + self.unconfirmed = [] + # Which unsynced writes a failure keeps: none, or a random prefix. + self.random = None if seed is None else random.Random(seed) + + def confirm(self): + for key, value in self.unconfirmed: + if value is None: + self.stored.pop(key, None) + else: + self.stored[key] = value + self.unconfirmed.clear() + + def lose(self): + """A failure: the store keeps a prefix of the unsynced writes, as Restate does.""" + kept = 0 if self.random is None else self.random.randint(0, len(self.unconfirmed)) + del self.unconfirmed[kept:] + self.confirm() + + def request(self, fails=False): + """A request's view: its own writes over the stored values.""" + return Request(self, fails) + + +class Request: + # A listing of the store's keys, if a test gives it one. + keys: typing.Any = None + + def __init__(self, backend, fails): + self.backend = backend + self.fails = fails + + def get(self, key): + # The read is issued now: what it sees is decided here. + for k, value in reversed(self.backend.unconfirmed): + if k == key: + return value + return self.backend.stored.get(key) + + def set(self, key, value): + self.backend.unconfirmed.append((key, bytes(value))) + + def clear(self, key): + self.backend.unconfirmed.append((key, None)) + + async def sync(self): + if self.fails: + self.backend.lose() + raise RuntimeError("the request failed") + self.backend.confirm() + + +_namespaces = itertools.count() + + +def fresh(): + return f"namespace-{next(_namespaces)}" + + +async def test_runs_sql_and_round_trips_values(): + backend = Backend() + ns = fresh() + db = await open(ns, backend.request()) + await db.execute("CREATE TABLE t(i INTEGER, r REAL, s TEXT, b BLOB, n INTEGER)") + await db.execute("INSERT INTO t VALUES(?, ?, ?, ?, ?)", [42, 1.5, "héllo", b"\x01\x02\x03", None]) + await db.execute("INSERT INTO t VALUES(?, ?, ?, ?, ?)", [2**62, -0.25, "", bytearray(), True]) + rows = await db.values("SELECT * FROM t ORDER BY i") + assert rows == [[42, 1.5, "héllo", b"\x01\x02\x03", None], [2**62, -0.25, "", b"", 1]] + assert await db.query_columns("SELECT i AS answer FROM t WHERE i = ?", [42]) == (["answer"], [[42]]) + assert await db.query("SELECT i, s FROM t ORDER BY i") == [{"i": 42, "s": "héllo"}, {"i": 2**62, "s": ""}] + assert await db.get("SELECT i FROM t WHERE i = ?", [42]) == {"i": 42} + assert await db.get("SELECT i FROM t WHERE i = 0") is None + await db.save() + + # Another session, and another process's view (no cached state), read it. + for name in [ns, fresh()]: + again = await open(name, backend.request()) + assert await again.values("SELECT count(*) FROM t") == [[2]] + await again.save() + + +async def test_a_column_named_twice_is_an_sql_error_for_query_but_not_for_values(): + db = await open(fresh(), Backend().request()) + with pytest.raises(SqlError, match="named twice"): + await db.query("SELECT 1 AS a, 2 AS a") + assert await db.values("SELECT 1 AS a, 2 AS a") == [[1, 2]] + await db.save() + + +async def test_values_it_cannot_store_are_sql_errors(): + db = await open(fresh(), Backend().request()) + with pytest.raises(SqlError, match="exceeds 64 bits"): + await db.values("SELECT ?", [2**64]) + with pytest.raises(SqlError, match="unsupported SQL value"): + await db.values("SELECT ?", [object()]) # type: ignore[list-item] + with pytest.raises(SqlError, match="a sequence of values"): + await db.values("SELECT ?", "x") # type: ignore[arg-type] + assert db.is_usable() + await db.save() + + +async def test_an_sql_error_leaves_the_session_usable_and_a_transaction_rolls_back(): + db = await open(fresh(), Backend().request()) + await db.execute("CREATE TABLE t(id INTEGER PRIMARY KEY)") + with pytest.raises(SqlError, match="no such table"): + await db.values("SELECT * FROM missing") + assert db.is_usable() + with pytest.raises(ValueError, match="abort"): + async with db.transaction() as tx: + await tx.execute("INSERT INTO t VALUES(1)") + raise ValueError("abort") + assert await db.values("SELECT count(*) FROM t") == [[0]] + async with db.transaction() as tx: + assert await tx.values("INSERT INTO t VALUES(2) RETURNING id") == [[2]] + assert await db.values("SELECT id FROM t") == [[2]] + await db.save() + + +async def consistent_after_loss(lost, seed): + """ + Requests insert a row each; request ``lost`` fails, its sync losing some of + its writes (none kept, or a prefix by ``seed``). The owner and a process + without cached state must then read the rows of the other requests, and the + lost one's only if both do. + """ + backend = Backend(seed) + ns = fresh() + db = await open(ns, backend.request()) + await db.execute("CREATE TABLE t(id INTEGER PRIMARY KEY, v BLOB)") + await db.save() + for id in range(20): + fails = id == lost + db = await open(ns, backend.request(fails)) + insert = db.execute("INSERT INTO t VALUES(?, randomblob(3000))", [id]) + if fails: + # The store's own error, as it raised it. + with pytest.raises(RuntimeError, match="the request failed"): + await insert + await db.close() + else: + await insert + await db.save() + seen = [] + for name in [ns, fresh()]: + view = await open(name, backend.request()) + [[rows, landed]] = await view.values("SELECT count(*), count(*) FILTER (WHERE id = ?) FROM t", [lost]) + [[integrity]] = await view.values("PRAGMA integrity_check") + await view.save() + if rows != 19 + typing.cast(int, landed) or integrity != "ok": + return False + seen.append(landed) + return seen[0] == seen[1] + + +async def test_lost_writes_leave_every_reader_consistent(): + for seed in [None, 1, 2, 3, 4, 5]: + for lost in range(1, 15): + assert await consistent_after_loss(lost, seed), f"lost request {lost}, seed {seed}" + + +async def test_a_write_the_store_rejects_fails_the_next_sync_with_its_error(): + store = Backend().request() + + class Closed(Exception): + pass + + def rejected(key, value): + raise Closed("state is closed") + + store.set = rejected + db = await open(fresh(), store) + with pytest.raises(Closed, match="state is closed"): + await db.execute("CREATE TABLE t(id INTEGER)") + assert not db.is_usable() + await db.close() + + +async def test_a_transaction_holds_the_session_other_calls_wait_until_it_ends(): + db = await open(fresh(), Backend().request()) + await db.execute("CREATE TABLE t(id INTEGER PRIMARY KEY)") + handles = [] + + async def rolled_back(): + async with db.transaction() as tx: + handles.append(tx) + await tx.execute("INSERT INTO t VALUES(1)") + await asyncio.sleep(0.01) + raise ValueError("abort") + + async def committed(): + async with db.transaction() as tx: + await tx.execute("INSERT INTO t VALUES(3)") + + first = asyncio.ensure_future(rolled_back()) + await asyncio.sleep(0) + # Started while the transaction waits: neither runs inside it. + write = asyncio.ensure_future(db.execute("INSERT INTO t VALUES(2)")) + last = asyncio.ensure_future(committed()) + with pytest.raises(ValueError, match="abort"): + await first + await write + await last + assert await db.values("SELECT id FROM t ORDER BY id") == [[2], [3]] + with pytest.raises(SqlError, match="the transaction has ended"): + await handles[0].values("SELECT 1") + await db.save() + + +async def test_the_calls_of_a_body_run_in_call_order_and_before_its_commit(): + db = await open(fresh(), Backend().request()) + await db.execute("CREATE TABLE t(id INTEGER)") + async with db.transaction() as tx: + # Made together, and one left unawaited: each runs, in order, before the commit. + pending = asyncio.ensure_future(tx.execute("INSERT INTO t VALUES(3)")) + await asyncio.gather(*(tx.execute("INSERT INTO t VALUES(?)", [i]) for i in (1, 2))) + assert pending.done() + assert await db.values("SELECT id FROM t ORDER BY rowid") == [[3], [1], [2]] + await db.save() + + +async def test_a_commit_that_fails_rolls_the_transaction_back(): + db = await open(fresh(), Backend().request()) + await db.execute("PRAGMA foreign_keys=ON") + await db.execute("CREATE TABLE p(id INTEGER PRIMARY KEY)") + await db.execute("CREATE TABLE c(pid INTEGER REFERENCES p(id) DEFERRABLE INITIALLY DEFERRED)") + # The deferred constraint fails at COMMIT, not at the INSERT. + with pytest.raises(SqlError, match="FOREIGN KEY constraint failed"): + async with db.transaction() as tx: + await tx.execute("INSERT INTO c VALUES(7)") + assert db.is_usable() + # Nothing of it is left, and the next transaction can begin. + assert await db.values("SELECT * FROM c") == [] + async with db.transaction() as tx: + await tx.execute("INSERT INTO p VALUES(7)") + assert await db.values("SELECT id FROM p") == [[7]] + await db.save() + + +async def test_a_statement_that_ends_the_transaction_of_a_body_ends_the_session_and_commits_nothing(): + backend = Backend() + ns = fresh() + db = await open(ns, backend.request()) + await db.execute("CREATE TABLE t(id INTEGER)") + for end in ["COMMIT", "ROLLBACK", "END"]: + session = await open(ns, backend.request()) + with pytest.raises(SqlError, match="ended the transaction"): + async with session.transaction() as tx: + await tx.execute("INSERT INTO t VALUES(1)") + await tx.execute(end) + assert not session.is_usable() + assert await db.values("SELECT count(*) FROM t") == [[0]] + await db.save() + + +async def test_save_refuses_a_transaction_left_open_whose_changes_are_lost(): + backend = Backend() + ns = fresh() + db = await open(ns, backend.request()) + await db.execute("CREATE TABLE t(id INTEGER)") + await db.execute("BEGIN") + await db.execute("INSERT INTO t VALUES(1)") + with pytest.raises(SqlError, match="left open"): + await db.save() + again = await open(ns, backend.request()) + assert await again.values("SELECT count(*) FROM t") == [[0]] + await again.save() + + +async def test_a_record_a_commit_cannot_hold_is_refused_and_the_session_goes_on(): + db = await open(fresh(), Backend().request()) + await db.execute("CREATE TABLE t(id INTEGER)") + with pytest.raises(SqlError): + async with db.transaction() as tx: + await tx.execute("INSERT INTO t VALUES(1)") + tx.set_record("outbox", bytes((16 << 20) + 1)) + assert db.is_usable() + assert await db.values("SELECT count(*) FROM t") == [[0]] + await db.save() + + +async def test_a_call_running_meanwhile_does_not_make_the_session_look_unusable(): + db = await open(fresh(), Backend().request()) + await db.execute("CREATE TABLE t(id INTEGER)") + running = asyncio.ensure_future( + db.values( + "WITH RECURSIVE n(i) AS (SELECT 1 UNION ALL SELECT i + 1 FROM n WHERE i < 200000) SELECT count(*) FROM n" + ) + ) + await asyncio.sleep(0) + assert db.is_usable() + await running + await db.save() + assert not db.is_usable() + + +async def test_records_commit_with_their_transaction_and_only_then(): + backend = Backend() + namespace = fresh() + db = await open(namespace, backend.request()) + await db.execute("CREATE TABLE t(id INTEGER)") + assert await db.record("outbox") == b"" + # Even with a transaction that changes nothing. + async with db.transaction() as tx: + await tx.values("SELECT * FROM t") + tx.set_record("outbox", b"sends 1") + assert await tx.record("outbox") == b"sends 1" + assert await db.record("outbox") == b"sends 1" + # A rollback commits none; nor does a commit that sets none change them. + with pytest.raises(ValueError, match="no"): + async with db.transaction() as tx: + await tx.execute("INSERT INTO t VALUES(2)") + tx.set_record("outbox", b"sends 2") + raise ValueError("no") + async with db.transaction() as tx: + await tx.execute("INSERT INTO t VALUES(3)") + assert await db.record("outbox") == b"sends 1" + await db.save() + # A later session opens with them; an empty value reads as never set. + again = await open(namespace, backend.request()) + assert await again.record("outbox") == b"sends 1" + async with again.transaction() as tx: + tx.set_record("outbox", b"") + assert await again.record("outbox") == b"" + await again.save() + + +async def test_with_commit_runs_after_the_commits_writes_before_its_sync_and_never_on_a_rollback(): + backend = Backend() + events = [] + + class Recording(Request): + def set(self, key, value): + events.append("set") + super().set(key, value) + + async def sync(self): + events.append("sync") + await super().sync() + + db = await open(fresh(), Recording(backend, False)) + await db.execute("CREATE TABLE t(id INTEGER)") + events.clear() + async with db.transaction() as tx: + await tx.execute("INSERT INTO t VALUES(1)") + tx.with_commit(lambda: events.append("sent")) + assert events[-2:] == ["sent", "sync"] and events[0] == "set" + + events.clear() + with pytest.raises(ValueError): + async with db.transaction() as tx: + await tx.execute("INSERT INTO t VALUES(2)") + tx.with_commit(lambda: events.append("sent")) + raise ValueError("no") + assert "sent" not in events + + # A transaction that writes nothing syncs nothing: it runs after the commit. + events.clear() + async with db.transaction() as tx: + tx.with_commit(lambda: events.append("sent")) + assert events == ["sent"] + await db.save() + + +async def test_a_call_to_the_session_inside_its_transaction_is_refused_not_left_waiting(): + db = await open(fresh(), Backend().request()) + await db.execute("CREATE TABLE t(id INTEGER)") + + async def after(delay): + await asyncio.sleep(delay) + await db.execute("INSERT INTO t VALUES(3)") + + async with db.transaction() as tx: + await tx.execute("INSERT INTO t VALUES(1)") + # It would wait for this transaction, which waits for this body. + with pytest.raises(SqlError, match="use the transaction handle"): + await db.execute("INSERT INTO t VALUES(2)") + with pytest.raises(SqlError, match="use the transaction handle"): + async with db.transaction(): + pass + # Once the body has ended, a task it started may use the session. + later = asyncio.ensure_future(after(0.005)) + await later + assert await db.values("SELECT id FROM t ORDER BY id") == [[1], [3]] + await db.save() + + +async def test_a_call_to_a_session_inside_a_transaction_of_another_inside_its_own_is_refused(): + a = await open(fresh(), Backend().request()) + b = await open(fresh(), Backend().request()) + async with a.transaction(): + async with b.transaction() as tx: + await tx.values("SELECT 1") + # `a` waits for its transaction, which waits for this body. + with pytest.raises(SqlError, match="use the transaction handle"): + await a.values("SELECT 1") + # Once b's transaction has ended, b may be used again from a's body. + assert await b.values("SELECT 1") == [[1]] + await a.save() + await b.save() + + +async def test_a_statement_that_commits_completes_once_its_writes_are_stored_however_many(): + backend = Backend() + db = await open(fresh(), backend.request()) + await db.execute("CREATE TABLE t(id INTEGER PRIMARY KEY, b BLOB)") + # Thousands of page writes: a tail entry of megabytes, in parts, each + # scheduled on the loop. + await db.execute( + "INSERT INTO t(b) SELECT randomblob(3000) FROM " + "(WITH RECURSIVE c(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM c WHERE x < 3000) SELECT x FROM c)" + ) + assert backend.unconfirmed == [] + assert sum(key.startswith("v22/tail/") for key in backend.stored) > 8 + await db.save() + + +async def test_an_awaitable_get_is_read_when_issued(): + """A store whose get returns an awaitable, as a Restate context's does.""" + backend = Backend() + plain = backend.request() + + class AwaitableGets(Request): + def get(self, key): + future = asyncio.get_running_loop().create_future() + future.set_result(plain.get(key)) + return future + + ns = fresh() + db = await open(ns, AwaitableGets(backend, False)) + await db.execute("CREATE TABLE t(id INTEGER)") + await db.execute("INSERT INTO t VALUES(1)") + await db.save() + for name in [ns, fresh()]: + again = await open(name, AwaitableGets(backend, False)) + assert await again.values("SELECT * FROM t") == [[1]] + await again.save() + + +class ReadOnly(Request): + """A request that refuses writes, as a shared handler's store does.""" + + def set(self, key, value): + raise AssertionError("a reader writes nothing") + + def clear(self, key): + raise AssertionError("a reader writes nothing") + + +async def test_a_reader_reads_the_state_it_opened_on_beside_the_writer(): + backend = Backend() + ns = fresh() + writer = await open(ns, backend.request()) + await writer.execute("CREATE TABLE t(n INTEGER)") + await writer.execute("INSERT INTO t VALUES (1)") + + reader = await open(ns, ReadOnly(backend, False), reader=True) + assert await reader.values("SELECT n FROM t") == [[1]] + # The writer commits again: the reader reads on, the state it opened on. + await writer.execute("INSERT INTO t VALUES (2)") + assert await reader.values("SELECT n FROM t") == [[1]] + await reader.close() + + again = await open(ns, ReadOnly(backend, False), reader=True) + assert await again.values("SELECT n FROM t ORDER BY n") == [[1], [2]] + await again.close() + await writer.save() + + +def listing(backend): + """As the object's ctx.state_keys(): every key, those not yet synced too.""" + return lambda: [*backend.stored, *(key for key, _ in backend.unconfirmed)] + + +async def packs_after_commits(keys): + """The pack parts left after 100 commits that each replace a row.""" + backend = Backend() + store = backend.request() + if keys: + store.keys = listing(backend) + db = await open(fresh(), store) + await db.execute("CREATE TABLE t(id INTEGER PRIMARY KEY, b BLOB)") + for _ in range(100): + await db.execute("INSERT OR REPLACE INTO t VALUES(1, randomblob(3000))") + await db.save() + return sum(key.startswith("v22/pack/") for key in backend.stored) + + +async def test_a_store_that_lists_its_keys_is_collected_and_one_that_cannot_keeps_its_packs(): + swept, kept = await packs_after_commits(True), await packs_after_commits(False) + assert swept < 5 and kept > 10, (swept, kept) + + +def raises(): + raise RuntimeError("listing failed") + + +async def rejects(): + raise RuntimeError("listing failed") + + +@pytest.mark.parametrize("keys", [raises, rejects]) +async def test_a_listing_that_fails_ends_the_session_after_the_commit_that_collects_which_landed(keys): + backend = Backend() + ns = fresh() + store = backend.request() + store.keys = keys + db = await open(ns, store) + await db.execute("CREATE TABLE t(b BLOB)") + committed = 0 + # A commit that checkpoints collects: once the tail holds a few rows. It + # returns, and the next call fails with the listing's error. + with pytest.raises(RuntimeError, match="listing failed"): + for _ in range(40): + await db.execute("INSERT INTO t VALUES (randomblob(3000))") + committed += 1 + assert not db.is_usable() + reopened = await open(ns, backend.request()) + assert await reopened.values("SELECT count(*) FROM t") == [[committed]] From fdb9e6f374e3054d223c92e15a36b935f6b236a6 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 14:37:46 +0100 Subject: [PATCH 04/45] Add the selector: one loop for an invocation's concurrent waits Without journal replay, a handler awaits and combines futures as it likes (asyncio.gather, asyncio.wait, tasks of its own), but the protocol has one await point per invocation: a suspension, the awaiting-on hint and a cancellation each take one unresolved future. The SDK drives each awaited future with a do_progress loop of its own, each reading the input; the selector registers every wait and asks the VM for progress on the first of all of them to complete, as sdk-sessions' selector.ts does. hook(context) makes a ServerInvocationContext's waits go through it, and starts each future when it is made. Ported from sqlexperiment's restate_fns, without its version pin (the context is of this package). Nothing uses it yet: session handlers will. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/sessions/_selector.py | 325 +++++++++++++++++++++++++++ tests/sessions_selector.py | 231 +++++++++++++++++++ 2 files changed, 556 insertions(+) create mode 100644 python/restate/sessions/_selector.py create mode 100644 tests/sessions_selector.py diff --git a/python/restate/sessions/_selector.py b/python/restate/sessions/_selector.py new file mode 100644 index 00000000..26619065 --- /dev/null +++ b/python/restate/sessions/_selector.py @@ -0,0 +1,325 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +One loop that drives every pending wait of an invocation without journal replay. + +Without journal replay, the handler's futures are plain awaitables that user +code awaits and combines as it likes: ``asyncio.gather``, ``asyncio.wait``, +its own tasks, or a library's. The protocol has one await point per +invocation, though: a suspension, the awaiting-on hint and a cancellation each +take one unresolved future. So every wait is registered with the selector, +which asks the VM for progress on the first of all of them to complete, and is +the only one reading the invocation's input. +""" + +from __future__ import annotations + +import asyncio +import logging +import typing +from collections.abc import Callable +from typing import Any, Protocol + +from restate.exceptions import SdkInternalException, SuspendedException, TerminalError +from restate.server_context import DisconnectedException, ServerInvocationContext +from restate.vm import ( + AllCompletedUnresolvedFuture, + DoProgressAnyCompleted, + DoProgressCancelSignalReceived, + DoProgressExecuteRun, + DoProgressWaitExternalProgress, + FirstCompletedUnresolvedFuture, + SingleUnresolvedFuture, + Suspended, + UnresolvedFuture, +) + +logger = logging.getLogger(__name__) + + +class SelectorDriver(Protocol): + """What the selector needs of the invocation.""" + + def do_progress(self, future: UnresolvedFuture) -> Any: + """The VM's progress on ``future``: a DoProgress* response, Suspended, or an exception.""" + ... + + def is_completed(self, handle: int) -> bool: + """Whether the VM has the result of ``handle``.""" + ... + + async def flush_output(self) -> None: + """Writes out what the VM has buffered.""" + ... + + async def receive(self) -> dict[str, Any]: + """The next input event: a chunk of the request body, a disconnect, or a restate.* event.""" + ... + + def wake(self) -> None: + """Makes ``receive`` return a restate.* event.""" + ... + + def execute_run(self, handle: int) -> None: + """Starts the run of ``handle``, which signals its completion through ``receive``.""" + ... + + def notify_input(self, chunk: dict[str, Any]) -> None: + """Passes a chunk of the request body to the VM.""" + ... + + def track(self, task: asyncio.Future) -> None: + """Cancels ``task`` at the end of the attempt.""" + ... + + def on_end(self) -> None: + """The attempt cannot go on: the selector has failed every wait.""" + ... + + +def is_ready(future: UnresolvedFuture, completed: Callable[[int], bool]) -> bool: + """Whether ``future`` needs no more progress.""" + if isinstance(future, SingleUnresolvedFuture): + return completed(future.handle) + if isinstance(future, FirstCompletedUnresolvedFuture): + return any(is_ready(child, completed) for child in future.children) + if isinstance(future, AllCompletedUnresolvedFuture): + return all(is_ready(child, completed) for child in future.children) + raise TypeError(f"Unknown UnresolvedFuture variant: {type(future).__name__}") + + +class FutureSelector: + """ + The one loop of an invocation without journal replay. + + A wait stays registered until it is ready. The loop settles whatever is + ready, then asks the VM for progress on the first of all the pending waits + to complete: so a suspension waits for any of them, a pending run starts + whichever wait reaches it first, and a cancellation reaches all of them. A + wait registered while the loop reads the input wakes it, so it is part of + the next check. The loop runs while something is pending. + """ + + def __init__(self, driver: SelectorDriver, disconnected: Callable[[], BaseException]) -> None: + self._driver = driver + self._disconnected = disconnected + self._waits: dict[asyncio.Future, UnresolvedFuture] = {} + self._task: asyncio.Task | None = None + # Between the VM's WaitExternalProgress and the next input. + self._waiting = False + # Why the attempt ended, once it has: each wait then raises a new one. + self._end: Callable[[], BaseException] | None = None + + @property + def ended(self) -> BaseException | None: + """Why the attempt ended (a suspension, an error, a disconnect), or None while it goes on.""" + return None if self._end is None else self._end() + + @property + def size(self) -> int: + """The waits still pending, for tests.""" + return len(self._waits) + + async def wait(self, future: UnresolvedFuture) -> None: + """Returns once ``future`` is ready; raises if the attempt ends first, or is cancelled.""" + if self._end is not None: + raise self._end() + if is_ready(future, self._driver.is_completed): + return + waiter = asyncio.get_running_loop().create_future() + self._waits[waiter] = future + self._start() + try: + await waiter + finally: + self._waits.pop(waiter, None) + + def _start(self) -> None: + if self._task is None: + self._task = asyncio.create_task(self._loop()) + self._driver.track(self._task) + elif self._waiting: + # Only a loop waiting for input needs waking, once: any other picks + # up the new wait on its next pass. + self._waiting = False + self._driver.wake() + + def _settle(self) -> None: + for waiter, future in self._waits.items(): + if not waiter.done() and is_ready(future, self._driver.is_completed): + waiter.set_result(None) + + def _fail_all(self, error: Callable[[], BaseException]) -> None: + for waiter in list(self._waits): + if not waiter.done(): + waiter.set_exception(error()) + + def _finish(self, end: Callable[[], BaseException]) -> None: + """The attempt cannot go on: every wait, now and later, raises ``end()``.""" + self._end = end + self._fail_all(end) + self._driver.on_end() + + async def _loop(self) -> None: + driver = self._driver + try: + while True: + self._settle() + # Let what that settled run, and register what it waits for + # next, so that one check covers them. An optimization: what + # registers later wakes the loop anyway. + await asyncio.sleep(0) + pending = [future for waiter, future in self._waits.items() if not waiter.done()] + if not pending: + return + await driver.flush_output() + response = driver.do_progress( + pending[0] if len(pending) == 1 else FirstCompletedUnresolvedFuture(pending) + ) + if isinstance(response, BaseException): + logger.exception("Exception in do_progress", exc_info=response) + self._finish(_internal(response)) + return + if isinstance(response, Suspended): + self._finish(SuspendedException) + return + if isinstance(response, DoProgressAnyCompleted): + continue + if isinstance(response, DoProgressCancelSignalReceived): + # Reaches every pending wait; what waits after it is not cancelled. + self._fail_all(lambda: TerminalError("cancelled", 409)) + continue + if isinstance(response, DoProgressExecuteRun): + driver.execute_run(response.handle) + continue + if isinstance(response, DoProgressWaitExternalProgress): + self._waiting = True + await driver.flush_output() + chunk = await driver.receive() + self._waiting = False + kind = chunk.get("type") + if kind == "http.disconnect": + self._finish(self._disconnected) + return + if kind == "http.request": + driver.notify_input(chunk) + except asyncio.CancelledError: + # The attempt is over. + self._end = self._end or SuspendedException + self._fail_all(self._end) + raise + except BaseException as error: # pylint: disable=broad-exception-caught + logger.exception("Exception in the invocation's selector") + self._finish(_internal(error)) + finally: + self._task = None + + +def _internal(cause: BaseException) -> Callable[[], BaseException]: + """New SdkInternalExceptions, caused by ``cause``.""" + + def internal() -> BaseException: + error = SdkInternalException() + error.__cause__ = cause + return error + + return internal + + +# The selector takes over an invocation of the SDK's (``hook``): the waits of +# its futures, which the SDK makes through ``create_poll_or_cancel_coroutine``, +# go to the selector, and each future starts when it is made. It reaches into +# ``ServerInvocationContext``, of this same package, whose tests cover it. +# TODO: once sessions run on it, the context could drive its waits through +# the selector itself, rather than be hooked. + + +class _ContextDriver: + """What the selector needs of the SDK's context.""" + + def __init__(self, context: ServerInvocationContext, on_end: Callable[[], None]) -> None: + self.context = context + self.on_end = on_end + + def do_progress(self, future: UnresolvedFuture) -> Any: + return self.context.vm.do_progress(future) + + def is_completed(self, handle: int) -> bool: + return self.context.vm.is_completed(handle) + + async def flush_output(self) -> None: + await self.context.take_and_send_output() + + async def receive(self) -> dict[str, Any]: + return typing.cast(dict[str, Any], await self.context.receive()) + + def wake(self) -> None: + # Any restate.* event: the loop takes it as a wake-up. + event = self.context.receive.enqueue_restate_event({"type": "restate.run_completed", "data": None}) + self.track(asyncio.ensure_future(event)) + + def execute_run(self, handle: int) -> None: + context = self.context + fn = context.run_coros_to_execute.pop(handle) + + async def run() -> None: + try: + await fn() + finally: + await context.receive.enqueue_restate_event({"type": "restate.run_completed", "data": None}) + + self.track(asyncio.create_task(run())) + + def notify_input(self, chunk: dict[str, Any]) -> None: + # Skip empty body frames, as the SDK does. + body = chunk.get("body", None) + if body is not None and len(body) > 0: + self.context.vm.notify_input(body) + if not chunk.get("more_body", False): + self.context.vm.notify_input_closed() + + def track(self, task: asyncio.Future) -> None: + self.context.tasks.add(task) + + +def hook(context: ServerInvocationContext, on_end: Callable[[], None]) -> FutureSelector: + """ + Make the waits of ``context``, an invocation without journal replay, go + through one selector, which calls ``on_end`` if it fails them all; each + future starts when it is made, whether or not anything awaits it. The + hook lasts as long as ``context``, one attempt. + """ + selector = FutureSelector(_ContextDriver(context, on_end), DisconnectedException) + started: set[asyncio.Task] = set() + + def starting(make: Callable[..., Any]) -> Callable[..., Any]: + def make_started(*args: Any, **kwargs: Any) -> Any: + future = make(*args, **kwargs) + lazy = future.future + if lazy.task is None: + lazy.task = asyncio.create_task(lazy.async_def()) + # Its error, if nothing awaits it, is dropped. + lazy.task.add_done_callback(_retrieve) + started.add(lazy.task) + lazy.task.add_done_callback(started.discard) + return future + + return make_started + + for name in ("create_future", "create_sleep_future", "create_call_future"): + setattr(context, name, starting(getattr(context, name))) + setattr(context, "create_poll_or_cancel_coroutine", selector.wait) + return selector + + +def _retrieve(task: asyncio.Future) -> None: + if not task.cancelled(): + task.exception() diff --git a/tests/sessions_selector.py b/tests/sessions_selector.py new file mode 100644 index 00000000..33c7254f --- /dev/null +++ b/tests/sessions_selector.py @@ -0,0 +1,231 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""The selector of an invocation without journal replay, against a fake VM.""" + +# pylint: disable=missing-function-docstring + +import asyncio +from typing import Any + +import pytest +from restate.exceptions import SdkInternalException, SuspendedException, TerminalError +from restate.vm import ( + AllCompletedUnresolvedFuture, + DoProgressAnyCompleted, + DoProgressCancelSignalReceived, + DoProgressExecuteRun, + DoProgressWaitExternalProgress, + FirstCompletedUnresolvedFuture, + SingleUnresolvedFuture, + Suspended, + UnresolvedFuture, +) +from restate.sessions._selector import FutureSelector, is_ready + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +pytestmark = [pytest.mark.anyio] + + +class Disconnected(Exception): + pass + + +class FakeVm: + """ + A VM whose handles complete as input names them: a body chunk b"3" completes + handle 3. + """ + + def __init__(self) -> None: + self.completed: set[int] = set() + self.input: asyncio.Queue[dict[str, Any]] = asyncio.Queue() + self.checks: list[UnresolvedFuture] = [] + # Responses that do_progress gives before looking at the handles. + self.script: list[Any] = [] + self.runs: list[int] = [] + self.wakes = 0 + self.ended = 0 + self.tasks: list[asyncio.Future] = [] + + def do_progress(self, future: UnresolvedFuture) -> Any: + self.checks.append(future) + if self.script: + return self.script.pop(0) + if is_ready(future, self.is_completed): + return DoProgressAnyCompleted() + return DoProgressWaitExternalProgress() + + def is_completed(self, handle: int) -> bool: + return handle in self.completed + + async def flush_output(self) -> None: + pass + + async def receive(self) -> dict[str, Any]: + return await self.input.get() + + def wake(self) -> None: + self.wakes += 1 + self.input.put_nowait({"type": "restate.wake"}) + + def execute_run(self, handle: int) -> None: + self.runs.append(handle) + self.completed.add(handle) + self.input.put_nowait({"type": "restate.run_completed"}) + + def notify_input(self, chunk: dict[str, Any]) -> None: + self.completed.add(int(chunk["body"])) + + def track(self, task: asyncio.Future) -> None: + self.tasks.append(task) + + def on_end(self) -> None: + self.ended += 1 + + def complete(self, handle: int) -> None: + self.input.put_nowait({"type": "http.request", "body": str(handle).encode(), "more_body": True}) + + +def single(handle: int) -> SingleUnresolvedFuture: + return SingleUnresolvedFuture(handle) + + +async def settle() -> None: + for _ in range(5): + await asyncio.sleep(0) + + +async def test_a_ready_future_needs_no_loop(): + vm = FakeVm() + vm.completed.add(1) + selector = FutureSelector(vm, Disconnected) + await selector.wait(single(1)) + assert vm.checks == [] + + +async def test_concurrent_waits_share_one_await_point(): + vm = FakeVm() + selector = FutureSelector(vm, Disconnected) + both = asyncio.ensure_future(asyncio.gather(selector.wait(single(1)), selector.wait(single(2)))) + await settle() + # One check, of the first of both to complete. + assert vm.checks == [FirstCompletedUnresolvedFuture([single(1), single(2)])] + vm.complete(2) + await settle() + assert not both.done() + vm.complete(1) + await both + assert selector.size == 0 + + +async def test_a_wait_made_while_the_loop_waits_for_input_joins_the_next_check(): + vm = FakeVm() + selector = FutureSelector(vm, Disconnected) + first = asyncio.ensure_future(selector.wait(single(1))) + await settle() + assert vm.checks == [single(1)] + second = asyncio.ensure_future(selector.wait(single(2))) + await settle() + assert vm.wakes == 1 + assert vm.checks[-1] == FirstCompletedUnresolvedFuture([single(1), single(2)]) + # A third wait while the loop waits again wakes it once more, not once per wait. + third = asyncio.ensure_future(selector.wait(single(3))) + await settle() + assert vm.wakes == 2 + for handle in (3, 2, 1): + vm.complete(handle) + await asyncio.gather(first, second, third) + + +async def test_combinators_wait_for_their_children(): + vm = FakeVm() + selector = FutureSelector(vm, Disconnected) + every = asyncio.ensure_future(selector.wait(AllCompletedUnresolvedFuture([single(1), single(2)]))) + vm.complete(1) + await settle() + assert not every.done() + vm.complete(2) + await every + + +async def test_a_cancellation_reaches_every_pending_wait_and_only_those(): + vm = FakeVm() + selector = FutureSelector(vm, Disconnected) + waits = [asyncio.ensure_future(selector.wait(single(h))) for h in (1, 2)] + await settle() + vm.script.append(DoProgressCancelSignalReceived()) + vm.complete(9) + for wait in waits: + with pytest.raises(TerminalError) as error: + await wait + assert error.value.status_code == 409 + # What waits after it is not cancelled. + later = asyncio.ensure_future(selector.wait(single(3))) + vm.complete(3) + await later + + +async def test_a_suspension_ends_every_wait_now_and_later(): + vm = FakeVm() + selector = FutureSelector(vm, Disconnected) + vm.script.append(Suspended()) + with pytest.raises(SuspendedException): + await selector.wait(single(1)) + assert isinstance(selector.ended, SuspendedException) + assert vm.ended == 1 + with pytest.raises(SuspendedException): + await selector.wait(single(2)) + + +async def test_an_error_of_the_vm_ends_the_attempt(): + vm = FakeVm() + selector = FutureSelector(vm, Disconnected) + vm.script.append(ValueError("journal mismatch")) + with pytest.raises(SdkInternalException) as error: + await selector.wait(single(1)) + assert isinstance(error.value.__cause__, ValueError) + + +async def test_a_disconnect_ends_the_attempt(): + vm = FakeVm() + selector = FutureSelector(vm, Disconnected) + wait = asyncio.ensure_future(selector.wait(single(1))) + await settle() + vm.input.put_nowait({"type": "http.disconnect"}) + with pytest.raises(Disconnected): + await wait + + +async def test_a_run_is_started_and_its_completion_wakes_the_loop(): + vm = FakeVm() + selector = FutureSelector(vm, Disconnected) + vm.script.append(DoProgressExecuteRun(7)) + await selector.wait(single(7)) + assert vm.runs == [7] + + +async def test_the_loop_stops_when_nothing_is_pending_and_starts_again(): + vm = FakeVm() + selector = FutureSelector(vm, Disconnected) + first = asyncio.ensure_future(selector.wait(single(1))) + vm.complete(1) + await first + await settle() + assert all(task.done() for task in vm.tasks) + second = asyncio.ensure_future(selector.wait(single(2))) + vm.complete(2) + await second + assert len(vm.tasks) == 2 From 6449697fd99b24b2b7064bfd762e7e8eadc729a3 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 14:52:24 +0100 Subject: [PATCH 05/45] Add the schema and the calls of session objects _schema.py: the steps that make each key's database, applied once per key, with how far each schema got in the database's record `schema`, so a step added later runs once on each database; drop_all, for reset. As sdk-sessions' schema.ts. _calls.py: ctx.rpc and ctx.send (and a transaction's tx.send), typed by the handler they target: ctx.rpc(handler, key=...)(input) returns the call's awaitable, its result read by the handler's output serde. An object's or a workflow's handler needs its key, a service's has none. generic() calls or sends by name, for a target known only at run time. Nothing uses them yet. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/sessions/_calls.py | 273 +++++++++++++++++++++++++++++ python/restate/sessions/_schema.py | 177 +++++++++++++++++++ 2 files changed, 450 insertions(+) create mode 100644 python/restate/sessions/_calls.py create mode 100644 python/restate/sessions/_schema.py diff --git a/python/restate/sessions/_calls.py b/python/restate/sessions/_calls.py new file mode 100644 index 00000000..825313b1 --- /dev/null +++ b/python/restate/sessions/_calls.py @@ -0,0 +1,273 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +Calls and sends from a session object's handler: ``ctx.rpc`` and ``ctx.send``, +and a transaction's ``tx.send``. A target is a handler (with its key for an +object or a workflow), typed by its signature; or, for one known only at run +time, its names, with ``generic``. Options go with the target: what it returns +takes only the input:: + + count = await ctx.rpc(increment, key="a", idempotency_key=f"{ctx.request.id}:inc")(5) + ctx.send(record, key="audit", delay=timedelta(hours=1))(entry) + answer = await ctx.rpc.generic("Counter", "get", key="a") +""" + +from __future__ import annotations + +import typing +from collections.abc import Awaitable, Callable, Generator +from dataclasses import dataclass +from datetime import timedelta +from typing import Any, Protocol + +from restate.context import RestateDurableCallFuture +from restate.exceptions import TerminalError +from restate.handler import Handler, handler_from_callable +from restate.serde import JsonSerde, Serde + +I = typing.TypeVar("I") +O = typing.TypeVar("O") + + +@dataclass(frozen=True) +class GenericSend: + """A send, its input serialized: as ``ctx.send`` makes it, or a transaction's outbox keeps it.""" + + service: str + handler: str + parameter: bytes + key: str | None = None + delay: timedelta | None = None + idempotency_key: str | None = None + headers: dict[str, str] | None = None + + +class CallHost(Protocol): + """What calls and sends need of the SDK's context.""" + + def generic_call( + self, + service: str, + handler: str, + arg: bytes, + key: str | None = None, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> RestateDurableCallFuture[bytes]: ... + + def generic_send( + self, + service: str, + handler: str, + arg: bytes, + key: str | None = None, + send_delay: timedelta | None = None, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> Any: ... + + +def _target(handler: Callable[..., Any], key: str | None) -> tuple[str, str, Serde[Any], Serde[Any]]: + """ + The service, handler name and serdes of ``handler``, called with ``key``. An + object's or a workflow's handler needs its key, and a service's has none: + a ``TerminalError`` otherwise, since a retry would fail the same way. + """ + try: + target: Handler[Any, Any] = handler_from_callable(handler) + except ValueError: + raise TerminalError(f"{handler!r} is not a handler of a service, an object or a workflow") from None + service = target.service_tag.name + keyed = target.service_tag.kind != "service" + if keyed and key is None: + what = "a workflow" if target.service_tag.kind == "workflow" else "an object" + raise TerminalError(f"{service} is {what}: call it with its key") + if not keyed and key is not None: + raise TerminalError(f"{service} is a service: call it without a key") + return service, target.name, target.handler_io.input_serde, target.handler_io.output_serde + + +class CallFuture(RestateDurableCallFuture[O]): + """The result of a call, as its output serde reads it: awaitable, with the call's invocation id.""" + + def __init__(self, call: RestateDurableCallFuture[bytes], serde: Serde[O]) -> None: + super().__init__() + self._call = call + self._serde = serde + + async def _result(self) -> O: + return typing.cast(O, self._serde.deserialize(typing.cast(bytes, await self._call))) + + def __await__(self) -> Generator[Any, Any, O]: + return self._result().__await__() + + async def invocation_id(self) -> str: + return await self._call.invocation_id() + + async def cancel_invocation(self) -> None: + await self._call.cancel_invocation() + + +class Rpc: + """ + Calls: ``ctx.rpc(handler, key=...)`` returns a function of the handler's + input, if it takes one, which makes the call and returns its result's + awaitable. A call is made when that function is called, awaited or not. + """ + + def __init__(self, host: CallHost) -> None: + self._host = host + + @typing.overload + def __call__( + self, + handler: Callable[[Any], Awaitable[O]], + key: str | None = None, + *, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> Callable[[], CallFuture[O]]: ... + + @typing.overload + def __call__( + self, + handler: Callable[[Any, I], Awaitable[O]], + key: str | None = None, + *, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> Callable[[I], CallFuture[O]]: ... + + def __call__( + self, + handler: Callable[..., Awaitable[Any]], + key: str | None = None, + *, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> Callable[..., CallFuture[Any]]: + """ + A caller of ``handler``, with ``key`` for an object's or a workflow's. + + ``idempotency_key`` is the same on every attempt, so that a new attempt + gets the first call's result; without one, every attempt makes the call + again, the one that resumes a suspension too. Restate scopes it to the + target handler. + """ + service, name, input_serde, output_serde = _target(handler, key) + + def call(arg: Any = None) -> CallFuture[Any]: + parameter = input_serde.serialize(arg) + future = self._host.generic_call(service, name, parameter, key, idempotency_key, headers) + return CallFuture(future, output_serde) + + return call + + def generic( + self, + service: str, + handler: str, + arg: Any = None, + *, + key: str | None = None, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + input_serde: Serde[Any] = JsonSerde(), + output_serde: Serde[Any] = JsonSerde(), + ) -> CallFuture[Any]: + """ + Call a target known only at run time, by its names: untyped, and JSON + unless serdes are given. ``key`` is the caller's to give, for an + object or a workflow. + """ + future = self._host.generic_call(service, handler, input_serde.serialize(arg), key, idempotency_key, headers) + return CallFuture(future, output_serde) + + +class Send: + """ + Sends, as ``Rpc``'s calls: ``ctx.send(handler, key=...)`` returns a function + of the handler's input, if it takes one, which makes the send at once + (``ctx.send``), or records it for its transaction's commit (``tx.send``). + """ + + def __init__(self, sink: Callable[[GenericSend], None]) -> None: + self._sink = sink + + @typing.overload + def __call__( + self, + handler: Callable[[Any], Awaitable[Any]], + key: str | None = None, + *, + delay: timedelta | None = None, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> Callable[[], None]: ... + + @typing.overload + def __call__( + self, + handler: Callable[[Any, I], Awaitable[Any]], + key: str | None = None, + *, + delay: timedelta | None = None, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> Callable[[I], None]: ... + + def __call__( + self, + handler: Callable[..., Awaitable[Any]], + key: str | None = None, + *, + delay: timedelta | None = None, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> Callable[..., None]: + """ + A sender to ``handler``, with ``key`` for an object's or a workflow's; + ``delay`` is how long after the send (for ``tx.send``, after the + commit) the target runs. See ``Rpc`` for ``idempotency_key``. + """ + service, name, input_serde, _ = _target(handler, key) + + def send(arg: Any = None) -> None: + parameter = input_serde.serialize(arg) + self._sink(GenericSend(service, name, parameter, key, delay, idempotency_key, headers)) + + return send + + def generic( + self, + service: str, + handler: str, + arg: Any = None, + *, + key: str | None = None, + delay: timedelta | None = None, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + input_serde: Serde[Any] = JsonSerde(), + ) -> None: + """Send to a target known only at run time, by its names: untyped, and JSON unless a serde is given.""" + self._sink(GenericSend(service, handler, input_serde.serialize(arg), key, delay, idempotency_key, headers)) + + +def send_now(host: CallHost) -> Callable[[GenericSend], None]: + """The sink of ``ctx.send``: each send is made at once.""" + + def sink(send: GenericSend) -> None: + host.generic_send( + send.service, send.handler, send.parameter, send.key, send.delay, send.idempotency_key, send.headers + ) + + return sink diff --git a/python/restate/sessions/_schema.py b/python/restate/sessions/_schema.py new file mode 100644 index 00000000..cda8829b --- /dev/null +++ b/python/restate/sessions/_schema.py @@ -0,0 +1,177 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +The schema of a session object's databases: each key's database is made by +the same steps, in order, applied once per key, as its first exclusive attempt +opens it. A database keeps how far each schema got in its record ``schema`` +(beside its tables, in none of them), committed with the steps, so a step +added later (a column, an index) runs on each key's database once, as that key +is next opened:: + + notes = sessions.SessionObject("notes", schema=[ + "CREATE TABLE notes (text TEXT NOT NULL)", + "ALTER TABLE notes ADD COLUMN at TEXT", + ]) + +A step is a statement, or an async function of the transaction (to move data, +say). Steps are never changed or removed once they ran: a change is a step of +its own. + +A view applies the steps its key's database lacks in its own session, which +stores nothing, then reads: a key never written reads as empty tables. +Libraries (``restate.agents``) keep their tables in a schema of their own +name, which ``schema.ensure(ctx.db)`` applies where they use them, in any +object. +""" + +from __future__ import annotations + +import json +import typing +from collections.abc import Awaitable, Callable, Sequence +from typing import Any + +from restate.sessions._database import Connection, ConnectionTransaction + +RECORD = "schema" +""" +The database's record of its schemas: how many steps of each it has, by name. +One record, so that ``reset`` clears all of them, a library's this process +never loaded too. +""" + +Step = typing.Union[str, Callable[[Any], Awaitable[None]]] +"""A statement, or an async function of the transaction.""" + +_ENSURE = "_restate_sessions_ensure" +"""The attribute of a database of this package's: how it applies schemas (see ``Schema.ensure``).""" + + +class Schema: + """A part of a database's schema: its ``name``, and the ``steps`` that make it, in order.""" + + def __init__(self, name: str, steps: Sequence[Step]) -> None: + self.name = name + self.steps: tuple[Step, ...] = tuple(steps) + + async def ensure(self, db: Any) -> None: + """ + Apply the steps ``db`` lacks: a library's tables, where it uses them, in + any object. On a handler's ``ctx.db``, once per attempt; a view's are + its own, vanishing with its session. + """ + ensure = ensure_of(db) + await (ensure([self]) if ensure is not None else apply(db, [self])) + + +Ensure = Callable[[Sequence[Schema]], Awaitable[None]] +"""How a database applies schemas: a handler's ``ctx.db``, its reader, or one of ``as_database``.""" + + +def ensuring(db: Any, ensure: Ensure) -> Any: + """``db``, which applies schemas with ``ensure`` (see ``Schema.ensure``).""" + object.__setattr__(db, _ENSURE, ensure) + return db + + +def ensure_of(db: Any) -> Ensure | None: + """How ``db`` applies schemas, if it is a database of this package's.""" + return getattr(db, _ENSURE, None) + + +async def _applied_in(db: Connection | ConnectionTransaction) -> dict[str, int]: + """How many steps of each schema the database has, by name.""" + record = await db.record(RECORD) + return json.loads(record.decode()) if record else {} + + +def _behind(applied: dict[str, int], schemas: Sequence[Schema]) -> list[Schema]: + """The schemas of ``schemas`` that ``applied`` lacks steps of.""" + return [schema for schema in schemas if applied.get(schema.name, 0) < len(schema.steps)] + + +async def apply(db: Connection, schemas: Sequence[Schema], *, reader: bool = False) -> None: + """Apply the steps of ``schemas`` that ``db`` lacks, in one transaction, with how far each got.""" + if not _behind(await _applied_in(db), schemas): + return + # Again in the transaction, which holds the session: a call before it may + # have applied them. + async with db.transaction() as tx: + if not reader: + await apply_in(tx, schemas) + return + # A reader's session is query-only, but for its schemas. + await tx.execute("PRAGMA query_only = 0") + try: + await apply_in(tx, schemas) + finally: + await tx.execute("PRAGMA query_only = 1") + + +async def apply_in(tx: ConnectionTransaction, schemas: Sequence[Schema]) -> None: + """Apply the steps of ``schemas`` that the transaction ``tx`` lacks, and record how far each got.""" + applied = await _applied_in(tx) + lacking = _behind(applied, schemas) + if not lacking: + return + for schema in lacking: + for step in schema.steps[applied.get(schema.name, 0) :]: + await (tx.execute(step) if isinstance(step, str) else step(tx)) + applied[schema.name] = len(schema.steps) + tx.set_record(RECORD, json.dumps(applied).encode()) + + +async def drop_all(tx: ConnectionTransaction, schemas: Sequence[Schema]) -> None: + """ + Empty the database in the transaction ``tx``: drop every table, view, index + and trigger, and forget every schema, then apply ``schemas`` again. Its + other records (the outbox) stay. It commits with ``tx``: what ``tx`` writes + after it are the new database's first rows. + """ + made = await tx.values("SELECT type, name FROM sqlite_schema WHERE substr(name, 1, 7) != 'sqlite_'") + # A trigger or view may name a table; an index goes with its table. + for kind in ("trigger", "view"): + for made_kind, name in made: + if made_kind == kind: + await _drop(tx, kind, typing.cast(str, name)) + tables = [typing.cast(str, name) for kind, name in made if kind == "table"] + for name in await _children_first(tx, tables): + await _drop(tx, "table", name) + tx.set_record(RECORD, b"") + await apply_in(tx, schemas) + + +async def _drop(tx: ConnectionTransaction, kind: str, name: str) -> None: + quoted = name.replace('"', '""') + await tx.execute(f'DROP {kind.upper()} IF EXISTS "{quoted}"') + + +async def _children_first(tx: ConnectionTransaction, tables: list[str]) -> list[str]: + """ + ``tables``, each after every table whose foreign keys name it: with foreign + keys on, dropping a table deletes its rows, which a child's rows may not + outlive. Tables in a cycle stay in their order. + """ + parents: dict[str, set[str]] = {} + for name in tables: + keys = await tx.values('SELECT "table" FROM pragma_foreign_key_list(?)', [name]) + # SQL names are case-insensitive. + own = name.lower() + parents[name] = {typing.cast(str, table).lower() for (table,) in keys} - {own} + ordered: list[str] = [] + left = tables + while left: + named = {parent for name in left for parent in parents[name]} + unnamed = [name for name in left if name.lower() not in named] + step = unnamed or left + ordered.extend(step) + left = [name for name in left if name not in step] + return ordered From c794d5c59f30e87236d84656b86015d50e4b51c1 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 14:52:24 +0100 Subject: [PATCH 06/45] Add session objects: restate.sessions.SessionObject The Python variant of @restatedev/sessions: a virtual object with a SQL database per key, ctx.db, kept in the object's state. notes = sessions.SessionObject("notes", schema=[...]) @notes.handler() # exclusive: SessionContext(key, request, db, rpc, send, cancel) @notes.view() # shared, read-only: ViewContext(key, request, db, rpc, cancel) Handlers run with explicit recovery: every attempt runs from its first line against what the attempts before it committed. Each attempt hooks the selector, so its waits (SQL, calls) can be awaited together. A transaction's sends (tx.send) are kept in the outbox record and made with its commit, each with an idempotency key of the transaction, and made again by the next attempt if a crash left them there. The outbox format is sdk-sessions'. _session.py is session.ts's Attempt; _object.py its object.ts, over the SDK's VirtualObject (lazy state, metadata hiding the database's entries). Not yet: preloading only the database's head, explicit cancellation, and the scope that @restatedev/agents keeps per attempt. Tests: attempts over an in-memory state (sessions_object.py), and a real Restate server (sessions_harness.py). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/sessions/__init__.py | 61 ++- python/restate/sessions/_object.py | 433 +++++++++++++++ python/restate/sessions/_session.py | 815 ++++++++++++++++++++++++++++ tests/sessions_harness.py | 146 +++++ tests/sessions_object.py | 403 ++++++++++++++ 5 files changed, 1857 insertions(+), 1 deletion(-) create mode 100644 python/restate/sessions/_object.py create mode 100644 python/restate/sessions/_session.py create mode 100644 tests/sessions_harness.py create mode 100644 tests/sessions_object.py diff --git a/python/restate/sessions/__init__.py b/python/restate/sessions/__init__.py index 7a6a8b62..3f13572e 100644 --- a/python/restate/sessions/__init__.py +++ b/python/restate/sessions/__init__.py @@ -10,5 +10,64 @@ # """ Session objects: virtual objects with a SQL database per key, kept in the -object's state (``ctx.db``). The Python variant of ``@restatedev/sessions``. +object's state (``ctx.db``). The Python variant of ``@restatedev/sessions``:: + + import restate + from restate import sessions + + notes = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) + + + @notes.handler() + async def add(ctx: sessions.SessionContext, text: str) -> None: + await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [text]) + + + @notes.view() + async def list(ctx: sessions.ViewContext) -> list[str]: + return [row["text"] for row in await ctx.db.query("SELECT text FROM notes ORDER BY rowid")] + + + app = restate.app([notes]) + +Session handlers don't replay a journal: a retried attempt runs again from its +first line, against what the attempts before it committed. A session object +serves beside the SDK's services, objects and workflows, which it calls with +``ctx.rpc`` for journaled work. """ + +from restate.sessions._calls import CallFuture, Rpc, Send +from restate.sessions._database import Row, Rows, SqlError, SqlParam, SqlValue +from restate.sessions._object import Cancel, SessionContext, SessionObject, ViewContext +from restate.sessions._schema import Schema, Step +from restate.sessions._session import ( + Database, + DatabaseReader, + DatabaseWriter, + Transaction, + database, + reset, +) + +__all__ = [ + "CallFuture", + "Cancel", + "Database", + "DatabaseReader", + "DatabaseWriter", + "Row", + "Rows", + "Rpc", + "Schema", + "Send", + "SessionContext", + "SessionObject", + "SqlError", + "SqlParam", + "SqlValue", + "Step", + "Transaction", + "ViewContext", + "database", + "reset", +] diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py new file mode 100644 index 00000000..17a30132 --- /dev/null +++ b/python/restate/sessions/_object.py @@ -0,0 +1,433 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +Session objects: Restate virtual objects whose state is a SQL database per +key. Its handlers run one at a time per key, its views beside them; each gets a +context of its own, made of the key, the request, the database and calls, and +nothing else of the SDK's context:: + + notes = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) + + @notes.handler() + async def add(ctx: sessions.SessionContext, text: str) -> None: + await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [text]) + + @notes.view() + async def list(ctx: sessions.ViewContext) -> list[str]: + return [row["text"] for row in await ctx.db.query("SELECT text FROM notes ORDER BY rowid")] +""" + +from __future__ import annotations + +import functools +import typing +from collections.abc import Awaitable, Callable, Sequence +from dataclasses import dataclass +from datetime import timedelta +from typing import Any + +from restate.context import Request +from restate.exceptions import SuspendedException, TerminalError +from restate.handler import RESTATE_UNIQUE_HANDLER_SYMBOL, Handler, handler_from_callable +from restate.object import VirtualObject +from restate.retry_policy import InvocationRetryPolicy +from restate.serde import BytesSerde, DefaultSerde, Serde +from restate.server_context import ServerInvocationContext +from restate.sessions import _selector +from restate.sessions._calls import CallFuture, Rpc, Send, send_now +from restate.sessions._database import Store +from restate.sessions._schema import Schema, Step +from restate.sessions._session import Attempt, Database, DatabaseReader, ObjectConfig, SqlHost, reader_of, with_attempt + +F = typing.TypeVar("F", bound=Callable[..., Awaitable[Any]]) + +HEAD_KEY = "v22/head" +"""The key of the object database's head, its first read as it opens (shared-core's state format).""" + +SQL_METADATA = { + "dev.restate.tooling.journal.hidden.names": "sync", + "dev.restate.tooling.journal.hidden.name.prefix": "v22/", +} +""" +The metadata every session object has: Restate's UI hides the database from +the journal view (its head, tail, manifests and packs in the object's state, +``v22/...``; and the ``sync`` run after its writes). +""" + +Cancel = Callable[[str], None] +""" +Cancel invocation ``invocation_id`` (such as ``inv_1f…``), as Restate's admin +API does: made at once, and not awaited. Cancelling an invocation that already +ended does nothing, so a new attempt that cancels it again is harmless. +""" + + +@dataclass(frozen=True) +class SessionContext: + """ + The context of a session object's handler. A handler runs in attempts: + Restate starts a new one after a crash, a retryable error or a suspension, + and each runs the handler from its first line, finding what the attempts + before it committed. Its operations can be awaited together + (``asyncio.gather``, ``asyncio.wait``). + """ + + key: str + """The session's key.""" + request: Request + """The request: its id (the same on every attempt), headers, and the attempt's.""" + db: Database + """ + The key's database. Each statement outside a transaction is its own commit, + and ``db.transaction()`` groups several; a commit is durable when its call + returns. Calls run one at a time, in call order. + """ + rpc: Rpc + """Calls, typed by the handler called: ``await ctx.rpc(handler, key="a")(input)``.""" + send: Send + """Sends, made at once; a transaction's (``tx.send``) are made once it commits.""" + cancel: Cancel + """Cancel an invocation, by its id: see ``Cancel``.""" + + +@dataclass(frozen=True) +class ViewContext: + """ + The context of a session object's view: it runs beside the handlers, and + reads the database as committed when its invocation started, whatever they + commit meanwhile. It writes nothing, and sends nothing, but it can cancel an + invocation: a handler of its own object's, say, as a Stop button does. + """ + + key: str + request: Request + db: DatabaseReader + rpc: Rpc + cancel: Cancel + + +class _ContextHost: + """ + An attempt's session's host (``SqlHost``) over the SDK's context: its waits + go through the selector, and its state commands are named after their keys, + so that tooling can hide the database by name. + """ + + def __init__(self, context: ServerInvocationContext, shared: bool) -> None: + self.context = context + self.shared = shared + self.service = context.handler.service_tag.name + self.key = context.key() + self._listeners: list[Callable[[], None]] = [] + self.selector = _selector.hook(context, self._notify_attempt_end) + attempt_finished = context.on_attempt_finished + + def on_attempt_finished() -> None: + self._notify_attempt_end() + attempt_finished() + + context.on_attempt_finished = on_attempt_finished # type: ignore[method-assign] + + def get(self, key: str) -> Awaitable[bytes | None]: + return self.context.get(key, serde=BytesSerde(), entry_name=key) + + def set(self, key: str, value: bytes) -> None: + self.context.set(key, value, serde=BytesSerde(), entry_name=key) + + def clear(self, key: str) -> None: + self.context.clear(key, entry_name=key) + + def state_keys(self) -> Awaitable[list[str]]: + # Named under the format's prefix, which tooling hides. + return self.context.state_keys(entry_name="v22/keys") + + def sync(self) -> Awaitable[Any]: + return self.context.run_typed("sync", _nothing) + + def generic_call( + self, + service: str, + handler: str, + arg: bytes, + key: str | None = None, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> Any: + return self.context.generic_call(service, handler, arg, key, idempotency_key, headers) + + def generic_send( + self, + service: str, + handler: str, + arg: bytes, + key: str | None = None, + send_delay: timedelta | None = None, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> Any: + return self.context.generic_send(service, handler, arg, key, send_delay, idempotency_key, headers) + + def attempt_end(self) -> BaseException | None: + ended = self.selector.ended + if ended is not None: + return ended + if self.context.request_finished_event.is_set(): + return SuspendedException() + return None + + def on_attempt_end(self, listener: Callable[[], None]) -> Callable[[], None]: + if self.attempt_end() is not None: + listener() + return lambda: None + self._listeners.append(listener) + + def cancel() -> None: + if listener in self._listeners: + self._listeners.remove(listener) + + return cancel + + def _notify_attempt_end(self) -> None: + listeners, self._listeners = self._listeners, [] + for listener in listeners: + listener() + + +async def _nothing() -> None: + """The action of an empty run.""" + + +def _cancel_of(host: Any) -> Cancel: + """``Cancel`` over the SDK's context.""" + + def cancel(invocation_id: str) -> None: + # Not an id: a retry would fail the same way. + if not invocation_id.startswith("inv"): + raise TerminalError(f"not an invocation id: {invocation_id}") + host.cancel_invocation(invocation_id) + + return cancel + + +@dataclass(frozen=True) +class _Entry: + """A handler or view of a session object: its function, kind, and object's config.""" + + fn: Callable[..., Awaitable[Any]] + view: bool + config: ObjectConfig + + +def _context_of(entry: _Entry, host: SqlHost, request: Request, attempt: Attempt) -> SessionContext | ViewContext: + """The context of an attempt of ``entry``.""" + rpc = Rpc(host) + cancel = _cancel_of(host) + if entry.view: + return ViewContext(key=host.key, request=request, db=reader_of(attempt.db), rpc=rpc, cancel=cancel) + return SessionContext( + key=host.key, request=request, db=attempt.db, rpc=rpc, send=Send(send_now(host)), cancel=cancel + ) + + +async def _attempt(entry: _Entry, host: SqlHost, request: Request, args: Sequence[Any]) -> Any: + """One attempt of ``entry`` over ``host``: its own context and session, which ends with it.""" + attempt = Attempt(host, entry.config) + attempt.open_now() + context = _context_of(entry, host, request, attempt) + return await with_attempt(lambda: entry.fn(context, *args), attempt) + + +def _sdk_handler(entry: _Entry) -> Callable[..., Awaitable[Any]]: + """``entry`` as the SDK's handler: one loop drives each attempt's awaits (the selector), until it ends.""" + + @functools.wraps(entry.fn) + async def handler(ctx: ServerInvocationContext, *args: Any) -> Any: + host = _ContextHost(ctx, shared=entry.view) + return await _attempt(entry, host, ctx.request(), args) + + return handler + + +async def invoke(obj: SessionObject, name: str, host: SqlHost, request: Request, *args: Any) -> Any: + """Run handler or view ``name`` of ``obj`` for one attempt over ``host``, a context of a test's own, without Restate.""" + entry = obj._entries.get(name) # pylint: disable=protected-access + if entry is None: + raise ValueError(f"{obj.name}.{name} is no handler or view of a session object") + return await _attempt(entry, host, request, args) + + +class SessionObject(VirtualObject): + """ + A virtual object whose state is a SQL database per key. + + - ``schema``: the steps that make each key's database, in order, applied + once per key by the first handler that uses it. A step that ran never + changes: a change is a step of its own. + - ``@obj.handler()``: handlers run one at a time per key, each with a + ``SessionContext``. + - ``@obj.view()``: views run beside the handlers, each with a + ``ViewContext``: they read, and write nothing. + + A name is a handler or a view, not both. The other options are the SDK's + virtual object's, but for its state options: a session object's state is + its database's. + """ + + def __init__( + self, + name: str, + *, + schema: Sequence[Step] = (), + description: str | None = None, + metadata: dict[str, str] | None = None, + inactivity_timeout: timedelta | None = None, + abort_timeout: timedelta | None = None, + journal_retention: timedelta | None = None, + idempotency_retention: timedelta | None = None, + ingress_private: bool | None = None, + invocation_retry_policy: InvocationRetryPolicy | None = None, + ) -> None: + # TODO: start an invocation with only the database's head + # (HEAD_KEY), once the SDK can say which keys to load eagerly: lazy + # state loads none, so the open's first read is a round trip. + super().__init__( + name, + description=description, + metadata={**SQL_METADATA, **(metadata or {})}, + inactivity_timeout=inactivity_timeout, + abort_timeout=abort_timeout, + journal_retention=journal_retention, + idempotency_retention=idempotency_retention, + enable_lazy_state=True, + ingress_private=ingress_private, + invocation_retry_policy=invocation_retry_policy, + ) + self.config = ObjectConfig(schema=Schema(name, schema) if schema else None) + self._entries: dict[str, _Entry] = {} + + def handler( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + self, + name: str | None = None, + *, + accept: str = "application/json", + content_type: str = "application/json", + input_serde: Serde[Any] = DefaultSerde(), + output_serde: Serde[Any] = DefaultSerde(), + metadata: dict[str, str] | None = None, + inactivity_timeout: timedelta | None = None, + abort_timeout: timedelta | None = None, + journal_retention: timedelta | None = None, + idempotency_retention: timedelta | None = None, + ingress_private: bool | None = None, + invocation_retry_policy: InvocationRetryPolicy | None = None, + ) -> Callable[[F], F]: + """ + A handler: it runs one at a time per key, and is the only one that + writes. Its function takes a ``SessionContext``, and its input, if any:: + + @notes.handler() + async def add(ctx: sessions.SessionContext, text: str) -> None: + await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [text]) + + The options are the SDK's handler's. + """ + return self._register( + False, + name, + dict( + accept=accept, + content_type=content_type, + input_serde=input_serde, + output_serde=output_serde, + metadata=metadata, + inactivity_timeout=inactivity_timeout, + abort_timeout=abort_timeout, + journal_retention=journal_retention, + idempotency_retention=idempotency_retention, + ingress_private=ingress_private, + invocation_retry_policy=invocation_retry_policy, + ), + ) + + def view( + self, + name: str | None = None, + *, + accept: str = "application/json", + content_type: str = "application/json", + input_serde: Serde[Any] = DefaultSerde(), + output_serde: Serde[Any] = DefaultSerde(), + metadata: dict[str, str] | None = None, + inactivity_timeout: timedelta | None = None, + abort_timeout: timedelta | None = None, + journal_retention: timedelta | None = None, + idempotency_retention: timedelta | None = None, + ingress_private: bool | None = None, + invocation_retry_policy: InvocationRetryPolicy | None = None, + ) -> Callable[[F], F]: + """ + A view: it runs beside the handlers, and reads the database as + committed when its invocation started. Its function takes a + ``ViewContext``, and its input, if any:: + + @notes.view() + async def count(ctx: sessions.ViewContext) -> int: + row = await ctx.db.get("SELECT count(*) AS n FROM notes") + return row["n"] if row else 0 + + A view that is still running after 16 commits of the handlers fails, + and Restate retries it. The options are the SDK's handler's. + """ + return self._register( + True, + name, + dict( + accept=accept, + content_type=content_type, + input_serde=input_serde, + output_serde=output_serde, + metadata=metadata, + inactivity_timeout=inactivity_timeout, + abort_timeout=abort_timeout, + journal_retention=journal_retention, + idempotency_retention=idempotency_retention, + ingress_private=ingress_private, + invocation_retry_policy=invocation_retry_policy, + ), + ) + + def _register(self, view: bool, name: str | None, options: dict[str, Any]) -> Callable[[F], F]: + def register(fn: F) -> F: + if RESTATE_UNIQUE_HANDLER_SYMBOL in vars(fn): + raise TypeError( + f"{self.name}.{fn.__name__} is an SDK handler: a session object's handlers are functions of its " + "context. Put journaled work in a restate.VirtualObject, Service or Workflow, and call it with ctx.rpc" + ) + entry = _Entry(fn=fn, view=view, config=self.config) + registered = VirtualObject.handler(self, name, kind="shared" if view else "exclusive", **options)( + _sdk_handler(entry) + ) + sdk: Handler[Any, Any] = handler_from_callable(registered) + # A later attempt makes different reads than a replayed journal + # would have: each runs from the start, against the state as it is. + sdk.explicit_recovery = True + self._entries[sdk.name] = entry + return typing.cast(F, registered) + + return register + + def with_store_wrapper(self, wrap: Callable[[Store], Store]) -> SessionObject: + """Wrap the stores of this object's attempts with ``wrap``: fault injection, in tests.""" + self.config.wrap_store = wrap + return self + + +__all__ = ["CallFuture", "Cancel", "SessionContext", "SessionObject", "ViewContext"] diff --git a/python/restate/sessions/_session.py b/python/restate/sessions/_session.py new file mode 100644 index 00000000..97eacbc2 --- /dev/null +++ b/python/restate/sessions/_session.py @@ -0,0 +1,815 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +An attempt's database: the session's database, kept in the object's state +(``_database.py``), opened as an exclusive attempt starts (a view's on its +first read) and saved when the handler ends; and its transactions, with the +sends each commits with (``tx.send``). +""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import re +import time +import typing +import uuid +from collections.abc import AsyncIterator, Awaitable, Callable, Sequence +from contextlib import asynccontextmanager +from dataclasses import dataclass, field +from datetime import timedelta +from typing import Any, Protocol +from urllib.parse import quote + +from restate.exceptions import SuspendedException, TerminalError +from restate.sessions._calls import CallHost, GenericSend, Send +from restate.sessions._database import ( + Connection, + ConnectionTransaction, + Row, + Rows, + SqlError, + SqlParam, + Store, + _objects, + open as open_connection, +) +from restate.sessions._schema import Schema, apply, drop_all, ensure_of, ensuring + +T = typing.TypeVar("T") + + +class SqlHost(CallHost, Protocol): + """What an attempt's session needs of the handler's context.""" + + @property + def shared(self) -> bool: + """A view's: it reads beside the exclusive handler, and stores nothing.""" + ... + + @property + def service(self) -> str: + """The object's name.""" + ... + + @property + def key(self) -> str: + """The object's key.""" + ... + + def get(self, key: str) -> Awaitable[bytes | None]: + """A state read, issued now.""" + ... + + def set(self, key: str, value: bytes) -> None: + """A state write.""" + ... + + def clear(self, key: str) -> None: + """A state clear.""" + ... + + def state_keys(self) -> Awaitable[list[str]]: + """The object's state keys, issued now.""" + ... + + def sync(self) -> Awaitable[Any]: + """An empty run: done once every command before it is stored.""" + ... + + def attempt_end(self) -> BaseException | None: + """Why the attempt ended (a suspension, say), or None while it goes on.""" + ... + + def on_attempt_end(self, listener: Callable[[], None]) -> Callable[[], None]: + """Call ``listener`` when the attempt ends; the returned function cancels it.""" + ... + + +@dataclass +class ObjectConfig: + """What a session object gives its attempts: its schema, and a wrapper of its store (fault injection, in tests).""" + + schema: Schema | None = None + wrap_store: Callable[[Store], Store] | None = field(default=None) + + +# --- Errors --------------------------------------------------------------------------------------- + + +def handler_error(error: BaseException) -> BaseException: + """ + What a failed call to the session fails the handler with: an ``SqlError``, + which a retry would repeat, becomes a ``TerminalError``, its cause the SQL + error. Anything else is as for any handler: a storage error, which ends the + session, or the handler's own (a failed call to another service, say) is + retried, unless it is a ``TerminalError`` itself. + """ + if not isinstance(error, SqlError): + return error + terminal = TerminalError(str(error)) + terminal.__cause__ = error + return terminal + + +def _raise_for_handler(error: BaseException) -> typing.NoReturn: + """Raise what ``error`` fails the handler with (see ``handler_error``).""" + failure = handler_error(error) + if failure is error: + raise error + raise failure + + +# --- The object's state as a store ---------------------------------------------------------------- + + +def _beside(name: str, key: str) -> str: + """ + The key of the database ``name``, apart from the object's own: each of the + format's keys (``v22/head``, ``v22/tail/...``) as ``v22/db//head``, + under the format's prefix still, which tooling hides, and under none of the + object's database's. + """ + at = key.index("/") + 1 + return f"{key[:at]}db/{name}/{key[at:]}" + + +def _unbeside(name: str, key: str) -> str | None: + """The format's key of a key of the database ``name`` (undoes ``_beside``); None for any other key of the object.""" + at = key.find("/") + 1 + prefix = f"db/{name}/" + if at > 0 and key.startswith(prefix, at): + return key[:at] + key[at + len(prefix) :] + return None + + +class _ObjectStore: + """ + The object's state as an ordered store: ``get`` is a state get, ``set`` and + ``clear`` are state writes, which return at once and which the runtime + stores in order; ``sync`` is an empty run, done once the commands before it + are stored. A ``get`` or ``sync`` the attempt's end leaves waiting raises, + so the session that made it ends instead of holding its lock. + """ + + def __init__(self, host: SqlHost, name: str | None) -> None: + self._host = host + self._name = name + + def _key(self, key: str) -> str: + return key if self._name is None else _beside(self._name, key) + + def get(self, key: str) -> Awaitable[bytes | None]: + # TODO: read with protocol v8's ephemeral get, left out of the journal, + # once a Restate release ships v8 and the SDK can make one. Restate 1.8 + # negotiates v7 at most: a journaled get. + return _until_attempt_end(self._host, self._host.get(self._key(key))) + + def set(self, key: str, value: bytes) -> None: + self._host.set(self._key(key), value) + + def clear(self, key: str) -> None: + self._host.clear(self._key(key)) + + async def keys(self) -> list[str]: + """The database's keys, for a garbage collection.""" + keys = await _until_attempt_end(self._host, self._host.state_keys()) + if self._name is None: + return keys + name = self._name + return [own for own in (_unbeside(name, key) for key in keys) if own is not None] + + async def sync(self) -> None: + await _until_attempt_end(self._host, self._host.sync()) + + +class _ReaderStore: + """The object's state as a view's store: it reads beside the exclusive handler, and writes nothing.""" + + def __init__(self, host: SqlHost, name: str | None) -> None: + self._host = host + self._name = name + + def get(self, key: str) -> Awaitable[bytes | None]: + key = key if self._name is None else _beside(self._name, key) + return _until_attempt_end(self._host, self._host.get(key)) + + def set(self, key: str, value: bytes) -> None: + raise RuntimeError("a view writes no state") + + def clear(self, key: str) -> None: + raise RuntimeError("a view writes no state") + + async def sync(self) -> None: + pass + + +async def _until_attempt_end(host: SqlHost, request: Awaitable[T]) -> T: + """ + ``request``'s result, or why the attempt ended, once it has: an attempt that + is suspended or aborted never completes what it was waiting for. + """ + ended = host.attempt_end() + if ended is not None: + raise ended + task = asyncio.ensure_future(request) + end = asyncio.get_running_loop().create_future() + + def ended_now() -> None: + if not end.done(): + end.set_result(None) + + cancel = host.on_attempt_end(ended_now) + try: + await asyncio.wait({task, end}, return_when=asyncio.FIRST_COMPLETED) + finally: + cancel() + if task.done(): + return task.result() + # Its error, if it ever has one, is dropped. + task.add_done_callback(lambda t: t.cancelled() or t.exception()) + raise host.attempt_end() or SuspendedException() + + +# --- The outbox ----------------------------------------------------------------------------------- + +# A transaction's sends commit with it, in the database's record `outbox` +# (beside its tables, in none of them), which holds each committed +# transaction's sends until they are stored. They are made with its commit +# (see `Attempt.transaction`); a later open makes again those still in the +# outbox, with the same idempotency keys (`transaction//`, with a +# random id of the transaction), so making them again is dropped. Restate +# keeps an attempt's commands in the order they were made, so a send made +# before a commit is stored whenever the commit is: the next commit leaves the +# sends this attempt made out of the outbox, and so does the end of the +# attempt. The format is sdk-sessions' (TypeScript). + +OUTBOX = "outbox" +"""The database's record of the outbox.""" + +Outgoing = dict[str, Any] +"""The sends of a committed transaction, in the outbox: ``{"id": ..., "sends": [...]}``.""" + + +def _encode_outbox(outbox: list[Outgoing]) -> bytes: + """The outbox as its record: none is empty, as a record never set reads.""" + return json.dumps(outbox).encode() if outbox else b"" + + +def _decode_outbox(record: bytes) -> list[Outgoing]: + return json.loads(record.decode()) if record else [] + + +def _now_ms() -> int: + return int(time.time() * 1000) + + +class _Sends: + """The sends of one transaction, as its body makes them (``tx.send``), until it ends.""" + + def __init__(self) -> None: + self._sends: list[dict[str, Any]] = [] + self._open = True + self.send = Send(self._record) + + def _record(self, send: GenericSend) -> None: + if not self._open: + raise TerminalError("the transaction has ended: send from inside its body") + stored: dict[str, Any] = { + "service": send.service, + "method": send.handler, + "parameter": base64.b64encode(send.parameter).decode("ascii"), + } + if send.key is not None: + stored["key"] = send.key + if send.delay is not None: + stored["delay"] = int(send.delay.total_seconds() * 1000) + if send.idempotency_key is not None: + stored["idempotencyKey"] = send.idempotency_key + if send.headers: + stored["headers"] = send.headers + self._sends.append(stored) + + def close(self) -> Outgoing | None: + """The body ended: the transaction's sends, if it made any, their delays counted from now.""" + self._open = False + if not self._sends: + return None + now = _now_ms() + for send in self._sends: + if "delay" in send: + send["due"] = now + send["delay"] + return {"id": str(uuid.uuid4()), "sends": self._sends} + + +def _refused_send(_: GenericSend) -> None: + """The sends of a transaction that has no outbox: a view's, or another database's.""" + raise TerminalError("a send of a transaction is an exclusive handler's, in the object's own database") + + +# --- The database, as a handler gets it ----------------------------------------------------------- + + +class DatabaseReader: + """Reads of a database, or of one transaction of it.""" + + async def _query_columns(self, sql: str, params: Sequence[SqlParam] | None) -> tuple[list[str], Rows]: + raise NotImplementedError + + async def query(self, sql: str, params: Sequence[SqlParam] | None = None) -> list[Row]: + """ + The rows of one statement, each a dict by column name; parameters bind + to a single statement. Columns are named as SQL names them (an + expression needs ``AS``), and two of the same name are an error: alias + them, or use ``values``. + """ + columns, rows = await self._query_columns(sql, params) + try: + return _objects(columns, rows) + except SqlError as error: + _raise_for_handler(error) + + async def get(self, sql: str, params: Sequence[SqlParam] | None = None) -> Row | None: + """The first row of ``query``, or ``None`` when there is none.""" + rows = await self.query(sql, params) + return rows[0] if rows else None + + async def values(self, sql: str, params: Sequence[SqlParam] | None = None) -> Rows: + """The rows of one statement as lists, in column order: for adapters, and for columns that share a name.""" + return (await self._query_columns(sql, params))[1] + + +class DatabaseWriter(DatabaseReader): + """SQL: a database, or one transaction of it.""" + + async def execute(self, sql: str, params: Sequence[SqlParam] | None = None) -> None: + """Run SQL (one or more statements), discarding rows; parameters only with a single statement.""" + await self._query_columns(sql, params) + + +class Transaction(DatabaseWriter): + """ + One transaction of a ``Database``: its SQL, and the sends it commits with. + + ``send`` makes its sends right after the transaction commits, and never if + it rolls back. Each has the idempotency key it is given, or else one of the + transaction (``transaction//``), so a send an attempt makes + again is dropped. + """ + + def __init__(self, handle: ConnectionTransaction, send: Send, reset: Callable[[], Awaitable[None]] | None) -> None: + self._handle = handle + self.send = send + self._reset = reset + + async def _query_columns(self, sql: str, params: Sequence[SqlParam] | None) -> tuple[list[str], Rows]: + try: + return await self._handle.query_columns(sql, params) + except SqlError as error: + _raise_for_handler(error) + + +class Database(DatabaseWriter): + """ + The session's database, ``ctx.db``, for the whole handler: each statement + outside ``transaction()`` is its own transaction. Calls run one at a time, + in call order. A commit is durable once its call returns. + """ + + def __init__(self, attempt: Attempt) -> None: + self._attempt = attempt + + async def _query_columns(self, sql: str, params: Sequence[SqlParam] | None) -> tuple[list[str], Rows]: + return await self._attempt.call(lambda db: db.query_columns(sql, params)) + + def transaction(self) -> typing.AsyncContextManager[Transaction]: + """ + One transaction, as the body of an ``async with``:: + + async with ctx.db.transaction() as tx: + await tx.execute("UPDATE accounts SET balance = balance - ? WHERE id = ?", [amount, src]) + tx.send(deposit, key=dst)(amount) + + It commits when the body ends, and rolls back when the body raises. It + holds the database: other calls wait until it ends, and a call to the + database from inside the body is refused, as it would wait for the body. + So the body uses the ``tx`` it gets, which is invalid once the body has + ended. + """ + return self._attempt.transaction() + + +class _ReaderOf(DatabaseReader): + """The reads of a ``Database``: a view's ``ctx.db``.""" + + def __init__(self, db: Database) -> None: + self._db = db + + async def _query_columns(self, sql: str, params: Sequence[SqlParam] | None) -> tuple[list[str], Rows]: + return await self._db._query_columns(sql, params) # pylint: disable=protected-access + + +def reader_of(db: Database) -> DatabaseReader: + """The reads of ``db``, which apply a library's schema as ``db`` does (``Schema.ensure``).""" + reader = _ReaderOf(db) + ensure = ensure_of(db) + return ensuring(reader, ensure) if ensure is not None else reader + + +_DATABASE_NAME = re.compile(r"^[\w.-]{1,64}$") + + +def database(db: Database, name: str) -> Database: + """ + The database ``name`` of the key, apart from the object's own (``ctx.db``, + of the attempt whose database ``db`` is): for a library whose tables must + be alone in their database. It opens on its first use, and ends with the + attempt, as ``ctx.db`` does. A name is 1 to 64 letters, digits, ``_``, + ``-`` or ``.``. + """ + if not _DATABASE_NAME.match(name): + raise TerminalError(f"database: {name!r} is not a database name") + if not isinstance(db, Database): + raise TerminalError("database takes a handler's ctx.db") + return db._attempt.database(name) # pylint: disable=protected-access + + +async def reset(tx: Transaction) -> None: + """ + Empty the key's database in ``tx``, a transaction of an exclusive handler's + ``ctx.db``: every table, view, index and trigger is dropped, the schemas are + applied again (the object's, and the libraries' the attempt applied), and + the outbox stays. It commits with ``tx``: what ``tx`` writes after it are + the new database's first rows. + """ + if not isinstance(tx, Transaction) or tx._reset is None: # pylint: disable=protected-access + raise TerminalError("reset takes a transaction of an exclusive handler's ctx.db") + try: + await tx._reset() # pylint: disable=protected-access + except SqlError as error: + _raise_for_handler(error) + + +# --- An attempt ----------------------------------------------------------------------------------- + + +def _close_soon(db: Connection) -> None: + task = asyncio.ensure_future(db.close()) + task.add_done_callback(lambda t: t.cancelled() or t.exception()) + + +class _Session: + """ + One session of an attempt's database, from its open on. A failure ends it: + its open's, or a request's since. It closes as the attempt ends, unless the + attempt ends it first (``end``). + """ + + def __init__(self, ready: asyncio.Task[Connection], host: SqlHost) -> None: + self.ready = ready + # A call failed on it once a failure ended it: the next call opens another. + self.retired = False + self._release: Callable[[], None] = lambda: None + ready.add_done_callback(lambda task: self._track(task, host)) + + def _track(self, task: asyncio.Task[Connection], host: SqlHost) -> None: + if task.cancelled() or task.exception() is not None: + return + db = task.result() + # A handler still waiting for something else when the attempt ends (a + # model's request, say) keeps the session no longer: it closes once no + # statement runs, and the handler's later calls fail. + self._release = host.on_attempt_end(lambda: _close_soon(db)) + + @property + def ended(self) -> bool: + """Whether a failure ended it.""" + if not self.ready.done(): + return False + if self.ready.cancelled() or self.ready.exception() is not None: + return True + return not self.ready.result().is_usable() + + async def end(self) -> Connection | None: + """Its connection, or none if the open failed; the attempt's end no longer closes it.""" + try: + db: Connection | None = await asyncio.shield(self.ready) + except BaseException: # pylint: disable=broad-exception-caught + db = None + self._release() + return db + + def close(self) -> None: + """Close it now.""" + + async def close() -> None: + db = await self.end() + if db is not None: + await db.close() + + task = asyncio.ensure_future(close()) + task.add_done_callback(lambda t: t.cancelled() or t.exception()) + + +class Attempt: + """ + An attempt's SQL session: opened on its first use, or as an exclusive + attempt starts (``open_now``), shared by everything the handler does, and + saved when the handler ends (``end``). With ``name``, of the key's database + ``name`` (``database``), which has neither the object's schema nor an + outbox. + """ + + def __init__(self, host: SqlHost, config: ObjectConfig | None = None, name: str | None = None) -> None: + self._host = host + self._config = config or ObjectConfig() + self._name = name + # The object's own database of an exclusive handler: the one with an outbox. + self._has_outbox = not host.shared and name is None + # The key's other databases this attempt opened, by name. + self._databases: dict[str, Attempt] = {} + self._session: _Session | None = None + # The ids of the outbox's transactions this attempt made the sends of. + self._made: set[str] = set() + # The schemas this attempt's database has all the steps of, by name. + self._ensured: dict[str, Schema] = {} + self._ended = False + self.db = Database(self) + """``ctx.db``.""" + ensuring(self.db, self._ensure) + + async def _ensure(self, schemas: Sequence[Schema]) -> None: + await self.call(lambda db: self._ensure_schemas(db, schemas)) + + def database(self, name: str) -> Database: + """The key's database ``name``, in this attempt: see ``database``.""" + if self._ended: + raise RuntimeError("the handler has ended: its SQL session is closed") + other = self._databases.get(name) + if other is None: + other = self._databases[name] = Attempt(self._host, self._config, name) + return other.db + + def open_now(self) -> None: + """ + Open the session now, as the attempt starts, rather than on its first + call: the open makes the sends a crashed attempt committed and left + unmade, which must not wait for a handler that uses the database. + """ + if self._has_outbox and self._session is None: + # Its failure goes to the first call, as it would had that call opened it. + self._current().ready.add_done_callback(lambda t: t.cancelled() or t.exception()) + + async def end(self, error: BaseException | None) -> None: + """ + The handler ended, with ``error`` or not: the key's other databases end, + then this one. It raises what the handler fails with (see + ``handler_error``), or a failure to end the session. + """ + failures: list[BaseException] = [] + for other in self._databases.values(): + try: + await other.end(None) + except BaseException as failure: # pylint: disable=broad-exception-caught + failures.append(failure) + await self._end_session(error) + if failures: + raise failures[0] + + async def call(self, call: Callable[[Connection], Awaitable[T]]) -> T: + """ + ``call`` on the session, opening it on the first; its failure as the + handler fails with it. The call that fails on a session that a failure + ended is its last: the next opens another, which reads what is stored. + """ + session = self._current() + try: + return await call(await session.ready) + except BaseException as error: + if session.ended: + session.retired = True + _raise_for_handler(error) + + @asynccontextmanager + async def transaction(self) -> AsyncIterator[Transaction]: + """ + One transaction. In the object's own database of an exclusive handler, + its sends commit with it, in the outbox, and are made with its commit, + right after its writes and before it waits for anything + (``with_commit``): Restate stores an attempt's commands in order, so + they are stored only if the commit is, and whatever fails the commit's + wait (a cancellation, say) comes after them. A second commit empties + the outbox before the transaction returns: left until the handler's + end, they would be made again by an attempt that resumes a long wait, + past the idempotency retention that drops a repeat. + """ + outgoing: Outgoing | None = None + session = self._current() + try: + db = await session.ready + async with db.transaction() as handle: + sends = _Sends() if self._has_outbox else None + # A view's writes stay in its session: it has nothing to reset. + resetter = None if self._host.shared else (lambda: self._reset(handle)) + tx = Transaction(handle, sends.send if sends is not None else Send(_refused_send), resetter) + try: + yield tx + finally: + outgoing = sends.close() if sends is not None else None + if self._has_outbox: + await self._commit_outbox(handle, outgoing) + except BaseException as error: + if session.ended: + session.retired = True + _raise_for_handler(error) + if outgoing is not None: + await self.call(self._flush) + + async def _commit_outbox(self, handle: ConnectionTransaction, outgoing: Outgoing | None) -> None: + """Commit with ``handle`` the outbox as committed, but the sends stored since, and with ``outgoing``, which the commit makes.""" + outbox = await self._unstored(handle) + if outgoing is not None: + outbox.append(outgoing) + sent = outgoing + handle.with_commit(lambda: self._make(sent)) + handle.set_record(OUTBOX, _encode_outbox(outbox)) + + def _current(self) -> _Session: + """The session, opening one if there is none, or the one there is was retired.""" + if self._ended: + raise RuntimeError("the handler has ended: its SQL session is closed") + if self._session is not None and self._session.retired: + self._session.close() + if self._session is None or self._session.retired: + self._session = self._open() + return self._session + + def _open(self) -> _Session: + return _Session(asyncio.ensure_future(self._start()), self._host) + + def _own_schema(self) -> Schema | None: + """The object's schema, which the key's other databases have none of.""" + return self._config.schema if self._name is None else None + + async def _start(self) -> Connection: + host = self._host + # The page cache's: the key escaped, so that no key reads as another's + # database (a name has no `/`). + key = f"{host.service}/{quote(host.key, safe='')}" + namespace = key if self._name is None else f"{key}/{self._name}" + store: Store = _ReaderStore(host, self._name) if host.shared else _ObjectStore(host, self._name) + if self._config.wrap_store is not None: + store = self._config.wrap_store(store) + db = await open_connection(namespace, store, reader=host.shared) + try: + # A view writes nothing: a write fails, rather than vanish with the attempt. + if host.shared: + await db.execute("PRAGMA query_only = 1") + if self._has_outbox: + # What a crash left unmade; an open again in this attempt finds + # what this attempt made, stored before it, as made. + left = _decode_outbox(await db.record(OUTBOX)) + for outgoing in left: + if outgoing["id"] not in self._made: + self._make(outgoing, again=True) + # Emptied at once, as a transaction's own sends are. + if left: + await self._flush(db) + # A view's too, whose writes are never stored: a key never written reads as empty tables. + own = self._own_schema() + if own is not None: + await self._ensure_schemas(db, [own]) + except BaseException: + await db.close() + raise + return db + + async def _end_session(self, error: BaseException | None) -> None: + """ + End the session, then raise as the handler did: an SQL error that + escaped it becomes terminal. The handler's error stands: what it + committed is stored, and an outbox left full is made by the next open. + A failure to end the session fails a handler that succeeded, which a + retry then ends (a transaction it left open, say). + """ + self._ended = True + failed: BaseException | None = None + try: + await self._close() + except BaseException as failure: # pylint: disable=broad-exception-caught + failed = failure + if error is not None: + _raise_for_handler(error) + if failed is not None: + _raise_for_handler(failed) + + async def _close(self) -> None: + """ + Close the session as the attempt ends. A view's closes: it read the + snapshot it opened on. Another is saved, its outbox emptied. One that a + failure ended closes, and the object's own database is opened again and + saved: it may hold sends unmade (a crash left them, and this attempt's + open failed) or made and not emptied (the commit's wait failed). + """ + if self._session is None: + return + db = await self._session.end() + if self._host.shared: + if db is not None: + await db.close() + return + if db is not None and db.is_usable(): + await self._save(db) + return + if db is not None: + try: + await db.close() + except Exception: # pylint: disable=broad-exception-caught + pass + if not self._has_outbox or self._host.attempt_end() is not None: + return + again = self._open() + try: + await self._save(await again.ready) + finally: + await again.end() + + async def _save(self, db: Connection) -> None: + """Empty the outbox and save the session, which closes if that fails.""" + try: + if self._has_outbox: + await self._flush(db) + await db.save() + except BaseException: + try: + await db.close() + except Exception: # pylint: disable=broad-exception-caught + pass + raise + + async def _reset(self, handle: ConnectionTransaction) -> None: + """Empty the database in the transaction ``handle`` (``reset``), and apply the schemas this attempt applied again.""" + own = self._own_schema() + schemas = ([own] if own is not None else []) + [s for s in self._ensured.values() if s is not own] + await drop_all(handle, schemas) + + async def _ensure_schemas(self, db: Connection, schemas: Sequence[Schema]) -> None: + """Apply the steps of ``schemas`` that the database lacks, each schema once per attempt.""" + lacking = [s for s in schemas if s.name not in self._ensured] + if not lacking: + return + await apply(db, lacking, reader=self._host.shared) + for schema in lacking: + self._ensured[schema.name] = schema + + def _make(self, outgoing: Outgoing, again: bool = False) -> None: + """Make the sends of a committed transaction: each with an idempotency key of the transaction (or its own), so that making them again is dropped.""" + for index, send in enumerate(outgoing["sends"]): + delay = send.get("delay") + due = send.get("due") + if delay is not None and again and due is not None: + # Made again, a delayed send waits only what remains of its delay. + delay = max(0, due - _now_ms()) + self._host.generic_send( + send["service"], + send["method"], + base64.b64decode(send["parameter"]), + send.get("key"), + None if delay is None else timedelta(milliseconds=delay), + send.get("idempotencyKey") or f"transaction/{outgoing['id']}/{index}", + send.get("headers"), + ) + self._made.add(outgoing["id"]) + + async def _unstored(self, db: Connection | ConnectionTransaction) -> list[Outgoing]: + """The outbox as committed, but the transactions this attempt made the sends of: a commit made after them stores them first.""" + return [outgoing for outgoing in _decode_outbox(await db.record(OUTBOX)) if outgoing["id"] not in self._made] + + async def _flush(self, db: Connection) -> None: + """Empty the outbox, so that a later open does not make its sends again: a commit after them, which stores them first.""" + if not await db.record(OUTBOX): + return + # Read in the transaction that writes it: one that committed sends since keeps them. + async with db.transaction() as tx: + tx.set_record(OUTBOX, _encode_outbox(await self._unstored(tx))) + + +async def with_attempt(handler: Callable[[], Awaitable[T]], attempt: Attempt) -> T: + """Run ``handler``, then end ``attempt``'s SQL session, if it opened one (see ``Attempt.end``).""" + try: + value = await handler() + except BaseException as error: + await attempt.end(error) + raise # `end` raises: this is for the type checker + await attempt.end(None) + return value diff --git a/tests/sessions_harness.py b/tests/sessions_harness.py new file mode 100644 index 00000000..5c3df314 --- /dev/null +++ b/tests/sessions_harness.py @@ -0,0 +1,146 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""Session objects against a real Restate server: no journal replay, concurrent awaits, and the outbox.""" + +# pylint: disable=missing-function-docstring + +import asyncio +import typing +import uuid + +import pytest + +import restate +from restate import HarnessEnvironment, sessions + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +pytestmark = [pytest.mark.anyio] + +# --- A regular object, which session objects call and send to ----------------------------------- + +inbox = restate.VirtualObject("inbox") + + +@inbox.handler() +async def deliver(ctx: restate.ObjectContext, message: str) -> None: + messages = await ctx.get("messages", type_hint=list[str]) or [] + ctx.set("messages", [*messages, message]) + + +@inbox.handler(kind="shared") +async def messages(ctx: restate.ObjectSharedContext) -> list[str]: + return await ctx.get("messages", type_hint=list[str]) or [] + + +pricing = restate.Service("pricing") + + +@pricing.handler() +async def price(ctx: restate.Context, sku: str) -> int: + await asyncio.sleep(0.05) + return len(sku) + + +# --- Session objects ---------------------------------------------------------------------------- + +notes = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) + + +@notes.handler() +async def add(ctx: sessions.SessionContext, text: str) -> int: + await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [text]) + row = await ctx.db.get("SELECT count(*) AS n FROM notes") + assert row is not None + return typing.cast(int, row["n"]) + + +@notes.view() +async def texts(ctx: sessions.ViewContext) -> list[str]: + return [typing.cast(str, row["text"]) for row in await ctx.db.query("SELECT text FROM notes ORDER BY rowid")] + + +@notes.handler() +async def flaky(ctx: sessions.SessionContext, text: str) -> int: + """Commits, then fails its first attempt: the retry runs from the start, and finds the commit.""" + seen = await ctx.db.get("SELECT count(*) AS n FROM notes WHERE text = ?", [text]) + if seen is not None and seen["n"] == 0: + await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [text]) + raise ValueError("crash after the commit") + return typing.cast(int, (await ctx.db.get("SELECT count(*) AS n FROM notes"))["n"]) # type: ignore[index] + + +@notes.handler() +async def priced(ctx: sessions.SessionContext, sku: str) -> list[int]: + """A call and SQL awaited together, and two calls at once: one await point, the selector's.""" + first, second, row = await asyncio.gather( + ctx.rpc(price)(sku), + ctx.rpc(price)(sku * 2), + ctx.db.get("SELECT count(*) AS n FROM notes"), + ) + assert row is not None + return [first, second, typing.cast(int, row["n"])] + + +@notes.handler() +async def forward(ctx: sessions.SessionContext, text: str) -> None: + """A note, and its delivery to the inbox of the same key, sent once the transaction commits.""" + async with ctx.db.transaction() as tx: + await tx.execute("INSERT INTO notes (text) VALUES (?)", [text]) + tx.send(deliver, key=ctx.key)(text) + + +@pytest.fixture(scope="session") +async def harness(): + async with restate.create_test_harness( + restate.app([inbox, pricing, notes]), restate_image="ghcr.io/restatedev/restate:latest" + ) as harness: + yield harness + + +async def test_a_handler_commits_and_a_view_reads_it(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + assert await harness.client.object_call(add, key=key, arg="milk") == 1 + assert await harness.client.object_call(add, key=key, arg="eggs") == 2 + assert await harness.client.object_call(texts, key=key, arg=None) == ["milk", "eggs"] + + +async def test_a_view_of_a_key_never_written_reads_empty_tables(harness: HarnessEnvironment): + assert await harness.client.object_call(texts, key=str(uuid.uuid4()), arg=None) == [] + + +async def test_a_retry_runs_the_handler_from_the_start_against_what_was_committed(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + assert await harness.client.object_call(flaky, key=key, arg="once") == 1 + assert await harness.client.object_call(texts, key=key, arg=None) == ["once"] + + +async def test_calls_and_sql_are_awaited_together(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + await harness.client.object_call(add, key=key, arg="x") + assert await harness.client.object_call(priced, key=key, arg="ab") == [2, 4, 1] + + +async def test_the_sends_of_a_transaction_are_made_once_it_commits(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + await harness.client.object_call(forward, key=key, arg="hello") + assert await harness.client.object_call(texts, key=key, arg=None) == ["hello"] + delivered: list[str] = [] + for _ in range(50): + delivered = await harness.client.object_call(messages, key=key, arg=None) + if delivered: + break + await asyncio.sleep(0.1) + assert delivered == ["hello"] diff --git a/tests/sessions_object.py b/tests/sessions_object.py new file mode 100644 index 00000000..233c8ad8 --- /dev/null +++ b/tests/sessions_object.py @@ -0,0 +1,403 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""Session objects, their attempts run over an in-memory copy of the object's state, without Restate.""" + +# pylint: disable=missing-function-docstring,missing-class-docstring + +import asyncio +import itertools +import typing +from dataclasses import dataclass, field +from typing import Any + +import pytest + +import restate +from restate import sessions +from restate.context import AttemptFinishedEvent, Request +from restate.exceptions import TerminalError +from restate.sessions._object import invoke +from sessions_database import Backend # pylint: disable=import-error + + +@pytest.fixture +def anyio_backend(): + return "asyncio" + + +class _NeverFinished(AttemptFinishedEvent): + def is_set(self) -> bool: + return False + + async def wait(self): + await asyncio.Event().wait() + + +_ids = itertools.count() + + +def request() -> Request: + return Request( + id=f"inv_test{next(_ids)}", headers={}, attempt_headers={}, body=b"", attempt_finished_event=_NeverFinished() + ) + + +@dataclass +class Send: + service: str + handler: str + arg: bytes + key: str | None + delay: Any + idempotency_key: str | None + + +class Answer(restate.RestateDurableCallFuture[bytes]): + """A call's result, at once.""" + + def __init__(self, value: bytes) -> None: + super().__init__() + self.value = value + + def __await__(self): + async def value() -> bytes: + return self.value + + return value().__await__() + + async def invocation_id(self) -> str: + return "inv_call" + + async def cancel_invocation(self) -> None: + pass + + +@dataclass +class Host: + """An attempt's context over ``backend``, the object's state, as the SDK's gives it.""" + + backend: Backend + service: str = "notes" + key: str = "alice" + shared: bool = False + sends: list[Send] = field(default_factory=list) + calls: list[tuple[str, str, bytes, str | None, str | None]] = field(default_factory=list) + answers: dict[str, bytes] = field(default_factory=dict) + cancelled: list[str] = field(default_factory=list) + # How many more syncs go through before the process crashes: every sync after fails, losing its writes. + _syncs_left: int | None = None + + def crash_after(self, syncs: int) -> None: + """Let ``syncs`` more syncs through, then crash: every later sync fails.""" + self._syncs_left = syncs + + def get(self, key: str): + future = asyncio.get_running_loop().create_future() + future.set_result(self.backend.request().get(key)) + return future + + def set(self, key: str, value: bytes) -> None: + self.backend.request().set(key, value) + + def clear(self, key: str) -> None: + self.backend.request().clear(key) + + async def state_keys(self) -> list[str]: + return [*self.backend.stored, *(key for key, _ in self.backend.unconfirmed)] + + async def sync(self) -> None: + if self._syncs_left is not None: + if self._syncs_left == 0: + self.backend.lose() + raise RuntimeError("the sync failed") + self._syncs_left -= 1 + self.backend.confirm() + + def generic_call(self, service, handler, arg, key=None, idempotency_key=None, headers=None): + self.calls.append((service, handler, arg, key, idempotency_key)) + return Answer(self.answers[handler]) + + def generic_send(self, service, handler, arg, key=None, send_delay=None, idempotency_key=None, headers=None): + self.sends.append(Send(service, handler, arg, key, send_delay, idempotency_key)) + + def cancel_invocation(self, invocation_id: str) -> None: + self.cancelled.append(invocation_id) + + def attempt_end(self): + return None + + def on_attempt_end(self, listener): + return lambda: None + + def view(self) -> "Host": + """A view's context of the same key.""" + return Host(self.backend, self.service, self.key, shared=True) + + +def run(obj: sessions.SessionObject, name: str, host: Host, *args: Any) -> typing.Awaitable[Any]: + return invoke(obj, name, host, request(), *args) + + +# --- The objects -------------------------------------------------------------------------------- + +notes = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) + + +@notes.handler() +async def add(ctx: sessions.SessionContext, text: str) -> None: + await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [text]) + + +@notes.view() +async def texts(ctx: sessions.ViewContext) -> list[str]: + return [typing.cast(str, row["text"]) for row in await ctx.db.query("SELECT text FROM notes ORDER BY rowid")] + + +audit = restate.VirtualObject("audit") + + +@audit.handler() +async def record(ctx: restate.ObjectContext, entry: str) -> None: + pass + + +pricing = restate.Service("pricing") + + +@pricing.handler() +async def quote(ctx: restate.Context, sku: str) -> int: + return 0 + + +# --- Tests -------------------------------------------------------------------------------------- + + +async def test_a_handler_commits_and_a_view_reads_what_it_committed(): + host = Host(Backend()) + await run(notes, "add", host, "milk") + await run(notes, "add", host, "eggs") + assert await run(notes, "texts", host.view()) == ["milk", "eggs"] + + +async def test_a_view_of_a_key_never_written_reads_empty_tables_and_stores_nothing(): + backend = Backend() + assert await run(notes, "texts", Host(backend).view()) == [] + assert backend.stored == {} + + +async def test_a_step_added_later_runs_once_on_each_database(): + backend = Backend() + v1 = sessions.SessionObject("evolving", schema=["CREATE TABLE t (a INTEGER)"]) + + @v1.handler() + async def put(ctx: sessions.SessionContext) -> None: + await ctx.db.execute("INSERT INTO t (a) VALUES (1)") + + await run(v1, "put", Host(backend, service="evolving")) + + v2 = sessions.SessionObject("evolving", schema=["CREATE TABLE t (a INTEGER)", "ALTER TABLE t ADD COLUMN b TEXT"]) + + @v2.handler() + async def columns(ctx: sessions.SessionContext) -> list[str]: + return [typing.cast(str, row["name"]) for row in await ctx.db.query("SELECT name FROM pragma_table_info('t')")] + + host = Host(backend, service="evolving") + assert await run(v2, "columns", host) == ["a", "b"] + assert await run(v2, "columns", host) == ["a", "b"] + + +async def test_the_sends_of_a_transaction_are_made_once_it_commits_and_never_if_it_rolls_back(): + shop = sessions.SessionObject("shop", schema=["CREATE TABLE orders (id TEXT PRIMARY KEY)"]) + + @shop.handler() + async def place(ctx: sessions.SessionContext, order: str) -> None: + async with ctx.db.transaction() as tx: + await tx.execute("INSERT INTO orders VALUES (?)", [order]) + tx.send(record, key="orders")(order) + if order == "bad": + raise ValueError("no") + + host = Host(Backend(), service="shop") + await run(shop, "place", host, "o1") + assert len(host.sends) == 1 + sent = host.sends[0] + assert (sent.service, sent.handler, sent.key, sent.arg) == ("audit", "record", "orders", b'"o1"') + assert sent.idempotency_key is not None and sent.idempotency_key.startswith("transaction/") + with pytest.raises(ValueError): + await run(shop, "place", host, "bad") + assert len(host.sends) == 1 + # The outbox is empty: a later attempt makes nothing again. + again = Host(host.backend, service="shop") + await run(shop, "place", again, "o2") + assert [send.arg for send in again.sends] == [b'"o2"'] + + +async def test_the_sends_a_crash_left_in_the_outbox_are_made_again_with_the_same_keys(): + shop = sessions.SessionObject("shop2", schema=["CREATE TABLE orders (id TEXT PRIMARY KEY)"]) + + @shop.handler() + async def place(ctx: sessions.SessionContext, order: str) -> None: + async with ctx.db.transaction() as tx: + await tx.execute("INSERT INTO orders VALUES (?)", [order]) + tx.send(record, key="orders")(order) + # The commit's sync goes through; then the process crashes, before the outbox is emptied. + first.crash_after(1) + + backend = Backend() + first = Host(backend, service="shop2") + with pytest.raises(RuntimeError, match="the sync failed"): + await run(shop, "place", first, "o1") + assert len(first.sends) == 1 + + # The next attempt opens the database as it starts, and makes them again. + second = Host(backend, service="shop2") + await run(shop, "place", second, "o2") + keys = [send.idempotency_key for send in second.sends] + assert keys[0] == first.sends[0].idempotency_key + assert [send.arg for send in second.sends] == [b'"o1"', b'"o2"'] + + third = Host(backend, service="shop2") + await run(shop, "place", third, "o3") + assert [send.arg for send in third.sends] == [b'"o3"'] + + +async def test_an_sql_error_fails_the_handler_with_a_terminal_error(): + host = Host(Backend()) + broken = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) + + @broken.handler() + async def bad(ctx: sessions.SessionContext) -> None: + await ctx.db.execute("SELECT * FROM missing") + + with pytest.raises(TerminalError, match="no such table") as error: + await run(broken, "bad", host) + assert isinstance(error.value.__cause__, sessions.SqlError) + + +async def test_a_call_to_ctx_db_inside_its_transaction_is_refused(): + nested = sessions.SessionObject("nested", schema=["CREATE TABLE t (a INTEGER)"]) + + @nested.handler() + async def inside(ctx: sessions.SessionContext) -> None: + async with ctx.db.transaction(): + await ctx.db.execute("INSERT INTO t VALUES (1)") + + with pytest.raises(TerminalError, match="use the transaction handle"): + await run(nested, "inside", Host(Backend(), service="nested")) + + +async def test_a_view_writes_nothing(): + host = Host(Backend()) + await run(notes, "add", host, "milk") + writer = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) + + @writer.view() + async def sneaky(ctx: sessions.ViewContext) -> None: + await ctx.db.query("INSERT INTO notes (text) VALUES ('x')") + + with pytest.raises(TerminalError): + await run(writer, "sneaky", host.view()) + assert await run(notes, "texts", host.view()) == ["milk"] + + +async def test_calls_are_typed_by_their_handler_and_checked_for_their_key(): + caller = sessions.SessionObject("caller") + + @caller.handler() + async def price(ctx: sessions.SessionContext, sku: str) -> int: + return await ctx.rpc(quote, idempotency_key="k")(sku) + + @caller.handler() + async def unkeyed(ctx: sessions.SessionContext) -> None: + ctx.send(record)("x") + + @caller.handler() + async def literal(ctx: sessions.SessionContext) -> Any: + return await ctx.rpc.generic("pricing", "quote", "pear") + + host = Host(Backend(), service="caller", answers={"quote": b"42"}) + assert await run(caller, "price", host, "apple") == 42 + assert host.calls[0] == ("pricing", "quote", b'"apple"', None, "k") + with pytest.raises(TerminalError, match="audit is an object: call it with its key"): + await run(caller, "unkeyed", host) + assert await run(caller, "literal", host) == 42 + + +async def test_ctx_send_is_made_at_once_and_cancel_takes_an_invocation_id(): + sender = sessions.SessionObject("sender") + + @sender.handler() + async def go(ctx: sessions.SessionContext) -> None: + ctx.send(record, key="k")("entry") + ctx.cancel("inv_1") + ctx.cancel("not-an-id") + + host = Host(Backend(), service="sender") + with pytest.raises(TerminalError, match="not an invocation id"): + await run(sender, "go", host) + assert [(s.handler, s.key, s.idempotency_key) for s in host.sends] == [("record", "k", None)] + assert host.cancelled == ["inv_1"] + + +async def test_another_database_of_the_key_is_apart_from_the_objects_own(): + multi = sessions.SessionObject("multi", schema=["CREATE TABLE own (a INTEGER)"]) + + @multi.handler() + async def both(ctx: sessions.SessionContext) -> list[str]: + other = sessions.database(ctx.db, "other") + await other.execute("CREATE TABLE apart (b INTEGER)") + rows = await other.query("SELECT name FROM sqlite_schema WHERE type = 'table'") + return [typing.cast(str, row["name"]) for row in rows] + + host = Host(Backend(), service="multi") + assert await run(multi, "both", host) == ["apart"] + assert any(key.startswith("v22/db/other/") for key in host.backend.stored) + + +async def test_reset_empties_the_database_and_applies_the_schema_again(): + resettable = sessions.SessionObject("resettable", schema=["CREATE TABLE t (a INTEGER)"]) + + @resettable.handler() + async def fill(ctx: sessions.SessionContext) -> None: + await ctx.db.execute("CREATE TABLE extra (b INTEGER)") + await ctx.db.execute("INSERT INTO t VALUES (1)") + + @resettable.handler() + async def empty(ctx: sessions.SessionContext) -> list[Any]: + async with ctx.db.transaction() as tx: + await sessions.reset(tx) + await tx.execute("INSERT INTO t VALUES (2)") + return await ctx.db.values("SELECT name FROM sqlite_schema WHERE type = 'table' UNION ALL SELECT a FROM t") + + host = Host(Backend(), service="resettable") + await run(resettable, "fill", host) + assert await run(resettable, "empty", host) == [["t"], [2]] + + +def test_an_sdk_handler_is_refused(): + plain = restate.VirtualObject("plain") + + @plain.handler() + async def journaled(ctx: restate.ObjectContext) -> None: + pass + + obj = sessions.SessionObject("refusing") + with pytest.raises(TypeError, match="is an SDK handler"): + obj.handler()(journaled) + + +def test_a_session_object_runs_its_handlers_with_explicit_recovery_and_lazy_state(): + assert notes.handlers["add"].explicit_recovery + assert notes.handlers["texts"].explicit_recovery + assert notes.handlers["texts"].kind == "shared" + assert notes.enable_lazy_state + assert notes.service_tag.metadata is not None + assert notes.service_tag.metadata["dev.restate.tooling.journal.hidden.name.prefix"] == "v22/" From fbbcb28ffda8427601bcf849dd8153b37f28a2fc Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 14:54:46 +0100 Subject: [PATCH 07/45] Let @obj.handler and @obj.view decorate bare, without parentheses @notes.handler async def add(ctx: sessions.SessionContext, text: str) -> None: ... @notes.handler() and @notes.handler(name=..., ...) work as before; the overloads keep both typed. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/sessions/__init__.py | 4 +- python/restate/sessions/_object.py | 75 ++++++++++++++++++++++++----- tests/sessions_object.py | 19 +++++++- 3 files changed, 81 insertions(+), 17 deletions(-) diff --git a/python/restate/sessions/__init__.py b/python/restate/sessions/__init__.py index 3f13572e..4c7aa309 100644 --- a/python/restate/sessions/__init__.py +++ b/python/restate/sessions/__init__.py @@ -18,12 +18,12 @@ notes = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) - @notes.handler() + @notes.handler async def add(ctx: sessions.SessionContext, text: str) -> None: await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [text]) - @notes.view() + @notes.view async def list(ctx: sessions.ViewContext) -> list[str]: return [row["text"] for row in await ctx.db.query("SELECT text FROM notes ORDER BY rowid")] diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index 17a30132..35d90274 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -16,11 +16,11 @@ notes = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) - @notes.handler() + @notes.handler async def add(ctx: sessions.SessionContext, text: str) -> None: await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [text]) - @notes.view() + @notes.view async def list(ctx: sessions.ViewContext) -> list[str]: return [row["text"] for row in await ctx.db.query("SELECT text FROM notes ORDER BY rowid")] """ @@ -271,9 +271,9 @@ class SessionObject(VirtualObject): - ``schema``: the steps that make each key's database, in order, applied once per key by the first handler that uses it. A step that ran never changes: a change is a step of its own. - - ``@obj.handler()``: handlers run one at a time per key, each with a + - ``@obj.handler``: handlers run one at a time per key, each with a ``SessionContext``. - - ``@obj.view()``: views run beside the handlers, each with a + - ``@obj.view``: views run beside the handlers, each with a ``ViewContext``: they read, and write nothing. A name is a handler or a view, not both. The other options are the SDK's @@ -313,7 +313,11 @@ def __init__( self.config = ObjectConfig(schema=Schema(name, schema) if schema else None) self._entries: dict[str, _Entry] = {} - def handler( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride] + @typing.overload # type: ignore[override] + def handler(self, fn: F, /) -> F: ... # pyright: ignore[reportIncompatibleMethodOverride] + + @typing.overload + def handler( self, name: str | None = None, *, @@ -328,18 +332,35 @@ def handler( # type: ignore[override] # pyright: ignore[reportIncompatibleMeth idempotency_retention: timedelta | None = None, ingress_private: bool | None = None, invocation_retry_policy: InvocationRetryPolicy | None = None, - ) -> Callable[[F], F]: + ) -> Callable[[F], F]: ... + + def handler( # pyright: ignore[reportIncompatibleMethodOverride] + self, + name: str | F | None = None, + *, + accept: str = "application/json", + content_type: str = "application/json", + input_serde: Serde[Any] = DefaultSerde(), + output_serde: Serde[Any] = DefaultSerde(), + metadata: dict[str, str] | None = None, + inactivity_timeout: timedelta | None = None, + abort_timeout: timedelta | None = None, + journal_retention: timedelta | None = None, + idempotency_retention: timedelta | None = None, + ingress_private: bool | None = None, + invocation_retry_policy: InvocationRetryPolicy | None = None, + ) -> F | Callable[[F], F]: """ A handler: it runs one at a time per key, and is the only one that writes. Its function takes a ``SessionContext``, and its input, if any:: - @notes.handler() + @notes.handler async def add(ctx: sessions.SessionContext, text: str) -> None: await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [text]) - The options are the SDK's handler's. + With options, ``@notes.handler(name=..., ...)``: the SDK's handler's. """ - return self._register( + return self._decorate( False, name, dict( @@ -357,6 +378,10 @@ async def add(ctx: sessions.SessionContext, text: str) -> None: ), ) + @typing.overload + def view(self, fn: F, /) -> F: ... + + @typing.overload def view( self, name: str | None = None, @@ -372,21 +397,39 @@ def view( idempotency_retention: timedelta | None = None, ingress_private: bool | None = None, invocation_retry_policy: InvocationRetryPolicy | None = None, - ) -> Callable[[F], F]: + ) -> Callable[[F], F]: ... + + def view( + self, + name: str | F | None = None, + *, + accept: str = "application/json", + content_type: str = "application/json", + input_serde: Serde[Any] = DefaultSerde(), + output_serde: Serde[Any] = DefaultSerde(), + metadata: dict[str, str] | None = None, + inactivity_timeout: timedelta | None = None, + abort_timeout: timedelta | None = None, + journal_retention: timedelta | None = None, + idempotency_retention: timedelta | None = None, + ingress_private: bool | None = None, + invocation_retry_policy: InvocationRetryPolicy | None = None, + ) -> F | Callable[[F], F]: """ A view: it runs beside the handlers, and reads the database as committed when its invocation started. Its function takes a ``ViewContext``, and its input, if any:: - @notes.view() + @notes.view async def count(ctx: sessions.ViewContext) -> int: row = await ctx.db.get("SELECT count(*) AS n FROM notes") return row["n"] if row else 0 A view that is still running after 16 commits of the handlers fails, - and Restate retries it. The options are the SDK's handler's. + and Restate retries it. With options, ``@notes.view(name=..., ...)``: + the SDK's handler's. """ - return self._register( + return self._decorate( True, name, dict( @@ -404,6 +447,12 @@ async def count(ctx: sessions.ViewContext) -> int: ), ) + def _decorate(self, view: bool, name: str | F | None, options: dict[str, Any]) -> F | Callable[[F], F]: + """The decorator of ``@obj.handler(...)``, or, bare (``@obj.handler``), the function it registers.""" + if callable(name): + return self._register(view, None, options)(name) + return self._register(view, name, options) + def _register(self, view: bool, name: str | None, options: dict[str, Any]) -> Callable[[F], F]: def register(fn: F) -> F: if RESTATE_UNIQUE_HANDLER_SYMBOL in vars(fn): diff --git a/tests/sessions_object.py b/tests/sessions_object.py index 233c8ad8..b447a859 100644 --- a/tests/sessions_object.py +++ b/tests/sessions_object.py @@ -151,12 +151,12 @@ def run(obj: sessions.SessionObject, name: str, host: Host, *args: Any) -> typin notes = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) -@notes.handler() +@notes.handler async def add(ctx: sessions.SessionContext, text: str) -> None: await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [text]) -@notes.view() +@notes.view async def texts(ctx: sessions.ViewContext) -> list[str]: return [typing.cast(str, row["text"]) for row in await ctx.db.query("SELECT text FROM notes ORDER BY rowid")] @@ -401,3 +401,18 @@ def test_a_session_object_runs_its_handlers_with_explicit_recovery_and_lazy_stat assert notes.enable_lazy_state assert notes.service_tag.metadata is not None assert notes.service_tag.metadata["dev.restate.tooling.journal.hidden.name.prefix"] == "v22/" + + +def test_a_handler_or_view_is_named_after_its_function_or_as_its_option_says(): + named = sessions.SessionObject("named") + + @named.handler(name="put") + async def put_one(ctx: sessions.SessionContext, value: int) -> None: + pass + + @named.view("peek") + async def peek_one(ctx: sessions.ViewContext) -> None: + pass + + assert sorted(named.handlers) == ["peek", "put"] + assert named.handlers["peek"].kind == "shared" From fa928f8b853459d89472f2aca76fa5025fd1ae78 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 14:58:14 +0100 Subject: [PATCH 08/45] Fixes from a self-review of the sessions code - A store that answers with the wrong type fails the request with why, rather than leaving it to fail as dropped when the reply is collected. - The object's store makes its sync (the empty run) and its listing when it is called, in order with the writes, as the Store contract says and sdk-sessions does; as coroutines, they were made a loop turn later. - Ending a session no longer swallows the CancelledError of the task that waits for its open. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/sessions/_session.py | 26 ++++++++++++++++++-------- src/sessions.rs | 19 ++++++++++++------- tests/sessions_database.py | 7 +++++++ tests/sessions_object.py | 21 +++++++++++++++++++++ 4 files changed, 58 insertions(+), 15 deletions(-) diff --git a/python/restate/sessions/_session.py b/python/restate/sessions/_session.py index 97eacbc2..e2eb879f 100644 --- a/python/restate/sessions/_session.py +++ b/python/restate/sessions/_session.py @@ -181,16 +181,21 @@ def set(self, key: str, value: bytes) -> None: def clear(self, key: str) -> None: self._host.clear(self._key(key)) - async def keys(self) -> list[str]: - """The database's keys, for a garbage collection.""" - keys = await _until_attempt_end(self._host, self._host.state_keys()) - if self._name is None: - return keys + def keys(self) -> Awaitable[list[str]]: + """The database's keys, for a garbage collection: listed now.""" + listing = _until_attempt_end(self._host, self._host.state_keys()) name = self._name - return [own for own in (_unbeside(name, key) for key in keys) if own is not None] + if name is None: + return listing - async def sync(self) -> None: - await _until_attempt_end(self._host, self._host.sync()) + async def own() -> list[str]: + return [own for own in (_unbeside(name, key) for key in await listing) if own is not None] + + return own() + + def sync(self) -> Awaitable[None]: + """Done once every write before it is stored: the run is made now, in order with them.""" + return _until_attempt_end(self._host, self._host.sync()) class _ReaderStore: @@ -504,6 +509,11 @@ async def end(self) -> Connection | None: """Its connection, or none if the open failed; the attempt's end no longer closes it.""" try: db: Connection | None = await asyncio.shield(self.ready) + except asyncio.CancelledError: + # Ours, as the task awaiting is cancelled, unless the open itself was. + if not self.ready.cancelled(): + raise + db = None except BaseException: # pylint: disable=broad-exception-caught db = None self._release() diff --git a/src/sessions.rs b/src/sessions.rs index 927b9569..aebb18ec 100644 --- a/src/sessions.rs +++ b/src/sessions.rs @@ -73,22 +73,27 @@ struct SqlReply { #[pymethods] impl SqlReply { + /// A value that is none of those fails the request, with why. #[pyo3(signature = (value, error=None))] - fn __call__(&self, value: Option>, error: Option) -> PyResult<()> { + fn __call__(&self, value: Option>, error: Option) { let answer = match (error, value) { (Some(error), _) => Err(error), (None, None) => Ok(Answer::Nothing), - (None, Some(value)) => Ok(if value.is_instance_of::() { - Answer::Keys(value.extract()?) - } else { - Answer::Value(Bytes::from(bytes_of(&value)?)) - }), + (None, Some(value)) => answer_of(&value).map_err(|e| e.to_string()), }; let sender = self.sender.lock().ok().and_then(|mut sender| sender.take()); if let Some(sender) = sender { let _ = sender.send(answer); } - Ok(()) + } +} + +/// A request's answer from the value the store gave: its keys (a list), or a value read. +fn answer_of(value: &Bound<'_, PyAny>) -> PyResult { + if value.is_instance_of::() { + Ok(Answer::Keys(value.extract()?)) + } else { + Ok(Answer::Value(Bytes::from(bytes_of(value)?))) } } diff --git a/tests/sessions_database.py b/tests/sessions_database.py index 8a8ac626..04454844 100644 --- a/tests/sessions_database.py +++ b/tests/sessions_database.py @@ -549,3 +549,10 @@ async def test_a_listing_that_fails_ends_the_session_after_the_commit_that_colle assert not db.is_usable() reopened = await open(ns, backend.request()) assert await reopened.values("SELECT count(*) FROM t") == [[committed]] + + +async def test_a_store_that_answers_with_the_wrong_type_fails_the_request_with_why(): + store = Backend().request() + store.get = lambda key: "not bytes" # type: ignore[method-assign,assignment,return-value] + with pytest.raises(RuntimeError, match="expected bytes, got str"): + await open(fresh(), store) diff --git a/tests/sessions_object.py b/tests/sessions_object.py index b447a859..6b649e61 100644 --- a/tests/sessions_object.py +++ b/tests/sessions_object.py @@ -416,3 +416,24 @@ async def peek_one(ctx: sessions.ViewContext) -> None: assert sorted(named.handlers) == ["peek", "put"] assert named.handlers["peek"].kind == "shared" + + +async def test_the_object_store_issues_its_sync_and_listing_when_called(): + from restate.sessions._session import _ObjectStore # pylint: disable=import-outside-toplevel + + issued: list[str] = [] + + class Counting(Host): + def state_keys(self): # type: ignore[override] + issued.append("keys") + return super().state_keys() + + def sync(self): # type: ignore[override] + issued.append("sync") + return super().sync() + + store = _ObjectStore(Counting(Backend()), "other") + sync, keys = store.sync(), store.keys() + assert issued == ["sync", "keys"] + await sync + assert await keys == [] From 6602f80e6ea6e2e4381070fbf52a206169f22964 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 15:12:54 +0100 Subject: [PATCH 09/45] Run the database on the loop's thread, and write to the VM directly The VM and the storage are in the same extension, so a request needs no trip between threads, nor, for a write, through Python: - The database's calls run on the thread of the asyncio loop that makes them (SqlTask): a call that waits returns an asyncio future, polled again on the loop as what it waits for wakes it. shared-core's database creates no thread, and runs on whatever polls it. The page cache's disk tier, on foyer's runtime, wakes the loop through a thread of its own, so that it never waits for the GIL. pyo3-async-runtimes goes, and with it the tokio threads every call crossed to and from. - The store's requests are calls of Python's dispatch on that thread, in order, rather than callbacks scheduled from another. - With the invocation's VM (the SDK's context gives it, unless a journal codec encodes the state), the store is the VM (VmStore): writes are its state commands, and a read it can answer at once is answered by it. Only a read the runtime must answer waits through the selector (WAIT), as do the sync and the listing. In-process, over the tests' in-memory store, a single-row commit goes from 182 to 126 us, the loop no longer sleeping between threads. Against a real server a commit is about 5 ms of wall time, its sync's round trip, and about 1.4 ms of CPU, most of it the SDK's HTTP/2 and run machinery: the two stores measure the same there. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/sessions/_database.py | 30 +- python/restate/sessions/_object.py | 14 + python/restate/sessions/_session.py | 23 +- src/sessions.rs | 498 ++++++++++++++++++++++----- 4 files changed, 466 insertions(+), 99 deletions(-) diff --git a/python/restate/sessions/_database.py b/python/restate/sessions/_database.py index 047e2f8c..1c4b83c9 100644 --- a/python/restate/sessions/_database.py +++ b/python/restate/sessions/_database.py @@ -36,6 +36,7 @@ SQL_KEYS, SQL_SET, SQL_SYNC, + SQL_WAIT, SqlRefusedException, sql_open, ) @@ -79,6 +80,13 @@ class Store(Protocol): A store may also have ``keys()``: every key, as ``ctx.state_keys()``, or an awaitable of them. Then a commit that checkpoints deletes the packs no checkpoint needs any more; a store without it keeps them. + + A store over a Restate context may give the invocation's VM, as ``vm`` + (``restate._internal.PyVM``): then the extension writes to it, and reads + from it, directly, without the store's ``get``, ``set`` and ``clear``; it + calls ``wait(handle)`` (an awaitable, done once the VM has the read whose + handle it is) only for a read it must wait for. ``database_name`` names the + key's database, apart from the object's own (see ``_session.database``). """ def get(self, key: str) -> bytes | None | Awaitable[bytes | None]: ... @@ -95,8 +103,10 @@ async def open(namespace: str, store: Store, *, reader: bool = False) -> Connect meanwhile, until the owner's second checkpoint after it; it writes nothing. """ boundary = _Boundary(store, asyncio.get_running_loop()) + vm = getattr(store, "vm", None) + name = getattr(store, "database_name", None) if vm is not None else None try: - native = await sql_open(namespace, boundary.dispatch, boundary.loop, boundary.lists_keys, reader) + native = await _settled(sql_open(namespace, boundary.dispatch, boundary.lists_keys, reader, vm, name)) except Exception as error: if boundary.failed is not None: raise boundary.failed from error @@ -104,9 +114,15 @@ async def open(namespace: str, store: Store, *, reader: bool = False) -> Connect return Connection(native, boundary) +async def _settled(answer: Any) -> Any: + """A call's answer: its value, as it gave it, or once the asyncio future it gave is done.""" + return await answer if inspect.isawaitable(answer) else answer + + class _Boundary: """ - The store as the extension calls it, through ``dispatch`` on ``loop``. An + The store as the extension calls it, through ``dispatch``, on the loop's + thread, in the order of the extension's requests. An error of the store's crosses the extension as text, which ends the session: the first is kept here, and the calls that fail on the session it ended raise it as it was raised (the SDK's suspension, say). A failed write @@ -131,8 +147,8 @@ def _keep(self, error: BaseException) -> None: if self.failed is None: self.failed = error - def dispatch(self, op: int, key: str, value: bytes | None, reply: Any) -> None: - """One request of the extension's, on the event loop, in the order it made them.""" + def dispatch(self, op: int, key: Any, value: bytes | None, reply: Any) -> None: + """One request of the extension's: ``key`` is a key of the store, or for ``SQL_WAIT`` a handle of the VM.""" if op in (SQL_SET, SQL_CLEAR): try: if op == SQL_SET: @@ -149,6 +165,8 @@ def dispatch(self, op: int, key: str, value: bytes | None, reply: Any) -> None: result: Any = self.store.get(key) elif op == SQL_KEYS: result = self.store.keys() # type: ignore[attr-defined] + elif op == SQL_WAIT: + result = self.store.wait(key) # type: ignore[attr-defined] elif self._write_failed is not None: reply(None, str(self._write_failed)) return @@ -183,7 +201,7 @@ def done(future: asyncio.Future) -> None: def _reply_value(op: int, result: Any) -> Any: """A request's answer as the extension reads it: a value, the keys as a list, or nothing for a sync.""" - if op == SQL_SYNC: + if op in (SQL_SYNC, SQL_WAIT): return None if op == SQL_KEYS: return list(result) @@ -362,7 +380,7 @@ async def save(self) -> None: async def close(self) -> None: """End the session now, whatever its state: queued calls fail.""" - await self._native.close() + await _settled(self._native.close()) async def _rollback(self) -> None: """Roll back, unless a storage error ended the session: the caller raises its own error, which says why.""" diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index 35d90274..aac6bbb4 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -41,6 +41,7 @@ async def list(ctx: sessions.ViewContext) -> list[str]: from restate.retry_policy import InvocationRetryPolicy from restate.serde import BytesSerde, DefaultSerde, Serde from restate.server_context import ServerInvocationContext +from restate.vm import SingleUnresolvedFuture from restate.sessions import _selector from restate.sessions._calls import CallFuture, Rpc, Send, send_now from restate.sessions._database import Store @@ -136,6 +137,19 @@ def on_attempt_finished() -> None: context.on_attempt_finished = on_attempt_finished # type: ignore[method-assign] + @property + def vm(self) -> Any: + """ + The invocation's VM, which the database's store writes and reads + directly (see ``_database.Store``); none with a journal codec, whose + values only the context's state commands encode. + """ + return None if self.context.journal_codec is not None else self.context.vm.vm + + def wait(self, handle: int) -> Awaitable[None]: + """Done once the VM has the read whose handle is ``handle``: through the selector, as every wait.""" + return self.selector.wait(SingleUnresolvedFuture(handle)) + def get(self, key: str) -> Awaitable[bytes | None]: return self.context.get(key, serde=BytesSerde(), entry_name=key) diff --git a/python/restate/sessions/_session.py b/python/restate/sessions/_session.py index e2eb879f..8aa52272 100644 --- a/python/restate/sessions/_session.py +++ b/python/restate/sessions/_session.py @@ -153,18 +153,31 @@ def _unbeside(name: str, key: str) -> str | None: return None +def _vm_of(host: SqlHost) -> Any: + """The invocation's VM, if ``host`` gives it (see ``Store``): the SDK's context does.""" + return getattr(host, "vm", None) + + class _ObjectStore: """ The object's state as an ordered store: ``get`` is a state get, ``set`` and ``clear`` are state writes, which return at once and which the runtime stores in order; ``sync`` is an empty run, done once the commands before it - are stored. A ``get`` or ``sync`` the attempt's end leaves waiting raises, - so the session that made it ends instead of holding its lock. + are stored. A ``get``, ``wait`` or ``sync`` the attempt's end leaves + waiting raises, so the session that made it ends instead of holding its + lock. With the invocation's VM (``vm``), the extension writes and reads it + directly, and waits for a read with ``wait``. """ def __init__(self, host: SqlHost, name: str | None) -> None: self._host = host self._name = name + self.vm = _vm_of(host) + self.database_name = name + + def wait(self, handle: int) -> Awaitable[None]: + """Done once the VM has the read whose handle is ``handle``.""" + return _until_attempt_end(self._host, self._host.wait(handle)) # type: ignore[attr-defined] def _key(self, key: str) -> str: return key if self._name is None else _beside(self._name, key) @@ -204,6 +217,12 @@ class _ReaderStore: def __init__(self, host: SqlHost, name: str | None) -> None: self._host = host self._name = name + self.vm = _vm_of(host) + self.database_name = name + + def wait(self, handle: int) -> Awaitable[None]: + """Done once the VM has the read whose handle is ``handle``.""" + return _until_attempt_end(self._host, self._host.wait(handle)) # type: ignore[attr-defined] def get(self, key: str) -> Awaitable[bytes | None]: key = key if self._name is None else _beside(self._name, key) diff --git a/src/sessions.rs b/src/sessions.rs index aebb18ec..f9fc8595 100644 --- a/src/sessions.rs +++ b/src/sessions.rs @@ -1,46 +1,66 @@ //! Sessions: an object's SQLite-compatible database, kept in its Restate -//! state, over a key/value store implemented in Python. +//! state. //! //! The storage and the SQL on it are shared-core's (`sql::Database`, with its //! feature `sql`); this module binds them to Python, as `native/src/node.rs` //! of sdk-sessions binds them to JavaScript. `restate/sessions/_database.py` -//! opens a database over a store, and orders the calls. +//! opens a database, and orders the calls. //! -//! `sql_open` takes the store as one dispatcher function, -//! `dispatch(op, key, value, reply)`, and the asyncio event loop it runs on. -//! Every request is scheduled on that loop with `call_soon_threadsafe`, which -//! runs callbacks in the order they were scheduled, so a read sees the writes -//! before it. +//! # Where a call runs //! -//! `set` and `clear` queue a write; `sync` completes once every write before -//! it is stored. The object's state keeps writes in issue order, so a crash -//! loses only a suffix of them (see [`KvStore`]): a statement that commits -//! queues its writes, then completes at one `sync`. With `keys`, a checkpoint -//! deletes the packs no checkpoint needs; a store without keeps its garbage. +//! On the thread of the asyncio loop that makes it, and nowhere else +//! ([`SqlTask`]): shared-core's database creates no thread, and runs on the +//! task that polls it. A call that needs nothing it must wait for, such as SQL +//! on cached pages or a statement inside a transaction, returns its value at +//! once, or raises its error ([`answer`]); one that waits returns an asyncio +//! future, and is polled again on the loop as what it waits for completes. The +//! one other thread is the page cache's disk tier (foyer's runtime): its +//! wake-ups reach the loop through a thread of their own, so that the cache +//! never waits for the GIL. +//! +//! # The store +//! +//! The state the database is kept in: `set` and `clear` queue a write; `sync` +//! completes once every write before it is stored. The object's state keeps +//! writes in issue order, so a crash loses only a suffix of them (see +//! [`KvStore`]): a statement that commits queues its writes, then completes at +//! one `sync`. With `keys`, a checkpoint deletes the packs no checkpoint needs; +//! a store without keeps its garbage. +//! +//! The store is the invocation's VM, as a handler's context has it +//! ([`VmStore`]): a write goes to it directly, as does a read that it can +//! answer at once (state preloaded, or written by the invocation); waiting +//! for a read, the sync (an empty run) and the listing of the keys go through +//! Python, whose selector drives the invocation's waits. Or, without a VM (in +//! tests, or with a journal codec), a store implemented in Python +//! ([`PyStore`]). Either way, Python's side is one function, +//! `dispatch(op, arg, value, reply)`, called directly, on the loop's thread, +//! in the order of the requests, so that a read sees the writes before it. //! //! Every database of the process is kept warm in one page cache //! ([`page_cache`]). -//! -//! A call runs on the calling thread as far as it can ([`answer`]): one that -//! needs nothing from the store, such as SQL on cached pages or a statement -//! inside a transaction, returns its value, or raises its error; one that -//! waits for the store goes on in the tokio runtime and returns an awaitable. -//! The caller awaits the latter. +use crate::PyVM; use bytes::Bytes; use futures::channel::oneshot; use futures::future::BoxFuture; use pyo3::create_exception; use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; use pyo3::prelude::*; +use pyo3::sync::PyOnceLock; use pyo3::types::{PyBool, PyByteArray, PyBytes, PyFloat, PyInt, PyList, PyString, PyTuple}; use pyo3::IntoPyObjectExt; use restate_sdk_shared_core::sql::{self, Numeric, Value}; use restate_sdk_shared_core::state::{self, KvStore, PageCache, PageCacheConfig}; +use restate_sdk_shared_core::{Value as VmValue, VM}; use std::borrow::Cow; use std::collections::HashMap; use std::future::Future; use std::io; -use std::sync::{Arc, Mutex}; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{mpsc, Arc, Mutex, OnceLock}; +use std::task::{Context, Poll, Wake, Waker}; +use std::thread::ThreadId; create_exception!( restate_sdk_python_core, @@ -49,12 +69,187 @@ create_exception!( "A failure of the database that a retry would repeat (shared-core's `sql::Refused`): it comes of what the host gave, its SQL, its records, a transaction it left open. It may have ended the session." ); -/// What a request to the store does: the store's method it calls. +/// What a request to Python's side does: the store's method it calls. pub(crate) const GET: u32 = 0; pub(crate) const SET: u32 = 1; pub(crate) const CLEAR: u32 = 2; pub(crate) const SYNC: u32 = 3; pub(crate) const KEYS: u32 = 4; +/// Wait until the VM completes the handle `arg`, a read's ([`VmStore`]). +pub(crate) const WAIT: u32 = 5; + +// --- Running calls on the loop's thread ------------------------------------------------------------ + +/// A call to the database, its value a Python object. +type Call = Pin>> + Send>>; + +/// A call that waits: polled on its loop's thread, as what it waits for wakes +/// it, until it settles its asyncio future. A cancelled future drops the call, +/// which ends its session (shared-core: a call that does not return leaves its +/// state unknown). +#[pyclass(frozen)] +struct SqlTask { + call: Mutex>, + future: Py, + event_loop: Py, + /// The loop's thread. + thread: ThreadId, + /// Whether a poll is scheduled on the loop already. + scheduled: AtomicBool, +} + +#[pymethods] +impl SqlTask { + /// Poll the call once more: what the waker schedules on the loop. + fn _poll(slf: Py, py: Python<'_>) -> PyResult<()> { + poll(py, &slf) + } + + /// The future's done-callback: a cancelled one drops its call. + fn _done(&self, future: &Bound<'_, PyAny>) -> PyResult<()> { + if future.call_method0("cancelled")?.is_truthy()? { + if let Ok(mut call) = self.call.try_lock() { + call.take(); + } + } + Ok(()) + } +} + +/// `asyncio.get_running_loop`. +static GET_RUNNING_LOOP: PyOnceLock> = PyOnceLock::new(); + +impl SqlTask { + /// Run `call`, which a first poll left pending, on the running loop: its + /// asyncio future. + fn spawn<'py>(py: Python<'py>, call: Call) -> PyResult> { + let event_loop = GET_RUNNING_LOOP + .import(py, "asyncio", "get_running_loop")? + .call0()?; + let future = event_loop.call_method0("create_future")?; + let task = Py::new( + py, + SqlTask { + call: Mutex::new(Some(call)), + future: future.clone().unbind(), + event_loop: event_loop.unbind(), + thread: std::thread::current().id(), + scheduled: AtomicBool::new(false), + }, + )?; + future.call_method1("add_done_callback", (task.bind(py).getattr("_done")?,))?; + poll(py, &task)?; + Ok(future) + } +} + +/// Poll `task`'s call, settling its future once it is ready. +fn poll(py: Python<'_>, task: &Py) -> PyResult<()> { + let state = task.get(); + state.scheduled.store(false, Ordering::Release); + let mut slot = state + .call + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let Some(call) = slot.as_mut() else { + return Ok(()); + }; + let future = state.future.bind(py); + if future.call_method0("done")?.is_truthy()? { + slot.take(); + return Ok(()); + } + let waker = Waker::from(Arc::new(TaskWaker(task.clone_ref(py)))); + let Poll::Ready(result) = call.as_mut().poll(&mut Context::from_waker(&waker)) else { + return Ok(()); + }; + slot.take(); + drop(slot); + match result { + Ok(value) => future.call_method1("set_result", (value,))?, + Err(error) => future.call_method1("set_exception", (error.into_value(py),))?, + }; + Ok(()) +} + +/// Wakes a [`SqlTask`]: a poll on its loop, scheduled once however often it +/// is woken before it runs. Woken on the loop's thread, it schedules it +/// there; woken on another (the page cache's disk tier), it hands the +/// scheduling to [`foreign_wakes`], so that the waking thread never waits +/// for the GIL. +struct TaskWaker(Py); + +impl Wake for TaskWaker { + fn wake(self: Arc) { + self.wake_by_ref(); + } + + fn wake_by_ref(self: &Arc) { + let task = self.0.get(); + if task.scheduled.swap(true, Ordering::AcqRel) { + return; + } + if std::thread::current().id() == task.thread { + Python::attach(|py| schedule(py, &self.0, false)); + } else if foreign_wakes().send(self.clone()).is_err() { + task.scheduled.store(false, Ordering::Release); + } + } +} + +/// The thread that schedules the wake-ups of other threads on their loops. +fn foreign_wakes() -> &'static mpsc::Sender> { + static WAKES: OnceLock>> = OnceLock::new(); + WAKES.get_or_init(|| { + let (wakes, woken) = mpsc::channel::>(); + std::thread::Builder::new() + .name("restate-sql-wake".into()) + .spawn(move || { + for waker in woken { + // None once the interpreter is gone: nothing more runs. + let _ = Python::try_attach(|py| schedule(py, &waker.0, true)); + } + }) + .expect("the thread of the page cache's wake-ups"); + wakes + }) +} + +/// Schedule a poll of `task` on its loop; a closed loop runs nothing more. +fn schedule(py: Python<'_>, task: &Py, threadsafe: bool) { + let state = task.get(); + let scheduled = task.bind(py).getattr("_poll").and_then(|poll| { + let method = if threadsafe { + "call_soon_threadsafe" + } else { + "call_soon" + }; + state.event_loop.call_method1(py, method, (poll,)) + }); + if scheduled.is_err() { + state.scheduled.store(false, Ordering::Release); + } +} + +/// `call`'s value if a first poll finds it ready, on the calling thread; else +/// an asyncio future of it ([`SqlTask`]). Ready, a call costs no future: SQL +/// on cached pages, or a statement inside a transaction. +fn answer<'py, T, F>(py: Python<'py>, call: F) -> PyResult> +where + F: Future> + Send + 'static, + T: for<'a> IntoPyObject<'a> + Send + 'static, +{ + let mut call: Call = Box::pin(async move { + let value = call.await?; + Python::attach(|py| value.into_py_any(py)) + }); + match call.as_mut().poll(&mut Context::from_waker(Waker::noop())) { + Poll::Ready(value) => value.map(|value| value.into_bound(py)), + Poll::Pending => SqlTask::spawn(py, call), + } +} + +// --- Python's side of the store -------------------------------------------------------------------- /// What a request answers: a `get`'s value, the store's keys, or nothing. enum Answer { @@ -113,67 +308,90 @@ fn bytes_of(value: &Bound<'_, PyAny>) -> PyResult> { } } -/// The Python store, behind its dispatcher on its event loop. -struct PyStore { - event_loop: Py, - dispatch: Py, - /// Whether the store has `keys`. - lists_keys: bool, -} - fn io_error(error: impl std::fmt::Display) -> io::Error { io::Error::other(error.to_string()) } -impl PyStore { - /// Schedule `dispatch(op, key, value, reply)` on the event loop; `reply` - /// is `None` for a write, which answers nothing. - fn schedule( +/// What a request names: a key of the store, or a handle of the VM. +enum Arg<'a> { + Key(&'a str), + Handle(u32), +} + +/// Python's side of a store: `dispatch(op, arg, value, reply)`, called on +/// the loop's thread, where every call runs. +struct Dispatch(Py); + +impl Dispatch { + /// `dispatch(op, arg, value, reply)`; `reply` is `None` for a write, which answers nothing. + fn call( &self, op: u32, - key: &str, + arg: Arg<'_>, value: Option, reply: Option, ) -> io::Result<()> { Python::attach(|py| { + let arg = match arg { + Arg::Key(key) => PyString::new(py, key).into_any(), + Arg::Handle(handle) => handle.into_bound_py_any(py)?, + }; let value = value.map(|value| PyBytes::new(py, &value).into_any().unbind()); let reply = match reply { Some(reply) => Py::new(py, reply)?.into_any(), None => py.None(), }; - self.event_loop.call_method1( - py, - "call_soon_threadsafe", - (self.dispatch.clone_ref(py), op, key, value, reply), - )?; + self.0.call1(py, (op, arg, value, reply))?; Ok::<_, PyErr>(()) }) .map_err(io_error) } - /// A request that answers: a read, a sync, or the keys. - async fn request(&self, op: u32, key: &str) -> io::Result { + /// A request that answers: a read, a wait, a sync, or the keys. + async fn request(&self, op: u32, arg: Arg<'_>) -> io::Result { let (sender, receiver) = oneshot::channel(); let reply = SqlReply { sender: Mutex::new(Some(sender)), }; - self.schedule(op, key, None, Some(reply))?; + self.call(op, arg, None, Some(reply))?; match receiver.await { Ok(Ok(answer)) => Ok(answer), Ok(Err(error)) => Err(io_error(error)), Err(_) => Err(io_error("the store dropped the request")), } } + + async fn sync(&self) -> io::Result<()> { + self.request(SYNC, Arg::Key("")).await.map(drop) + } + + async fn keys(&self) -> io::Result> { + match self.request(KEYS, Arg::Key("")).await? { + Answer::Keys(keys) => Ok(keys), + _ => Err(unexpected("keys")), + } + } } fn unexpected(method: &str) -> io::Error { io_error(format!("the store's {method} answered with the wrong type")) } +fn no_keys() -> io::Error { + io::Error::new(io::ErrorKind::Unsupported, "the store has no keys") +} + +/// A store implemented in Python: every request goes to `dispatch`. +struct PyStore { + dispatch: Dispatch, + /// Whether the store has `keys`. + lists_keys: bool, +} + impl KvStore for PyStore { fn get<'a>(&'a self, key: &'a str) -> BoxFuture<'a, io::Result>> { Box::pin(async move { - match self.request(GET, key).await? { + match self.dispatch.request(GET, Arg::Key(key)).await? { Answer::Nothing => Ok(None), Answer::Value(value) => Ok(Some(value)), Answer::Keys(_) => Err(unexpected("get")), @@ -182,33 +400,143 @@ impl KvStore for PyStore { } fn put(&self, key: &str, value: Bytes) -> io::Result<()> { - self.schedule(SET, key, Some(value), None) + self.dispatch.call(SET, Arg::Key(key), Some(value), None) } fn delete(&self, key: &str) -> io::Result<()> { - self.schedule(CLEAR, key, None, None) + self.dispatch.call(CLEAR, Arg::Key(key), None, None) } fn sync(&self) -> BoxFuture<'_, io::Result<()>> { - Box::pin(async { self.request(SYNC, "").await.map(drop) }) + Box::pin(self.dispatch.sync()) } fn keys(&self) -> BoxFuture<'_, io::Result>> { Box::pin(async move { - if !self.lists_keys { - return Err(io::Error::new( - io::ErrorKind::Unsupported, - "the store has no keys", - )); + match self.lists_keys { + true => self.dispatch.keys().await, + false => Err(no_keys()), } - match self.request(KEYS, "").await? { - Answer::Keys(keys) => Ok(keys), - _ => Err(unexpected("keys")), + }) + } +} + +// --- The invocation's VM as the store -------------------------------------------------------------- + +/// The object's state through the invocation's VM, as a handler's context +/// reaches it, but without Python on the way: a write is a state command of +/// the VM's, a read that the VM can answer at once (state preloaded, or +/// written by this invocation) is answered by it. A read the runtime must +/// answer waits through Python's selector (`WAIT`), as do the sync (`SYNC`, an +/// empty run) and the listing of the keys (`KEYS`): `dispatch`'s. +/// +/// Each command is named after its key, as the context's are, so that +/// tooling can hide the database by name. With `name`, the database `name` of +/// the key, apart from the object's own: each key under `db//` (see +/// `_beside` in `restate/sessions/_session.py`, which lists its keys). A +/// reader (a view's) writes nothing. +struct VmStore { + vm: Py, + dispatch: Dispatch, + name: Option, + reader: bool, +} + +impl VmStore { + /// The state key of the database's `key`. + fn key(&self, key: &str) -> String { + match &self.name { + None => key.to_owned(), + Some(name) => { + let at = key.find('/').map_or(0, |at| at + 1); + format!("{}db/{name}/{}", &key[..at], &key[at..]) } + } + } + + /// `f` of the VM, borrowed for that long: never across an await. + fn with_vm(&self, f: impl FnOnce(&mut PyVM) -> io::Result) -> io::Result { + Python::attach(|py| { + let mut vm = self.vm.bind(py).try_borrow_mut().map_err(io_error)?; + f(&mut vm) }) } + + fn refuse_write(&self) -> io::Result<()> { + match self.reader { + true => Err(io_error("a view writes no state")), + false => Ok(()), + } + } } +impl KvStore for VmStore { + fn get<'a>(&'a self, key: &'a str) -> BoxFuture<'a, io::Result>> { + Box::pin(async move { + let key = self.key(key); + let handle = self.with_vm(|vm| { + (vm.vm + .sys_state_get(key.clone(), Some(key), Default::default())) + .map_err(io_error) + })?; + let ready = self.with_vm(|vm| Ok(vm.vm.is_completed(handle)))?; + if !ready { + self.dispatch + .request(WAIT, Arg::Handle(u32::from(handle))) + .await?; + } + self.with_vm(|vm| match vm.vm.take_notification(handle) { + Ok(Some(VmValue::Success(value))) => Ok(Some(value)), + Ok(Some(VmValue::Void)) => Ok(None), + Ok(Some(VmValue::Failure(failure))) => Err(io_error(failure.message)), + Ok(_) => Err(io_error("the state read completed with no value")), + Err(error) => Err(io_error(error)), + }) + }) + } + + fn put(&self, key: &str, value: Bytes) -> io::Result<()> { + self.refuse_write()?; + let key = self.key(key); + self.with_vm(|vm| { + (vm.vm + .sys_state_set(key.clone(), value, Some(key), Default::default())) + .map_err(io_error) + }) + } + + fn delete(&self, key: &str) -> io::Result<()> { + self.refuse_write()?; + let key = self.key(key); + self.with_vm(|vm| { + vm.vm + .sys_state_clear(key.clone(), Some(key)) + .map_err(io_error) + }) + } + + fn sync(&self) -> BoxFuture<'_, io::Result<()>> { + Box::pin(async move { + match self.reader { + // A reader wrote nothing to sync. + true => Ok(()), + false => self.dispatch.sync().await, + } + }) + } + + fn keys(&self) -> BoxFuture<'_, io::Result>> { + Box::pin(async move { + match self.reader { + true => Err(no_keys()), + false => self.dispatch.keys().await, + } + }) + } +} + +// --- The database ---------------------------------------------------------------------------------- + /// A failure of the database as Python gets it: one a retry would repeat /// ([`sql::Refused`]) is an `SqlRefusedException`, which /// `restate/sessions/_database.py` tells apart, as it may have ended the @@ -220,29 +548,6 @@ fn error(error: Box) -> PyErr { } } -/// `call`'s value if it is ready at once, polled here on the calling thread; -/// else an awaitable of it, as it goes on in the tokio runtime. Ready, a call -/// spares two thread hops: what a call costs otherwise, several times the SQL -/// it runs. -fn answer<'py, T, F>(py: Python<'py>, call: F) -> PyResult> -where - F: Future> + Send + 'static, - T: for<'a> IntoPyObject<'a> + Send + 'static, -{ - use std::task::{Context, Poll, Waker}; - let mut call = Box::pin(call); - let poll = { - // What the call reaches of tokio (the page cache's disk tier) finds - // the runtime it goes on in when pending. - let _runtime = pyo3_async_runtimes::tokio::get_runtime().enter(); - call.as_mut().poll(&mut Context::from_waker(Waker::noop())) - }; - match poll { - Poll::Ready(result) => result?.into_bound_py_any(py), - Poll::Pending => pyo3_async_runtimes::tokio::future_into_py(py, call), - } -} - /// The page cache shared by the process's databases: 256 MiB of memory and /// 1 GiB of disk, in a file under the temporary directory, or the file /// `RESTATE_SQL_CACHE_FILE` names. `RESTATE_SQL_CACHE_MEMORY` and @@ -289,26 +594,35 @@ fn cache_config() -> io::Result> { })) } -/// Open the database of `namespace` over the store behind `dispatch`, which -/// runs on `event_loop` (see the module's doc), kept warm in the process's -/// page cache; `lists_keys` if the store has `keys`. See -/// [`sql::Open::reader`] for a `reader`. +/// Open the database of `namespace`, kept warm in the process's page cache, +/// over the invocation's `vm` if given (with `name`, the key's database +/// `name`), or else over the store behind `dispatch` (`lists_keys` if it has +/// `keys`): see the module's doc. See [`sql::Open::reader`] for a `reader`. #[pyfunction] -#[pyo3(signature = (namespace, dispatch, event_loop, lists_keys = false, reader = false))] +#[pyo3(signature = (namespace, dispatch, lists_keys = false, reader = false, vm = None, name = None))] pub(crate) fn sql_open<'py>( py: Python<'py>, namespace: String, dispatch: Py, - event_loop: Py, lists_keys: bool, reader: bool, + vm: Option>, + name: Option, ) -> PyResult> { - let store = Arc::new(PyStore { - event_loop, - dispatch, - lists_keys, - }); - pyo3_async_runtimes::tokio::future_into_py(py, async move { + let dispatch = Dispatch(dispatch); + let store: Arc = match vm { + Some(vm) => Arc::new(VmStore { + vm, + dispatch, + name, + reader, + }), + None => Arc::new(PyStore { + dispatch, + lists_keys, + }), + }; + answer(py, async move { let mut open = sql::Database::open(store); let cache = page_cache() .await @@ -325,7 +639,7 @@ pub(crate) fn sql_open<'py>( /// One session of a namespace's database (see [`sql::Database`]): every SQL /// commit is durable when its call completes. Each call returns its value, or -/// an awaitable of it (see [`answer`]). +/// an asyncio future of it (see [`answer`]). #[pyclass(frozen)] pub(crate) struct PySqlDatabase(Arc); @@ -395,10 +709,10 @@ impl PySqlDatabase { answer(py, async move { db.save().await.map_err(error) }) } - /// See [`sql::Database::close`]: an awaitable. + /// See [`sql::Database::close`]. fn close<'py>(&self, py: Python<'py>) -> PyResult> { let db = self.0.clone(); - pyo3_async_runtimes::tokio::future_into_py(py, async move { + answer(py, async move { db.close().await; Ok(()) }) @@ -484,6 +798,7 @@ fn python_value(py: Python<'_>, value: Value) -> PyResult> { pub(crate) fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_function(wrap_pyfunction!(sql_open, m)?)?; m.add( "SqlRefusedException", @@ -494,5 +809,6 @@ pub(crate) fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add("SQL_CLEAR", CLEAR)?; m.add("SQL_SYNC", SYNC)?; m.add("SQL_KEYS", KEYS)?; + m.add("SQL_WAIT", WAIT)?; Ok(()) } From 975d2b9920b6d3a577ab28c5f42b8504895a24c5 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 15:17:26 +0100 Subject: [PATCH 10/45] Run the real-server session tests on both stores Each test runs on the invocation's VM, which the extension writes and reads directly, and on the Python store over the context's state commands, which an endpoint with a journal codec uses; one test checks that the writes went through the store asked for. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- tests/sessions_harness.py | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/tests/sessions_harness.py b/tests/sessions_harness.py index 5c3df314..546077d8 100644 --- a/tests/sessions_harness.py +++ b/tests/sessions_harness.py @@ -20,6 +20,7 @@ import restate from restate import HarnessEnvironment, sessions +from restate.sessions import _object, _session @pytest.fixture(scope="session") @@ -29,6 +30,33 @@ def anyio_backend(): pytestmark = [pytest.mark.anyio] +# The writes that went through the Python store, rather than straight to the VM. +python_writes = 0 + + +@pytest.fixture(autouse=True, params=["vm", "python"]) +def store(request, monkeypatch): + """ + Each test runs on both stores: the invocation's VM, which the extension + writes and reads directly, and the Python store over the context's state + commands, which an endpoint with a journal codec uses. The app serves in + this process, so the patch reaches its handlers. + """ + global python_writes # pylint: disable=global-statement + python_writes = 0 + if request.param == "python": + monkeypatch.setattr(_object._ContextHost, "vm", property(lambda self: None)) + set_state = _session._ObjectStore.set + + def counted(self, key: str, value: bytes) -> None: + global python_writes # pylint: disable=global-statement + python_writes += 1 + set_state(self, key, value) + + monkeypatch.setattr(_session._ObjectStore, "set", counted) + return request.param + + # --- A regular object, which session objects call and send to ----------------------------------- inbox = restate.VirtualObject("inbox") @@ -144,3 +172,8 @@ async def test_the_sends_of_a_transaction_are_made_once_it_commits(harness: Harn break await asyncio.sleep(0.1) assert delivered == ["hello"] + + +async def test_the_writes_go_through_the_store_asked_for(harness: HarnessEnvironment, store: str): + await harness.client.object_call(add, key=str(uuid.uuid4()), arg="x") + assert (python_writes > 0) == (store == "python") From dfaa339faf7fca6d5a4617b9544b139c26201424 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 15:22:13 +0100 Subject: [PATCH 11/45] Explicit cancellation for session handlers @obj.handler(explicit_cancellation=True), with sessions.cancelled(ctx), which returns when Restate cancels the invocation, and sessions.cancel_calls(ctx), which cancels the calls the attempt made: as sdk-sessions' cancelled and cancelCalls, which the agents' chat handler uses to stop a turn, a model's request included. The cancellation is one more wait of the selector's, awaited with the handler's others. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/sessions/__init__.py | 4 +- python/restate/sessions/_object.py | 64 ++++++++++++++++++++++++++--- tests/sessions_harness.py | 59 +++++++++++++++++++++++++- tests/sessions_object.py | 18 ++++++++ 4 files changed, 138 insertions(+), 7 deletions(-) diff --git a/python/restate/sessions/__init__.py b/python/restate/sessions/__init__.py index 4c7aa309..138ea327 100644 --- a/python/restate/sessions/__init__.py +++ b/python/restate/sessions/__init__.py @@ -38,7 +38,7 @@ async def list(ctx: sessions.ViewContext) -> list[str]: from restate.sessions._calls import CallFuture, Rpc, Send from restate.sessions._database import Row, Rows, SqlError, SqlParam, SqlValue -from restate.sessions._object import Cancel, SessionContext, SessionObject, ViewContext +from restate.sessions._object import Cancel, SessionContext, SessionObject, ViewContext, cancel_calls, cancelled from restate.sessions._schema import Schema, Step from restate.sessions._session import ( Database, @@ -68,6 +68,8 @@ async def list(ctx: sessions.ViewContext) -> list[str]: "Step", "Transaction", "ViewContext", + "cancel_calls", + "cancelled", "database", "reset", ] diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index aac6bbb4..0ffed759 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -30,10 +30,11 @@ async def list(ctx: sessions.ViewContext) -> list[str]: import functools import typing from collections.abc import Awaitable, Callable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import timedelta from typing import Any +from restate import _preview from restate.context import Request from restate.exceptions import SuspendedException, TerminalError from restate.handler import RESTATE_UNIQUE_HANDLER_SYMBOL, Handler, handler_from_callable @@ -97,6 +98,10 @@ class SessionContext: """Sends, made at once; a transaction's (``tx.send``) are made once it commits.""" cancel: Cancel """Cancel an invocation, by its id: see ``Cancel``.""" + _sdk: Any = field(default=None, repr=False, compare=False) + """The SDK's context of the attempt, for ``cancelled`` and ``cancel_calls``.""" + _explicit_cancellation: bool = field(default=False, repr=False, compare=False) + """Whether its handler has ``explicit_cancellation=True``.""" @dataclass(frozen=True) @@ -238,6 +243,7 @@ class _Entry: fn: Callable[..., Awaitable[Any]] view: bool config: ObjectConfig + explicit_cancellation: bool = False def _context_of(entry: _Entry, host: SqlHost, request: Request, attempt: Attempt) -> SessionContext | ViewContext: @@ -247,7 +253,14 @@ def _context_of(entry: _Entry, host: SqlHost, request: Request, attempt: Attempt if entry.view: return ViewContext(key=host.key, request=request, db=reader_of(attempt.db), rpc=rpc, cancel=cancel) return SessionContext( - key=host.key, request=request, db=attempt.db, rpc=rpc, send=Send(send_now(host)), cancel=cancel + key=host.key, + request=request, + db=attempt.db, + rpc=rpc, + send=Send(send_now(host)), + cancel=cancel, + _sdk=getattr(host, "context", None), + _explicit_cancellation=entry.explicit_cancellation, ) @@ -270,6 +283,37 @@ async def handler(ctx: ServerInvocationContext, *args: Any) -> Any: return handler +def _explicit_of(ctx: SessionContext, operation: str) -> Any: + """The SDK's context behind ``ctx``, a handler's with ``explicit_cancellation=True``.""" + if not isinstance(ctx, SessionContext) or ctx._sdk is None: # pylint: disable=protected-access + raise TerminalError(f"{operation} takes a session object's ctx") + if not ctx._explicit_cancellation: # pylint: disable=protected-access + raise TerminalError(f"{operation} takes the ctx of a handler with explicit_cancellation=True") + return ctx._sdk # pylint: disable=protected-access + + +async def cancelled(ctx: SessionContext) -> None: + """ + Returns when Restate cancels the invocation of ``ctx``, a handler's with + ``explicit_cancellation=True``, whose awaits are left as they are: for + code that stops its own work then (a model's request, say):: + + async def stop_when_cancelled(): + await sessions.cancelled(ctx) + await sessions.cancel_calls(ctx) + """ + await _preview.cancellation(_explicit_of(ctx, "cancelled")) + + +async def cancel_calls(ctx: SessionContext) -> list[str]: + """ + Cancel the calls the attempt of ``ctx`` made and that have not completed, + as a cancelled invocation's would be without explicit cancellation: a + handler's with ``explicit_cancellation=True``. Returns their invocation ids. + """ + return await _preview.cancel_previous_calls(_explicit_of(ctx, "cancel_calls")) + + async def invoke(obj: SessionObject, name: str, host: SqlHost, request: Request, *args: Any) -> Any: """Run handler or view ``name`` of ``obj`` for one attempt over ``host``, a context of a test's own, without Restate.""" entry = obj._entries.get(name) # pylint: disable=protected-access @@ -346,6 +390,7 @@ def handler( idempotency_retention: timedelta | None = None, ingress_private: bool | None = None, invocation_retry_policy: InvocationRetryPolicy | None = None, + explicit_cancellation: bool = False, ) -> Callable[[F], F]: ... def handler( # pyright: ignore[reportIncompatibleMethodOverride] @@ -363,6 +408,7 @@ def handler( # pyright: ignore[reportIncompatibleMethodOverride] idempotency_retention: timedelta | None = None, ingress_private: bool | None = None, invocation_retry_policy: InvocationRetryPolicy | None = None, + explicit_cancellation: bool = False, ) -> F | Callable[[F], F]: """ A handler: it runs one at a time per key, and is the only one that @@ -373,11 +419,16 @@ async def add(ctx: sessions.SessionContext, text: str) -> None: await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [text]) With options, ``@notes.handler(name=..., ...)``: the SDK's handler's. + With ``explicit_cancellation=True``, Restate's cancellation of its + invocation fails no await, nor cancels its calls: the handler learns of + it with ``sessions.cancelled(ctx)``, and decides what to stop (its calls, + with ``sessions.cancel_calls(ctx)``). """ return self._decorate( False, name, dict( + explicit_cancellation=explicit_cancellation, accept=accept, content_type=content_type, input_serde=input_serde, @@ -474,14 +525,17 @@ def register(fn: F) -> F: f"{self.name}.{fn.__name__} is an SDK handler: a session object's handlers are functions of its " "context. Put journaled work in a restate.VirtualObject, Service or Workflow, and call it with ctx.rpc" ) - entry = _Entry(fn=fn, view=view, config=self.config) - registered = VirtualObject.handler(self, name, kind="shared" if view else "exclusive", **options)( + explicit = bool(options.get("explicit_cancellation")) + sdk_options = {key: value for key, value in options.items() if key != "explicit_cancellation"} + entry = _Entry(fn=fn, view=view, config=self.config, explicit_cancellation=explicit) + registered = VirtualObject.handler(self, name, kind="shared" if view else "exclusive", **sdk_options)( _sdk_handler(entry) ) sdk: Handler[Any, Any] = handler_from_callable(registered) # A later attempt makes different reads than a replayed journal # would have: each runs from the start, against the state as it is. sdk.explicit_recovery = True + sdk.explicit_cancellation = explicit self._entries[sdk.name] = entry return typing.cast(F, registered) @@ -493,4 +547,4 @@ def with_store_wrapper(self, wrap: Callable[[Store], Store]) -> SessionObject: return self -__all__ = ["CallFuture", "Cancel", "SessionContext", "SessionObject", "ViewContext"] +__all__ = ["CallFuture", "Cancel", "SessionContext", "SessionObject", "ViewContext", "cancel_calls", "cancelled"] diff --git a/tests/sessions_harness.py b/tests/sessions_harness.py index 546077d8..c01844a5 100644 --- a/tests/sessions_harness.py +++ b/tests/sessions_harness.py @@ -15,6 +15,9 @@ import asyncio import typing import uuid +from datetime import timedelta + +import httpx import pytest @@ -130,10 +133,31 @@ async def forward(ctx: sessions.SessionContext, text: str) -> None: tx.send(deliver, key=ctx.key)(text) +sleeper = restate.Service("sleeper") + + +@sleeper.handler() +async def nap(ctx: restate.Context) -> str: + await ctx.sleep(timedelta(minutes=5)) + return "rested" + + +@notes.handler(explicit_cancellation=True) +async def stoppable(ctx: sessions.SessionContext, text: str) -> list[str]: + """Waits for a nap; cancelled, it cancels the nap, records that it stopped, and answers the calls it cancelled.""" + napping = asyncio.ensure_future(ctx.rpc(nap)()) + stopping = asyncio.ensure_future(sessions.cancelled(ctx)) + await asyncio.wait({napping, stopping}, return_when=asyncio.FIRST_COMPLETED) + assert stopping.done() + cancelled = await sessions.cancel_calls(ctx) + await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [f"stopped: {text}"]) + return cancelled + + @pytest.fixture(scope="session") async def harness(): async with restate.create_test_harness( - restate.app([inbox, pricing, notes]), restate_image="ghcr.io/restatedev/restate:latest" + restate.app([inbox, pricing, sleeper, notes]), restate_image="ghcr.io/restatedev/restate:latest" ) as harness: yield harness @@ -177,3 +201,36 @@ async def test_the_sends_of_a_transaction_are_made_once_it_commits(harness: Harn async def test_the_writes_go_through_the_store_asked_for(harness: HarnessEnvironment, store: str): await harness.client.object_call(add, key=str(uuid.uuid4()), arg="x") assert (python_writes > 0) == (store == "python") + + +async def cancel_when_running(harness: HarnessEnvironment, invocation_id: str) -> None: + """Cancel the invocation, once Restate runs it.""" + async with httpx.AsyncClient() as admin: + for _ in range(100): + status = await admin.post( + f"{harness.admin_api_url}/query", + json={"query": f"SELECT status FROM sys_invocation WHERE id = '{invocation_id}'"}, + headers={"Accept": "application/json"}, + ) + rows = status.json().get("rows", []) + if rows and rows[0]["status"] in ("running", "suspended"): + break + await asyncio.sleep(0.05) + response = await admin.patch(f"{harness.admin_api_url}/invocations/{invocation_id}/cancel") + response.raise_for_status() + + +async def attach(harness: HarnessEnvironment, invocation_id: str) -> httpx.Response: + async with httpx.AsyncClient(timeout=30) as ingress: + return await ingress.get(f"{harness.ingress_url}/restate/invocation/{invocation_id}/attach") + + +async def test_a_handler_with_explicit_cancellation_stops_its_own_work(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + sent = await harness.client.object_send(stoppable, key=key, arg="nap") + await cancel_when_running(harness, sent.invocation_id) + answer = await attach(harness, sent.invocation_id) + assert answer.status_code == 200, answer.text + [nap_id] = answer.json() + assert (await attach(harness, nap_id)).status_code == 409 + assert await harness.client.object_call(texts, key=key, arg=None) == ["stopped: nap"] diff --git a/tests/sessions_object.py b/tests/sessions_object.py index 6b649e61..c83f8042 100644 --- a/tests/sessions_object.py +++ b/tests/sessions_object.py @@ -437,3 +437,21 @@ def sync(self): # type: ignore[override] assert issued == ["sync", "keys"] await sync assert await keys == [] + + +async def test_cancelled_takes_the_ctx_of_a_handler_with_explicit_cancellation(): + obj = sessions.SessionObject("cancelling") + + @obj.handler + async def implicit(ctx: sessions.SessionContext) -> None: + await sessions.cancelled(ctx) + + @obj.handler(explicit_cancellation=True) + async def explicit(ctx: sessions.SessionContext) -> None: + await sessions.cancel_calls(ctx) + + assert not obj.handlers["implicit"].explicit_cancellation + assert obj.handlers["explicit"].explicit_cancellation + # Over a test's own host there is no SDK context to cancel. + with pytest.raises(TerminalError, match="takes a session object's ctx"): + await run(obj, "explicit", Host(Backend(), service="cancelling")) From 6cd1175b2ea5af5f4d57996ef2bc0aad5c82e31f Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 19:00:47 +0100 Subject: [PATCH 12/45] Fix from a self-review: sessions.cancelled ends cancelled, not failed sessions.cancelled(ctx), waited for by a task all along a turn (the agents' way), was failed with the SuspendedException of the attempt's end whenever the turn ended uncancelled, and asyncio logged "Task exception was never retrieved" for every such turn. The attempt's end now ends it with asyncio.CancelledError: its task ends cancelled, quietly. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/sessions/_object.py | 16 +++++++++++++--- tests/sessions_object.py | 25 +++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index 0ffed759..ed6b55a7 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -27,6 +27,7 @@ async def list(ctx: sessions.ViewContext) -> list[str]: from __future__ import annotations +import asyncio import functools import typing from collections.abc import Awaitable, Callable, Sequence @@ -36,12 +37,12 @@ async def list(ctx: sessions.ViewContext) -> list[str]: from restate import _preview from restate.context import Request -from restate.exceptions import SuspendedException, TerminalError +from restate.exceptions import SdkInternalBaseException, SuspendedException, TerminalError from restate.handler import RESTATE_UNIQUE_HANDLER_SYMBOL, Handler, handler_from_callable from restate.object import VirtualObject from restate.retry_policy import InvocationRetryPolicy from restate.serde import BytesSerde, DefaultSerde, Serde -from restate.server_context import ServerInvocationContext +from restate.server_context import DisconnectedException, ServerInvocationContext from restate.vm import SingleUnresolvedFuture from restate.sessions import _selector from restate.sessions._calls import CallFuture, Rpc, Send, send_now @@ -301,8 +302,17 @@ async def cancelled(ctx: SessionContext) -> None: async def stop_when_cancelled(): await sessions.cancelled(ctx) await sessions.cancel_calls(ctx) + + asyncio.ensure_future(stop_when_cancelled()) + + If the attempt ends first (the handler returned, or Restate suspended + it), it raises ``asyncio.CancelledError``: a task that waits for it ends + cancelled, quietly, as a task whose work is over does. """ - await _preview.cancellation(_explicit_of(ctx, "cancelled")) + try: + await _preview.cancellation(_explicit_of(ctx, "cancelled")) + except (SdkInternalBaseException, DisconnectedException) as ended: + raise asyncio.CancelledError() from ended async def cancel_calls(ctx: SessionContext) -> list[str]: diff --git a/tests/sessions_object.py b/tests/sessions_object.py index c83f8042..d5ae05d0 100644 --- a/tests/sessions_object.py +++ b/tests/sessions_object.py @@ -455,3 +455,28 @@ async def explicit(ctx: sessions.SessionContext) -> None: # Over a test's own host there is no SDK context to cancel. with pytest.raises(TerminalError, match="takes a session object's ctx"): await run(obj, "explicit", Host(Backend(), service="cancelling")) + + +async def test_cancelled_ends_its_waiter_cancelled_when_the_attempt_ends_first(monkeypatch): + from restate.exceptions import SuspendedException # pylint: disable=import-outside-toplevel + from restate.sessions import _object # pylint: disable=import-outside-toplevel + + async def attempt_ended(_sdk: Any) -> None: + raise SuspendedException() + + monkeypatch.setattr(_object._preview, "cancellation", attempt_ended) + host = Host(Backend()) + ctx = sessions.SessionContext( + key="k", + request=request(), + db=typing.cast(Any, None), + rpc=sessions.Rpc(host), + send=sessions.Send(lambda send: None), + cancel=lambda invocation_id: None, + _sdk=object(), + _explicit_cancellation=True, + ) + watcher = asyncio.ensure_future(sessions.cancelled(ctx)) + await asyncio.wait({watcher}) + # Cancelled, not failed: asyncio reports no exception nobody retrieved. + assert watcher.cancelled() From 06656efaf9044dd92e85ec51ecf803504f7bf308 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 15:42:03 +0100 Subject: [PATCH 13/45] Type restate.sessions strictly, and lint it with more rules The packages written to be fully typed, restate.sessions and (to come) restate.agents, are strict in both type checkers: pyright's strict mode, and mypy with annotations required, no Any returned, no bare generics, no unused ignores. Ruff checks restate.sessions with more rules (python/restate/sessions/ruff.toml): annotations, bugbear, pyupgrade for Python 3.10, sorted imports, simplify, ruff's own, async pitfalls; Any and serde defaults (the SDK's API style) are allowed. What it took: - python/restate/_internal.pyi types the native sessions API (the database, sql_open, the store's ops); the VM's bindings stay untyped behind a module __getattr__, as without the stub. - The internals a package object keeps for another (a context's attempt, a database's attempt, a transaction's reset, an object's entries) are Links (_link.py, as sdk-sessions' link.ts), not hidden fields: the contexts compare by identity. Reads have query_columns, public. - Generic parameters, a typed dispatch, and the SDK's VirtualObject and Workflow name parameters annotated. The rest of the SDK keeps its configuration. All checks pass on Python 3.10 and 3.14. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- pyproject.toml | 14 ++++ python/restate/_internal.pyi | 74 +++++++++++++++++ python/restate/object.py | 2 +- python/restate/server_context.py | 4 +- python/restate/sessions/_calls.py | 2 +- python/restate/sessions/_database.py | 89 ++++++++++++--------- python/restate/sessions/_link.py | 44 ++++++++++ python/restate/sessions/_object.py | 115 +++++++++++++++------------ python/restate/sessions/_schema.py | 10 +-- python/restate/sessions/_selector.py | 21 +++-- python/restate/sessions/_session.py | 58 ++++++++------ python/restate/sessions/ruff.toml | 15 ++++ python/restate/workflow.py | 2 +- tests/sessions_object.py | 6 +- 14 files changed, 319 insertions(+), 137 deletions(-) create mode 100644 python/restate/_internal.pyi create mode 100644 python/restate/sessions/_link.py create mode 100644 python/restate/sessions/ruff.toml diff --git a/pyproject.toml b/pyproject.toml index 1970cec1..bc9cb220 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,20 @@ cache-keys = [ { file = "src/**/*.rs" }, ] +# The packages written to be fully typed: strict in both type checkers. +[tool.pyright] +strict = ["python/restate/sessions", "python/restate/agents"] + +[[tool.mypy.overrides]] +module = ["restate.sessions.*", "restate.agents.*"] +disallow_untyped_defs = true +disallow_incomplete_defs = true +disallow_any_generics = true +warn_return_any = true +warn_unused_ignores = true +no_implicit_reexport = true +strict_equality = true + [tool.ruff] line-length = 120 target-version = "py310" diff --git a/python/restate/_internal.pyi b/python/restate/_internal.pyi new file mode 100644 index 00000000..fbe66e73 --- /dev/null +++ b/python/restate/_internal.pyi @@ -0,0 +1,74 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +The native extension (``src/``), as far as it is typed: the sessions' +database (``src/sessions.rs``) is; the VM's bindings (``src/lib.rs``), which +``restate.vm`` wraps, are left to ``__getattr__``, as untyped as without this +stub. +""" + +from collections.abc import Awaitable, Callable, Mapping, Sequence +from typing import Any, TypeVar, Union + +_T = TypeVar("_T") + +_Answer = Union[_T, Awaitable[_T]] +"""A call's answer: its value, if it is ready at once, or an asyncio future of it.""" + +_SqlValue = Union[None, int, float, str, bytes] + +def __getattr__(name: str) -> Any: ... + +class PyVM: + """The invocation's VM (``restate.vm.VMWrapper`` wraps it).""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: ... + def __getattr__(self, name: str) -> Any: ... + +# --- Sessions (src/sessions.rs) ----------------------------------------------------------------- + +SQL_GET: int +SQL_SET: int +SQL_CLEAR: int +SQL_SYNC: int +SQL_KEYS: int +SQL_WAIT: int + +class SqlRefusedException(Exception): + """A failure of the database that a retry would repeat (shared-core's ``sql::Refused``).""" + +class PySqlDatabase: + """One session of a namespace's database.""" + + def query_columns( + self, sql: str, params: Sequence[Any] | None = None + ) -> _Answer[tuple[list[str], list[list[_SqlValue]]]]: ... + def begin(self) -> _Answer[None]: ... + def commit(self, records: Mapping[str, bytes] | None = None) -> _Answer[None]: ... + def rollback(self) -> _Answer[None]: ... + def is_usable(self) -> bool: ... + def record(self, name: str) -> _Answer[bytes]: ... + def save(self) -> _Answer[None]: ... + def close(self) -> _Answer[None]: ... + +class SqlReply: + """The callable ``dispatch`` answers a request with: ``reply(value)`` or ``reply(None, error)``.""" + + def __call__(self, value: object | None, error: str | None = None) -> None: ... + +def sql_open( + namespace: str, + dispatch: Callable[[int, Any, bytes | None, SqlReply | None], None], + lists_keys: bool = False, + reader: bool = False, + vm: PyVM | None = None, + name: str | None = None, +) -> _Answer[PySqlDatabase]: ... diff --git a/python/restate/object.py b/python/restate/object.py index a8d726aa..04241871 100644 --- a/python/restate/object.py +++ b/python/restate/object.py @@ -85,7 +85,7 @@ class VirtualObject: def __init__( self, - name, + name: str, description: typing.Optional[str] = None, metadata: typing.Optional[typing.Dict[str, str]] = None, inactivity_timeout: typing.Optional[timedelta] = None, diff --git a/python/restate/server_context.py b/python/restate/server_context.py index e7e0dc5f..aa22cfc8 100644 --- a/python/restate/server_context.py +++ b/python/restate/server_context.py @@ -418,9 +418,9 @@ class Tasks: """ def __init__(self) -> None: - self.tasks: set[asyncio.Future] = set() + self.tasks: set[asyncio.Future[Any]] = set() - def add(self, task: asyncio.Future): + def add(self, task: asyncio.Future[Any]) -> None: """Add a task to the list.""" self.tasks.add(task) diff --git a/python/restate/sessions/_calls.py b/python/restate/sessions/_calls.py index 825313b1..fd16f0fe 100644 --- a/python/restate/sessions/_calls.py +++ b/python/restate/sessions/_calls.py @@ -104,7 +104,7 @@ def __init__(self, call: RestateDurableCallFuture[bytes], serde: Serde[O]) -> No self._serde = serde async def _result(self) -> O: - return typing.cast(O, self._serde.deserialize(typing.cast(bytes, await self._call))) + return typing.cast(O, self._serde.deserialize(await self._call)) def __await__(self) -> Generator[Any, Any, O]: return self._result().__await__() diff --git a/python/restate/sessions/_database.py b/python/restate/sessions/_database.py index 1c4b83c9..4306c615 100644 --- a/python/restate/sessions/_database.py +++ b/python/restate/sessions/_database.py @@ -23,6 +23,7 @@ from __future__ import annotations import asyncio +import contextlib import contextvars import inspect import typing @@ -37,14 +38,16 @@ SQL_SET, SQL_SYNC, SQL_WAIT, + PySqlDatabase, SqlRefusedException, + SqlReply, sql_open, ) -SqlValue = typing.Union[None, int, float, str, bytes] +SqlValue = None | int | float | str | bytes """An SQL value: ``None``, ``int`` (64 bits), ``float``, ``str``, or ``bytes`` (BLOB).""" -SqlParam = typing.Union[SqlValue, bool, bytearray, memoryview] +SqlParam = SqlValue | bool | bytearray | memoryview """An SQL parameter: an ``SqlValue``, a ``bool`` (stored as 0 or 1), or a BLOB as any bytes-like object.""" Row = dict[str, SqlValue] @@ -95,6 +98,18 @@ def clear(self, key: str) -> None: ... def sync(self) -> None | Awaitable[None]: ... +class _ListsKeys(Protocol): + """A store that lists its keys (see ``Store``).""" + + def keys(self) -> Sequence[str] | Awaitable[Sequence[str]]: ... + + +class _Waits(Protocol): + """A store over the invocation's VM, which waits for a read (see ``Store``).""" + + def wait(self, handle: int) -> Awaitable[None]: ... + + async def open(namespace: str, store: Store, *, reader: bool = False) -> Connection: # pylint: disable=redefined-builtin """ Open the database of ``namespace`` over ``store``. One process-wide page @@ -114,9 +129,14 @@ async def open(namespace: str, store: Store, *, reader: bool = False) -> Connect return Connection(native, boundary) -async def _settled(answer: Any) -> Any: +T = typing.TypeVar("T") + + +async def _settled(answer: T | Awaitable[T]) -> T: """A call's answer: its value, as it gave it, or once the asyncio future it gave is done.""" - return await answer if inspect.isawaitable(answer) else answer + if inspect.isawaitable(answer): + return await answer + return answer class _Boundary: @@ -147,7 +167,7 @@ def _keep(self, error: BaseException) -> None: if self.failed is None: self.failed = error - def dispatch(self, op: int, key: Any, value: bytes | None, reply: Any) -> None: + def dispatch(self, op: int, key: Any, value: bytes | None, reply: SqlReply | None) -> None: """One request of the extension's: ``key`` is a key of the store, or for ``SQL_WAIT`` a handle of the VM.""" if op in (SQL_SET, SQL_CLEAR): try: @@ -160,13 +180,15 @@ def dispatch(self, op: int, key: Any, value: bytes | None, reply: Any) -> None: if self._write_failed is None: self._write_failed = error return + # Every request but a write has a reply (src/sessions.rs). + assert reply is not None try: if op == SQL_GET: - result: Any = self.store.get(key) + result: object = self.store.get(key) elif op == SQL_KEYS: - result = self.store.keys() # type: ignore[attr-defined] + result = typing.cast(_ListsKeys, self.store).keys() elif op == SQL_WAIT: - result = self.store.wait(key) # type: ignore[attr-defined] + result = typing.cast(_Waits, self.store).wait(key) elif self._write_failed is not None: reply(None, str(self._write_failed)) return @@ -184,8 +206,8 @@ def dispatch(self, op: int, key: Any, value: bytes | None, reply: Any) -> None: else: reply(_reply_value(op, result)) - def _answer(self, op: int, reply: Any, awaitable: Awaitable[Any]) -> None: - def done(future: asyncio.Future) -> None: + def _answer(self, op: int, reply: SqlReply, awaitable: Awaitable[object]) -> None: + def done(future: asyncio.Future[object]) -> None: if future.cancelled(): self._keep(asyncio.CancelledError()) reply(None, "the store's request was cancelled") @@ -199,27 +221,27 @@ def done(future: asyncio.Future) -> None: asyncio.ensure_future(awaitable, loop=self.loop).add_done_callback(done) -def _reply_value(op: int, result: Any) -> Any: +def _reply_value(op: int, result: object) -> object: """A request's answer as the extension reads it: a value, the keys as a list, or nothing for a sync.""" if op in (SQL_SYNC, SQL_WAIT): return None if op == SQL_KEYS: - return list(result) + return list(typing.cast("Sequence[str]", result)) return result -def _objects(columns: list[str], rows: Rows) -> list[Row]: +def objects_of(columns: list[str], rows: Rows) -> list[Row]: """The rows as dicts by column name; a column name twice is an ``SqlError``.""" if len(set(columns)) != len(columns): twice = next(name for i, name in enumerate(columns) if name in columns[:i]) raise SqlError(f"the column {twice} is named twice: alias one (AS), or read the rows with values()") - return [dict(zip(columns, row)) for row in rows] + return [dict(zip(columns, row, strict=True)) for row in rows] class _Sql: """SQL on a session of the extension: ``query``, ``get``, ``values`` and ``execute`` by way of ``query_columns``.""" - def __init__(self, native: Any, boundary: _Boundary) -> None: + def __init__(self, native: PySqlDatabase, boundary: _Boundary) -> None: self._native = native self._boundary = boundary @@ -234,7 +256,7 @@ async def query(self, sql: str, params: Sequence[SqlParam] | None = None) -> lis expression needs ``AS``), and two of the same name are an ``SqlError``: alias them, or use ``values``. """ - return _objects(*await self.query_columns(sql, params)) + return objects_of(*await self.query_columns(sql, params)) async def get(self, sql: str, params: Sequence[SqlParam] | None = None) -> Row | None: """The first row of ``query``, or ``None`` when there is none.""" @@ -249,7 +271,7 @@ async def execute(self, sql: str, params: Sequence[SqlParam] | None = None) -> N """Run SQL (one or more statements), discarding rows; parameters only with a single statement.""" await self.query_columns(sql, params) - async def _request(self, call: Callable[[], Any]) -> Any: + async def _request(self, call: Callable[[], T | Awaitable[T]]) -> T: """ ``call``'s result, or its ``_failure``: ``call`` makes the request, as the extension raises as it is called (arguments it cannot convert, and @@ -257,10 +279,7 @@ async def _request(self, call: Callable[[], Any]) -> Any: awaitable of it. """ try: - result = call() - if inspect.isawaitable(result): - result = await result - return result + return await _settled(call()) except Exception as error: raise self._failure(error) from None @@ -297,7 +316,7 @@ class Connection(_Sql): as it would wait for that body: the body uses its ``tx``. """ - def __init__(self, native: Any, boundary: _Boundary) -> None: + def __init__(self, native: PySqlDatabase, boundary: _Boundary) -> None: super().__init__(native, boundary) # FIFO: calls run in the order they were made. self._lock = asyncio.Lock() @@ -336,19 +355,17 @@ async def transaction(self) -> AsyncIterator[ConnectionTransaction]: try: yield tx except BaseException: - await tx._end() # pylint: disable=protected-access + await tx.end() await self._rollback() raise finally: _running.reset(token) - records = await tx._end() # pylint: disable=protected-access + records = await tx.end() # Its commit's sync runs what the body gave `with_commit`; one that # wrote nothing synced nothing, and it runs after. - self._boundary.with_sync = tx._commit_runs() # pylint: disable=protected-access + self._boundary.with_sync = tx.commit_runs() try: - result = self._native.commit(records) - if inspect.isawaitable(result): - await result + await _settled(self._native.commit(records)) except Exception as error: # A COMMIT that fails (a deferred constraint) leaves the # transaction open: without the rollback, the calls after it @@ -384,12 +401,8 @@ async def close(self) -> None: async def _rollback(self) -> None: """Roll back, unless a storage error ended the session: the caller raises its own error, which says why.""" - try: - result = self._native.rollback() - if inspect.isawaitable(result): - await result - except Exception: # pylint: disable=broad-exception-caught - pass + with contextlib.suppress(Exception): + await _settled(self._native.rollback()) class ConnectionTransaction(_Sql): @@ -400,7 +413,9 @@ class ConnectionTransaction(_Sql): before its commit or rollback, so none outlives it. """ - def __init__(self, native: Any, boundary: _Boundary, connection: Connection, outer: ConnectionTransaction | None): + def __init__( + self, native: PySqlDatabase, boundary: _Boundary, connection: Connection, outer: ConnectionTransaction | None + ) -> None: super().__init__(native, boundary) self._connection = connection # The transaction whose body this one's runs in, if any: of another @@ -457,14 +472,14 @@ def with_commit(self, run: Callable[[], None]) -> None: self._refuse_ended() self._with_commits.append(run) - async def _end(self) -> dict[str, bytes]: + async def end(self) -> dict[str, bytes]: """End it, refusing calls from now on, once the calls made before have completed; then the records it set.""" self._ended = True async with self._lock: pass return dict(self._records) - def _commit_runs(self) -> Callable[[], None] | None: + def commit_runs(self) -> Callable[[], None] | None: """What runs with its commit, once it ended: see ``with_commit``.""" runs = self._with_commits if not runs: diff --git a/python/restate/sessions/_link.py b/python/restate/sessions/_link.py new file mode 100644 index 00000000..6e024186 --- /dev/null +++ b/python/restate/sessions/_link.py @@ -0,0 +1,44 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +A private link on the objects this package makes: set once, as it makes one, +and read by its code alone (a context's attempt, a database's attempt, a +transaction's reset, an object's entries). The object reads as its fields +alone, and its link goes with it. As sdk-sessions' ``link.ts``. +""" + +from __future__ import annotations + +import typing +import weakref + +T = typing.TypeVar("T") +O = typing.TypeVar("O") + + +class Link(typing.Generic[T]): + """A link of one kind, on the objects that have one: compared by identity, kept as long as the object is.""" + + def __init__(self, what: str) -> None: + self.what = what + self._links: weakref.WeakKeyDictionary[object, T] = weakref.WeakKeyDictionary() + + def on(self, owner: O, value: T) -> O: + """``owner``, linked to ``value``: once, as it is made.""" + self._links[owner] = value + return owner + + def of(self, owner: object) -> T | None: + """What ``owner`` is linked to, if it has this link (and None for anything that cannot have one).""" + try: + return self._links.get(owner) + except TypeError: + return None diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index ed6b55a7..36198702 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -31,7 +31,7 @@ async def list(ctx: sessions.ViewContext) -> list[str]: import functools import typing from collections.abc import Awaitable, Callable, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass from datetime import timedelta from typing import Any @@ -43,12 +43,13 @@ async def list(ctx: sessions.ViewContext) -> list[str]: from restate.retry_policy import InvocationRetryPolicy from restate.serde import BytesSerde, DefaultSerde, Serde from restate.server_context import DisconnectedException, ServerInvocationContext -from restate.vm import SingleUnresolvedFuture from restate.sessions import _selector from restate.sessions._calls import CallFuture, Rpc, Send, send_now from restate.sessions._database import Store +from restate.sessions._link import Link from restate.sessions._schema import Schema, Step from restate.sessions._session import Attempt, Database, DatabaseReader, ObjectConfig, SqlHost, reader_of, with_attempt +from restate.vm import SingleUnresolvedFuture F = typing.TypeVar("F", bound=Callable[..., Awaitable[Any]]) @@ -73,7 +74,7 @@ async def list(ctx: sessions.ViewContext) -> list[str]: """ -@dataclass(frozen=True) +@dataclass(frozen=True, eq=False) class SessionContext: """ The context of a session object's handler. A handler runs in attempts: @@ -99,13 +100,9 @@ class SessionContext: """Sends, made at once; a transaction's (``tx.send``) are made once it commits.""" cancel: Cancel """Cancel an invocation, by its id: see ``Cancel``.""" - _sdk: Any = field(default=None, repr=False, compare=False) - """The SDK's context of the attempt, for ``cancelled`` and ``cancel_calls``.""" - _explicit_cancellation: bool = field(default=False, repr=False, compare=False) - """Whether its handler has ``explicit_cancellation=True``.""" -@dataclass(frozen=True) +@dataclass(frozen=True, eq=False) class ViewContext: """ The context of a session object's view: it runs beside the handlers, and @@ -237,6 +234,20 @@ def cancel(invocation_id: str) -> None: return cancel +@dataclass(frozen=True) +class Scope: + """An attempt's own, as this package made it, on its context (``SCOPE``).""" + + sdk: ServerInvocationContext | None + """The SDK's context of the attempt: none over a test's own host.""" + explicit_cancellation: bool = False + """Whether its handler has ``explicit_cancellation=True``.""" + + +SCOPE: Link[Scope] = Link("scope") +"""The attempt's scope, on its context.""" + + @dataclass(frozen=True) class _Entry: """A handler or view of a session object: its function, kind, and object's config.""" @@ -247,22 +258,21 @@ class _Entry: explicit_cancellation: bool = False +ENTRIES: Link[dict[str, _Entry]] = Link("entries") +"""A session object's handlers and views, by name: for ``invoke``.""" + + def _context_of(entry: _Entry, host: SqlHost, request: Request, attempt: Attempt) -> SessionContext | ViewContext: """The context of an attempt of ``entry``.""" rpc = Rpc(host) cancel = _cancel_of(host) if entry.view: return ViewContext(key=host.key, request=request, db=reader_of(attempt.db), rpc=rpc, cancel=cancel) - return SessionContext( - key=host.key, - request=request, - db=attempt.db, - rpc=rpc, - send=Send(send_now(host)), - cancel=cancel, - _sdk=getattr(host, "context", None), - _explicit_cancellation=entry.explicit_cancellation, + context = SessionContext( + key=host.key, request=request, db=attempt.db, rpc=rpc, send=Send(send_now(host)), cancel=cancel ) + sdk = typing.cast("ServerInvocationContext | None", getattr(host, "context", None)) + return SCOPE.on(context, Scope(sdk=sdk, explicit_cancellation=entry.explicit_cancellation)) async def _attempt(entry: _Entry, host: SqlHost, request: Request, args: Sequence[Any]) -> Any: @@ -284,13 +294,14 @@ async def handler(ctx: ServerInvocationContext, *args: Any) -> Any: return handler -def _explicit_of(ctx: SessionContext, operation: str) -> Any: +def _explicit_of(ctx: SessionContext, operation: str) -> ServerInvocationContext: """The SDK's context behind ``ctx``, a handler's with ``explicit_cancellation=True``.""" - if not isinstance(ctx, SessionContext) or ctx._sdk is None: # pylint: disable=protected-access + scope = SCOPE.of(ctx) + if scope is None or scope.sdk is None: raise TerminalError(f"{operation} takes a session object's ctx") - if not ctx._explicit_cancellation: # pylint: disable=protected-access + if not scope.explicit_cancellation: raise TerminalError(f"{operation} takes the ctx of a handler with explicit_cancellation=True") - return ctx._sdk # pylint: disable=protected-access + return scope.sdk async def cancelled(ctx: SessionContext) -> None: @@ -326,7 +337,7 @@ async def cancel_calls(ctx: SessionContext) -> list[str]: async def invoke(obj: SessionObject, name: str, host: SqlHost, request: Request, *args: Any) -> Any: """Run handler or view ``name`` of ``obj`` for one attempt over ``host``, a context of a test's own, without Restate.""" - entry = obj._entries.get(name) # pylint: disable=protected-access + entry = (ENTRIES.of(obj) or {}).get(name) if entry is None: raise ValueError(f"{obj.name}.{name} is no handler or view of a session object") return await _attempt(entry, host, request, args) @@ -379,7 +390,7 @@ def __init__( invocation_retry_policy=invocation_retry_policy, ) self.config = ObjectConfig(schema=Schema(name, schema) if schema else None) - self._entries: dict[str, _Entry] = {} + ENTRIES.on(self, {}) @typing.overload # type: ignore[override] def handler(self, fn: F, /) -> F: ... # pyright: ignore[reportIncompatibleMethodOverride] @@ -437,20 +448,20 @@ async def add(ctx: sessions.SessionContext, text: str) -> None: return self._decorate( False, name, - dict( - explicit_cancellation=explicit_cancellation, - accept=accept, - content_type=content_type, - input_serde=input_serde, - output_serde=output_serde, - metadata=metadata, - inactivity_timeout=inactivity_timeout, - abort_timeout=abort_timeout, - journal_retention=journal_retention, - idempotency_retention=idempotency_retention, - ingress_private=ingress_private, - invocation_retry_policy=invocation_retry_policy, - ), + { + "explicit_cancellation": explicit_cancellation, + "accept": accept, + "content_type": content_type, + "input_serde": input_serde, + "output_serde": output_serde, + "metadata": metadata, + "inactivity_timeout": inactivity_timeout, + "abort_timeout": abort_timeout, + "journal_retention": journal_retention, + "idempotency_retention": idempotency_retention, + "ingress_private": ingress_private, + "invocation_retry_policy": invocation_retry_policy, + }, ) @typing.overload @@ -507,19 +518,19 @@ async def count(ctx: sessions.ViewContext) -> int: return self._decorate( True, name, - dict( - accept=accept, - content_type=content_type, - input_serde=input_serde, - output_serde=output_serde, - metadata=metadata, - inactivity_timeout=inactivity_timeout, - abort_timeout=abort_timeout, - journal_retention=journal_retention, - idempotency_retention=idempotency_retention, - ingress_private=ingress_private, - invocation_retry_policy=invocation_retry_policy, - ), + { + "accept": accept, + "content_type": content_type, + "input_serde": input_serde, + "output_serde": output_serde, + "metadata": metadata, + "inactivity_timeout": inactivity_timeout, + "abort_timeout": abort_timeout, + "journal_retention": journal_retention, + "idempotency_retention": idempotency_retention, + "ingress_private": ingress_private, + "invocation_retry_policy": invocation_retry_policy, + }, ) def _decorate(self, view: bool, name: str | F | None, options: dict[str, Any]) -> F | Callable[[F], F]: @@ -546,7 +557,9 @@ def register(fn: F) -> F: # would have: each runs from the start, against the state as it is. sdk.explicit_recovery = True sdk.explicit_cancellation = explicit - self._entries[sdk.name] = entry + entries = ENTRIES.of(self) + assert entries is not None + entries[sdk.name] = entry return typing.cast(F, registered) return register diff --git a/python/restate/sessions/_schema.py b/python/restate/sessions/_schema.py index cda8829b..f24a2a38 100644 --- a/python/restate/sessions/_schema.py +++ b/python/restate/sessions/_schema.py @@ -48,7 +48,7 @@ never loaded too. """ -Step = typing.Union[str, Callable[[Any], Awaitable[None]]] +Step = str | Callable[[Any], Awaitable[None]] """A statement, or an async function of the transaction.""" _ENSURE = "_restate_sessions_ensure" @@ -167,11 +167,11 @@ async def _children_first(tx: ConnectionTransaction, tables: list[str]) -> list[ own = name.lower() parents[name] = {typing.cast(str, table).lower() for (table,) in keys} - {own} ordered: list[str] = [] - left = tables + left: list[str] = tables while left: - named = {parent for name in left for parent in parents[name]} - unnamed = [name for name in left if name.lower() not in named] + named = {parent for table in left for parent in parents[table]} + unnamed = [table for table in left if table.lower() not in named] step = unnamed or left ordered.extend(step) - left = [name for name in left if name not in step] + left = [table for table in left if table not in step] return ordered diff --git a/python/restate/sessions/_selector.py b/python/restate/sessions/_selector.py index 26619065..2c0fd458 100644 --- a/python/restate/sessions/_selector.py +++ b/python/restate/sessions/_selector.py @@ -31,7 +31,6 @@ from restate.exceptions import SdkInternalException, SuspendedException, TerminalError from restate.server_context import DisconnectedException, ServerInvocationContext from restate.vm import ( - AllCompletedUnresolvedFuture, DoProgressAnyCompleted, DoProgressCancelSignalReceived, DoProgressExecuteRun, @@ -76,7 +75,7 @@ def notify_input(self, chunk: dict[str, Any]) -> None: """Passes a chunk of the request body to the VM.""" ... - def track(self, task: asyncio.Future) -> None: + def track(self, task: asyncio.Future[Any]) -> None: """Cancels ``task`` at the end of the attempt.""" ... @@ -91,9 +90,7 @@ def is_ready(future: UnresolvedFuture, completed: Callable[[int], bool]) -> bool return completed(future.handle) if isinstance(future, FirstCompletedUnresolvedFuture): return any(is_ready(child, completed) for child in future.children) - if isinstance(future, AllCompletedUnresolvedFuture): - return all(is_ready(child, completed) for child in future.children) - raise TypeError(f"Unknown UnresolvedFuture variant: {type(future).__name__}") + return all(is_ready(child, completed) for child in future.children) class FutureSelector: @@ -111,8 +108,8 @@ class FutureSelector: def __init__(self, driver: SelectorDriver, disconnected: Callable[[], BaseException]) -> None: self._driver = driver self._disconnected = disconnected - self._waits: dict[asyncio.Future, UnresolvedFuture] = {} - self._task: asyncio.Task | None = None + self._waits: dict[asyncio.Future[None], UnresolvedFuture] = {} + self._task: asyncio.Task[None] | None = None # Between the VM's WaitExternalProgress and the next input. self._waiting = False # Why the attempt ended, once it has: each wait then raises a new one. @@ -280,13 +277,13 @@ async def run() -> None: def notify_input(self, chunk: dict[str, Any]) -> None: # Skip empty body frames, as the SDK does. - body = chunk.get("body", None) + body = chunk.get("body") if body is not None and len(body) > 0: self.context.vm.notify_input(body) if not chunk.get("more_body", False): self.context.vm.notify_input_closed() - def track(self, task: asyncio.Future) -> None: + def track(self, task: asyncio.Future[Any]) -> None: self.context.tasks.add(task) @@ -298,7 +295,7 @@ def hook(context: ServerInvocationContext, on_end: Callable[[], None]) -> Future hook lasts as long as ``context``, one attempt. """ selector = FutureSelector(_ContextDriver(context, on_end), DisconnectedException) - started: set[asyncio.Task] = set() + started: set[asyncio.Task[Any]] = set() def starting(make: Callable[..., Any]) -> Callable[..., Any]: def make_started(*args: Any, **kwargs: Any) -> Any: @@ -316,10 +313,10 @@ def make_started(*args: Any, **kwargs: Any) -> Any: for name in ("create_future", "create_sleep_future", "create_call_future"): setattr(context, name, starting(getattr(context, name))) - setattr(context, "create_poll_or_cancel_coroutine", selector.wait) + context.create_poll_or_cancel_coroutine = selector.wait # type: ignore[method-assign,assignment] return selector -def _retrieve(task: asyncio.Future) -> None: +def _retrieve(task: asyncio.Future[Any]) -> None: if not task.cancelled(): task.exception() diff --git a/python/restate/sessions/_session.py b/python/restate/sessions/_session.py index 8aa52272..a039ba92 100644 --- a/python/restate/sessions/_session.py +++ b/python/restate/sessions/_session.py @@ -19,6 +19,7 @@ import asyncio import base64 +import contextlib import json import re import time @@ -41,9 +42,12 @@ SqlError, SqlParam, Store, - _objects, + objects_of, +) +from restate.sessions._database import ( open as open_connection, ) +from restate.sessions._link import Link from restate.sessions._schema import Schema, apply, drop_all, ensure_of, ensuring T = typing.TypeVar("T") @@ -347,7 +351,8 @@ def _refused_send(_: GenericSend) -> None: class DatabaseReader: """Reads of a database, or of one transaction of it.""" - async def _query_columns(self, sql: str, params: Sequence[SqlParam] | None) -> tuple[list[str], Rows]: + async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: + """The rows of one statement as lists, with the column names of the statement: ``(columns, rows)``.""" raise NotImplementedError async def query(self, sql: str, params: Sequence[SqlParam] | None = None) -> list[Row]: @@ -357,9 +362,9 @@ async def query(self, sql: str, params: Sequence[SqlParam] | None = None) -> lis expression needs ``AS``), and two of the same name are an error: alias them, or use ``values``. """ - columns, rows = await self._query_columns(sql, params) + columns, rows = await self.query_columns(sql, params) try: - return _objects(columns, rows) + return objects_of(columns, rows) except SqlError as error: _raise_for_handler(error) @@ -370,7 +375,7 @@ async def get(self, sql: str, params: Sequence[SqlParam] | None = None) -> Row | async def values(self, sql: str, params: Sequence[SqlParam] | None = None) -> Rows: """The rows of one statement as lists, in column order: for adapters, and for columns that share a name.""" - return (await self._query_columns(sql, params))[1] + return (await self.query_columns(sql, params))[1] class DatabaseWriter(DatabaseReader): @@ -378,7 +383,7 @@ class DatabaseWriter(DatabaseReader): async def execute(self, sql: str, params: Sequence[SqlParam] | None = None) -> None: """Run SQL (one or more statements), discarding rows; parameters only with a single statement.""" - await self._query_columns(sql, params) + await self.query_columns(sql, params) class Transaction(DatabaseWriter): @@ -391,12 +396,11 @@ class Transaction(DatabaseWriter): again is dropped. """ - def __init__(self, handle: ConnectionTransaction, send: Send, reset: Callable[[], Awaitable[None]] | None) -> None: + def __init__(self, handle: ConnectionTransaction, send: Send) -> None: self._handle = handle self.send = send - self._reset = reset - async def _query_columns(self, sql: str, params: Sequence[SqlParam] | None) -> tuple[list[str], Rows]: + async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: try: return await self._handle.query_columns(sql, params) except SqlError as error: @@ -412,8 +416,9 @@ class Database(DatabaseWriter): def __init__(self, attempt: Attempt) -> None: self._attempt = attempt + _ATTEMPT.on(self, attempt) - async def _query_columns(self, sql: str, params: Sequence[SqlParam] | None) -> tuple[list[str], Rows]: + async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: return await self._attempt.call(lambda db: db.query_columns(sql, params)) def transaction(self) -> typing.AsyncContextManager[Transaction]: @@ -439,8 +444,8 @@ class _ReaderOf(DatabaseReader): def __init__(self, db: Database) -> None: self._db = db - async def _query_columns(self, sql: str, params: Sequence[SqlParam] | None) -> tuple[list[str], Rows]: - return await self._db._query_columns(sql, params) # pylint: disable=protected-access + async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: + return await self._db.query_columns(sql, params) def reader_of(db: Database) -> DatabaseReader: @@ -452,6 +457,12 @@ def reader_of(db: Database) -> DatabaseReader: _DATABASE_NAME = re.compile(r"^[\w.-]{1,64}$") +_ATTEMPT: Link[Attempt] = Link("database") +"""The attempt of a ``ctx.db``: for ``database``.""" + +_RESET: Link[Callable[[], Awaitable[None]]] = Link("reset") +"""How a transaction of an exclusive handler's ``ctx.db`` empties the database: for ``reset``.""" + def database(db: Database, name: str) -> Database: """ @@ -463,9 +474,10 @@ def database(db: Database, name: str) -> Database: """ if not _DATABASE_NAME.match(name): raise TerminalError(f"database: {name!r} is not a database name") - if not isinstance(db, Database): + attempt = _ATTEMPT.of(db) + if attempt is None: raise TerminalError("database takes a handler's ctx.db") - return db._attempt.database(name) # pylint: disable=protected-access + return attempt.database(name) async def reset(tx: Transaction) -> None: @@ -476,10 +488,11 @@ async def reset(tx: Transaction) -> None: the outbox stays. It commits with ``tx``: what ``tx`` writes after it are the new database's first rows. """ - if not isinstance(tx, Transaction) or tx._reset is None: # pylint: disable=protected-access + empty = _RESET.of(tx) + if empty is None: raise TerminalError("reset takes a transaction of an exclusive handler's ctx.db") try: - await tx._reset() # pylint: disable=protected-access + await empty() except SqlError as error: _raise_for_handler(error) @@ -649,8 +662,9 @@ async def transaction(self) -> AsyncIterator[Transaction]: async with db.transaction() as handle: sends = _Sends() if self._has_outbox else None # A view's writes stay in its session: it has nothing to reset. - resetter = None if self._host.shared else (lambda: self._reset(handle)) - tx = Transaction(handle, sends.send if sends is not None else Send(_refused_send), resetter) + tx = Transaction(handle, sends.send if sends is not None else Send(_refused_send)) + if not self._host.shared: + _RESET.on(tx, lambda: self._reset(handle)) try: yield tx finally: @@ -761,10 +775,8 @@ async def _close(self) -> None: await self._save(db) return if db is not None: - try: + with contextlib.suppress(Exception): await db.close() - except Exception: # pylint: disable=broad-exception-caught - pass if not self._has_outbox or self._host.attempt_end() is not None: return again = self._open() @@ -780,10 +792,8 @@ async def _save(self, db: Connection) -> None: await self._flush(db) await db.save() except BaseException: - try: + with contextlib.suppress(Exception): await db.close() - except Exception: # pylint: disable=broad-exception-caught - pass raise async def _reset(self, handle: ConnectionTransaction) -> None: diff --git a/python/restate/sessions/ruff.toml b/python/restate/sessions/ruff.toml new file mode 100644 index 00000000..6b09b194 --- /dev/null +++ b/python/restate/sessions/ruff.toml @@ -0,0 +1,15 @@ +# The project's ruff settings, and more rules for this package, which is +# written to them: annotations everywhere, likely bugs, modern syntax for +# Python 3.10, sorted imports, simpler code, async pitfalls. +extend = "../../../pyproject.toml" + +[lint] +extend-select = ["ANN", "B", "C4", "I", "UP", "SIM", "RUF", "ASYNC", "PIE", "RET"] +extend-ignore = [ + # Any is what a handler's input or a serialized value is; the type + # checkers (pyright strict) catch unknown types where they matter. + "ANN401", + # A serde instance as a default (input_serde=DefaultSerde()) is the + # SDK's API style, which SessionObject.handler mirrors. + "B008", +] diff --git a/python/restate/workflow.py b/python/restate/workflow.py index 7a6a1d93..a3523264 100644 --- a/python/restate/workflow.py +++ b/python/restate/workflow.py @@ -91,7 +91,7 @@ class Workflow: def __init__( self, - name, + name: str, description: typing.Optional[str] = None, metadata: typing.Optional[typing.Dict[str, str]] = None, inactivity_timeout: typing.Optional[timedelta] = None, diff --git a/tests/sessions_object.py b/tests/sessions_object.py index d5ae05d0..50557006 100644 --- a/tests/sessions_object.py +++ b/tests/sessions_object.py @@ -458,13 +458,14 @@ async def explicit(ctx: sessions.SessionContext) -> None: async def test_cancelled_ends_its_waiter_cancelled_when_the_attempt_ends_first(monkeypatch): + from restate import _preview # pylint: disable=import-outside-toplevel from restate.exceptions import SuspendedException # pylint: disable=import-outside-toplevel from restate.sessions import _object # pylint: disable=import-outside-toplevel async def attempt_ended(_sdk: Any) -> None: raise SuspendedException() - monkeypatch.setattr(_object._preview, "cancellation", attempt_ended) + monkeypatch.setattr(_preview, "cancellation", attempt_ended) host = Host(Backend()) ctx = sessions.SessionContext( key="k", @@ -473,9 +474,8 @@ async def attempt_ended(_sdk: Any) -> None: rpc=sessions.Rpc(host), send=sessions.Send(lambda send: None), cancel=lambda invocation_id: None, - _sdk=object(), - _explicit_cancellation=True, ) + _object.SCOPE.on(ctx, _object.Scope(sdk=typing.cast(Any, object()), explicit_cancellation=True)) watcher = asyncio.ensure_future(sessions.cancelled(ctx)) await asyncio.wait({watcher}) # Cancelled, not failed: asyncio reports no exception nobody retrieved. From fae4abc8e48de32b943d38455e8fbfbe67e165ff Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 15:42:03 +0100 Subject: [PATCH 14/45] Lint the native extension in just verify, and so in CI `just lint-rust`: cargo fmt --check and clippy with warnings as errors, which CI did not run. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- Justfile | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/Justfile b/Justfile index 49e24421..6645a9f2 100644 --- a/Justfile +++ b/Justfile @@ -16,6 +16,11 @@ lint: uv run ruff format --check uv run ruff check +# The native extension (src/): formatted, and clean of clippy's warnings. +lint-rust: + cargo fmt --check + cargo clippy --all-targets -- -D warnings + typecheck-pyright: PYRIGHT_PYTHON_IGNORE_WARNINGS=1 uv run pyright python/ PYRIGHT_PYTHON_IGNORE_WARNINGS=1 uv run pyright examples/ @@ -34,7 +39,7 @@ test: # Recipe to run both mypy and pylint -verify: format lint typecheck test +verify: format lint lint-rust typecheck test @echo "Type checking and linting completed successfully." # Recipe to build the project From 74d448cf914b44f41d05d2e7c91c470d9116ea81 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 15:47:35 +0100 Subject: [PATCH 15/45] Clean restate.sessions under basedpyright, and report deprecations basedpyright, an editor's, flagged what pyright and mypy let through: - @asynccontextmanager functions annotated -> AsyncIterator[...] (deprecated): -> AsyncGenerator[..., None], here and in the SDK's pydantic integration and tests. - What `# type: ignore` hid from basedpyright: the host's VM and its wait for a read are now part of SqlHost; the store's optional keys and wait are looked up once rather than cast to; the selector's events are mappings, not cast to dicts. - Instance attributes annotated; the abstract reads declared abstract; the store's protocol positional-only; the selector imported as a module of its own (no cycle through the package). [tool.pyright] sets the standard mode, which pyright has by default and basedpyright reads from it (it refuses a [tool.basedpyright] beside it), and reports deprecations: restate.sessions has no error nor warning in either, nor in mypy; Any, serde defaults and @override (typing 3.12) are left to recommended mode. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- pyproject.toml | 7 +++ python/restate/ext/pydantic/_agent.py | 6 +-- python/restate/ext/pydantic/_model.py | 4 +- python/restate/sessions/_calls.py | 8 +-- python/restate/sessions/_database.py | 60 ++++++++++----------- python/restate/sessions/_link.py | 2 +- python/restate/sessions/_object.py | 23 ++++---- python/restate/sessions/_schema.py | 2 +- python/restate/sessions/_selector.py | 25 +++++---- python/restate/sessions/_session.py | 78 ++++++++++++++------------- tests/entry_codec.py | 2 +- tests/servercontext.py | 4 +- tests/sessions_object.py | 6 +++ tests/sessions_selector.py | 3 +- 14 files changed, 123 insertions(+), 107 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index bc9cb220..15a2e26a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,8 +58,15 @@ cache-keys = [ ] # The packages written to be fully typed: strict in both type checkers. +# basedpyright (an editor's, say) reads this section too, as there is no +# [tool.basedpyright] (it refuses both): the mode is set, standard, for it to +# check what pyright does rather than its stricter default, whose Any rules +# do not fit an SDK (Any is what a handler's input or a serialized value is). +# Deprecations are reported, by both. [tool.pyright] +typeCheckingMode = "standard" strict = ["python/restate/sessions", "python/restate/agents"] +reportDeprecated = "warning" [[tool.mypy.overrides]] module = ["restate.sessions.*", "restate.agents.*"] diff --git a/python/restate/ext/pydantic/_agent.py b/python/restate/ext/pydantic/_agent.py index cc7888cf..c6d65288 100644 --- a/python/restate/ext/pydantic/_agent.py +++ b/python/restate/ext/pydantic/_agent.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import AsyncIterable, AsyncIterator, Iterator, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Generator, Sequence from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager from typing import Any, overload @@ -133,7 +133,7 @@ def set_context(toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDe self._toolsets = [toolset.visit_and_replace(set_context) for toolset in wrapped.toolsets] @contextmanager - def _restate_overrides(self) -> Iterator[None]: + def _restate_overrides(self) -> Generator[None, None, None]: with ( super().override(model=self._model, toolsets=self._toolsets, tools=[], native_tools=[]), ): @@ -378,7 +378,7 @@ async def run_stream( event_stream_handler: EventStreamHandler[AgentDepsT] | None = None, capabilities: Sequence[AgentCapability[AgentDepsT]] | None = None, spec: dict[str, Any] | AgentSpec | None = None, - ) -> AsyncIterator[StreamedRunResult[AgentDepsT, Any]]: + ) -> AsyncGenerator[StreamedRunResult[AgentDepsT, Any], None]: """Run the agent with a user prompt in async mode, returning a streamed response. Example: diff --git a/python/restate/ext/pydantic/_model.py b/python/restate/ext/pydantic/_model.py index fd162da8..aa830ce1 100644 --- a/python/restate/ext/pydantic/_model.py +++ b/python/restate/ext/pydantic/_model.py @@ -1,4 +1,4 @@ -from collections.abc import AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator from contextlib import asynccontextmanager from datetime import datetime from typing import Any @@ -91,7 +91,7 @@ async def request_stream( model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters, run_context: RunContext[AgentDepsT] | None = None, - ) -> AsyncIterator[StreamedResponse]: + ) -> AsyncGenerator[StreamedResponse, None]: if run_context is None: raise UserError( "A model cannot be used with `pydantic_ai.direct.model_request_stream()` as it requires a `run_context`. Set an `event_stream_handler` on the agent and use `agent.run()` instead." diff --git a/python/restate/sessions/_calls.py b/python/restate/sessions/_calls.py index fd16f0fe..8a7f0768 100644 --- a/python/restate/sessions/_calls.py +++ b/python/restate/sessions/_calls.py @@ -100,8 +100,8 @@ class CallFuture(RestateDurableCallFuture[O]): def __init__(self, call: RestateDurableCallFuture[bytes], serde: Serde[O]) -> None: super().__init__() - self._call = call - self._serde = serde + self._call: RestateDurableCallFuture[bytes] = call + self._serde: Serde[O] = serde async def _result(self) -> O: return typing.cast(O, self._serde.deserialize(await self._call)) @@ -124,7 +124,7 @@ class Rpc: """ def __init__(self, host: CallHost) -> None: - self._host = host + self._host: CallHost = host @typing.overload def __call__( @@ -200,7 +200,7 @@ class Send: """ def __init__(self, sink: Callable[[GenericSend], None]) -> None: - self._sink = sink + self._sink: Callable[[GenericSend], None] = sink @typing.overload def __call__( diff --git a/python/restate/sessions/_database.py b/python/restate/sessions/_database.py index 4306c615..cbb9d8a0 100644 --- a/python/restate/sessions/_database.py +++ b/python/restate/sessions/_database.py @@ -22,12 +22,13 @@ from __future__ import annotations +import abc import asyncio import contextlib import contextvars import inspect import typing -from collections.abc import AsyncIterator, Awaitable, Callable, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence from contextlib import asynccontextmanager from typing import Any, Protocol @@ -92,24 +93,12 @@ class Store(Protocol): key's database, apart from the object's own (see ``_session.database``). """ - def get(self, key: str) -> bytes | None | Awaitable[bytes | None]: ... - def set(self, key: str, value: bytes) -> None: ... - def clear(self, key: str) -> None: ... + def get(self, key: str, /) -> bytes | None | Awaitable[bytes | None]: ... + def set(self, key: str, value: bytes, /) -> None: ... + def clear(self, key: str, /) -> None: ... def sync(self) -> None | Awaitable[None]: ... -class _ListsKeys(Protocol): - """A store that lists its keys (see ``Store``).""" - - def keys(self) -> Sequence[str] | Awaitable[Sequence[str]]: ... - - -class _Waits(Protocol): - """A store over the invocation's VM, which waits for a read (see ``Store``).""" - - def wait(self, handle: int) -> Awaitable[None]: ... - - async def open(namespace: str, store: Store, *, reader: bool = False) -> Connection: # pylint: disable=redefined-builtin """ Open the database of ``namespace`` over ``store``. One process-wide page @@ -153,9 +142,12 @@ class _Boundary: """ def __init__(self, store: Store, loop: asyncio.AbstractEventLoop) -> None: - self.store = store - self.loop = loop - self.lists_keys = callable(getattr(store, "keys", None)) + self.store: Store = store + self.loop: asyncio.AbstractEventLoop = loop + # Its optional methods (see `Store`). + self._keys: Callable[[], Sequence[str] | Awaitable[Sequence[str]]] | None = getattr(store, "keys", None) + self._wait: Callable[[int], Awaitable[None]] | None = getattr(store, "wait", None) + self.lists_keys: bool = callable(self._keys) # The first error of the store's. self.failed: BaseException | None = None # What the next sync runs first: the committing transaction's. @@ -186,9 +178,11 @@ def dispatch(self, op: int, key: Any, value: bytes | None, reply: SqlReply | Non if op == SQL_GET: result: object = self.store.get(key) elif op == SQL_KEYS: - result = typing.cast(_ListsKeys, self.store).keys() + assert self._keys is not None + result = self._keys() elif op == SQL_WAIT: - result = typing.cast(_Waits, self.store).wait(key) + assert self._wait is not None + result = self._wait(key) elif self._write_failed is not None: reply(None, str(self._write_failed)) return @@ -238,16 +232,16 @@ def objects_of(columns: list[str], rows: Rows) -> list[Row]: return [dict(zip(columns, row, strict=True)) for row in rows] -class _Sql: +class _Sql(abc.ABC): """SQL on a session of the extension: ``query``, ``get``, ``values`` and ``execute`` by way of ``query_columns``.""" def __init__(self, native: PySqlDatabase, boundary: _Boundary) -> None: - self._native = native - self._boundary = boundary + self._native: PySqlDatabase = native + self._boundary: _Boundary = boundary + @abc.abstractmethod async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: """The rows of one statement, with its column names: ``(columns, rows)``.""" - raise NotImplementedError async def query(self, sql: str, params: Sequence[SqlParam] | None = None) -> list[Row]: """ @@ -319,10 +313,10 @@ class Connection(_Sql): def __init__(self, native: PySqlDatabase, boundary: _Boundary) -> None: super().__init__(native, boundary) # FIFO: calls run in the order they were made. - self._lock = asyncio.Lock() + self._lock: asyncio.Lock = asyncio.Lock() @asynccontextmanager - async def _serial(self) -> AsyncIterator[None]: + async def _serial(self) -> AsyncGenerator[None, None]: """Run the body once the calls before it have, unless it comes from the body of a transaction it would wait for.""" running = _running.get() if running is not None and running.holds(self): @@ -335,7 +329,7 @@ async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None return await self._request(lambda: self._native.query_columns(sql, params)) @asynccontextmanager - async def transaction(self) -> AsyncIterator[ConnectionTransaction]: + async def transaction(self) -> AsyncGenerator[ConnectionTransaction, None]: """ One SQL transaction, as the body of an ``async with``:: @@ -417,14 +411,14 @@ def __init__( self, native: PySqlDatabase, boundary: _Boundary, connection: Connection, outer: ConnectionTransaction | None ) -> None: super().__init__(native, boundary) - self._connection = connection + self._connection: Connection = connection # The transaction whose body this one's runs in, if any: of another # connection, since a connection's own are refused. - self._outer = outer + self._outer: ConnectionTransaction | None = outer self._records: dict[str, bytes] = {} self._with_commits: list[Callable[[], None]] = [] - self._lock = asyncio.Lock() - self._ended = False + self._lock: asyncio.Lock = asyncio.Lock() + self._ended: bool = False @property def open(self) -> bool: @@ -492,7 +486,7 @@ def run_all() -> None: return run_all @asynccontextmanager - async def _serial(self) -> AsyncIterator[None]: + async def _serial(self) -> AsyncGenerator[None, None]: """Run the body, a call made while the transaction is open, once the calls before it have.""" self._refuse_ended() async with self._lock: diff --git a/python/restate/sessions/_link.py b/python/restate/sessions/_link.py index 6e024186..2affce0a 100644 --- a/python/restate/sessions/_link.py +++ b/python/restate/sessions/_link.py @@ -28,7 +28,7 @@ class Link(typing.Generic[T]): """A link of one kind, on the objects that have one: compared by identity, kept as long as the object is.""" def __init__(self, what: str) -> None: - self.what = what + self.what: str = what self._links: weakref.WeakKeyDictionary[object, T] = weakref.WeakKeyDictionary() def on(self, owner: O, value: T) -> O: diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index 36198702..16240f83 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -43,11 +43,11 @@ async def list(ctx: sessions.ViewContext) -> list[str]: from restate.retry_policy import InvocationRetryPolicy from restate.serde import BytesSerde, DefaultSerde, Serde from restate.server_context import DisconnectedException, ServerInvocationContext -from restate.sessions import _selector from restate.sessions._calls import CallFuture, Rpc, Send, send_now from restate.sessions._database import Store from restate.sessions._link import Link from restate.sessions._schema import Schema, Step +from restate.sessions._selector import FutureSelector, hook from restate.sessions._session import Attempt, Database, DatabaseReader, ObjectConfig, SqlHost, reader_of, with_attempt from restate.vm import SingleUnresolvedFuture @@ -126,12 +126,12 @@ class _ContextHost: """ def __init__(self, context: ServerInvocationContext, shared: bool) -> None: - self.context = context - self.shared = shared - self.service = context.handler.service_tag.name - self.key = context.key() + self.context: ServerInvocationContext = context + self.shared: bool = shared + self.service: str = context.handler.service_tag.name + self.key: str = context.key() self._listeners: list[Callable[[], None]] = [] - self.selector = _selector.hook(context, self._notify_attempt_end) + self.selector: FutureSelector = hook(context, self._notify_attempt_end) attempt_finished = context.on_attempt_finished def on_attempt_finished() -> None: @@ -389,11 +389,11 @@ def __init__( ingress_private=ingress_private, invocation_retry_policy=invocation_retry_policy, ) - self.config = ObjectConfig(schema=Schema(name, schema) if schema else None) + self.config: ObjectConfig = ObjectConfig(schema=Schema(name, schema) if schema else None) ENTRIES.on(self, {}) @typing.overload # type: ignore[override] - def handler(self, fn: F, /) -> F: ... # pyright: ignore[reportIncompatibleMethodOverride] + def handler(self, fn: F, /) -> F: ... @typing.overload def handler( @@ -542,10 +542,11 @@ def _decorate(self, view: bool, name: str | F | None, options: dict[str, Any]) - def _register(self, view: bool, name: str | None, options: dict[str, Any]) -> Callable[[F], F]: def register(fn: F) -> F: if RESTATE_UNIQUE_HANDLER_SYMBOL in vars(fn): - raise TypeError( - f"{self.name}.{fn.__name__} is an SDK handler: a session object's handlers are functions of its " - "context. Put journaled work in a restate.VirtualObject, Service or Workflow, and call it with ctx.rpc" + what = f"{self.name}.{fn.__name__} is an SDK handler: a session object's handlers are functions of its context." + instead = ( + "Put journaled work in a restate.VirtualObject, Service or Workflow, and call it with ctx.rpc." ) + raise TypeError(f"{what} {instead}") explicit = bool(options.get("explicit_cancellation")) sdk_options = {key: value for key, value in options.items() if key != "explicit_cancellation"} entry = _Entry(fn=fn, view=view, config=self.config, explicit_cancellation=explicit) diff --git a/python/restate/sessions/_schema.py b/python/restate/sessions/_schema.py index f24a2a38..d7790c30 100644 --- a/python/restate/sessions/_schema.py +++ b/python/restate/sessions/_schema.py @@ -59,7 +59,7 @@ class Schema: """A part of a database's schema: its ``name``, and the ``steps`` that make it, in order.""" def __init__(self, name: str, steps: Sequence[Step]) -> None: - self.name = name + self.name: str = name self.steps: tuple[Step, ...] = tuple(steps) async def ensure(self, db: Any) -> None: diff --git a/python/restate/sessions/_selector.py b/python/restate/sessions/_selector.py index 2c0fd458..5b1f7777 100644 --- a/python/restate/sessions/_selector.py +++ b/python/restate/sessions/_selector.py @@ -24,8 +24,7 @@ import asyncio import logging -import typing -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import Any, Protocol from restate.exceptions import SdkInternalException, SuspendedException, TerminalError @@ -59,7 +58,7 @@ async def flush_output(self) -> None: """Writes out what the VM has buffered.""" ... - async def receive(self) -> dict[str, Any]: + async def receive(self) -> Mapping[str, Any]: """The next input event: a chunk of the request body, a disconnect, or a restate.* event.""" ... @@ -71,7 +70,7 @@ def execute_run(self, handle: int) -> None: """Starts the run of ``handle``, which signals its completion through ``receive``.""" ... - def notify_input(self, chunk: dict[str, Any]) -> None: + def notify_input(self, chunk: Mapping[str, Any]) -> None: """Passes a chunk of the request body to the VM.""" ... @@ -106,12 +105,12 @@ class FutureSelector: """ def __init__(self, driver: SelectorDriver, disconnected: Callable[[], BaseException]) -> None: - self._driver = driver - self._disconnected = disconnected + self._driver: SelectorDriver = driver + self._disconnected: Callable[[], BaseException] = disconnected self._waits: dict[asyncio.Future[None], UnresolvedFuture] = {} self._task: asyncio.Task[None] | None = None # Between the VM's WaitExternalProgress and the next input. - self._waiting = False + self._waiting: bool = False # Why the attempt ended, once it has: each wait then raises a new one. self._end: Callable[[], BaseException] | None = None @@ -243,8 +242,8 @@ class _ContextDriver: """What the selector needs of the SDK's context.""" def __init__(self, context: ServerInvocationContext, on_end: Callable[[], None]) -> None: - self.context = context - self.on_end = on_end + self.context: ServerInvocationContext = context + self.on_end: Callable[[], None] = on_end def do_progress(self, future: UnresolvedFuture) -> Any: return self.context.vm.do_progress(future) @@ -255,8 +254,8 @@ def is_completed(self, handle: int) -> bool: async def flush_output(self) -> None: await self.context.take_and_send_output() - async def receive(self) -> dict[str, Any]: - return typing.cast(dict[str, Any], await self.context.receive()) + async def receive(self) -> Mapping[str, Any]: + return await self.context.receive() def wake(self) -> None: # Any restate.* event: the loop takes it as a wake-up. @@ -275,7 +274,7 @@ async def run() -> None: self.track(asyncio.create_task(run())) - def notify_input(self, chunk: dict[str, Any]) -> None: + def notify_input(self, chunk: Mapping[str, Any]) -> None: # Skip empty body frames, as the SDK does. body = chunk.get("body") if body is not None and len(body) > 0: @@ -313,7 +312,7 @@ def make_started(*args: Any, **kwargs: Any) -> Any: for name in ("create_future", "create_sleep_future", "create_call_future"): setattr(context, name, starting(getattr(context, name))) - context.create_poll_or_cancel_coroutine = selector.wait # type: ignore[method-assign,assignment] + context.create_poll_or_cancel_coroutine = selector.wait # type: ignore[method-assign,assignment] # pyright: ignore[reportAttributeAccessIssue] return selector diff --git a/python/restate/sessions/_session.py b/python/restate/sessions/_session.py index a039ba92..22af952a 100644 --- a/python/restate/sessions/_session.py +++ b/python/restate/sessions/_session.py @@ -17,6 +17,7 @@ from __future__ import annotations +import abc import asyncio import base64 import contextlib @@ -25,13 +26,14 @@ import time import typing import uuid -from collections.abc import AsyncIterator, Awaitable, Callable, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass, field from datetime import timedelta from typing import Any, Protocol from urllib.parse import quote +from restate._internal import PyVM # pylint: disable=no-name-in-module from restate.exceptions import SuspendedException, TerminalError from restate.sessions._calls import CallHost, GenericSend, Send from restate.sessions._database import ( @@ -99,6 +101,15 @@ def on_attempt_end(self, listener: Callable[[], None]) -> Callable[[], None]: """Call ``listener`` when the attempt ends; the returned function cancels it.""" ... + @property + def vm(self) -> PyVM | None: + """The invocation's VM, which the store writes and reads directly; none without one (a test's host).""" + ... + + def wait(self, handle: int) -> Awaitable[None]: + """Done once the VM has the read whose handle is ``handle``.""" + ... + @dataclass class ObjectConfig: @@ -157,11 +168,6 @@ def _unbeside(name: str, key: str) -> str | None: return None -def _vm_of(host: SqlHost) -> Any: - """The invocation's VM, if ``host`` gives it (see ``Store``): the SDK's context does.""" - return getattr(host, "vm", None) - - class _ObjectStore: """ The object's state as an ordered store: ``get`` is a state get, ``set`` and @@ -174,14 +180,14 @@ class _ObjectStore: """ def __init__(self, host: SqlHost, name: str | None) -> None: - self._host = host - self._name = name - self.vm = _vm_of(host) - self.database_name = name + self._host: SqlHost = host + self._name: str | None = name + self.vm: PyVM | None = host.vm + self.database_name: str | None = name def wait(self, handle: int) -> Awaitable[None]: """Done once the VM has the read whose handle is ``handle``.""" - return _until_attempt_end(self._host, self._host.wait(handle)) # type: ignore[attr-defined] + return _until_attempt_end(self._host, self._host.wait(handle)) def _key(self, key: str) -> str: return key if self._name is None else _beside(self._name, key) @@ -219,23 +225,23 @@ class _ReaderStore: """The object's state as a view's store: it reads beside the exclusive handler, and writes nothing.""" def __init__(self, host: SqlHost, name: str | None) -> None: - self._host = host - self._name = name - self.vm = _vm_of(host) - self.database_name = name + self._host: SqlHost = host + self._name: str | None = name + self.vm: PyVM | None = host.vm + self.database_name: str | None = name def wait(self, handle: int) -> Awaitable[None]: """Done once the VM has the read whose handle is ``handle``.""" - return _until_attempt_end(self._host, self._host.wait(handle)) # type: ignore[attr-defined] + return _until_attempt_end(self._host, self._host.wait(handle)) def get(self, key: str) -> Awaitable[bytes | None]: key = key if self._name is None else _beside(self._name, key) return _until_attempt_end(self._host, self._host.get(key)) - def set(self, key: str, value: bytes) -> None: + def set(self, _key: str, _value: bytes, /) -> None: raise RuntimeError("a view writes no state") - def clear(self, key: str) -> None: + def clear(self, _key: str, /) -> None: raise RuntimeError("a view writes no state") async def sync(self) -> None: @@ -307,8 +313,8 @@ class _Sends: def __init__(self) -> None: self._sends: list[dict[str, Any]] = [] - self._open = True - self.send = Send(self._record) + self._open: bool = True + self.send: Send = Send(self._record) def _record(self, send: GenericSend) -> None: if not self._open: @@ -348,12 +354,12 @@ def _refused_send(_: GenericSend) -> None: # --- The database, as a handler gets it ----------------------------------------------------------- -class DatabaseReader: +class DatabaseReader(abc.ABC): """Reads of a database, or of one transaction of it.""" + @abc.abstractmethod async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: """The rows of one statement as lists, with the column names of the statement: ``(columns, rows)``.""" - raise NotImplementedError async def query(self, sql: str, params: Sequence[SqlParam] | None = None) -> list[Row]: """ @@ -378,7 +384,7 @@ async def values(self, sql: str, params: Sequence[SqlParam] | None = None) -> Ro return (await self.query_columns(sql, params))[1] -class DatabaseWriter(DatabaseReader): +class DatabaseWriter(DatabaseReader, abc.ABC): """SQL: a database, or one transaction of it.""" async def execute(self, sql: str, params: Sequence[SqlParam] | None = None) -> None: @@ -397,8 +403,8 @@ class Transaction(DatabaseWriter): """ def __init__(self, handle: ConnectionTransaction, send: Send) -> None: - self._handle = handle - self.send = send + self._handle: ConnectionTransaction = handle + self.send: Send = send async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: try: @@ -415,7 +421,7 @@ class Database(DatabaseWriter): """ def __init__(self, attempt: Attempt) -> None: - self._attempt = attempt + self._attempt: Attempt = attempt _ATTEMPT.on(self, attempt) async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: @@ -442,7 +448,7 @@ class _ReaderOf(DatabaseReader): """The reads of a ``Database``: a view's ``ctx.db``.""" def __init__(self, db: Database) -> None: - self._db = db + self._db: Database = db async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: return await self._db.query_columns(sql, params) @@ -513,9 +519,9 @@ class _Session: """ def __init__(self, ready: asyncio.Task[Connection], host: SqlHost) -> None: - self.ready = ready + self.ready: asyncio.Task[Connection] = ready # A call failed on it once a failure ended it: the next call opens another. - self.retired = False + self.retired: bool = False self._release: Callable[[], None] = lambda: None ready.add_done_callback(lambda task: self._track(task, host)) @@ -573,11 +579,11 @@ class Attempt: """ def __init__(self, host: SqlHost, config: ObjectConfig | None = None, name: str | None = None) -> None: - self._host = host - self._config = config or ObjectConfig() - self._name = name + self._host: SqlHost = host + self._config: ObjectConfig = config or ObjectConfig() + self._name: str | None = name # The object's own database of an exclusive handler: the one with an outbox. - self._has_outbox = not host.shared and name is None + self._has_outbox: bool = not host.shared and name is None # The key's other databases this attempt opened, by name. self._databases: dict[str, Attempt] = {} self._session: _Session | None = None @@ -585,8 +591,8 @@ def __init__(self, host: SqlHost, config: ObjectConfig | None = None, name: str self._made: set[str] = set() # The schemas this attempt's database has all the steps of, by name. self._ensured: dict[str, Schema] = {} - self._ended = False - self.db = Database(self) + self._ended: bool = False + self.db: Database = Database(self) """``ctx.db``.""" ensuring(self.db, self._ensure) @@ -643,7 +649,7 @@ async def call(self, call: Callable[[Connection], Awaitable[T]]) -> T: _raise_for_handler(error) @asynccontextmanager - async def transaction(self) -> AsyncIterator[Transaction]: + async def transaction(self) -> AsyncGenerator[Transaction, None]: """ One transaction. In the object's own database of an exclusive handler, its sends commit with it, in the outbox, and are made with its commit, diff --git a/tests/entry_codec.py b/tests/entry_codec.py index 5c25d3bf..8ff9c553 100644 --- a/tests/entry_codec.py +++ b/tests/entry_codec.py @@ -169,7 +169,7 @@ async def test_client_send_skips_response_decode(): @asynccontextmanager async def codec_harness( service: typing.Union[Service, restate.VirtualObject, restate.Workflow], codec: JournalValueCodec -) -> typing.AsyncIterator[restate.RestateClient]: +) -> typing.AsyncGenerator[restate.RestateClient, None]: """Spin up a harness where BOTH the endpoint and the ingress client share the same codec.""" async with restate.create_test_harness( restate.app([service], journal_value_codec=codec), diff --git a/tests/servercontext.py b/tests/servercontext.py index 4a30de4a..1a3b92a2 100644 --- a/tests/servercontext.py +++ b/tests/servercontext.py @@ -46,7 +46,9 @@ def ohoh(): @asynccontextmanager -async def simple_harness(service: Service | VirtualObject | Workflow) -> typing.AsyncIterator[restate.RestateClient]: +async def simple_harness( + service: Service | VirtualObject | Workflow, +) -> typing.AsyncGenerator[restate.RestateClient, None]: async with restate.create_test_harness( restate.app([service]), restate_image="ghcr.io/restatedev/restate:latest" ) as restate_test_harness: diff --git a/tests/sessions_object.py b/tests/sessions_object.py index 50557006..bf07cd4b 100644 --- a/tests/sessions_object.py +++ b/tests/sessions_object.py @@ -131,6 +131,12 @@ def generic_send(self, service, handler, arg, key=None, send_delay=None, idempot def cancel_invocation(self, invocation_id: str) -> None: self.cancelled.append(invocation_id) + # No VM: the database's store goes through get, set and clear. + vm = None + + def wait(self, handle: int): + raise AssertionError("a host without a VM waits for no read") + def attempt_end(self): return None diff --git a/tests/sessions_selector.py b/tests/sessions_selector.py index 33c7254f..8fd95639 100644 --- a/tests/sessions_selector.py +++ b/tests/sessions_selector.py @@ -13,6 +13,7 @@ # pylint: disable=missing-function-docstring import asyncio +from collections.abc import Mapping from typing import Any import pytest @@ -86,7 +87,7 @@ def execute_run(self, handle: int) -> None: self.completed.add(handle) self.input.put_nowait({"type": "restate.run_completed"}) - def notify_input(self, chunk: dict[str, Any]) -> None: + def notify_input(self, chunk: Mapping[str, Any]) -> None: self.completed.add(int(chunk["body"])) def track(self, task: asyncio.Future) -> None: From e8d793978e2db53200df983320dfcd7fead8bdf7 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 15:57:08 +0100 Subject: [PATCH 16/45] Separate restate.sessions' API from its implementation, as protocols What handlers and views code against is declared in _api.py, as protocols: SessionContext, ViewContext, Database, DatabaseReader, DatabaseWriter, Transaction, Rpc and Send (and Cancel). restate.sessions exports them under the names it had; the implementations are private classes of the modules beside it (AttemptContext, AttemptDatabase, AttemptTransaction, Caller, Sender, ...), which the strict type checkers check against the protocols wherever one is given for the other. What a handler sees is what is declared, and a test can give a handler a fake of its own: @obj.handler and @obj.view now return the function itself, marked as the handler it is (ctx.rpc and the client find it), so a test calls it with a fake context. SessionObject stays a class, as the SDK's VirtualObject; CallFuture, Schema and SqlError stay values. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/sessions/__init__.py | 25 ++- python/restate/sessions/_api.py | 328 ++++++++++++++++++++++++++++ python/restate/sessions/_calls.py | 18 +- python/restate/sessions/_object.py | 62 ++---- python/restate/sessions/_session.py | 61 ++---- tests/sessions_object.py | 62 +++++- 6 files changed, 446 insertions(+), 110 deletions(-) create mode 100644 python/restate/sessions/_api.py diff --git a/python/restate/sessions/__init__.py b/python/restate/sessions/__init__.py index 138ea327..ec1c431a 100644 --- a/python/restate/sessions/__init__.py +++ b/python/restate/sessions/__init__.py @@ -31,23 +31,32 @@ async def list(ctx: sessions.ViewContext) -> list[str]: app = restate.app([notes]) Session handlers don't replay a journal: a retried attempt runs again from its -first line, against what the attempts before it committed. A session object +first line, against what the attempts before it committed. + +The types a handler gets (``SessionContext``, ``Database``, ``Transaction``, +``Rpc``, ``Send``, ...) are protocols: a decorated handler is still the plain +function, so a test can call it with a fake context of its own. A session object serves beside the SDK's services, objects and workflows, which it calls with ``ctx.rpc`` for journaled work. """ -from restate.sessions._calls import CallFuture, Rpc, Send -from restate.sessions._database import Row, Rows, SqlError, SqlParam, SqlValue -from restate.sessions._object import Cancel, SessionContext, SessionObject, ViewContext, cancel_calls, cancelled -from restate.sessions._schema import Schema, Step -from restate.sessions._session import ( +# The API: protocols (_api.py), which handlers and views code against. +from restate.sessions._api import ( + Cancel, Database, DatabaseReader, DatabaseWriter, + Rpc, + Send, + SessionContext, Transaction, - database, - reset, + ViewContext, ) +from restate.sessions._calls import CallFuture +from restate.sessions._database import Row, Rows, SqlError, SqlParam, SqlValue +from restate.sessions._object import SessionObject, cancel_calls, cancelled +from restate.sessions._schema import Schema, Step +from restate.sessions._session import database, reset __all__ = [ "CallFuture", diff --git a/python/restate/sessions/_api.py b/python/restate/sessions/_api.py new file mode 100644 index 00000000..0ed93375 --- /dev/null +++ b/python/restate/sessions/_api.py @@ -0,0 +1,328 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +The API of session objects, as handlers and views code against it: the +contexts, the database and its transactions, calls and sends. Protocols, so +that what a handler sees is what is declared here, and a test can give a +handler a fake of its own; the implementations are this package's private +classes (``_object.py``, ``_session.py``, ``_calls.py``). +""" + +from __future__ import annotations + +import typing +from collections.abc import Awaitable, Callable, Sequence +from datetime import timedelta +from typing import Any, Protocol + +from restate.context import Request +from restate.serde import JsonSerde, Serde +from restate.sessions._calls import CallFuture +from restate.sessions._database import Row, Rows, SqlParam + +I = typing.TypeVar("I") +O = typing.TypeVar("O") + +Cancel = Callable[[str], None] +""" +Cancel invocation ``invocation_id`` (such as ``inv_1f…``), as Restate's admin +API does: made at once, and not awaited. Cancelling an invocation that already +ended does nothing, so a new attempt that cancels it again is harmless. +""" + + +# --- The database --------------------------------------------------------------------------------- + + +class DatabaseReader(Protocol): + """Reads of a database, or of one transaction of it: a view's ``ctx.db``.""" + + async def query(self, sql: str, params: Sequence[SqlParam] | None = None) -> list[Row]: + """ + The rows of one statement, each a dict by column name; parameters bind + to a single statement. Columns are named as SQL names them (an + expression needs ``AS``), and two of the same name are an error: alias + them, or use ``values``. + """ + ... + + async def get(self, sql: str, params: Sequence[SqlParam] | None = None) -> Row | None: + """The first row of ``query``, or ``None`` when there is none.""" + ... + + async def values(self, sql: str, params: Sequence[SqlParam] | None = None) -> Rows: + """The rows of one statement as lists, in column order: for adapters, and for columns that share a name.""" + ... + + async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: + """The rows of one statement as lists, with the column names of the statement: ``(columns, rows)``.""" + ... + + +class DatabaseWriter(DatabaseReader, Protocol): + """SQL: a database, or one transaction of it.""" + + async def execute(self, sql: str, params: Sequence[SqlParam] | None = None) -> None: + """Run SQL (one or more statements), discarding rows; parameters only with a single statement.""" + ... + + +class Transaction(DatabaseWriter, Protocol): + """One transaction of a ``Database``: its SQL, and the sends it commits with.""" + + @property + def send(self) -> Send: + """ + Sends made right after the transaction commits, and never if it rolls + back. Each has the idempotency key it is given, or else one of the + transaction (``transaction//``), so a send an attempt makes + again is dropped. + """ + ... + + +class Database(DatabaseWriter, Protocol): + """ + The session's database, ``ctx.db``, for the whole handler: each statement + outside ``transaction()`` is its own transaction. Calls run one at a time, + in call order. A commit is durable once its call returns. + """ + + def transaction(self) -> typing.AsyncContextManager[Transaction]: + """ + One transaction, as the body of an ``async with``:: + + async with ctx.db.transaction() as tx: + await tx.execute("UPDATE accounts SET balance = balance - ? WHERE id = ?", [amount, src]) + tx.send(deposit, key=dst)(amount) + + It commits when the body ends, and rolls back when the body raises. It + holds the database: other calls wait until it ends, and a call to the + database from inside the body is refused, as it would wait for the body. + So the body uses the ``tx`` it gets, which is invalid once the body has + ended. + """ + ... + + +# --- Calls and sends ------------------------------------------------------------------------------ + + +class Rpc(Protocol): + """ + Calls: ``ctx.rpc(handler, key=...)`` returns a function of the handler's + input, if it takes one, which makes the call and returns its result's + awaitable. A call is made when that function is called, awaited or not:: + + count = await ctx.rpc(increment, key="a", idempotency_key=f"{ctx.request.id}:inc")(5) + """ + + @typing.overload + def __call__( + self, + handler: Callable[[Any], Awaitable[O]], + key: str | None = None, + *, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> Callable[[], CallFuture[O]]: ... + + @typing.overload + def __call__( + self, + handler: Callable[[Any, I], Awaitable[O]], + key: str | None = None, + *, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> Callable[[I], CallFuture[O]]: ... + + def __call__( + self, + handler: Callable[..., Awaitable[Any]], + key: str | None = None, + *, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> Callable[..., CallFuture[Any]]: + """ + A caller of ``handler``, with ``key`` for an object's or a workflow's + (a ``TerminalError`` otherwise, as a retry would fail the same way). + + ``idempotency_key`` is the same on every attempt, so that a new attempt + gets the first call's result; without one, every attempt makes the call + again, the one that resumes a suspension too. Restate scopes it to the + target handler. + """ + ... + + def generic( + self, + service: str, + handler: str, + arg: Any = None, + *, + key: str | None = None, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + input_serde: Serde[Any] = JsonSerde(), + output_serde: Serde[Any] = JsonSerde(), + ) -> CallFuture[Any]: + """ + Call a target known only at run time, by its names: untyped, and JSON + unless serdes are given. ``key`` is the caller's to give, for an + object or a workflow. + """ + ... + + +class Send(Protocol): + """ + Sends, as ``Rpc``'s calls: ``ctx.send(handler, key=...)`` returns a function + of the handler's input, if it takes one, which makes the send at once + (``ctx.send``), or records it for its transaction's commit (``tx.send``):: + + ctx.send(record, key="audit", delay=timedelta(hours=1))(entry) + """ + + @typing.overload + def __call__( + self, + handler: Callable[[Any], Awaitable[Any]], + key: str | None = None, + *, + delay: timedelta | None = None, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> Callable[[], None]: ... + + @typing.overload + def __call__( + self, + handler: Callable[[Any, I], Awaitable[Any]], + key: str | None = None, + *, + delay: timedelta | None = None, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> Callable[[I], None]: ... + + def __call__( + self, + handler: Callable[..., Awaitable[Any]], + key: str | None = None, + *, + delay: timedelta | None = None, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + ) -> Callable[..., None]: + """ + A sender to ``handler``, with ``key`` for an object's or a workflow's; + ``delay`` is how long after the send (for ``tx.send``, after the + commit) the target runs. See ``Rpc`` for ``idempotency_key``. + """ + ... + + def generic( + self, + service: str, + handler: str, + arg: Any = None, + *, + key: str | None = None, + delay: timedelta | None = None, + idempotency_key: str | None = None, + headers: dict[str, str] | None = None, + input_serde: Serde[Any] = JsonSerde(), + ) -> None: + """Send to a target known only at run time, by its names: untyped, and JSON unless a serde is given.""" + ... + + +# --- The contexts --------------------------------------------------------------------------------- + + +class SessionContext(Protocol): + """ + The context of a session object's handler. A handler runs in attempts: + Restate starts a new one after a crash, a retryable error or a suspension, + and each runs the handler from its first line, finding what the attempts + before it committed. Its operations can be awaited together + (``asyncio.gather``, ``asyncio.wait``). + """ + + @property + def key(self) -> str: + """The session's key.""" + ... + + @property + def request(self) -> Request: + """The request: its id (the same on every attempt), headers, and the attempt's.""" + ... + + @property + def db(self) -> Database: + """ + The key's database. Each statement outside a transaction is its own + commit, and ``db.transaction()`` groups several; a commit is durable + when its call returns. Calls run one at a time, in call order. + """ + ... + + @property + def rpc(self) -> Rpc: + """Calls, typed by the handler called: ``await ctx.rpc(handler, key="a")(input)``.""" + ... + + @property + def send(self) -> Send: + """Sends, made at once; a transaction's (``tx.send``) are made once it commits.""" + ... + + @property + def cancel(self) -> Cancel: + """Cancel an invocation, by its id: see ``Cancel``.""" + ... + + +class ViewContext(Protocol): + """ + The context of a session object's view: it runs beside the handlers, and + reads the database as committed when its invocation started, whatever they + commit meanwhile. It writes nothing, and sends nothing, but it can cancel an + invocation: a handler of its own object's, say, as a Stop button does. + """ + + @property + def key(self) -> str: + """The session's key.""" + ... + + @property + def request(self) -> Request: + """The request: its id, headers, and the attempt's.""" + ... + + @property + def db(self) -> DatabaseReader: + """The key's database as committed when the view's invocation started: reads only.""" + ... + + @property + def rpc(self) -> Rpc: + """Calls, typed by the handler called.""" + ... + + @property + def cancel(self) -> Cancel: + """Cancel an invocation, by its id: see ``Cancel``.""" + ... diff --git a/python/restate/sessions/_calls.py b/python/restate/sessions/_calls.py index 8a7f0768..1b855ed2 100644 --- a/python/restate/sessions/_calls.py +++ b/python/restate/sessions/_calls.py @@ -116,12 +116,8 @@ async def cancel_invocation(self) -> None: await self._call.cancel_invocation() -class Rpc: - """ - Calls: ``ctx.rpc(handler, key=...)`` returns a function of the handler's - input, if it takes one, which makes the call and returns its result's - awaitable. A call is made when that function is called, awaited or not. - """ +class Caller: + """``Rpc`` (see ``_api.py``) over the SDK's context: each call a generic call, its result read by its serde.""" def __init__(self, host: CallHost) -> None: self._host: CallHost = host @@ -192,12 +188,8 @@ def generic( return CallFuture(future, output_serde) -class Send: - """ - Sends, as ``Rpc``'s calls: ``ctx.send(handler, key=...)`` returns a function - of the handler's input, if it takes one, which makes the send at once - (``ctx.send``), or records it for its transaction's commit (``tx.send``). - """ +class Sender: + """``Send`` (see ``_api.py``) into a sink: ``send_now``'s for ``ctx.send``, a transaction's outbox for ``tx.send``.""" def __init__(self, sink: Callable[[GenericSend], None]) -> None: self._sink: Callable[[GenericSend], None] = sink @@ -236,7 +228,7 @@ def __call__( """ A sender to ``handler``, with ``key`` for an object's or a workflow's; ``delay`` is how long after the send (for ``tx.send``, after the - commit) the target runs. See ``Rpc`` for ``idempotency_key``. + commit) the target runs. See ``Caller`` for ``idempotency_key``. """ service, name, input_serde, _ = _target(handler, key) diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index 16240f83..b19c8bb9 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -43,12 +43,13 @@ async def list(ctx: sessions.ViewContext) -> list[str]: from restate.retry_policy import InvocationRetryPolicy from restate.serde import BytesSerde, DefaultSerde, Serde from restate.server_context import DisconnectedException, ServerInvocationContext -from restate.sessions._calls import CallFuture, Rpc, Send, send_now +from restate.sessions._api import Cancel, Database, DatabaseReader, Rpc, Send, SessionContext, ViewContext +from restate.sessions._calls import Caller, Sender, send_now from restate.sessions._database import Store from restate.sessions._link import Link from restate.sessions._schema import Schema, Step from restate.sessions._selector import FutureSelector, hook -from restate.sessions._session import Attempt, Database, DatabaseReader, ObjectConfig, SqlHost, reader_of, with_attempt +from restate.sessions._session import Attempt, ObjectConfig, SqlHost, reader_of, with_attempt from restate.vm import SingleUnresolvedFuture F = typing.TypeVar("F", bound=Callable[..., Awaitable[Any]]) @@ -66,50 +67,22 @@ async def list(ctx: sessions.ViewContext) -> list[str]: ``v22/...``; and the ``sync`` run after its writes). """ -Cancel = Callable[[str], None] -""" -Cancel invocation ``invocation_id`` (such as ``inv_1f…``), as Restate's admin -API does: made at once, and not awaited. Cancelling an invocation that already -ended does nothing, so a new attempt that cancels it again is harmless. -""" - @dataclass(frozen=True, eq=False) -class SessionContext: - """ - The context of a session object's handler. A handler runs in attempts: - Restate starts a new one after a crash, a retryable error or a suspension, - and each runs the handler from its first line, finding what the attempts - before it committed. Its operations can be awaited together - (``asyncio.gather``, ``asyncio.wait``). - """ +class AttemptContext: + """``SessionContext`` (see ``_api.py``): one attempt's, made as it starts. Compared by identity.""" key: str - """The session's key.""" request: Request - """The request: its id (the same on every attempt), headers, and the attempt's.""" db: Database - """ - The key's database. Each statement outside a transaction is its own commit, - and ``db.transaction()`` groups several; a commit is durable when its call - returns. Calls run one at a time, in call order. - """ rpc: Rpc - """Calls, typed by the handler called: ``await ctx.rpc(handler, key="a")(input)``.""" send: Send - """Sends, made at once; a transaction's (``tx.send``) are made once it commits.""" cancel: Cancel - """Cancel an invocation, by its id: see ``Cancel``.""" @dataclass(frozen=True, eq=False) -class ViewContext: - """ - The context of a session object's view: it runs beside the handlers, and - reads the database as committed when its invocation started, whatever they - commit meanwhile. It writes nothing, and sends nothing, but it can cancel an - invocation: a handler of its own object's, say, as a Stop button does. - """ +class AttemptViewContext: + """``ViewContext`` (see ``_api.py``): one attempt's, made as it starts. Compared by identity.""" key: str request: Request @@ -264,12 +237,12 @@ class _Entry: def _context_of(entry: _Entry, host: SqlHost, request: Request, attempt: Attempt) -> SessionContext | ViewContext: """The context of an attempt of ``entry``.""" - rpc = Rpc(host) + rpc = Caller(host) cancel = _cancel_of(host) if entry.view: - return ViewContext(key=host.key, request=request, db=reader_of(attempt.db), rpc=rpc, cancel=cancel) - context = SessionContext( - key=host.key, request=request, db=attempt.db, rpc=rpc, send=Send(send_now(host)), cancel=cancel + return AttemptViewContext(key=host.key, request=request, db=reader_of(attempt.db), rpc=rpc, cancel=cancel) + context = AttemptContext( + key=host.key, request=request, db=attempt.db, rpc=rpc, send=Sender(send_now(host)), cancel=cancel ) sdk = typing.cast("ServerInvocationContext | None", getattr(host, "context", None)) return SCOPE.on(context, Scope(sdk=sdk, explicit_cancellation=entry.explicit_cancellation)) @@ -542,9 +515,10 @@ def _decorate(self, view: bool, name: str | F | None, options: dict[str, Any]) - def _register(self, view: bool, name: str | None, options: dict[str, Any]) -> Callable[[F], F]: def register(fn: F) -> F: if RESTATE_UNIQUE_HANDLER_SYMBOL in vars(fn): - what = f"{self.name}.{fn.__name__} is an SDK handler: a session object's handlers are functions of its context." + what = f"{self.name}.{fn.__name__} is a handler already, of the SDK's or of a session object's." instead = ( - "Put journaled work in a restate.VirtualObject, Service or Workflow, and call it with ctx.rpc." + "A session object's handlers are functions of its context, of their own; put journaled work " + "in a restate.VirtualObject, Service or Workflow, and call it with ctx.rpc." ) raise TypeError(f"{what} {instead}") explicit = bool(options.get("explicit_cancellation")) @@ -561,7 +535,11 @@ def register(fn: F) -> F: entries = ENTRIES.of(self) assert entries is not None entries[sdk.name] = entry - return typing.cast(F, registered) + # The function itself, marked as the handler it is: `ctx.rpc(fn, ...)` + # and the client find it, and a test calls it with a context of + # its own (a fake of the protocols). + vars(fn)[RESTATE_UNIQUE_HANDLER_SYMBOL] = sdk + return fn return register @@ -571,4 +549,4 @@ def with_store_wrapper(self, wrap: Callable[[Store], Store]) -> SessionObject: return self -__all__ = ["CallFuture", "Cancel", "SessionContext", "SessionObject", "ViewContext", "cancel_calls", "cancelled"] +__all__ = ["SessionObject", "cancel_calls", "cancelled"] diff --git a/python/restate/sessions/_session.py b/python/restate/sessions/_session.py index 22af952a..3b0c3c5b 100644 --- a/python/restate/sessions/_session.py +++ b/python/restate/sessions/_session.py @@ -35,7 +35,8 @@ from restate._internal import PyVM # pylint: disable=no-name-in-module from restate.exceptions import SuspendedException, TerminalError -from restate.sessions._calls import CallHost, GenericSend, Send +from restate.sessions._api import Database, DatabaseReader, Send, Transaction +from restate.sessions._calls import CallHost, GenericSend, Sender from restate.sessions._database import ( Connection, ConnectionTransaction, @@ -314,7 +315,7 @@ class _Sends: def __init__(self) -> None: self._sends: list[dict[str, Any]] = [] self._open: bool = True - self.send: Send = Send(self._record) + self.send: Send = Sender(self._record) def _record(self, send: GenericSend) -> None: if not self._open: @@ -354,8 +355,8 @@ def _refused_send(_: GenericSend) -> None: # --- The database, as a handler gets it ----------------------------------------------------------- -class DatabaseReader(abc.ABC): - """Reads of a database, or of one transaction of it.""" +class Reads(abc.ABC): + """``DatabaseReader`` (see ``_api.py``) by way of ``query_columns``: an SQL error that escapes is terminal.""" @abc.abstractmethod async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: @@ -384,23 +385,16 @@ async def values(self, sql: str, params: Sequence[SqlParam] | None = None) -> Ro return (await self.query_columns(sql, params))[1] -class DatabaseWriter(DatabaseReader, abc.ABC): - """SQL: a database, or one transaction of it.""" +class Writes(Reads, abc.ABC): + """``DatabaseWriter`` (see ``_api.py``) by way of ``query_columns``.""" async def execute(self, sql: str, params: Sequence[SqlParam] | None = None) -> None: """Run SQL (one or more statements), discarding rows; parameters only with a single statement.""" await self.query_columns(sql, params) -class Transaction(DatabaseWriter): - """ - One transaction of a ``Database``: its SQL, and the sends it commits with. - - ``send`` makes its sends right after the transaction commits, and never if - it rolls back. Each has the idempotency key it is given, or else one of the - transaction (``transaction//``), so a send an attempt makes - again is dropped. - """ +class AttemptTransaction(Writes): + """``Transaction`` (see ``_api.py``): the SQL of a connection's transaction, and the sends it commits with.""" def __init__(self, handle: ConnectionTransaction, send: Send) -> None: self._handle: ConnectionTransaction = handle @@ -413,12 +407,8 @@ async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None _raise_for_handler(error) -class Database(DatabaseWriter): - """ - The session's database, ``ctx.db``, for the whole handler: each statement - outside ``transaction()`` is its own transaction. Calls run one at a time, - in call order. A commit is durable once its call returns. - """ +class AttemptDatabase(Writes): + """``Database`` (see ``_api.py``): an attempt's session, opened on its first use.""" def __init__(self, attempt: Attempt) -> None: self._attempt: Attempt = attempt @@ -428,33 +418,20 @@ async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None return await self._attempt.call(lambda db: db.query_columns(sql, params)) def transaction(self) -> typing.AsyncContextManager[Transaction]: - """ - One transaction, as the body of an ``async with``:: - - async with ctx.db.transaction() as tx: - await tx.execute("UPDATE accounts SET balance = balance - ? WHERE id = ?", [amount, src]) - tx.send(deposit, key=dst)(amount) - - It commits when the body ends, and rolls back when the body raises. It - holds the database: other calls wait until it ends, and a call to the - database from inside the body is refused, as it would wait for the body. - So the body uses the ``tx`` it gets, which is invalid once the body has - ended. - """ return self._attempt.transaction() -class _ReaderOf(DatabaseReader): - """The reads of a ``Database``: a view's ``ctx.db``.""" +class _ReaderOf(Reads): + """The reads of a database: a view's ``ctx.db``.""" - def __init__(self, db: Database) -> None: - self._db: Database = db + def __init__(self, db: AttemptDatabase) -> None: + self._db: AttemptDatabase = db async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: return await self._db.query_columns(sql, params) -def reader_of(db: Database) -> DatabaseReader: +def reader_of(db: AttemptDatabase) -> DatabaseReader: """The reads of ``db``, which apply a library's schema as ``db`` does (``Schema.ensure``).""" reader = _ReaderOf(db) ensure = ensure_of(db) @@ -592,7 +569,7 @@ def __init__(self, host: SqlHost, config: ObjectConfig | None = None, name: str # The schemas this attempt's database has all the steps of, by name. self._ensured: dict[str, Schema] = {} self._ended: bool = False - self.db: Database = Database(self) + self.db: AttemptDatabase = AttemptDatabase(self) """``ctx.db``.""" ensuring(self.db, self._ensure) @@ -649,7 +626,7 @@ async def call(self, call: Callable[[Connection], Awaitable[T]]) -> T: _raise_for_handler(error) @asynccontextmanager - async def transaction(self) -> AsyncGenerator[Transaction, None]: + async def transaction(self) -> AsyncGenerator[AttemptTransaction, None]: """ One transaction. In the object's own database of an exclusive handler, its sends commit with it, in the outbox, and are made with its commit, @@ -668,7 +645,7 @@ async def transaction(self) -> AsyncGenerator[Transaction, None]: async with db.transaction() as handle: sends = _Sends() if self._has_outbox else None # A view's writes stay in its session: it has nothing to reset. - tx = Transaction(handle, sends.send if sends is not None else Send(_refused_send)) + tx = AttemptTransaction(handle, sends.send if sends is not None else Sender(_refused_send)) if not self._host.shared: _RESET.on(tx, lambda: self._reset(handle)) try: diff --git a/tests/sessions_object.py b/tests/sessions_object.py index bf07cd4b..35021947 100644 --- a/tests/sessions_object.py +++ b/tests/sessions_object.py @@ -24,6 +24,7 @@ from restate import sessions from restate.context import AttemptFinishedEvent, Request from restate.exceptions import TerminalError +from restate.handler import handler_from_callable from restate.sessions._object import invoke from sessions_database import Backend # pylint: disable=import-error @@ -396,8 +397,11 @@ async def journaled(ctx: restate.ObjectContext) -> None: pass obj = sessions.SessionObject("refusing") - with pytest.raises(TypeError, match="is an SDK handler"): + with pytest.raises(TypeError, match="is a handler already"): obj.handler()(journaled) + # A session object's handler too: it is registered once. + with pytest.raises(TypeError, match="is a handler already"): + obj.handler()(add) def test_a_session_object_runs_its_handlers_with_explicit_recovery_and_lazy_state(): @@ -466,19 +470,19 @@ async def explicit(ctx: sessions.SessionContext) -> None: async def test_cancelled_ends_its_waiter_cancelled_when_the_attempt_ends_first(monkeypatch): from restate import _preview # pylint: disable=import-outside-toplevel from restate.exceptions import SuspendedException # pylint: disable=import-outside-toplevel - from restate.sessions import _object # pylint: disable=import-outside-toplevel + from restate.sessions import _calls, _object # pylint: disable=import-outside-toplevel async def attempt_ended(_sdk: Any) -> None: raise SuspendedException() monkeypatch.setattr(_preview, "cancellation", attempt_ended) host = Host(Backend()) - ctx = sessions.SessionContext( + ctx = _object.AttemptContext( key="k", request=request(), db=typing.cast(Any, None), - rpc=sessions.Rpc(host), - send=sessions.Send(lambda send: None), + rpc=_calls.Caller(host), + send=_calls.Sender(lambda send: None), cancel=lambda invocation_id: None, ) _object.SCOPE.on(ctx, _object.Scope(sdk=typing.cast(Any, object()), explicit_cancellation=True)) @@ -486,3 +490,51 @@ async def attempt_ended(_sdk: Any) -> None: await asyncio.wait({watcher}) # Cancelled, not failed: asyncio reports no exception nobody retrieved. assert watcher.cancelled() + + +class FakeDatabase: + """A test's own database, as ``sessions.Database`` declares one: it records the SQL it is given.""" + + def __init__(self) -> None: + self.statements: list[tuple[str, Any]] = [] + + async def query(self, sql: str, params: Any = None) -> list[sessions.Row]: + self.statements.append((sql, params)) + return [] + + async def get(self, sql: str, params: Any = None) -> sessions.Row | None: + return None + + async def values(self, sql: str, params: Any = None) -> sessions.Rows: + return [] + + async def query_columns(self, sql: str, params: Any = None) -> tuple[list[str], sessions.Rows]: + return [], [] + + async def execute(self, sql: str, params: Any = None) -> None: + self.statements.append((sql, params)) + + def transaction(self) -> typing.AsyncContextManager[sessions.Transaction]: + raise NotImplementedError + + +@dataclass +class FakeContext: + """A test's own handler context, as ``sessions.SessionContext`` declares one.""" + + db: FakeDatabase + key: str = "alice" + request: Request = field(default_factory=request) + rpc: sessions.Rpc = field(default_factory=lambda: typing.cast(sessions.Rpc, None)) + send: sessions.Send = field(default_factory=lambda: typing.cast(sessions.Send, None)) + cancel: sessions.Cancel = lambda invocation_id: None + + +async def test_a_handler_runs_over_a_fake_context_of_a_test(): + fake = FakeContext(db=FakeDatabase()) + # The protocols are structural: the fake is a SessionContext to the type checkers. + ctx: sessions.SessionContext = fake + await add(ctx, "milk") + assert fake.db.statements == [("INSERT INTO notes (text) VALUES (?)", ["milk"])] + # And the function is still the handler that calls and the client name. + assert handler_from_callable(add).name == "add" From c89c303917448d41334efe473a15dfe2182538bf Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 16:07:00 +0100 Subject: [PATCH 17/45] Add attempt scopes, served dependencies, and restate.sessions.internal What restate.agents needs of sessions, as sdk-sessions has it: - ScopeKey (_scope.py, as scope.ts): what a library keeps for an attempt (an agent's turn), by its context; a retry starts without it. Views have a scope too, and definition_of(ctx) gives the session object of a handler or a view, where a library keeps what it keeps per object. - needing(definition, services) and served(...) (_serve.py, as serve.ts): the services a definition needs served beside it (an agent's signals). sessions.app(*services) is restate.app's with them, each once; RESTATE_SERVE_ONLY serves only those it names. The SDK has no serve: the program serves the app (hypercorn, uvicorn). - restate.sessions.internal (as @restatedev/sessions/internal): Schema, database, reset, cancelled, cancel_calls, ScopeKey, definition_of, needing, served and invoke, which libraries and tests use; the package exports what handlers code against, and app. Also a fix, found in this review: a Link holds its value, and an attempt holds its database, so the link from a database to its attempt kept every attempt alive for good. Link(weak=True) holds the value weakly, for the owner to keep it; a test checks an attempt is freed once it ends. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/sessions/__init__.py | 18 +++-- python/restate/sessions/_link.py | 21 ++++-- python/restate/sessions/_object.py | 50 +++++++------ python/restate/sessions/_scope.py | 87 ++++++++++++++++++++++ python/restate/sessions/_serve.py | 93 ++++++++++++++++++++++++ python/restate/sessions/_session.py | 4 +- python/restate/sessions/internal.py | 48 +++++++++++++ tests/sessions_harness.py | 7 +- tests/sessions_object.py | 108 ++++++++++++++++++++++++++-- 9 files changed, 385 insertions(+), 51 deletions(-) create mode 100644 python/restate/sessions/_scope.py create mode 100644 python/restate/sessions/_serve.py create mode 100644 python/restate/sessions/internal.py diff --git a/python/restate/sessions/__init__.py b/python/restate/sessions/__init__.py index ec1c431a..874d7849 100644 --- a/python/restate/sessions/__init__.py +++ b/python/restate/sessions/__init__.py @@ -28,14 +28,16 @@ async def list(ctx: sessions.ViewContext) -> list[str]: return [row["text"] for row in await ctx.db.query("SELECT text FROM notes ORDER BY rowid")] - app = restate.app([notes]) + app = sessions.app(notes) # with what it needs beside it; hypercorn module:app Session handlers don't replay a journal: a retried attempt runs again from its first line, against what the attempts before it committed. The types a handler gets (``SessionContext``, ``Database``, ``Transaction``, ``Rpc``, ``Send``, ...) are protocols: a decorated handler is still the plain -function, so a test can call it with a fake context of its own. A session object +function, so a test can call it with a fake context of its own. What libraries +built on session objects, and tests, need beside this is in +``restate.sessions.internal``. A session object serves beside the SDK's services, objects and workflows, which it calls with ``ctx.rpc`` for journaled work. """ @@ -54,9 +56,9 @@ async def list(ctx: sessions.ViewContext) -> list[str]: ) from restate.sessions._calls import CallFuture from restate.sessions._database import Row, Rows, SqlError, SqlParam, SqlValue -from restate.sessions._object import SessionObject, cancel_calls, cancelled -from restate.sessions._schema import Schema, Step -from restate.sessions._session import database, reset +from restate.sessions._object import SessionObject +from restate.sessions._schema import Step +from restate.sessions._serve import app __all__ = [ "CallFuture", @@ -67,7 +69,6 @@ async def list(ctx: sessions.ViewContext) -> list[str]: "Row", "Rows", "Rpc", - "Schema", "Send", "SessionContext", "SessionObject", @@ -77,8 +78,5 @@ async def list(ctx: sessions.ViewContext) -> list[str]: "Step", "Transaction", "ViewContext", - "cancel_calls", - "cancelled", - "database", - "reset", + "app", ] diff --git a/python/restate/sessions/_link.py b/python/restate/sessions/_link.py index 2affce0a..db69f3ae 100644 --- a/python/restate/sessions/_link.py +++ b/python/restate/sessions/_link.py @@ -25,20 +25,29 @@ class Link(typing.Generic[T]): - """A link of one kind, on the objects that have one: compared by identity, kept as long as the object is.""" - - def __init__(self, what: str) -> None: + """ + A link of one kind, on the objects that have one: compared by identity, + kept as long as the object is. A value that refers to its owner (an + attempt, to its database) would keep both forever, as a link holds its + value: ``weak`` holds it weakly, for the owner to keep it alive. + """ + + def __init__(self, what: str, *, weak: bool = False) -> None: self.what: str = what - self._links: weakref.WeakKeyDictionary[object, T] = weakref.WeakKeyDictionary() + self._weak: bool = weak + self._links: weakref.WeakKeyDictionary[object, T | weakref.ref[typing.Any]] = weakref.WeakKeyDictionary() def on(self, owner: O, value: T) -> O: """``owner``, linked to ``value``: once, as it is made.""" - self._links[owner] = value + self._links[owner] = weakref.ref(value) if self._weak else value return owner def of(self, owner: object) -> T | None: """What ``owner`` is linked to, if it has this link (and None for anything that cannot have one).""" try: - return self._links.get(owner) + value = self._links.get(owner) except TypeError: return None + if self._weak and isinstance(value, weakref.ref): + return typing.cast("T | None", value()) + return typing.cast("T | None", value) diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index b19c8bb9..b0847ea0 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -30,6 +30,7 @@ async def list(ctx: sessions.ViewContext) -> list[str]: import asyncio import functools import typing +import weakref from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass from datetime import timedelta @@ -48,6 +49,7 @@ async def list(ctx: sessions.ViewContext) -> list[str]: from restate.sessions._database import Store from restate.sessions._link import Link from restate.sessions._schema import Schema, Step +from restate.sessions._scope import SCOPE, Scope from restate.sessions._selector import FutureSelector, hook from restate.sessions._session import Attempt, ObjectConfig, SqlHost, reader_of, with_attempt from restate.vm import SingleUnresolvedFuture @@ -207,20 +209,6 @@ def cancel(invocation_id: str) -> None: return cancel -@dataclass(frozen=True) -class Scope: - """An attempt's own, as this package made it, on its context (``SCOPE``).""" - - sdk: ServerInvocationContext | None - """The SDK's context of the attempt: none over a test's own host.""" - explicit_cancellation: bool = False - """Whether its handler has ``explicit_cancellation=True``.""" - - -SCOPE: Link[Scope] = Link("scope") -"""The attempt's scope, on its context.""" - - @dataclass(frozen=True) class _Entry: """A handler or view of a session object: its function, kind, and object's config.""" @@ -228,6 +216,8 @@ class _Entry: fn: Callable[..., Awaitable[Any]] view: bool config: ObjectConfig + definition: weakref.ref[SessionObject] + """Its session object, weakly: the object's entries, which it holds, hold this.""" explicit_cancellation: bool = False @@ -239,13 +229,15 @@ def _context_of(entry: _Entry, host: SqlHost, request: Request, attempt: Attempt """The context of an attempt of ``entry``.""" rpc = Caller(host) cancel = _cancel_of(host) + sdk = typing.cast("ServerInvocationContext | None", getattr(host, "context", None)) + scope = Scope(sdk=sdk, explicit_cancellation=entry.explicit_cancellation, definition=entry.definition()) if entry.view: - return AttemptViewContext(key=host.key, request=request, db=reader_of(attempt.db), rpc=rpc, cancel=cancel) + view = AttemptViewContext(key=host.key, request=request, db=reader_of(attempt.db), rpc=rpc, cancel=cancel) + return SCOPE.on(view, scope) context = AttemptContext( key=host.key, request=request, db=attempt.db, rpc=rpc, send=Sender(send_now(host)), cancel=cancel ) - sdk = typing.cast("ServerInvocationContext | None", getattr(host, "context", None)) - return SCOPE.on(context, Scope(sdk=sdk, explicit_cancellation=entry.explicit_cancellation)) + return SCOPE.on(context, scope) async def _attempt(entry: _Entry, host: SqlHost, request: Request, args: Sequence[Any]) -> Any: @@ -284,8 +276,8 @@ async def cancelled(ctx: SessionContext) -> None: code that stops its own work then (a model's request, say):: async def stop_when_cancelled(): - await sessions.cancelled(ctx) - await sessions.cancel_calls(ctx) + await internal.cancelled(ctx) + await internal.cancel_calls(ctx) asyncio.ensure_future(stop_when_cancelled()) @@ -308,6 +300,16 @@ async def cancel_calls(ctx: SessionContext) -> list[str]: return await _preview.cancel_previous_calls(_explicit_of(ctx, "cancel_calls")) +def definition_of(ctx: SessionContext | ViewContext) -> SessionObject | None: + """ + The session object whose handler or view ``ctx`` is the context of: what + a library keeps per object (an agent's options, say) lives on it. + """ + scope = SCOPE.of(ctx) + definition = scope.definition if scope is not None else None + return definition if isinstance(definition, SessionObject) else None + + async def invoke(obj: SessionObject, name: str, host: SqlHost, request: Request, *args: Any) -> Any: """Run handler or view ``name`` of ``obj`` for one attempt over ``host``, a context of a test's own, without Restate.""" entry = (ENTRIES.of(obj) or {}).get(name) @@ -415,8 +417,8 @@ async def add(ctx: sessions.SessionContext, text: str) -> None: With options, ``@notes.handler(name=..., ...)``: the SDK's handler's. With ``explicit_cancellation=True``, Restate's cancellation of its invocation fails no await, nor cancels its calls: the handler learns of - it with ``sessions.cancelled(ctx)``, and decides what to stop (its calls, - with ``sessions.cancel_calls(ctx)``). + it with ``internal.cancelled(ctx)``, and decides what to stop (its calls, + with ``internal.cancel_calls(ctx)``), of ``restate.sessions.internal``. """ return self._decorate( False, @@ -523,7 +525,9 @@ def register(fn: F) -> F: raise TypeError(f"{what} {instead}") explicit = bool(options.get("explicit_cancellation")) sdk_options = {key: value for key, value in options.items() if key != "explicit_cancellation"} - entry = _Entry(fn=fn, view=view, config=self.config, explicit_cancellation=explicit) + entry = _Entry( + fn=fn, view=view, config=self.config, definition=weakref.ref(self), explicit_cancellation=explicit + ) registered = VirtualObject.handler(self, name, kind="shared" if view else "exclusive", **sdk_options)( _sdk_handler(entry) ) @@ -549,4 +553,4 @@ def with_store_wrapper(self, wrap: Callable[[Store], Store]) -> SessionObject: return self -__all__ = ["SessionObject", "cancel_calls", "cancelled"] +__all__ = ["SessionObject", "cancel_calls", "cancelled", "definition_of"] diff --git a/python/restate/sessions/_scope.py b/python/restate/sessions/_scope.py new file mode 100644 index 00000000..b2dafac1 --- /dev/null +++ b/python/restate/sessions/_scope.py @@ -0,0 +1,87 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +An attempt's scope: what its code keeps for the attempt (an agent's turn, the +calls of its tools...), which its context holds and which ends with it. A +retry gets a context, and so a scope, of its own. A module holds only the key +of what it keeps (``ScopeKey``), never the state. As sdk-sessions' ``scope.ts``. +""" + +from __future__ import annotations + +import typing +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +from restate.server_context import ServerInvocationContext +from restate.sessions._api import SessionContext, ViewContext +from restate.sessions._link import Link + +T = typing.TypeVar("T") + + +@dataclass +class Scope: + """An attempt's own, as this package made it, on its context (``SCOPE``), and what its code keeps for it.""" + + sdk: ServerInvocationContext | None + """The SDK's context of the attempt: none over a test's own host.""" + explicit_cancellation: bool = False + """Whether its handler has ``explicit_cancellation=True``.""" + definition: object | None = None + """The session object whose handler or view it runs (``definition_of``).""" + values: dict[ScopeKey[Any], Any] = field(default_factory=dict[Any, Any]) + """What its code keeps for it, by key: only ``ScopeKey`` reaches it.""" + + +SCOPE: Link[Scope] = Link("scope") +"""The attempt's scope, on its context.""" + + +class ScopeKey(typing.Generic[T]): + """ + A kind of state an attempt's code keeps, by its context: its key in the + attempt's scope. Kept per attempt, so a retry starts without it:: + + TURN: ScopeKey[Turn] = ScopeKey("an agent's turn") + + turn = TURN.of(ctx, lambda: Turn(...)) + """ + + def __init__(self, what: str) -> None: + self.what: str = what + """What it is, for an error.""" + + def get(self, ctx: SessionContext | ViewContext) -> T | None: + """What the attempt of ``ctx`` keeps under this key, if anything.""" + return typing.cast("T | None", self._values(ctx).get(self)) + + def set(self, ctx: SessionContext | ViewContext, value: T) -> None: + """Keep ``value`` under this key, for the attempt of ``ctx``.""" + self._values(ctx)[self] = value + + def delete(self, ctx: SessionContext | ViewContext) -> None: + """Keep nothing under this key, for the attempt of ``ctx``.""" + self._values(ctx).pop(self, None) + + def of(self, ctx: SessionContext | ViewContext, make: Callable[[], T]) -> T: + """What the attempt of ``ctx`` keeps under this key, made by ``make`` the first time.""" + values = self._values(ctx) + if self not in values: + values[self] = make() + return typing.cast(T, values[self]) + + def _values(self, ctx: SessionContext | ViewContext) -> dict[ScopeKey[Any], Any]: + scope = SCOPE.of(ctx) + if scope is None: + raise TypeError(f"{self.what} takes the context of a session object's handler or view") + return scope.values diff --git a/python/restate/sessions/_serve.py b/python/restate/sessions/_serve.py new file mode 100644 index 00000000..322005b0 --- /dev/null +++ b/python/restate/sessions/_serve.py @@ -0,0 +1,93 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +``app``: the ASGI app of a program's services, with what they need beside them +(a library's objects, such as an agent's ``signals``). As sdk-sessions' +``serve.ts``; serving it is the program's (hypercorn, uvicorn...), as for +``restate.app``. +""" + +from __future__ import annotations + +import os +import typing +from collections.abc import Callable, Iterable, Sequence + +import restate +from restate.entry_codec import JournalValueCodec, JournalValueCodecProvider +from restate.server_types import RestateAppT +from restate.sessions._link import Link + +Service = restate.Service | restate.VirtualObject | restate.Workflow +"""A service, virtual object (a session object too) or workflow.""" + +D = typing.TypeVar("D") + +_NEEDS: Link[Callable[[], Sequence[Service]]] = Link("needs") +"""The services a definition needs served beside it: see ``needing``.""" + + +def needing(definition: D, services: Callable[[], Sequence[Service]]) -> D: + """ + ``definition``, which needs ``services`` served beside it: ``app`` serves + them too. For a library whose objects call one of their own (an agent's + ``signals``, say). ``services`` is called once ``app`` asks. + """ + return _NEEDS.on(definition, services) + + +def served(services: Iterable[Service], names: str | None = None) -> list[Service]: + """ + What ``app`` serves of ``services``: each with what it needs beside it, + once, of those ``names`` names (comma separated; all of them without). A + name of none is an error. + """ + every: dict[str, Service] = {} + + def add(service: Service) -> None: + if service.name in every: + return + every[service.name] = service + needs = _NEEDS.of(service) + for needed in needs() if needs is not None else (): + add(needed) + + for service in services: + add(service) + wanted = [name.strip() for name in (names or "").split(",") if name.strip()] + if not wanted: + return list(every.values()) + unknown = [name for name in wanted if name not in every] + if unknown: + raise ValueError(f"RESTATE_SERVE_ONLY names {', '.join(unknown)}: the services are {', '.join(every)}") + return [service for name, service in every.items() if name in wanted] + + +def app( + *services: Service, + protocol: typing.Literal["bidi", "request_response"] | None = None, + identity_keys: list[str] | None = None, + journal_value_codec: JournalValueCodec | JournalValueCodecProvider | None = None, +) -> RestateAppT: + """ + The ASGI app of ``services``, with the services they need beside them, + each once: ``restate.app``'s, whose options it takes. ``RESTATE_SERVE_ONLY`` + (names separated by commas, such as ``triage,planner,signals``) serves only + those of them: one program (or image), deployed in parts:: + + app = sessions.app(notes, orders) # hypercorn module:app + """ + return restate.app( + served(services, os.environ.get("RESTATE_SERVE_ONLY")), + protocol=protocol, + identity_keys=identity_keys, + journal_value_codec=journal_value_codec, + ) diff --git a/python/restate/sessions/_session.py b/python/restate/sessions/_session.py index 3b0c3c5b..77ad05da 100644 --- a/python/restate/sessions/_session.py +++ b/python/restate/sessions/_session.py @@ -440,8 +440,8 @@ def reader_of(db: AttemptDatabase) -> DatabaseReader: _DATABASE_NAME = re.compile(r"^[\w.-]{1,64}$") -_ATTEMPT: Link[Attempt] = Link("database") -"""The attempt of a ``ctx.db``: for ``database``.""" +_ATTEMPT: Link[Attempt] = Link("database", weak=True) +"""The attempt of a ``ctx.db``, which keeps it alive (and refers to it): for ``database``.""" _RESET: Link[Callable[[], Awaitable[None]]] = Link("reset") """How a transaction of an exclusive handler's ``ctx.db`` empties the database: for ``reset``.""" diff --git a/python/restate/sessions/internal.py b/python/restate/sessions/internal.py new file mode 100644 index 00000000..0209366c --- /dev/null +++ b/python/restate/sessions/internal.py @@ -0,0 +1,48 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +What libraries built on session objects (``restate.agents``) and tests need, +beside the API handlers code against: as sdk-sessions' +``@restatedev/sessions/internal``. + +- ``Schema``: tables a library owns, applied where it uses them + (``schema.ensure(ctx.db)``); ``database``, another database of the same + key; ``reset``, the key's database emptied in a transaction. +- ``cancelled`` and ``cancel_calls``: explicit cancellation, for a handler + with ``explicit_cancellation=True``. +- ``ScopeKey`` and ``definition_of``: what a library keeps for an attempt, or + for an object. +- ``needing`` and ``served``: the services a definition needs served beside + it, which ``restate.sessions.app`` serves. +- ``invoke``: a handler's attempt over a host of a test's own, without + Restate (see ``SqlHost``). +""" + +from restate.sessions._object import cancel_calls, cancelled, definition_of, invoke +from restate.sessions._schema import Schema +from restate.sessions._scope import ScopeKey +from restate.sessions._serve import needing, served +from restate.sessions._session import ObjectConfig, SqlHost, database, reset + +__all__ = [ + "ObjectConfig", + "Schema", + "ScopeKey", + "SqlHost", + "cancel_calls", + "cancelled", + "database", + "definition_of", + "invoke", + "needing", + "reset", + "served", +] diff --git a/tests/sessions_harness.py b/tests/sessions_harness.py index c01844a5..2ea69cb7 100644 --- a/tests/sessions_harness.py +++ b/tests/sessions_harness.py @@ -23,6 +23,7 @@ import restate from restate import HarnessEnvironment, sessions +from restate.sessions import internal from restate.sessions import _object, _session @@ -146,10 +147,10 @@ async def nap(ctx: restate.Context) -> str: async def stoppable(ctx: sessions.SessionContext, text: str) -> list[str]: """Waits for a nap; cancelled, it cancels the nap, records that it stopped, and answers the calls it cancelled.""" napping = asyncio.ensure_future(ctx.rpc(nap)()) - stopping = asyncio.ensure_future(sessions.cancelled(ctx)) + stopping = asyncio.ensure_future(internal.cancelled(ctx)) await asyncio.wait({napping, stopping}, return_when=asyncio.FIRST_COMPLETED) assert stopping.done() - cancelled = await sessions.cancel_calls(ctx) + cancelled = await internal.cancel_calls(ctx) await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [f"stopped: {text}"]) return cancelled @@ -157,7 +158,7 @@ async def stoppable(ctx: sessions.SessionContext, text: str) -> list[str]: @pytest.fixture(scope="session") async def harness(): async with restate.create_test_harness( - restate.app([inbox, pricing, sleeper, notes]), restate_image="ghcr.io/restatedev/restate:latest" + sessions.app(inbox, pricing, sleeper, notes), restate_image="ghcr.io/restatedev/restate:latest" ) as harness: yield harness diff --git a/tests/sessions_object.py b/tests/sessions_object.py index 35021947..c80ef9b4 100644 --- a/tests/sessions_object.py +++ b/tests/sessions_object.py @@ -22,6 +22,7 @@ import restate from restate import sessions +from restate.sessions import internal from restate.context import AttemptFinishedEvent, Request from restate.exceptions import TerminalError from restate.handler import handler_from_callable @@ -359,7 +360,7 @@ async def test_another_database_of_the_key_is_apart_from_the_objects_own(): @multi.handler() async def both(ctx: sessions.SessionContext) -> list[str]: - other = sessions.database(ctx.db, "other") + other = internal.database(ctx.db, "other") await other.execute("CREATE TABLE apart (b INTEGER)") rows = await other.query("SELECT name FROM sqlite_schema WHERE type = 'table'") return [typing.cast(str, row["name"]) for row in rows] @@ -380,7 +381,7 @@ async def fill(ctx: sessions.SessionContext) -> None: @resettable.handler() async def empty(ctx: sessions.SessionContext) -> list[Any]: async with ctx.db.transaction() as tx: - await sessions.reset(tx) + await internal.reset(tx) await tx.execute("INSERT INTO t VALUES (2)") return await ctx.db.values("SELECT name FROM sqlite_schema WHERE type = 'table' UNION ALL SELECT a FROM t") @@ -454,11 +455,11 @@ async def test_cancelled_takes_the_ctx_of_a_handler_with_explicit_cancellation() @obj.handler async def implicit(ctx: sessions.SessionContext) -> None: - await sessions.cancelled(ctx) + await internal.cancelled(ctx) @obj.handler(explicit_cancellation=True) async def explicit(ctx: sessions.SessionContext) -> None: - await sessions.cancel_calls(ctx) + await internal.cancel_calls(ctx) assert not obj.handlers["implicit"].explicit_cancellation assert obj.handlers["explicit"].explicit_cancellation @@ -470,7 +471,7 @@ async def explicit(ctx: sessions.SessionContext) -> None: async def test_cancelled_ends_its_waiter_cancelled_when_the_attempt_ends_first(monkeypatch): from restate import _preview # pylint: disable=import-outside-toplevel from restate.exceptions import SuspendedException # pylint: disable=import-outside-toplevel - from restate.sessions import _calls, _object # pylint: disable=import-outside-toplevel + from restate.sessions import _calls, _object, _scope # pylint: disable=import-outside-toplevel async def attempt_ended(_sdk: Any) -> None: raise SuspendedException() @@ -485,8 +486,8 @@ async def attempt_ended(_sdk: Any) -> None: send=_calls.Sender(lambda send: None), cancel=lambda invocation_id: None, ) - _object.SCOPE.on(ctx, _object.Scope(sdk=typing.cast(Any, object()), explicit_cancellation=True)) - watcher = asyncio.ensure_future(sessions.cancelled(ctx)) + _scope.SCOPE.on(ctx, _scope.Scope(sdk=typing.cast(Any, object()), explicit_cancellation=True)) + watcher = asyncio.ensure_future(internal.cancelled(ctx)) await asyncio.wait({watcher}) # Cancelled, not failed: asyncio reports no exception nobody retrieved. assert watcher.cancelled() @@ -538,3 +539,96 @@ async def test_a_handler_runs_over_a_fake_context_of_a_test(): assert fake.db.statements == [("INSERT INTO notes (text) VALUES (?)", ["milk"])] # And the function is still the handler that calls and the client name. assert handler_from_callable(add).name == "add" + + +TURNS: internal.ScopeKey[list[str]] = internal.ScopeKey("the turns of a test") + + +async def test_a_scope_key_keeps_its_value_for_the_attempt_and_a_new_attempt_starts_without_it(): + scoped = sessions.SessionObject("scoped") + + @scoped.handler + async def turn(ctx: sessions.SessionContext, text: str) -> list[str]: + seen_before = TURNS.get(ctx) + turns = TURNS.of(ctx, list) + turns.append(text) + # `of` makes it once: the same list, for the attempt. + assert TURNS.of(ctx, list) is turns + return [*(seen_before or []), *turns] + + @scoped.view + async def peek(ctx: sessions.ViewContext) -> list[str]: + TURNS.set(ctx, ["view"]) + value = TURNS.get(ctx) + TURNS.delete(ctx) + assert TURNS.get(ctx) is None + return value or [] + + host = Host(Backend(), service="scoped") + assert await run(scoped, "turn", host, "a") == ["a"] + # Another attempt, another scope: it starts without what the first kept. + assert await run(scoped, "turn", host, "b") == ["b"] + assert await run(scoped, "peek", host.view()) == ["view"] + + +async def test_definition_of_gives_the_session_object_of_a_handler_or_a_view(): + defined = sessions.SessionObject("defined") + + @defined.handler + async def handled(ctx: sessions.SessionContext) -> bool: + return internal.definition_of(ctx) is defined + + @defined.view + async def viewed(ctx: sessions.ViewContext) -> bool: + return internal.definition_of(ctx) is defined + + host = Host(Backend(), service="defined") + assert await run(defined, "handled", host) + assert await run(defined, "viewed", host.view()) + + +def test_a_scope_key_takes_the_context_of_a_session_object(): + with pytest.raises(TypeError, match="takes the context of a session object's handler or view"): + TURNS.get(FakeContext(db=FakeDatabase())) + assert internal.definition_of(FakeContext(db=FakeDatabase())) is None + + +def test_served_adds_what_each_service_needs_once_and_serve_only_picks(): + companion = restate.VirtualObject("companion") + helper = restate.Service("helper") + first = internal.needing(sessions.SessionObject("first"), lambda: [companion, helper]) + second = internal.needing(sessions.SessionObject("second"), lambda: [companion]) + internal.needing(companion, lambda: [helper]) + + assert [s.name for s in internal.served([first, second])] == ["first", "companion", "helper", "second"] + assert [s.name for s in internal.served([first, second], "second, helper")] == ["helper", "second"] + with pytest.raises(ValueError, match="RESTATE_SERVE_ONLY names nowhere"): + internal.served([first], "nowhere") + + +def test_app_serves_what_restate_serve_only_names(monkeypatch): + companion = restate.VirtualObject("companion2") + first = internal.needing(sessions.SessionObject("first2"), lambda: [companion]) + monkeypatch.setenv("RESTATE_SERVE_ONLY", "companion2") + assert callable(sessions.app(first)) + monkeypatch.setenv("RESTATE_SERVE_ONLY", "unknown") + with pytest.raises(ValueError, match="RESTATE_SERVE_ONLY names unknown"): + sessions.app(first) + + +async def test_an_attempt_is_freed_once_it_ends(): + import gc # pylint: disable=import-outside-toplevel + import weakref # pylint: disable=import-outside-toplevel + + kept: list[weakref.ref[Any]] = [] + freed = sessions.SessionObject("freed", schema=["CREATE TABLE t (a INTEGER)"]) + + @freed.handler + async def use(ctx: sessions.SessionContext) -> None: + await internal.database(ctx.db, "other").execute("CREATE TABLE IF NOT EXISTS u (b INTEGER)") + kept.append(weakref.ref(ctx.db)) + + await run(freed, "use", Host(Backend(), service="freed")) + gc.collect() + # Nothing of this package's holds an attempt's database (and so the attempt) past it. + assert kept[0]() is None From c0fb67c6947f53ad6e3a3c1641ad3ad7b63b71cf Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 16:10:18 +0100 Subject: [PATCH 18/45] restate.agents: the conversation's types, Transcript, and its log The first part of the agents' core, ported from sdk-sessions' agents package: - _api.py: what clients read and send, as JSON (ChatMessage and its parts, Messages and MessagesRequest, Turn), with sdk-sessions' field names, so that a client follows an agent of either SDK. - _transcript.py: Transcript, the conversation as a framework builds it (say, call, result); a message given out never changes, so the log stores only what is new. - _log.py: the log in the object's database (sdk-sessions' tables): each message an id that never changes and a seq, the cursor of its last change, so a client reads each change once; a replaced transcript leaves its messages in the log; restart, for clear, goes on from where the log got to. The agents' package is linted and typed as restate.sessions is. Nothing uses it yet. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/agents/_api.py | 120 ++++++++++++++ python/restate/agents/_log.py | 238 +++++++++++++++++++++++++++ python/restate/agents/_transcript.py | 96 +++++++++++ python/restate/agents/ruff.toml | 15 ++ tests/agents_log.py | 133 +++++++++++++++ 5 files changed, 602 insertions(+) create mode 100644 python/restate/agents/_api.py create mode 100644 python/restate/agents/_log.py create mode 100644 python/restate/agents/_transcript.py create mode 100644 python/restate/agents/ruff.toml create mode 100644 tests/agents_log.py diff --git a/python/restate/agents/_api.py b/python/restate/agents/_api.py new file mode 100644 index 00000000..ed0580b2 --- /dev/null +++ b/python/restate/agents/_api.py @@ -0,0 +1,120 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +What clients of an agent read and send, as JSON: its conversation's messages +and pages of them, and the message it answers. The field names are +sdk-sessions' (``invocationId``), so that a client follows an agent of either +SDK. +""" + +from __future__ import annotations + +import typing +from typing import Any, Literal, TypedDict + +Role = Literal["user", "assistant"] +"""Who said a message.""" + + +class TextPart(TypedDict): + """Text a message says.""" + + type: Literal["text"] + text: str + + +class _ToolCall(TypedDict): + type: Literal["tool"] + id: str + name: str + input: Any + + +class ToolPart(_ToolCall, total=False): + """ + A tool call a message makes, with its output once it ran, and ``error: + true`` if the output is the tool's error (``Error: `` of a + ``restate.TerminalError``, say). + """ + + output: str + error: Literal[True] + + +Part = TextPart | ToolPart +"""What a message says: text, or a tool call.""" + + +class Said(TypedDict): + """A message as a framework's transcript gives it: its role and parts, without an id.""" + + role: Role + parts: list[Part] + + +class ChatMessage(Said): + """A message of a conversation, as an agent's ``messages`` view gives it.""" + + id: int + """Its id, from 0, in the order of the conversation; the same however it changes. ``before`` takes it.""" + + +class MessagesRequest(TypedDict, total=False): + """ + Which messages to read: the latest ``limit``, or those before a message, + for a page of the conversation; or those new or changed since a + ``cursor``, to follow it. + """ + + limit: int + """At most this many messages (default 50).""" + before: int + """The messages before this one (its ``id``): older ones, such as for scrolling back.""" + cursor: int + """The messages new or changed since this cursor (a page's ``cursor``), by when they changed.""" + + +class _Messages(TypedDict): + messages: list[ChatMessage] + more: bool + cursor: int + + +class Messages(_Messages, total=False): + """ + A page of messages: oldest first; or, with a ``cursor``, those new or + changed since it, by when they changed. Keep them by ``id``: a message + comes again once it changed (its tool call's output came, say). ``more`` + says whether there are more (older ones, or with a ``cursor``, newer + ones); ``cursor`` is the one to read on from. + """ + + first: int + """The id of the conversation's first message, once it was cleared: a client drops the messages it kept from before.""" + + +class Turn(TypedDict): + """The message a ``chat`` handler is answering: its invocation, and when it started.""" + + invocationId: str + """The invocation's id: ``/restate/invocation//attach`` waits for its answer.""" + since: str + """When it started, as an ISO date.""" + + +def text_part(text: str) -> TextPart: + """A text part.""" + return {"type": "text", "text": text} + + +def is_tool_part(part: Part) -> typing.TypeGuard[ToolPart]: + """Whether ``part`` is a tool call's.""" + return part["type"] == "tool" diff --git a/python/restate/agents/_log.py b/python/restate/agents/_log.py new file mode 100644 index 00000000..2dc92f72 --- /dev/null +++ b/python/restate/agents/_log.py @@ -0,0 +1,238 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +An agent's conversation as clients read it: a log in the object's database, +kept in step with the framework's transcript each time the agent publishes a +change, and read by ``messages`` and ``sync``. As sdk-sessions' ``log.ts``, +whose tables it shares. + +Each message has an ``id`` that never changes, and a ``seq``: the sequence +number of its last change, which is the cursor. A message that grows in place +(a tool call's output comes) keeps its id and takes a new ``seq``; a +transcript replaced (compacted, reset) leaves its messages in the log, and the +messages after them are new ones. So a client reads from a cursor, and gets +each message that is new or changed since, once. + +``clear`` forgets the log with the rest of the conversation, but not where it +got to: the next conversation's ids and cursors go on from there (its start, a +row of ``restate_agent_log_start``), so a client following it reads on, and +drops the messages before ``first``. +""" + +from __future__ import annotations + +import json +import typing +from dataclasses import dataclass + +from restate.agents._api import ChatMessage, Messages, MessagesRequest, Said +from restate.sessions import DatabaseReader, DatabaseWriter, SessionContext +from restate.sessions.internal import Schema, ScopeKey + +TABLE = "restate_agent_log" + +SCHEMA = Schema( + "restate-agents.log", + [ + f"""CREATE TABLE {TABLE} ( + id INTEGER PRIMARY KEY, + seq INTEGER NOT NULL UNIQUE, + position INTEGER UNIQUE, + message TEXT NOT NULL + )""", + # Where the log starts, which `clear` sets: no row before. + f"""CREATE TABLE {TABLE}_start ( + row INTEGER PRIMARY KEY CHECK (row = 0), + seq INTEGER NOT NULL, + id INTEGER NOT NULL + )""", + ], +) + + +@dataclass(frozen=True) +class Start: + """Where the log starts: past the ids and cursors of the conversations ``clear`` forgot.""" + + seq: int + """The cursor it starts at.""" + id: int + """Its first message's id.""" + + +async def _start_of(db: DatabaseReader) -> Start | None: + """Where the log of ``db`` starts, if ``clear`` set it.""" + row = await db.get(f"SELECT seq, id FROM {TABLE}_start") + return Start(seq=int(typing.cast(int, row["seq"])), id=int(typing.cast(int, row["id"]))) if row else None + + +async def restart(tx: DatabaseWriter, start: Start) -> None: + """ + Start the log at ``start``, in the transaction ``tx`` of ``clear`` that + emptied the database: the next conversation's ids and cursors go on from + there. + """ + await tx.execute(f"INSERT OR REPLACE INTO {TABLE}_start (row, seq, id) VALUES (0, ?, ?)", [start.seq, start.id]) + + +@dataclass +class _Row: + """A message of the log's head, the transcript as last stored: its id, and the message, also as JSON.""" + + id: int + message: Said + json: str + + +@dataclass +class _Head: + """The log's head as a session stored it last, with the last cursor and the next id.""" + + rows: list[_Row] + seq: int + next: int + + +_HEAD: ScopeKey[_Head] = ScopeKey("the log") +"""The log's head, as the attempt's last ``sync`` stored it: the next compares the transcript with it.""" + + +def _json(message: Said) -> str: + return json.dumps(message, separators=(",", ":")) + + +async def _head_of(ctx: SessionContext) -> _Head: + """The log's head in the attempt of ``ctx``: as it stored it last, or else as the log has it.""" + cached = _HEAD.get(ctx) + if cached is not None: + return cached + db = ctx.db + rows = [ + _Row( + id=int(typing.cast(int, id)), message=json.loads(typing.cast(str, message)), json=typing.cast(str, message) + ) + for id, message in await db.values( + f"SELECT id, message FROM {TABLE} WHERE position IS NOT NULL ORDER BY position" + ) + ] + [[last, nxt]] = await db.values(f"SELECT MAX(seq), MAX(id) + 1 FROM {TABLE}") + if last is not None: + return _Head(rows=rows, seq=int(typing.cast(int, last)), next=int(typing.cast(int, nxt))) + start = await _start_of(db) or Start(seq=0, id=0) + return _Head(rows=rows, seq=start.seq, next=start.id) + + +async def sync(ctx: SessionContext, messages: list[Said]) -> None: + """Store the transcript's ``messages``, in order, in the log: those that changed since the last time, in one transaction.""" + db = ctx.db + await SCHEMA.ensure(db) + head = await _head_of(ctx) + # A message the transcript did not change is the one stored (see Transcript): only the others are compared, as JSON. + jsons: dict[int, str] = {} + + def json_of(position: int) -> str: + if position not in jsons: + jsons[position] = _json(messages[position]) + return jsons[position] + + def unchanged(position: int) -> bool: + row = head.rows[position] + return row.message is messages[position] or row.json == json_of(position) + + both = min(len(head.rows), len(messages)) + same = 0 + while same < both and unchanged(same): + same += 1 + grown = same + while grown < both and _grown_from(messages[grown], head.rows[grown].message): + grown += 1 + if same == len(head.rows) == len(messages): + return + seq, next_id = head.seq, head.next + rows = head.rows[:same] + try: + async with db.transaction() as tx: + for position in range(same, grown): + row = _Row(id=head.rows[position].id, message=messages[position], json=json_of(position)) + seq += 1 + await tx.execute(f"UPDATE {TABLE} SET seq = ?, message = ? WHERE id = ?", [seq, row.json, row.id]) + rows.append(row) + # The rest of the transcript as it was is replaced: its messages stay in the log, out of it. + await tx.execute(f"UPDATE {TABLE} SET position = NULL WHERE position >= ?", [grown]) + for position in range(grown, len(messages)): + row = _Row(id=next_id, message=messages[position], json=json_of(position)) + next_id += 1 + seq += 1 + await tx.execute(f"INSERT INTO {TABLE} VALUES (?, ?, ?, ?)", [row.id, seq, position, row.json]) + rows.append(row) + except BaseException: + # Rolled back: the log is as it was, which the next sync reads again. + _HEAD.delete(ctx) + raise + _HEAD.set(ctx, _Head(rows=rows, seq=seq, next=next_id)) + + +async def page(db: DatabaseReader, request: MessagesRequest | None = None) -> Messages: + """ + A page of the log. With ``cursor``: the messages new or changed since it, + by when they changed, and the cursor to read on from. Without: the latest + ``limit`` messages (or those before ``before``, an id), oldest first, and + the cursor to follow the conversation from. + """ + request = request or {} + limit = request.get("limit", 50) + await SCHEMA.ensure(db) + start = await _start_of(db) + cursor = request.get("cursor") + if cursor is not None: + rows = await db.values( + f"SELECT id, message, seq FROM {TABLE} WHERE seq > ? ORDER BY seq LIMIT ?", [cursor, limit + 1] + ) + shown = rows[:limit] + result: Messages = { + "messages": [_message(row) for row in shown], + "more": len(rows) > limit, + "cursor": int(typing.cast(int, shown[-1][2])) if shown else cursor, + } + else: + [[now]] = await db.values(f"SELECT COALESCE(MAX(seq), ?) FROM {TABLE}", [start.seq if start else 0]) + before = request.get("before") + rows = await db.values( + f"SELECT id, message FROM {TABLE} WHERE id < ? ORDER BY id DESC LIMIT ?", + [before if before is not None else 2**53 - 1, limit + 1], + ) + result = { + "messages": [_message(row) for row in reversed(rows[:limit])], + "more": len(rows) > limit, + "cursor": int(typing.cast(int, now)), + } + if start is not None: + result["first"] = start.id + return result + + +def _message(row: list[typing.Any]) -> ChatMessage: + """A row of the log, as ``messages`` gives it.""" + said: Said = json.loads(typing.cast(str, row[1])) + return {"id": int(row[0]), "role": said["role"], "parts": said["parts"]} + + +def _grown_from(now: Said, was: Said) -> bool: + """Whether ``now`` is ``was`` grown in place: the same role and parts, with more (a tool call's output, say).""" + if was["role"] != now["role"] or len(now["parts"]) < len(was["parts"]): + return False + for part, other in zip(was["parts"], now["parts"], strict=False): + if part["type"] == "text": + if other["type"] != "text" or not other["text"].startswith(part["text"]): + return False + elif other["type"] != "tool" or other["id"] != part["id"]: + return False + return True diff --git a/python/restate/agents/_transcript.py b/python/restate/agents/_transcript.py new file mode 100644 index 00000000..7a88699a --- /dev/null +++ b/python/restate/agents/_transcript.py @@ -0,0 +1,96 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +A conversation as ``messages`` gives it, built from a framework's transcript. +As sdk-sessions' ``Transcript`` (``common.ts``). +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +from restate.agents._api import Part, Role, Said, ToolPart, is_tool_part, text_part + + +class Transcript: + """ + A conversation as ``messages`` gives it, built from a framework's + transcript: each message is started by ``say``, a tool call goes to the + assistant's message, and its result to its call. Messages without parts + are left out. + + A message, once given out by ``messages()``, is never changed: a change + makes a new one in its place. So a framework can keep its transcript and + add to it as its conversation grows, and the log stores only the messages + that are not the ones it stored last (see ``_log.py``). + """ + + def __init__(self) -> None: + self._messages: list[Said] = [] + # The message of each tool call, by the call's id: its index. + self._calls: dict[str, int] = {} + + def say(self, role: Role, text: str = "") -> None: + """A new message of ``role``, with ``text`` if there is any.""" + self._messages.append({"role": role, "parts": [text_part(text)] if text.strip() else []}) + + def call(self, call_id: str, name: str, input: Any) -> None: + """Tool call ``call_id`` of the assistant's message.""" + if not self._messages or self._messages[-1]["role"] != "assistant": + self.say("assistant") + at = len(self._messages) - 1 + call: ToolPart = {"type": "tool", "id": call_id, "name": name, "input": input} + self._change(at, lambda parts: [*parts, call]) + self._calls[call_id] = at + + def result(self, call_id: str, output: str, error: bool = False) -> None: + """The output of tool call ``call_id``; an error's if the framework says so, or it is a tool's ``Error: ``.""" + at = self._calls.get(call_id) + if at is None: + return + failed = error or output.startswith("Error: ") + + def with_output(parts: list[Part]) -> list[Part]: + changed: list[Part] = [] + for part in parts: + if is_tool_part(part) and part["id"] == call_id: + answered: ToolPart = {**part, "output": output} + if failed: + answered["error"] = True + changed.append(answered) + else: + changed.append(part) + return changed + + self._change(at, with_output) + + def answered(self, call_id: str, failed: str | None = None) -> bool: + """Whether tool call ``call_id`` has its output: not an error's, unless it is ``failed``, the error the call answered with.""" + at = self._calls.get(call_id) + if at is None: + return False + return any( + is_tool_part(part) + and part["id"] == call_id + and "output" in part + and (not part.get("error") or part["output"] == failed) + for part in self._messages[at]["parts"] + ) + + def messages(self) -> list[Said]: + """Its messages, in order: those with parts.""" + return [message for message in self._messages if message["parts"]] + + def _change(self, at: int, change: Callable[[list[Part]], list[Part]]) -> None: + """Message ``at`` with its parts as ``change`` makes them: a new message, in its place.""" + message = self._messages[at] + self._messages[at] = {"role": message["role"], "parts": change(message["parts"])} diff --git a/python/restate/agents/ruff.toml b/python/restate/agents/ruff.toml new file mode 100644 index 00000000..6b09b194 --- /dev/null +++ b/python/restate/agents/ruff.toml @@ -0,0 +1,15 @@ +# The project's ruff settings, and more rules for this package, which is +# written to them: annotations everywhere, likely bugs, modern syntax for +# Python 3.10, sorted imports, simpler code, async pitfalls. +extend = "../../../pyproject.toml" + +[lint] +extend-select = ["ANN", "B", "C4", "I", "UP", "SIM", "RUF", "ASYNC", "PIE", "RET"] +extend-ignore = [ + # Any is what a handler's input or a serialized value is; the type + # checkers (pyright strict) catch unknown types where they matter. + "ANN401", + # A serde instance as a default (input_serde=DefaultSerde()) is the + # SDK's API style, which SessionObject.handler mirrors. + "B008", +] diff --git a/tests/agents_log.py b/tests/agents_log.py new file mode 100644 index 00000000..b3d0535e --- /dev/null +++ b/tests/agents_log.py @@ -0,0 +1,133 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""An agent's conversation log: what a transcript stores, and what clients read of it.""" + +# pylint: disable=missing-function-docstring + +import typing + +import pytest + +from restate import sessions +from restate.agents import _log +from restate.agents._api import Messages, MessagesRequest +from restate.agents._transcript import Transcript +from sessions_database import Backend # pylint: disable=import-error +from sessions_object import Host, run # pylint: disable=import-error + + +@pytest.fixture +def anyio_backend(): + return "asyncio" + + +chat = sessions.SessionObject("logged") +transcripts: dict[str, Transcript] = {} + + +@chat.handler +async def store(ctx: sessions.SessionContext, name: str) -> None: + await _log.sync(ctx, transcripts[name].messages()) + + +@chat.handler +async def start_again(ctx: sessions.SessionContext, start: list[int]) -> None: + async with ctx.db.transaction() as tx: + await tx.execute("DELETE FROM restate_agent_log") + await _log.restart(tx, _log.Start(seq=start[0], id=start[1])) + + +@chat.view +async def read(ctx: sessions.ViewContext, request: MessagesRequest) -> Messages: + return await _log.page(ctx.db, request) + + +def said(name: str) -> Transcript: + transcripts[name] = Transcript() + return transcripts[name] + + +def texts(page: Messages) -> list[tuple[int, str]]: + return [ + (message["id"], " ".join(typing.cast(str, part.get("text", part.get("name"))) for part in message["parts"])) + for message in page["messages"] + ] + + +async def test_a_transcript_is_stored_and_read_in_pages_and_from_a_cursor(): + host = Host(Backend(), service="logged") + transcript = said("first") + transcript.say("user", "What is the weather in Lisbon?") + transcript.call("call-1", "get_weather", {"city": "Lisbon"}) + await run(chat, "store", host, "first") + page = await run(chat, "read", host.view(), {}) + assert texts(page) == [(0, "What is the weather in Lisbon?"), (1, "get_weather")] + cursor = page["cursor"] + + # The tool call's output comes: its message grows in place, keeps its id, and moves past the cursor. + transcript.result("call-1", "Sunny") + transcript.say("assistant", "It is sunny in Lisbon.") + await run(chat, "store", host, "first") + changed = await run(chat, "read", host.view(), {"cursor": cursor}) + assert [message["id"] for message in changed["messages"]] == [1, 2] + tool = changed["messages"][0]["parts"][0] + assert tool == { + "type": "tool", + "id": "call-1", + "name": "get_weather", + "input": {"city": "Lisbon"}, + "output": "Sunny", + } + + # Nothing changed: nothing is written, and the cursor stays. + await run(chat, "store", host, "first") + assert (await run(chat, "read", host.view(), {}))["cursor"] == changed["cursor"] + assert (await run(chat, "read", host.view(), {"cursor": changed["cursor"]}))["messages"] == [] + + +async def test_an_error_output_is_marked_and_answered_only_by_its_own_error(): + transcript = Transcript() + transcript.call("call-1", "charge", {"amount": 5}) + assert not transcript.answered("call-1") + transcript.result("call-1", "Error: card declined") + assert transcript.messages()[0]["parts"][0].get("error") is True + # An error is the call's result only if the call answered with it (see the calls' bookkeeping). + assert not transcript.answered("call-1") + assert transcript.answered("call-1", failed="Error: card declined") + + +async def test_a_replaced_transcript_leaves_its_messages_in_the_log(): + host = Host(Backend(), service="logged") + transcript = said("long") + for n in range(3): + transcript.say("user", f"message {n}") + await run(chat, "store", host, "long") + compacted = said("long") + compacted.say("assistant", "A summary of it all") + await run(chat, "store", host, "long") + page = await run(chat, "read", host.view(), {}) + assert texts(page) == [(0, "message 0"), (1, "message 1"), (2, "message 2"), (3, "A summary of it all")] + assert texts(await run(chat, "read", host.view(), {"limit": 2})) == [(2, "message 2"), (3, "A summary of it all")] + assert texts(await run(chat, "read", host.view(), {"before": 2})) == [(0, "message 0"), (1, "message 1")] + + +async def test_a_restarted_log_goes_on_from_where_the_cleared_one_got(): + host = Host(Backend(), service="logged") + transcript = said("cleared") + transcript.say("user", "before") + await run(chat, "store", host, "cleared") + await run(chat, "start_again", host, [10, 5]) + after = said("cleared") + after.say("user", "after") + await run(chat, "store", host, "cleared") + page = await run(chat, "read", host.view(), {}) + assert texts(page) == [(5, "after")] + assert page["first"] == 5 and page["cursor"] == 11 From ebd96dd69031efabdaea7e9b5649aed2594e4145 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 16:12:11 +0100 Subject: [PATCH 19/45] restate.agents: the companion object, signals The second part of the agents' core: sdk-sessions' signals.ts, a journaled virtual object per agent key (`:`), with the same handlers, so that it is the same service whichever SDK serves it: - wait / resolve / pending (and listen, forget): a tool's wait for a signal (an approval, say), resolved by its id or the oldest of a name, or kept for the next wait; an agent's own handler cannot wait on an awakeable, so it calls wait, which can, holding no process. - start / awaitTurn / close / status: the message the agent answers; awaitTurn attaches to the turn's invocation, and closes it if it failed without closing itself (killed, out of retries). - publish / notifications / watch (subscribe, unsubscribe): the conversation's revision, which clients wait on for a newer one. - with_companion, companion_of, deliver: how an agent reaches its companion, and sessions.app serves it beside the agent. Tested against a real Restate server. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/agents/_api.py | 66 +++++ python/restate/agents/_signals.py | 392 ++++++++++++++++++++++++++++++ tests/agents_signals.py | 156 ++++++++++++ 3 files changed, 614 insertions(+) create mode 100644 python/restate/agents/_signals.py create mode 100644 tests/agents_signals.py diff --git a/python/restate/agents/_api.py b/python/restate/agents/_api.py index ed0580b2..f980aa15 100644 --- a/python/restate/agents/_api.py +++ b/python/restate/agents/_api.py @@ -118,3 +118,69 @@ def text_part(text: str) -> TextPart: def is_tool_part(part: Part) -> typing.TypeGuard[ToolPart]: """Whether ``part`` is a tool call's.""" return part["type"] == "tool" + + +class _Waiting(TypedDict): + id: str + name: str + since: str + + +class Waiting(_Waiting, total=False): + """ + A wait for a signal, as ``pending`` lists it: its ``id``, which ``signal`` + takes to resolve this wait and no other; the ``name`` of the signal it + waits for (such as ``approval``); and ``since``, when it started, as an + ISO date. + """ + + input: Any + """What it waits for, as the model said it (such as the refund to approve), if the tool takes an input.""" + + +class Signal(TypedDict, total=False): + """Which wait a signal resolves, and with what.""" + + id: str + """The wait of this id (from ``pending``).""" + name: str + """Else the oldest wait for the signal of this name.""" + value: Any + """Its value: what the wait answers the model with.""" + keep: bool + """With a ``name`` and no one waiting for it, keep the signal for the next wait for it (by default it is dropped).""" + + +class Resolved(TypedDict): + """What ``signal`` did: the id of the wait it resolved, or None if none waited (and it was dropped, or kept).""" + + resolved: str | None + + +class Notifications(TypedDict): + """The revision of an agent's conversation: each change of what a client reads increases it.""" + + revision: int + + +class _Watch(TypedDict): + revision: int + + +class Watch(_Watch, total=False): + """What ``watch`` waits for: a revision newer than ``revision``, for at most ``timeoutSeconds`` (default 50, at most 60).""" + + timeoutSeconds: float + + +class Status(TypedDict): + """What an agent is doing: the message it answers, if any, and what its tools wait for, oldest first.""" + + answering: Turn | None + waiting: list[Waiting] + + +class Stopped(TypedDict): + """What ``stop`` did: the invocation of the message it stopped, or None if none was answered.""" + + stopped: str | None diff --git a/python/restate/agents/_signals.py b/python/restate/agents/_signals.py new file mode 100644 index 00000000..27ebe837 --- /dev/null +++ b/python/restate/agents/_signals.py @@ -0,0 +1,392 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +Signals to agents: a virtual object per agent, keyed ``:`` (such as +``assistant:alice``), whose ``wait`` returns once ``resolve`` gives the wait its +value; the message the agent answers; and the revision of the agent's +conversation, which its ``watch`` waits on. As sdk-sessions' ``signals.ts``, +with the same handlers. + +An agent's own handler cannot wait on an awakeable (it has no journal to replay +one from), so it calls ``wait``, which can: a journaled handler of this object, +which waits without holding a process. +""" + +from __future__ import annotations + +import typing +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Any, TypedDict + +import restate +from restate import ObjectContext, ObjectSharedContext, TerminalError +from restate.agents._api import Notifications, Resolved, Signal, Status, Turn, Waiting, Watch +from restate.sessions import SessionContext, ViewContext +from restate.sessions._link import Link +from restate.sessions.internal import definition_of, needing + +T = typing.TypeVar("T") +D = typing.TypeVar("D") + +signals = restate.VirtualObject("signals") +""" +The signals of every agent, keyed ``:`` (``signals_of``). +``sessions.app`` serves it beside the agents; its ingress-private handlers are +called by the agents' own handlers. +""" + +WATCH = timedelta(hours=1) +"""How long one ``awaitTurn`` waits for its turn to end, before a new one takes over.""" + + +class _WaitRequest(TypedDict): + id: str + name: str + + +class WaitRequest(_WaitRequest, total=False): + """A tool's wait: its id, the signal it waits for, and what for (the model's input), if anything.""" + + input: Any + + +class _WaitingOn(Waiting): + """A wait, with its awakeable.""" + + awakeable: str + + +class _Watching(TypedDict): + """A ``watch`` waiting for a revision newer than ``revision``, by its awakeable.""" + + revision: int + awakeable: str + + +Context = ObjectContext | ObjectSharedContext + + +async def _get(ctx: Context, key: str, default: T) -> T: + """The state ``key``, or ``default`` while it has none.""" + value = await ctx.get(key, type_hint=object) + return typing.cast(T, value) if value is not None else default + + +def _public(wait: _WaitingOn) -> Waiting: + """A wait, as ``pending`` lists it: without its awakeable.""" + shown: Waiting = {"id": wait["id"], "name": wait["name"], "since": wait["since"]} + if "input" in wait: + shown["input"] = wait["input"] + return shown + + +def iso(seconds: float) -> str: + """``seconds`` since the epoch as an ISO date, as JavaScript's ``toISOString`` gives it.""" + moment = datetime.fromtimestamp(seconds, tz=timezone.utc) + return moment.isoformat(timespec="milliseconds").replace("+00:00", "Z") + + +async def _publish(ctx: ObjectContext) -> Notifications: + """Increase the revision, and answer each ``watch`` waiting for it: the new revision.""" + revision = await _get(ctx, "revision", 0) + 1 + ctx.set("revision", revision) + notified: Notifications = {"revision": revision} + for watching in await _get(ctx, "watching", typing.cast(list[_Watching], [])): + ctx.resolve_awakeable(watching["awakeable"], notified) + ctx.clear("watching") + return notified + + +async def _closed(ctx: ObjectContext, turn: str) -> None: + """ + ``turn`` (an invocation id) is over: forget it, if the agent answers it, + and its tools' waits, which no one waits for any more, ending their + ``wait``s (a turn that ends as it should has none left, but a killed one + leaves them). It publishes the change; a close again changes nothing, and + publishes nothing. + """ + answering = await _get(ctx, "turn", typing.cast(Turn | None, None)) + is_answered = answering is not None and answering["invocationId"] == turn + if is_answered: + ctx.clear("turn") + waiting = await _get(ctx, "waiting", typing.cast(list[_WaitingOn], [])) + # A wait's id starts with the invocation id of the turn whose tool waits. + over = [wait for wait in waiting if wait["id"].startswith(f"{turn}:")] + for wait in over: + ctx.reject_awakeable(wait["awakeable"], "its turn is over") + if over: + ctx.set("waiting", [wait for wait in waiting if not wait["id"].startswith(f"{turn}:")]) + if is_answered or over: + await _publish(ctx) + + +@signals.handler(kind="shared") +async def wait(ctx: ObjectSharedContext, request: WaitRequest) -> Any: + """Wait for the signal ``name``, as wait ``id`` (for ``input``), and return its value.""" + awakeable, value = ctx.awakeable(type_hint=object) + since = iso(await ctx.time()) + waiting: _WaitingOn = {"id": request["id"], "name": request["name"], "since": since, "awakeable": awakeable} + if "input" in request: + waiting["input"] = request["input"] + ctx.object_send(listen, ctx.key(), waiting) + try: + return await value + except TerminalError as error: + # The agent's invocation was cancelled: no one waits any more. + if error.status_code == 409: + ctx.object_send(forget, ctx.key(), request["id"]) + raise + + +@signals.handler() +async def resolve(ctx: ObjectContext, signal: Signal) -> Resolved: + """Resolve the wait ``id``, or else the oldest wait for ``name``, with ``value``.""" + wait_id, name = signal.get("id"), signal.get("name") + if wait_id is None and name is None: + raise TerminalError("a signal needs an id or a name") + value = signal.get("value") + waiting = await _get(ctx, "waiting", typing.cast(list[_WaitingOn], [])) + found = next( + (at for at, wait in enumerate(waiting) if (wait["name"] == name if wait_id is None else wait["id"] == wait_id)), + None, + ) + if found is None: + if signal.get("keep", False) and name is not None and wait_id is None: + kept = await _get(ctx, "kept", typing.cast(dict[str, list[Any]], {})) + ctx.set("kept", {**kept, name: [*kept.get(name, []), value]}) + return {"resolved": None} + resolved = waiting.pop(found) + ctx.resolve_awakeable(resolved["awakeable"], value) + ctx.set("waiting", waiting) + await _publish(ctx) + return {"resolved": resolved["id"]} + + +@signals.handler(kind="shared") +async def pending(ctx: ObjectSharedContext) -> list[Waiting]: + """The waits, oldest first.""" + return [_public(wait) for wait in await _get(ctx, "waiting", typing.cast(list[_WaitingOn], []))] + + +@signals.handler(ingress_private=True) +async def listen(ctx: ObjectContext, waiting: _WaitingOn) -> None: + """A wait: resolved now with a signal kept for it, or else when one comes.""" + kept = await _get(ctx, "kept", typing.cast(dict[str, list[Any]], {})) + values = kept.get(waiting["name"], []) + if values: + ctx.resolve_awakeable(waiting["awakeable"], values[0]) + ctx.set("kept", {**kept, waiting["name"]: values[1:]}) + return + ctx.set("waiting", [*await _get(ctx, "waiting", typing.cast(list[_WaitingOn], [])), waiting]) + await _publish(ctx) + + +@signals.handler(ingress_private=True) +async def start(ctx: ObjectContext, turn: Turn) -> None: + """ + The agent's ``chat`` started answering ``turn``: the same invocation again + (a retry) keeps when it started. A turn before it that never closed itself + (its invocation was killed) closes now. ``awaitTurn`` watches the new turn, + so that it closes even if no turn comes after it. + """ + before = await _get(ctx, "turn", typing.cast(Turn | None, None)) + if before is not None and before["invocationId"] == turn["invocationId"]: + return + if before is not None: + await _closed(ctx, before["invocationId"]) + ctx.set("turn", turn) + ctx.object_send(await_turn, ctx.key(), turn["invocationId"]) + + +@signals.handler( + name="awaitTurn", + kind="shared", + ingress_private=True, + # It leaves nothing behind once done: a turn has one at a time. + journal_retention=timedelta(0), + idempotency_retention=timedelta(0), +) +async def await_turn(ctx: ObjectSharedContext, turn: str) -> None: + """ + Wait for ``turn`` to end, without holding a process, and close it if it + failed: its waits end, and ``status`` no longer says the agent answers it. + A turn that never closed itself failed (killed, out of retries, lost); one + that ended as it should, or was cancelled, closed itself already. + + It waits ``WATCH`` at most: a turn still answered then gets a new watch, + which runs on the deployment of the day. A cancelled watch (an + operator's) closes nothing. + """ + attached = ctx.attach_invocation(turn, type_hint=object) + # This watch's own cancellation raises here, and closes nothing; the turn's failure does not. + completed, _ = await restate.wait_completed(attached, ctx.sleep(WATCH)) + if attached in completed: + try: + await attached + except TerminalError: + ctx.object_send(close, ctx.key(), turn) + return + answering = await _get(ctx, "turn", typing.cast(Turn | None, None)) + if answering is not None and answering["invocationId"] == turn: + ctx.object_send(await_turn, ctx.key(), turn) + + +@signals.handler(kind="shared") +async def status(ctx: ObjectSharedContext) -> Status: + """What the agent does: the message it answers, if any, and the waits of its tools.""" + return { + "answering": await _get(ctx, "turn", typing.cast(Turn | None, None)), + "waiting": [_public(wait) for wait in await _get(ctx, "waiting", typing.cast(list[_WaitingOn], []))], + } + + +@signals.handler(kind="shared") +async def watch(ctx: ObjectSharedContext, request: Watch) -> Notifications: + """ + The revision, once there is one newer than ``revision``, or at the timeout + (50 s by default, at most 60 s), whichever comes first. It waits in + Restate, holding no process; a caller that retries it with the same + idempotency key joins the same wait. + """ + awakeable, changed = ctx.awakeable(type_hint=object) + watching: _Watching = {"revision": request["revision"], "awakeable": awakeable} + ctx.object_send(subscribe, ctx.key(), watching) + seconds = min(max(float(request.get("timeoutSeconds", 50)), 0.0), 60.0) + which, value = await restate.select(changed=changed, timeout=ctx.sleep(timedelta(seconds=seconds))) + if which == "changed": + return typing.cast(Notifications, value) + ctx.object_send(unsubscribe, ctx.key(), awakeable) + return {"revision": await _get(ctx, "revision", 0)} + + +@signals.handler(ingress_private=True) +async def close(ctx: ObjectContext, turn: str) -> None: + """The message ``turn`` (an invocation id) is answered, or its turn ended otherwise: see ``_closed``.""" + await _closed(ctx, turn) + + +@signals.handler(ingress_private=True) +async def reset(ctx: ObjectContext) -> None: + """The agent forgot its conversation (``clear``): forget its waits and kept signals, and publish the change.""" + ctx.clear("waiting") + ctx.clear("kept") + await _publish(ctx) + + +@signals.handler(ingress_private=True) +async def forget(ctx: ObjectContext, wait_id: str) -> None: + """Forget wait ``wait_id``, whose agent no longer waits.""" + waiting = await _get(ctx, "waiting", typing.cast(list[_WaitingOn], [])) + ctx.set("waiting", [wait for wait in waiting if wait["id"] != wait_id]) + await _publish(ctx) + + +@signals.handler(ingress_private=True) +async def publish(ctx: ObjectContext) -> Notifications: + """The conversation changed: increase its revision, and answer the ``watch``es.""" + return await _publish(ctx) + + +@signals.handler(kind="shared") +async def notifications(ctx: ObjectSharedContext) -> Notifications: + """The conversation's revision.""" + return {"revision": await _get(ctx, "revision", 0)} + + +@signals.handler(ingress_private=True) +async def subscribe(ctx: ObjectContext, watching: _Watching) -> None: + """Answer ``awakeable`` once there is a revision newer than ``revision``: at once, if there is.""" + revision = await _get(ctx, "revision", 0) + if revision > watching["revision"]: + ctx.resolve_awakeable(watching["awakeable"], {"revision": revision}) + return + ctx.set("watching", [*await _get(ctx, "watching", typing.cast(list[_Watching], [])), watching]) + + +@signals.handler(ingress_private=True) +async def unsubscribe(ctx: ObjectContext, awakeable: str) -> None: + """Forget the ``watch`` of ``awakeable``, which no longer waits.""" + watching = await _get(ctx, "watching", typing.cast(list[_Watching], [])) + ctx.set("watching", [one for one in watching if one["awakeable"] != awakeable]) + + +# --- An agent's companion ------------------------------------------------------------------------- + + +def key_of(service: str, key: str) -> str: + """The key of the signals of agent ``service``'s key ``key``: ``:``.""" + return f"{service}:{key}" + + +def signals_of(ctx: SessionContext | ViewContext) -> str: + """The key of the signals of the agent of ``ctx``.""" + definition = definition_of(ctx) + if definition is None: + raise TypeError("signals_of takes the context of an agent's handler or view") + return key_of(definition.name, ctx.key) + + +@dataclass(frozen=True) +class Companion: + """ + An agent's companion, for one of its keys: its ``signals`` object, which + knows what the agent is doing (the message it answers, what its tools wait + for), and how to reach it. + """ + + key: str + """The key of the agent's key there: ``:``.""" + agent: str + """The agent's name.""" + agent_key: str + """The agent's key.""" + + +_COMPANION: Link[Callable[[str], Companion]] = Link("restate-agents.companion") +"""How an agent reaches its companion, for one of its keys.""" + + +def with_companion(definition: D) -> D: + """``definition`` (an agent), pointing to its companion; ``sessions.app`` serves its ``signals`` object beside it.""" + name = typing.cast(str, getattr(definition, "name")) # noqa: B009 (a definition of the SDK's) + _COMPANION.on(definition, lambda key: Companion(key=key_of(name, key), agent=name, agent_key=key)) + return needing(definition, lambda: [signals]) + + +def companion_of(agent: object, key: str) -> Companion: + """The companion of ``agent``'s key ``key``: for anyone holding an agent to reach what it is doing.""" + companion = _COMPANION.of(agent) + if companion is None: + raise TypeError(f"{agent!r} is not an agent of restate.agents") + return companion(key) + + +def has_companion(agent: object) -> bool: + """Whether ``agent`` is an object of restate.agents, with a companion.""" + return _COMPANION.of(agent) is not None + + +def deliver(ctx: SessionContext, agent: object, key: str, text: str, *, idempotency_key: str | None = None) -> None: + """ + Deliver ``text`` to ``agent``'s key ``key``, from a session handler: a + message of its own (as ``chat``), which it answers after the messages + before it. It is sent, not awaited: ``idempotency_key`` makes it + delivered once:: + + deliver(ctx, triage, "ENG-1", "The issue was closed", idempotency_key=delivery) + """ + companion = companion_of(agent, key) # an agent of restate.agents, with a `chat` handler + ctx.send.generic(companion.agent, "chat", text, key=key, idempotency_key=idempotency_key) + + +__all__: Sequence[str] = ["Companion", "companion_of", "deliver", "signals"] diff --git a/tests/agents_signals.py b/tests/agents_signals.py new file mode 100644 index 00000000..0d5999d1 --- /dev/null +++ b/tests/agents_signals.py @@ -0,0 +1,156 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""The agents' companion object, signals, against a real Restate server.""" + +# pylint: disable=missing-function-docstring + +import asyncio +import typing +import uuid +from collections.abc import Awaitable, Callable + +import httpx +import pytest + +import restate +from restate import HarnessEnvironment +from restate.agents import _signals +from restate.agents._api import Status, Turn + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +pytestmark = [pytest.mark.anyio] + +# The agent's side: its handlers call the companion's internal handlers, which the ingress refuses. +driver = restate.Service("driver") + + +@driver.handler() +async def start_turn(ctx: restate.Context, request: dict[str, str]) -> None: + turn: Turn = {"invocationId": request["invocation"], "since": "2026-11-04T00:00:00.000Z"} + await ctx.object_call(_signals.start, key=request["key"], arg=turn) + + +@driver.handler() +async def close_turn(ctx: restate.Context, request: dict[str, str]) -> None: + await ctx.object_call(_signals.close, key=request["key"], arg=request["invocation"]) + + +failing = restate.Service("failing") + + +@failing.handler() +async def fail(ctx: restate.Context) -> None: + raise restate.TerminalError("this turn was killed") + + +@pytest.fixture(scope="session") +async def harness(): + async with restate.create_test_harness( + restate.app([_signals.signals, driver, failing]), restate_image="ghcr.io/restatedev/restate:latest" + ) as harness: + yield harness + + +async def attach(harness: HarnessEnvironment, invocation_id: str) -> httpx.Response: + async with httpx.AsyncClient(timeout=60) as ingress: + return await ingress.get(f"{harness.ingress_url}/restate/invocation/{invocation_id}/attach") + + +async def until(check: Callable[[], Awaitable[bool]]) -> None: + for _ in range(100): + if await check(): + return + await asyncio.sleep(0.05) + raise AssertionError("never") + + +def key() -> str: + return f"assistant:{uuid.uuid4()}" + + +async def test_a_wait_returns_the_value_of_the_signal_that_resolves_it(harness: HarnessEnvironment): + k = key() + waiting = await harness.client.object_send(_signals.wait, key=k, arg={"id": "inv:1", "name": "approval"}) + + async def listed() -> bool: + return len(await harness.client.object_call(_signals.pending, key=k, arg=None)) == 1 + + await until(listed) + [pending] = await harness.client.object_call(_signals.pending, key=k, arg=None) + assert pending["id"] == "inv:1" and pending["name"] == "approval" and pending["since"].endswith("Z") + resolved = await harness.client.object_call(_signals.resolve, key=k, arg={"id": "inv:1", "value": "approved"}) + assert resolved == {"resolved": "inv:1"} + answer = await attach(harness, waiting.invocation_id) + assert answer.json() == "approved" + assert await harness.client.object_call(_signals.pending, key=k, arg=None) == [] + + +async def test_a_signal_resolves_the_oldest_wait_of_its_name_or_is_dropped_or_kept(harness: HarnessEnvironment): + k = key() + # No one waits: dropped, unless kept for the next wait. + assert await harness.client.object_call(_signals.resolve, key=k, arg={"name": "go", "value": 1}) == { + "resolved": None + } + await harness.client.object_call(_signals.resolve, key=k, arg={"name": "go", "value": 2, "keep": True}) + kept = await harness.client.object_send(_signals.wait, key=k, arg={"id": "inv:2", "name": "go"}) + assert (await attach(harness, kept.invocation_id)).json() == 2 + + first = await harness.client.object_send(_signals.wait, key=k, arg={"id": "inv:3", "name": "go"}) + second = await harness.client.object_send(_signals.wait, key=k, arg={"id": "inv:4", "name": "go"}) + + async def both() -> bool: + return len(await harness.client.object_call(_signals.pending, key=k, arg=None)) == 2 + + await until(both) + assert await harness.client.object_call(_signals.resolve, key=k, arg={"name": "go", "value": "a"}) in ( + {"resolved": "inv:3"}, + {"resolved": "inv:4"}, + ) + await harness.client.object_call(_signals.resolve, key=k, arg={"name": "go", "value": "b"}) + answers = {(await attach(harness, w.invocation_id)).json() for w in (first, second)} + assert answers == {"a", "b"} + + +async def test_watch_returns_a_newer_revision_or_the_current_one_at_its_timeout(harness: HarnessEnvironment): + k = key() + now = (await harness.client.object_call(_signals.notifications, key=k, arg=None))["revision"] + timed_out = await harness.client.object_call(_signals.watch, key=k, arg={"revision": now, "timeoutSeconds": 0.2}) + assert timed_out == {"revision": now} + watching = asyncio.ensure_future( + harness.client.object_call(_signals.watch, key=k, arg={"revision": now, "timeoutSeconds": 30}) + ) + await asyncio.sleep(0.3) + # A wait that starts publishes a change. + await harness.client.object_send(_signals.wait, key=k, arg={"id": "inv:5", "name": "x"}) + assert (await watching)["revision"] > now + + +async def test_a_turn_is_answered_until_it_closes_and_a_failed_one_is_closed_for_it(harness: HarnessEnvironment): + k = key() + await harness.client.service_call(start_turn, arg={"key": k, "invocation": "inv_running"}) + status = typing.cast(Status, await harness.client.object_call(_signals.status, key=k, arg=None)) + assert status["answering"] is not None and status["answering"]["invocationId"] == "inv_running" + await harness.client.service_call(close_turn, arg={"key": k, "invocation": "inv_running"}) + assert (await harness.client.object_call(_signals.status, key=k, arg=None))["answering"] is None + + # A turn whose invocation failed without closing itself: awaitTurn closes it. + killed = await harness.client.service_send(fail, arg=None) + await harness.client.service_call(start_turn, arg={"key": k, "invocation": killed.invocation_id}) + + async def closed() -> bool: + return (await harness.client.object_call(_signals.status, key=k, arg=None))["answering"] is None + + await until(closed) From f3a06aeadc37bdd1711ac09182215d5b57619bde Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 16:13:06 +0100 Subject: [PATCH 20/45] Type the companion tests' arguments for mypy mypy takes a dict literal passed to the client for a dict, not the handler's TypedDict: the waits, signals and watches are built by typed helpers. (The previous commit went in with these 15 mypy errors.) Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- tests/agents_signals.py | 47 ++++++++++++++++++++++++++--------------- 1 file changed, 30 insertions(+), 17 deletions(-) diff --git a/tests/agents_signals.py b/tests/agents_signals.py index 0d5999d1..df24c1b4 100644 --- a/tests/agents_signals.py +++ b/tests/agents_signals.py @@ -23,7 +23,8 @@ import restate from restate import HarnessEnvironment from restate.agents import _signals -from restate.agents._api import Status, Turn +from restate.agents._api import Signal, Status, Turn, Watch +from restate.agents._signals import WaitRequest @pytest.fixture(scope="session") @@ -81,9 +82,21 @@ def key() -> str: return f"assistant:{uuid.uuid4()}" +def waiting(wait_id: str, name: str) -> WaitRequest: + return {"id": wait_id, "name": name} + + +def watching(revision: int, timeout: float) -> Watch: + return {"revision": revision, "timeoutSeconds": timeout} + + +def signal(**fields: typing.Any) -> Signal: + return typing.cast(Signal, fields) + + async def test_a_wait_returns_the_value_of_the_signal_that_resolves_it(harness: HarnessEnvironment): k = key() - waiting = await harness.client.object_send(_signals.wait, key=k, arg={"id": "inv:1", "name": "approval"}) + sent = await harness.client.object_send(_signals.wait, key=k, arg=waiting("inv:1", "approval")) async def listed() -> bool: return len(await harness.client.object_call(_signals.pending, key=k, arg=None)) == 1 @@ -91,9 +104,11 @@ async def listed() -> bool: await until(listed) [pending] = await harness.client.object_call(_signals.pending, key=k, arg=None) assert pending["id"] == "inv:1" and pending["name"] == "approval" and pending["since"].endswith("Z") - resolved = await harness.client.object_call(_signals.resolve, key=k, arg={"id": "inv:1", "value": "approved"}) + resolved = await harness.client.object_call( + _signals.resolve, key=k, arg=signal(**{"id": "inv:1", "value": "approved"}) + ) assert resolved == {"resolved": "inv:1"} - answer = await attach(harness, waiting.invocation_id) + answer = await attach(harness, sent.invocation_id) assert answer.json() == "approved" assert await harness.client.object_call(_signals.pending, key=k, arg=None) == [] @@ -101,25 +116,25 @@ async def listed() -> bool: async def test_a_signal_resolves_the_oldest_wait_of_its_name_or_is_dropped_or_kept(harness: HarnessEnvironment): k = key() # No one waits: dropped, unless kept for the next wait. - assert await harness.client.object_call(_signals.resolve, key=k, arg={"name": "go", "value": 1}) == { + assert await harness.client.object_call(_signals.resolve, key=k, arg=signal(**{"name": "go", "value": 1})) == { "resolved": None } - await harness.client.object_call(_signals.resolve, key=k, arg={"name": "go", "value": 2, "keep": True}) - kept = await harness.client.object_send(_signals.wait, key=k, arg={"id": "inv:2", "name": "go"}) + await harness.client.object_call(_signals.resolve, key=k, arg=signal(**{"name": "go", "value": 2, "keep": True})) + kept = await harness.client.object_send(_signals.wait, key=k, arg=waiting("inv:2", "go")) assert (await attach(harness, kept.invocation_id)).json() == 2 - first = await harness.client.object_send(_signals.wait, key=k, arg={"id": "inv:3", "name": "go"}) - second = await harness.client.object_send(_signals.wait, key=k, arg={"id": "inv:4", "name": "go"}) + first = await harness.client.object_send(_signals.wait, key=k, arg=waiting("inv:3", "go")) + second = await harness.client.object_send(_signals.wait, key=k, arg=waiting("inv:4", "go")) async def both() -> bool: return len(await harness.client.object_call(_signals.pending, key=k, arg=None)) == 2 await until(both) - assert await harness.client.object_call(_signals.resolve, key=k, arg={"name": "go", "value": "a"}) in ( + assert await harness.client.object_call(_signals.resolve, key=k, arg=signal(**{"name": "go", "value": "a"})) in ( {"resolved": "inv:3"}, {"resolved": "inv:4"}, ) - await harness.client.object_call(_signals.resolve, key=k, arg={"name": "go", "value": "b"}) + await harness.client.object_call(_signals.resolve, key=k, arg=signal(**{"name": "go", "value": "b"})) answers = {(await attach(harness, w.invocation_id)).json() for w in (first, second)} assert answers == {"a", "b"} @@ -127,15 +142,13 @@ async def both() -> bool: async def test_watch_returns_a_newer_revision_or_the_current_one_at_its_timeout(harness: HarnessEnvironment): k = key() now = (await harness.client.object_call(_signals.notifications, key=k, arg=None))["revision"] - timed_out = await harness.client.object_call(_signals.watch, key=k, arg={"revision": now, "timeoutSeconds": 0.2}) + timed_out = await harness.client.object_call(_signals.watch, key=k, arg=watching(now, 0.2)) assert timed_out == {"revision": now} - watching = asyncio.ensure_future( - harness.client.object_call(_signals.watch, key=k, arg={"revision": now, "timeoutSeconds": 30}) - ) + watcher = asyncio.ensure_future(harness.client.object_call(_signals.watch, key=k, arg=watching(now, 30))) await asyncio.sleep(0.3) # A wait that starts publishes a change. - await harness.client.object_send(_signals.wait, key=k, arg={"id": "inv:5", "name": "x"}) - assert (await watching)["revision"] > now + await harness.client.object_send(_signals.wait, key=k, arg=waiting("inv:5", "x")) + assert (await watcher)["revision"] > now async def test_a_turn_is_answered_until_it_closes_and_a_failed_one_is_closed_for_it(harness: HarnessEnvironment): From 973b1a8141b047a1750f7c3615d4f71cc0ff6e76 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 16:21:13 +0100 Subject: [PATCH 21/45] restate.agents: AgentObject and its chat turns The third part of the agents' core, from sdk-sessions' object.ts and common.ts: - AgentObject, a SessionObject set up for agents (weeks of retries, its companion served beside it), with the views every agent has: signal, pending, status, stop, notifications and watch, each asking the companion. chat(answer) makes it an agent: the chat handler (a message in, as text or a JSON string, the answer out), the messages view (its log), clear (stop, then the internal erase empties the database, and the log's ids go on), and sync, which a follower reads. - The turn (_turn.py): it starts at the companion, publishes its steps, keeps the framework's transcript in the log (transcribing), and closes; Restate's cancellation (explicit) and the attempt's end stop it (Stop, as AbortSignal, with guard() for a framework's request), and a model provider's refusal ends it with 424 (_common.py), in any handler of an agent. - The tool calls' bookkeeping (_calls.py), which publishing settles. Found by its tests, a fix in sessions: ctx.cancel called a cancel_invocation the SDK's host lacked (typed Any, only the tests' fake host had it); it is now in SqlHost, implemented, and tested from a view (a Stop button). SessionObject gains a protected hook, _attempt_fn, for a subclass to wrap what an attempt runs. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/agents/_calls.py | 163 +++++++++++++++ python/restate/agents/_common.py | 210 +++++++++++++++++++ python/restate/agents/_object.py | 311 ++++++++++++++++++++++++++++ python/restate/agents/_turn.py | 301 +++++++++++++++++++++++++++ python/restate/sessions/_object.py | 20 +- python/restate/sessions/_session.py | 4 + tests/agents_object.py | 173 ++++++++++++++++ tests/sessions_harness.py | 26 +++ 8 files changed, 1206 insertions(+), 2 deletions(-) create mode 100644 python/restate/agents/_calls.py create mode 100644 python/restate/agents/_common.py create mode 100644 python/restate/agents/_object.py create mode 100644 python/restate/agents/_turn.py create mode 100644 tests/agents_object.py diff --git a/python/restate/agents/_calls.py b/python/restate/agents/_calls.py new file mode 100644 index 00000000..b3e6952c --- /dev/null +++ b/python/restate/agents/_calls.py @@ -0,0 +1,163 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +The Restate calls of an invocation's tools (``call``, ``signal``, +``subagent``...), kept in the object's database (``restate_agent_calls``) so +that a call is made once, however often its tool call runs. As sdk-sessions' +``calls.ts``. + +A crash or a suspension restarts the handler. A framework then runs a tool call +that was cut off again, with its id; or tells the model it was interrupted, and +the model makes it again with a new id: made with the same tool and input, it +is the call an earlier attempt made, so it takes that call's id, and so the +same idempotency key: the same Restate call, which it waits for again, or +whose answer it gets. + +A call is kept until the framework's transcript has its result under the last +id it was made with (``settle``, at each publish, and before each tool call). +An error there is not its result, unless the call answered with that error. +""" + +from __future__ import annotations + +import json +import typing +import uuid +from collections.abc import Callable +from dataclasses import dataclass, field + +from restate.sessions import SessionContext +from restate.sessions.internal import Schema, ScopeKey + +SCHEMA = Schema( + "restate-agents.calls", + [ + """CREATE TABLE restate_agent_calls ( + row INTEGER PRIMARY KEY CHECK (row = 0), + invocation TEXT NOT NULL, + calls TEXT NOT NULL + )""" + ], +) + + +class _Made(typing.TypedDict): + """A tool's call: the tool, its input as JSON, the ids it was made with (the first is its own), and its attempt.""" + + tool: str + input: str + ids: list[str] + attempt: str + + +class _Call(_Made, total=False): + """A tool's call, and whether it answered.""" + + done: bool + """Answered, until the transcript has its result.""" + failed: str + """It answered with this error (``Error: ``), a terminal one: its result, as the transcript has it.""" + + +@dataclass +class _Calls: + """The calls an invocation's tools made that the transcript has no result of yet.""" + + invocation: str + calls: list[_Call] = field(default_factory=list[_Call]) + + +async def _store(ctx: SessionContext, state: _Calls) -> None: + """Store ``state``, committed before the call it records is made.""" + await ctx.db.execute( + "INSERT OR REPLACE INTO restate_agent_calls (row, invocation, calls) VALUES (0, ?, ?)", + [state.invocation, json.dumps(state.calls)], + ) + + +async def _load(ctx: SessionContext, invocation: str) -> _Calls: + """The calls of ``invocation``, as stored: none if the row is another invocation's.""" + await SCHEMA.ensure(ctx.db) + row = await ctx.db.get("SELECT invocation, calls FROM restate_agent_calls") + if row is None or row["invocation"] != invocation: + return _Calls(invocation=invocation) + return _Calls(invocation=invocation, calls=json.loads(typing.cast(str, row["calls"]))) + + +@dataclass +class _Attempt: + id: str + calls: _Calls | None = None + + +_ATTEMPT: ScopeKey[_Attempt] = ScopeKey("the calls of an attempt") + + +async def _calls_of(ctx: SessionContext) -> tuple[_Attempt, _Calls]: + """The attempt of ``ctx``, and the calls of its invocation.""" + attempt = _ATTEMPT.of(ctx, lambda: _Attempt(id=str(uuid.uuid4()))) + if attempt.calls is None: + attempt.calls = await _load(ctx, ctx.request.id) + return attempt, attempt.calls + + +@dataclass +class Call: + """The call a tool call makes: its id (an earlier attempt's, maybe), and how to say it answered.""" + + id: str + _ctx: SessionContext + _state: _Calls + _call: _Call + + async def answered(self, failed: str | None = None) -> None: + """It answered, with ``failed`` if a terminal error answered it: kept until the transcript has its result.""" + self._call["done"] = True + if failed is not None: + self._call["failed"] = failed + await _store(self._ctx, self._state) + + +async def call_of(ctx: SessionContext, tool: str, input: object, call_id: str) -> Call: + """ + The call of tool ``tool`` with ``input`` that tool call ``call_id`` makes: + the one it made before (a run again), or the one an earlier attempt made + with the same tool and input (made again under a new id), or else a new one. + """ + attempt, state = await _calls_of(ctx) + encoded = json.dumps(input, sort_keys=True) + call = next((c for c in state.calls if call_id in c["ids"]), None) or next( + (c for c in state.calls if c["attempt"] != attempt.id and c["tool"] == tool and c["input"] == encoded), + None, + ) + if call is not None: + call["attempt"] = attempt.id + if call_id not in call["ids"]: + call["ids"].append(call_id) + else: + call = _Call(tool=tool, input=encoded, ids=[call_id], attempt=attempt.id) + state.calls.append(call) + await _store(ctx, state) + return Call(id=call["ids"][0], _ctx=ctx, _state=state, _call=call) + + +async def settle(ctx: SessionContext, stored: Callable[[str, str | None], bool]) -> None: + """ + Forget the calls that answered whose result the framework stored: those of + which ``stored(id, failed)`` is true for the last id they were made with, + ``failed`` the error that answered the call, if one did. + """ + _, state = await _calls_of(ctx) + kept = [c for c in state.calls if not c.get("done") or not stored(c["ids"][-1], c.get("failed"))] + if len(kept) == len(state.calls): + return + state.calls = kept + await _store(ctx, state) diff --git a/python/restate/agents/_common.py b/python/restate/agents/_common.py new file mode 100644 index 00000000..f3d70552 --- /dev/null +++ b/python/restate/agents/_common.py @@ -0,0 +1,210 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +What the frameworks share: their options (``Configure``, the model), the ids of +an invocation's prompts, a model provider's refusal, and the ``chat`` message. +As sdk-sessions' ``common.ts``. +""" + +from __future__ import annotations + +import inspect +import json +import os +import re +import typing +import weakref +from collections.abc import Awaitable, Callable + +from restate.exceptions import TerminalError +from restate.serde import Serde +from restate.sessions import SessionContext +from restate.sessions.internal import ScopeKey + +T = typing.TypeVar("T") + +Configure: typing.TypeAlias = T | Callable[[SessionContext], T | Awaitable[T]] +""" +A value, or a function of the invocation's context that gives it: the same for +every message, or each user's (from the object's database, say):: + + claude.agent("assistant", model="anthropic/claude-sonnet-5-5") + claude.agent("assistant", model=lambda ctx: user_model(ctx)) +""" + + +async def configured(value: Configure[T], ctx: SessionContext) -> T: + """``value`` for the invocation of ``ctx``.""" + if not callable(value): + return value + made = typing.cast(Callable[[SessionContext], T | Awaitable[T]], value)(ctx) + return await made if inspect.isawaitable(made) else typing.cast(T, made) + + +_PROMPTS: ScopeKey[int] = ScopeKey("next_prompt_id") + + +def next_prompt_id(ctx: SessionContext) -> str: + """ + The id of the next prompt of the invocation of ``ctx``: the invocation's + id, then ``:2``, ``:3``... A retried attempt makes its prompts in + the same order, so each gets the same id again, and its answer. + """ + n = (_PROMPTS.get(ctx) or 0) + 1 + _PROMPTS.set(ctx, n) + return ctx.request.id if n == 1 else f"{ctx.request.id}:{n}" + + +DEFAULT_MODEL = "openai/gpt-6-sol" +"""The default model of the frameworks that take any provider's, with OPENAI_API_KEY.""" + +CLAUDE_MODEL = "anthropic/claude-sonnet-5-5" +"""The default model of Claude agents, with ANTHROPIC_API_KEY.""" + + +def model_of(model: str | None, fallback: str = DEFAULT_MODEL) -> str: + """The model, as ``provider/model-id``: ``model``, or else MODEL, or else the framework's default.""" + return model or os.environ.get("MODEL") or fallback + + +def provider_of(model: str) -> tuple[str, str]: + """``provider/model-id`` as its provider and model id, split at the first ``/``.""" + provider, slash, model_id = model.partition("/") + return (provider, model_id) if slash else ("", model) + + +REFUSED = 424 +""" +The error code of a message the model's provider refused (a missing or wrong +API key, an unknown model): Failed Dependency, as its HTTP status. +""" + +CHECK_MODEL = "check the agent's model (its `model` option, or MODEL) and its provider's API key" +"""What to check when a model's provider refuses a request.""" + + +class ModelError(TerminalError): + """ + An agent's model that it cannot take: not ``provider/model-id``, another + provider's, or one its framework does not know. Terminal (``REFUSED``), + wherever it is raised: a retry takes the same model. + """ + + def __init__(self, message: str) -> None: + super().__init__(message, status_code=REFUSED) + + +def check_model_name(framework: str, model: str | None) -> None: + """``model``, the explicit ``model`` option of ``framework``'s agent, if it is not ``provider/model-id``.""" + if model is not None and not re.fullmatch(r"[^/\s]+/\S+", model): + raise ModelError(f'{framework}: the model "{model}" is not provider/model-id, such as {CLAUDE_MODEL}') + + +def model_id_of(model: str, provider: str) -> str: + """The model id of ``model``, for a framework with only ``provider``'s models: ``provider/model-id``, or the id alone.""" + named, model_id = provider_of(model) if "/" in model else (provider, model) + if named != provider or not model_id: + from_env = " (from MODEL)" if os.environ.get("MODEL") == model else "" + raise ModelError(f'the model "{model}"{from_env} is not {provider}\'s: this agent takes {provider}/') + return model_id + + +class MessageSerde(Serde[str]): + """A chat message: the request's body as text, or the string it is as JSON (``"Hi"``).""" + + def deserialize(self, buf: bytes) -> str | None: + text = buf.decode() + try: + value = json.loads(text) + except ValueError: + return text + return value if isinstance(value, str) else text + + def serialize(self, obj: str | None) -> bytes: + return (obj or "").encode() + + +CANCELLED = "Cancelled: the turn was stopped before it finished." +"""What a cancelled turn answers, and each tool call it cut off: the model sees it in the conversation that goes on.""" + +UNANSWERED = "Unanswered: the model's provider refused the request." +""" +What a prompt answers that the model's provider refused for good: the prompt is +over, and the model sees it in the conversation that goes on, rather than +answer it with the next message. +""" + + +class ToolError(Exception): + """A tool's error, which fails the attempt (Restate retries it): the error of tool ``tool``.""" + + def __init__(self, tool: str, cause: BaseException) -> None: + super().__init__(f"the tool {tool} failed: {cause}") + self.tool: str = tool + self.__cause__ = cause + + +_TOOL_ERRORS: ScopeKey[weakref.WeakSet[BaseException]] = ScopeKey("from_tool") +"""The errors the attempt's framework tools (not ``tool()``'s) failed with: see ``from_tool``.""" + + +def from_tool(ctx: SessionContext, error: BaseException) -> None: + """ + ``error`` is a tool's, in the attempt of ``ctx``: a framework's own tool + that failed, whose error is never taken for the model's provider refusing + (a tool that calls a model API may fail as one does). + """ + _TOOL_ERRORS.of(ctx, weakref.WeakSet[BaseException]).add(error) + + +_FOR_GOOD = frozenset({"AuthenticationError", "PermissionDeniedError", "NotFoundError", "BadRequestError"}) +"""The model SDKs' errors (Anthropic's, OpenAI's) of a request that fails again however often it is made.""" + +_REFUSED_TEXT = re.compile( + r"API key not found|no API key|Could not resolve authentication method|Provider is not configured" + r'|"code":"(AUTH|MISSING_CREDENTIAL)"', + re.IGNORECASE, +) + + +def _refused(error: BaseException) -> bool: + """Whether ``error`` is a model provider's refusal that a retry does not change: a missing key, an unknown model...""" + if isinstance(error, ToolError): + return False + if type(error).__name__ in _FOR_GOOD: + return True + return bool(_REFUSED_TEXT.search(str(error))) + + +def refusal(error: BaseException) -> TerminalError | None: + """``error`` as a terminal error (``REFUSED``), if a retry does not change it: a model the agent cannot take, or its provider's refusal.""" + if isinstance(error, ModelError): + return error + if isinstance(error, TerminalError) or not _refused(error): + return None + terminal = TerminalError(f"the model's provider refused the request: {CHECK_MODEL}. {error}", status_code=REFUSED) + terminal.__cause__ = error + return terminal + + +def refusal_of(ctx: SessionContext, error: BaseException) -> TerminalError | None: + """``refusal``, but not of an error a framework's tool failed with in the attempt of ``ctx`` (``from_tool``).""" + failed_tool = _TOOL_ERRORS.get(ctx) + if failed_tool is not None and error in failed_tool: + return None + return refusal(error) + + +def stored_answer(answer: str) -> str: + """A prompt's stored answer, as the prompt returns it: ``UNANSWERED`` fails as the prompt did then (``REFUSED``).""" + if answer != UNANSWERED: + return answer + raise TerminalError(f"the model's provider refused the request: {CHECK_MODEL}", status_code=REFUSED) diff --git a/python/restate/agents/_object.py b/python/restate/agents/_object.py new file mode 100644 index 00000000..efadd8cd --- /dev/null +++ b/python/restate/agents/_object.py @@ -0,0 +1,311 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +Session objects set up for agents: ``AgentObject``, whose ``chat`` makes it an +agent. As sdk-sessions' ``sessionObject`` of ``@restatedev/agents`` +(``object.ts``). +""" + +from __future__ import annotations + +import functools +import typing +from collections.abc import Awaitable, Callable, Sequence +from datetime import timedelta +from typing import Any, TypedDict + +from restate.agents import _log, _signals +from restate.agents._api import ( + Messages, + MessagesRequest, + Notifications, + Resolved, + Signal, + Status, + Stopped, + Waiting, + Watch, +) +from restate.agents._common import MessageSerde, refusal_of +from restate.agents._turn import AGENT, AgentInfo, chat_turn, cleared +from restate.retry_policy import InvocationRetryPolicy +from restate.sessions import SessionContext, SessionObject, Step, ViewContext +from restate.sessions.internal import reset + +T = typing.TypeVar("T") + +RETRIES = InvocationRetryPolicy( + initial_interval=timedelta(milliseconds=50), + exponentiation_factor=2, + max_interval=timedelta(seconds=10), + max_attempts=100_000, +) +""" +An agent's retries, for weeks, since each attempt resumes from the database: +soon after a crash, then twice as long each time, up to 10 seconds apart (a +provider's rate limit, say, is not asked again twice a second). +""" + +WAITS = timedelta(minutes=5) +""" +The inactivity timeout of the views that wait (``watch``, ``sync``, ``clear``): +Restate suspends one that waits past it, which restarts it, so it is longer +than the wait. +""" + + +class SyncRequest(TypedDict, total=False): + """ + What ``sync`` reads: what is new or changed since ``cursor`` (default 0, + all of it), at most ``limit`` messages (default 100); once the revision is + past ``revision``, if given, or at ``timeoutSeconds`` (default 50, at most + 60). + """ + + cursor: int + revision: int + timeoutSeconds: float + limit: int + + +class Synced(Messages): + """What ``sync`` answers: the messages since the cursor, as ``messages`` gives them, the status and the revision.""" + + revision: int + status: Status + + +def _start_after(page: Messages) -> _log.Start: + """Where the log starts after the latest message of ``page``: past its ids and cursor.""" + latest = page["messages"][0] if page["messages"] else None + return _log.Start(seq=page["cursor"], id=latest["id"] + 1 if latest else page.get("first", 0)) + + +class AgentObject(SessionObject): + """ + A Restate virtual object with a SQL database per key, ``ctx.db``, set up + for agents. Besides its own handlers and views, it has the views + ``status``, what it is doing (the message it answers, and what its tools + wait for), ``stop``, which stops that message, ``pending``, the waits of + its tools, ``signal({id, value})``, which resolves one, and + ``notifications`` and ``watch({revision})``, to follow its conversation. + + ``chat(answer)`` makes it an agent: a ``chat`` handler, a message in and + the answer out, the ``messages`` view (its conversation), ``clear``, which + forgets it (its database too), and ``sync``, which a client following it + reads. A framework's ``agent()`` makes one. + + ``schema``: tables of your own, beside the agent's (named ``restate_...``). + ``background``: it works unwatched (a subagent, say): its turn publishes at + its end only, not at each step. + """ + + def __init__( + self, + name: str, + *, + description: str | None = None, + schema: Sequence[Step] = (), + background: bool = False, + metadata: dict[str, str] | None = None, + inactivity_timeout: timedelta | None = None, + abort_timeout: timedelta | None = None, + journal_retention: timedelta | None = None, + idempotency_retention: timedelta | None = None, + ingress_private: bool | None = None, + invocation_retry_policy: InvocationRetryPolicy | None = RETRIES, + ) -> None: + super().__init__( + name, + schema=schema, + description=description, + metadata=metadata, + inactivity_timeout=inactivity_timeout, + abort_timeout=abort_timeout, + journal_retention=journal_retention, + idempotency_retention=idempotency_retention, + ingress_private=ingress_private, + invocation_retry_policy=invocation_retry_policy, + ) + self.background: bool = background + self._clear: Callable[[SessionContext], Awaitable[None]] | None = None + AGENT.on(self, AgentInfo(background=background)) + _signals.with_companion(self) + self._built_in() + + def _attempt_fn(self, fn: Callable[..., Awaitable[Any]], view: bool) -> Callable[..., Awaitable[Any]]: + """A handler's attempt: a model provider's refusal it fails with ends its invocation, as ``chat``'s does.""" + if view: + return fn + + # With fn's signature, from which the SDK takes whether it has an input. + @functools.wraps(fn) + async def answering(ctx: SessionContext, *args: Any) -> Any: + try: + return await fn(ctx, *args) + except Exception as error: + terminal = refusal_of(ctx, error) + if terminal is None: + raise + raise terminal from error + + return answering + + def _built_in(self) -> None: + """The views every agent has: each asks the companion, which knows what the object does.""" + + async def signal(ctx: ViewContext, signal: Signal) -> Resolved: + """Resolve a wait of the object's tools: by its ``id``, or the oldest of a ``name``, with ``value``.""" + return await _companion(ctx, _signals.resolve, "signal")(signal) + + async def pending(ctx: ViewContext) -> list[Waiting]: + """The waits of the object's tools, oldest first.""" + return await _companion(ctx, _signals.pending)() + + async def status(ctx: ViewContext) -> Status: + """What the object is doing: the message it answers, if any, and what its tools wait for.""" + return await _companion(ctx, _signals.status)() + + async def stop(ctx: ViewContext) -> Stopped: + """Cancel the invocation of the message being answered (``status``'s), as Restate's admin API cancels one.""" + return await _stopped(ctx, "stop") + + async def notifications(ctx: ViewContext) -> Notifications: + """The revision of the conversation: each change of what ``messages`` and ``status`` give increases it.""" + return await _companion(ctx, _signals.notifications)() + + async def watch(ctx: ViewContext, request: Watch) -> Notifications: + """ + The revision, once there is one newer than ``revision``, or at the + timeout (50 s by default, at most 60 s): a client then reads what + changed. It waits in Restate, holding no process. + """ + return await _companion(ctx, _signals.watch, "watch")(request) + + for fn in (signal, pending, status, stop, notifications): + self.view(fn) + self.view(inactivity_timeout=WAITS)(watch) + + def chat( + self, + answer: Callable[[SessionContext, str], Awaitable[T]], + *, + clear: Callable[[SessionContext], Awaitable[None]] | None = None, + ) -> None: + """ + Make this object an agent, which answers each message with ``answer``: + a framework's turn, which keeps the conversation in the log + (``transcribing``) and stops with ``stop_signal_of(ctx)``. ``clear``, + what the agent keeps beside its database, which ``clear`` empties too. + + A message comes as the request's body, text or a JSON string:: + + curl localhost:8080/assistant/alice/chat -d 'Hi' + + A key answers one message at a time; one its model's provider refused + (a missing API key, an unknown model) answers with an error (424). + """ + AGENT.on(self, AgentInfo(background=self.background, clear=clear)) + + async def chat(ctx: SessionContext, text: str) -> T: + return await chat_turn(ctx, text, answer) + + self.handler(name="chat", accept="*/*", input_serde=MessageSerde(), explicit_cancellation=True)(chat) + _followed(self) + + +def _companion( + ctx: ViewContext, handler: Callable[..., Awaitable[T]], step: str | None = None +) -> Callable[..., Awaitable[T]]: + """ + A call of the companion's of ``ctx``; with ``step``, a change, keyed by the + request and ``step``: a restarted view joins the same call. A read has no + key: asked again, it answers again. + """ + key = None if step is None else f"{ctx.request.id}:{step}" + return typing.cast(Callable[..., Awaitable[T]], ctx.rpc(handler, key=_signals.signals_of(ctx), idempotency_key=key)) + + +async def _stopped(ctx: ViewContext, step: str) -> Stopped: + """ + Cancel the invocation of the message being answered, if any, and say + which. The companion's ``status`` says which, with a key of the request and + ``step``: a restarted view cancels the same invocation again (which does + nothing), not the message that started once it was cancelled. + """ + turn = (await _companion(ctx, _signals.status, step)())["answering"] + if turn is not None: + ctx.cancel(turn["invocationId"]) + return {"stopped": turn["invocationId"] if turn is not None else None} + + +def _followed(agent: AgentObject) -> None: + """ + An agent's views beside those of every object: ``messages``, its + conversation; ``clear``, which forgets it; and ``sync``, the conversation + since ``cursor``, what the agent is doing, and the revision, in one + request. And ``erase``, ``clear``'s work, which only its views call. + """ + + async def messages(ctx: ViewContext, request: MessagesRequest | None = None) -> Messages: + """ + The conversation as clients read it, from its log: a page of it + (``{limit, before}``), or what is new or changed since a cursor + (``{cursor}``). It answers while a message is answered, with what was + stored so far. + """ + return await _log.page(ctx.db, request) + + async def erase(ctx: SessionContext) -> None: + """``clear``'s work, once the handlers before it are done: the key's database empties, and its signals'.""" + start = _start_after(await _log.page(ctx.db, {"limit": 1})) + async with ctx.db.transaction() as tx: + await reset(tx) + await _log.restart(tx, start) + await cleared(ctx) + await ctx.rpc(_signals.reset, key=_signals.signals_of(ctx), idempotency_key=f"{ctx.request.id}:reset")() + + async def clear(ctx: ViewContext) -> Stopped: + """ + Forget the conversation: stop the message being answered (as + ``stop``), then empty the key's database, and what its ``signals`` + hold. The log's ids and cursors go on, so a follower reads on, and + drops the messages before ``first``. A message sent before it is + answered first, then forgotten too. + """ + stopped = await _stopped(ctx, "clear") + await ctx.rpc(erase, key=ctx.key, idempotency_key=f"{ctx.request.id}:erase")() + return stopped + + async def sync(ctx: ViewContext, request: SyncRequest | None = None) -> Synced: + """ + The conversation since ``cursor`` (by default, all of it), what the + agent is doing, and the revision; with ``revision``, once there is a + newer one, or at the timeout. A client follows the conversation by + sending each answer's ``revision`` and ``cursor`` back, at once while + it has ``more``. + """ + request = request or {} + revision = request.get("revision") + if revision is not None: + watch: Watch = {"revision": revision, "timeoutSeconds": request.get("timeoutSeconds", 50)} + await _companion(ctx, _signals.watch, "watch")(watch) + # The revision first: a change after it shows in the next sync. + now = await _companion(ctx, _signals.notifications)() + status = await _companion(ctx, _signals.status)() + page = await _log.page(ctx.db, {"cursor": request.get("cursor", 0), "limit": request.get("limit", 100)}) + synced: Synced = {**page, "revision": now["revision"], "status": status} + return synced + + agent.view(messages) + agent.handler(ingress_private=True)(erase) + agent.view(inactivity_timeout=WAITS)(clear) + agent.view(inactivity_timeout=WAITS)(sync) diff --git a/python/restate/agents/_turn.py b/python/restate/agents/_turn.py new file mode 100644 index 00000000..efa483ae --- /dev/null +++ b/python/restate/agents/_turn.py @@ -0,0 +1,301 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +The turn of a ``chat`` handler: its status, its cancellation, and publishing +for the clients that follow the conversation, with the framework's transcript +kept in the log as it goes. As sdk-sessions' ``chatHandler`` and +``transcribing`` (``common.ts``). +""" + +from __future__ import annotations + +import asyncio +import contextlib +import time +import typing +from collections.abc import Awaitable, Callable +from dataclasses import dataclass + +from restate.agents import _log, _signals +from restate.agents._api import Turn +from restate.agents._calls import settle +from restate.agents._common import refusal_of +from restate.agents._transcript import Transcript +from restate.exceptions import TerminalError +from restate.sessions import SessionContext +from restate.sessions._link import Link +from restate.sessions.internal import ScopeKey, cancel_calls, cancelled, definition_of + +T = typing.TypeVar("T") + +Read = Callable[[], Awaitable[Transcript]] +"""A framework's transcript as it stands, while it answers: a new one, or the one it keeps and adds to.""" + + +class AttemptEnded(Exception): + """ + Why a turn stopped when its attempt ended before it did: Restate suspended + the attempt, gave up on it, or lost its connection. Nothing the attempt does + after is stored, and the next attempt goes on with the turn, so it stops + without answering: unlike a cancellation, nothing is cancelled. + """ + + def __init__(self) -> None: + super().__init__("the attempt ended before its turn did") + + +def cancellation() -> TerminalError: + """What a cancelled turn fails with: Restate's cancellation, 409.""" + return TerminalError("cancelled", status_code=409) + + +class Stop: + """ + What stops a turn: a framework stops its work with it (a model's request, + say). It stops with Restate's cancellation (``409``) once Restate cancels + the invocation, the turn then ending as cancelled; or with an + ``AttemptEnded`` once the attempt ends first, the work stopping and + answering nothing. + """ + + def __init__(self) -> None: + self._event = asyncio.Event() + self.reason: BaseException | None = None + """Why it stopped, once it has.""" + + @property + def stopped(self) -> bool: + """Whether it stopped.""" + return self.reason is not None + + def stop(self, reason: BaseException) -> None: + """Stop, with ``reason``; once.""" + if self.reason is None: + self.reason = reason + self._event.set() + + async def wait(self) -> BaseException: + """Return once it stopped: why.""" + await self._event.wait() + assert self.reason is not None + return self.reason + + async def guard(self, work: Awaitable[T]) -> T: + """ + ``work``, unless it stops first: then ``work`` is cancelled, and this + raises why it stopped. For a framework's request that the turn's stop + must cut short. + """ + task = asyncio.ensure_future(work) + stopping = asyncio.ensure_future(self._event.wait()) + try: + await asyncio.wait({task, stopping}, return_when=asyncio.FIRST_COMPLETED) + finally: + stopping.cancel() + if task.done(): + return task.result() + task.cancel() + with contextlib.suppress(BaseException): + await task + assert self.reason is not None + raise self.reason + + +_STOP: ScopeKey[Stop] = ScopeKey("stop_signal_of") + + +def stop_signal_of(ctx: SessionContext) -> Stop | None: + """The ``Stop`` of the turn of ``ctx``, a ``chat`` handler's: see ``Stop``.""" + return _STOP.get(ctx) + + +@dataclass(frozen=True) +class AgentInfo: + """An agent's own: whether it works unwatched (``background``), and what its ``clear`` also empties.""" + + background: bool = False + clear: Callable[[SessionContext], Awaitable[None]] | None = None + + +AGENT: Link[AgentInfo] = Link("an agent") +"""An agent's own, on its definition.""" + + +def _agent_of(ctx: SessionContext) -> AgentInfo | None: + definition = definition_of(ctx) + return AGENT.of(definition) if definition is not None else None + + +async def cleared(ctx: SessionContext) -> None: + """Clear what the agent of ``ctx`` keeps beside its database.""" + agent = _agent_of(ctx) + if agent is not None and agent.clear is not None: + await agent.clear(ctx) + + +def _unfollowed(ctx: SessionContext) -> bool: + """Whether the agent of ``ctx`` runs in the background: it publishes the turn's end only.""" + agent = _agent_of(ctx) + return agent is not None and agent.background + + +_READING: ScopeKey[Read] = ScopeKey("transcribing") +"""The transcript of the attempt's turn, while it answers.""" + +_WRITING: ScopeKey[asyncio.Future[None]] = ScopeKey("the log's writes") +"""The attempt's last log write: they run one after another (tool calls in parallel publish at once).""" + + +def _logged(ctx: SessionContext, read: Read) -> Awaitable[None]: + """The log of the conversation of ``ctx`` takes its transcript as ``read`` gives it.""" + before = _WRITING.get(ctx) + + async def write() -> None: + if before is not None: + with contextlib.suppress(BaseException): + await before + transcript = await read() + await _log.sync(ctx, transcript.messages()) + # The calls whose results the framework stored are made: see _calls.py. + await settle(ctx, transcript.answered) + + done = asyncio.ensure_future(write()) + _WRITING.set(ctx, done) + return done + + +async def publish(ctx: SessionContext) -> None: + """ + The conversation of ``ctx`` changed: its log takes the change (while a + framework answers: ``transcribing``), its revision increases, and the + clients that follow it read it again. + """ + read = _READING.get(ctx) + if read is not None: + await _logged(ctx, read) + await ctx.rpc(_signals.publish, key=_signals.signals_of(ctx))() + + +async def stepped(ctx: SessionContext) -> None: + """A step of a turn of ``ctx``: ``publish`` it, unless the agent runs in the background.""" + if not _unfollowed(ctx): + await publish(ctx) + + +async def tool_call_started(ctx: SessionContext) -> None: + """ + A tool call of the turn of ``ctx`` starts: ``publish`` it, which forgets + the calls whose results are stored before the tool makes its own. In the + background, it forgets them alone. + """ + if not _unfollowed(ctx): + await publish(ctx) + return + read = _READING.get(ctx) + if read is not None: + transcript = await read() + await settle(ctx, transcript.answered) + + +async def tool_call_ended(ctx: SessionContext) -> None: + """A tool call of the turn of ``ctx`` ended: ``publish`` it, unless the agent runs in the background.""" + await stepped(ctx) + + +async def transcribing(ctx: SessionContext, read: Read, work: Callable[[], Awaitable[T]]) -> T: + """ + Run ``work``, a framework's turn (or a compaction, a reset), with ``read`` + giving the framework's transcript as it stands: each publish meanwhile (a + tool call's start and end) takes it into the conversation's log, and so + does ``work``'s end. ``messages`` and ``sync`` read the log. + """ + outer = _READING.get(ctx) + _READING.set(ctx, read) + failed = False + try: + return await work() + except BaseException as error: + failed = not isinstance(error, TerminalError) + raise + finally: + try: + # A failed attempt stores nothing more: its retry runs again. + if not failed: + await _logged(ctx, read) + finally: + if outer is not None: + _READING.set(ctx, outer) + else: + _READING.delete(ctx) + + +def _now() -> str: + return _signals.iso(time.time()) + + +async def chat_turn(ctx: SessionContext, text: str, answer: Callable[[SessionContext, str], Awaitable[T]]) -> T: + """ + ``answer`` as a chat handler's turn. While it answers, the object's + ``status`` says so (its ``Turn``). Its invocation's cancellation (Restate's, + explicit) stops the turn (``stop_signal_of``) whenever it comes, and cancels + the calls its tools made: the turn ends, and the handler fails with the + cancellation. An attempt that ends first stops it too, with an + ``AttemptEnded``: the work stops, and the next attempt goes on with it. It + publishes when it starts and ends, for the clients that follow it. + """ + invocation_id = ctx.request.id + companion = _signals.signals_of(ctx) + turn: Turn = {"invocationId": invocation_id, "since": _now()} + # The companion knows the turn from here to its close: `status` and `stop` ask it. + await ctx.rpc(_signals.start, key=companion)(turn) + await stepped(ctx) + + async def end() -> None: + # The log takes the turn's last change, then the companion closes the + # turn, which publishes it: one revision. The close is a call the turn + # waits for, so that it is done before the turn ends. + read = _READING.get(ctx) + if read is not None: + await _logged(ctx, read) + await ctx.rpc(_signals.close, key=companion)(invocation_id) + + stop = Stop() + _STOP.set(ctx, stop) + + async def on_cancel() -> None: + await cancelled(ctx) + stop.stop(cancellation()) + await cancel_calls(ctx) + + async def on_attempt_end() -> None: + await ctx.request.attempt_finished_event.wait() + stop.stop(AttemptEnded()) + + watchers = [asyncio.ensure_future(on_cancel()), asyncio.ensure_future(on_attempt_end())] + try: + try: + answered = await answer(ctx, text) + # Cancelled once answered: the answer stands. + except TerminalError: + await end() + raise + except Exception as error: + # Any other error fails the attempt, and Restate retries the turn, but a refusal. + terminal = refusal_of(ctx, error) + if terminal is None: + raise + await end() + raise terminal from error + await end() + return answered + finally: + for watcher in watchers: + watcher.cancel() diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index b0847ea0..bf3ba3e0 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -128,6 +128,9 @@ def wait(self, handle: int) -> Awaitable[None]: """Done once the VM has the read whose handle is ``handle``: through the selector, as every wait.""" return self.selector.wait(SingleUnresolvedFuture(handle)) + def cancel_invocation(self, invocation_id: str) -> None: + self.context.cancel_invocation(invocation_id) + def get(self, key: str) -> Awaitable[bytes | None]: return self.context.get(key, serde=BytesSerde(), entry_name=key) @@ -197,7 +200,7 @@ async def _nothing() -> None: """The action of an empty run.""" -def _cancel_of(host: Any) -> Cancel: +def _cancel_of(host: SqlHost) -> Cancel: """``Cancel`` over the SDK's context.""" def cancel(invocation_id: str) -> None: @@ -514,6 +517,15 @@ def _decorate(self, view: bool, name: str | F | None, options: dict[str, Any]) - return self._register(view, None, options)(name) return self._register(view, name, options) + def _attempt_fn(self, fn: Callable[..., Awaitable[Any]], view: bool) -> Callable[..., Awaitable[Any]]: + """ + What an attempt of handler (or view, with ``view``) ``fn`` runs: ``fn`` + itself; a subclass's may wrap it (an agent's turns a model provider's + refusal into a terminal error). + """ + del view + return fn + def _register(self, view: bool, name: str | None, options: dict[str, Any]) -> Callable[[F], F]: def register(fn: F) -> F: if RESTATE_UNIQUE_HANDLER_SYMBOL in vars(fn): @@ -526,7 +538,11 @@ def register(fn: F) -> F: explicit = bool(options.get("explicit_cancellation")) sdk_options = {key: value for key, value in options.items() if key != "explicit_cancellation"} entry = _Entry( - fn=fn, view=view, config=self.config, definition=weakref.ref(self), explicit_cancellation=explicit + fn=self._attempt_fn(fn, view), + view=view, + config=self.config, + definition=weakref.ref(self), + explicit_cancellation=explicit, ) registered = VirtualObject.handler(self, name, kind="shared" if view else "exclusive", **sdk_options)( _sdk_handler(entry) diff --git a/python/restate/sessions/_session.py b/python/restate/sessions/_session.py index 77ad05da..a0a78152 100644 --- a/python/restate/sessions/_session.py +++ b/python/restate/sessions/_session.py @@ -111,6 +111,10 @@ def wait(self, handle: int) -> Awaitable[None]: """Done once the VM has the read whose handle is ``handle``.""" ... + def cancel_invocation(self, invocation_id: str) -> None: + """Cancel the invocation ``invocation_id``, as Restate's admin API does.""" + ... + @dataclass class ObjectConfig: diff --git a/tests/agents_object.py b/tests/agents_object.py new file mode 100644 index 00000000..23e8996a --- /dev/null +++ b/tests/agents_object.py @@ -0,0 +1,173 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""An agent object and its turns, with a scripted framework, against a real Restate server.""" + +# pylint: disable=missing-function-docstring,missing-class-docstring + +import asyncio +import typing +import uuid +from collections.abc import Awaitable, Callable + +import httpx +import pytest + +import restate +from restate import HarnessEnvironment, sessions +from restate.agents._api import Messages, Status +from restate.agents._object import AgentObject +from restate.agents._transcript import Transcript +from restate.agents._turn import stop_signal_of, transcribing + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +pytestmark = [pytest.mark.anyio] + + +class AuthenticationError(Exception): + """As a model SDK names the error of a missing or wrong API key.""" + + +# A scripted framework: it keeps each key's transcript in this process, and answers with an echo. +transcripts: dict[str, Transcript] = {} +gate = asyncio.Event() + + +async def answer(ctx: sessions.SessionContext, text: str) -> str: + transcript = transcripts.setdefault(ctx.key, Transcript()) + transcript.say("user", text) + + async def read() -> Transcript: + return transcript + + async def work() -> str: + if text == "wait": + stop = stop_signal_of(ctx) + assert stop is not None + await stop.guard(gate.wait()) + if text == "refuse": + raise AuthenticationError("invalid x-api-key") + transcript.say("assistant", f"echo: {text}") + return f"echo: {text}" + + return await transcribing(ctx, read, work) + + +async def forget(ctx: sessions.SessionContext) -> None: + transcripts.pop(ctx.key, None) + + +assistant = AgentObject("assistant") +assistant.chat(answer, clear=forget) + + +@pytest.fixture(scope="session") +async def harness(): + async with restate.create_test_harness( + sessions.app(assistant), restate_image="ghcr.io/restatedev/restate:latest" + ) as harness: + yield harness + + +TEXT = {"content-type": "text/plain"} +"""A message as text, as `curl -d` sends one.""" + + +async def call(harness: HarnessEnvironment, handler: str, key: str, body: object = None) -> httpx.Response: + async with httpx.AsyncClient(timeout=60) as ingress: + return await ingress.post(f"{harness.ingress_url}/assistant/{key}/{handler}", json=body) + + +async def send(harness: HarnessEnvironment, key: str, text: str) -> str: + async with httpx.AsyncClient(timeout=60) as ingress: + sent = await ingress.post(f"{harness.ingress_url}/assistant/{key}/chat/send", content=text, headers=TEXT) + return typing.cast(str, sent.json()["invocationId"]) + + +async def attach(harness: HarnessEnvironment, invocation_id: str) -> httpx.Response: + async with httpx.AsyncClient(timeout=60) as ingress: + return await ingress.get(f"{harness.ingress_url}/restate/invocation/{invocation_id}/attach") + + +async def until(check: Callable[[], Awaitable[bool]]) -> None: + for _ in range(200): + if await check(): + return + await asyncio.sleep(0.05) + raise AssertionError("never") + + +def texts(page: Messages) -> list[tuple[int, str, str]]: + return [ + (message["id"], message["role"], typing.cast(str, message["parts"][0].get("text"))) + for message in page["messages"] + ] + + +async def test_a_message_is_answered_and_the_conversation_read_and_followed(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + async with httpx.AsyncClient(timeout=60) as ingress: + answered = await ingress.post(f"{harness.ingress_url}/assistant/{key}/chat", content="Hi", headers=TEXT) + assert answered.json() == "echo: Hi" + page = (await call(harness, "messages", key, {})).json() + assert texts(page) == [(0, "user", "Hi"), (1, "assistant", "echo: Hi")] + synced = (await call(harness, "sync", key, {})).json() + assert synced["revision"] > 0 and synced["status"] == {"answering": None, "waiting": []} + assert texts(synced) == texts(page) + # Followed from its cursor, the conversation has nothing newer; the next message is. + later = (await call(harness, "sync", key, {"cursor": synced["cursor"]})).json() + assert later["messages"] == [] + await call(harness, "chat", key, "Again") + newer = (await call(harness, "sync", key, {"cursor": synced["cursor"]})).json() + assert [m["id"] for m in newer["messages"]] == [2, 3] + + +async def test_status_says_which_message_is_answered_and_stop_cancels_it(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + gate.clear() + invocation = await send(harness, key, "wait") + + async def answering() -> bool: + status = typing.cast(Status, (await call(harness, "status", key)).json()) + return status["answering"] is not None and status["answering"]["invocationId"] == invocation + + await until(answering) + assert (await call(harness, "stop", key)).json() == {"stopped": invocation} + assert (await attach(harness, invocation)).status_code == 409 + assert (await call(harness, "status", key)).json()["answering"] is None + # The next message is answered. + gate.set() + assert (await call(harness, "chat", key, "next")).json() == "echo: next" + + +async def test_clear_forgets_the_conversation_and_its_ids_go_on(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + await call(harness, "chat", key, "one") + assert (await call(harness, "clear", key)).json() == {"stopped": None} + cleared = (await call(harness, "messages", key, {})).json() + assert cleared["messages"] == [] and cleared["first"] == 2 + await call(harness, "chat", key, "two") + assert texts((await call(harness, "messages", key, {})).json()) == [ + (2, "user", "two"), + (3, "assistant", "echo: two"), + ] + + +async def test_a_refused_message_fails_with_424_and_the_next_is_answered(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + refused = await call(harness, "chat", key, "refuse") + assert refused.status_code == 424 + assert (await call(harness, "status", key)).json()["answering"] is None + assert (await call(harness, "chat", key, "fine")).json() == "echo: fine" diff --git a/tests/sessions_harness.py b/tests/sessions_harness.py index 2ea69cb7..3f5e0a99 100644 --- a/tests/sessions_harness.py +++ b/tests/sessions_harness.py @@ -235,3 +235,29 @@ async def test_a_handler_with_explicit_cancellation_stops_its_own_work(harness: [nap_id] = answer.json() assert (await attach(harness, nap_id)).status_code == 409 assert await harness.client.object_call(texts, key=key, arg=None) == ["stopped: nap"] + + +@notes.view +async def stop_it(ctx: sessions.ViewContext, invocation_id: str) -> None: + """A Stop button: a view cancels a running handler of its own object.""" + ctx.cancel(invocation_id) + + +async def test_a_view_cancels_a_handler_of_its_own_object(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + sent = await harness.client.object_send(stoppable, key=key, arg="nap") + async with httpx.AsyncClient() as admin: + for _ in range(100): + status = await admin.post( + f"{harness.admin_api_url}/query", + json={"query": f"SELECT status FROM sys_invocation WHERE id = '{sent.invocation_id}'"}, + headers={"Accept": "application/json"}, + ) + rows = status.json().get("rows", []) + if rows and rows[0]["status"] in ("running", "suspended"): + break + await asyncio.sleep(0.05) + await harness.client.object_call(stop_it, key=key, arg=sent.invocation_id) + answer = await attach(harness, sent.invocation_id) + assert answer.status_code == 200, answer.text + assert await harness.client.object_call(texts, key=key, arg=None) == ["stopped: nap"] From 1ae3731f6ee384cebd183745cd9b6a38afebef78 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 16:25:49 +0100 Subject: [PATCH 22/45] restate.agents: tools, and the package's API The fourth part of the agents' core, from sdk-sessions' tools.ts: - tool: a tool from a plain function, its input from the type hints and its description from the docstring; ctx, call_id and idempotency_key are given, not the model's. A TerminalError is the tool's answer, any other error fails the attempt (ToolError), and arguments that do not fit the input answer the model with what is wrong. - The tools that are Restate itself: call and send (a handler, the model choosing an object's key unless given), workflow, subagent (a task to another agent, keyed by this key and the task's name), and signal (a wait at the companion, which the agent's signal view resolves). Each is made once per tool call (_calls.py): the same Restate call however often a crash runs its tool call again. - run_tool, for a framework's integration: the tool call's start and end are published. Tools take their schemas from pydantic, imported when a tool is made: restate.sessions and restate.agents import without it (a test checks, in a process without pydantic), and an agent framework's extra brings it. restate.agents exports what agents are built with; restate.agents.internal, what a framework's integration needs. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/agents/__init__.py | 66 +++- python/restate/agents/_tools.py | 519 ++++++++++++++++++++++++++++++ python/restate/agents/internal.py | 76 +++++ tests/agents_imports.py | 37 +++ tests/agents_object.py | 39 +++ tests/agents_tools.py | 153 +++++++++ 6 files changed, 889 insertions(+), 1 deletion(-) create mode 100644 python/restate/agents/_tools.py create mode 100644 python/restate/agents/internal.py create mode 100644 tests/agents_imports.py create mode 100644 tests/agents_tools.py diff --git a/python/restate/agents/__init__.py b/python/restate/agents/__init__.py index ef1b198d..645b0000 100644 --- a/python/restate/agents/__init__.py +++ b/python/restate/agents/__init__.py @@ -10,5 +10,69 @@ # """ AI agents in session objects, each conversation in its session's database. -The Python variant of ``@restatedev/agents``. +The Python variant of ``@restatedev/agents``: what the frameworks share +(session objects set up for agents, tools for any of them, and the tools that +are Restate itself). Each framework is a module of its own, with the extra of +its packages; what their integrations need beside this is in +``restate.agents.internal``. + +A crash, a restart or a person who takes a day to approve changes nothing: +each message is answered once, and each tool call is made once. """ + +from restate.agents._api import ( + ChatMessage, + Messages, + MessagesRequest, + Notifications, + Part, + Resolved, + Signal, + Status, + Stopped, + TextPart, + ToolPart, + Turn, + Waiting, + Watch, +) +from restate.agents._common import CANCELLED, REFUSED, UNANSWERED, Configure +from restate.agents._object import AgentObject, Synced, SyncRequest +from restate.agents._signals import Companion, companion_of, deliver, signals +from restate.agents._tools import Tool, ToolCall, call, send, signal, subagent, tool, workflow + +__all__ = [ + "CANCELLED", + "REFUSED", + "UNANSWERED", + "AgentObject", + "ChatMessage", + "Companion", + "Configure", + "Messages", + "MessagesRequest", + "Notifications", + "Part", + "Resolved", + "Signal", + "Status", + "Stopped", + "SyncRequest", + "Synced", + "TextPart", + "Tool", + "ToolCall", + "ToolPart", + "Turn", + "Waiting", + "Watch", + "call", + "companion_of", + "deliver", + "send", + "signal", + "signals", + "subagent", + "tool", + "workflow", +] diff --git a/python/restate/agents/_tools.py b/python/restate/agents/_tools.py new file mode 100644 index 00000000..1b202c89 --- /dev/null +++ b/python/restate/agents/_tools.py @@ -0,0 +1,519 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +Tools for any of the frameworks: defined once with ``tool``, given to any +agent's ``tools``, which turns them into the framework's own. And the tools +that are Restate itself: calling a handler, sending to one, running a workflow, +handing a task to another agent, and waiting for a signal. As sdk-sessions' +``tools.ts``. + +A tool runs with the agent's context, and with the id of the model's tool +call: Restate calls carry an idempotency key made of the invocation id and that +id, so a tool call run again after a crash makes the same call, once. + +A tool's input schema comes from type hints, by pydantic, imported when a tool +is made: an agent framework's extra brings it. +""" + +from __future__ import annotations + +import inspect +import json +import typing +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from datetime import timedelta +from typing import Any + +from restate.agents import _signals +from restate.agents._calls import call_of +from restate.agents._common import Configure, ToolError, configured +from restate.agents._turn import stop_signal_of, tool_call_ended, tool_call_started +from restate.exceptions import TerminalError +from restate.handler import Handler, handler_from_callable +from restate.sessions import SessionContext + +T = typing.TypeVar("T") + +INJECTED = ("ctx", "call_id", "idempotency_key") +"""The parameters a tool's function gets from where it runs, rather than from the model.""" + + +def _pydantic() -> Any: + """pydantic, which tools take their inputs' schemas from.""" + try: + import pydantic + except ImportError as error: + raise ImportError( + "restate.agents' tools take their inputs' schemas from pydantic: install it, or an agent framework's " + "extra, which brings it (such as restate-sdk[anthropic])" + ) from error + return pydantic + + +@dataclass(frozen=True) +class ToolCall: + """A tool call: the agent's context, and the id of the model's call (the same when a crash runs it again).""" + + ctx: SessionContext + call_id: str + + @property + def idempotency_key(self) -> str: + """ + The tool call's idempotency key, ``:``: for a + call the tool makes once, however often it runs (a Restate call's, or a + payment API's). + """ + return f"{self.ctx.request.id}:{self.call_id}" + + +@dataclass(frozen=True) +class Tool: + """ + A tool of any framework: its ``name`` and ``description`` for the model, + its input's JSON schema, which a framework offers the model, and how it + runs: ``run(args, call)``, its result as text for the model, of the + model's arguments as they came (arguments that do not fit the input answer + with what is wrong, for the model to call it again). + """ + + name: str + description: str + input_schema: dict[str, Any] + _run: Callable[[dict[str, Any], ToolCall], Awaitable[Any]] + _validate: Callable[[dict[str, Any]], Any] + + async def run(self, args: object, call: ToolCall) -> str: + """The tool's result for the model, as text: see ``tool`` for how it fails.""" + try: + arguments = self._validate(typing.cast(dict[str, Any], args if isinstance(args, dict) else {})) + except Exception as error: # pydantic's ValidationError: the model's to fix + return f"Error: the arguments do not fit the tool's input: {error}" + try: + return _text(await self._run(arguments, call)) + except BaseException as error: + # The turn stopped (its invocation was cancelled, or its attempt + # ended): it ends, rather than the model going on with the tool's + # error as its answer (such as a call the cancellation cancelled). + stop = stop_signal_of(call.ctx) + if stop is not None and stop.reason is not None: + raise stop.reason from error + if isinstance(error, TerminalError): + # A terminal error is the tool's answer, which the model can act on. + return f"Error: {error.message}" + if not isinstance(error, Exception): + raise + # Any other error fails the attempt, which Restate retries. + raise ToolError(self.name, error) from error + + +async def run_tool(tool: Tool, args: object, call: ToolCall) -> str: + """ + Run ``tool`` for tool call ``call``, as a framework's integration runs one: + the conversation publishes its start and its end, for the clients that + follow it, and the calls whose results are stored are settled first. + """ + await tool_call_started(call.ctx) + try: + return await tool.run(args, call) + finally: + await tool_call_ended(call.ctx) + + +def _text(value: object) -> str: + """A result as text: a string as it is, nothing as "done", anything else as JSON.""" + if isinstance(value, str): + return value + if value is None: + return "done" + try: + return json.dumps(value) + except TypeError: + return typing.cast(bytes, _pydantic().TypeAdapter(type(value)).dump_json(value)).decode() + + +# --- tool: a function of your own ----------------------------------------------------------------- + + +@typing.overload +def tool(fn: Callable[..., Any], /, *, name: str | None = None, description: str | None = None) -> Tool: ... + + +@typing.overload +def tool(*, name: str | None = None, description: str | None = None) -> Callable[[Callable[..., Any]], Tool]: ... + + +def tool( + fn: Callable[..., Any] | None = None, *, name: str | None = None, description: str | None = None +) -> Tool | Callable[[Callable[..., Any]], Tool]: + """ + A tool, from a function: the model gives its parameters, described by + their types and the docstring:: + + @tool + async def get_weather(city: str) -> str: + '''The current weather in a city.''' + return f"Sunny in {city}" + + Three parameters are not the model's: ``ctx``, the agent's context + (``ctx.db``, calls, sends); ``call_id``, the tool call's id; and + ``idempotency_key``, made of both, for what the tool must do once. A crash + runs an unfinished tool call again with the same id. A + ``restate.TerminalError`` the function raises (or a Restate call it makes + fails with) is its answer, ``Error: ``, for the model to act on; + any other error fails the attempt, which Restate retries. + """ + if fn is None: + + def decorate(fn: Callable[..., Any]) -> Tool: + return tool(fn, name=name, description=description) + + return decorate + pydantic = _pydantic() + parameters = inspect.signature(fn, eval_str=True).parameters.values() + fields: dict[str, Any] = { + p.name: (p.annotation if p.annotation is not p.empty else Any, ... if p.default is p.empty else p.default) + for p in parameters + if p.name not in INJECTED + } + model = pydantic.create_model(f"{fn.__name__}_input", **fields) + injected = [p.name for p in parameters if p.name in INJECTED] + + async def run(arguments: dict[str, Any], call: ToolCall) -> Any: + values = {"ctx": call.ctx, "call_id": call.call_id, "idempotency_key": call.idempotency_key} + result = fn(**arguments, **{name: values[name] for name in injected}) + return await result if inspect.isawaitable(result) else result + + def validate(args: dict[str, Any]) -> dict[str, Any]: + validated = model.model_validate(args) + return {field: getattr(validated, field) for field in fields} + + return Tool( + name=name or fn.__name__, + description=description or inspect.getdoc(fn) or "", + input_schema=typing.cast(dict[str, Any], model.model_json_schema()), + _run=run, + _validate=validate, + ) + + +# --- The tools that are Restate itself ------------------------------------------------------------ + + +async def _once(name: str, input: object, call: ToolCall, body: Callable[[ToolCall], Awaitable[T]]) -> T: + """ + ``body``, a Restate call of tool ``name`` with ``input``, made once however + often the tool call runs (see _calls.py): it gets the id of the call, which + an earlier attempt may have made, and so its idempotency key. + """ + made = await call_of(call.ctx, name, input, call.call_id) + try: + result = await body(ToolCall(call.ctx, made.id)) + except TerminalError as error: + # A terminal error answers the call; any other fails the attempt, and the retry makes it again. + await made.answered(f"Error: {error.message}") + raise + await made.answered() + return result + + +def _untyped(caller: object) -> Callable[..., Awaitable[Any]]: + """A caller of a handler whose input is known only at run time (a tool's): it takes the input, if any.""" + return typing.cast(Callable[..., Awaitable[Any]], caller) + + +def _untyped_send(sender: object) -> Callable[..., None]: + """A sender to a handler whose input is known only at run time (a tool's): it takes the input, if any.""" + return typing.cast(Callable[..., None], sender) + + +def _handler(target: Callable[..., Any]) -> Handler[Any, Any]: + try: + return handler_from_callable(target) + except ValueError: + raise TypeError(f"{target!r} is not a handler of a service, an object or a workflow") from None + + +def _input_type(handler: Handler[Any, Any]) -> Any: + """The type of a handler's input, or None if it takes none.""" + if handler.arity < 2: + return None + parameters = list(inspect.signature(handler.fn, eval_str=True).parameters.values()) + annotation = parameters[1].annotation + return Any if annotation is inspect.Parameter.empty else annotation + + +@dataclass(frozen=True) +class _Input: + """A Restate tool's input: the handler's (flattened if an object), and the key if the model chooses it.""" + + schema: dict[str, Any] + validate: Callable[[dict[str, Any]], tuple[str | None, Any]] + + +def _input_of(input_type: Any, model_keys: str | None) -> _Input: + """ + The tool's input: ``input_type``'s, with ``key`` beside it if the model + chooses the key of ``model_keys`` (an object's, a workflow's). An object's + fields are the tool's own; any other input is the tool's field ``input``. + """ + pydantic = _pydantic() + adapter = pydantic.TypeAdapter(input_type) if input_type is not None else None + inner: dict[str, Any] = adapter.json_schema() if adapter is not None else {"type": "object", "properties": {}} + flat = inner.get("type") == "object" and "properties" in inner + properties: dict[str, Any] = {} + required: list[str] = [] + if model_keys is not None: + properties["key"] = {"type": "string", "description": f"Which {model_keys}: its key"} + required.append("key") + if flat: + properties.update(inner["properties"]) + required.extend(inner.get("required", [])) + elif adapter is not None: + properties["input"] = inner + required.append("input") + schema: dict[str, Any] = {"type": "object", "properties": properties, "required": required} + if "$defs" in inner: + schema["$defs"] = inner["$defs"] + + def validate(args: dict[str, Any]) -> tuple[str | None, Any]: + key = None + if model_keys is not None: + key = args.get("key") + if not isinstance(key, str): + raise ValueError("key: a string is required") + if adapter is None: + return key, None + value = {k: v for k, v in args.items() if k != "key"} if flat else args.get("input") + return key, adapter.validate_python(value) + + return _Input(schema=schema, validate=validate) + + +def _keyed(handler: Handler[Any, Any]) -> bool: + return handler.service_tag.kind != "service" + + +def call( + target: Callable[..., Any], + *, + description: str, + name: str | None = None, + key: Configure[str] | None = None, +) -> Tool: + """ + A tool that calls handler ``target`` and answers with its result:: + + call(order_status, description="The status of an order.") + + The model fills in the handler's input; of a virtual object, it chooses the + key too, unless ``key`` gives it. + """ + handler = _handler(target) + service = handler.service_tag.name + model_keys = service if _keyed(handler) and key is None else None + shape = _input_of(_input_type(handler), model_keys) + + async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: + chosen, value = arguments["key"], arguments["value"] + + async def made(c: ToolCall) -> Any: + target_key = chosen if key is None else await configured(key, c.ctx) + caller = _untyped( + c.ctx.rpc(target, key=target_key if _keyed(handler) else None, idempotency_key=c.idempotency_key) + ) + return await (caller(value) if handler.arity >= 2 else caller()) + + return await _once(f"{service}/{handler.name}", arguments["raw"], tool_call, made) + + return Tool( + name=name or f"{service}_{handler.name}", + description=description, + input_schema=shape.schema, + _run=run, + _validate=lambda args: _split(shape, args), + ) + + +def send( + target: Callable[..., Any], + *, + description: str, + name: str | None = None, + key: Configure[str] | None = None, + delay: timedelta | None = None, +) -> Tool: + """A tool that sends to handler ``target``, without waiting for it: a one-way call, made once, after ``delay``.""" + handler = _handler(target) + service = handler.service_tag.name + model_keys = service if _keyed(handler) and key is None else None + shape = _input_of(_input_type(handler), model_keys) + + async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: + chosen, value = arguments["key"], arguments["value"] + + async def made(c: ToolCall) -> str: + target_key = chosen if key is None else await configured(key, c.ctx) + sender = _untyped_send( + c.ctx.send( + target, key=target_key if _keyed(handler) else None, delay=delay, idempotency_key=c.idempotency_key + ) + ) + if handler.arity >= 2: + sender(value) + else: + sender() + return f"sent to {service}{'' if target_key is None else f'/{target_key}'}/{handler.name}" + + return await _once(f"{service}/{handler.name}", arguments["raw"], tool_call, made) + + return Tool( + name=name or f"{service}_{handler.name}", + description=description, + input_schema=shape.schema, + _run=run, + _validate=lambda args: _split(shape, args), + ) + + +def workflow( + target: Callable[..., Any], + *, + description: str, + name: str | None = None, + key: Configure[str] | None = None, +) -> Tool: + """ + A tool that runs the workflow of ``target`` (its main handler) to the end, + and answers with its result. Each tool call is a workflow of its own, with + an id made of the agent's key and the tool call's id, unless ``key`` gives + one. + """ + handler = _handler(target) + service = handler.service_tag.name + shape = _input_of(_input_type(handler), None) + + async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: + value = arguments["value"] + + async def made(c: ToolCall) -> Any: + workflow_id = f"{c.ctx.key}-{c.call_id}" if key is None else await configured(key, c.ctx) + caller = _untyped(c.ctx.rpc(target, key=workflow_id, idempotency_key=c.idempotency_key)) + return await (caller(value) if handler.arity >= 2 else caller()) + + return await _once(service, arguments["raw"], tool_call, made) + + return Tool( + name=name or service, + description=description, + input_schema=shape.schema, + _run=run, + _validate=lambda args: _split(shape, args), + ) + + +def _split(shape: _Input, args: dict[str, Any]) -> dict[str, Any]: + """The tool call's key and the handler's input, validated, and the raw arguments (what makes the call the same).""" + key, value = shape.validate(args) + return {"key": key, "value": value, "raw": args} + + +_TASK_NAME = r"^[a-z0-9-]+$" + + +def subagent( + agent: object, + *, + description: str | None = None, + name: str | None = None, + handler: str = "chat", +) -> Tool: + """ + A tool, named after ``agent`` (an agent, or any virtual object) unless + ``name`` names it, that hands a task to it: to the object keyed + ``.``, through its ``handler`` (``chat``), and answers + with its answer. Tasks of different names run at once; the same name + continues the same task. The model reads what the agent is from its + description:: + + coder = claude.agent("coder", description="Builds and tests code in Python") + lead = claude.agent("lead", tools=[subagent(coder)]) + """ + agent_name = typing.cast(str, getattr(agent, "name")) # noqa: B009 (an SDK's definition) + tag = getattr(agent, "service_tag", None) + about = description if description is not None else typing.cast("str | None", getattr(tag, "description", None)) + pydantic = _pydantic() + + class Task(pydantic.BaseModel): # type: ignore[name-defined] # pyright: ignore[reportUntypedBaseClass] + name: str = pydantic.Field(pattern=_TASK_NAME, description='A short name for the task, such as "fizzbuzz"') + task: str = pydantic.Field(description="The task in full: the agent sees nothing else") + + async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: + task_name, task = typing.cast(str, arguments["name"]), typing.cast(str, arguments["task"]) + + async def made(c: ToolCall) -> Any: + return await c.ctx.rpc.generic( + agent_name, handler, task, key=f"{c.ctx.key}.{task_name}", idempotency_key=c.idempotency_key + ) + + return await _once(agent_name, {"name": task_name, "task": task}, tool_call, made) + + def validate(args: dict[str, Any]) -> dict[str, Any]: + task = Task.model_validate(args) + return {"name": task.name, "task": task.task} + + return Tool( + name=name or agent_name, + description=( + f"Hand a task to the {agent_name} agent{f' ({about})' if about else ''}, and get its answer. " + "Tasks of different names run at once; the same name continues the same task." + ), + input_schema=typing.cast(dict[str, Any], Task.model_json_schema()), + _run=run, + _validate=validate, + ) + + +def signal(name: str, *, description: str, input: Any = None, tool_name: str | None = None) -> Tool: + """ + A tool that waits for the signal ``name`` to the agent, and answers with its + value: a person's approval, say, or another system's result. The agent + waits without holding a process. With an ``input`` type, the model says + what it waits for (such as the refund to approve), which ``pending`` shows:: + + signal("approval", description="Wait for a manager to approve a refund.", input=Refund) + + Each wait has an id: the agent's ``pending`` lists them, and its ``signal`` + resolves one, by its id, or the oldest wait for a name:: + + curl localhost:8080/support/alice/signal --json '{"name": "approval", "value": "approved"}' + """ + shape = _input_of(input, None) + + async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: + async def made(c: ToolCall) -> Any: + wait_id = c.idempotency_key + request: _signals.WaitRequest = {"id": wait_id, "name": name} + if input is not None: + request["input"] = arguments["raw"] + waiting = c.ctx.rpc(_signals.wait, key=_signals.signals_of(c.ctx), idempotency_key=wait_id) + return await waiting(request) + + return await _once(f"signal/{name}", arguments["raw"], tool_call, made) + + return Tool( + name=tool_name or f"wait_for_{name}", + description=description, + input_schema=shape.schema, + _run=run, + _validate=lambda args: _split(shape, args), + ) diff --git a/python/restate/agents/internal.py b/python/restate/agents/internal.py new file mode 100644 index 00000000..db06cb29 --- /dev/null +++ b/python/restate/agents/internal.py @@ -0,0 +1,76 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +What a framework's integration needs beside ``restate.agents``: its turn's +transcript kept in the log (``transcribing``, ``Transcript``), its stop +(``stop_signal_of``), its tool calls (``run_tool``), its model (``model_of`` +and the rest), and a model provider's refusal (``refusal_of``, ``from_tool``). +""" + +from restate.agents._common import ( + CHECK_MODEL, + CLAUDE_MODEL, + DEFAULT_MODEL, + MessageSerde, + ModelError, + ToolError, + check_model_name, + configured, + from_tool, + model_id_of, + model_of, + next_prompt_id, + provider_of, + refusal, + refusal_of, + stored_answer, +) +from restate.agents._tools import run_tool +from restate.agents._transcript import Transcript +from restate.agents._turn import ( + AttemptEnded, + Read, + Stop, + publish, + stop_signal_of, + tool_call_ended, + tool_call_started, + transcribing, +) + +__all__ = [ + "CHECK_MODEL", + "CLAUDE_MODEL", + "DEFAULT_MODEL", + "AttemptEnded", + "MessageSerde", + "ModelError", + "Read", + "Stop", + "ToolError", + "Transcript", + "check_model_name", + "configured", + "from_tool", + "model_id_of", + "model_of", + "next_prompt_id", + "provider_of", + "publish", + "refusal", + "refusal_of", + "run_tool", + "stop_signal_of", + "stored_answer", + "tool_call_ended", + "tool_call_started", + "transcribing", +] diff --git a/tests/agents_imports.py b/tests/agents_imports.py new file mode 100644 index 00000000..0da90074 --- /dev/null +++ b/tests/agents_imports.py @@ -0,0 +1,37 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +restate.sessions and restate.agents need no package beyond the SDK's own: a +framework's comes with its extra, and pydantic, which tools take their schemas +from, with any of them. +""" + +# pylint: disable=missing-function-docstring + +import subprocess +import sys + +WITHOUT_PYDANTIC = """ +import sys +sys.modules["pydantic"] = None +sys.modules["pydantic_core"] = None +import restate, restate.sessions, restate.sessions.internal, restate.agents, restate.agents.internal +try: + restate.agents.tool(lambda city: city) +except ImportError as error: + print(error) +""" + + +def test_sessions_and_agents_import_without_pydantic_and_a_tool_says_what_it_needs(): + ran = subprocess.run([sys.executable, "-c", WITHOUT_PYDANTIC], capture_output=True, text=True, check=False) + assert ran.returncode == 0, ran.stderr + assert "tools take their inputs' schemas from pydantic" in ran.stdout diff --git a/tests/agents_object.py b/tests/agents_object.py index 23e8996a..aa8706c8 100644 --- a/tests/agents_object.py +++ b/tests/agents_object.py @@ -18,12 +18,14 @@ from collections.abc import Awaitable, Callable import httpx +import pydantic import pytest import restate from restate import HarnessEnvironment, sessions from restate.agents._api import Messages, Status from restate.agents._object import AgentObject +from restate.agents._tools import ToolCall, run_tool, signal from restate.agents._transcript import Transcript from restate.agents._turn import stop_signal_of, transcribing @@ -40,6 +42,13 @@ class AuthenticationError(Exception): """As a model SDK names the error of a missing or wrong API key.""" +class Refund(pydantic.BaseModel): + order: str + amount: float + + +approval = signal("approval", description="Wait for a manager to approve a refund.", input=Refund) + # A scripted framework: it keeps each key's transcript in this process, and answers with an echo. transcripts: dict[str, Transcript] = {} gate = asyncio.Event() @@ -59,6 +68,15 @@ async def work() -> str: await stop.guard(gate.wait()) if text == "refuse": raise AuthenticationError("invalid x-api-key") + if text.startswith("refund "): + # A tool call: the model asks for an approval, and waits for it. + call_id = f"call-{text.split()[1]}" + args = {"order": text.split()[1], "amount": 5} + transcript.call(call_id, approval.name, args) + result = await run_tool(approval, args, ToolCall(ctx, call_id)) + transcript.result(call_id, result) + transcript.say("assistant", f"refund: {result}") + return f"refund: {result}" transcript.say("assistant", f"echo: {text}") return f"echo: {text}" @@ -171,3 +189,24 @@ async def test_a_refused_message_fails_with_424_and_the_next_is_answered(harness assert refused.status_code == 424 assert (await call(harness, "status", key)).json()["answering"] is None assert (await call(harness, "chat", key, "fine")).json() == "echo: fine" + + +async def test_a_tool_waits_for_a_signal_which_the_agents_signal_view_resolves(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + invocation = await send(harness, key, "refund o7") + + async def waiting() -> bool: + return len((await call(harness, "pending", key)).json()) == 1 + + await until(waiting) + [wait] = (await call(harness, "pending", key)).json() + assert wait["name"] == "approval" and wait["input"] == {"order": "o7", "amount": 5} + assert wait["id"] == f"{invocation}:call-o7" + assert (await call(harness, "signal", key, {"name": "approval", "value": "approved"})).json() == { + "resolved": wait["id"] + } + assert (await attach(harness, invocation)).json() == "refund: approved" + page = (await call(harness, "messages", key, {})).json() + [tool_call] = page["messages"][1]["parts"] + assert tool_call["name"] == "wait_for_approval" and tool_call["output"] == "approved" + assert (await call(harness, "pending", key)).json() == [] diff --git a/tests/agents_tools.py b/tests/agents_tools.py new file mode 100644 index 00000000..23d59b26 --- /dev/null +++ b/tests/agents_tools.py @@ -0,0 +1,153 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""Tools: from functions, and the tools that are Restate itself, over an in-memory copy of an agent's state.""" + +# pylint: disable=missing-function-docstring,missing-class-docstring + +import json +import typing +from typing import Any + +import pydantic +import pytest + +import restate +from restate import sessions +from restate.sessions import internal +from restate.agents._common import ToolError +from restate.agents._object import AgentObject +from restate.agents._tools import Tool, ToolCall, call, signal, subagent, tool +from restate.context import Request +from sessions_database import Backend # pylint: disable=import-error +from sessions_object import Host, request # pylint: disable=import-error + + +@pytest.fixture +def anyio_backend(): + return "asyncio" + + +# --- The targets the tools call ------------------------------------------------------------------- + + +class Query(pydantic.BaseModel): + order: str + + +orders = restate.VirtualObject("orders") + + +@orders.handler() +async def status(ctx: restate.ObjectContext, query: Query) -> dict[str, str]: + return {"order": query.order, "status": "shipped"} + + +helper = AgentObject("helper", description="Looks things up") + + +# --- An agent whose handler runs a tool, for a tool call ------------------------------------------ + +runner = AgentObject("runner") +tools: dict[str, Tool] = {} + + +@runner.handler +async def use(ctx: sessions.SessionContext, step: dict[str, Any]) -> str: + return await tools[step["tool"]].run(step["args"], ToolCall(ctx, step["call"])) + + +async def run_tool(host: Host, name: str, args: dict[str, Any], call_id: str, req: Request | None = None) -> str: + return await internal.invoke(runner, "use", host, req or request(), {"tool": name, "args": args, "call": call_id}) + + +# --- tool ----------------------------------------------------------------------------------------- + + +@tool +async def get_weather(ctx: sessions.SessionContext, call_id: str, city: str, unit: str = "C") -> str: + """The current weather in a city.""" + if city == "Atlantis": + raise restate.TerminalError("no such city") + if city == "Crash": + raise RuntimeError("the weather API is down") + return f"Sunny, 21°{unit} in {city} ({ctx.key}, {call_id})" + + +def test_a_tool_takes_its_input_from_the_function_and_its_description_from_the_docstring(): + assert get_weather.name == "get_weather" + assert get_weather.description == "The current weather in a city." + assert set(get_weather.input_schema["properties"]) == {"city", "unit"} + assert get_weather.input_schema["required"] == ["city"] + + +async def test_a_tool_answers_with_text_and_its_errors_go_to_the_model_or_fail_the_attempt(): + tools["weather"] = get_weather + host = Host(Backend(), service="runner", key="alice") + assert await run_tool(host, "weather", {"city": "Lisbon"}, "c1") == "Sunny, 21°C in Lisbon (alice, c1)" + bad = await run_tool(host, "weather", {"town": "Lisbon"}, "c2") + assert bad.startswith("Error: the arguments do not fit the tool's input") + assert await run_tool(host, "weather", {"city": "Atlantis"}, "c3") == "Error: no such city" + with pytest.raises(ToolError, match="the tool get_weather failed: the weather API is down"): + await run_tool(host, "weather", {"city": "Crash"}, "c4") + + +# --- call ----------------------------------------------------------------------------------------- + + +async def test_call_lets_the_model_choose_the_key_and_makes_the_call_once_per_tool_call(): + tools["status"] = call(status, description="The status of an order.") + assert tools["status"].name == "orders_status" + schema = tools["status"].input_schema + assert set(schema["properties"]) == {"key", "order"} and set(schema["required"]) == {"key", "order"} + + host = Host(Backend(), service="runner", key="alice", answers={"status": b'{"order": "o1", "status": "shipped"}'}) + req = request() + answer = await run_tool(host, "status", {"key": "eu", "order": "o1"}, "t1", req) + assert json.loads(answer) == {"order": "o1", "status": "shipped"} + [(service, handler, arg, key, idempotency_key)] = host.calls + assert (service, handler, key, idempotency_key) == ("orders", "status", "eu", f"{req.id}:t1") + assert json.loads(arg) == {"order": "o1"} + + # Another attempt of the invocation, whose framework makes the call again under a new id: the same call. + await run_tool(host, "status", {"key": "eu", "order": "o1"}, "t2", req) + assert host.calls[-1][4] == f"{req.id}:t1" + + +# --- signal and subagent -------------------------------------------------------------------------- + + +class Refund(pydantic.BaseModel): + order: str + amount: float + + +async def test_signal_waits_at_the_companion_of_the_agent_with_the_tool_calls_key(): + tools["approval"] = signal("approval", description="Wait for a manager's approval.", input=Refund) + assert tools["approval"].name == "wait_for_approval" + assert set(tools["approval"].input_schema["properties"]) == {"order", "amount"} + host = Host(Backend(), service="runner", key="alice", answers={"wait": b'"approved"'}) + req = request() + assert await run_tool(host, "approval", {"order": "o1", "amount": 5}, "w1", req) == "approved" + [(service, handler, arg, key, idempotency_key)] = host.calls + assert (service, handler, key, idempotency_key) == ("signals", "wait", "runner:alice", f"{req.id}:w1") + wait = typing.cast(dict[str, Any], json.loads(arg)) + assert wait == {"id": f"{req.id}:w1", "name": "approval", "input": {"order": "o1", "amount": 5}} + + +async def test_subagent_hands_a_task_to_the_agent_keyed_by_this_key_and_the_tasks_name(): + tools["helper"] = subagent(helper) + assert "Hand a task to the helper agent (Looks things up)" in tools["helper"].description + host = Host(Backend(), service="runner", key="alice", answers={"chat": b'"found it"'}) + assert await run_tool(host, "helper", {"name": "lookup", "task": "Find it"}, "s1") == "found it" + [(service, handler, arg, key, _)] = host.calls + assert (service, handler, key, json.loads(arg)) == ("helper", "chat", "alice.lookup", "Find it") + bad = await run_tool(host, "helper", {"name": "Not A Name", "task": "x"}, "s2") + assert bad.startswith("Error: the arguments do not fit") From a8726f3ba813ebc9200650c5e1957dff96d340ab Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 16:41:09 +0100 Subject: [PATCH 23/45] Keep the ASGI server's SIGTERM handler when the endpoint adds its own loop.add_signal_handler keeps one handler per signal, so the endpoint's replaced the server's: on Python 3.12+ hypercorn no longer stopped on SIGTERM. The endpoint now chains a signal-module handler to the previous one; the loop's own handlers still run through its wakeup fd. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/server.py | 40 +++++++++++++++++++++++++++------ tests/server_sigterm.py | 48 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 7 deletions(-) create mode 100644 tests/server_sigterm.py diff --git a/python/restate/server.py b/python/restate/server.py index 77cc1830..f4208bdf 100644 --- a/python/restate/server.py +++ b/python/restate/server.py @@ -15,7 +15,7 @@ import logging import signal import weakref -from typing import Callable, Dict, List, Optional, Set, TypedDict, Literal +from typing import Any, Callable, Dict, List, Optional, Set, TypedDict, Literal from restate.discovery import compute_discovery_json from restate.endpoint import Endpoint @@ -39,13 +39,18 @@ def install_shutdown_hook( - loop: asyncio.AbstractEventLoop, hook: Callable[[], None], signum: int = signal.SIGTERM + loop: asyncio.AbstractEventLoop, hook: Callable[[], None], signum: int = signal.SIGTERM, *, chain: bool = False ) -> None: """ Run ``hook`` when ``signum`` (default SIGTERM) is received while ``loop`` runs. - Raises what ``loop.add_signal_handler`` raises when signals cannot be handled here - (Windows, or a non-main thread). + With ``chain``, the handler the process already has for ``signum`` still runs too: the ASGI + server's (hypercorn's, uvicorn's), which stops it. ``loop.add_signal_handler`` keeps a single + handler per signal, and would replace it. Without, ``hook``'s dispatcher replaces it, for a + program that owns its signals (``restate.connect``). The first hook of a loop and signal decides. + + Raises what ``loop.add_signal_handler`` (or ``signal.signal``, with ``chain``) raises when + signals cannot be handled here (Windows, or a non-main thread). """ hooks = _shutdown_hooks.setdefault(loop, {}) if signum not in hooks: @@ -58,11 +63,33 @@ def dispatch() -> None: except Exception: # pylint: disable=broad-except logger.exception("Shutdown hook failed") - loop.add_signal_handler(signum, dispatch) + if chain: + _chain_signal(loop, signum, dispatch) + else: + loop.add_signal_handler(signum, dispatch) hooks[signum] = registered hooks[signum].append(hook) +def _chain_signal(loop: asyncio.AbstractEventLoop, signum: int, on_signal: Callable[[], None]) -> None: + """ + Call ``on_signal`` on ``loop`` when ``signum`` arrives, then the signal's previous handler. + + A signal-module handler leaves the loop's own handlers working too (they are woken by its wakeup + fd). Python's default handlers are not chained: the default SIGINT handler would raise + KeyboardInterrupt into whatever runs. + """ + previous = signal.getsignal(signum) + + def handler(received: int, frame: Any) -> None: + if not loop.is_closed(): + loop.call_soon_threadsafe(on_signal) + if callable(previous) and previous is not signal.default_int_handler: + previous(received, frame) + + signal.signal(signum, handler) + + X_RESTATE_SERVER = header_to_binary([("x-restate-server", f"restate-sdk-python/{SDK_VERSION}")]) @@ -295,9 +322,8 @@ def _on_sigterm() -> None: async def app(scope: Scope, receive: Receive, send: Send): nonlocal sigterm_installed if not sigterm_installed: - loop = asyncio.get_running_loop() try: - install_shutdown_hook(loop, _on_sigterm, signal.SIGTERM) + install_shutdown_hook(asyncio.get_running_loop(), _on_sigterm, signal.SIGTERM, chain=True) except (NotImplementedError, RuntimeError, ValueError): pass # Windows or non-main thread sigterm_installed = True diff --git a/tests/server_sigterm.py b/tests/server_sigterm.py new file mode 100644 index 00000000..55c03062 --- /dev/null +++ b/tests/server_sigterm.py @@ -0,0 +1,48 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""The endpoint's SIGTERM handler leaves the ASGI server's own in place.""" + +import asyncio +import os +import signal +import sys + +import pytest + +from restate.server import install_shutdown_hook + +pytestmark = pytest.mark.skipif(sys.platform == "win32", reason="POSIX signals") + + +def test_sigterm_reaches_the_server_and_the_endpoint(): + async def main() -> list[str]: + loop = asyncio.get_running_loop() + heard: list[str] = [] + both = asyncio.Event() + + def hear(who: str) -> None: + heard.append(who) + if len(heard) == 2: + both.set() + + # The server's handler, as hypercorn installs it, and then the endpoint's, on its first request. + loop.add_signal_handler(signal.SIGTERM, hear, "server") + install_shutdown_hook(loop, lambda: hear("endpoint"), signal.SIGTERM, chain=True) + os.kill(os.getpid(), signal.SIGTERM) + await asyncio.wait_for(both.wait(), 5) + loop.remove_signal_handler(signal.SIGTERM) + return heard + + previous = signal.getsignal(signal.SIGTERM) + try: + assert sorted(asyncio.run(main())) == ["endpoint", "server"] + finally: + signal.signal(signal.SIGTERM, previous) From b263b010cfca06a66ea7cae43addea6066a22aa7 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 16:41:09 +0100 Subject: [PATCH 24/45] examples-agents: the wallet tutorial (01-08) and notes, and sessions.serve Ports sdk-sessions' template lessons and sqlexperiment's notes. Each file runs on its own with sessions.serve (hypercorn, imported lazily), which stops on SIGINT/SIGTERM within hypercorn's graceful timeout. tests/examples_agents.py runs every example's commands against a real Restate server; the directory is linted and typechecked with the rest. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- Justfile | 2 + examples-agents/README.md | 74 +++++++++++ examples-agents/basic/notes.py | 33 +++++ examples-agents/tutorial/01_object.py | 34 +++++ examples-agents/tutorial/02_database.py | 32 +++++ examples-agents/tutorial/03_views.py | 44 +++++++ examples-agents/tutorial/04_schema.py | 57 +++++++++ examples-agents/tutorial/05_transaction.py | 72 +++++++++++ examples-agents/tutorial/06_retries.py | 100 +++++++++++++++ examples-agents/tutorial/07_calls.py | 88 +++++++++++++ examples-agents/tutorial/08_sends.py | 90 +++++++++++++ pyproject.toml | 1 + python/restate/sessions/__init__.py | 3 +- python/restate/sessions/_serve.py | 69 +++++++++- tests/examples_agents.py | 142 +++++++++++++++++++++ 15 files changed, 839 insertions(+), 2 deletions(-) create mode 100644 examples-agents/README.md create mode 100644 examples-agents/basic/notes.py create mode 100644 examples-agents/tutorial/01_object.py create mode 100644 examples-agents/tutorial/02_database.py create mode 100644 examples-agents/tutorial/03_views.py create mode 100644 examples-agents/tutorial/04_schema.py create mode 100644 examples-agents/tutorial/05_transaction.py create mode 100644 examples-agents/tutorial/06_retries.py create mode 100644 examples-agents/tutorial/07_calls.py create mode 100644 examples-agents/tutorial/08_sends.py create mode 100644 tests/examples_agents.py diff --git a/Justfile b/Justfile index 6645a9f2..b226cca9 100644 --- a/Justfile +++ b/Justfile @@ -24,12 +24,14 @@ lint-rust: typecheck-pyright: PYRIGHT_PYTHON_IGNORE_WARNINGS=1 uv run pyright python/ PYRIGHT_PYTHON_IGNORE_WARNINGS=1 uv run pyright examples/ + PYRIGHT_PYTHON_IGNORE_WARNINGS=1 uv run pyright examples-agents/ PYRIGHT_PYTHON_IGNORE_WARNINGS=1 uv run pyright tests PYRIGHT_PYTHON_IGNORE_WARNINGS=1 uv run pyright test-services/ typecheck-mypy: uv run -m mypy --check-untyped-defs --ignore-missing-imports --implicit-optional python/ uv run -m mypy --check-untyped-defs --ignore-missing-imports --implicit-optional examples/ + uv run -m mypy --check-untyped-defs --ignore-missing-imports --implicit-optional examples-agents/ uv run -m mypy --check-untyped-defs --ignore-missing-imports --implicit-optional tests/ typecheck: typecheck-pyright typecheck-mypy diff --git a/examples-agents/README.md b/examples-agents/README.md new file mode 100644 index 00000000..1ae4a9ae --- /dev/null +++ b/examples-agents/README.md @@ -0,0 +1,74 @@ +# Session objects and agents: examples + +Examples of [`restate.sessions`](../python/restate/sessions/__init__.py), +virtual objects with a SQL database per key, and +[`restate.agents`](../python/restate/agents/__init__.py), AI agents in them. +They are the Python counterparts of +[sdk-sessions](https://github.com/restatedev/sdk-sessions)' tutorial and +examples. Each is a file that runs on its own. + +## Run an example + +From the repository's root, once (it builds the SDK, with its native +extension): + +```sh +uv sync --all-extras --all-packages +``` + +Then run an example, start Restate, and tell Restate where the example is: + +```sh +uv run examples-agents/tutorial/01_object.py # the example, on port 9080 +npx @restatedev/restate-server # Restate, in another terminal +npx @restatedev/restate -y deployments register --force http://localhost:9080 # tell Restate where it is +curl localhost:8080/wallet/alice/hello +``` + +Every example serves on port 9080 (or `PORT`): stop one before you start the +next, and register again. The comment at the top of each file lists the +commands to try it. + +## Tutorial: a wallet, from a database to agents + +One application, built one idea per file: a wallet for every user. It starts +as a Restate virtual object, gets its own SQL database, and learns to move +money safely between wallets. Keep Restate running from one lesson to the +next: the wallets you filled in one lesson are still there in the next, as a +real application keeps its data across releases. + +### Part 1 · A database per key + +| Lesson | You learn | +| -------------------------------------------- | ------------------------------------------------------ | +| [`01_object.py`](tutorial/01_object.py) | A virtual object: handlers that Restate calls by key | +| [`02_database.py`](tutorial/02_database.py) | `ctx.db`: every wallet has its own SQL database | +| [`03_views.py`](tutorial/03_views.py) | Views: reading the database while a handler is running | +| [`04_schema.py`](tutorial/04_schema.py) | Changing the tables of databases that already exist | + +### Part 2 · Changes that happen once + +| Lesson | You learn | +| -------------------------------------------------- | ------------------------------------------------------------- | +| [`05_transaction.py`](tutorial/05_transaction.py) | Transactions: spending fails as a whole, leaving no trace | +| [`06_retries.py`](tutorial/06_retries.py) | Retries run a handler again: an id makes a deposit count once | +| [`07_calls.py`](tutorial/07_calls.py) | Calling another object, with an idempotency key | +| [`08_sends.py`](tutorial/08_sends.py) | Transfers: a send that is delivered only if the debit commits | + +Part 3, agents in the wallet (a prompt, an assistant, its tools, a person in +the loop, a subagent), comes with the agents' framework integrations. + +## Basic + +| File | What | +| ----------------------------- | ----------------------------------------------- | +| [`notes.py`](basic/notes.py) | SQL in a virtual object, with `ctx.db`, and no agent | + +## Tips + +- Restate's UI, at , shows every request and what it + did, including the retries of lesson 06. +- To start over with empty wallets, stop Restate and delete the + `restate-data` directory it created. +- `tests/examples_agents.py` runs each example's commands against a real + Restate server. diff --git a/examples-agents/basic/notes.py b/examples-agents/basic/notes.py new file mode 100644 index 00000000..47d37b76 --- /dev/null +++ b/examples-agents/basic/notes.py @@ -0,0 +1,33 @@ +""" +Notes: SQL in a virtual object, with ``ctx.db``, and no agent. + +Each key (``alice``, ``bob``) has a SQLite-compatible database of its own, kept +in the object's state: no database to run. ``add`` writes, one at a time per +key; ``list`` is a view, which reads while ``add`` writes:: + + uv run examples-agents/basic/notes.py # on port 9080 + npx @restatedev/restate-server # Restate, in another terminal + npx @restatedev/restate -y deployments register --force http://localhost:9080 + curl localhost:8080/notes/alice/add --json '"Buy milk"' + curl localhost:8080/notes/alice/list # ["Buy milk"] +""" + +from restate import sessions + +# The schema of each key's database: its steps, applied once per key (a step added later runs then). +notes = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) + + +@notes.handler +async def add(ctx: sessions.SessionContext, note: str) -> None: + await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [note]) + + +@notes.view(name="list") +async def list_notes(ctx: sessions.ViewContext) -> list[str]: + """A view: it reads while ``add`` writes; a key with no notes yet has none.""" + return [str(row["text"]) for row in await ctx.db.query("SELECT text FROM notes ORDER BY rowid")] + + +if __name__ == "__main__": + sessions.serve(notes) diff --git a/examples-agents/tutorial/01_object.py b/examples-agents/tutorial/01_object.py new file mode 100644 index 00000000..fc4b8d2c --- /dev/null +++ b/examples-agents/tutorial/01_object.py @@ -0,0 +1,34 @@ +""" +01 · A virtual object + +Throughout this tutorial we build a wallet. Every user has one, and Restate +finds it by key: ``wallet/alice`` and ``wallet/bob`` are two instances of the +same object. Restate runs the requests for one key one at a time, so a wallet +never handles two of them at once. + +This process serves the wallet's handlers. Restate calls them over HTTP, and +you call them through Restate:: + + uv run examples-agents/tutorial/01_object.py # serves the wallet on port 9080 + npx @restatedev/restate-server # Restate, in another terminal + npx @restatedev/restate -y deployments register --force http://localhost:9080 # tells Restate where the wallet is + curl localhost:8080/wallet/alice/hello # "Hello, alice!" + +Every lesson serves on port 9080, so stop this one before you start the next, +and register again. +""" + +from restate import sessions + +# A session object: a virtual object that gets a database in the next lesson. +wallet = sessions.SessionObject("wallet") + + +@wallet.handler +async def hello(ctx: sessions.SessionContext) -> str: + # `ctx.key` is the key the request was sent to: the wallet's owner. + return f"Hello, {ctx.key}!" + + +if __name__ == "__main__": + sessions.serve(wallet) diff --git a/examples-agents/tutorial/02_database.py b/examples-agents/tutorial/02_database.py new file mode 100644 index 00000000..04521cc6 --- /dev/null +++ b/examples-agents/tutorial/02_database.py @@ -0,0 +1,32 @@ +""" +02 · A database per key + +Each key of a session object has its own SQLite database, ``ctx.db``. It is +stored in Restate together with the object's state, so there is no database +server to run. ``schema`` lists the statements that create the tables; they +run once per key, the first time that key is used. + +The wallet records each deposit as a row, and the balance is their sum. Bob's +wallet is another key, so it has a database of its own:: + + uv run examples-agents/tutorial/02_database.py # and register it (see 01) + curl localhost:8080/wallet/alice/deposit --json 10 # 10 + curl localhost:8080/wallet/alice/deposit --json 5 # 15 + curl localhost:8080/wallet/bob/deposit --json 7 # 7 +""" + +from restate import sessions + +wallet = sessions.SessionObject("wallet", schema=["CREATE TABLE entries (amount INTEGER NOT NULL)"]) + + +@wallet.handler +async def deposit(ctx: sessions.SessionContext, amount: int) -> int: + await ctx.db.execute("INSERT INTO entries (amount) VALUES (?)", [amount]) + # `get` returns the first row as a dict, with a key per column. + row = await ctx.db.get("SELECT sum(amount) AS balance FROM entries") + return int(row["balance"] or 0) if row else 0 + + +if __name__ == "__main__": + sessions.serve(wallet) diff --git a/examples-agents/tutorial/03_views.py b/examples-agents/tutorial/03_views.py new file mode 100644 index 00000000..479393a2 --- /dev/null +++ b/examples-agents/tutorial/03_views.py @@ -0,0 +1,44 @@ +""" +03 · Views + +Handlers of one key run one at a time, so their writes never race. The price +is that a slow handler makes every other request to that key wait. Reading +doesn't need to wait: a view runs as soon as it is called, even while a handler +of the same key is running, and reads what has been committed so far. A view +can't write:: + + uv run examples-agents/tutorial/03_views.py # and register it (see 01) + curl localhost:8080/wallet/alice/deposit --json 20 + curl localhost:8080/wallet/alice/balance # 20, plus what lesson 02 deposited + curl localhost:8080/wallet/alice/history # every deposit, oldest first +""" + +from restate import sessions + +wallet = sessions.SessionObject("wallet", schema=["CREATE TABLE entries (amount INTEGER NOT NULL)"]) + + +async def balance_of(db: sessions.DatabaseReader) -> int: + """The sum of the entries. A ``DatabaseReader`` is any database you can read: a handler's, a view's or a transaction's.""" + row = await db.get("SELECT coalesce(sum(amount), 0) AS balance FROM entries") + return int(row["balance"] or 0) if row else 0 + + +@wallet.handler +async def deposit(ctx: sessions.SessionContext, amount: int) -> int: + await ctx.db.execute("INSERT INTO entries (amount) VALUES (?)", [amount]) + return await balance_of(ctx.db) + + +@wallet.view +async def balance(ctx: sessions.ViewContext) -> int: + return await balance_of(ctx.db) + + +@wallet.view +async def history(ctx: sessions.ViewContext) -> list[int]: + return [int(row["amount"] or 0) for row in await ctx.db.query("SELECT amount FROM entries ORDER BY rowid")] + + +if __name__ == "__main__": + sessions.serve(wallet) diff --git a/examples-agents/tutorial/04_schema.py b/examples-agents/tutorial/04_schema.py new file mode 100644 index 00000000..ce02a58c --- /dev/null +++ b/examples-agents/tutorial/04_schema.py @@ -0,0 +1,57 @@ +""" +04 · Changing the schema + +Alice's database already exists, so editing its CREATE TABLE statement would +change nothing for her. That's why ``schema`` is a list of steps: to change the +tables, add a step at the end. Each key's database runs the steps it hasn't run +yet, the next time it is used. Never edit or remove a step that may have run +somewhere; add a new one instead. + +Here, deposits get a note:: + + uv run examples-agents/tutorial/04_schema.py # and register it (see 01) + curl localhost:8080/wallet/alice/deposit --json '{"amount": 50, "note": "birthday"}' + curl localhost:8080/wallet/alice/history # the earlier deposits have no note +""" + +from pydantic import BaseModel + +from restate import sessions + +wallet = sessions.SessionObject( + "wallet", + schema=[ + "CREATE TABLE entries (amount INTEGER NOT NULL)", # from lesson 02 + "ALTER TABLE entries ADD COLUMN note TEXT", # new: runs once on each wallet's database + ], +) + + +class Deposit(BaseModel): + amount: int + note: str | None = None + + +async def balance_of(db: sessions.DatabaseReader) -> int: + row = await db.get("SELECT coalesce(sum(amount), 0) AS balance FROM entries") + return int(row["balance"] or 0) if row else 0 + + +@wallet.handler +async def deposit(ctx: sessions.SessionContext, money: Deposit) -> int: + await ctx.db.execute("INSERT INTO entries (amount, note) VALUES (?, ?)", [money.amount, money.note]) + return await balance_of(ctx.db) + + +@wallet.view +async def balance(ctx: sessions.ViewContext) -> int: + return await balance_of(ctx.db) + + +@wallet.view +async def history(ctx: sessions.ViewContext) -> list[sessions.Row]: + return await ctx.db.query("SELECT amount, note FROM entries ORDER BY rowid") + + +if __name__ == "__main__": + sessions.serve(wallet) diff --git a/examples-agents/tutorial/05_transaction.py b/examples-agents/tutorial/05_transaction.py new file mode 100644 index 00000000..7ef51ac8 --- /dev/null +++ b/examples-agents/tutorial/05_transaction.py @@ -0,0 +1,72 @@ +""" +05 · Transactions + +Each statement on ``ctx.db`` commits on its own. When several changes must +happen together or not at all, put them in ``ctx.db.transaction()``: it commits +when its body ends, and rolls everything back if it raises. + +Spending records the payment first, then checks the balance. If the balance +went below zero, it raises, and the payment is rolled back as if it never +happened. + +It raises a ``TerminalError``, which tells Restate not to retry: the request +fails, and the caller gets the error. Restate retries any other error, which +suits failures that may go away, such as a lost connection:: + + uv run examples-agents/tutorial/05_transaction.py # and register it (see 01) + curl localhost:8080/wallet/alice/deposit --json '{"amount": 30, "note": "salary"}' + curl localhost:8080/wallet/alice/spend --json '{"amount": 12, "note": "lunch"}' # the balance after lunch + curl localhost:8080/wallet/alice/spend --json '{"amount": 1000, "note": "a boat"}' # fails: not enough money + curl localhost:8080/wallet/alice/history # lunch, but no boat +""" + +from pydantic import BaseModel + +import restate +from restate import sessions + +wallet = sessions.SessionObject( + "wallet", + schema=["CREATE TABLE entries (amount INTEGER NOT NULL)", "ALTER TABLE entries ADD COLUMN note TEXT"], +) + + +class Money(BaseModel): + amount: int + note: str | None = None + + +async def balance_of(db: sessions.DatabaseReader) -> int: + row = await db.get("SELECT coalesce(sum(amount), 0) AS balance FROM entries") + return int(row["balance"] or 0) if row else 0 + + +@wallet.handler +async def deposit(ctx: sessions.SessionContext, money: Money) -> int: + await ctx.db.execute("INSERT INTO entries (amount, note) VALUES (?, ?)", [money.amount, money.note]) + return await balance_of(ctx.db) + + +@wallet.handler +async def spend(ctx: sessions.SessionContext, money: Money) -> int: + async with ctx.db.transaction() as tx: + await tx.execute("INSERT INTO entries (amount, note) VALUES (?, ?)", [-money.amount, money.note]) + balance = await balance_of(tx) + # Raising rolls back the insert above. + if balance < 0: + raise restate.TerminalError("Not enough money") + return balance + + +@wallet.view +async def balance(ctx: sessions.ViewContext) -> int: + return await balance_of(ctx.db) + + +@wallet.view +async def history(ctx: sessions.ViewContext) -> list[sessions.Row]: + return await ctx.db.query("SELECT amount, note FROM entries ORDER BY rowid") + + +if __name__ == "__main__": + sessions.serve(wallet) diff --git a/examples-agents/tutorial/06_retries.py b/examples-agents/tutorial/06_retries.py new file mode 100644 index 00000000..f42b9dc3 --- /dev/null +++ b/examples-agents/tutorial/06_retries.py @@ -0,0 +1,100 @@ +""" +06 · Retries + +When a handler fails partway, because the process crashed or a connection +dropped, Restate runs it again from the first line. Whatever the failed attempt +committed is still in the database, so every handler must be safe to run twice. + +The way to do that is an id that is the same on every attempt. A deposit now +comes with an id from the caller, and ``INSERT OR IGNORE`` skips a deposit whose +id is already there: running it twice, or sending it twice, counts it once. +Spending has no id from the caller, so it uses ``ctx.request.id``, the id +Restate gives each request, which also stays the same across attempts. + +To show it, this lesson fails the first attempt of every deposit right after +the insert, the way a crash would. Restate retries it, and the balance still +counts the deposit once:: + + uv run examples-agents/tutorial/06_retries.py # and register it (see 01) + curl localhost:8080/wallet/alice/deposit --json '{"id": "d1", "amount": 10, "note": "refund"}' # counted once, though the first attempt failed + curl localhost:8080/wallet/alice/deposit --json '{"id": "d1", "amount": 10, "note": "refund"}' # sent again: the balance stays the same +""" + +from pydantic import BaseModel + +import restate +from restate import sessions + +wallet = sessions.SessionObject( + "wallet", + schema=[ + "CREATE TABLE entries (amount INTEGER NOT NULL)", + "ALTER TABLE entries ADD COLUMN note TEXT", + # New: an id per entry. Entries from earlier lessons have none, which a unique index allows. + "ALTER TABLE entries ADD COLUMN id TEXT", + "CREATE UNIQUE INDEX entries_by_id ON entries (id)", + ], +) + + +class Deposit(BaseModel): + id: str + amount: int + note: str | None = None + + +class Money(BaseModel): + amount: int + note: str | None = None + + +failed: set[str] = set() + + +def fail_first_attempt(ctx: sessions.SessionContext) -> None: + """Fail the first attempt of a request, as a crash would. Restate then retries it.""" + if ctx.request.id in failed: + return + failed.add(ctx.request.id) + raise RuntimeError("Simulated crash, after the deposit was committed") + + +async def balance_of(db: sessions.DatabaseReader) -> int: + row = await db.get("SELECT coalesce(sum(amount), 0) AS balance FROM entries") + return int(row["balance"] or 0) if row else 0 + + +@wallet.handler +async def deposit(ctx: sessions.SessionContext, money: Deposit) -> int: + await ctx.db.execute( + "INSERT OR IGNORE INTO entries (id, amount, note) VALUES (?, ?, ?)", [money.id, money.amount, money.note] + ) + fail_first_attempt(ctx) + return await balance_of(ctx.db) + + +@wallet.handler +async def spend(ctx: sessions.SessionContext, money: Money) -> int: + async with ctx.db.transaction() as tx: + await tx.execute( + "INSERT OR IGNORE INTO entries (id, amount, note) VALUES (?, ?, ?)", + [ctx.request.id, -money.amount, money.note], + ) + balance = await balance_of(tx) + if balance < 0: + raise restate.TerminalError("Not enough money") + return balance + + +@wallet.view +async def balance(ctx: sessions.ViewContext) -> int: + return await balance_of(ctx.db) + + +@wallet.view +async def history(ctx: sessions.ViewContext) -> list[sessions.Row]: + return await ctx.db.query("SELECT amount, note FROM entries ORDER BY rowid") + + +if __name__ == "__main__": + sessions.serve(wallet) diff --git a/examples-agents/tutorial/07_calls.py b/examples-agents/tutorial/07_calls.py new file mode 100644 index 00000000..5cc8520f --- /dev/null +++ b/examples-agents/tutorial/07_calls.py @@ -0,0 +1,88 @@ +""" +07 · Calling another object + +Friends should be able to pay alice without knowing her wallet's key, so she +picks a handle, such as @ace. Handles are a second session object: each handle +is a key with its own database, which records the wallet that owns it. The +first wallet to claim a handle gets it. + +``ctx.rpc(claim_handle, key=name)(wallet)`` calls that object through Restate, +and waits for the answer. A retried handler (06) makes its calls again, so a +call that changes something gets an idempotency key that is the same on every +attempt. Restate then answers a repeated call with the first call's result, +instead of running it again:: + + uv run examples-agents/tutorial/07_calls.py # and register it (see 01) + curl localhost:8080/wallet/alice/claim --json '"ace"' # "alice is now @ace" + curl localhost:8080/wallet/bob/claim --json '"ace"' # fails: @ace belongs to alice + curl localhost:8080/handle/ace/owner # "alice" +""" + +from pydantic import BaseModel + +import restate +from restate import sessions + +handle = sessions.SessionObject("handle", schema=["CREATE TABLE owner (wallet TEXT NOT NULL)"]) + + +@handle.handler(name="claim") +async def claim_handle(ctx: sessions.SessionContext, wallet: str) -> str: + """Give the handle to ``wallet``, unless another wallet has it, and answer with its owner.""" + await ctx.db.execute("INSERT INTO owner (wallet) SELECT ? WHERE NOT EXISTS (SELECT 1 FROM owner)", [wallet]) + owner = await ctx.db.get("SELECT wallet FROM owner") + assert owner is not None # the insert above made sure there is one + return str(owner["wallet"]) + + +@handle.view +async def owner(ctx: sessions.ViewContext) -> str | None: + row = await ctx.db.get("SELECT wallet FROM owner") + return str(row["wallet"]) if row else None + + +wallet = sessions.SessionObject( + "wallet", + schema=[ + "CREATE TABLE entries (amount INTEGER NOT NULL)", + "ALTER TABLE entries ADD COLUMN note TEXT", + "ALTER TABLE entries ADD COLUMN id TEXT", + "CREATE UNIQUE INDEX entries_by_id ON entries (id)", + ], +) + + +class Deposit(BaseModel): + id: str + amount: int + note: str | None = None + + +async def balance_of(db: sessions.DatabaseReader) -> int: + row = await db.get("SELECT coalesce(sum(amount), 0) AS balance FROM entries") + return int(row["balance"] or 0) if row else 0 + + +@wallet.handler +async def deposit(ctx: sessions.SessionContext, money: Deposit) -> int: + await ctx.db.execute( + "INSERT OR IGNORE INTO entries (id, amount, note) VALUES (?, ?, ?)", [money.id, money.amount, money.note] + ) + return await balance_of(ctx.db) + + +@wallet.handler +async def claim(ctx: sessions.SessionContext, name: str) -> str: + owner = await ctx.rpc(claim_handle, key=name, idempotency_key=f"{ctx.request.id}:claim")(ctx.key) + if owner != ctx.key: + raise restate.TerminalError(f"@{name} belongs to {owner}") + return f"{ctx.key} is now @{name}" + + +@wallet.view +async def balance(ctx: sessions.ViewContext) -> int: + return await balance_of(ctx.db) + + +if __name__ == "__main__": + sessions.serve(wallet, handle) diff --git a/examples-agents/tutorial/08_sends.py b/examples-agents/tutorial/08_sends.py new file mode 100644 index 00000000..addaa38f --- /dev/null +++ b/examples-agents/tutorial/08_sends.py @@ -0,0 +1,90 @@ +""" +08 · Sends that commit with the transaction + +A transfer takes money out of alice's wallet and puts it into bob's. Alice's +wallet doesn't need an answer from bob's, so instead of calling it (07), it +sends a message, which doesn't wait. Waiting could even get stuck: if bob paid +alice at the same moment, each wallet would wait for the other, since each +runs one request at a time. + +A send made with a transaction's ``tx.send`` belongs to that transaction: +Restate delivers it once the transaction commits, and never if it rolls back. +So the money leaves alice and arrives at bob together, or not at all:: + + uv run examples-agents/tutorial/08_sends.py # and register it (see 01) + curl localhost:8080/wallet/alice/deposit --json '{"id": "d2", "amount": 100}' + curl localhost:8080/wallet/alice/transfer --json '{"to": "bob", "amount": 30}' # alice's balance after the transfer + curl localhost:8080/wallet/bob/history # ends with 30, from alice + curl localhost:8080/wallet/alice/transfer --json '{"to": "bob", "amount": 5000}' # fails, and bob gets nothing +""" + +from pydantic import BaseModel + +import restate +from restate import sessions + +wallet = sessions.SessionObject( + "wallet", + schema=[ + "CREATE TABLE entries (amount INTEGER NOT NULL)", + "ALTER TABLE entries ADD COLUMN note TEXT", + "ALTER TABLE entries ADD COLUMN id TEXT", + "CREATE UNIQUE INDEX entries_by_id ON entries (id)", + ], +) + + +class Deposit(BaseModel): + id: str + amount: int + note: str | None = None + + +class Transfer(BaseModel): + to: str + amount: int + + +async def balance_of(db: sessions.DatabaseReader) -> int: + row = await db.get("SELECT coalesce(sum(amount), 0) AS balance FROM entries") + return int(row["balance"] or 0) if row else 0 + + +@wallet.handler +async def deposit(ctx: sessions.SessionContext, money: Deposit) -> int: + await ctx.db.execute( + "INSERT OR IGNORE INTO entries (id, amount, note) VALUES (?, ?, ?)", [money.id, money.amount, money.note] + ) + return await balance_of(ctx.db) + + +@wallet.handler +async def transfer(ctx: sessions.SessionContext, money: Transfer) -> int: + async with ctx.db.transaction() as tx: + # The transfer's id, the same on every attempt (06). Bob's wallet records it too. + transfer_id = ctx.request.id + # A retry finds the transfer already committed, with its send: it does nothing more. + if await tx.get("SELECT 1 FROM entries WHERE id = ?", [transfer_id]) is None: + await tx.execute( + "INSERT INTO entries (id, amount, note) VALUES (?, ?, ?)", + [transfer_id, -money.amount, f"to {money.to}"], + ) + if await balance_of(tx) < 0: + raise restate.TerminalError("Not enough money") + # Delivered once this transaction commits, and never if it rolls back. + tx.send(deposit, key=money.to)(Deposit(id=transfer_id, amount=money.amount, note=f"from {ctx.key}")) + return await balance_of(tx) + + +@wallet.view +async def balance(ctx: sessions.ViewContext) -> int: + return await balance_of(ctx.db) + + +@wallet.view +async def history(ctx: sessions.ViewContext) -> list[sessions.Row]: + return await ctx.db.query("SELECT amount, note FROM entries ORDER BY rowid") + + +if __name__ == "__main__": + sessions.serve(wallet) diff --git a/pyproject.toml b/pyproject.toml index 15a2e26a..36827385 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -84,6 +84,7 @@ target-version = "py310" include = [ "python/**/*.py", "examples/**/*.py", + "examples-agents/**/*.py", "tests/**/*.py", "test-services//**/*.py", ] diff --git a/python/restate/sessions/__init__.py b/python/restate/sessions/__init__.py index 874d7849..f9b8d77b 100644 --- a/python/restate/sessions/__init__.py +++ b/python/restate/sessions/__init__.py @@ -58,7 +58,7 @@ async def list(ctx: sessions.ViewContext) -> list[str]: from restate.sessions._database import Row, Rows, SqlError, SqlParam, SqlValue from restate.sessions._object import SessionObject from restate.sessions._schema import Step -from restate.sessions._serve import app +from restate.sessions._serve import app, serve __all__ = [ "CallFuture", @@ -79,4 +79,5 @@ async def list(ctx: sessions.ViewContext) -> list[str]: "Transaction", "ViewContext", "app", + "serve", ] diff --git a/python/restate/sessions/_serve.py b/python/restate/sessions/_serve.py index 322005b0..595dac19 100644 --- a/python/restate/sessions/_serve.py +++ b/python/restate/sessions/_serve.py @@ -17,9 +17,14 @@ from __future__ import annotations +import asyncio +import contextlib +import importlib import os +import signal import typing -from collections.abc import Callable, Iterable, Sequence +from collections.abc import Callable, Coroutine, Iterable, Sequence +from typing import Any import restate from restate.entry_codec import JournalValueCodec, JournalValueCodecProvider @@ -91,3 +96,65 @@ def app( identity_keys=identity_keys, journal_value_codec=journal_value_codec, ) + + +def serve( + *services: Service, + port: int | None = None, + identity_keys: list[str] | None = None, + journal_value_codec: JournalValueCodec | JournalValueCodecProvider | None = None, +) -> None: + """ + Serve ``services`` (``app``'s: with what they need beside them, of those + ``RESTATE_SERVE_ONLY`` names) on ``port``, or PORT, or 9080, until the + process stops: a program of session objects and agents in one line:: + + if __name__ == "__main__": + sessions.serve(wallet, handle) + + It runs them on hypercorn, imported here: install it, or serve ``app``'s + ASGI app with a server of your own. + """ + try: + hypercorn_asyncio = importlib.import_module("hypercorn.asyncio") + hypercorn_config = importlib.import_module("hypercorn.config") + except ImportError as error: + raise ImportError( + "sessions.serve runs on hypercorn: install it, or serve sessions.app(...) with an ASGI server of your own" + ) from error + config: Any = hypercorn_config.Config() + config.bind = [f"0.0.0.0:{port or int(os.environ.get('PORT', '9080'))}"] + # Restate keeps its connections open and streams many invocations over each. + config.h2_max_concurrent_streams = 2147483647 + config.keep_alive_max_requests = 2147483647 + config.keep_alive_timeout = 2147483647 + served_app = app(*services, identity_keys=identity_keys, journal_value_codec=journal_value_codec) + run = typing.cast(Callable[..., Coroutine[Any, Any, None]], hypercorn_asyncio.serve) + asyncio.run(_until_stopped(run, served_app, config)) + + +async def _until_stopped(run: Callable[..., Coroutine[Any, Any, None]], served_app: Any, config: Any) -> None: + """ + Serve until SIGINT or SIGTERM, then shut down: gracefully, but not for + longer than hypercorn's graceful timeout and a little more. A graceful + shutdown waits for every open connection to close, and Restate keeps its + own open. + """ + loop = asyncio.get_running_loop() + stop = asyncio.Event() + for name in ("SIGINT", "SIGTERM"): + with contextlib.suppress(NotImplementedError, AttributeError): + loop.add_signal_handler(getattr(signal, name), stop.set) + serving = asyncio.ensure_future(run(served_app, config, shutdown_trigger=stop.wait)) + stopping = asyncio.ensure_future(stop.wait()) + await asyncio.wait({serving, stopping}, return_when=asyncio.FIRST_COMPLETED) + stopping.cancel() + if serving.done(): + serving.result() + return + try: + await asyncio.wait_for(asyncio.shield(serving), timeout=float(config.graceful_timeout) + 2) + except asyncio.TimeoutError: + serving.cancel() + with contextlib.suppress(asyncio.CancelledError): + await serving diff --git a/tests/examples_agents.py b/tests/examples_agents.py new file mode 100644 index 00000000..367a9901 --- /dev/null +++ b/tests/examples_agents.py @@ -0,0 +1,142 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""The examples of examples-agents/, each with the commands its comment lists, against a real Restate server.""" + +# pylint: disable=missing-function-docstring + +import asyncio +import importlib.util +import pathlib +import typing +import uuid +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from types import ModuleType + +import httpx +import pytest + +import restate +from restate import sessions + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +pytestmark = [pytest.mark.anyio] + +EXAMPLES = pathlib.Path(__file__).parent.parent / "examples-agents" + + +def load(path: str) -> ModuleType: + """An example's module, from its file: what it serves, without serving it.""" + spec = importlib.util.spec_from_file_location(f"example_{uuid.uuid4().hex}", EXAMPLES / path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class Ingress: + """Restate's ingress, as the examples' curl commands call it.""" + + def __init__(self, url: str) -> None: + self.url = url + + async def post(self, path: str, body: object = None) -> httpx.Response: + async with httpx.AsyncClient(timeout=60) as client: + return await client.post(f"{self.url}/{path}", json=body) + + +@asynccontextmanager +async def serving(path: str) -> AsyncGenerator[Ingress, None]: + """The session objects of the example at ``path``, served to a Restate server of their own.""" + module = load(path) + objects = [value for value in vars(module).values() if isinstance(value, sessions.SessionObject)] + async with restate.create_test_harness( + sessions.app(*objects), restate_image="ghcr.io/restatedev/restate:latest" + ) as harness: + yield Ingress(harness.ingress_url) + + +async def test_01_object(): + async with serving("tutorial/01_object.py") as ingress: + assert (await ingress.post("wallet/alice/hello")).json() == "Hello, alice!" + + +async def test_02_database(): + async with serving("tutorial/02_database.py") as ingress: + assert (await ingress.post("wallet/alice/deposit", 10)).json() == 10 + assert (await ingress.post("wallet/alice/deposit", 5)).json() == 15 + assert (await ingress.post("wallet/bob/deposit", 7)).json() == 7 + + +async def test_03_views(): + async with serving("tutorial/03_views.py") as ingress: + await ingress.post("wallet/alice/deposit", 20) + assert (await ingress.post("wallet/alice/balance")).json() == 20 + assert (await ingress.post("wallet/alice/history")).json() == [20] + + +async def test_04_schema(): + async with serving("tutorial/04_schema.py") as ingress: + await ingress.post("wallet/alice/deposit", {"amount": 50, "note": "birthday"}) + assert (await ingress.post("wallet/alice/history")).json() == [{"amount": 50, "note": "birthday"}] + + +async def test_05_transaction(): + async with serving("tutorial/05_transaction.py") as ingress: + await ingress.post("wallet/alice/deposit", {"amount": 30, "note": "salary"}) + assert (await ingress.post("wallet/alice/spend", {"amount": 12, "note": "lunch"})).json() == 18 + boat = await ingress.post("wallet/alice/spend", {"amount": 1000, "note": "a boat"}) + assert boat.status_code == 500 and "Not enough money" in boat.text + history = (await ingress.post("wallet/alice/history")).json() + assert history == [{"amount": 30, "note": "salary"}, {"amount": -12, "note": "lunch"}] + + +async def test_06_retries(): + async with serving("tutorial/06_retries.py") as ingress: + refund = {"id": "d1", "amount": 10, "note": "refund"} + # Counted once, though the first attempt failed after the insert; and once when sent again. + assert (await ingress.post("wallet/alice/deposit", refund)).json() == 10 + assert (await ingress.post("wallet/alice/deposit", refund)).json() == 10 + + +async def test_07_calls(): + async with serving("tutorial/07_calls.py") as ingress: + assert (await ingress.post("wallet/alice/claim", "ace")).json() == "alice is now @ace" + taken = await ingress.post("wallet/bob/claim", "ace") + assert taken.status_code == 500 and "@ace belongs to alice" in taken.text + assert (await ingress.post("handle/ace/owner")).json() == "alice" + + +async def test_08_sends(): + async with serving("tutorial/08_sends.py") as ingress: + await ingress.post("wallet/alice/deposit", {"id": "d2", "amount": 100}) + assert (await ingress.post("wallet/alice/transfer", {"to": "bob", "amount": 30})).json() == 70 + history: list[typing.Any] = [] + for _ in range(100): + history = (await ingress.post("wallet/bob/history")).json() + if history: + break + await asyncio.sleep(0.05) + assert history == [{"amount": 30, "note": "from alice"}] + too_much = await ingress.post("wallet/alice/transfer", {"to": "bob", "amount": 5000}) + assert too_much.status_code == 500 + assert (await ingress.post("wallet/bob/history")).json() == [{"amount": 30, "note": "from alice"}] + + +async def test_basic_notes(): + async with serving("basic/notes.py") as ingress: + await ingress.post("notes/alice/add", "Buy milk") + assert (await ingress.post("notes/alice/list")).json() == ["Buy milk"] From 90460d366c3d39a89487c44d3e7d3a3cf4a2ba94 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 16:47:46 +0100 Subject: [PATCH 25/45] sessions.serve runs on uvicorn's HTTP/2, the serve extra uvicorn 0.53 serves HTTP/2 (cleartext, with prior knowledge, as Restate calls endpoints) through zttp. The new serve extra installs both, on Python 3.11+ (zttp's floor). serve uses uvicorn's own signal handling, with a bounded graceful shutdown. A test runs an example as a program, registers it with Restate, checks Restate speaks HTTP/2.0 to it, and stops it with SIGTERM. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- examples-agents/README.md | 3 +- python/restate/sessions/_serve.py | 82 +++++++++++++------------------ tests/examples_agents.py | 46 +++++++++++++++++ 3 files changed, 81 insertions(+), 50 deletions(-) diff --git a/examples-agents/README.md b/examples-agents/README.md index 1ae4a9ae..958c6617 100644 --- a/examples-agents/README.md +++ b/examples-agents/README.md @@ -10,7 +10,8 @@ examples. Each is a file that runs on its own. ## Run an example From the repository's root, once (it builds the SDK, with its native -extension): +extension, and installs the `serve` extra the examples run on: uvicorn's +HTTP/2, on Python 3.11 or later): ```sh uv sync --all-extras --all-packages diff --git a/python/restate/sessions/_serve.py b/python/restate/sessions/_serve.py index 595dac19..7890322c 100644 --- a/python/restate/sessions/_serve.py +++ b/python/restate/sessions/_serve.py @@ -12,19 +12,16 @@ ``app``: the ASGI app of a program's services, with what they need beside them (a library's objects, such as an agent's ``signals``). As sdk-sessions' ``serve.ts``; serving it is the program's (hypercorn, uvicorn...), as for -``restate.app``. +``restate.app``, or ``serve``'s. """ from __future__ import annotations -import asyncio -import contextlib import importlib import os -import signal +import re import typing -from collections.abc import Callable, Coroutine, Iterable, Sequence -from typing import Any +from collections.abc import Callable, Iterable, Sequence import restate from restate.entry_codec import JournalValueCodec, JournalValueCodecProvider @@ -112,49 +109,36 @@ def serve( if __name__ == "__main__": sessions.serve(wallet, handle) - It runs them on hypercorn, imported here: install it, or serve ``app``'s - ASGI app with a server of your own. + It runs them on uvicorn's HTTP/2 (cleartext, as Restate calls them), + imported here: install the ``serve`` extra (uvicorn 0.53 or later, and + zttp; Python 3.11 or later), or serve ``app``'s ASGI app with a server of + your own. It stops on SIGINT or SIGTERM, within a few seconds. """ try: - hypercorn_asyncio = importlib.import_module("hypercorn.asyncio") - hypercorn_config = importlib.import_module("hypercorn.config") + uvicorn = importlib.import_module("uvicorn") + importlib.import_module("zttp") except ImportError as error: - raise ImportError( - "sessions.serve runs on hypercorn: install it, or serve sessions.app(...) with an ASGI server of your own" - ) from error - config: Any = hypercorn_config.Config() - config.bind = [f"0.0.0.0:{port or int(os.environ.get('PORT', '9080'))}"] - # Restate keeps its connections open and streams many invocations over each. - config.h2_max_concurrent_streams = 2147483647 - config.keep_alive_max_requests = 2147483647 - config.keep_alive_timeout = 2147483647 - served_app = app(*services, identity_keys=identity_keys, journal_value_codec=journal_value_codec) - run = typing.cast(Callable[..., Coroutine[Any, Any, None]], hypercorn_asyncio.serve) - asyncio.run(_until_stopped(run, served_app, config)) - - -async def _until_stopped(run: Callable[..., Coroutine[Any, Any, None]], served_app: Any, config: Any) -> None: - """ - Serve until SIGINT or SIGTERM, then shut down: gracefully, but not for - longer than hypercorn's graceful timeout and a little more. A graceful - shutdown waits for every open connection to close, and Restate keeps its - own open. - """ - loop = asyncio.get_running_loop() - stop = asyncio.Event() - for name in ("SIGINT", "SIGTERM"): - with contextlib.suppress(NotImplementedError, AttributeError): - loop.add_signal_handler(getattr(signal, name), stop.set) - serving = asyncio.ensure_future(run(served_app, config, shutdown_trigger=stop.wait)) - stopping = asyncio.ensure_future(stop.wait()) - await asyncio.wait({serving, stopping}, return_when=asyncio.FIRST_COMPLETED) - stopping.cancel() - if serving.done(): - serving.result() - return - try: - await asyncio.wait_for(asyncio.shield(serving), timeout=float(config.graceful_timeout) + 2) - except asyncio.TimeoutError: - serving.cancel() - with contextlib.suppress(asyncio.CancelledError): - await serving + raise ImportError(_NO_SERVER) from error + version = re.match(r"(\d+)\.(\d+)", str(getattr(uvicorn, "__version__", ""))) + if version is None or (int(version[1]), int(version[2])) < (0, 53): + raise ImportError(_NO_SERVER) + config = uvicorn.Config( + app(*services, identity_keys=identity_keys, journal_value_codec=journal_value_codec), + host="0.0.0.0", + port=port or int(os.environ.get("PORT", "9080")), + http="zttp", + http2=True, + # The endpoint has no lifespan; and a graceful shutdown would wait for Restate's open connections. + lifespan="off", + timeout_graceful_shutdown=_GRACEFUL_SHUTDOWN, + ) + uvicorn.Server(config).run() + + +_NO_SERVER = ( + "sessions.serve runs on uvicorn's HTTP/2: install restate_sdk[serve] (uvicorn 0.53 or later, and zttp; " + "Python 3.11 or later), or serve sessions.app(...) with an ASGI server of your own" +) + +_GRACEFUL_SHUTDOWN = 3 +"""Seconds a stopping ``serve`` waits for the requests in flight, which Restate retries elsewhere.""" diff --git a/tests/examples_agents.py b/tests/examples_agents.py index 367a9901..e92ed864 100644 --- a/tests/examples_agents.py +++ b/tests/examples_agents.py @@ -14,7 +14,11 @@ import asyncio import importlib.util +import os import pathlib +import signal +import subprocess +import sys import typing import uuid from collections.abc import AsyncGenerator @@ -26,6 +30,7 @@ import restate from restate import sessions +from restate.harness import RestateContainer, find_free_port @pytest.fixture(scope="session") @@ -140,3 +145,44 @@ async def test_basic_notes(): async with serving("basic/notes.py") as ingress: await ingress.post("notes/alice/add", "Buy milk") assert (await ingress.post("notes/alice/list")).json() == ["Buy milk"] + + +def serves_http2() -> bool: + """Whether ``sessions.serve`` can run here: uvicorn 0.53+ with zttp.""" + return importlib.util.find_spec("uvicorn") is not None and importlib.util.find_spec("zttp") is not None + + +@pytest.mark.skipif(not serves_http2(), reason="sessions.serve needs the serve extra (Python 3.11+)") +async def test_an_example_as_a_program(): + """``uv run`` an example: it serves on PORT, Restate calls it over HTTP/2, and it stops on SIGTERM.""" + port = find_free_port() + program = subprocess.Popen( + [sys.executable, str(EXAMPLES / "basic/notes.py")], + env={**os.environ, "PORT": str(port)}, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + container = RestateContainer("ghcr.io/restatedev/restate:latest", always_replay=False, disable_retries=False) + try: + container.start() + async with httpx.AsyncClient(base_url=container.admin_url(), timeout=30) as admin: + uri = f"http://host.docker.internal:{port}" + registered = await admin.post("/deployments", json={"uri": uri, "force": True}) + for _ in range(100): # until the program serves + if registered.status_code < 300: + break + await asyncio.sleep(0.1) + registered = await admin.post("/deployments", json={"uri": uri, "force": True}) + assert registered.status_code < 300, registered.text + deployment = (await admin.get(f"/deployments/{registered.json()['id']}")).json() + assert deployment.get("http_version") == "HTTP/2.0", deployment + ingress = Ingress(container.ingress_url()) + await ingress.post("notes/alice/add", "Buy milk") + assert (await ingress.post("notes/alice/list")).json() == ["Buy milk"] + finally: + container.stop() + program.send_signal(signal.SIGTERM) + try: + program.wait(timeout=10) + finally: + program.kill() From e0bb1062114c90045e1f3c1cdffa676030f87084 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 16:59:59 +0100 Subject: [PATCH 26/45] Fit main's Cloud connection: its native names in the stub, quiet fake peer The stub declares ConnectError and TunnelInfo, which restate.connection re-exports. tests/connect.py's fake peer fails futures no test may await any more: their exceptions surfaced, once garbage-collected, in whichever test ran next on the shared loop (tests sorting before connect.py made that disconnect_hotloop). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/_internal.pyi | 10 ++++++++++ tests/connect.py | 12 +++++++++--- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/python/restate/_internal.pyi b/python/restate/_internal.pyi index fbe66e73..45cfae3f 100644 --- a/python/restate/_internal.pyi +++ b/python/restate/_internal.pyi @@ -33,6 +33,16 @@ class PyVM: def __init__(self, *args: Any, **kwargs: Any) -> None: ... def __getattr__(self, name: str) -> Any: ... +# Declared, as names ``restate.connection`` re-exports must be (still untyped). + +class ConnectError(Exception): + """Connecting to Restate Cloud failed for good (``src/tunnel``).""" + +class TunnelInfo: + """A Restate Cloud connection's details.""" + + def __getattr__(self, name: str) -> Any: ... + # --- Sessions (src/sessions.rs) ----------------------------------------------------------------- SQL_GET: int diff --git a/tests/connect.py b/tests/connect.py index 75e1a0fb..29ec9446 100644 --- a/tests/connect.py +++ b/tests/connect.py @@ -177,6 +177,12 @@ def identity_headers(audience: str) -> typing.List[typing.Tuple[str, str]]: TEST_TIMEOUT = 5.0 +def _fail(future: "asyncio.Future[None]", error: BaseException) -> None: + """Fail ``future``, which a test may no longer await: its exception is not one the loop must report.""" + future.set_exception(error) + future.exception() + + class Response: def __init__(self) -> None: self.stream_id = 0 @@ -284,9 +290,9 @@ async def _read(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter self._closed.set_result(None) for response in self._responses.values(): if not response.ended.done(): - response.ended.set_exception(ConnectionError("session closed")) + _fail(response.ended, ConnectionError("session closed")) if not response.started.done(): - response.started.set_exception(ConnectionError("session closed")) + _fail(response.started, ConnectionError("session closed")) def _dispatch(self, event: typing.Any) -> None: assert self._conn is not None @@ -311,7 +317,7 @@ def _dispatch(self, event: typing.Any) -> None: elif isinstance(event, StreamReset): reset = self._responses.get(event.stream_id) if reset is not None and not reset.ended.done(): - reset.ended.set_exception(ConnectionError(f"stream reset: {event.error_code}")) + _fail(reset.ended, ConnectionError(f"stream reset: {event.error_code}")) elif isinstance(event, ConnectionTerminated): if not self._goaway.done(): self._goaway.set_result(None) From f27f1bdec87390cafa057386e20b1cf460004f1f Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 16:59:59 +0100 Subject: [PATCH 27/45] Session objects preload only their database's head state_preload=[HEAD_KEY], as sdk-sessions' sqlOptions: an invocation starts with the head, the open's first read, instead of a round trip for it; the rest of the state is read as needed. Needs restate-server 1.8 (endpoint manifest v5). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- examples-agents/README.md | 2 +- python/restate/sessions/_object.py | 8 ++++---- tests/sessions_object.py | 12 +++++++++--- 3 files changed, 14 insertions(+), 8 deletions(-) diff --git a/examples-agents/README.md b/examples-agents/README.md index 958c6617..ebe12864 100644 --- a/examples-agents/README.md +++ b/examples-agents/README.md @@ -26,7 +26,7 @@ npx @restatedev/restate -y deployments register --force http://localhost:9080 curl localhost:8080/wallet/alice/hello ``` -Every example serves on port 9080 (or `PORT`): stop one before you start the +Session objects need Restate 1.8 or later. Every example serves on port 9080 (or `PORT`): stop one before you start the next, and register again. The comment at the top of each file lists the commands to try it. diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index bf3ba3e0..226bf1e9 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -352,9 +352,6 @@ def __init__( ingress_private: bool | None = None, invocation_retry_policy: InvocationRetryPolicy | None = None, ) -> None: - # TODO: start an invocation with only the database's head - # (HEAD_KEY), once the SDK can say which keys to load eagerly: lazy - # state loads none, so the open's first read is a round trip. super().__init__( name, description=description, @@ -363,7 +360,10 @@ def __init__( abort_timeout=abort_timeout, journal_retention=journal_retention, idempotency_retention=idempotency_retention, - enable_lazy_state=True, + # An invocation starts with the database's head only, its first read as it opens: the rest it + # reads as it needs it (its tail entries and packs change keys with each flush: only the head + # can be named). As sdk-sessions' sqlOptions; restate-server 1.8 or later. + state_preload=[HEAD_KEY], ingress_private=ingress_private, invocation_retry_policy=invocation_retry_policy, ) diff --git a/tests/sessions_object.py b/tests/sessions_object.py index c80ef9b4..81e038ec 100644 --- a/tests/sessions_object.py +++ b/tests/sessions_object.py @@ -24,9 +24,11 @@ from restate import sessions from restate.sessions import internal from restate.context import AttemptFinishedEvent, Request +from restate.discovery import compute_discovery +from restate.endpoint import Endpoint from restate.exceptions import TerminalError from restate.handler import handler_from_callable -from restate.sessions._object import invoke +from restate.sessions._object import HEAD_KEY, invoke from sessions_database import Backend # pylint: disable=import-error @@ -405,11 +407,15 @@ async def journaled(ctx: restate.ObjectContext) -> None: obj.handler()(add) -def test_a_session_object_runs_its_handlers_with_explicit_recovery_and_lazy_state(): +def test_a_session_object_runs_its_handlers_with_explicit_recovery_and_its_head_preloaded(): assert notes.handlers["add"].explicit_recovery assert notes.handlers["texts"].explicit_recovery assert notes.handlers["texts"].kind == "shared" - assert notes.enable_lazy_state + # Restate sends the database's head with each invocation, and the rest of the state as it is read. + manifest = compute_discovery(Endpoint().bind(notes), "bidi").services[0] + assert manifest.enableLazyState + assert manifest.eagerStateKeysWhitelist == [HEAD_KEY] + assert all(handler.eagerStateKeysWhitelist is None for handler in manifest.handlers) assert notes.service_tag.metadata is not None assert notes.service_tag.metadata["dev.restate.tooling.journal.hidden.name.prefix"] == "v22/" From 1840b9622ada625075669122106b7b2aed9da7ef Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 17:06:40 +0100 Subject: [PATCH 28/45] restate.agents.langgraph: SqlCheckpointSaver, LangGraph's checkpoints in ctx.db As sdk-sessions' langgraph/checkpointer.ts (and sqlexperiment's Python port): checkpoints, channel values (stored when a step changes them, at unique versions, so forks never clobber) and pending writes in three tables; each put is one transaction; history="latest" keeps only the thread's last checkpoint. The langgraph extra installs LangGraph. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- pyproject.toml | 1 + python/restate/agents/langgraph/__init__.py | 19 + .../restate/agents/langgraph/_checkpointer.py | 378 ++++++++++++++++++ tests/agents_langgraph_checkpointer.py | 193 +++++++++ uv.lock | 8 +- 5 files changed, 598 insertions(+), 1 deletion(-) create mode 100644 python/restate/agents/langgraph/__init__.py create mode 100644 python/restate/agents/langgraph/_checkpointer.py create mode 100644 tests/agents_langgraph_checkpointer.py diff --git a/pyproject.toml b/pyproject.toml index 36827385..35a79030 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ adk = ["google-adk>=1.20.0"] openai = ["openai-agents>=0.6.1"] pydantic_ai = ["pydantic-ai-slim>=2.0,<3"] langchain = ["langchain>=1.0.0", "langgraph>=1.0.0", "langchain-core>=1.0.0"] +langgraph = ["langgraph>=1.0.0", "langchain-core>=1.0.0"] tracing = ["opentelemetry-api>=1.36.0"] zstd = [ "backports.zstd ; python_version < '3.14'", diff --git a/python/restate/agents/langgraph/__init__.py b/python/restate/agents/langgraph/__init__.py new file mode 100644 index 00000000..f4fd8523 --- /dev/null +++ b/python/restate/agents/langgraph/__init__.py @@ -0,0 +1,19 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +LangGraph agents in session objects: their checkpoints in the object's +database (``SqlCheckpointSaver``). As sdk-sessions' ``@restatedev/agents/langgraph``. +It needs the ``langgraph`` extra (LangGraph and LangChain's core). +""" + +from restate.agents.langgraph._checkpointer import SqlCheckpointSaver + +__all__ = ["SqlCheckpointSaver"] diff --git a/python/restate/agents/langgraph/_checkpointer.py b/python/restate/agents/langgraph/_checkpointer.py new file mode 100644 index 00000000..4d74c63f --- /dev/null +++ b/python/restate/agents/langgraph/_checkpointer.py @@ -0,0 +1,378 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +A LangGraph checkpointer on a session object's database: ``SqlCheckpointSaver``, +over a database session ``db`` (in a handler, ``ctx.db``). As sdk-sessions' +``langgraph/checkpointer.ts``. + +Three tables, with values as LangGraph's serializer writes them: + +- ``restate_langgraph_checkpoints``: each checkpoint, without its channel values. +- ``restate_langgraph_channel_values``: each channel's value at each version, + stored when a step changes it. A checkpoint's values are those at its + channel versions, so a step stores only what it changed, not the whole + state again (an agent's message list, say). Versions are unique (a count + and a random part): a thread forked from an earlier checkpoint takes the + same counts again, and its values must not replace those of the branch it + forked from. +- ``restate_langgraph_writes``: the pending writes of each checkpoint's tasks. + +Each ``aput`` and ``aput_writes`` is one SQL transaction, so a step is stored +whole, in order with everything else the handler does. Operations run one at +a time, in call order: LangGraph writes the results of parallel tasks at once. +It is asynchronous only, as ``ctx.db`` is. + +By default a thread keeps every checkpoint (LangGraph's history: ``alist``, +time travel, forks). With ``history="latest"``, each ``aput`` drops the +thread's earlier checkpoints, their writes and the values only they had: a +thread then takes the space of its state once, not once per step (a message +list stored again at each step grows with the square of its length). +""" + +from __future__ import annotations + +import asyncio +import typing +import uuid +from collections.abc import AsyncIterator, Sequence +from typing import Any, Literal + +from langchain_core.runnables import RunnableConfig +from langgraph.checkpoint.base import ( + WRITES_IDX_MAP, + BaseCheckpointSaver, + ChannelVersions, + Checkpoint, + CheckpointMetadata, + CheckpointTuple, + get_checkpoint_id, + get_checkpoint_metadata, +) +from langgraph.checkpoint.serde.base import SerializerProtocol + +from restate.sessions import Database, DatabaseWriter, SqlParam, SqlValue +from restate.sessions.internal import Schema + +SCHEMA = Schema( + "restate_agents.langgraph", + [ + """CREATE TABLE restate_langgraph_checkpoints ( + thread_id TEXT NOT NULL, + checkpoint_ns TEXT NOT NULL, + checkpoint_id TEXT NOT NULL, + parent_checkpoint_id TEXT, + type TEXT NOT NULL, + checkpoint BLOB NOT NULL, + metadata BLOB NOT NULL, + PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id) + )""", + """CREATE TABLE restate_langgraph_channel_values ( + thread_id TEXT NOT NULL, + checkpoint_ns TEXT NOT NULL, + channel TEXT NOT NULL, + version TEXT NOT NULL, + type TEXT NOT NULL, + value BLOB, + PRIMARY KEY (thread_id, checkpoint_ns, channel, version) + )""", + """CREATE TABLE restate_langgraph_writes ( + thread_id TEXT NOT NULL, + checkpoint_ns TEXT NOT NULL, + checkpoint_id TEXT NOT NULL, + task_id TEXT NOT NULL, + task_path TEXT NOT NULL, + idx INTEGER NOT NULL, + channel TEXT NOT NULL, + type TEXT NOT NULL, + value BLOB NOT NULL, + PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) + )""", + ], +) +"""The checkpointer's tables, a schema of their own, applied where it is used.""" + +_TABLES = ("restate_langgraph_checkpoints", "restate_langgraph_channel_values", "restate_langgraph_writes") + +_COLUMNS = "thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata" + +_EMPTY = "empty" +"""The type of a channel a step emptied.""" + + +class SqlCheckpointSaver(BaseCheckpointSaver[str]): + """ + LangGraph's checkpoints in a database session, such as ``ctx.db``: see + the module. ``history``: the checkpoints a thread keeps, ``"all"`` (the + default) or the ``"latest"`` only. + """ + + def __init__( + self, + db: Database, + *, + history: Literal["all", "latest"] = "all", + serde: SerializerProtocol | None = None, + ) -> None: + super().__init__(serde=serde) + self.db = db + self.latest = history == "latest" + self._lock = asyncio.Lock() + self._ready = False + + def get_next_version(self, current: str | None, channel: None = None) -> str: + """ + The version after ``current``: its count plus one, zero-padded so that + versions compare as their counts do, and a random part, so that + branches of a thread never share one. + """ + if current is None: + count = 0 + elif isinstance(current, int): # pyright: ignore[reportUnnecessaryIsInstance] # LangGraph's own start at 1 + count = current + else: + count = int(current.split(".")[0]) + return f"{count + 1:032}.{uuid.uuid4().hex}" + + async def _setup(self) -> None: + if not self._ready: + await SCHEMA.ensure(self.db) + self._ready = True + + def _load(self, type_: SqlValue, value: SqlValue) -> Any: + return self.serde.loads_typed((str(type_), _bytes(value))) + + async def _channel_values(self, thread_id: str, ns: str, checkpoint: Checkpoint) -> dict[str, Any]: + """The values of ``checkpoint``'s channels, at its versions.""" + versions = list(checkpoint.get("channel_versions", {}).items()) + if not versions: + return {} + rows = await self.db.values( + "SELECT channel, type, value FROM restate_langgraph_channel_values " + "WHERE thread_id = ? AND checkpoint_ns = ? AND (" + + " OR ".join("(channel = ? AND version = ?)" for _ in versions) + + ")", + [thread_id, ns, *(item for channel, version in versions for item in (channel, str(version)))], + ) + return {str(channel): self._load(type_, value) for channel, type_, value in rows if type_ != _EMPTY} + + async def _tuple(self, row: list[SqlValue]) -> CheckpointTuple: + thread_id, ns, checkpoint_id, parent_id, type_, stored, metadata = ( + str(row[0]), + str(row[1]), + str(row[2]), + row[3], + row[4], + row[5], + row[6], + ) + checkpoint = typing.cast(Checkpoint, self._load(type_, stored)) + checkpoint = typing.cast( + Checkpoint, {**checkpoint, "channel_values": await self._channel_values(thread_id, ns, checkpoint)} + ) + writes = await self.db.values( + "SELECT task_id, channel, type, value FROM restate_langgraph_writes " + "WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? ORDER BY task_path, task_id, idx", + [thread_id, ns, checkpoint_id], + ) + + def config(id_: str) -> RunnableConfig: + return {"configurable": {"thread_id": thread_id, "checkpoint_ns": ns, "checkpoint_id": id_}} + + return CheckpointTuple( + config=config(checkpoint_id), + checkpoint=checkpoint, + metadata=typing.cast(CheckpointMetadata, self._load(type_, metadata)), + parent_config=None if parent_id is None else config(str(parent_id)), + pending_writes=[ + (str(task_id), str(channel), self._load(write_type, value)) + for task_id, channel, write_type, value in writes + ], + ) + + async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: + async with self._lock: + await self._setup() + thread_id, ns = _thread(config) + checkpoint_id = get_checkpoint_id(config) + if checkpoint_id: + rows = await self.db.values( + f"SELECT {_COLUMNS} FROM restate_langgraph_checkpoints " + "WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?", + [thread_id, ns, checkpoint_id], + ) + else: + rows = await self.db.values( + f"SELECT {_COLUMNS} FROM restate_langgraph_checkpoints WHERE thread_id = ? AND checkpoint_ns = ? " + "ORDER BY checkpoint_id DESC LIMIT 1", + [thread_id, ns], + ) + return await self._tuple(rows[0]) if rows else None + + async def alist( + self, + config: RunnableConfig | None, + *, + filter: dict[str, Any] | None = None, # LangGraph's name + before: RunnableConfig | None = None, + limit: int | None = None, + ) -> AsyncIterator[CheckpointTuple]: + async with self._lock: + await self._setup() + configurable: dict[str, Any] = dict((config or {}).get("configurable", {})) + conditions: list[tuple[str, Any]] = [ + ("thread_id = ?", configurable.get("thread_id")), + ("checkpoint_ns = ?", configurable.get("checkpoint_ns")), + ("checkpoint_id = ?", configurable.get("checkpoint_id")), + ("checkpoint_id < ?", get_checkpoint_id(before) if before else None), + ] + where = [condition for condition, value in conditions if value is not None] + params: list[SqlParam] = [value for _, value in conditions if value is not None] + rows = await self.db.values( + f"SELECT {_COLUMNS} FROM restate_langgraph_checkpoints " + f"{'WHERE ' + ' AND '.join(where) if where else ''} ORDER BY checkpoint_id DESC", + params, + ) + found: list[CheckpointTuple] = [] + for row in rows: + if limit is not None and len(found) >= limit: + break + checkpoint = await self._tuple(row) + if all(_matches(checkpoint.metadata.get(key), value) for key, value in (filter or {}).items()): + found.append(checkpoint) + for checkpoint in found: + yield checkpoint + + async def aput( + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + new_versions: ChannelVersions, + ) -> RunnableConfig: + async with self._lock: + await self._setup() + thread_id, ns = _thread(config) + parent_id: str | None = _configurable(config).get("checkpoint_id") + stored: dict[str, Any] = dict(checkpoint) + values: dict[str, Any] = stored.pop("channel_values") + type_, serialized = self.serde.dumps_typed(stored) + metadata_type, metadata_value = self.serde.dumps_typed(get_checkpoint_metadata(config, metadata)) + if type_ != metadata_type: + raise ValueError("Failed to serialize checkpoint and metadata to the same type.") + # The channels this step changed, at their new versions. + changed: list[list[SqlParam]] = [ + [ + channel, + str(version), + *(self.serde.dumps_typed(values[channel]) if channel in values else (_EMPTY, None)), + ] + for channel, version in new_versions.items() + ] + async with self.db.transaction() as tx: + for row in changed: + await tx.execute( + "INSERT OR REPLACE INTO restate_langgraph_channel_values " + "(thread_id, checkpoint_ns, channel, version, type, value) VALUES (?, ?, ?, ?, ?, ?)", + [thread_id, ns, *row], + ) + await tx.execute( + f"INSERT OR REPLACE INTO restate_langgraph_checkpoints ({_COLUMNS}) VALUES (?, ?, ?, ?, ?, ?, ?)", + [thread_id, ns, checkpoint["id"], parent_id, type_, serialized, metadata_value], + ) + if self.latest: + await _drop_earlier(tx, thread_id, ns, checkpoint) + return {"configurable": {"thread_id": thread_id, "checkpoint_ns": ns, "checkpoint_id": checkpoint["id"]}} + + async def aput_writes( + self, + config: RunnableConfig, + writes: Sequence[tuple[str, Any]], + task_id: str, + task_path: str = "", + ) -> None: + async with self._lock: + await self._setup() + thread_id, ns = _thread(config) + checkpoint_id = str(_configurable(config)["checkpoint_id"]) + # Special writes (errors, interrupts) replace; others keep the first. + verb = ( + "INSERT OR REPLACE" if all(channel in WRITES_IDX_MAP for channel, _ in writes) else "INSERT OR IGNORE" + ) + rows: list[list[SqlParam]] = [ + [ + thread_id, + ns, + checkpoint_id, + task_id, + task_path, + WRITES_IDX_MAP.get(channel, idx), + channel, + *self.serde.dumps_typed(value), + ] + for idx, (channel, value) in enumerate(writes) + ] + async with self.db.transaction() as tx: + for row in rows: + await tx.execute( + f"{verb} INTO restate_langgraph_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, " + "task_path, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + row, + ) + + async def adelete_thread(self, thread_id: str) -> None: + async with self._lock: + await self._setup() + async with self.db.transaction() as tx: + for table in _TABLES: + await tx.execute(f"DELETE FROM {table} WHERE thread_id = ?", [thread_id]) + + +def _thread(config: RunnableConfig) -> tuple[str, str]: + """The thread of ``config``, and its namespace.""" + configurable = _configurable(config) + return str(configurable["thread_id"]), str(configurable.get("checkpoint_ns", "")) + + +def _configurable(config: RunnableConfig) -> dict[str, Any]: + return dict(config.get("configurable", {})) + + +def _bytes(value: SqlValue) -> bytes: + return value if isinstance(value, bytes) else b"" if value is None else str(value).encode() + + +def _matches(actual: object, expected: object) -> bool: + """Whether ``actual`` has every field of ``expected``, recursively (a metadata filter).""" + if isinstance(expected, dict): + if not isinstance(actual, dict): + return False + fields = typing.cast(dict[Any, Any], actual) + return all(_matches(fields.get(key), value) for key, value in typing.cast(dict[Any, Any], expected).items()) + return actual == expected + + +async def _drop_earlier(tx: DatabaseWriter, thread_id: str, ns: str, checkpoint: Checkpoint) -> None: + """ + Drop what the thread keeps from before ``checkpoint``: its other + checkpoints, their pending writes, and the channel values at versions + other than ``checkpoint``'s. + """ + for table in ("restate_langgraph_checkpoints", "restate_langgraph_writes"): + await tx.execute( + f"DELETE FROM {table} WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id <> ?", + [thread_id, ns, checkpoint["id"]], + ) + versions = list(checkpoint.get("channel_versions", {}).items()) + kept = " OR ".join("(channel = ? AND version = ?)" for _ in versions) or "0" + await tx.execute( + f"DELETE FROM restate_langgraph_channel_values WHERE thread_id = ? AND checkpoint_ns = ? AND NOT ({kept})", + [thread_id, ns, *(item for channel, version in versions for item in (channel, str(version)))], + ) diff --git a/tests/agents_langgraph_checkpointer.py b/tests/agents_langgraph_checkpointer.py new file mode 100644 index 00000000..e6b7015f --- /dev/null +++ b/tests/agents_langgraph_checkpointer.py @@ -0,0 +1,193 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""restate.agents.langgraph's checkpointer: LangGraph's checkpoints in a session object's database.""" + +# pylint: disable=missing-function-docstring,missing-class-docstring + +import operator +import typing +from collections.abc import Awaitable, Callable +from typing import Annotated, Any, TypedDict + +import pytest + +pytest.importorskip("langgraph") + +from langchain_core.runnables import RunnableConfig # noqa: E402 +from langgraph.graph import END, START, StateGraph # noqa: E402 + +from restate import sessions # noqa: E402 +from restate.agents.langgraph import SqlCheckpointSaver # noqa: E402 +from sessions_database import Backend # noqa: E402 # pylint: disable=import-error +from sessions_object import Host, run # noqa: E402 # pylint: disable=import-error + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +pytestmark = [pytest.mark.anyio] + + +class State(TypedDict): + log: Annotated[list[str], operator.add] + + +# Each test's work, run by the handler `work` with the attempt's context. +Work = Callable[[sessions.SessionContext], Awaitable[None]] +job: list[Work] = [] + +graphs = sessions.SessionObject("graphs") + + +@graphs.handler +async def work(ctx: sessions.SessionContext) -> None: + await job[-1](ctx) + + +async def attempt(backend: Backend, work: Work) -> None: + """An attempt of ``work`` on ``backend``'s key.""" + job.append(work) + await run(graphs, "work", Host(backend, service="graphs")) + + +THREAD: RunnableConfig = {"configurable": {"thread_id": "t"}} + + +def steps(ran: list[str], fail: set[str] | None = None) -> Any: + """a, then b and c at once, then d: each records that it ran; those in ``fail`` fail once.""" + failing = set(fail or ()) + + def node(name: str) -> Callable[[State], dict[str, list[str]]]: + def step(_: State) -> dict[str, list[str]]: + ran.append(name) + if name in failing: + failing.discard(name) + raise RuntimeError(f"{name} crashed") + return {"log": [name]} + + return step + + graph = StateGraph(State) + for name in "abcd": + graph.add_node(name, typing.cast(Any, node(name))) + graph.add_edge(START, "a") + graph.add_edge("a", "b") + graph.add_edge("a", "c") + graph.add_edge(["b", "c"], "d") + graph.add_edge("d", END) + return graph + + +async def test_a_retried_attempt_resumes_from_the_last_stored_step_and_its_tasks_writes(): + backend = Backend() + ran: list[str] = [] + graph = steps(ran, fail={"c"}) + results: list[Any] = [] + + async def go(ctx: sessions.SessionContext) -> None: + app = graph.compile(checkpointer=SqlCheckpointSaver(ctx.db)) + state = await app.aget_state(THREAD) + results.append(await app.ainvoke(None if state.next else {"log": []}, THREAD, durability="sync")) + + with pytest.raises(RuntimeError, match="c crashed"): + await attempt(backend, go) + await attempt(backend, go) + # a ran once; b's write, stored while c failed, is not made again; c runs again. + assert ran == ["a", "b", "c", "c", "d"] or ran == ["a", "c", "b", "c", "d"] + assert sorted(results[-1]["log"]) == ["a", "b", "c", "d"] + + +async def test_history_lists_the_checkpoints_newest_first_by_limit_filter_and_before(): + backend = Backend() + listed: dict[str, list[Any]] = {} + + async def go(ctx: sessions.SessionContext) -> None: + saver = SqlCheckpointSaver(ctx.db) + app = steps([]).compile(checkpointer=saver) + await app.ainvoke({"log": []}, THREAD, durability="sync") + listed["all"] = [c async for c in saver.alist(THREAD)] + listed["two"] = [c async for c in saver.alist(THREAD, limit=2)] + listed["step 1"] = [c async for c in saver.alist(THREAD, filter={"step": 1})] + listed["before"] = [c async for c in saver.alist(THREAD, before=listed["all"][1].config)] + listed["other"] = [c async for c in saver.alist({"configurable": {"thread_id": "other"}})] + + await attempt(backend, go) + stepped = [c.metadata.get("step") for c in listed["all"]] + assert stepped == sorted(stepped, reverse=True) and len(stepped) == 5 + assert [c.config for c in listed["two"]] == [c.config for c in listed["all"][:2]] + assert [c.metadata.get("step") for c in listed["step 1"]] == [1] + assert [c.config for c in listed["before"]] == [c.config for c in listed["all"][2:]] + assert listed["other"] == [] + # Each checkpoint names its parent. + assert listed["all"][0].parent_config == listed["all"][1].config + + +async def test_latest_history_keeps_one_checkpoint_and_only_its_values(): + backend = Backend() + counted: dict[str, Any] = {} + + async def go(ctx: sessions.SessionContext) -> None: + saver = SqlCheckpointSaver(ctx.db, history="latest") + app = steps([]).compile(checkpointer=saver) + await app.ainvoke({"log": []}, THREAD, durability="sync") + await app.ainvoke({"log": []}, THREAD, durability="sync") + for table in ("checkpoints", "channel_values", "writes"): + row = await ctx.db.get(f"SELECT count(*) AS n FROM restate_langgraph_{table}") + counted[table] = row["n"] if row else None + latest = await saver.aget_tuple(THREAD) + assert latest is not None + counted["versions"] = len(latest.checkpoint["channel_versions"]) + counted["log"] = latest.checkpoint["channel_values"]["log"] + + await attempt(backend, go) + assert counted["checkpoints"] == 1 + assert counted["channel_values"] == counted["versions"] + assert sorted(counted["log"]) == sorted(["a", "b", "c", "d"] * 2) + + +async def test_a_fork_from_an_earlier_checkpoint_leaves_the_branch_it_forked_from(): + backend = Backend() + seen: dict[str, Any] = {} + + async def go(ctx: sessions.SessionContext) -> None: + saver = SqlCheckpointSaver(ctx.db) + app = steps([]).compile(checkpointer=saver) + await app.ainvoke({"log": []}, THREAD, durability="sync") + history = [c async for c in saver.alist(THREAD)] + after_a = next(c for c in history if c.metadata.get("step") == 1) + forked = await app.aupdate_state(after_a.config, {"log": ["fork"]}) + seen["fork"] = (await app.aget_state(forked)).values["log"] + seen["main"] = (await app.aget_state(history[0].config)).values["log"] + + await attempt(backend, go) + assert seen["fork"] == ["a", "fork"] + assert sorted(seen["main"]) == ["a", "b", "c", "d"] + + +async def test_delete_thread_forgets_it_and_no_other(): + backend = Backend() + seen: dict[str, Any] = {} + + async def go(ctx: sessions.SessionContext) -> None: + saver = SqlCheckpointSaver(ctx.db) + app = steps([]).compile(checkpointer=saver) + other: RunnableConfig = {"configurable": {"thread_id": "other"}} + await app.ainvoke({"log": []}, THREAD, durability="sync") + await app.ainvoke({"log": []}, other, durability="sync") + await saver.adelete_thread("t") + seen["t"] = await saver.aget_tuple(THREAD) + seen["other"] = await saver.aget_tuple(other) + + await attempt(backend, go) + assert seen["t"] is None + assert typing.cast(Any, seen["other"]).checkpoint["channel_values"]["log"] diff --git a/uv.lock b/uv.lock index 4958dfee..bbc54b3a 100644 --- a/uv.lock +++ b/uv.lock @@ -3310,6 +3310,10 @@ langchain = [ { name = "langchain-core" }, { name = "langgraph" }, ] +langgraph = [ + { name = "langchain-core" }, + { name = "langgraph" }, +] lint = [ { name = "mypy" }, { name = "pyright" }, @@ -3354,7 +3358,9 @@ requires-dist = [ { name = "hypercorn", marker = "extra == 'test'" }, { name = "langchain", marker = "extra == 'langchain'", specifier = ">=1.0.0" }, { name = "langchain-core", marker = "extra == 'langchain'", specifier = ">=1.0.0" }, + { name = "langchain-core", marker = "extra == 'langgraph'", specifier = ">=1.0.0" }, { name = "langgraph", marker = "extra == 'langchain'", specifier = ">=1.0.0" }, + { name = "langgraph", marker = "extra == 'langgraph'", specifier = ">=1.0.0" }, { name = "msgspec", marker = "extra == 'serde'" }, { name = "mypy", marker = "extra == 'lint'", specifier = ">=1.11.2" }, { name = "openai-agents", marker = "extra == 'openai'", specifier = ">=0.6.1" }, @@ -3368,7 +3374,7 @@ requires-dist = [ { name = "uvicorn", marker = "python_full_version >= '3.11' and extra == 'serve'", specifier = ">=0.53" }, { name = "zttp", marker = "python_full_version >= '3.11' and extra == 'serve'", specifier = ">=0.0.34" }, ] -provides-extras = ["adk", "client", "harness", "langchain", "lint", "openai", "pydantic-ai", "serde", "serve", "test", "tracing", "zstd"] +provides-extras = ["adk", "client", "harness", "langchain", "langgraph", "lint", "openai", "pydantic-ai", "serde", "serve", "test", "tracing", "zstd"] [[package]] name = "rpds-py" From 5d692dd0f993ee04cf947c3fafeed58f7f522487 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 17:13:46 +0100 Subject: [PATCH 29/45] restate.agents.langgraph: agents, chat turns and prompts on LangGraph As sdk-sessions' langgraph/index.ts: the usual tool-calling graph (model, tools, answer), each step stored before the next (durability=sync) with the latest checkpoint only; each prompt answered once by id; a crash in a tool resumes it with the same call id, without asking the model again; a stop closes the run as cancelled; a provider's refusal fails with 424. A LangChain tool's terminal error also closes the run, so the next message does not meet it again. Dependencies stay opt-in: the langgraph extra (LangGraph, LangChain's core, pydantic); importing restate.agents.langgraph without it says which extra to install; openai/ and anthropic/ models import langchain-openai or langchain-anthropic only when used, and a missing one is a 424 naming the package. Example: examples-agents/basic/langgraph_agent.py. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- examples-agents/README.md | 5 +- examples-agents/basic/langgraph_agent.py | 31 ++ pyproject.toml | 2 +- python/restate/agents/langgraph/__init__.py | 35 +- python/restate/agents/langgraph/_agent.py | 383 ++++++++++++++++++++ tests/agents_imports.py | 18 + tests/agents_langgraph.py | 294 +++++++++++++++ tests/examples_agents.py | 8 + uv.lock | 2 + 9 files changed, 772 insertions(+), 6 deletions(-) create mode 100644 examples-agents/basic/langgraph_agent.py create mode 100644 python/restate/agents/langgraph/_agent.py create mode 100644 tests/agents_langgraph.py diff --git a/examples-agents/README.md b/examples-agents/README.md index ebe12864..c82b3fc2 100644 --- a/examples-agents/README.md +++ b/examples-agents/README.md @@ -26,7 +26,9 @@ npx @restatedev/restate -y deployments register --force http://localhost:9080 curl localhost:8080/wallet/alice/hello ``` -Session objects need Restate 1.8 or later. Every example serves on port 9080 (or `PORT`): stop one before you start the +Session objects need Restate 1.8 or later. An agent's framework comes with its extra, and its model's provider with +its package: the comment at the top of each agent's file says which (`uv run +--extra langgraph --with langchain-openai ...`, say). Every example serves on port 9080 (or `PORT`): stop one before you start the next, and register again. The comment at the top of each file lists the commands to try it. @@ -64,6 +66,7 @@ the loop, a subagent), comes with the agents' framework integrations. | File | What | | ----------------------------- | ----------------------------------------------- | | [`notes.py`](basic/notes.py) | SQL in a virtual object, with `ctx.db`, and no agent | +| [`langgraph_agent.py`](basic/langgraph_agent.py) | A LangGraph agent: its checkpoints in the key's database | ## Tips diff --git a/examples-agents/basic/langgraph_agent.py b/examples-agents/basic/langgraph_agent.py new file mode 100644 index 00000000..a3758fae --- /dev/null +++ b/examples-agents/basic/langgraph_agent.py @@ -0,0 +1,31 @@ +""" +A LangGraph agent: a weather assistant, its conversation (and LangGraph's +checkpoints) in each key's database. A crash mid-turn resumes from LangGraph's +last stored step: a tool call already made is not made again:: + + OPENAI_API_KEY=... uv run --extra langgraph --extra serve --with langchain-openai examples-agents/basic/langgraph_agent.py + npx @restatedev/restate-server # Restate, in another terminal + npx @restatedev/restate -y deployments register --force http://localhost:9080 + curl localhost:8080/assistant/alice/chat -d 'What is the weather in Lisbon?' + curl localhost:8080/assistant/alice/messages --json '{}' # the conversation + +The model is ``openai/gpt-6-sol`` unless MODEL names another: +``anthropic/`` takes ``langchain-anthropic`` (and +ANTHROPIC_API_KEY) instead of ``langchain-openai``. +""" + +from restate import sessions +from restate.agents import langgraph, tool + + +@tool +async def get_weather(city: str) -> str: + """The current weather in a city.""" + return f"Sunny and 21°C in {city}" + + +assistant = langgraph.agent("assistant", system="Answer briefly.", tools=[get_weather]) + + +if __name__ == "__main__": + sessions.serve(assistant) diff --git a/pyproject.toml b/pyproject.toml index 35a79030..17730858 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,7 @@ adk = ["google-adk>=1.20.0"] openai = ["openai-agents>=0.6.1"] pydantic_ai = ["pydantic-ai-slim>=2.0,<3"] langchain = ["langchain>=1.0.0", "langgraph>=1.0.0", "langchain-core>=1.0.0"] -langgraph = ["langgraph>=1.0.0", "langchain-core>=1.0.0"] +langgraph = ["langgraph>=1.0.0", "langchain-core>=1.0.0", "pydantic>=2"] tracing = ["opentelemetry-api>=1.36.0"] zstd = [ "backports.zstd ; python_version < '3.14'", diff --git a/python/restate/agents/langgraph/__init__.py b/python/restate/agents/langgraph/__init__.py index f4fd8523..92fe9dd2 100644 --- a/python/restate/agents/langgraph/__init__.py +++ b/python/restate/agents/langgraph/__init__.py @@ -9,11 +9,38 @@ # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # """ -LangGraph agents in session objects: their checkpoints in the object's -database (``SqlCheckpointSaver``). As sdk-sessions' ``@restatedev/agents/langgraph``. -It needs the ``langgraph`` extra (LangGraph and LangChain's core). +LangGraph agents in session objects: a tool-calling graph whose checkpoints +are in the object's database. As sdk-sessions' ``@restatedev/agents/langgraph``:: + + from restate import sessions + from restate.agents import langgraph, tool + + @tool + async def get_weather(city: str) -> str: + '''The current weather in a city.''' + return f"Sunny in {city}" + + assistant = langgraph.agent("assistant", system="Answer briefly.", tools=[get_weather]) + sessions.serve(assistant) + +It needs the ``langgraph`` extra (``pip install restate-sdk[langgraph]``); +a model named ``openai/...`` needs ``langchain-openai`` beside it, and +``anthropic/...`` ``langchain-anthropic``. Or bring a LangChain chat model of +your own (``model=``). ``agent`` makes an agent; ``chat`` is the turn of an +``AgentObject`` of your own, ``prompt`` a prompt of a handler of your own, +and ``SqlCheckpointSaver`` the checkpointer, for a graph of your own. """ +try: + import langchain_core # noqa: F401 # pyright: ignore[reportUnusedImport] + + import langgraph # noqa: F401 # pyright: ignore[reportUnusedImport] +except ImportError as error: + raise ImportError( + "restate.agents.langgraph needs LangGraph: install the langgraph extra, pip install restate-sdk[langgraph]" + ) from error + +from restate.agents.langgraph._agent import Options, agent, chat, prompt from restate.agents.langgraph._checkpointer import SqlCheckpointSaver -__all__ = ["SqlCheckpointSaver"] +__all__ = ["Options", "SqlCheckpointSaver", "agent", "chat", "prompt"] diff --git a/python/restate/agents/langgraph/_agent.py b/python/restate/agents/langgraph/_agent.py new file mode 100644 index 00000000..bc6a5d82 --- /dev/null +++ b/python/restate/agents/langgraph/_agent.py @@ -0,0 +1,383 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +# LangGraph's graph API (langgraph.graph, its StateGraph and compiled graphs) is untyped: no py.typed. +# pyright: reportMissingTypeStubs=false, reportUnknownMemberType=false +""" +LangGraph agents: a tool-calling graph whose checkpoints are in the object's +database (``_checkpointer.py``). As sdk-sessions' ``langgraph/index.ts``. + +The graph is the usual tool-calling loop: a model node, a tools node, and back +to the model until it answers. The handler is not replayed: a retried attempt +(a crash, a lost connection) runs again from the start, and LangGraph carries +on from the last step it stored: + +- Each step is stored before the next starts (``durability="sync"``), in order + with everything else the handler does. So the tool calls a model made are + stored before any tool runs, and a crash during the tools resumes them with + the same tool call ids, which a tool passes on to stay idempotent. +- A tool's error fails the attempt, which Restate retries (LangGraph's + ToolNode would hand it to the model as the tool's result); a call of a tool + the agent lacks, or with arguments that do not fit, is the model's to fix. +- A message is answered once per invocation: the graph records each answer in + its state, by prompt id, and a retried attempt returns it. +""" + +from __future__ import annotations + +import importlib +import os +import typing +from collections.abc import Awaitable, Callable, Sequence +from dataclasses import dataclass +from datetime import timedelta +from typing import Annotated, Any, TypedDict + +from langchain_core.language_models import BaseChatModel +from langchain_core.messages import AIMessage, AnyMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage +from langchain_core.runnables import RunnableConfig +from langchain_core.tools import BaseTool +from langgraph.graph import END, START, StateGraph +from langgraph.graph.message import add_messages +from langgraph.graph.state import CompiledStateGraph +from langgraph.prebuilt import ToolNode, tools_condition +from pydantic import PrivateAttr + +from restate.agents._common import ( + CANCELLED, + UNANSWERED, + Configure, + ModelError, + configured, + from_tool, + model_of, + next_prompt_id, + provider_of, + refusal_of, + stored_answer, +) +from restate.agents._object import RETRIES, AgentObject +from restate.agents._tools import Tool, ToolCall +from restate.agents._transcript import Transcript +from restate.agents._turn import AttemptEnded, stop_signal_of, tool_call_ended, tool_call_started, transcribing +from restate.agents.langgraph._checkpointer import SqlCheckpointSaver +from restate.exceptions import TerminalError +from restate.retry_policy import InvocationRetryPolicy +from restate.sessions import SessionContext, Step + +_CANCELLATION = 409 +"""The status of Restate's cancellation, as a turn that was stopped fails with it.""" + + +def _merged(old: dict[str, str], new: dict[str, str]) -> dict[str, str]: + return {**old, **new} + + +class State(TypedDict, total=False): + """The agent's graph's state: its conversation, the prompt it answers, and each prompt's answer.""" + + messages: Annotated[list[AnyMessage], add_messages] + prompt: str + """The id of the prompt being answered.""" + answers: Annotated[dict[str, str], _merged] + """The answer to each prompt, by id.""" + + +@dataclass(frozen=True) +class Options: + """ + A LangGraph agent's options. ``model``: ``openai/model-id`` or + ``anthropic/model-id`` (MODEL, or else ``openai/gpt-6-sol``), or a + LangChain chat model that calls tools. ``system``: the system prompt. + ``tools``: ``restate.agents`` tools (``tool``, ``call``, ``signal``...), or + LangChain's. + """ + + model: str | BaseChatModel | None = None + system: str | None = None + tools: Sequence[Tool | BaseTool] = () + + +_PROVIDERS = { + "openai": ("langchain_openai", "ChatOpenAI", "langchain-openai"), + "anthropic": ("langchain_anthropic", "ChatAnthropic", "langchain-anthropic"), +} +"""The providers the usual graph takes by name: LangChain's integration of each, its class, and its package.""" + + +def _provider_of_model(named: str) -> tuple[str, str]: + """``provider/model-id`` as its provider and model id, for a provider the usual graph takes.""" + provider, model_id = provider_of(named) + if provider not in _PROVIDERS or not model_id: + from_env = " (from MODEL)" if os.environ.get("MODEL") == named else "" + raise ModelError(f'the model "{named}"{from_env} is neither openai/ nor anthropic/') + return provider, model_id + + +def _chat_model(named: str) -> BaseChatModel: + """ + The chat model ``provider/model-id`` names, through LangChain's + integration of the provider, imported when a model of it is used: only + that one need be installed, beside the ``langgraph`` extra. + """ + provider, model_id = _provider_of_model(named) + module, class_name, package = _PROVIDERS[provider] + try: + integration = importlib.import_module(module) + except ImportError as error: + raise ModelError( + f"{provider}/ models need {package}: install it beside restate-sdk[langgraph]" + ) from error + return typing.cast(BaseChatModel, getattr(integration, class_name)(model=model_id)) + + +class _RestateTool(BaseTool): + """A ``restate.agents`` tool as a LangChain tool: it gets its tool call's id, and checks the model's arguments itself.""" + + _tool: Tool = PrivateAttr() + _ctx: SessionContext = PrivateAttr() + + def _run(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError(f"{self.name} runs asynchronously only") + + async def ainvoke(self, input: Any, config: RunnableConfig | None = None, **kwargs: Any) -> ToolMessage: + call = typing.cast(dict[str, Any], input) + call_id = call.get("id") if isinstance(input, dict) else None + if not call_id: + raise ValueError(f"{self.name}: a tool call without an id") + text = await self._tool.run(call.get("args", {}), ToolCall(self._ctx, str(call_id))) + return ToolMessage(content=text, tool_call_id=str(call_id), name=self.name) + + +def _native(tool: Tool, ctx: SessionContext) -> BaseTool: + native = _RestateTool(name=tool.name, description=tool.description, args_schema=tool.input_schema) + native._tool = tool # pyright: ignore[reportPrivateUsage] + native._ctx = ctx # pyright: ignore[reportPrivateUsage] + return native + + +def _tools_node(ctx: SessionContext, tools: list[BaseTool]) -> Callable[[State, RunnableConfig], Awaitable[Any]]: + """ + The tools node: each tool call of the model's last message, run by a + ToolNode, which answers the model with what is wrong when a call names a + tool the agent lacks or arguments that do not fit; any other error of a + tool fails the attempt. The round is published as it starts and ends, for + the clients that follow the conversation. + """ + node = ToolNode(tools) + + async def run_tools(state: State, config: RunnableConfig) -> Any: + await tool_call_started(ctx) + try: + ran = await node.ainvoke(state, config) + except Exception as error: + # A tool's error, never taken for the model's provider refusing. + from_tool(ctx, error) + raise + await tool_call_ended(ctx) + return ran + + return run_tools + + +def _text(message: BaseMessage) -> str: + return str(message.text) + + +def _graph(ctx: SessionContext, model: BaseChatModel, tools: list[BaseTool], system: str | None) -> Any: + """The agent's graph: ``model``, calling ``tools``, until it answers the prompt.""" + bound = model.bind_tools(tools) if tools else model + + async def call_model(state: State) -> dict[str, Any]: + messages = state.get("messages", []) + reply = await bound.ainvoke([SystemMessage(system), *messages] if system else messages) + return {"messages": [reply]} + + def answer(state: State) -> dict[str, Any]: + messages = state.get("messages", []) + return {"answers": {state.get("prompt", ""): _text(messages[-1]) if messages else ""}} + + graph = StateGraph(State) + graph.add_node("model", call_model) + graph.add_node("tools", typing.cast(Any, _tools_node(ctx, tools))) + graph.add_node("answer", answer) + graph.add_edge(START, "model") + graph.add_conditional_edges("model", tools_condition, {"tools": "tools", END: "answer"}) + graph.add_edge("tools", "model") + graph.add_edge("answer", END) + return graph + + +Compiled = CompiledStateGraph[Any, Any, Any, Any] + + +async def _answer_once(ctx: SessionContext, text: str, options: Options) -> str: + model = options.model if isinstance(options.model, BaseChatModel) else _chat_model(model_of(options.model)) + tools = [_native(t, ctx) if isinstance(t, Tool) else t for t in options.tools] + # The latest checkpoint is all an agent resumes from: earlier ones would store the conversation again at each step. + app: Compiled = _graph(ctx, model, tools, options.system).compile( + checkpointer=SqlCheckpointSaver(ctx.db, history="latest") + ) + config: RunnableConfig = {"configurable": {"thread_id": ctx.key}} + prompt_id = next_prompt_id(ctx) + while True: + values = typing.cast(State, (await app.aget_state(config)).values) + answers = values.get("answers", {}) + if prompt_id in answers: + return stored_answer(answers[prompt_id]) + # Finish an interrupted run first (this prompt's, after a crash), then start this prompt's. + # A run is unfinished while its prompt has no answer. + last = values.get("prompt") + interrupted = last is not None and last not in answers + given: State | None = None if interrupted else {"messages": [HumanMessage(text)], "prompt": prompt_id} + stop = stop_signal_of(ctx) + run = app.ainvoke(given, config, durability="sync") + try: + await (stop.guard(run) if stop is not None else run) + except AttemptEnded: + # The attempt ended: it stops, and the next attempt goes on with the run. + raise + except TerminalError as error: + # Cancelled: the run ends, and the next prompt starts afresh. + if error.status_code == _CANCELLATION: + await _close(app, config, CANCELLED) + raise + # Failed for good (a LangChain tool's terminal error): the run ends with it, rather than the next + # prompt meeting it again. An earlier prompt's, this one goes on. + await _close(app, config, f"Error: {error.message}") + if not interrupted or last == prompt_id: + raise + except Exception as error: + # Refused for good (a wrong API key): the run ends unanswered, and fails as terminal. + # An earlier prompt's, this one goes on. + refused = refusal_of(ctx, error) + if refused is None: + raise + await _close(app, config, UNANSWERED) + if not interrupted or last == prompt_id: + raise refused from error + + +async def _close(app: Compiled, config: RunnableConfig, text: str) -> None: + """ + End the run of ``app`` that was cancelled, refused or failed, as its ``answer`` + node: a result of ``text`` for each tool call of the last reply that has + none, then ``text`` as the reply and the prompt's answer. A run that + stored no input yet (cancelled at once) has nothing to end: the prompt the + state has is an earlier one, which has its answer. + """ + values = typing.cast(State, (await app.aget_state(config)).values) + prompt = values.get("prompt") + if prompt is None or prompt in values.get("answers", {}): + return + messages = values.get("messages", []) + answered = {m.tool_call_id for m in messages if isinstance(m, ToolMessage)} + last = next((m for m in reversed(messages) if isinstance(m, AIMessage)), None) + calls = [c for c in (last.tool_calls if last else []) if c.get("id") and c.get("id") not in answered] + closing: list[BaseMessage] = [ + *(ToolMessage(text, tool_call_id=str(c.get("id")), status="error") for c in calls), + AIMessage(text), + ] + await app.aupdate_state(config, {"messages": closing, "answers": {prompt: text}}, as_node="answer") + + +async def _transcript(ctx: SessionContext) -> Transcript: + """The thread (the object's key) from its last checkpoint, as ``messages`` gives it.""" + transcript = Transcript() + saved = await SqlCheckpointSaver(ctx.db).aget_tuple({"configurable": {"thread_id": ctx.key}}) + messages: list[BaseMessage] = saved.checkpoint["channel_values"].get("messages", []) if saved else [] + for m in messages: + if isinstance(m, HumanMessage): + transcript.say("user", _text(m)) + elif isinstance(m, AIMessage): + transcript.say("assistant", _text(m)) + for call in m.tool_calls: + transcript.call(call.get("id") or "", call["name"], call["args"]) + elif isinstance(m, ToolMessage): + transcript.result(m.tool_call_id, _text(m), m.status == "error") + return transcript + + +async def prompt(ctx: SessionContext, text: str, options: Options | None = None) -> str: + """ + Answer ``text`` with a LangGraph agent of ``options``, once: a retried + attempt resumes the same prompt. Each prompt of an invocation is a prompt + of its own. The thread is the object's key. In a handler of your own:: + + @assistant.handler + async def summarize(ctx: sessions.SessionContext, topic: str) -> str: + return await langgraph.prompt(ctx, f"Summarize what we said about {topic}") + """ + chosen = options or Options() + + async def read() -> Transcript: + return await _transcript(ctx) + + async def work() -> str: + return await _answer_once(ctx, text, chosen) + + return await transcribing(ctx, read, work) + + +def chat(options: Configure[Options] | None = None) -> Callable[[SessionContext, str], Awaitable[str]]: + """ + A chat turn for ``AgentObject.chat``: each message answered by ``prompt`` + with ``options``, or the options a function of the invocation's context + gives (each user's system prompt, say):: + + assistant = AgentObject("assistant") + assistant.chat(langgraph.chat(lambda ctx: langgraph.Options(system=system_of(ctx)))) + """ + given: Configure[Options] = options if options is not None else Options() + if isinstance(given, Options) and not isinstance(given.model, BaseChatModel): + _provider_of_model(model_of(given.model)) + + async def answer(ctx: SessionContext, text: str) -> str: + return await prompt(ctx, text, await configured(given, ctx)) + + return answer + + +def agent( + name: str, + *, + model: str | BaseChatModel | None = None, + system: str | None = None, + tools: Sequence[Tool | BaseTool] = (), + description: str | None = None, + schema: Sequence[Step] = (), + background: bool = False, + metadata: dict[str, str] | None = None, + inactivity_timeout: timedelta | None = None, + abort_timeout: timedelta | None = None, + ingress_private: bool | None = None, + invocation_retry_policy: InvocationRetryPolicy | None = RETRIES, +) -> AgentObject: + """ + A LangGraph agent: an ``AgentObject`` ``name``, one conversation per key, + whose ``chat`` answers with the usual graph of ``model``, ``system`` and + ``tools`` (see ``Options``):: + + assistant = langgraph.agent("assistant", system="Answer briefly.", tools=[get_weather]) + sessions.serve(assistant) + """ + obj = AgentObject( + name, + description=description, + schema=schema, + background=background, + metadata=metadata, + inactivity_timeout=inactivity_timeout, + abort_timeout=abort_timeout, + ingress_private=ingress_private, + invocation_retry_policy=invocation_retry_policy, + ) + obj.chat(chat(Options(model=model, system=system, tools=tools))) + return obj diff --git a/tests/agents_imports.py b/tests/agents_imports.py index 0da90074..78c9bc4a 100644 --- a/tests/agents_imports.py +++ b/tests/agents_imports.py @@ -35,3 +35,21 @@ def test_sessions_and_agents_import_without_pydantic_and_a_tool_says_what_it_nee ran = subprocess.run([sys.executable, "-c", WITHOUT_PYDANTIC], capture_output=True, text=True, check=False) assert ran.returncode == 0, ran.stderr assert "tools take their inputs' schemas from pydantic" in ran.stdout + + +WITHOUT_LANGGRAPH = """ +import sys +for name in ("langgraph", "langchain_core"): + sys.modules[name] = None +import restate.agents, restate.agents.internal +try: + import restate.agents.langgraph +except ImportError as error: + print(error) +""" + + +def test_agents_import_without_langgraph_and_its_integration_names_its_extra(): + ran = subprocess.run([sys.executable, "-c", WITHOUT_LANGGRAPH], capture_output=True, text=True, check=False) + assert ran.returncode == 0, ran.stderr + assert "install the langgraph extra, pip install restate-sdk[langgraph]" in ran.stdout diff --git a/tests/agents_langgraph.py b/tests/agents_langgraph.py new file mode 100644 index 00000000..47cb15ce --- /dev/null +++ b/tests/agents_langgraph.py @@ -0,0 +1,294 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""A LangGraph agent on a scripted chat model, against a real Restate server: each point of _agent.py's docstring.""" + +# pylint: disable=missing-function-docstring,missing-class-docstring + +import asyncio +import collections +import sys +import typing +import uuid +from collections.abc import Awaitable, Callable, Sequence +from typing import Any + +import httpx +import pytest + +pytest.importorskip("langgraph") + +from langchain_core.callbacks import CallbackManagerForLLMRun # noqa: E402 +from langchain_core.language_models import BaseChatModel # noqa: E402 +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage # noqa: E402 +from langchain_core.outputs import ChatGeneration, ChatResult # noqa: E402 +from langchain_core.tools import tool as langchain_tool # noqa: E402 + +import restate # noqa: E402 +from restate import HarnessEnvironment, sessions # noqa: E402 +from restate.agents import CANCELLED, AgentObject, tool # noqa: E402 +from restate.agents import langgraph # noqa: E402 +from restate.agents._api import Messages # noqa: E402 + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +pytestmark = [pytest.mark.anyio] + + +class AuthenticationError(Exception): + """As a model SDK names the error of a missing or wrong API key.""" + + +asked: collections.Counter[str] = collections.Counter() +"""How often the model was asked, by the first word of the conversation's first message (each test's own).""" + + +def reply(messages: Sequence[BaseMessage]) -> AIMessage: + """ + The scripted model: it answers the last prompt from the conversation alone, + so that a retried attempt that asks again gets the same reply. + """ + prompt = next(m for m in reversed(messages) if isinstance(m, HumanMessage)) + text = str(prompt.content) + after = messages[list(messages).index(prompt) + 1 :] + results = [m for m in after if isinstance(m, ToolMessage)] + word, _, rest = text.partition(" ") + if word == "system": + system = next((m for m in messages if isinstance(m, SystemMessage)), None) + return AIMessage(f"system: {system.content if system else None}") + if word == "refuse": + raise AuthenticationError("invalid x-api-key") + if word in ("weather", "crash", "wait", "wrong", "unknown", "forbidden") and not results: + name = { + "weather": "get_weather", + "crash": "flaky", + "wait": "blocking", + "wrong": "get_weather", + "forbidden": "forbidden", + }.get(word, "nonexistent") + args: dict[str, Any] = {"town": 3} if word == "wrong" else {"city": rest} + return AIMessage("", tool_calls=[{"id": f"call-{rest}", "name": name, "args": args}]) + if results: + return AIMessage(f"tool said: {results[-1].content}") + return AIMessage(f"echo: {text}") + + +class Scripted(BaseChatModel): + @property + def _llm_type(self) -> str: + return "scripted" + + def bind_tools(self, tools: Any, **kwargs: Any) -> Any: # pyright: ignore[reportIncompatibleMethodOverride] + return self + + def _generate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + first = next(m for m in messages if isinstance(m, HumanMessage)) + asked[str(first.content).split(" ")[-1]] += 1 + return ChatResult(generations=[ChatGeneration(message=reply(messages))]) + + +@tool +async def get_weather(city: str) -> str: + """The current weather in a city.""" + return f"sunny in {city}" + + +crashes: collections.Counter[str] = collections.Counter() +calls_seen: list[str] = [] + + +@tool +async def flaky(city: str, call_id: str) -> str: + """Fails its first run, as a crash would.""" + calls_seen.append(call_id) + crashes[city] += 1 + if crashes[city] == 1: + raise RuntimeError("the process crashed") + return f"recovered in {city}" + + +gate = asyncio.Event() +blocked: set[str] = set() + + +@tool +async def blocking(city: str) -> str: + """Waits until it is stopped.""" + blocked.add(city) + await gate.wait() + return "never" + + +@langchain_tool +async def forbidden(city: str) -> str: + """A LangChain tool of the agent's that fails for good.""" + raise restate.TerminalError(f"no access to {city}") + + +assistant = langgraph.agent( + "lg_assistant", + model=Scripted(), + system="Answer briefly.", + tools=[get_weather, flaky, blocking, forbidden], +) + + +@assistant.handler +async def twice(ctx: sessions.SessionContext, text: str) -> list[str]: + """Two prompts of one invocation: each its own.""" + options = langgraph.Options(model=Scripted()) + return [await langgraph.prompt(ctx, text, options), await langgraph.prompt(ctx, f"again {text}", options)] + + +@pytest.fixture(scope="session") +async def harness(): + async with restate.create_test_harness( + sessions.app(assistant), restate_image="ghcr.io/restatedev/restate:latest" + ) as harness: + yield harness + + +TEXT = {"content-type": "text/plain"} + + +async def post(harness: HarnessEnvironment, path: str, **kwargs: Any) -> httpx.Response: + async with httpx.AsyncClient(timeout=60) as ingress: + return await ingress.post(f"{harness.ingress_url}/lg_assistant/{path}", **kwargs) + + +async def chat(harness: HarnessEnvironment, key: str, text: str) -> httpx.Response: + return await post(harness, f"{key}/chat", content=text, headers=TEXT) + + +async def messages(harness: HarnessEnvironment, key: str) -> list[tuple[str, str]]: + page = typing.cast(Messages, (await post(harness, f"{key}/messages", json={})).json()) + return [ + (message["role"], " ".join(str(part.get("text", part.get("output", ""))) for part in message["parts"])) + for message in page["messages"] + ] + + +async def until(check: Callable[[], Awaitable[bool]]) -> None: + for _ in range(200): + if await check(): + return + await asyncio.sleep(0.05) + raise AssertionError("never") + + +def key() -> str: + return uuid.uuid4().hex + + +async def test_a_message_is_answered_with_the_system_prompt_and_kept_in_the_thread(harness: HarnessEnvironment): + k = key() + assert (await chat(harness, k, f"hi {k}")).json() == f"echo: hi {k}" + assert (await chat(harness, k, f"system {k}")).json() == "system: Answer briefly." + assert await messages(harness, k) == [ + ("user", f"hi {k}"), + ("assistant", f"echo: hi {k}"), + ("user", f"system {k}"), + ("assistant", "system: Answer briefly."), + ] + + +async def test_a_tool_call_is_made_and_its_result_answered(harness: HarnessEnvironment): + k = key() + assert (await chat(harness, k, f"weather {k}")).json() == f"tool said: sunny in {k}" + assert ("assistant", f"sunny in {k}") in await messages(harness, k) + + +async def test_a_crash_in_a_tool_resumes_it_with_the_same_call_id_without_asking_the_model_again( + harness: HarnessEnvironment, +): + k = key() + assert (await chat(harness, k, f"crash {k}")).json() == f"tool said: recovered in {k}" + assert crashes[k] == 2 + assert calls_seen.count(f"call-{k}") == 2 + # Asked for the tool call once, then for the answer once: the retry resumed at the tools. + assert asked[k] == 2 + + +async def test_arguments_that_do_not_fit_and_an_unknown_tool_are_the_models_to_fix(harness: HarnessEnvironment): + k = key() + wrong = (await chat(harness, k, f"wrong {k}")).json() + assert wrong.startswith("tool said: Error: the arguments do not fit") + unknown = (await chat(harness, key(), f"unknown {k}")).json() + assert "nonexistent is not a valid tool" in unknown + + +async def test_a_refused_message_fails_with_424_and_the_next_is_answered(harness: HarnessEnvironment): + k = key() + refused = await chat(harness, k, f"refuse {k}") + assert refused.status_code == 424, refused.text + assert (await chat(harness, k, f"hi {k}")).json() == f"echo: hi {k}" + + +async def test_a_stopped_turn_ends_cancelled_and_the_next_message_starts_afresh(harness: HarnessEnvironment): + k = key() + async with httpx.AsyncClient(timeout=60) as ingress: + sent = await ingress.post( + f"{harness.ingress_url}/lg_assistant/{k}/chat/send", content=f"wait {k}", headers=TEXT + ) + invocation = sent.json()["invocationId"] + + async def in_the_tool() -> bool: + return k in blocked + + await until(in_the_tool) + await post(harness, f"{k}/stop", json=None) + ended = await ingress.get(f"{harness.ingress_url}/restate/invocation/{invocation}/attach") + assert ended.status_code == 409, ended.text + assert (await chat(harness, k, f"hi {k}")).json() == f"echo: hi {k}" + thread = await messages(harness, k) + assert ("assistant", CANCELLED) in thread + assert thread[-1] == ("assistant", f"echo: hi {k}") + + +async def test_each_prompt_of_an_invocation_is_its_own(harness: HarnessEnvironment): + k = key() + answered = await post(harness, f"{k}/twice", json=f"hi {k}") + assert answered.json() == [f"echo: hi {k}", f"echo: again hi {k}"] + + +def test_an_agent_refuses_a_model_it_cannot_name(): + with pytest.raises(restate.TerminalError, match="neither openai/ nor anthropic/"): + langgraph.agent("bad", model="mistral/large") + with pytest.raises(restate.TerminalError, match="neither"): + AgentObject("bad2").chat(langgraph.chat(langgraph.Options(model="gpt"))) + + +def test_a_named_model_whose_provider_package_is_missing_says_which_to_install(monkeypatch: pytest.MonkeyPatch): + from restate.agents.langgraph import _agent # pylint: disable=import-outside-toplevel + + monkeypatch.setitem(sys.modules, "langchain_anthropic", None) + with pytest.raises(restate.TerminalError, match="need langchain-anthropic: install it beside") as refused: + _agent._chat_model("anthropic/claude-sonnet-5-5") # pyright: ignore[reportPrivateUsage] + assert refused.value.status_code == 424 + + +async def test_a_langchain_tools_terminal_error_ends_the_turn_and_the_next_message_is_answered( + harness: HarnessEnvironment, +): + k = key() + failed = await chat(harness, k, f"forbidden {k}") + assert failed.status_code == 500 and f"no access to {k}" in failed.text + assert (await chat(harness, k, f"hi {k}")).json() == f"echo: hi {k}" + assert ("assistant", f"Error: no access to {k}") in await messages(harness, k) diff --git a/tests/examples_agents.py b/tests/examples_agents.py index e92ed864..595dcc55 100644 --- a/tests/examples_agents.py +++ b/tests/examples_agents.py @@ -147,6 +147,14 @@ async def test_basic_notes(): assert (await ingress.post("notes/alice/list")).json() == ["Buy milk"] +async def test_basic_langgraph_agent(): + """The agent serves, with its companion; a message needs a model's provider, and is left to a real run.""" + pytest.importorskip("langgraph") + async with serving("basic/langgraph_agent.py") as ingress: + assert (await ingress.post("assistant/alice/messages", {})).json()["messages"] == [] + assert (await ingress.post("assistant/alice/status")).json() == {"answering": None, "waiting": []} + + def serves_http2() -> bool: """Whether ``sessions.serve`` can run here: uvicorn 0.53+ with zttp.""" return importlib.util.find_spec("uvicorn") is not None and importlib.util.find_spec("zttp") is not None diff --git a/uv.lock b/uv.lock index bbc54b3a..5206f024 100644 --- a/uv.lock +++ b/uv.lock @@ -3313,6 +3313,7 @@ langchain = [ langgraph = [ { name = "langchain-core" }, { name = "langgraph" }, + { name = "pydantic" }, ] lint = [ { name = "mypy" }, @@ -3365,6 +3366,7 @@ requires-dist = [ { name = "mypy", marker = "extra == 'lint'", specifier = ">=1.11.2" }, { name = "openai-agents", marker = "extra == 'openai'", specifier = ">=0.6.1" }, { name = "opentelemetry-api", marker = "extra == 'tracing'", specifier = ">=1.36.0" }, + { name = "pydantic", marker = "extra == 'langgraph'", specifier = ">=2" }, { name = "pydantic", marker = "extra == 'serde'" }, { name = "pydantic-ai-slim", marker = "extra == 'pydantic-ai'", specifier = ">=2.0,<3" }, { name = "pyright", marker = "extra == 'lint'", specifier = ">=1.1.390" }, From b1d3a093545d8795061dd1ca9fd216fbb6f57f51 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 17:24:17 +0100 Subject: [PATCH 30/45] restate.agents.langgraph: keyword arguments, options as functions of ctx agent(name, model=, system=, tools=) and prompt(ctx, text, model=, system=, tools=); each of agent()'s options may be a value or a (sync or async) function of the invocation's context. Options and chat() are gone: one way to make an agent. Google-style docstrings. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/agents/langgraph/__init__.py | 17 +-- python/restate/agents/langgraph/_agent.py | 142 ++++++++++++-------- tests/agents_langgraph.py | 28 +++- 3 files changed, 120 insertions(+), 67 deletions(-) diff --git a/python/restate/agents/langgraph/__init__.py b/python/restate/agents/langgraph/__init__.py index 92fe9dd2..04504f9e 100644 --- a/python/restate/agents/langgraph/__init__.py +++ b/python/restate/agents/langgraph/__init__.py @@ -23,12 +23,13 @@ async def get_weather(city: str) -> str: assistant = langgraph.agent("assistant", system="Answer briefly.", tools=[get_weather]) sessions.serve(assistant) -It needs the ``langgraph`` extra (``pip install restate-sdk[langgraph]``); -a model named ``openai/...`` needs ``langchain-openai`` beside it, and -``anthropic/...`` ``langchain-anthropic``. Or bring a LangChain chat model of -your own (``model=``). ``agent`` makes an agent; ``chat`` is the turn of an -``AgentObject`` of your own, ``prompt`` a prompt of a handler of your own, -and ``SqlCheckpointSaver`` the checkpointer, for a graph of your own. +``agent`` creates an agent; ``prompt`` answers a prompt in a handler of your +own; ``SqlCheckpointSaver`` keeps a graph of your own in ``ctx.db``. + +Install the ``langgraph`` extra (``pip install restate-sdk[langgraph]``). A +model named ``openai/...`` also needs ``langchain-openai``, and +``anthropic/...`` needs ``langchain-anthropic``; or pass any LangChain chat +model as ``model``. """ try: @@ -40,7 +41,7 @@ async def get_weather(city: str) -> str: "restate.agents.langgraph needs LangGraph: install the langgraph extra, pip install restate-sdk[langgraph]" ) from error -from restate.agents.langgraph._agent import Options, agent, chat, prompt +from restate.agents.langgraph._agent import AnyTool, ChatModel, agent, prompt from restate.agents.langgraph._checkpointer import SqlCheckpointSaver -__all__ = ["Options", "SqlCheckpointSaver", "agent", "chat", "prompt"] +__all__ = ["AnyTool", "ChatModel", "SqlCheckpointSaver", "agent", "prompt"] diff --git a/python/restate/agents/langgraph/_agent.py b/python/restate/agents/langgraph/_agent.py index bc6a5d82..cbf0e0fe 100644 --- a/python/restate/agents/langgraph/_agent.py +++ b/python/restate/agents/langgraph/_agent.py @@ -36,7 +36,6 @@ import os import typing from collections.abc import Awaitable, Callable, Sequence -from dataclasses import dataclass from datetime import timedelta from typing import Annotated, Any, TypedDict @@ -90,19 +89,11 @@ class State(TypedDict, total=False): """The answer to each prompt, by id.""" -@dataclass(frozen=True) -class Options: - """ - A LangGraph agent's options. ``model``: ``openai/model-id`` or - ``anthropic/model-id`` (MODEL, or else ``openai/gpt-6-sol``), or a - LangChain chat model that calls tools. ``system``: the system prompt. - ``tools``: ``restate.agents`` tools (``tool``, ``call``, ``signal``...), or - LangChain's. - """ +ChatModel: typing.TypeAlias = str | BaseChatModel +"""A model: ``"openai/"`` or ``"anthropic/"``, or a LangChain chat model that calls tools.""" - model: str | BaseChatModel | None = None - system: str | None = None - tools: Sequence[Tool | BaseTool] = () +AnyTool: typing.TypeAlias = Tool | BaseTool +"""A tool: one of ``restate.agents`` (``tool``, ``call``, ``signal``...), or a LangChain tool.""" _PROVIDERS = { @@ -218,11 +209,13 @@ def answer(state: State) -> dict[str, Any]: Compiled = CompiledStateGraph[Any, Any, Any, Any] -async def _answer_once(ctx: SessionContext, text: str, options: Options) -> str: - model = options.model if isinstance(options.model, BaseChatModel) else _chat_model(model_of(options.model)) - tools = [_native(t, ctx) if isinstance(t, Tool) else t for t in options.tools] +async def _answer_once( + ctx: SessionContext, text: str, model: ChatModel | None, system: str | None, given: Sequence[AnyTool] +) -> str: + chat_model = model if isinstance(model, BaseChatModel) else _chat_model(model_of(model)) + tools = [_native(t, ctx) if isinstance(t, Tool) else t for t in given] # The latest checkpoint is all an agent resumes from: earlier ones would store the conversation again at each step. - app: Compiled = _graph(ctx, model, tools, options.system).compile( + app: Compiled = _graph(ctx, chat_model, tools, system).compile( checkpointer=SqlCheckpointSaver(ctx.db, history="latest") ) config: RunnableConfig = {"configurable": {"thread_id": ctx.key}} @@ -236,9 +229,9 @@ async def _answer_once(ctx: SessionContext, text: str, options: Options) -> str: # A run is unfinished while its prompt has no answer. last = values.get("prompt") interrupted = last is not None and last not in answers - given: State | None = None if interrupted else {"messages": [HumanMessage(text)], "prompt": prompt_id} + start: State | None = None if interrupted else {"messages": [HumanMessage(text)], "prompt": prompt_id} stop = stop_signal_of(ctx) - run = app.ainvoke(given, config, durability="sync") + run = app.ainvoke(start, config, durability="sync") try: await (stop.guard(run) if stop is not None else run) except AttemptEnded: @@ -305,52 +298,52 @@ async def _transcript(ctx: SessionContext) -> Transcript: return transcript -async def prompt(ctx: SessionContext, text: str, options: Options | None = None) -> str: - """ - Answer ``text`` with a LangGraph agent of ``options``, once: a retried - attempt resumes the same prompt. Each prompt of an invocation is a prompt - of its own. The thread is the object's key. In a handler of your own:: +async def prompt( + ctx: SessionContext, + text: str, + *, + model: ChatModel | None = None, + system: str | None = None, + tools: Sequence[AnyTool] = (), +) -> str: + """Answer a prompt with a LangGraph agent, exactly once. + + For a handler of your own; the conversation is the object's key, shared + with its ``chat``. A retried attempt resumes the same prompt rather than + asking again, and each prompt of an invocation is answered on its own:: @assistant.handler async def summarize(ctx: sessions.SessionContext, topic: str) -> str: return await langgraph.prompt(ctx, f"Summarize what we said about {topic}") + + Args: + ctx: The handler's context. + text: The prompt. + model: ``"openai/"`` or ``"anthropic/"``, or a + LangChain chat model that calls tools. Defaults to ``$MODEL``, + else ``"openai/gpt-6-sol"``. + system: The system prompt. + tools: ``restate.agents`` tools, or LangChain tools. + + Returns: + The agent's answer. """ - chosen = options or Options() async def read() -> Transcript: return await _transcript(ctx) async def work() -> str: - return await _answer_once(ctx, text, chosen) + return await _answer_once(ctx, text, model, system, tools) return await transcribing(ctx, read, work) -def chat(options: Configure[Options] | None = None) -> Callable[[SessionContext, str], Awaitable[str]]: - """ - A chat turn for ``AgentObject.chat``: each message answered by ``prompt`` - with ``options``, or the options a function of the invocation's context - gives (each user's system prompt, say):: - - assistant = AgentObject("assistant") - assistant.chat(langgraph.chat(lambda ctx: langgraph.Options(system=system_of(ctx)))) - """ - given: Configure[Options] = options if options is not None else Options() - if isinstance(given, Options) and not isinstance(given.model, BaseChatModel): - _provider_of_model(model_of(given.model)) - - async def answer(ctx: SessionContext, text: str) -> str: - return await prompt(ctx, text, await configured(given, ctx)) - - return answer - - def agent( name: str, *, - model: str | BaseChatModel | None = None, - system: str | None = None, - tools: Sequence[Tool | BaseTool] = (), + model: Configure[ChatModel | None] = None, + system: Configure[str | None] = None, + tools: Configure[Sequence[AnyTool]] = (), description: str | None = None, schema: Sequence[Step] = (), background: bool = False, @@ -360,14 +353,45 @@ def agent( ingress_private: bool | None = None, invocation_retry_policy: InvocationRetryPolicy | None = RETRIES, ) -> AgentObject: - """ - A LangGraph agent: an ``AgentObject`` ``name``, one conversation per key, - whose ``chat`` answers with the usual graph of ``model``, ``system`` and - ``tools`` (see ``Options``):: + """Create a LangGraph agent: a virtual object with one conversation per key. + + Its ``chat`` handler answers each message with LangGraph's tool-calling + loop; LangGraph's checkpoints are kept in the key's database. Add + handlers of your own with ``@assistant.handler``:: assistant = langgraph.agent("assistant", system="Answer briefly.", tools=[get_weather]) - sessions.serve(assistant) + + ``model``, ``system`` and ``tools`` may each be a function of the + invocation's context, sync or async, for values that differ per user:: + + assistant = langgraph.agent("assistant", system=lambda ctx: f"You help {ctx.key}.") + + Args: + name: The object's name, as Restate serves it. + model: ``"openai/"`` (needs ``langchain-openai``) or + ``"anthropic/"`` (needs ``langchain-anthropic``), or a + LangChain chat model that calls tools. Defaults to ``$MODEL``, + else ``"openai/gpt-6-sol"``. + system: The system prompt. + tools: ``restate.agents`` tools, or LangChain tools. + description: The object's description, shown in Restate's UI. + schema: Tables of your own, beside the agent's. + background: Whether it works unwatched (a subagent, say): it then + publishes its progress once a turn ends, not at each step. + metadata: The object's metadata. + inactivity_timeout: Restate's inactivity timeout for its handlers. + abort_timeout: Restate's abort timeout for its handlers. + ingress_private: Whether only other services may call it. + invocation_retry_policy: How Restate retries a failed turn. + + Returns: + The agent, to serve with ``sessions.serve`` or ``sessions.app``. + + Raises: + TerminalError: If ``model`` names a provider other than OpenAI or Anthropic. """ + if model is None or isinstance(model, str): + _provider_of_model(model_of(model)) obj = AgentObject( name, description=description, @@ -379,5 +403,15 @@ def agent( ingress_private=ingress_private, invocation_retry_policy=invocation_retry_policy, ) - obj.chat(chat(Options(model=model, system=system, tools=tools))) + + async def answer(ctx: SessionContext, text: str) -> str: + return await prompt( + ctx, + text, + model=await configured(model, ctx), + system=await configured(system, ctx), + tools=await configured(tools, ctx), + ) + + obj.chat(answer) return obj diff --git a/tests/agents_langgraph.py b/tests/agents_langgraph.py index 47cb15ce..cffc3b58 100644 --- a/tests/agents_langgraph.py +++ b/tests/agents_langgraph.py @@ -33,7 +33,7 @@ import restate # noqa: E402 from restate import HarnessEnvironment, sessions # noqa: E402 -from restate.agents import CANCELLED, AgentObject, tool # noqa: E402 +from restate.agents import CANCELLED, tool # noqa: E402 from restate.agents import langgraph # noqa: E402 from restate.agents._api import Messages # noqa: E402 @@ -153,14 +153,24 @@ async def forbidden(city: str) -> str: @assistant.handler async def twice(ctx: sessions.SessionContext, text: str) -> list[str]: """Two prompts of one invocation: each its own.""" - options = langgraph.Options(model=Scripted()) - return [await langgraph.prompt(ctx, text, options), await langgraph.prompt(ctx, f"again {text}", options)] + model = Scripted() + return [await langgraph.prompt(ctx, text, model=model), await langgraph.prompt(ctx, f"again {text}", model=model)] + + +async def tools_of(ctx: sessions.SessionContext) -> list[Any]: + """Each user's tools, as an async function of the context gives them.""" + return [get_weather] + + +per_user = langgraph.agent( + "lg_per_user", model=lambda ctx: Scripted(), system=lambda ctx: f"You help {ctx.key}.", tools=tools_of +) @pytest.fixture(scope="session") async def harness(): async with restate.create_test_harness( - sessions.app(assistant), restate_image="ghcr.io/restatedev/restate:latest" + sessions.app(assistant, per_user), restate_image="ghcr.io/restatedev/restate:latest" ) as harness: yield harness @@ -272,7 +282,7 @@ def test_an_agent_refuses_a_model_it_cannot_name(): with pytest.raises(restate.TerminalError, match="neither openai/ nor anthropic/"): langgraph.agent("bad", model="mistral/large") with pytest.raises(restate.TerminalError, match="neither"): - AgentObject("bad2").chat(langgraph.chat(langgraph.Options(model="gpt"))) + langgraph.agent("bad2", model="gpt") def test_a_named_model_whose_provider_package_is_missing_says_which_to_install(monkeypatch: pytest.MonkeyPatch): @@ -292,3 +302,11 @@ async def test_a_langchain_tools_terminal_error_ends_the_turn_and_the_next_messa assert failed.status_code == 500 and f"no access to {k}" in failed.text assert (await chat(harness, k, f"hi {k}")).json() == f"echo: hi {k}" assert ("assistant", f"Error: no access to {k}") in await messages(harness, k) + + +async def test_options_may_be_functions_of_the_context(harness: HarnessEnvironment): + k = key() + async with httpx.AsyncClient(timeout=60) as ingress: + url = f"{harness.ingress_url}/lg_per_user/{k}/chat" + assert (await ingress.post(url, content=f"system {k}", headers=TEXT)).json() == f"system: You help {k}." + assert (await ingress.post(url, content=f"weather {k}", headers=TEXT)).json() == f"tool said: sunny in {k}" From bd3255e286294f43398f32604b939d8614bd2359 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 17:24:17 +0100 Subject: [PATCH 31/45] restate.agents.pydantic_ai: Pydantic AI agents, their conversation in ctx.db sqlexperiment's integration on this package: the conversation stored before each model request and with each response (capabilities hooks), each prompt answered once by id, a crash in a tool resumed with the same call id without asking the model again, a stop closing the run as cancelled, a tool function's TerminalError its answer. agent(name, [Agent], model=, system=, tools=) and prompt(ctx, text, ...), as in langgraph; this package's tools are bound to the turn per run (toolsets). Models resolve before the turn: a missing provider extra, API key or unknown provider is a 424 naming it. A provider's ModelHTTPError 400/ 401/403/404 is a refusal (424) too. The pydantic_ai extra names pydantic; importing the integration without it says which extra to install. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- examples-agents/README.md | 1 + examples-agents/basic/pydantic_ai_agent.py | 31 ++ pyproject.toml | 2 +- python/restate/agents/_common.py | 6 + python/restate/agents/pydantic_ai/__init__.py | 43 ++ python/restate/agents/pydantic_ai/_agent.py | 482 ++++++++++++++++++ tests/agents_imports.py | 17 + tests/agents_pydantic_ai.py | 329 ++++++++++++ tests/examples_agents.py | 8 + uv.lock | 2 + 10 files changed, 920 insertions(+), 1 deletion(-) create mode 100644 examples-agents/basic/pydantic_ai_agent.py create mode 100644 python/restate/agents/pydantic_ai/__init__.py create mode 100644 python/restate/agents/pydantic_ai/_agent.py create mode 100644 tests/agents_pydantic_ai.py diff --git a/examples-agents/README.md b/examples-agents/README.md index c82b3fc2..0bf39533 100644 --- a/examples-agents/README.md +++ b/examples-agents/README.md @@ -67,6 +67,7 @@ the loop, a subagent), comes with the agents' framework integrations. | ----------------------------- | ----------------------------------------------- | | [`notes.py`](basic/notes.py) | SQL in a virtual object, with `ctx.db`, and no agent | | [`langgraph_agent.py`](basic/langgraph_agent.py) | A LangGraph agent: its checkpoints in the key's database | +| [`pydantic_ai_agent.py`](basic/pydantic_ai_agent.py) | A Pydantic AI agent: its conversation in the key's database | ## Tips diff --git a/examples-agents/basic/pydantic_ai_agent.py b/examples-agents/basic/pydantic_ai_agent.py new file mode 100644 index 00000000..f6051154 --- /dev/null +++ b/examples-agents/basic/pydantic_ai_agent.py @@ -0,0 +1,31 @@ +""" +A Pydantic AI agent: a weather assistant, each conversation in its key's +database, stored a step at a time. A crash mid-turn resumes from the last +stored step: a tool call already made is not made again:: + + OPENAI_API_KEY=... uv run --extra pydantic_ai --extra serve --with 'pydantic-ai-slim[openai]' examples-agents/basic/pydantic_ai_agent.py + npx @restatedev/restate-server # Restate, in another terminal + npx @restatedev/restate -y deployments register --force http://localhost:9080 + curl localhost:8080/assistant/alice/chat -d 'What is the weather in Lisbon?' + curl localhost:8080/assistant/alice/messages --json '{}' # the conversation + +The model is ``openai/gpt-6-sol`` unless MODEL names another, of any of +Pydantic AI's providers: ``anthropic/`` takes +``pydantic-ai-slim[anthropic]`` (and ANTHROPIC_API_KEY) instead. +""" + +from restate import sessions +from restate.agents import pydantic_ai, tool + + +@tool +async def get_weather(city: str) -> str: + """The current weather in a city.""" + return f"Sunny and 21°C in {city}" + + +assistant = pydantic_ai.agent("assistant", system="Answer briefly.", tools=[get_weather]) + + +if __name__ == "__main__": + sessions.serve(assistant) diff --git a/pyproject.toml b/pyproject.toml index 17730858..7c13b70a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ serve = [ ] adk = ["google-adk>=1.20.0"] openai = ["openai-agents>=0.6.1"] -pydantic_ai = ["pydantic-ai-slim>=2.0,<3"] +pydantic_ai = ["pydantic-ai-slim>=2.0,<3", "pydantic>=2"] langchain = ["langchain>=1.0.0", "langgraph>=1.0.0", "langchain-core>=1.0.0"] langgraph = ["langgraph>=1.0.0", "langchain-core>=1.0.0", "pydantic>=2"] tracing = ["opentelemetry-api>=1.36.0"] diff --git a/python/restate/agents/_common.py b/python/restate/agents/_common.py index f3d70552..6830dbc0 100644 --- a/python/restate/agents/_common.py +++ b/python/restate/agents/_common.py @@ -175,12 +175,18 @@ def from_tool(ctx: SessionContext, error: BaseException) -> None: ) +_FOR_GOOD_STATUS = frozenset({400, 401, 403, 404}) +"""The HTTP statuses of a model request that fails again however often it is made (Pydantic AI's ``ModelHTTPError``'s).""" + + def _refused(error: BaseException) -> bool: """Whether ``error`` is a model provider's refusal that a retry does not change: a missing key, an unknown model...""" if isinstance(error, ToolError): return False if type(error).__name__ in _FOR_GOOD: return True + if type(error).__name__ == "ModelHTTPError" and getattr(error, "status_code", None) in _FOR_GOOD_STATUS: + return True return bool(_REFUSED_TEXT.search(str(error))) diff --git a/python/restate/agents/pydantic_ai/__init__.py b/python/restate/agents/pydantic_ai/__init__.py new file mode 100644 index 00000000..d08ac0e9 --- /dev/null +++ b/python/restate/agents/pydantic_ai/__init__.py @@ -0,0 +1,43 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +Pydantic AI agents in session objects: each conversation in the object's +database, stored a step at a time:: + + from restate import sessions + from restate.agents import pydantic_ai, tool + + @tool + async def get_weather(city: str) -> str: + '''The current weather in a city.''' + return f"Sunny in {city}" + + assistant = pydantic_ai.agent("assistant", system="Answer briefly.", tools=[get_weather]) + sessions.serve(assistant) + +``agent`` creates an agent; ``prompt`` answers a prompt in a handler of your own. + +Install the ``pydantic_ai`` extra (``pip install restate-sdk[pydantic_ai]``). +A model's provider also needs Pydantic AI's extra for it, such as +``pydantic-ai-slim[openai]`` for ``openai/...`` models; or pass any Pydantic +AI model as ``model``. +""" + +try: + import pydantic_ai # noqa: F401 # pyright: ignore[reportUnusedImport] +except ImportError as error: + raise ImportError( + "restate.agents.pydantic_ai needs Pydantic AI: install the pydantic_ai extra, pip install restate-sdk[pydantic_ai]" + ) from error + +from restate.agents.pydantic_ai._agent import AnyModel, AnyTool, agent, prompt + +__all__ = ["AnyModel", "AnyTool", "agent", "prompt"] diff --git a/python/restate/agents/pydantic_ai/_agent.py b/python/restate/agents/pydantic_ai/_agent.py new file mode 100644 index 00000000..dc2272e6 --- /dev/null +++ b/python/restate/agents/pydantic_ai/_agent.py @@ -0,0 +1,482 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +""" +Pydantic AI agents: each conversation in the object's database, stored a step +at a time (sqlexperiment's ``restate_agents.pydantic_ai``, on this package). + +The handler is not replayed: a retried attempt (a crash, a lost connection) +runs again from the start, and the run carries on from the last step stored: + +- The conversation is stored before each model request, and with each model + response, in order with everything else the handler does. So the tool calls + a model made are stored before any tool runs, and a crash during the tools + runs them again with the same tool call ids, which a tool passes on to stay + idempotent. +- A tool's error fails the attempt, which Restate retries; a + ``restate.TerminalError``, which a retry would repeat, is the tool's answer + to the model; and so are arguments that do not fit (Pydantic AI asks the + model again). +- A message is answered once per invocation: its request is marked with the + prompt's id, its answer too, and a retried attempt that finds them returns + the answer, or finishes the run first. +""" + +from __future__ import annotations + +import typing +from collections.abc import Callable, Sequence +from datetime import timedelta +from typing import Any + +from pydantic import TypeAdapter +from pydantic_ai.capabilities import AbstractCapability +from pydantic_ai.exceptions import UserError +from pydantic_ai.messages import ( + ModelMessage, + ModelRequest, + ModelResponse, + RetryPromptPart, + TextPart, + ToolCallPart, + ToolReturnPart, + UserPromptPart, +) +from pydantic_ai.models import Model, ModelRequestContext, infer_model +from pydantic_ai.tools import ToolDefinition +from pydantic_ai.toolsets import FunctionToolset + +from pydantic_ai import Agent, RunContext +from pydantic_ai import Tool as PydanticTool +from restate.agents._common import ( + CANCELLED, + UNANSWERED, + Configure, + ModelError, + configured, + from_tool, + model_of, + next_prompt_id, + refusal_of, + stored_answer, +) +from restate.agents._object import RETRIES, AgentObject +from restate.agents._tools import Tool, ToolCall +from restate.agents._transcript import Transcript +from restate.agents._turn import AttemptEnded, stop_signal_of, tool_call_ended, tool_call_started, transcribing +from restate.exceptions import TerminalError +from restate.retry_policy import InvocationRetryPolicy +from restate.sessions import Database, SessionContext, Step +from restate.sessions.internal import Schema + +_CANCELLATION = 409 +"""The status of Restate's cancellation, as a turn that was stopped fails with it.""" + +_PROMPT = "restate_agents_prompt" +"""The metadata of a prompt's request (which models do not read): the prompt's id.""" + +_ANSWERS = "restate_agents_answers" +"""The metadata of the response that answers a prompt: the prompt's id.""" + +AnyModel: typing.TypeAlias = str | Model +"""A model: ``"/"`` (or Pydantic AI's ``":"``), or a Pydantic AI model.""" + +AnyTool: typing.TypeAlias = Tool | PydanticTool[Any] | Callable[..., Any] +"""A tool: one of ``restate.agents`` (``tool``, ``call``, ``signal``...), a Pydantic AI tool, or a tool function.""" + + +def _check(own: Agent[Any, str] | None, system: object) -> None: + if own is not None and system is not None: + raise TypeError("an Agent of your own has its own instructions: pass no system") + + +def _toolset(given: Sequence[AnyTool], ctx: SessionContext) -> FunctionToolset[Any]: + """ + ``given`` as a Pydantic AI toolset for a run of the turn of ``ctx``: + ``restate.agents`` tools, each with its tool call's id (it checks the + model's arguments itself); Pydantic AI's tools and tool functions as they are. + """ + + def native(tool: Tool) -> PydanticTool[Any]: + async def run(run_ctx: RunContext[Any], **args: Any) -> str: + return await tool.run(args, ToolCall(ctx, run_ctx.tool_call_id or "")) + + return PydanticTool.from_schema(run, tool.name, tool.description, tool.input_schema, takes_ctx=True) + + return FunctionToolset([native(t) if isinstance(t, Tool) else t for t in given]) + + +_SCHEMA = Schema( + "restate_agents.pydantic_ai", + [ + """CREATE TABLE restate_pydantic_ai_messages ( + thread TEXT NOT NULL, + n INTEGER NOT NULL, + message TEXT NOT NULL, + PRIMARY KEY (thread, n) + )""" + ], +) + +# ModelMessage is an Annotated union (discriminated), which TypeAdapter takes though its type says otherwise. +_MESSAGE: TypeAdapter[ModelMessage] = TypeAdapter(typing.cast(Any, ModelMessage)) + + +class _Conversation: + """ + A conversation in the object's database, a message per row. Saving it + writes the messages that changed since it was loaded or saved, in one + transaction. + """ + + def __init__(self, db: Database, thread: str) -> None: + self.db = db + self.thread = thread + self._saved: list[str] = [] + + async def load(self) -> list[ModelMessage]: + await _SCHEMA.ensure(self.db) + rows = await self.db.values( + "SELECT message FROM restate_pydantic_ai_messages WHERE thread = ? ORDER BY n", [self.thread] + ) + self._saved = [str(row[0]) for row in rows] + return [_MESSAGE.validate_json(m) for m in self._saved] + + async def save(self, messages: Sequence[ModelMessage]) -> None: + dumped = [_MESSAGE.dump_json(m).decode() for m in messages] + first = next( + (i for i, (a, b) in enumerate(zip(dumped, self._saved, strict=False)) if a != b), + min(len(dumped), len(self._saved)), + ) + if first == len(dumped) == len(self._saved): + return + async with self.db.transaction() as tx: + await tx.execute( + "DELETE FROM restate_pydantic_ai_messages WHERE thread = ? AND n >= ?", [self.thread, first] + ) + for n in range(first, len(dumped)): + await tx.execute( + "INSERT INTO restate_pydantic_ai_messages VALUES (?, ?, ?)", [self.thread, n, dumped[n]] + ) + self._saved = dumped + + +class _Restate(AbstractCapability[Any]): + """ + A run in a Restate handler: its conversation stored before each model + request and with each response; each tool call published as it starts and + ends, for the clients that follow the conversation; a tool's + ``restate.TerminalError`` its answer, and any other error the attempt's. + """ + + def __init__(self, ctx: SessionContext, conversation: _Conversation) -> None: + self.ctx = ctx + self.conversation = conversation + + async def before_model_request( + self, ctx: RunContext[Any], request_context: ModelRequestContext + ) -> ModelRequestContext: + await self.conversation.save(ctx.messages) + return request_context + + async def after_model_request( + self, ctx: RunContext[Any], *, request_context: ModelRequestContext, response: ModelResponse + ) -> ModelResponse: + await self.conversation.save([*ctx.messages, response]) + return response + + async def before_tool_execute( + self, ctx: RunContext[Any], *, call: ToolCallPart, tool_def: ToolDefinition, args: Any + ) -> Any: + await tool_call_started(self.ctx) + return args + + async def after_tool_execute( + self, ctx: RunContext[Any], *, call: ToolCallPart, tool_def: ToolDefinition, args: Any, result: Any + ) -> Any: + await tool_call_ended(self.ctx) + return result + + async def on_tool_execute_error( + self, ctx: RunContext[Any], *, call: ToolCallPart, tool_def: ToolDefinition, args: Any, error: Exception + ) -> Any: + stop = stop_signal_of(self.ctx) + if stop is not None and stop.reason is not None: + raise stop.reason from error + if isinstance(error, TerminalError) and error.status_code != _CANCELLATION: + await tool_call_ended(self.ctx) + return f"Error: {error.message}" + # A tool's error, never taken for the model's provider refusing. + from_tool(self.ctx, error) + raise error + + +def _model(model: str | Model | None) -> Model: + """ + The model, as Pydantic AI resolves it: ``provider/model-id`` (or + ``provider:model-id``). A provider whose package is missing, or without + its API key, refuses every message alike: a ``ModelError``. + """ + if isinstance(model, Model): + return model + named = model_of(model) + try: + return infer_model(named if ":" in named else named.replace("/", ":", 1)) + except (ImportError, UserError, ValueError) as error: + raise ModelError(f'the model "{named}": {error}') from error + + +def _prompt_of(message: ModelMessage) -> str | None: + """The id of the prompt that ``message`` is, if it is one.""" + return (message.metadata or {}).get(_PROMPT) if isinstance(message, ModelRequest) else None + + +def _answer_of(history: list[ModelMessage], asked: int) -> int | None: + """The index of the response that answers the prompt at ``asked``, if it was answered.""" + return next((i for i in range(asked, len(history)) if (history[i].metadata or {}).get(_ANSWERS)), None) + + +def _result(part: ToolReturnPart | RetryPromptPart) -> str: + """A tool call's result, as the model reads it: its return, or what to fix.""" + return part.model_response() if isinstance(part, RetryPromptPart) else part.model_response_str() + + +def _text(response: ModelMessage) -> str: + return "".join(part.content for part in response.parts if isinstance(part, TextPart)) + + +async def _close(conversation: _Conversation, prompt_id: str, text: str) -> None: + """ + End the run of prompt ``prompt_id``, which was cancelled, refused or + failed: a result of ``text`` for each tool call of its last response that + has none, then ``text`` as its answer. A run that stored no prompt yet has + nothing to end. + """ + messages = await conversation.load() + if not any(_prompt_of(m) == prompt_id for m in messages): + return + last = messages[-1] + if isinstance(last, ModelResponse) and last.tool_calls: + returns = [ToolReturnPart(call.tool_name, text, call.tool_call_id) for call in last.tool_calls] + messages.append(ModelRequest(parts=returns)) + messages.append(ModelResponse(parts=[TextPart(text)], metadata={_ANSWERS: prompt_id})) + await conversation.save(messages) + + +async def _answer_once( + ctx: SessionContext, + text: str, + own: Agent[Any, str] | None, + named: AnyModel | None, + system: str | None, + given: Sequence[AnyTool], +) -> str: + # An Agent of your own's model too, resolved here: refused as any other (a missing key or package). + model = _model(named if named is not None else own.model if own is not None else None) + agent: Agent[Any, str] = own or Agent(instructions=system) + toolsets = [_toolset(given, ctx)] if given else [] + conversation = _Conversation(ctx.db, ctx.key) + prompt_id = next_prompt_id(ctx) + while True: + history = await conversation.load() + asked = next((i for i, m in enumerate(history) if _prompt_of(m) == prompt_id), None) + if asked is not None and (answer := _answer_of(history, asked)) is not None: + return stored_answer(_text(history[answer])) + # Finish an interrupted run first (this prompt's, after a crash), then start this prompt's. + last = next((i for i in reversed(range(len(history))) if _prompt_of(history[i]) is not None), None) + interrupted = last is not None and _answer_of(history, last) is None + running = _prompt_of(history[last]) if interrupted and last is not None else prompt_id + assert running is not None + if not interrupted: + history = [*history, ModelRequest(parts=[UserPromptPart(text)], metadata={_PROMPT: prompt_id})] + stop = stop_signal_of(ctx) + run = agent.run( + message_history=history, model=model, toolsets=toolsets, capabilities=[_Restate(ctx, conversation)] + ) + try: + result = await (stop.guard(run) if stop is not None else run) + except AttemptEnded: + # The attempt ended: it stops, and the next attempt goes on with the run. + raise + except TerminalError as error: + if error.status_code == _CANCELLATION: + # Cancelled: the run ends, and the next prompt starts afresh. + await _close(conversation, running, CANCELLED) + raise + # Failed for good: the run ends with it, rather than the next prompt meeting it again. + await _close(conversation, running, f"Error: {error.message}") + if running == prompt_id: + raise + continue + except Exception as error: + # Refused for good (a wrong API key): the run ends unanswered, and fails as terminal. + # An earlier prompt's, this one goes on. + refused = refusal_of(ctx, error) + if refused is None: + raise + await _close(conversation, running, UNANSWERED) + if running == prompt_id: + raise refused from error + continue + messages = result.all_messages() + response = next(m for m in reversed(messages) if isinstance(m, ModelResponse)) + response.metadata = {**(response.metadata or {}), _ANSWERS: running} + await conversation.save(messages) + + +async def _transcript(ctx: SessionContext) -> Transcript: + """The conversation (the object's key), as ``messages`` gives it.""" + transcript = Transcript() + for message in await _Conversation(ctx.db, ctx.key).load(): + for part in message.parts: + if isinstance(part, UserPromptPart): + transcript.say("user", part.content if isinstance(part.content, str) else "") + elif isinstance(part, TextPart): + transcript.say("assistant", part.content) + elif isinstance(part, ToolCallPart): + transcript.call(part.tool_call_id, part.tool_name, part.args_as_dict()) + elif isinstance(part, ToolReturnPart | RetryPromptPart) and part.tool_call_id: + transcript.result(part.tool_call_id, _result(part), isinstance(part, RetryPromptPart)) + return transcript + + +async def prompt( + ctx: SessionContext, + text: str, + *, + agent: Agent[Any, str] | None = None, + model: AnyModel | None = None, + system: str | None = None, + tools: Sequence[AnyTool] = (), +) -> str: + """Answer a prompt with a Pydantic AI agent, exactly once. + + For a handler of your own; the conversation is the object's key, shared + with its ``chat``. A retried attempt resumes the same prompt rather than + asking again, and each prompt of an invocation is answered on its own:: + + @assistant.handler + async def summarize(ctx: sessions.SessionContext, topic: str) -> str: + return await pydantic_ai.prompt(ctx, f"Summarize what we said about {topic}") + + Args: + ctx: The handler's context. + text: The prompt. + agent: A Pydantic AI ``Agent`` of your own, whose output is text. + Defaults to one with ``system`` as its instructions. + model: ``"/"`` of any of Pydantic AI's providers, + or a Pydantic AI model. Defaults to ``agent``'s own, else + ``$MODEL``, else ``"openai/gpt-6-sol"``. + system: The instructions, when there is no ``agent``. + tools: ``restate.agents`` tools, Pydantic AI tools or tool functions, + beside ``agent``'s own. + + Returns: + The agent's answer. + + Raises: + TypeError: If both ``agent`` and ``system`` are given. + """ + _check(agent, system) + + async def read() -> Transcript: + return await _transcript(ctx) + + async def work() -> str: + return await _answer_once(ctx, text, agent, model, system, tools) + + return await transcribing(ctx, read, work) + + +def agent( + name: str, + pydantic_agent: Agent[Any, str] | None = None, + /, + *, + model: Configure[AnyModel | None] = None, + system: Configure[str | None] = None, + tools: Configure[Sequence[AnyTool]] = (), + description: str | None = None, + schema: Sequence[Step] = (), + background: bool = False, + metadata: dict[str, str] | None = None, + inactivity_timeout: timedelta | None = None, + abort_timeout: timedelta | None = None, + ingress_private: bool | None = None, + invocation_retry_policy: InvocationRetryPolicy | None = RETRIES, +) -> AgentObject: + """Create a Pydantic AI agent: a virtual object with one conversation per key. + + Its ``chat`` handler answers each message with Pydantic AI; the + conversation is kept in the key's database, a step at a time. Add + handlers of your own with ``@assistant.handler``:: + + assistant = pydantic_ai.agent("assistant", system="Answer briefly.", tools=[get_weather]) + + Or bring an ``Agent`` of your own, with this package's tools beside its own:: + + assistant = pydantic_ai.agent("assistant", Agent("anthropic:claude-sonnet-5-5"), tools=[refund]) + + ``model``, ``system`` and ``tools`` may each be a function of the + invocation's context, sync or async, for values that differ per user:: + + assistant = pydantic_ai.agent("assistant", system=lambda ctx: f"You help {ctx.key}.") + + Args: + name: The object's name, as Restate serves it. + pydantic_agent: A Pydantic AI ``Agent`` of your own, whose output is + text. Defaults to one with ``system`` as its instructions. + model: ``"/"`` of any of Pydantic AI's providers + (each needs its extra, such as ``pydantic-ai-slim[openai]``), or a + Pydantic AI model. Defaults to ``pydantic_agent``'s own, else + ``$MODEL``, else ``"openai/gpt-6-sol"``. + system: The instructions, when there is no ``pydantic_agent``. + tools: ``restate.agents`` tools, Pydantic AI tools or tool functions. + description: The object's description, shown in Restate's UI. + schema: Tables of your own, beside the agent's. + background: Whether it works unwatched (a subagent, say): it then + publishes its progress once a turn ends, not at each step. + metadata: The object's metadata. + inactivity_timeout: Restate's inactivity timeout for its handlers. + abort_timeout: Restate's abort timeout for its handlers. + ingress_private: Whether only other services may call it. + invocation_retry_policy: How Restate retries a failed turn. + + Returns: + The agent, to serve with ``sessions.serve`` or ``sessions.app``. + + Raises: + TypeError: If both ``pydantic_agent`` and ``system`` are given. + """ + _check(pydantic_agent, system) + obj = AgentObject( + name, + description=description, + schema=schema, + background=background, + metadata=metadata, + inactivity_timeout=inactivity_timeout, + abort_timeout=abort_timeout, + ingress_private=ingress_private, + invocation_retry_policy=invocation_retry_policy, + ) + + async def answer(ctx: SessionContext, text: str) -> str: + return await prompt( + ctx, + text, + agent=pydantic_agent, + model=await configured(model, ctx), + system=await configured(system, ctx), + tools=await configured(tools, ctx), + ) + + obj.chat(answer) + return obj diff --git a/tests/agents_imports.py b/tests/agents_imports.py index 78c9bc4a..e568adb2 100644 --- a/tests/agents_imports.py +++ b/tests/agents_imports.py @@ -53,3 +53,20 @@ def test_agents_import_without_langgraph_and_its_integration_names_its_extra(): ran = subprocess.run([sys.executable, "-c", WITHOUT_LANGGRAPH], capture_output=True, text=True, check=False) assert ran.returncode == 0, ran.stderr assert "install the langgraph extra, pip install restate-sdk[langgraph]" in ran.stdout + + +WITHOUT_PYDANTIC_AI = """ +import sys +sys.modules["pydantic_ai"] = None +import restate.agents, restate.agents.internal +try: + import restate.agents.pydantic_ai +except ImportError as error: + print(error) +""" + + +def test_agents_import_without_pydantic_ai_and_its_integration_names_its_extra(): + ran = subprocess.run([sys.executable, "-c", WITHOUT_PYDANTIC_AI], capture_output=True, text=True, check=False) + assert ran.returncode == 0, ran.stderr + assert "install the pydantic_ai extra, pip install restate-sdk[pydantic_ai]" in ran.stdout diff --git a/tests/agents_pydantic_ai.py b/tests/agents_pydantic_ai.py new file mode 100644 index 00000000..ab96cc07 --- /dev/null +++ b/tests/agents_pydantic_ai.py @@ -0,0 +1,329 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""A Pydantic AI agent on a scripted model, against a real Restate server: each point of _agent.py's docstring.""" + +# pylint: disable=missing-function-docstring,missing-class-docstring + +import asyncio +import collections +import typing +import uuid +from collections.abc import Awaitable, Callable +from typing import Any + +import httpx +import pytest + +pytest.importorskip("pydantic_ai") + +from pydantic_ai.messages import ( # noqa: E402 + ModelMessage, + ModelRequest, + ModelResponse, + RetryPromptPart, + TextPart, + ToolCallPart, + ToolReturnPart, + UserPromptPart, +) +from pydantic_ai.models.function import AgentInfo, FunctionModel # noqa: E402 + +import restate # noqa: E402 +from restate import HarnessEnvironment, sessions # noqa: E402 +from restate.agents import CANCELLED, tool # noqa: E402 +from restate.agents import pydantic_ai # noqa: E402 +from restate.agents._api import Messages # noqa: E402 + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +pytestmark = [pytest.mark.anyio] + + +class AuthenticationError(Exception): + """As a model SDK names the error of a missing or wrong API key.""" + + +asked: collections.Counter[str] = collections.Counter() +"""How often the model was asked, by the key word that ends each test's first prompt.""" + + +def _prompts(messages: list[ModelMessage]) -> list[str]: + return [ + str(part.content) + for m in messages + if isinstance(m, ModelRequest) + for part in m.parts + if isinstance(part, UserPromptPart) + ] + + +def reply(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + """The scripted model: it answers the last prompt from the conversation alone, as a retry would ask again.""" + text = _prompts(messages)[-1] + asked[_prompts(messages)[0].split(" ")[-1]] += 1 + last = messages[-1] + results = ( + [p for p in last.parts if isinstance(p, ToolReturnPart | RetryPromptPart)] + if isinstance(last, ModelRequest) + else [] + ) + word, _, rest = text.partition(" ") + if word == "system": + return ModelResponse(parts=[TextPart(f"system: {info.instructions}")]) + if word == "refuse": + raise AuthenticationError("invalid x-api-key") + names = {"weather": "get_weather", "crash": "flaky", "wait": "blocking", "wrong": "get_weather"} + names |= {"unknown": "nonexistent", "forbidden": "forbidden"} + if word in names and not results: + args: dict[str, Any] = {"town": 3} if word == "wrong" else {"city": rest} + return ModelResponse(parts=[ToolCallPart(names[word], args, tool_call_id=f"call-{rest}")]) + if results: + result = results[-1] + said = result.model_response() if isinstance(result, RetryPromptPart) else result.model_response_str() + return ModelResponse(parts=[TextPart(f"tool said: {said}")]) + return ModelResponse(parts=[TextPart(f"echo: {text}")]) + + +@tool +async def get_weather(city: str) -> str: + """The current weather in a city.""" + return f"sunny in {city}" + + +crashes: collections.Counter[str] = collections.Counter() +calls_seen: list[str] = [] + + +@tool +async def flaky(city: str, call_id: str) -> str: + """Fails its first run, as a crash would.""" + calls_seen.append(call_id) + crashes[city] += 1 + if crashes[city] == 1: + raise RuntimeError("the process crashed") + return f"recovered in {city}" + + +gate = asyncio.Event() +blocked: set[str] = set() + + +@tool +async def blocking(city: str) -> str: + """Waits until it is stopped.""" + blocked.add(city) + await gate.wait() + return "never" + + +async def forbidden(city: str) -> str: + """A Pydantic AI tool function of the agent's that fails for good.""" + raise restate.TerminalError(f"no access to {city}") + + +assistant = pydantic_ai.agent( + "pai_assistant", + model=FunctionModel(reply), + system="Answer briefly.", + tools=[get_weather, flaky, blocking, forbidden], +) + + +@assistant.handler +async def twice(ctx: sessions.SessionContext, text: str) -> list[str]: + """Two prompts of one invocation: each its own.""" + model = FunctionModel(reply) + return [ + await pydantic_ai.prompt(ctx, text, model=model), + await pydantic_ai.prompt(ctx, f"again {text}", model=model), + ] + + +async def tools_of(ctx: sessions.SessionContext) -> list[Any]: + """Each user's tools, as an async function of the context gives them.""" + return [get_weather] + + +per_user = pydantic_ai.agent( + "pai_per_user", model=lambda ctx: FunctionModel(reply), system=lambda ctx: f"You help {ctx.key}.", tools=tools_of +) + + +@assistant.handler +async def own(ctx: sessions.SessionContext, text: str) -> str: + """An Agent of your own, with its instructions and model, and this package's tools beside its own.""" + from pydantic_ai import Agent # pylint: disable=import-outside-toplevel + + mine: Agent[None, str] = Agent(FunctionModel(reply), instructions="Mine.") + return await pydantic_ai.prompt(ctx, text, agent=mine, tools=[get_weather]) + + +@pytest.fixture(scope="session") +async def harness(): + async with restate.create_test_harness( + sessions.app(assistant, per_user), restate_image="ghcr.io/restatedev/restate:latest" + ) as harness: + yield harness + + +TEXT = {"content-type": "text/plain"} + + +async def post(harness: HarnessEnvironment, path: str, **kwargs: Any) -> httpx.Response: + async with httpx.AsyncClient(timeout=60) as ingress: + return await ingress.post(f"{harness.ingress_url}/pai_assistant/{path}", **kwargs) + + +async def chat(harness: HarnessEnvironment, key: str, text: str) -> httpx.Response: + return await post(harness, f"{key}/chat", content=text, headers=TEXT) + + +async def messages(harness: HarnessEnvironment, key: str) -> list[tuple[str, str]]: + page = typing.cast(Messages, (await post(harness, f"{key}/messages", json={})).json()) + return [ + (message["role"], " ".join(str(part.get("text", part.get("output", ""))) for part in message["parts"])) + for message in page["messages"] + ] + + +async def until(check: Callable[[], Awaitable[bool]]) -> None: + for _ in range(200): + if await check(): + return + await asyncio.sleep(0.05) + raise AssertionError("never") + + +def key() -> str: + return uuid.uuid4().hex + + +async def test_a_message_is_answered_with_the_instructions_and_kept_in_the_conversation(harness: HarnessEnvironment): + k = key() + assert (await chat(harness, k, f"hi {k}")).json() == f"echo: hi {k}" + assert (await chat(harness, k, f"system {k}")).json() == "system: Answer briefly." + assert await messages(harness, k) == [ + ("user", f"hi {k}"), + ("assistant", f"echo: hi {k}"), + ("user", f"system {k}"), + ("assistant", "system: Answer briefly."), + ] + + +async def test_a_tool_call_is_made_and_its_result_answered(harness: HarnessEnvironment): + k = key() + assert (await chat(harness, k, f"weather {k}")).json() == f"tool said: sunny in {k}" + assert ("assistant", f"sunny in {k}") in await messages(harness, k) + + +async def test_a_crash_in_a_tool_resumes_it_with_the_same_call_id_without_asking_the_model_again( + harness: HarnessEnvironment, +): + k = key() + assert (await chat(harness, k, f"crash {k}")).json() == f"tool said: recovered in {k}" + assert crashes[k] == 2 + assert calls_seen.count(f"call-{k}") == 2 + # Asked for the tool call once, then for the answer once: the retry resumed at the tools. + assert asked[k] == 2 + + +async def test_arguments_that_do_not_fit_and_an_unknown_tool_are_the_models_to_fix(harness: HarnessEnvironment): + k = key() + wrong = (await chat(harness, k, f"wrong {k}")).json() + assert wrong.startswith("tool said: Error: the arguments do not fit") + unknown = (await chat(harness, key(), f"unknown {k}")).json() + assert "nonexistent" in unknown + + +async def test_a_tool_functions_terminal_error_is_its_answer(harness: HarnessEnvironment): + k = key() + assert (await chat(harness, k, f"forbidden {k}")).json() == f"tool said: Error: no access to {k}" + + +async def test_a_refused_message_fails_with_424_and_the_next_is_answered(harness: HarnessEnvironment): + k = key() + refused = await chat(harness, k, f"refuse {k}") + assert refused.status_code == 424, refused.text + assert (await chat(harness, k, f"hi {k}")).json() == f"echo: hi {k}" + + +async def test_a_stopped_turn_ends_cancelled_and_the_next_message_starts_afresh(harness: HarnessEnvironment): + k = key() + async with httpx.AsyncClient(timeout=60) as ingress: + sent = await ingress.post( + f"{harness.ingress_url}/pai_assistant/{k}/chat/send", content=f"wait {k}", headers=TEXT + ) + invocation = sent.json()["invocationId"] + + async def in_the_tool() -> bool: + return k in blocked + + await until(in_the_tool) + await post(harness, f"{k}/stop", json=None) + ended = await ingress.get(f"{harness.ingress_url}/restate/invocation/{invocation}/attach") + assert ended.status_code == 409, ended.text + assert (await chat(harness, k, f"hi {k}")).json() == f"echo: hi {k}" + thread = await messages(harness, k) + assert ("assistant", CANCELLED) in thread + assert thread[-1] == ("assistant", f"echo: hi {k}") + + +async def test_each_prompt_of_an_invocation_is_its_own(harness: HarnessEnvironment): + k = key() + answered = await post(harness, f"{k}/twice", json=f"hi {k}") + assert answered.json() == [f"echo: hi {k}", f"echo: again hi {k}"] + + +def test_a_model_whose_provider_cannot_run_is_refused_with_what_to_install(monkeypatch: pytest.MonkeyPatch): + from restate.agents.pydantic_ai import _agent # pylint: disable=import-outside-toplevel + + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with pytest.raises(restate.TerminalError, match="OPENAI_API_KEY") as no_key: + _agent._model("openai/gpt-6-sol") # pyright: ignore[reportPrivateUsage] + assert no_key.value.status_code == 424 + with pytest.raises(restate.TerminalError, match="Unknown provider"): + _agent._model("nope/x") # pyright: ignore[reportPrivateUsage] + + +def test_an_agent_of_your_own_takes_no_system(): + from pydantic_ai import Agent # pylint: disable=import-outside-toplevel + + with pytest.raises(TypeError, match="pass no system"): + pydantic_ai.agent("mine", Agent(instructions="mine"), system="theirs") + + +def test_a_providers_http_refusal_is_terminal_and_a_rate_limit_is_retried(): + from pydantic_ai.exceptions import ModelHTTPError # pylint: disable=import-outside-toplevel + + from restate.agents.internal import refusal # pylint: disable=import-outside-toplevel + + unauthorized = refusal(ModelHTTPError(401, "gpt-6-sol", {"error": "invalid_api_key"})) + assert unauthorized is not None and unauthorized.status_code == 424 + assert refusal(ModelHTTPError(429, "gpt-6-sol", {"error": "rate_limited"})) is None + assert refusal(ModelHTTPError(500, "gpt-6-sol", None)) is None + + +async def test_an_agent_of_your_own_runs_with_its_instructions_and_this_packages_tools(harness: HarnessEnvironment): + k = key() + assert (await post(harness, f"{k}/own", json=f"system {k}")).json() == "system: Mine." + assert (await post(harness, f"{k}/own", json=f"weather {k}")).json() == f"tool said: sunny in {k}" + + +async def test_options_may_be_functions_of_the_context(harness: HarnessEnvironment): + k = key() + async with httpx.AsyncClient(timeout=60) as ingress: + url = f"{harness.ingress_url}/pai_per_user/{k}/chat" + assert (await ingress.post(url, content=f"system {k}", headers=TEXT)).json() == f"system: You help {k}." + assert (await ingress.post(url, content=f"weather {k}", headers=TEXT)).json() == f"tool said: sunny in {k}" diff --git a/tests/examples_agents.py b/tests/examples_agents.py index 595dcc55..d248872d 100644 --- a/tests/examples_agents.py +++ b/tests/examples_agents.py @@ -155,6 +155,14 @@ async def test_basic_langgraph_agent(): assert (await ingress.post("assistant/alice/status")).json() == {"answering": None, "waiting": []} +async def test_basic_pydantic_ai_agent(): + """The agent serves, with its companion; a message needs a model's provider, and is left to a real run.""" + pytest.importorskip("pydantic_ai") + async with serving("basic/pydantic_ai_agent.py") as ingress: + assert (await ingress.post("assistant/alice/messages", {})).json()["messages"] == [] + assert (await ingress.post("assistant/alice/status")).json() == {"answering": None, "waiting": []} + + def serves_http2() -> bool: """Whether ``sessions.serve`` can run here: uvicorn 0.53+ with zttp.""" return importlib.util.find_spec("uvicorn") is not None and importlib.util.find_spec("zttp") is not None diff --git a/uv.lock b/uv.lock index 5206f024..04f86b2e 100644 --- a/uv.lock +++ b/uv.lock @@ -3324,6 +3324,7 @@ openai = [ { name = "openai-agents" }, ] pydantic-ai = [ + { name = "pydantic" }, { name = "pydantic-ai-slim" }, ] serde = [ @@ -3367,6 +3368,7 @@ requires-dist = [ { name = "openai-agents", marker = "extra == 'openai'", specifier = ">=0.6.1" }, { name = "opentelemetry-api", marker = "extra == 'tracing'", specifier = ">=1.36.0" }, { name = "pydantic", marker = "extra == 'langgraph'", specifier = ">=2" }, + { name = "pydantic", marker = "extra == 'pydantic-ai'", specifier = ">=2" }, { name = "pydantic", marker = "extra == 'serde'" }, { name = "pydantic-ai-slim", marker = "extra == 'pydantic-ai'", specifier = ">=2.0,<3" }, { name = "pyright", marker = "extra == 'lint'", specifier = ">=1.1.390" }, From 526acf6f44dd8e48384fa00d855750fc90db0224 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 17:44:37 +0100 Subject: [PATCH 32/45] Review fixes: catchable SqlError, early refusals, clear errors From a developer-experience review of restate.sessions and restate.agents: - SqlError is a TerminalError: still terminal, and now catchable around a statement (a UNIQUE violation, say). - Refused when defined: a schema given as one string (it became a step per character), and a sync function as a handler or view. - A second handler of a name says so; an agent's built-in names (status, stop, chat...) say they are built-in and what to name yours instead. - AgentObject.chat returns its answer function, and works as a decorator (@agent.chat, @agent.chat(clear=...)). - Models name their provider with / or : ("anthropic:claude-..."), as the frameworks do. - SessionObject.config and with_store_wrapper are private; objects have a repr. - No references to modules and extras that do not exist (claude.agent, restate-sdk[anthropic], restate_sdk[serve]). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/agents/.DS_Store | Bin 0 -> 6148 bytes python/restate/agents/_common.py | 26 ++++-- python/restate/agents/_object.py | 52 ++++++++++- python/restate/agents/_tools.py | 8 +- python/restate/agents/pydantic_ai/_agent.py | 4 +- python/restate/sessions/_database.py | 25 ++++-- python/restate/sessions/_object.py | 26 ++++-- python/restate/sessions/_schema.py | 11 ++- python/restate/sessions/_serve.py | 2 +- python/restate/sessions/_session.py | 57 +++---------- tests/agents_definitions.py | 90 ++++++++++++++++++++ tests/sessions_object.py | 18 +++- 12 files changed, 235 insertions(+), 84 deletions(-) create mode 100644 python/restate/agents/.DS_Store create mode 100644 tests/agents_definitions.py diff --git a/python/restate/agents/.DS_Store b/python/restate/agents/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..9aca61ea22f85cc5e7db46bcca792f4d5bb07410 GIT binary patch literal 6148 zcmeHK%Sr=55Ukc50=eYqael!+7()Dld_WWf5|C9yh1~Y!cll|lejtX82M=DP8oFzy zr)!q2!}c}+TfObBfhB-B-4P!irsn7F6T7I45$U{Rz$^B+#};>!>hlTbF7S*4cI15L z@835s!{gKNkhhC5X)v#9lZiGDEYZ5|e5*Vp!6dZ&lX|$Hb(=YWT2vvekrQ@pRtbq8!!} z6{Ua_I91>>w@dH;H}oIo|5K87Qa}p)D+O$}-mF)ArRuGdm-AlR=(lvQ`J}sX9TbLW l$HZvI+;}^_ilVG*zUFx^920}ie9(#d8E{=>QsA!@I0A<&73Kf{ literal 0 HcmV?d00001 diff --git a/python/restate/agents/_common.py b/python/restate/agents/_common.py index 6830dbc0..20c3ee72 100644 --- a/python/restate/agents/_common.py +++ b/python/restate/agents/_common.py @@ -36,8 +36,8 @@ A value, or a function of the invocation's context that gives it: the same for every message, or each user's (from the object's database, say):: - claude.agent("assistant", model="anthropic/claude-sonnet-5-5") - claude.agent("assistant", model=lambda ctx: user_model(ctx)) + langgraph.agent("assistant", model="anthropic/claude-sonnet-5-5") + langgraph.agent("assistant", model=lambda ctx: user_model(ctx)) """ @@ -67,7 +67,7 @@ def next_prompt_id(ctx: SessionContext) -> str: """The default model of the frameworks that take any provider's, with OPENAI_API_KEY.""" CLAUDE_MODEL = "anthropic/claude-sonnet-5-5" -"""The default model of Claude agents, with ANTHROPIC_API_KEY.""" +"""An Anthropic model, as a framework that takes only Anthropic's defaults to (with ANTHROPIC_API_KEY).""" def model_of(model: str | None, fallback: str = DEFAULT_MODEL) -> str: @@ -76,9 +76,17 @@ def model_of(model: str | None, fallback: str = DEFAULT_MODEL) -> str: def provider_of(model: str) -> tuple[str, str]: - """``provider/model-id`` as its provider and model id, split at the first ``/``.""" - provider, slash, model_id = model.partition("/") - return (provider, model_id) if slash else ("", model) + """Split ``"provider/model-id"`` (or the frameworks' ``"provider:model-id"``) into its provider and model id. + + It splits at the first ``/``, else at the first ``:``, so model ids that + contain a colon (``"openai/ft:gpt-4o:acme"``) keep it. A name with neither + has no provider: ``("", model)``. + """ + for separator in ("/", ":"): + provider, found, model_id = model.partition(separator) + if found: + return provider, model_id + return "", model REFUSED = 424 @@ -103,14 +111,14 @@ def __init__(self, message: str) -> None: def check_model_name(framework: str, model: str | None) -> None: - """``model``, the explicit ``model`` option of ``framework``'s agent, if it is not ``provider/model-id``.""" - if model is not None and not re.fullmatch(r"[^/\s]+/\S+", model): + """Raise ``ModelError`` if ``model``, given explicitly to ``framework``'s agent, names no provider.""" + if model is not None and not re.fullmatch(r"[^/:\s]+[/:]\S+", model): raise ModelError(f'{framework}: the model "{model}" is not provider/model-id, such as {CLAUDE_MODEL}') def model_id_of(model: str, provider: str) -> str: """The model id of ``model``, for a framework with only ``provider``'s models: ``provider/model-id``, or the id alone.""" - named, model_id = provider_of(model) if "/" in model else (provider, model) + named, model_id = provider_of(model) if "/" in model or ":" in model else (provider, model) if named != provider or not model_id: from_env = " (from MODEL)" if os.environ.get("MODEL") == model else "" raise ModelError(f'the model "{model}"{from_env} is not {provider}\'s: this agent takes {provider}/') diff --git a/python/restate/agents/_object.py b/python/restate/agents/_object.py index efadd8cd..540373dd 100644 --- a/python/restate/agents/_object.py +++ b/python/restate/agents/_object.py @@ -41,6 +41,8 @@ from restate.sessions.internal import reset T = typing.TypeVar("T") +A = typing.TypeVar("A", bound=Callable[[SessionContext, str], Awaitable[Any]]) +"""A chat turn's answer function: a message in, the answer out.""" RETRIES = InvocationRetryPolicy( initial_interval=timedelta(milliseconds=50), @@ -89,6 +91,12 @@ def _start_after(page: Messages) -> _log.Start: return _log.Start(seq=page["cursor"], id=latest["id"] + 1 if latest else page.get("first", 0)) +BUILT_IN = frozenset( + {"signal", "pending", "status", "stop", "notifications", "watch", "chat", "messages", "clear", "sync", "erase"} +) +"""The names of the handlers and views every agent has (``chat`` and the rest once it chats).""" + + class AgentObject(SessionObject): """ A Restate virtual object with a SQL database per key, ``ctx.db``, set up @@ -141,6 +149,14 @@ def __init__( _signals.with_companion(self) self._built_in() + def _taken(self, name: str) -> str: + if name in BUILT_IN: + return ( + f"{name!r} is a built-in handler of every agent ({', '.join(sorted(BUILT_IN))}): " + f"give yours another name, such as @{self.name}.handler(name='my_{name}')" + ) + return super()._taken(name) + def _attempt_fn(self, fn: Callable[..., Awaitable[Any]], view: bool) -> Callable[..., Awaitable[Any]]: """A handler's attempt: a model provider's refusal it fails with ends its invocation, as ``chat``'s does.""" if view: @@ -194,12 +210,30 @@ async def watch(ctx: ViewContext, request: Watch) -> Notifications: self.view(fn) self.view(inactivity_timeout=WAITS)(watch) + @typing.overload def chat( self, - answer: Callable[[SessionContext, str], Awaitable[T]], + answer: A, + /, *, clear: Callable[[SessionContext], Awaitable[None]] | None = None, - ) -> None: + ) -> A: ... + + @typing.overload + def chat( + self, + /, + *, + clear: Callable[[SessionContext], Awaitable[None]] | None = None, + ) -> Callable[[A], A]: ... + + def chat( + self, + answer: Callable[[SessionContext, str], Awaitable[Any]] | None = None, + /, + *, + clear: Callable[[SessionContext], Awaitable[None]] | None = None, + ) -> Any: """ Make this object an agent, which answers each message with ``answer``: a framework's turn, which keeps the conversation in the log @@ -213,13 +247,23 @@ def chat( A key answers one message at a time; one its model's provider refused (a missing API key, an unknown model) answers with an error (424). """ + if answer is None: + + def decorate( + fn: Callable[[SessionContext, str], Awaitable[Any]], + ) -> Callable[[SessionContext, str], Awaitable[Any]]: + return self.chat(fn, clear=clear) + + return decorate AGENT.on(self, AgentInfo(background=self.background, clear=clear)) + answering = answer - async def chat(ctx: SessionContext, text: str) -> T: - return await chat_turn(ctx, text, answer) + async def chat(ctx: SessionContext, text: str) -> Any: + return await chat_turn(ctx, text, answering) self.handler(name="chat", accept="*/*", input_serde=MessageSerde(), explicit_cancellation=True)(chat) _followed(self) + return answer def _companion( diff --git a/python/restate/agents/_tools.py b/python/restate/agents/_tools.py index 1b202c89..36099eb9 100644 --- a/python/restate/agents/_tools.py +++ b/python/restate/agents/_tools.py @@ -53,8 +53,8 @@ def _pydantic() -> Any: import pydantic except ImportError as error: raise ImportError( - "restate.agents' tools take their inputs' schemas from pydantic: install it, or an agent framework's " - "extra, which brings it (such as restate-sdk[anthropic])" + "restate.agents' tools take their inputs' schemas from pydantic: install pydantic, or a framework's " + "extra that includes it, such as restate-sdk[pydantic_ai] or restate-sdk[langgraph]" ) from error return pydantic @@ -445,8 +445,8 @@ def subagent( continues the same task. The model reads what the agent is from its description:: - coder = claude.agent("coder", description="Builds and tests code in Python") - lead = claude.agent("lead", tools=[subagent(coder)]) + coder = langgraph.agent("coder", description="Builds and tests code in Python") + lead = langgraph.agent("lead", tools=[subagent(coder)]) """ agent_name = typing.cast(str, getattr(agent, "name")) # noqa: B009 (an SDK's definition) tag = getattr(agent, "service_tag", None) diff --git a/python/restate/agents/pydantic_ai/_agent.py b/python/restate/agents/pydantic_ai/_agent.py index dc2272e6..9a0d645c 100644 --- a/python/restate/agents/pydantic_ai/_agent.py +++ b/python/restate/agents/pydantic_ai/_agent.py @@ -64,6 +64,7 @@ from_tool, model_of, next_prompt_id, + provider_of, refusal_of, stored_answer, ) @@ -228,7 +229,8 @@ def _model(model: str | Model | None) -> Model: return model named = model_of(model) try: - return infer_model(named if ":" in named else named.replace("/", ":", 1)) + provider, model_id = provider_of(named) + return infer_model(f"{provider}:{model_id}" if provider else named) except (ImportError, UserError, ValueError) as error: raise ModelError(f'the model "{named}": {error}') from error diff --git a/python/restate/sessions/_database.py b/python/restate/sessions/_database.py index cbb9d8a0..924ca26e 100644 --- a/python/restate/sessions/_database.py +++ b/python/restate/sessions/_database.py @@ -44,6 +44,7 @@ SqlReply, sql_open, ) +from restate.exceptions import TerminalError SqlValue = None | int | float | str | bytes """An SQL value: ``None``, ``int`` (64 bits), ``float``, ``str``, or ``bytes`` (BLOB).""" @@ -58,14 +59,22 @@ """Rows as lists, in column order.""" -class SqlError(Exception): - """ - A failure a retry would repeat: SQL the engine rejected (a syntax error, a - constraint, a missing table), which undid just that statement and leaves - the session usable; a parameter it cannot bind; or a call the session - refused, which may have ended it (a ``COMMIT`` in a transaction's body, a - transaction left open). A storage error, which ends the session, raises - as it is. +class SqlError(TerminalError): + """A SQL statement failed, and retrying it would fail the same way. + + Raised for SQL the engine rejects (a syntax error, a constraint violation, + a missing table), which undoes just that statement; for a parameter it + cannot bind; and for a call the database refuses (a ``COMMIT`` inside a + transaction's body, say). It is a ``TerminalError``: uncaught, it fails + the invocation instead of being retried. Catch it to handle, say, a + ``UNIQUE`` violation:: + + try: + await ctx.db.execute("INSERT INTO users (email) VALUES (?)", [email]) + except sessions.SqlError: + return "taken" + + Storage errors are not ``SqlError``: they are retried. """ diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index 226bf1e9..2b9c82a8 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -29,6 +29,7 @@ async def list(ctx: sessions.ViewContext) -> list[str]: import asyncio import functools +import inspect import typing import weakref from collections.abc import Awaitable, Callable, Sequence @@ -46,9 +47,8 @@ async def list(ctx: sessions.ViewContext) -> list[str]: from restate.server_context import DisconnectedException, ServerInvocationContext from restate.sessions._api import Cancel, Database, DatabaseReader, Rpc, Send, SessionContext, ViewContext from restate.sessions._calls import Caller, Sender, send_now -from restate.sessions._database import Store from restate.sessions._link import Link -from restate.sessions._schema import Schema, Step +from restate.sessions._schema import Schema, Step, steps_of from restate.sessions._scope import SCOPE, Scope from restate.sessions._selector import FutureSelector, hook from restate.sessions._session import Attempt, ObjectConfig, SqlHost, reader_of, with_attempt @@ -367,9 +367,16 @@ def __init__( ingress_private=ingress_private, invocation_retry_policy=invocation_retry_policy, ) - self.config: ObjectConfig = ObjectConfig(schema=Schema(name, schema) if schema else None) + self._config: ObjectConfig = ObjectConfig(schema=Schema(name, schema) if steps_of(schema) else None) ENTRIES.on(self, {}) + def __repr__(self) -> str: + return f"{type(self).__name__}({self.name!r})" + + def _taken(self, name: str) -> str: + """Why ``name`` is taken already, for the error of a second handler of that name.""" + return f"{self.name} has a handler or view named {name!r} already: give this one another name" + @typing.overload # type: ignore[override] def handler(self, fn: F, /) -> F: ... @@ -535,12 +542,18 @@ def register(fn: F) -> F: "in a restate.VirtualObject, Service or Workflow, and call it with ctx.rpc." ) raise TypeError(f"{what} {instead}") + if not inspect.iscoroutinefunction(inspect.unwrap(fn)): + kind = "view" if view else "handler" + raise TypeError(f"{self.name}.{fn.__name__}: a {kind} must be a coroutine function (async def)") + chosen = name or fn.__name__ + if chosen in self.handlers: + raise ValueError(self._taken(chosen)) explicit = bool(options.get("explicit_cancellation")) sdk_options = {key: value for key, value in options.items() if key != "explicit_cancellation"} entry = _Entry( fn=self._attempt_fn(fn, view), view=view, - config=self.config, + config=self._config, definition=weakref.ref(self), explicit_cancellation=explicit, ) @@ -563,10 +576,5 @@ def register(fn: F) -> F: return register - def with_store_wrapper(self, wrap: Callable[[Store], Store]) -> SessionObject: - """Wrap the stores of this object's attempts with ``wrap``: fault injection, in tests.""" - self.config.wrap_store = wrap - return self - __all__ = ["SessionObject", "cancel_calls", "cancelled", "definition_of"] diff --git a/python/restate/sessions/_schema.py b/python/restate/sessions/_schema.py index d7790c30..1103a2ec 100644 --- a/python/restate/sessions/_schema.py +++ b/python/restate/sessions/_schema.py @@ -55,12 +55,21 @@ """The attribute of a database of this package's: how it applies schemas (see ``Schema.ensure``).""" +def steps_of(steps: Sequence[Step]) -> tuple[Step, ...]: + """``steps`` as a tuple; a lone string (a sequence of characters) is a mistake, not a list of statements.""" + if isinstance(steps, str): + raise TypeError( + f"schema takes a list of steps, not a string: schema=[{steps[:40]!r}{'...' if len(steps) > 40 else ''}]" + ) + return tuple(steps) + + class Schema: """A part of a database's schema: its ``name``, and the ``steps`` that make it, in order.""" def __init__(self, name: str, steps: Sequence[Step]) -> None: self.name: str = name - self.steps: tuple[Step, ...] = tuple(steps) + self.steps: tuple[Step, ...] = steps_of(steps) async def ensure(self, db: Any) -> None: """ diff --git a/python/restate/sessions/_serve.py b/python/restate/sessions/_serve.py index 7890322c..80288094 100644 --- a/python/restate/sessions/_serve.py +++ b/python/restate/sessions/_serve.py @@ -136,7 +136,7 @@ def serve( _NO_SERVER = ( - "sessions.serve runs on uvicorn's HTTP/2: install restate_sdk[serve] (uvicorn 0.53 or later, and zttp; " + "sessions.serve runs on uvicorn's HTTP/2: install restate-sdk[serve] (uvicorn 0.53 or later, and zttp; " "Python 3.11 or later), or serve sessions.app(...) with an ASGI server of your own" ) diff --git a/python/restate/sessions/_session.py b/python/restate/sessions/_session.py index a0a78152..afce043b 100644 --- a/python/restate/sessions/_session.py +++ b/python/restate/sessions/_session.py @@ -42,7 +42,6 @@ ConnectionTransaction, Row, Rows, - SqlError, SqlParam, Store, objects_of, @@ -127,29 +126,6 @@ class ObjectConfig: # --- Errors --------------------------------------------------------------------------------------- -def handler_error(error: BaseException) -> BaseException: - """ - What a failed call to the session fails the handler with: an ``SqlError``, - which a retry would repeat, becomes a ``TerminalError``, its cause the SQL - error. Anything else is as for any handler: a storage error, which ends the - session, or the handler's own (a failed call to another service, say) is - retried, unless it is a ``TerminalError`` itself. - """ - if not isinstance(error, SqlError): - return error - terminal = TerminalError(str(error)) - terminal.__cause__ = error - return terminal - - -def _raise_for_handler(error: BaseException) -> typing.NoReturn: - """Raise what ``error`` fails the handler with (see ``handler_error``).""" - failure = handler_error(error) - if failure is error: - raise error - raise failure - - # --- The object's state as a store ---------------------------------------------------------------- @@ -374,10 +350,7 @@ async def query(self, sql: str, params: Sequence[SqlParam] | None = None) -> lis them, or use ``values``. """ columns, rows = await self.query_columns(sql, params) - try: - return objects_of(columns, rows) - except SqlError as error: - _raise_for_handler(error) + return objects_of(columns, rows) async def get(self, sql: str, params: Sequence[SqlParam] | None = None) -> Row | None: """The first row of ``query``, or ``None`` when there is none.""" @@ -405,10 +378,7 @@ def __init__(self, handle: ConnectionTransaction, send: Send) -> None: self.send: Send = send async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: - try: - return await self._handle.query_columns(sql, params) - except SqlError as error: - _raise_for_handler(error) + return await self._handle.query_columns(sql, params) class AttemptDatabase(Writes): @@ -478,10 +448,7 @@ async def reset(tx: Transaction) -> None: empty = _RESET.of(tx) if empty is None: raise TerminalError("reset takes a transaction of an exclusive handler's ctx.db") - try: - await empty() - except SqlError as error: - _raise_for_handler(error) + await empty() # --- An attempt ----------------------------------------------------------------------------------- @@ -602,8 +569,7 @@ def open_now(self) -> None: async def end(self, error: BaseException | None) -> None: """ The handler ended, with ``error`` or not: the key's other databases end, - then this one. It raises what the handler fails with (see - ``handler_error``), or a failure to end the session. + then this one. It raises what the handler failed with, or a failure to end the session. """ failures: list[BaseException] = [] for other in self._databases.values(): @@ -624,10 +590,10 @@ async def call(self, call: Callable[[Connection], Awaitable[T]]) -> T: session = self._current() try: return await call(await session.ready) - except BaseException as error: + except BaseException: if session.ended: session.retired = True - _raise_for_handler(error) + raise @asynccontextmanager async def transaction(self) -> AsyncGenerator[AttemptTransaction, None]: @@ -658,10 +624,10 @@ async def transaction(self) -> AsyncGenerator[AttemptTransaction, None]: outgoing = sends.close() if sends is not None else None if self._has_outbox: await self._commit_outbox(handle, outgoing) - except BaseException as error: + except BaseException: if session.ended: session.retired = True - _raise_for_handler(error) + raise if outgoing is not None: await self.call(self._flush) @@ -726,8 +692,7 @@ async def _start(self) -> Connection: async def _end_session(self, error: BaseException | None) -> None: """ - End the session, then raise as the handler did: an SQL error that - escaped it becomes terminal. The handler's error stands: what it + End the session, then raise as the handler did. The handler's error stands: what it committed is stored, and an outbox left full is made by the next open. A failure to end the session fails a handler that succeeded, which a retry then ends (a transaction it left open, say). @@ -739,9 +704,9 @@ async def _end_session(self, error: BaseException | None) -> None: except BaseException as failure: # pylint: disable=broad-exception-caught failed = failure if error is not None: - _raise_for_handler(error) + raise error if failed is not None: - _raise_for_handler(failed) + raise failed async def _close(self) -> None: """ diff --git a/tests/agents_definitions.py b/tests/agents_definitions.py new file mode 100644 index 00000000..b7b45114 --- /dev/null +++ b/tests/agents_definitions.py @@ -0,0 +1,90 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""Defining session objects and agents: what is refused at definition time, and how, before anything runs.""" + +# pylint: disable=missing-function-docstring + +import pytest + +from restate import sessions +from restate.agents import AgentObject +from restate.agents.internal import provider_of + + +def test_a_schema_given_as_one_string_is_refused(): + with pytest.raises(TypeError, match="a list of steps, not a string"): + sessions.SessionObject("p", schema="CREATE TABLE x (a)") # type: ignore[arg-type] + + +def test_a_sync_function_is_refused_as_a_handler_or_view(): + obj = sessions.SessionObject("sync") + + def handler(ctx: sessions.SessionContext) -> None: ... + + with pytest.raises(TypeError, match="must be a coroutine function"): + obj.handler(handler) # type: ignore[type-var] + with pytest.raises(TypeError, match="must be a coroutine function"): + obj.view(handler) # type: ignore[type-var] + + +def test_a_second_handler_of_a_name_is_refused_with_the_name(): + obj = sessions.SessionObject("twice") + + @obj.handler + async def add(ctx: sessions.SessionContext) -> None: ... + + async def listing(ctx: sessions.ViewContext) -> None: ... + + with pytest.raises(ValueError, match="'add' already"): + obj.view(name="add")(listing) + + +def test_an_agents_built_in_names_are_refused_with_what_to_do(): + agent = AgentObject("helper") + + async def status(ctx: sessions.SessionContext) -> str: + return "mine" + + with pytest.raises(ValueError, match=r"'status' is a built-in handler of every agent.*name='my_status'"): + agent.handler(status) + + +def test_chat_is_a_decorator_too_and_returns_the_function(): + agent = AgentObject("decorated") + + @agent.chat + async def answer(ctx: sessions.SessionContext, text: str) -> str: + return text + + assert answer.__name__ == "answer" + assert "chat" in agent.handlers and "messages" in agent.handlers + + other = AgentObject("decorated_with_clear") + + async def forget(ctx: sessions.SessionContext) -> None: ... + + @other.chat(clear=forget) + async def reply(ctx: sessions.SessionContext, text: str) -> str: + return text + + assert reply.__name__ == "reply" and "chat" in other.handlers + + +def test_objects_have_a_readable_repr(): + assert repr(sessions.SessionObject("notes")) == "SessionObject('notes')" + assert repr(AgentObject("assistant")) == "AgentObject('assistant')" + + +def test_a_model_names_its_provider_with_a_slash_or_a_colon(): + assert provider_of("openai/gpt-6-sol") == ("openai", "gpt-6-sol") + assert provider_of("anthropic:claude-sonnet-5-5") == ("anthropic", "claude-sonnet-5-5") + assert provider_of("openai/ft:gpt-4o:acme") == ("openai", "ft:gpt-4o:acme") + assert provider_of("gpt") == ("", "gpt") diff --git a/tests/sessions_object.py b/tests/sessions_object.py index 81e038ec..872672ba 100644 --- a/tests/sessions_object.py +++ b/tests/sessions_object.py @@ -289,7 +289,23 @@ async def bad(ctx: sessions.SessionContext) -> None: with pytest.raises(TerminalError, match="no such table") as error: await run(broken, "bad", host) - assert isinstance(error.value.__cause__, sessions.SqlError) + assert isinstance(error.value, sessions.SqlError) + + +async def test_a_handler_catches_an_sql_error_and_goes_on(): + host = Host(Backend()) + users = sessions.SessionObject("users", schema=["CREATE TABLE users (email TEXT PRIMARY KEY)"]) + + @users.handler + async def register(ctx: sessions.SessionContext, email: str) -> str: + try: + await ctx.db.execute("INSERT INTO users (email) VALUES (?)", [email]) + except sessions.SqlError: + return "taken" + return "registered" + + assert await run(users, "register", host, "a@example.com") == "registered" + assert await run(users, "register", Host(host.backend), "a@example.com") == "taken" async def test_a_call_to_ctx_db_inside_its_transaction_is_refused(): From f6784d791eacb7d6e7dcdf21519850c6687f244e Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 18:09:13 +0100 Subject: [PATCH 33/45] restate.sessions: flat calls and sends, ctx.call(handler, arg, key=...) Calls and sends take their input as an argument instead of returning a function of it: await ctx.call(add, 5, key="a"), ctx.send(remind, note, key=k, delay=...), tx.send(deposit, money, key=to). A send can no longer be forgotten by leaving off a trailing call. Overloads keep the typing: the result is the handler's, a wrong input is flagged, and a handler without input takes none. call, generic_call, send, generic_send and cancel are methods of the contexts (and send, generic_send of a transaction), declared by the Caller and Sender protocols; Rpc, Send and Cancel are gone. _api.py's docstrings are PEP 257 / Google style. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- examples-agents/tutorial/07_calls.py | 4 +- examples-agents/tutorial/08_sends.py | 2 +- python/restate/agents/_object.py | 12 +- python/restate/agents/_signals.py | 2 +- python/restate/agents/_tools.py | 42 ++-- python/restate/agents/_turn.py | 6 +- python/restate/sessions/__init__.py | 41 ++-- python/restate/sessions/_api.py | 309 ++++++++++++++------------- python/restate/sessions/_calls.py | 123 +++++------ python/restate/sessions/_object.py | 44 ++-- python/restate/sessions/_session.py | 15 +- tests/sessions_harness.py | 8 +- tests/sessions_object.py | 39 ++-- 13 files changed, 327 insertions(+), 320 deletions(-) diff --git a/examples-agents/tutorial/07_calls.py b/examples-agents/tutorial/07_calls.py index 5cc8520f..f1a1d172 100644 --- a/examples-agents/tutorial/07_calls.py +++ b/examples-agents/tutorial/07_calls.py @@ -6,7 +6,7 @@ is a key with its own database, which records the wallet that owns it. The first wallet to claim a handle gets it. -``ctx.rpc(claim_handle, key=name)(wallet)`` calls that object through Restate, +``ctx.call(claim_handle, wallet, key=name)`` calls that object through Restate, and waits for the answer. A retried handler (06) makes its calls again, so a call that changes something gets an idempotency key that is the same on every attempt. Restate then answers a repeated call with the first call's result, @@ -73,7 +73,7 @@ async def deposit(ctx: sessions.SessionContext, money: Deposit) -> int: @wallet.handler async def claim(ctx: sessions.SessionContext, name: str) -> str: - owner = await ctx.rpc(claim_handle, key=name, idempotency_key=f"{ctx.request.id}:claim")(ctx.key) + owner = await ctx.call(claim_handle, ctx.key, key=name, idempotency_key=f"{ctx.request.id}:claim") if owner != ctx.key: raise restate.TerminalError(f"@{name} belongs to {owner}") return f"{ctx.key} is now @{name}" diff --git a/examples-agents/tutorial/08_sends.py b/examples-agents/tutorial/08_sends.py index addaa38f..96cb8ac6 100644 --- a/examples-agents/tutorial/08_sends.py +++ b/examples-agents/tutorial/08_sends.py @@ -72,7 +72,7 @@ async def transfer(ctx: sessions.SessionContext, money: Transfer) -> int: if await balance_of(tx) < 0: raise restate.TerminalError("Not enough money") # Delivered once this transaction commits, and never if it rolls back. - tx.send(deposit, key=money.to)(Deposit(id=transfer_id, amount=money.amount, note=f"from {ctx.key}")) + tx.send(deposit, Deposit(id=transfer_id, amount=money.amount, note=f"from {ctx.key}"), key=money.to) return await balance_of(tx) diff --git a/python/restate/agents/_object.py b/python/restate/agents/_object.py index 540373dd..3be253b2 100644 --- a/python/restate/agents/_object.py +++ b/python/restate/agents/_object.py @@ -275,7 +275,13 @@ def _companion( key: asked again, it answers again. """ key = None if step is None else f"{ctx.request.id}:{step}" - return typing.cast(Callable[..., Awaitable[T]], ctx.rpc(handler, key=_signals.signals_of(ctx), idempotency_key=key)) + companion = _signals.signals_of(ctx) + untyped = typing.cast(Any, ctx) + + def call(*args: Any) -> Awaitable[T]: + return typing.cast(Awaitable[T], untyped.call(handler, *args, key=companion, idempotency_key=key)) + + return call async def _stopped(ctx: ViewContext, step: str) -> Stopped: @@ -315,7 +321,7 @@ async def erase(ctx: SessionContext) -> None: await reset(tx) await _log.restart(tx, start) await cleared(ctx) - await ctx.rpc(_signals.reset, key=_signals.signals_of(ctx), idempotency_key=f"{ctx.request.id}:reset")() + await ctx.call(_signals.reset, key=_signals.signals_of(ctx), idempotency_key=f"{ctx.request.id}:reset") async def clear(ctx: ViewContext) -> Stopped: """ @@ -326,7 +332,7 @@ async def clear(ctx: ViewContext) -> Stopped: answered first, then forgotten too. """ stopped = await _stopped(ctx, "clear") - await ctx.rpc(erase, key=ctx.key, idempotency_key=f"{ctx.request.id}:erase")() + await ctx.call(erase, key=ctx.key, idempotency_key=f"{ctx.request.id}:erase") return stopped async def sync(ctx: ViewContext, request: SyncRequest | None = None) -> Synced: diff --git a/python/restate/agents/_signals.py b/python/restate/agents/_signals.py index 27ebe837..0fc3014b 100644 --- a/python/restate/agents/_signals.py +++ b/python/restate/agents/_signals.py @@ -386,7 +386,7 @@ def deliver(ctx: SessionContext, agent: object, key: str, text: str, *, idempote deliver(ctx, triage, "ENG-1", "The issue was closed", idempotency_key=delivery) """ companion = companion_of(agent, key) # an agent of restate.agents, with a `chat` handler - ctx.send.generic(companion.agent, "chat", text, key=key, idempotency_key=idempotency_key) + ctx.generic_send(companion.agent, "chat", text, key=key, idempotency_key=idempotency_key) __all__: Sequence[str] = ["Companion", "companion_of", "deliver", "signals"] diff --git a/python/restate/agents/_tools.py b/python/restate/agents/_tools.py index 36099eb9..ffa30a56 100644 --- a/python/restate/agents/_tools.py +++ b/python/restate/agents/_tools.py @@ -226,14 +226,14 @@ async def _once(name: str, input: object, call: ToolCall, body: Callable[[ToolCa return result -def _untyped(caller: object) -> Callable[..., Awaitable[Any]]: - """A caller of a handler whose input is known only at run time (a tool's): it takes the input, if any.""" - return typing.cast(Callable[..., Awaitable[Any]], caller) +def _untyped(ctx: SessionContext) -> Any: + """``ctx``, to call and send to a handler whose input a tool knows only at run time.""" + return ctx -def _untyped_send(sender: object) -> Callable[..., None]: - """A sender to a handler whose input is known only at run time (a tool's): it takes the input, if any.""" - return typing.cast(Callable[..., None], sender) +def _inputs(handler: Handler[Any, Any], value: Any) -> tuple[Any, ...]: + """The positional input of a call of ``handler``: ``value``, if it takes an input.""" + return (value,) if handler.arity >= 2 else () def _handler(target: Callable[..., Any]) -> Handler[Any, Any]: @@ -328,10 +328,10 @@ async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: async def made(c: ToolCall) -> Any: target_key = chosen if key is None else await configured(key, c.ctx) - caller = _untyped( - c.ctx.rpc(target, key=target_key if _keyed(handler) else None, idempotency_key=c.idempotency_key) + target_key = target_key if _keyed(handler) else None + return await _untyped(c.ctx).call( + target, *_inputs(handler, value), key=target_key, idempotency_key=c.idempotency_key ) - return await (caller(value) if handler.arity >= 2 else caller()) return await _once(f"{service}/{handler.name}", arguments["raw"], tool_call, made) @@ -363,15 +363,13 @@ async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: async def made(c: ToolCall) -> str: target_key = chosen if key is None else await configured(key, c.ctx) - sender = _untyped_send( - c.ctx.send( - target, key=target_key if _keyed(handler) else None, delay=delay, idempotency_key=c.idempotency_key - ) + _untyped(c.ctx).send( + target, + *_inputs(handler, value), + key=target_key if _keyed(handler) else None, + delay=delay, + idempotency_key=c.idempotency_key, ) - if handler.arity >= 2: - sender(value) - else: - sender() return f"sent to {service}{'' if target_key is None else f'/{target_key}'}/{handler.name}" return await _once(f"{service}/{handler.name}", arguments["raw"], tool_call, made) @@ -407,8 +405,9 @@ async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: async def made(c: ToolCall) -> Any: workflow_id = f"{c.ctx.key}-{c.call_id}" if key is None else await configured(key, c.ctx) - caller = _untyped(c.ctx.rpc(target, key=workflow_id, idempotency_key=c.idempotency_key)) - return await (caller(value) if handler.arity >= 2 else caller()) + return await _untyped(c.ctx).call( + target, *_inputs(handler, value), key=workflow_id, idempotency_key=c.idempotency_key + ) return await _once(service, arguments["raw"], tool_call, made) @@ -461,7 +460,7 @@ async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: task_name, task = typing.cast(str, arguments["name"]), typing.cast(str, arguments["task"]) async def made(c: ToolCall) -> Any: - return await c.ctx.rpc.generic( + return await c.ctx.generic_call( agent_name, handler, task, key=f"{c.ctx.key}.{task_name}", idempotency_key=c.idempotency_key ) @@ -505,8 +504,7 @@ async def made(c: ToolCall) -> Any: request: _signals.WaitRequest = {"id": wait_id, "name": name} if input is not None: request["input"] = arguments["raw"] - waiting = c.ctx.rpc(_signals.wait, key=_signals.signals_of(c.ctx), idempotency_key=wait_id) - return await waiting(request) + return await c.ctx.call(_signals.wait, request, key=_signals.signals_of(c.ctx), idempotency_key=wait_id) return await _once(f"signal/{name}", arguments["raw"], tool_call, made) diff --git a/python/restate/agents/_turn.py b/python/restate/agents/_turn.py index efa483ae..62f2decd 100644 --- a/python/restate/agents/_turn.py +++ b/python/restate/agents/_turn.py @@ -181,7 +181,7 @@ async def publish(ctx: SessionContext) -> None: read = _READING.get(ctx) if read is not None: await _logged(ctx, read) - await ctx.rpc(_signals.publish, key=_signals.signals_of(ctx))() + await ctx.call(_signals.publish, key=_signals.signals_of(ctx)) async def stepped(ctx: SessionContext) -> None: @@ -255,7 +255,7 @@ async def chat_turn(ctx: SessionContext, text: str, answer: Callable[[SessionCon companion = _signals.signals_of(ctx) turn: Turn = {"invocationId": invocation_id, "since": _now()} # The companion knows the turn from here to its close: `status` and `stop` ask it. - await ctx.rpc(_signals.start, key=companion)(turn) + await ctx.call(_signals.start, turn, key=companion) await stepped(ctx) async def end() -> None: @@ -265,7 +265,7 @@ async def end() -> None: read = _READING.get(ctx) if read is not None: await _logged(ctx, read) - await ctx.rpc(_signals.close, key=companion)(invocation_id) + await ctx.call(_signals.close, invocation_id, key=companion) stop = Stop() _STOP.set(ctx, stop) diff --git a/python/restate/sessions/__init__.py b/python/restate/sessions/__init__.py index f9b8d77b..fffb26fc 100644 --- a/python/restate/sessions/__init__.py +++ b/python/restate/sessions/__init__.py @@ -8,11 +8,11 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -Session objects: virtual objects with a SQL database per key, kept in the -object's state (``ctx.db``). The Python variant of ``@restatedev/sessions``:: +"""Session objects: Restate virtual objects with a SQL database per key. + +Each key's database (``ctx.db``) is kept in the object's state in Restate, +so there is no database to run:: - import restate from restate import sessions notes = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) @@ -23,33 +23,31 @@ async def add(ctx: sessions.SessionContext, text: str) -> None: await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [text]) - @notes.view - async def list(ctx: sessions.ViewContext) -> list[str]: + @notes.view(name="list") + async def list_notes(ctx: sessions.ViewContext) -> list[str]: return [row["text"] for row in await ctx.db.query("SELECT text FROM notes ORDER BY rowid")] - app = sessions.app(notes) # with what it needs beside it; hypercorn module:app + if __name__ == "__main__": + sessions.serve(notes) -Session handlers don't replay a journal: a retried attempt runs again from its -first line, against what the attempts before it committed. +Handlers run one at a time per key; views run beside them and read what was +committed. A handler is not replayed from a journal: a retried attempt runs +again from the start, against what earlier attempts committed. The types a handler gets (``SessionContext``, ``Database``, ``Transaction``, -``Rpc``, ``Send``, ...) are protocols: a decorated handler is still the plain -function, so a test can call it with a fake context of its own. What libraries -built on session objects, and tests, need beside this is in -``restate.sessions.internal``. A session object -serves beside the SDK's services, objects and workflows, which it calls with -``ctx.rpc`` for journaled work. +...) are protocols, and a decorated handler is still the plain function, so a +test can call it with a fake context. Session objects are served beside the +SDK's services, objects and workflows, and call them with ``ctx.call``. +``restate.sessions.internal`` has what libraries built on session objects need. """ -# The API: protocols (_api.py), which handlers and views code against. from restate.sessions._api import ( - Cancel, + Caller, Database, DatabaseReader, DatabaseWriter, - Rpc, - Send, + Sender, SessionContext, Transaction, ViewContext, @@ -62,14 +60,13 @@ async def list(ctx: sessions.ViewContext) -> list[str]: __all__ = [ "CallFuture", - "Cancel", + "Caller", "Database", "DatabaseReader", "DatabaseWriter", "Row", "Rows", - "Rpc", - "Send", + "Sender", "SessionContext", "SessionObject", "SqlError", diff --git a/python/restate/sessions/_api.py b/python/restate/sessions/_api.py index 0ed93375..87742120 100644 --- a/python/restate/sessions/_api.py +++ b/python/restate/sessions/_api.py @@ -8,12 +8,13 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -The API of session objects, as handlers and views code against it: the -contexts, the database and its transactions, calls and sends. Protocols, so -that what a handler sees is what is declared here, and a test can give a -handler a fake of its own; the implementations are this package's private -classes (``_object.py``, ``_session.py``, ``_calls.py``). +"""The API that session objects' handlers and views code against. + +The contexts, the database and its transactions, and calls and sends are +declared as protocols: what a handler sees is exactly what is declared here, +and a test can pass a handler a fake of its own. The implementations are +private classes of this package (``_object.py``, ``_session.py``, +``_calls.py``). """ from __future__ import annotations @@ -31,84 +32,45 @@ I = typing.TypeVar("I") O = typing.TypeVar("O") -Cancel = Callable[[str], None] -""" -Cancel invocation ``invocation_id`` (such as ``inv_1f…``), as Restate's admin -API does: made at once, and not awaited. Cancelling an invocation that already -ended does nothing, so a new attempt that cancels it again is harmless. -""" - # --- The database --------------------------------------------------------------------------------- class DatabaseReader(Protocol): - """Reads of a database, or of one transaction of it: a view's ``ctx.db``.""" + """Read access to a database, or to one transaction of it; a view's ``ctx.db``.""" async def query(self, sql: str, params: Sequence[SqlParam] | None = None) -> list[Row]: - """ - The rows of one statement, each a dict by column name; parameters bind - to a single statement. Columns are named as SQL names them (an - expression needs ``AS``), and two of the same name are an error: alias - them, or use ``values``. + """Run one statement and return its rows, each a dict keyed by column name. + + Columns are named as SQL names them, so an expression needs ``AS``; + two columns of the same name raise ``SqlError`` (alias one, or use + ``values``):: + + rows = await ctx.db.query("SELECT text FROM notes WHERE done = ?", [False]) """ ... async def get(self, sql: str, params: Sequence[SqlParam] | None = None) -> Row | None: - """The first row of ``query``, or ``None`` when there is none.""" + """Return the first row of ``query``, or ``None`` if there is none.""" ... async def values(self, sql: str, params: Sequence[SqlParam] | None = None) -> Rows: - """The rows of one statement as lists, in column order: for adapters, and for columns that share a name.""" + """Run one statement and return its rows as lists, in column order.""" ... async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: - """The rows of one statement as lists, with the column names of the statement: ``(columns, rows)``.""" + """Run one statement and return its column names and its rows as lists: ``(columns, rows)``.""" ... class DatabaseWriter(DatabaseReader, Protocol): - """SQL: a database, or one transaction of it.""" + """Read and write access to a database, or to one transaction of it.""" async def execute(self, sql: str, params: Sequence[SqlParam] | None = None) -> None: - """Run SQL (one or more statements), discarding rows; parameters only with a single statement.""" - ... - - -class Transaction(DatabaseWriter, Protocol): - """One transaction of a ``Database``: its SQL, and the sends it commits with.""" - - @property - def send(self) -> Send: - """ - Sends made right after the transaction commits, and never if it rolls - back. Each has the idempotency key it is given, or else one of the - transaction (``transaction//``), so a send an attempt makes - again is dropped. - """ - ... - - -class Database(DatabaseWriter, Protocol): - """ - The session's database, ``ctx.db``, for the whole handler: each statement - outside ``transaction()`` is its own transaction. Calls run one at a time, - in call order. A commit is durable once its call returns. - """ - - def transaction(self) -> typing.AsyncContextManager[Transaction]: - """ - One transaction, as the body of an ``async with``:: + """Run SQL and discard any rows. - async with ctx.db.transaction() as tx: - await tx.execute("UPDATE accounts SET balance = balance - ? WHERE id = ?", [amount, src]) - tx.send(deposit, key=dst)(amount) - - It commits when the body ends, and rolls back when the body raises. It - holds the database: other calls wait until it ends, and a call to the - database from inside the body is refused, as it would wait for the body. - So the body uses the ``tx`` it gets, which is invalid once the body has - ended. + ``sql`` may hold several statements, but parameters bind only to a + single one. """ ... @@ -116,55 +78,65 @@ def transaction(self) -> typing.AsyncContextManager[Transaction]: # --- Calls and sends ------------------------------------------------------------------------------ -class Rpc(Protocol): - """ - Calls: ``ctx.rpc(handler, key=...)`` returns a function of the handler's - input, if it takes one, which makes the call and returns its result's - awaitable. A call is made when that function is called, awaited or not:: - - count = await ctx.rpc(increment, key="a", idempotency_key=f"{ctx.request.id}:inc")(5) - """ +class Caller(Protocol): + """Calls to other handlers, through Restate.""" @typing.overload - def __call__( + def call( self, handler: Callable[[Any], Awaitable[O]], - key: str | None = None, + /, *, + key: str | None = None, idempotency_key: str | None = None, headers: dict[str, str] | None = None, - ) -> Callable[[], CallFuture[O]]: ... + ) -> CallFuture[O]: ... @typing.overload - def __call__( + def call( self, handler: Callable[[Any, I], Awaitable[O]], - key: str | None = None, + arg: I, + /, *, + key: str | None = None, idempotency_key: str | None = None, headers: dict[str, str] | None = None, - ) -> Callable[[I], CallFuture[O]]: ... + ) -> CallFuture[O]: ... - def __call__( + def call( self, handler: Callable[..., Awaitable[Any]], - key: str | None = None, + arg: Any = None, + /, *, + key: str | None = None, idempotency_key: str | None = None, headers: dict[str, str] | None = None, - ) -> Callable[..., CallFuture[Any]]: - """ - A caller of ``handler``, with ``key`` for an object's or a workflow's - (a ``TerminalError`` otherwise, as a retry would fail the same way). + ) -> CallFuture[Any]: + """Call ``handler`` with ``arg``, and return a future of its result. + + The call is typed by the handler's own signature, and is made right + away, whether or not it is awaited:: + + balance = await ctx.call(deposit, 10, key="alice", idempotency_key=f"{ctx.request.id}:deposit") - ``idempotency_key`` is the same on every attempt, so that a new attempt - gets the first call's result; without one, every attempt makes the call - again, the one that resumes a suspension too. Restate scopes it to the - target handler. + Args: + handler: A handler of a service, a virtual object (a session + object's too) or a workflow. + arg: Its input, if it takes one. + key: The object's or workflow's key; a service takes none. + idempotency_key: The same on every attempt, so that a retried + attempt gets the first call's result instead of calling again. + headers: Headers to send with the call. + + Raises: + TerminalError: If ``handler`` is not a handler, or ``key`` is + missing for an object or a workflow, or given for a service. """ ... - def generic( + def generic_call( self, service: str, handler: str, @@ -176,62 +148,73 @@ def generic( input_serde: Serde[Any] = JsonSerde(), output_serde: Serde[Any] = JsonSerde(), ) -> CallFuture[Any]: - """ - Call a target known only at run time, by its names: untyped, and JSON - unless serdes are given. ``key`` is the caller's to give, for an - object or a workflow. + """Call a handler known only at run time, by its service's and its own name. + + Untyped: the input and result are JSON unless serdes are given. """ ... -class Send(Protocol): - """ - Sends, as ``Rpc``'s calls: ``ctx.send(handler, key=...)`` returns a function - of the handler's input, if it takes one, which makes the send at once - (``ctx.send``), or records it for its transaction's commit (``tx.send``):: - - ctx.send(record, key="audit", delay=timedelta(hours=1))(entry) - """ +class Sender(Protocol): + """Sends to other handlers, through Restate: one-way calls, which do not wait for the target.""" @typing.overload - def __call__( + def send( self, handler: Callable[[Any], Awaitable[Any]], - key: str | None = None, + /, *, + key: str | None = None, delay: timedelta | None = None, idempotency_key: str | None = None, headers: dict[str, str] | None = None, - ) -> Callable[[], None]: ... + ) -> None: ... @typing.overload - def __call__( + def send( self, handler: Callable[[Any, I], Awaitable[Any]], - key: str | None = None, + arg: I, + /, *, + key: str | None = None, delay: timedelta | None = None, idempotency_key: str | None = None, headers: dict[str, str] | None = None, - ) -> Callable[[I], None]: ... + ) -> None: ... - def __call__( + def send( self, handler: Callable[..., Awaitable[Any]], - key: str | None = None, + arg: Any = None, + /, *, + key: str | None = None, delay: timedelta | None = None, idempotency_key: str | None = None, headers: dict[str, str] | None = None, - ) -> Callable[..., None]: - """ - A sender to ``handler``, with ``key`` for an object's or a workflow's; - ``delay`` is how long after the send (for ``tx.send``, after the - commit) the target runs. See ``Rpc`` for ``idempotency_key``. + ) -> None: + """Send ``arg`` to ``handler``, without waiting for it to run:: + + ctx.send(remind, note, key=ctx.key, delay=timedelta(days=1)) + + Args: + handler: A handler of a service, a virtual object (a session + object's too) or a workflow. + arg: Its input, if it takes one. + key: The object's or workflow's key; a service takes none. + delay: How long to wait before the target runs. + idempotency_key: The same on every attempt, so that a retried + attempt does not send again. + headers: Headers to send with the request. + + Raises: + TerminalError: If ``handler`` is not a handler, or ``key`` is + missing for an object or a workflow, or given for a service. """ ... - def generic( + def generic_send( self, service: str, handler: str, @@ -243,86 +226,104 @@ def generic( headers: dict[str, str] | None = None, input_serde: Serde[Any] = JsonSerde(), ) -> None: - """Send to a target known only at run time, by its names: untyped, and JSON unless a serde is given.""" + """Send to a handler known only at run time, by its service's and its own name (JSON unless a serde is given).""" ... -# --- The contexts --------------------------------------------------------------------------------- +# --- Transactions --------------------------------------------------------------------------------- -class SessionContext(Protocol): +class Transaction(DatabaseWriter, Sender, Protocol): + """One transaction of a ``Database``: its SQL, and the sends it commits with. + + ``tx.send`` sends right after the transaction commits, and never if it + rolls back. A send without an ``idempotency_key`` gets one from the + transaction, so a send that a retried attempt makes again is dropped. """ - The context of a session object's handler. A handler runs in attempts: - Restate starts a new one after a crash, a retryable error or a suspension, - and each runs the handler from its first line, finding what the attempts - before it committed. Its operations can be awaited together - (``asyncio.gather``, ``asyncio.wait``). + + +class Database(DatabaseWriter, Protocol): + """A session's database, ``ctx.db``. + + Each statement outside ``transaction()`` is a transaction of its own, + durable once its call returns. Calls run one at a time, in call order. + """ + + def transaction(self) -> typing.AsyncContextManager[Transaction]: + """Start a transaction, for the body of an ``async with``:: + + async with ctx.db.transaction() as tx: + await tx.execute("UPDATE accounts SET balance = balance - ? WHERE id = ?", [amount, ctx.key]) + tx.send(deposit, amount, key=to) + + It commits when the body ends, and rolls back if the body raises. It + holds the database until then: use the ``tx`` it gives, as a call to + ``ctx.db`` inside the body is refused. + """ + ... + + +# --- The contexts --------------------------------------------------------------------------------- + + +class SessionContext(Caller, Sender, Protocol): + """The context of a session object's handler. + + A handler runs in attempts: after a crash, a retryable error or a + suspension, Restate starts a new attempt, which runs the handler from the + start and finds what earlier attempts committed. Its operations may be + awaited together, with ``asyncio.gather`` or ``asyncio.wait``. """ @property def key(self) -> str: - """The session's key.""" + """The key of this session.""" ... @property def request(self) -> Request: - """The request: its id (the same on every attempt), headers, and the attempt's.""" + """The request: its id (the same on every attempt), its headers, and the attempt's.""" ... @property def db(self) -> Database: - """ - The key's database. Each statement outside a transaction is its own - commit, and ``db.transaction()`` groups several; a commit is durable - when its call returns. Calls run one at a time, in call order. - """ + """The key's database.""" ... - @property - def rpc(self) -> Rpc: - """Calls, typed by the handler called: ``await ctx.rpc(handler, key="a")(input)``.""" - ... + def cancel(self, invocation_id: str) -> None: + """Cancel the invocation ``invocation_id`` (such as ``"inv_1f…"``). - @property - def send(self) -> Send: - """Sends, made at once; a transaction's (``tx.send``) are made once it commits.""" + The cancellation is sent right away, and not awaited. Cancelling an + invocation that has ended does nothing, so a retried attempt may + cancel it again. + """ ... - @property - def cancel(self) -> Cancel: - """Cancel an invocation, by its id: see ``Cancel``.""" - ... +class ViewContext(Caller, Protocol): + """The context of a session object's view. -class ViewContext(Protocol): - """ - The context of a session object's view: it runs beside the handlers, and - reads the database as committed when its invocation started, whatever they - commit meanwhile. It writes nothing, and sends nothing, but it can cancel an - invocation: a handler of its own object's, say, as a Stop button does. + A view runs beside the handlers, and reads the database as committed + when it started, whatever they commit meanwhile. It writes and sends + nothing, but it may call other handlers, and cancel an invocation (a + handler of its own object's, as a Stop button does). """ @property def key(self) -> str: - """The session's key.""" + """The key of this session.""" ... @property def request(self) -> Request: - """The request: its id, headers, and the attempt's.""" + """The request: its id, its headers, and the attempt's.""" ... @property def db(self) -> DatabaseReader: - """The key's database as committed when the view's invocation started: reads only.""" + """The key's database, as committed when the view started: read-only.""" ... - @property - def rpc(self) -> Rpc: - """Calls, typed by the handler called.""" - ... - - @property - def cancel(self) -> Cancel: - """Cancel an invocation, by its id: see ``Cancel``.""" + def cancel(self, invocation_id: str) -> None: + """Cancel the invocation ``invocation_id``: see ``SessionContext.cancel``.""" ... diff --git a/python/restate/sessions/_calls.py b/python/restate/sessions/_calls.py index 1b855ed2..8c05a1fb 100644 --- a/python/restate/sessions/_calls.py +++ b/python/restate/sessions/_calls.py @@ -9,15 +9,14 @@ # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # """ -Calls and sends from a session object's handler: ``ctx.rpc`` and ``ctx.send``, -and a transaction's ``tx.send``. A target is a handler (with its key for an -object or a workflow), typed by its signature; or, for one known only at run -time, its names, with ``generic``. Options go with the target: what it returns -takes only the input:: - - count = await ctx.rpc(increment, key="a", idempotency_key=f"{ctx.request.id}:inc")(5) - ctx.send(record, key="audit", delay=timedelta(hours=1))(entry) - answer = await ctx.rpc.generic("Counter", "get", key="a") +Calls and sends from a session object's handler: ``ctx.call`` and +``ctx.send``, and a transaction's ``tx.send``. A target is a handler (with +its key, for an object or a workflow), typed by its signature; a target known +only at run time is named, with ``generic_call`` and ``generic_send``:: + + count = await ctx.call(increment, 5, key="a", idempotency_key=f"{ctx.request.id}:inc") + ctx.send(record, entry, key="audit", delay=timedelta(hours=1)) + answer = await ctx.generic_call("Counter", "get", key="a") """ from __future__ import annotations @@ -116,58 +115,49 @@ async def cancel_invocation(self) -> None: await self._call.cancel_invocation() -class Caller: - """``Rpc`` (see ``_api.py``) over the SDK's context: each call a generic call, its result read by its serde.""" +class Calls: + """``call`` and ``generic_call`` (see ``Caller`` in ``_api.py``), over ``_call_host``, the SDK's context.""" - def __init__(self, host: CallHost) -> None: - self._host: CallHost = host + _call_host: CallHost @typing.overload - def __call__( + def call( self, handler: Callable[[Any], Awaitable[O]], - key: str | None = None, + /, *, + key: str | None = None, idempotency_key: str | None = None, headers: dict[str, str] | None = None, - ) -> Callable[[], CallFuture[O]]: ... + ) -> CallFuture[O]: ... @typing.overload - def __call__( + def call( self, handler: Callable[[Any, I], Awaitable[O]], - key: str | None = None, + arg: I, + /, *, + key: str | None = None, idempotency_key: str | None = None, headers: dict[str, str] | None = None, - ) -> Callable[[I], CallFuture[O]]: ... + ) -> CallFuture[O]: ... - def __call__( + def call( self, handler: Callable[..., Awaitable[Any]], - key: str | None = None, + arg: Any = None, + /, *, + key: str | None = None, idempotency_key: str | None = None, headers: dict[str, str] | None = None, - ) -> Callable[..., CallFuture[Any]]: - """ - A caller of ``handler``, with ``key`` for an object's or a workflow's. - - ``idempotency_key`` is the same on every attempt, so that a new attempt - gets the first call's result; without one, every attempt makes the call - again, the one that resumes a suspension too. Restate scopes it to the - target handler. - """ + ) -> CallFuture[Any]: service, name, input_serde, output_serde = _target(handler, key) + future = self._call_host.generic_call(service, name, input_serde.serialize(arg), key, idempotency_key, headers) + return CallFuture(future, output_serde) - def call(arg: Any = None) -> CallFuture[Any]: - parameter = input_serde.serialize(arg) - future = self._host.generic_call(service, name, parameter, key, idempotency_key, headers) - return CallFuture(future, output_serde) - - return call - - def generic( + def generic_call( self, service: str, handler: str, @@ -179,66 +169,57 @@ def generic( input_serde: Serde[Any] = JsonSerde(), output_serde: Serde[Any] = JsonSerde(), ) -> CallFuture[Any]: - """ - Call a target known only at run time, by its names: untyped, and JSON - unless serdes are given. ``key`` is the caller's to give, for an - object or a workflow. - """ - future = self._host.generic_call(service, handler, input_serde.serialize(arg), key, idempotency_key, headers) + parameter = input_serde.serialize(arg) + future = self._call_host.generic_call(service, handler, parameter, key, idempotency_key, headers) return CallFuture(future, output_serde) -class Sender: - """``Send`` (see ``_api.py``) into a sink: ``send_now``'s for ``ctx.send``, a transaction's outbox for ``tx.send``.""" +class Sends: + """``send`` and ``generic_send`` (see ``Sender`` in ``_api.py``), into ``_send_sink``: made at once, or at a commit.""" - def __init__(self, sink: Callable[[GenericSend], None]) -> None: - self._sink: Callable[[GenericSend], None] = sink + _send_sink: Callable[[GenericSend], None] @typing.overload - def __call__( + def send( self, handler: Callable[[Any], Awaitable[Any]], - key: str | None = None, + /, *, + key: str | None = None, delay: timedelta | None = None, idempotency_key: str | None = None, headers: dict[str, str] | None = None, - ) -> Callable[[], None]: ... + ) -> None: ... @typing.overload - def __call__( + def send( self, handler: Callable[[Any, I], Awaitable[Any]], - key: str | None = None, + arg: I, + /, *, + key: str | None = None, delay: timedelta | None = None, idempotency_key: str | None = None, headers: dict[str, str] | None = None, - ) -> Callable[[I], None]: ... + ) -> None: ... - def __call__( + def send( self, handler: Callable[..., Awaitable[Any]], - key: str | None = None, + arg: Any = None, + /, *, + key: str | None = None, delay: timedelta | None = None, idempotency_key: str | None = None, headers: dict[str, str] | None = None, - ) -> Callable[..., None]: - """ - A sender to ``handler``, with ``key`` for an object's or a workflow's; - ``delay`` is how long after the send (for ``tx.send``, after the - commit) the target runs. See ``Caller`` for ``idempotency_key``. - """ + ) -> None: service, name, input_serde, _ = _target(handler, key) + parameter = input_serde.serialize(arg) + self._send_sink(GenericSend(service, name, parameter, key, delay, idempotency_key, headers)) - def send(arg: Any = None) -> None: - parameter = input_serde.serialize(arg) - self._sink(GenericSend(service, name, parameter, key, delay, idempotency_key, headers)) - - return send - - def generic( + def generic_send( self, service: str, handler: str, @@ -250,8 +231,8 @@ def generic( headers: dict[str, str] | None = None, input_serde: Serde[Any] = JsonSerde(), ) -> None: - """Send to a target known only at run time, by its names: untyped, and JSON unless a serde is given.""" - self._sink(GenericSend(service, handler, input_serde.serialize(arg), key, delay, idempotency_key, headers)) + parameter = input_serde.serialize(arg) + self._send_sink(GenericSend(service, handler, parameter, key, delay, idempotency_key, headers)) def send_now(host: CallHost) -> Callable[[GenericSend], None]: diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index 2b9c82a8..7e696a94 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -45,8 +45,8 @@ async def list(ctx: sessions.ViewContext) -> list[str]: from restate.retry_policy import InvocationRetryPolicy from restate.serde import BytesSerde, DefaultSerde, Serde from restate.server_context import DisconnectedException, ServerInvocationContext -from restate.sessions._api import Cancel, Database, DatabaseReader, Rpc, Send, SessionContext, ViewContext -from restate.sessions._calls import Caller, Sender, send_now +from restate.sessions._api import Database, DatabaseReader, SessionContext, ViewContext +from restate.sessions._calls import CallHost, Calls, GenericSend, Sends, send_now from restate.sessions._link import Link from restate.sessions._schema import Schema, Step, steps_of from restate.sessions._scope import SCOPE, Scope @@ -71,26 +71,32 @@ async def list(ctx: sessions.ViewContext) -> list[str]: @dataclass(frozen=True, eq=False) -class AttemptContext: +class AttemptContext(Calls, Sends): """``SessionContext`` (see ``_api.py``): one attempt's, made as it starts. Compared by identity.""" key: str request: Request db: Database - rpc: Rpc - send: Send - cancel: Cancel + _call_host: CallHost + _send_sink: Callable[[GenericSend], None] + _cancel: Callable[[str], None] + + def cancel(self, invocation_id: str) -> None: + self._cancel(invocation_id) @dataclass(frozen=True, eq=False) -class AttemptViewContext: +class AttemptViewContext(Calls): """``ViewContext`` (see ``_api.py``): one attempt's, made as it starts. Compared by identity.""" key: str request: Request db: DatabaseReader - rpc: Rpc - cancel: Cancel + _call_host: CallHost + _cancel: Callable[[str], None] + + def cancel(self, invocation_id: str) -> None: + self._cancel(invocation_id) class _ContextHost: @@ -200,8 +206,8 @@ async def _nothing() -> None: """The action of an empty run.""" -def _cancel_of(host: SqlHost) -> Cancel: - """``Cancel`` over the SDK's context.""" +def _cancel_of(host: SqlHost) -> Callable[[str], None]: + """``cancel`` over the SDK's context.""" def cancel(invocation_id: str) -> None: # Not an id: a retry would fail the same way. @@ -230,15 +236,21 @@ class _Entry: def _context_of(entry: _Entry, host: SqlHost, request: Request, attempt: Attempt) -> SessionContext | ViewContext: """The context of an attempt of ``entry``.""" - rpc = Caller(host) cancel = _cancel_of(host) sdk = typing.cast("ServerInvocationContext | None", getattr(host, "context", None)) scope = Scope(sdk=sdk, explicit_cancellation=entry.explicit_cancellation, definition=entry.definition()) if entry.view: - view = AttemptViewContext(key=host.key, request=request, db=reader_of(attempt.db), rpc=rpc, cancel=cancel) + view = AttemptViewContext( + key=host.key, request=request, db=reader_of(attempt.db), _call_host=host, _cancel=cancel + ) return SCOPE.on(view, scope) context = AttemptContext( - key=host.key, request=request, db=attempt.db, rpc=rpc, send=Sender(send_now(host)), cancel=cancel + key=host.key, + request=request, + db=attempt.db, + _call_host=host, + _send_sink=send_now(host), + _cancel=cancel, ) return SCOPE.on(context, scope) @@ -539,7 +551,7 @@ def register(fn: F) -> F: what = f"{self.name}.{fn.__name__} is a handler already, of the SDK's or of a session object's." instead = ( "A session object's handlers are functions of its context, of their own; put journaled work " - "in a restate.VirtualObject, Service or Workflow, and call it with ctx.rpc." + "in a restate.VirtualObject, Service or Workflow, and call it with ctx.call." ) raise TypeError(f"{what} {instead}") if not inspect.iscoroutinefunction(inspect.unwrap(fn)): @@ -568,7 +580,7 @@ def register(fn: F) -> F: entries = ENTRIES.of(self) assert entries is not None entries[sdk.name] = entry - # The function itself, marked as the handler it is: `ctx.rpc(fn, ...)` + # The function itself, marked as the handler it is: `ctx.call(fn, ...)` # and the client find it, and a test calls it with a context of # its own (a fake of the protocols). vars(fn)[RESTATE_UNIQUE_HANDLER_SYMBOL] = sdk diff --git a/python/restate/sessions/_session.py b/python/restate/sessions/_session.py index afce043b..336ee2b2 100644 --- a/python/restate/sessions/_session.py +++ b/python/restate/sessions/_session.py @@ -35,8 +35,8 @@ from restate._internal import PyVM # pylint: disable=no-name-in-module from restate.exceptions import SuspendedException, TerminalError -from restate.sessions._api import Database, DatabaseReader, Send, Transaction -from restate.sessions._calls import CallHost, GenericSend, Sender +from restate.sessions._api import Database, DatabaseReader, Transaction +from restate.sessions._calls import CallHost, GenericSend, Sends from restate.sessions._database import ( Connection, ConnectionTransaction, @@ -295,9 +295,8 @@ class _Sends: def __init__(self) -> None: self._sends: list[dict[str, Any]] = [] self._open: bool = True - self.send: Send = Sender(self._record) - def _record(self, send: GenericSend) -> None: + def record(self, send: GenericSend) -> None: if not self._open: raise TerminalError("the transaction has ended: send from inside its body") stored: dict[str, Any] = { @@ -370,12 +369,12 @@ async def execute(self, sql: str, params: Sequence[SqlParam] | None = None) -> N await self.query_columns(sql, params) -class AttemptTransaction(Writes): +class AttemptTransaction(Writes, Sends): """``Transaction`` (see ``_api.py``): the SQL of a connection's transaction, and the sends it commits with.""" - def __init__(self, handle: ConnectionTransaction, send: Send) -> None: + def __init__(self, handle: ConnectionTransaction, sink: Callable[[GenericSend], None]) -> None: self._handle: ConnectionTransaction = handle - self.send: Send = send + self._send_sink = sink async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: return await self._handle.query_columns(sql, params) @@ -615,7 +614,7 @@ async def transaction(self) -> AsyncGenerator[AttemptTransaction, None]: async with db.transaction() as handle: sends = _Sends() if self._has_outbox else None # A view's writes stay in its session: it has nothing to reset. - tx = AttemptTransaction(handle, sends.send if sends is not None else Sender(_refused_send)) + tx = AttemptTransaction(handle, sends.record if sends is not None else _refused_send) if not self._host.shared: _RESET.on(tx, lambda: self._reset(handle)) try: diff --git a/tests/sessions_harness.py b/tests/sessions_harness.py index 3f5e0a99..03df29e4 100644 --- a/tests/sessions_harness.py +++ b/tests/sessions_harness.py @@ -118,8 +118,8 @@ async def flaky(ctx: sessions.SessionContext, text: str) -> int: async def priced(ctx: sessions.SessionContext, sku: str) -> list[int]: """A call and SQL awaited together, and two calls at once: one await point, the selector's.""" first, second, row = await asyncio.gather( - ctx.rpc(price)(sku), - ctx.rpc(price)(sku * 2), + ctx.call(price, sku), + ctx.call(price, sku * 2), ctx.db.get("SELECT count(*) AS n FROM notes"), ) assert row is not None @@ -131,7 +131,7 @@ async def forward(ctx: sessions.SessionContext, text: str) -> None: """A note, and its delivery to the inbox of the same key, sent once the transaction commits.""" async with ctx.db.transaction() as tx: await tx.execute("INSERT INTO notes (text) VALUES (?)", [text]) - tx.send(deliver, key=ctx.key)(text) + tx.send(deliver, text, key=ctx.key) sleeper = restate.Service("sleeper") @@ -146,7 +146,7 @@ async def nap(ctx: restate.Context) -> str: @notes.handler(explicit_cancellation=True) async def stoppable(ctx: sessions.SessionContext, text: str) -> list[str]: """Waits for a nap; cancelled, it cancels the nap, records that it stopped, and answers the calls it cancelled.""" - napping = asyncio.ensure_future(ctx.rpc(nap)()) + napping = asyncio.ensure_future(ctx.call(nap)) stopping = asyncio.ensure_future(internal.cancelled(ctx)) await asyncio.wait({napping, stopping}, return_when=asyncio.FIRST_COMPLETED) assert stopping.done() diff --git a/tests/sessions_object.py b/tests/sessions_object.py index 872672ba..a05917e5 100644 --- a/tests/sessions_object.py +++ b/tests/sessions_object.py @@ -231,7 +231,7 @@ async def test_the_sends_of_a_transaction_are_made_once_it_commits_and_never_if_ async def place(ctx: sessions.SessionContext, order: str) -> None: async with ctx.db.transaction() as tx: await tx.execute("INSERT INTO orders VALUES (?)", [order]) - tx.send(record, key="orders")(order) + tx.send(record, order, key="orders") if order == "bad": raise ValueError("no") @@ -257,7 +257,7 @@ async def test_the_sends_a_crash_left_in_the_outbox_are_made_again_with_the_same async def place(ctx: sessions.SessionContext, order: str) -> None: async with ctx.db.transaction() as tx: await tx.execute("INSERT INTO orders VALUES (?)", [order]) - tx.send(record, key="orders")(order) + tx.send(record, order, key="orders") # The commit's sync goes through; then the process crashes, before the outbox is emptied. first.crash_after(1) @@ -339,15 +339,15 @@ async def test_calls_are_typed_by_their_handler_and_checked_for_their_key(): @caller.handler() async def price(ctx: sessions.SessionContext, sku: str) -> int: - return await ctx.rpc(quote, idempotency_key="k")(sku) + return await ctx.call(quote, sku, idempotency_key="k") @caller.handler() async def unkeyed(ctx: sessions.SessionContext) -> None: - ctx.send(record)("x") + ctx.send(record, "x") @caller.handler() async def literal(ctx: sessions.SessionContext) -> Any: - return await ctx.rpc.generic("pricing", "quote", "pear") + return await ctx.generic_call("pricing", "quote", "pear") host = Host(Backend(), service="caller", answers={"quote": b"42"}) assert await run(caller, "price", host, "apple") == 42 @@ -362,7 +362,7 @@ async def test_ctx_send_is_made_at_once_and_cancel_takes_an_invocation_id(): @sender.handler() async def go(ctx: sessions.SessionContext) -> None: - ctx.send(record, key="k")("entry") + ctx.send(record, "entry", key="k") ctx.cancel("inv_1") ctx.cancel("not-an-id") @@ -493,7 +493,7 @@ async def explicit(ctx: sessions.SessionContext) -> None: async def test_cancelled_ends_its_waiter_cancelled_when_the_attempt_ends_first(monkeypatch): from restate import _preview # pylint: disable=import-outside-toplevel from restate.exceptions import SuspendedException # pylint: disable=import-outside-toplevel - from restate.sessions import _calls, _object, _scope # pylint: disable=import-outside-toplevel + from restate.sessions import _object, _scope # pylint: disable=import-outside-toplevel async def attempt_ended(_sdk: Any) -> None: raise SuspendedException() @@ -504,9 +504,9 @@ async def attempt_ended(_sdk: Any) -> None: key="k", request=request(), db=typing.cast(Any, None), - rpc=_calls.Caller(host), - send=_calls.Sender(lambda send: None), - cancel=lambda invocation_id: None, + _call_host=host, + _send_sink=lambda send: None, + _cancel=lambda invocation_id: None, ) _scope.SCOPE.on(ctx, _scope.Scope(sdk=typing.cast(Any, object()), explicit_cancellation=True)) watcher = asyncio.ensure_future(internal.cancelled(ctx)) @@ -548,9 +548,22 @@ class FakeContext: db: FakeDatabase key: str = "alice" request: Request = field(default_factory=request) - rpc: sessions.Rpc = field(default_factory=lambda: typing.cast(sessions.Rpc, None)) - send: sessions.Send = field(default_factory=lambda: typing.cast(sessions.Send, None)) - cancel: sessions.Cancel = lambda invocation_id: None + sent: list[tuple[Any, Any]] = field(default_factory=list) + + def call(self, handler: Any, arg: Any = None, /, **options: Any) -> Any: + raise NotImplementedError + + def generic_call(self, service: str, handler: str, arg: Any = None, **options: Any) -> Any: + raise NotImplementedError + + def send(self, handler: Any, arg: Any = None, /, **options: Any) -> None: + self.sent.append((handler, arg)) + + def generic_send(self, service: str, handler: str, arg: Any = None, **options: Any) -> None: + self.sent.append((f"{service}/{handler}", arg)) + + def cancel(self, invocation_id: str) -> None: + pass async def test_a_handler_runs_over_a_fake_context_of_a_test(): From 248b1b99ebae9697344c61e07c65011b742362e4 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 18:10:58 +0100 Subject: [PATCH 34/45] restate.agents: a smaller namespace; wait_for_signal, PerSession - The signal tool is wait_for_signal(signal, description=, input=, name=): no longer shadowing the stdlib's signal next to Signal and signals, and name= is the tool's name, as for every other tool. - Configure is PerSession: a value, or a (sync or async) function of the invocation's context giving it. - Companion, companion_of, deliver, signals, CANCELLED, REFUSED and UNANSWERED move to restate.agents.internal: autocomplete on restate.agents shows what an agent's author uses. - The package's docstring says what it is, with an example, and no references to the TypeScript library. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/agents/__init__.py | 51 ++++++++++++--------- python/restate/agents/_common.py | 14 +++--- python/restate/agents/_tools.py | 41 ++++++++++------- python/restate/agents/internal.py | 11 +++++ python/restate/agents/langgraph/_agent.py | 8 ++-- python/restate/agents/pydantic_ai/_agent.py | 8 ++-- tests/agents_langgraph.py | 3 +- tests/agents_object.py | 4 +- tests/agents_pydantic_ai.py | 3 +- tests/agents_tools.py | 4 +- 10 files changed, 89 insertions(+), 58 deletions(-) diff --git a/python/restate/agents/__init__.py b/python/restate/agents/__init__.py index 645b0000..f0c069ea 100644 --- a/python/restate/agents/__init__.py +++ b/python/restate/agents/__init__.py @@ -8,16 +8,33 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -AI agents in session objects, each conversation in its session's database. -The Python variant of ``@restatedev/agents``: what the frameworks share -(session objects set up for agents, tools for any of them, and the tools that -are Restate itself). Each framework is a module of its own, with the extra of -its packages; what their integrations need beside this is in -``restate.agents.internal``. +"""AI agents in session objects, each conversation kept in its key's database. + +A crash, a restart, or a person who takes a day to approve changes nothing: +each message is answered once, and each tool call is made once:: + + from restate import sessions + from restate.agents import langgraph, tool + + + @tool + async def get_weather(city: str) -> str: + '''The current weather in a city.''' + return f"Sunny in {city}" + + + assistant = langgraph.agent("assistant", system="Answer briefly.", tools=[get_weather]) + + if __name__ == "__main__": + sessions.serve(assistant) -A crash, a restart or a person who takes a day to approve changes nothing: -each message is answered once, and each tool call is made once. +This package has what every framework shares: ``AgentObject``, session +objects set up for agents; tools for any framework (``tool``, and the tools +that are Restate itself: ``call``, ``send``, ``workflow``, ``subagent``, +``wait_for_signal``); and the types of the conversation that clients read. +Each framework is a subpackage with an extra of its own: +``restate.agents.langgraph`` and ``restate.agents.pydantic_ai``. What +integrations need beside this is in ``restate.agents.internal``. """ from restate.agents._api import ( @@ -36,23 +53,18 @@ Waiting, Watch, ) -from restate.agents._common import CANCELLED, REFUSED, UNANSWERED, Configure +from restate.agents._common import PerSession from restate.agents._object import AgentObject, Synced, SyncRequest -from restate.agents._signals import Companion, companion_of, deliver, signals -from restate.agents._tools import Tool, ToolCall, call, send, signal, subagent, tool, workflow +from restate.agents._tools import Tool, ToolCall, call, send, subagent, tool, wait_for_signal, workflow __all__ = [ - "CANCELLED", - "REFUSED", - "UNANSWERED", "AgentObject", "ChatMessage", - "Companion", - "Configure", "Messages", "MessagesRequest", "Notifications", "Part", + "PerSession", "Resolved", "Signal", "Status", @@ -67,12 +79,9 @@ "Waiting", "Watch", "call", - "companion_of", - "deliver", "send", - "signal", - "signals", "subagent", "tool", + "wait_for_signal", "workflow", ] diff --git a/python/restate/agents/_common.py b/python/restate/agents/_common.py index 20c3ee72..2bfd2f98 100644 --- a/python/restate/agents/_common.py +++ b/python/restate/agents/_common.py @@ -9,7 +9,7 @@ # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # """ -What the frameworks share: their options (``Configure``, the model), the ids of +What the frameworks share: their options (``PerSession``, the model), the ids of an invocation's prompts, a model provider's refusal, and the ``chat`` message. As sdk-sessions' ``common.ts``. """ @@ -31,17 +31,17 @@ T = typing.TypeVar("T") -Configure: typing.TypeAlias = T | Callable[[SessionContext], T | Awaitable[T]] -""" -A value, or a function of the invocation's context that gives it: the same for -every message, or each user's (from the object's database, say):: +PerSession: typing.TypeAlias = T | Callable[[SessionContext], T | Awaitable[T]] +"""A value, or a function of the invocation's context (sync or async) that returns it. + +For options that are the same for every message, or that differ per user:: langgraph.agent("assistant", model="anthropic/claude-sonnet-5-5") - langgraph.agent("assistant", model=lambda ctx: user_model(ctx)) + langgraph.agent("assistant", system=lambda ctx: f"You help {ctx.key}.") """ -async def configured(value: Configure[T], ctx: SessionContext) -> T: +async def configured(value: PerSession[T], ctx: SessionContext) -> T: """``value`` for the invocation of ``ctx``.""" if not callable(value): return value diff --git a/python/restate/agents/_tools.py b/python/restate/agents/_tools.py index ffa30a56..8680c3cb 100644 --- a/python/restate/agents/_tools.py +++ b/python/restate/agents/_tools.py @@ -35,7 +35,7 @@ from restate.agents import _signals from restate.agents._calls import call_of -from restate.agents._common import Configure, ToolError, configured +from restate.agents._common import PerSession, ToolError, configured from restate.agents._turn import stop_signal_of, tool_call_ended, tool_call_started from restate.exceptions import TerminalError from restate.handler import Handler, handler_from_callable @@ -308,7 +308,7 @@ def call( *, description: str, name: str | None = None, - key: Configure[str] | None = None, + key: PerSession[str] | None = None, ) -> Tool: """ A tool that calls handler ``target`` and answers with its result:: @@ -349,7 +349,7 @@ def send( *, description: str, name: str | None = None, - key: Configure[str] | None = None, + key: PerSession[str] | None = None, delay: timedelta | None = None, ) -> Tool: """A tool that sends to handler ``target``, without waiting for it: a one-way call, made once, after ``delay``.""" @@ -388,7 +388,7 @@ def workflow( *, description: str, name: str | None = None, - key: Configure[str] | None = None, + key: PerSession[str] | None = None, ) -> Tool: """ A tool that runs the workflow of ``target`` (its main handler) to the end, @@ -482,34 +482,43 @@ def validate(args: dict[str, Any]) -> dict[str, Any]: ) -def signal(name: str, *, description: str, input: Any = None, tool_name: str | None = None) -> Tool: - """ - A tool that waits for the signal ``name`` to the agent, and answers with its - value: a person's approval, say, or another system's result. The agent - waits without holding a process. With an ``input`` type, the model says - what it waits for (such as the refund to approve), which ``pending`` shows:: +def wait_for_signal(signal: str, *, description: str, input: Any = None, name: str | None = None) -> Tool: + """Create a tool that waits for a signal to the agent, and answers with its value. + + A person's approval, say, or another system's result. The agent waits + without holding a process, for as long as it takes:: - signal("approval", description="Wait for a manager to approve a refund.", input=Refund) + approval = wait_for_signal("approval", description="Wait for a manager to approve a refund.", input=Refund) - Each wait has an id: the agent's ``pending`` lists them, and its ``signal`` - resolves one, by its id, or the oldest wait for a name:: + Each wait has an id. The agent's ``pending`` view lists the waits, and its + ``signal`` view resolves one, by its id or as the oldest wait for a name:: curl localhost:8080/support/alice/signal --json '{"name": "approval", "value": "approved"}' + + Args: + signal: The signal's name. + description: What the tool waits for, for the model. + input: The type of what the model says it waits for (such as the + refund to approve), which ``pending`` shows. None if it says nothing. + name: The tool's name. Defaults to ``"wait_for_"``. + + Returns: + The tool, for any framework's agent. """ shape = _input_of(input, None) async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: async def made(c: ToolCall) -> Any: wait_id = c.idempotency_key - request: _signals.WaitRequest = {"id": wait_id, "name": name} + request: _signals.WaitRequest = {"id": wait_id, "name": signal} if input is not None: request["input"] = arguments["raw"] return await c.ctx.call(_signals.wait, request, key=_signals.signals_of(c.ctx), idempotency_key=wait_id) - return await _once(f"signal/{name}", arguments["raw"], tool_call, made) + return await _once(f"signal/{signal}", arguments["raw"], tool_call, made) return Tool( - name=tool_name or f"wait_for_{name}", + name=name or f"wait_for_{signal}", description=description, input_schema=shape.schema, _run=run, diff --git a/python/restate/agents/internal.py b/python/restate/agents/internal.py index db06cb29..a4976332 100644 --- a/python/restate/agents/internal.py +++ b/python/restate/agents/internal.py @@ -16,9 +16,12 @@ """ from restate.agents._common import ( + CANCELLED, CHECK_MODEL, CLAUDE_MODEL, DEFAULT_MODEL, + REFUSED, + UNANSWERED, MessageSerde, ModelError, ToolError, @@ -33,6 +36,7 @@ refusal_of, stored_answer, ) +from restate.agents._signals import Companion, companion_of, deliver, signals from restate.agents._tools import run_tool from restate.agents._transcript import Transcript from restate.agents._turn import ( @@ -47,10 +51,14 @@ ) __all__ = [ + "CANCELLED", "CHECK_MODEL", "CLAUDE_MODEL", "DEFAULT_MODEL", + "REFUSED", + "UNANSWERED", "AttemptEnded", + "Companion", "MessageSerde", "ModelError", "Read", @@ -58,7 +66,9 @@ "ToolError", "Transcript", "check_model_name", + "companion_of", "configured", + "deliver", "from_tool", "model_id_of", "model_of", @@ -68,6 +78,7 @@ "refusal", "refusal_of", "run_tool", + "signals", "stop_signal_of", "stored_answer", "tool_call_ended", diff --git a/python/restate/agents/langgraph/_agent.py b/python/restate/agents/langgraph/_agent.py index cbf0e0fe..68f7ce9d 100644 --- a/python/restate/agents/langgraph/_agent.py +++ b/python/restate/agents/langgraph/_agent.py @@ -52,8 +52,8 @@ from restate.agents._common import ( CANCELLED, UNANSWERED, - Configure, ModelError, + PerSession, configured, from_tool, model_of, @@ -341,9 +341,9 @@ async def work() -> str: def agent( name: str, *, - model: Configure[ChatModel | None] = None, - system: Configure[str | None] = None, - tools: Configure[Sequence[AnyTool]] = (), + model: PerSession[ChatModel | None] = None, + system: PerSession[str | None] = None, + tools: PerSession[Sequence[AnyTool]] = (), description: str | None = None, schema: Sequence[Step] = (), background: bool = False, diff --git a/python/restate/agents/pydantic_ai/_agent.py b/python/restate/agents/pydantic_ai/_agent.py index 9a0d645c..665f0c53 100644 --- a/python/restate/agents/pydantic_ai/_agent.py +++ b/python/restate/agents/pydantic_ai/_agent.py @@ -58,8 +58,8 @@ from restate.agents._common import ( CANCELLED, UNANSWERED, - Configure, ModelError, + PerSession, configured, from_tool, model_of, @@ -402,9 +402,9 @@ def agent( pydantic_agent: Agent[Any, str] | None = None, /, *, - model: Configure[AnyModel | None] = None, - system: Configure[str | None] = None, - tools: Configure[Sequence[AnyTool]] = (), + model: PerSession[AnyModel | None] = None, + system: PerSession[str | None] = None, + tools: PerSession[Sequence[AnyTool]] = (), description: str | None = None, schema: Sequence[Step] = (), background: bool = False, diff --git a/tests/agents_langgraph.py b/tests/agents_langgraph.py index cffc3b58..b8a8b0b9 100644 --- a/tests/agents_langgraph.py +++ b/tests/agents_langgraph.py @@ -33,7 +33,8 @@ import restate # noqa: E402 from restate import HarnessEnvironment, sessions # noqa: E402 -from restate.agents import CANCELLED, tool # noqa: E402 +from restate.agents import tool # noqa: E402 +from restate.agents.internal import CANCELLED # noqa: E402 from restate.agents import langgraph # noqa: E402 from restate.agents._api import Messages # noqa: E402 diff --git a/tests/agents_object.py b/tests/agents_object.py index aa8706c8..6ed621c0 100644 --- a/tests/agents_object.py +++ b/tests/agents_object.py @@ -25,7 +25,7 @@ from restate import HarnessEnvironment, sessions from restate.agents._api import Messages, Status from restate.agents._object import AgentObject -from restate.agents._tools import ToolCall, run_tool, signal +from restate.agents._tools import ToolCall, run_tool, wait_for_signal from restate.agents._transcript import Transcript from restate.agents._turn import stop_signal_of, transcribing @@ -47,7 +47,7 @@ class Refund(pydantic.BaseModel): amount: float -approval = signal("approval", description="Wait for a manager to approve a refund.", input=Refund) +approval = wait_for_signal("approval", description="Wait for a manager to approve a refund.", input=Refund) # A scripted framework: it keeps each key's transcript in this process, and answers with an echo. transcripts: dict[str, Transcript] = {} diff --git a/tests/agents_pydantic_ai.py b/tests/agents_pydantic_ai.py index ab96cc07..e7c6bd60 100644 --- a/tests/agents_pydantic_ai.py +++ b/tests/agents_pydantic_ai.py @@ -38,7 +38,8 @@ import restate # noqa: E402 from restate import HarnessEnvironment, sessions # noqa: E402 -from restate.agents import CANCELLED, tool # noqa: E402 +from restate.agents import tool # noqa: E402 +from restate.agents.internal import CANCELLED # noqa: E402 from restate.agents import pydantic_ai # noqa: E402 from restate.agents._api import Messages # noqa: E402 diff --git a/tests/agents_tools.py b/tests/agents_tools.py index 23d59b26..a93983e4 100644 --- a/tests/agents_tools.py +++ b/tests/agents_tools.py @@ -24,7 +24,7 @@ from restate.sessions import internal from restate.agents._common import ToolError from restate.agents._object import AgentObject -from restate.agents._tools import Tool, ToolCall, call, signal, subagent, tool +from restate.agents._tools import Tool, ToolCall, call, subagent, tool, wait_for_signal from restate.context import Request from sessions_database import Backend # pylint: disable=import-error from sessions_object import Host, request # pylint: disable=import-error @@ -130,7 +130,7 @@ class Refund(pydantic.BaseModel): async def test_signal_waits_at_the_companion_of_the_agent_with_the_tool_calls_key(): - tools["approval"] = signal("approval", description="Wait for a manager's approval.", input=Refund) + tools["approval"] = wait_for_signal("approval", description="Wait for a manager's approval.", input=Refund) assert tools["approval"].name == "wait_for_approval" assert set(tools["approval"].input_schema["properties"]) == {"order", "amount"} host = Host(Backend(), service="runner", key="alice", answers={"wait": b'"approved"'}) From 53b90a32a675319d4718b78820626d84401e2f52 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 18:14:19 +0100 Subject: [PATCH 35/45] restate.sessions: migrations= and Migration, for what schema= was A session object's migrations=[...] is an append-only list of steps (SQL, or async functions of the transaction), each run once per key; calling it migrations, and its type Migration, makes the rule (never edit a step, add one) obvious. SessionObject, AgentObject and the frameworks' agent() take it; tutorial lesson 04 is 04_migrations.py. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- examples-agents/README.md | 2 +- examples-agents/basic/notes.py | 2 +- examples-agents/tutorial/02_database.py | 4 +-- examples-agents/tutorial/03_views.py | 2 +- .../{04_schema.py => 04_migrations.py} | 14 +++++----- examples-agents/tutorial/05_transaction.py | 2 +- examples-agents/tutorial/06_retries.py | 2 +- examples-agents/tutorial/07_calls.py | 4 +-- examples-agents/tutorial/08_sends.py | 2 +- python/restate/agents/_object.py | 8 +++--- python/restate/agents/langgraph/_agent.py | 8 +++--- python/restate/agents/pydantic_ai/_agent.py | 8 +++--- python/restate/sessions/__init__.py | 6 ++--- python/restate/sessions/_object.py | 10 +++---- python/restate/sessions/_schema.py | 19 +++++++------- tests/agents_definitions.py | 4 +-- tests/examples_agents.py | 2 +- tests/sessions_harness.py | 2 +- tests/sessions_object.py | 26 ++++++++++--------- 19 files changed, 64 insertions(+), 63 deletions(-) rename examples-agents/tutorial/{04_schema.py => 04_migrations.py} (79%) diff --git a/examples-agents/README.md b/examples-agents/README.md index 0bf39533..7fe93f8d 100644 --- a/examples-agents/README.md +++ b/examples-agents/README.md @@ -47,7 +47,7 @@ real application keeps its data across releases. | [`01_object.py`](tutorial/01_object.py) | A virtual object: handlers that Restate calls by key | | [`02_database.py`](tutorial/02_database.py) | `ctx.db`: every wallet has its own SQL database | | [`03_views.py`](tutorial/03_views.py) | Views: reading the database while a handler is running | -| [`04_schema.py`](tutorial/04_schema.py) | Changing the tables of databases that already exist | +| [`04_migrations.py`](tutorial/04_migrations.py) | Changing the tables of databases that already exist | ### Part 2 · Changes that happen once diff --git a/examples-agents/basic/notes.py b/examples-agents/basic/notes.py index 47d37b76..06ce4712 100644 --- a/examples-agents/basic/notes.py +++ b/examples-agents/basic/notes.py @@ -15,7 +15,7 @@ from restate import sessions # The schema of each key's database: its steps, applied once per key (a step added later runs then). -notes = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) +notes = sessions.SessionObject("notes", migrations=["CREATE TABLE notes (text TEXT NOT NULL)"]) @notes.handler diff --git a/examples-agents/tutorial/02_database.py b/examples-agents/tutorial/02_database.py index 04521cc6..f329d20f 100644 --- a/examples-agents/tutorial/02_database.py +++ b/examples-agents/tutorial/02_database.py @@ -3,7 +3,7 @@ Each key of a session object has its own SQLite database, ``ctx.db``. It is stored in Restate together with the object's state, so there is no database -server to run. ``schema`` lists the statements that create the tables; they +server to run. ``migrations`` lists the statements that create the tables; they run once per key, the first time that key is used. The wallet records each deposit as a row, and the balance is their sum. Bob's @@ -17,7 +17,7 @@ from restate import sessions -wallet = sessions.SessionObject("wallet", schema=["CREATE TABLE entries (amount INTEGER NOT NULL)"]) +wallet = sessions.SessionObject("wallet", migrations=["CREATE TABLE entries (amount INTEGER NOT NULL)"]) @wallet.handler diff --git a/examples-agents/tutorial/03_views.py b/examples-agents/tutorial/03_views.py index 479393a2..1ece278c 100644 --- a/examples-agents/tutorial/03_views.py +++ b/examples-agents/tutorial/03_views.py @@ -15,7 +15,7 @@ from restate import sessions -wallet = sessions.SessionObject("wallet", schema=["CREATE TABLE entries (amount INTEGER NOT NULL)"]) +wallet = sessions.SessionObject("wallet", migrations=["CREATE TABLE entries (amount INTEGER NOT NULL)"]) async def balance_of(db: sessions.DatabaseReader) -> int: diff --git a/examples-agents/tutorial/04_schema.py b/examples-agents/tutorial/04_migrations.py similarity index 79% rename from examples-agents/tutorial/04_schema.py rename to examples-agents/tutorial/04_migrations.py index ce02a58c..5b83185a 100644 --- a/examples-agents/tutorial/04_schema.py +++ b/examples-agents/tutorial/04_migrations.py @@ -1,15 +1,15 @@ """ -04 · Changing the schema +04 · Migrations: changing the tables Alice's database already exists, so editing its CREATE TABLE statement would -change nothing for her. That's why ``schema`` is a list of steps: to change the -tables, add a step at the end. Each key's database runs the steps it hasn't run -yet, the next time it is used. Never edit or remove a step that may have run -somewhere; add a new one instead. +change nothing for her. That's why ``migrations`` is a list of steps: to change +the tables, add a step at the end. Each key's database runs the steps it hasn't +run yet, the next time it is used. Never edit or remove a step that may have +run somewhere; add a new one instead. Here, deposits get a note:: - uv run examples-agents/tutorial/04_schema.py # and register it (see 01) + uv run examples-agents/tutorial/04_migrations.py # and register it (see 01) curl localhost:8080/wallet/alice/deposit --json '{"amount": 50, "note": "birthday"}' curl localhost:8080/wallet/alice/history # the earlier deposits have no note """ @@ -20,7 +20,7 @@ wallet = sessions.SessionObject( "wallet", - schema=[ + migrations=[ "CREATE TABLE entries (amount INTEGER NOT NULL)", # from lesson 02 "ALTER TABLE entries ADD COLUMN note TEXT", # new: runs once on each wallet's database ], diff --git a/examples-agents/tutorial/05_transaction.py b/examples-agents/tutorial/05_transaction.py index 7ef51ac8..a79720e7 100644 --- a/examples-agents/tutorial/05_transaction.py +++ b/examples-agents/tutorial/05_transaction.py @@ -27,7 +27,7 @@ wallet = sessions.SessionObject( "wallet", - schema=["CREATE TABLE entries (amount INTEGER NOT NULL)", "ALTER TABLE entries ADD COLUMN note TEXT"], + migrations=["CREATE TABLE entries (amount INTEGER NOT NULL)", "ALTER TABLE entries ADD COLUMN note TEXT"], ) diff --git a/examples-agents/tutorial/06_retries.py b/examples-agents/tutorial/06_retries.py index f42b9dc3..37c25aca 100644 --- a/examples-agents/tutorial/06_retries.py +++ b/examples-agents/tutorial/06_retries.py @@ -27,7 +27,7 @@ wallet = sessions.SessionObject( "wallet", - schema=[ + migrations=[ "CREATE TABLE entries (amount INTEGER NOT NULL)", "ALTER TABLE entries ADD COLUMN note TEXT", # New: an id per entry. Entries from earlier lessons have none, which a unique index allows. diff --git a/examples-agents/tutorial/07_calls.py b/examples-agents/tutorial/07_calls.py index f1a1d172..685ebfd4 100644 --- a/examples-agents/tutorial/07_calls.py +++ b/examples-agents/tutorial/07_calls.py @@ -23,7 +23,7 @@ import restate from restate import sessions -handle = sessions.SessionObject("handle", schema=["CREATE TABLE owner (wallet TEXT NOT NULL)"]) +handle = sessions.SessionObject("handle", migrations=["CREATE TABLE owner (wallet TEXT NOT NULL)"]) @handle.handler(name="claim") @@ -43,7 +43,7 @@ async def owner(ctx: sessions.ViewContext) -> str | None: wallet = sessions.SessionObject( "wallet", - schema=[ + migrations=[ "CREATE TABLE entries (amount INTEGER NOT NULL)", "ALTER TABLE entries ADD COLUMN note TEXT", "ALTER TABLE entries ADD COLUMN id TEXT", diff --git a/examples-agents/tutorial/08_sends.py b/examples-agents/tutorial/08_sends.py index 96cb8ac6..eaa71bfe 100644 --- a/examples-agents/tutorial/08_sends.py +++ b/examples-agents/tutorial/08_sends.py @@ -25,7 +25,7 @@ wallet = sessions.SessionObject( "wallet", - schema=[ + migrations=[ "CREATE TABLE entries (amount INTEGER NOT NULL)", "ALTER TABLE entries ADD COLUMN note TEXT", "ALTER TABLE entries ADD COLUMN id TEXT", diff --git a/python/restate/agents/_object.py b/python/restate/agents/_object.py index 3be253b2..a13f175e 100644 --- a/python/restate/agents/_object.py +++ b/python/restate/agents/_object.py @@ -37,7 +37,7 @@ from restate.agents._common import MessageSerde, refusal_of from restate.agents._turn import AGENT, AgentInfo, chat_turn, cleared from restate.retry_policy import InvocationRetryPolicy -from restate.sessions import SessionContext, SessionObject, Step, ViewContext +from restate.sessions import Migration, SessionContext, SessionObject, ViewContext from restate.sessions.internal import reset T = typing.TypeVar("T") @@ -111,7 +111,7 @@ class AgentObject(SessionObject): forgets it (its database too), and ``sync``, which a client following it reads. A framework's ``agent()`` makes one. - ``schema``: tables of your own, beside the agent's (named ``restate_...``). + ``migrations``: tables of your own, beside the agent's (named ``restate_...``). ``background``: it works unwatched (a subagent, say): its turn publishes at its end only, not at each step. """ @@ -121,7 +121,7 @@ def __init__( name: str, *, description: str | None = None, - schema: Sequence[Step] = (), + migrations: Sequence[Migration] = (), background: bool = False, metadata: dict[str, str] | None = None, inactivity_timeout: timedelta | None = None, @@ -133,7 +133,7 @@ def __init__( ) -> None: super().__init__( name, - schema=schema, + migrations=migrations, description=description, metadata=metadata, inactivity_timeout=inactivity_timeout, diff --git a/python/restate/agents/langgraph/_agent.py b/python/restate/agents/langgraph/_agent.py index 68f7ce9d..c3041ea8 100644 --- a/python/restate/agents/langgraph/_agent.py +++ b/python/restate/agents/langgraph/_agent.py @@ -69,7 +69,7 @@ from restate.agents.langgraph._checkpointer import SqlCheckpointSaver from restate.exceptions import TerminalError from restate.retry_policy import InvocationRetryPolicy -from restate.sessions import SessionContext, Step +from restate.sessions import Migration, SessionContext _CANCELLATION = 409 """The status of Restate's cancellation, as a turn that was stopped fails with it.""" @@ -345,7 +345,7 @@ def agent( system: PerSession[str | None] = None, tools: PerSession[Sequence[AnyTool]] = (), description: str | None = None, - schema: Sequence[Step] = (), + migrations: Sequence[Migration] = (), background: bool = False, metadata: dict[str, str] | None = None, inactivity_timeout: timedelta | None = None, @@ -375,7 +375,7 @@ def agent( system: The system prompt. tools: ``restate.agents`` tools, or LangChain tools. description: The object's description, shown in Restate's UI. - schema: Tables of your own, beside the agent's. + migrations: Tables of your own, beside the agent's. background: Whether it works unwatched (a subagent, say): it then publishes its progress once a turn ends, not at each step. metadata: The object's metadata. @@ -395,7 +395,7 @@ def agent( obj = AgentObject( name, description=description, - schema=schema, + migrations=migrations, background=background, metadata=metadata, inactivity_timeout=inactivity_timeout, diff --git a/python/restate/agents/pydantic_ai/_agent.py b/python/restate/agents/pydantic_ai/_agent.py index 665f0c53..31750069 100644 --- a/python/restate/agents/pydantic_ai/_agent.py +++ b/python/restate/agents/pydantic_ai/_agent.py @@ -74,7 +74,7 @@ from restate.agents._turn import AttemptEnded, stop_signal_of, tool_call_ended, tool_call_started, transcribing from restate.exceptions import TerminalError from restate.retry_policy import InvocationRetryPolicy -from restate.sessions import Database, SessionContext, Step +from restate.sessions import Database, Migration, SessionContext from restate.sessions.internal import Schema _CANCELLATION = 409 @@ -406,7 +406,7 @@ def agent( system: PerSession[str | None] = None, tools: PerSession[Sequence[AnyTool]] = (), description: str | None = None, - schema: Sequence[Step] = (), + migrations: Sequence[Migration] = (), background: bool = False, metadata: dict[str, str] | None = None, inactivity_timeout: timedelta | None = None, @@ -442,7 +442,7 @@ def agent( system: The instructions, when there is no ``pydantic_agent``. tools: ``restate.agents`` tools, Pydantic AI tools or tool functions. description: The object's description, shown in Restate's UI. - schema: Tables of your own, beside the agent's. + migrations: Tables of your own, beside the agent's. background: Whether it works unwatched (a subagent, say): it then publishes its progress once a turn ends, not at each step. metadata: The object's metadata. @@ -461,7 +461,7 @@ def agent( obj = AgentObject( name, description=description, - schema=schema, + migrations=migrations, background=background, metadata=metadata, inactivity_timeout=inactivity_timeout, diff --git a/python/restate/sessions/__init__.py b/python/restate/sessions/__init__.py index fffb26fc..d78e0f0f 100644 --- a/python/restate/sessions/__init__.py +++ b/python/restate/sessions/__init__.py @@ -15,7 +15,7 @@ from restate import sessions - notes = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) + notes = sessions.SessionObject("notes", migrations=["CREATE TABLE notes (text TEXT NOT NULL)"]) @notes.handler @@ -55,7 +55,7 @@ async def list_notes(ctx: sessions.ViewContext) -> list[str]: from restate.sessions._calls import CallFuture from restate.sessions._database import Row, Rows, SqlError, SqlParam, SqlValue from restate.sessions._object import SessionObject -from restate.sessions._schema import Step +from restate.sessions._schema import Migration from restate.sessions._serve import app, serve __all__ = [ @@ -64,6 +64,7 @@ async def list_notes(ctx: sessions.ViewContext) -> list[str]: "Database", "DatabaseReader", "DatabaseWriter", + "Migration", "Row", "Rows", "Sender", @@ -72,7 +73,6 @@ async def list_notes(ctx: sessions.ViewContext) -> list[str]: "SqlError", "SqlParam", "SqlValue", - "Step", "Transaction", "ViewContext", "app", diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index 7e696a94..5a166559 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -14,7 +14,7 @@ context of its own, made of the key, the request, the database and calls, and nothing else of the SDK's context:: - notes = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) + notes = sessions.SessionObject("notes", migrations=["CREATE TABLE notes (text TEXT NOT NULL)"]) @notes.handler async def add(ctx: sessions.SessionContext, text: str) -> None: @@ -48,7 +48,7 @@ async def list(ctx: sessions.ViewContext) -> list[str]: from restate.sessions._api import Database, DatabaseReader, SessionContext, ViewContext from restate.sessions._calls import CallHost, Calls, GenericSend, Sends, send_now from restate.sessions._link import Link -from restate.sessions._schema import Schema, Step, steps_of +from restate.sessions._schema import Migration, Schema, steps_of from restate.sessions._scope import SCOPE, Scope from restate.sessions._selector import FutureSelector, hook from restate.sessions._session import Attempt, ObjectConfig, SqlHost, reader_of, with_attempt @@ -337,7 +337,7 @@ class SessionObject(VirtualObject): """ A virtual object whose state is a SQL database per key. - - ``schema``: the steps that make each key's database, in order, applied + - ``migrations``: the steps that make each key's database, in order, applied once per key by the first handler that uses it. A step that ran never changes: a change is a step of its own. - ``@obj.handler``: handlers run one at a time per key, each with a @@ -354,7 +354,7 @@ def __init__( self, name: str, *, - schema: Sequence[Step] = (), + migrations: Sequence[Migration] = (), description: str | None = None, metadata: dict[str, str] | None = None, inactivity_timeout: timedelta | None = None, @@ -379,7 +379,7 @@ def __init__( ingress_private=ingress_private, invocation_retry_policy=invocation_retry_policy, ) - self._config: ObjectConfig = ObjectConfig(schema=Schema(name, schema) if steps_of(schema) else None) + self._config: ObjectConfig = ObjectConfig(schema=Schema(name, migrations) if steps_of(migrations) else None) ENTRIES.on(self, {}) def __repr__(self) -> str: diff --git a/python/restate/sessions/_schema.py b/python/restate/sessions/_schema.py index 1103a2ec..7e034e28 100644 --- a/python/restate/sessions/_schema.py +++ b/python/restate/sessions/_schema.py @@ -16,7 +16,7 @@ added later (a column, an index) runs on each key's database once, as that key is next opened:: - notes = sessions.SessionObject("notes", schema=[ + notes = sessions.SessionObject("notes", migrations=[ "CREATE TABLE notes (text TEXT NOT NULL)", "ALTER TABLE notes ADD COLUMN at TEXT", ]) @@ -48,28 +48,27 @@ never loaded too. """ -Step = str | Callable[[Any], Awaitable[None]] -"""A statement, or an async function of the transaction.""" +Migration = str | Callable[[Any], Awaitable[None]] +"""One step of a database's migrations: a SQL statement, or an async function of the transaction.""" _ENSURE = "_restate_sessions_ensure" """The attribute of a database of this package's: how it applies schemas (see ``Schema.ensure``).""" -def steps_of(steps: Sequence[Step]) -> tuple[Step, ...]: - """``steps`` as a tuple; a lone string (a sequence of characters) is a mistake, not a list of statements.""" +def steps_of(steps: Sequence[Migration]) -> tuple[Migration, ...]: + """Return ``steps`` as a tuple; raise ``TypeError`` for a lone string, which is a sequence of characters.""" if isinstance(steps, str): - raise TypeError( - f"schema takes a list of steps, not a string: schema=[{steps[:40]!r}{'...' if len(steps) > 40 else ''}]" - ) + shown = f"{steps[:40]!r}{'...' if len(steps) > 40 else ''}" + raise TypeError(f"migrations takes a list of steps, not a string: migrations=[{shown}]") return tuple(steps) class Schema: """A part of a database's schema: its ``name``, and the ``steps`` that make it, in order.""" - def __init__(self, name: str, steps: Sequence[Step]) -> None: + def __init__(self, name: str, steps: Sequence[Migration]) -> None: self.name: str = name - self.steps: tuple[Step, ...] = steps_of(steps) + self.steps: tuple[Migration, ...] = steps_of(steps) async def ensure(self, db: Any) -> None: """ diff --git a/tests/agents_definitions.py b/tests/agents_definitions.py index b7b45114..6db95aaf 100644 --- a/tests/agents_definitions.py +++ b/tests/agents_definitions.py @@ -19,9 +19,9 @@ from restate.agents.internal import provider_of -def test_a_schema_given_as_one_string_is_refused(): +def test_migrations_given_as_one_string_are_refused(): with pytest.raises(TypeError, match="a list of steps, not a string"): - sessions.SessionObject("p", schema="CREATE TABLE x (a)") # type: ignore[arg-type] + sessions.SessionObject("p", migrations="CREATE TABLE x (a)") # type: ignore[arg-type] def test_a_sync_function_is_refused_as_a_handler_or_view(): diff --git a/tests/examples_agents.py b/tests/examples_agents.py index d248872d..f7ec05a1 100644 --- a/tests/examples_agents.py +++ b/tests/examples_agents.py @@ -94,7 +94,7 @@ async def test_03_views(): async def test_04_schema(): - async with serving("tutorial/04_schema.py") as ingress: + async with serving("tutorial/04_migrations.py") as ingress: await ingress.post("wallet/alice/deposit", {"amount": 50, "note": "birthday"}) assert (await ingress.post("wallet/alice/history")).json() == [{"amount": 50, "note": "birthday"}] diff --git a/tests/sessions_harness.py b/tests/sessions_harness.py index 03df29e4..ba23844c 100644 --- a/tests/sessions_harness.py +++ b/tests/sessions_harness.py @@ -88,7 +88,7 @@ async def price(ctx: restate.Context, sku: str) -> int: # --- Session objects ---------------------------------------------------------------------------- -notes = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) +notes = sessions.SessionObject("notes", migrations=["CREATE TABLE notes (text TEXT NOT NULL)"]) @notes.handler() diff --git a/tests/sessions_object.py b/tests/sessions_object.py index a05917e5..361f8551 100644 --- a/tests/sessions_object.py +++ b/tests/sessions_object.py @@ -158,7 +158,7 @@ def run(obj: sessions.SessionObject, name: str, host: Host, *args: Any) -> typin # --- The objects -------------------------------------------------------------------------------- -notes = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) +notes = sessions.SessionObject("notes", migrations=["CREATE TABLE notes (text TEXT NOT NULL)"]) @notes.handler @@ -205,7 +205,7 @@ async def test_a_view_of_a_key_never_written_reads_empty_tables_and_stores_nothi async def test_a_step_added_later_runs_once_on_each_database(): backend = Backend() - v1 = sessions.SessionObject("evolving", schema=["CREATE TABLE t (a INTEGER)"]) + v1 = sessions.SessionObject("evolving", migrations=["CREATE TABLE t (a INTEGER)"]) @v1.handler() async def put(ctx: sessions.SessionContext) -> None: @@ -213,7 +213,9 @@ async def put(ctx: sessions.SessionContext) -> None: await run(v1, "put", Host(backend, service="evolving")) - v2 = sessions.SessionObject("evolving", schema=["CREATE TABLE t (a INTEGER)", "ALTER TABLE t ADD COLUMN b TEXT"]) + v2 = sessions.SessionObject( + "evolving", migrations=["CREATE TABLE t (a INTEGER)", "ALTER TABLE t ADD COLUMN b TEXT"] + ) @v2.handler() async def columns(ctx: sessions.SessionContext) -> list[str]: @@ -225,7 +227,7 @@ async def columns(ctx: sessions.SessionContext) -> list[str]: async def test_the_sends_of_a_transaction_are_made_once_it_commits_and_never_if_it_rolls_back(): - shop = sessions.SessionObject("shop", schema=["CREATE TABLE orders (id TEXT PRIMARY KEY)"]) + shop = sessions.SessionObject("shop", migrations=["CREATE TABLE orders (id TEXT PRIMARY KEY)"]) @shop.handler() async def place(ctx: sessions.SessionContext, order: str) -> None: @@ -251,7 +253,7 @@ async def place(ctx: sessions.SessionContext, order: str) -> None: async def test_the_sends_a_crash_left_in_the_outbox_are_made_again_with_the_same_keys(): - shop = sessions.SessionObject("shop2", schema=["CREATE TABLE orders (id TEXT PRIMARY KEY)"]) + shop = sessions.SessionObject("shop2", migrations=["CREATE TABLE orders (id TEXT PRIMARY KEY)"]) @shop.handler() async def place(ctx: sessions.SessionContext, order: str) -> None: @@ -281,7 +283,7 @@ async def place(ctx: sessions.SessionContext, order: str) -> None: async def test_an_sql_error_fails_the_handler_with_a_terminal_error(): host = Host(Backend()) - broken = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) + broken = sessions.SessionObject("notes", migrations=["CREATE TABLE notes (text TEXT NOT NULL)"]) @broken.handler() async def bad(ctx: sessions.SessionContext) -> None: @@ -294,7 +296,7 @@ async def bad(ctx: sessions.SessionContext) -> None: async def test_a_handler_catches_an_sql_error_and_goes_on(): host = Host(Backend()) - users = sessions.SessionObject("users", schema=["CREATE TABLE users (email TEXT PRIMARY KEY)"]) + users = sessions.SessionObject("users", migrations=["CREATE TABLE users (email TEXT PRIMARY KEY)"]) @users.handler async def register(ctx: sessions.SessionContext, email: str) -> str: @@ -309,7 +311,7 @@ async def register(ctx: sessions.SessionContext, email: str) -> str: async def test_a_call_to_ctx_db_inside_its_transaction_is_refused(): - nested = sessions.SessionObject("nested", schema=["CREATE TABLE t (a INTEGER)"]) + nested = sessions.SessionObject("nested", migrations=["CREATE TABLE t (a INTEGER)"]) @nested.handler() async def inside(ctx: sessions.SessionContext) -> None: @@ -323,7 +325,7 @@ async def inside(ctx: sessions.SessionContext) -> None: async def test_a_view_writes_nothing(): host = Host(Backend()) await run(notes, "add", host, "milk") - writer = sessions.SessionObject("notes", schema=["CREATE TABLE notes (text TEXT NOT NULL)"]) + writer = sessions.SessionObject("notes", migrations=["CREATE TABLE notes (text TEXT NOT NULL)"]) @writer.view() async def sneaky(ctx: sessions.ViewContext) -> None: @@ -374,7 +376,7 @@ async def go(ctx: sessions.SessionContext) -> None: async def test_another_database_of_the_key_is_apart_from_the_objects_own(): - multi = sessions.SessionObject("multi", schema=["CREATE TABLE own (a INTEGER)"]) + multi = sessions.SessionObject("multi", migrations=["CREATE TABLE own (a INTEGER)"]) @multi.handler() async def both(ctx: sessions.SessionContext) -> list[str]: @@ -389,7 +391,7 @@ async def both(ctx: sessions.SessionContext) -> list[str]: async def test_reset_empties_the_database_and_applies_the_schema_again(): - resettable = sessions.SessionObject("resettable", schema=["CREATE TABLE t (a INTEGER)"]) + resettable = sessions.SessionObject("resettable", migrations=["CREATE TABLE t (a INTEGER)"]) @resettable.handler() async def fill(ctx: sessions.SessionContext) -> None: @@ -656,7 +658,7 @@ async def test_an_attempt_is_freed_once_it_ends(): import weakref # pylint: disable=import-outside-toplevel kept: list[weakref.ref[Any]] = [] - freed = sessions.SessionObject("freed", schema=["CREATE TABLE t (a INTEGER)"]) + freed = sessions.SessionObject("freed", migrations=["CREATE TABLE t (a INTEGER)"]) @freed.handler async def use(ctx: sessions.SessionContext) -> None: From 1b0e7cc17862c61e83d0207fae5f24714d95ec9c Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 18:14:23 +0100 Subject: [PATCH 36/45] Untrack a .DS_Store committed by mistake, and ignore them Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- .gitignore | 3 +++ python/restate/agents/.DS_Store | Bin 6148 -> 0 bytes 2 files changed, 3 insertions(+) delete mode 100644 python/restate/agents/.DS_Store diff --git a/.gitignore b/.gitignore index 7e6a605c..4b0b03c8 100644 --- a/.gitignore +++ b/.gitignore @@ -174,3 +174,6 @@ restate-data # Test reports, test suite test-services/test_report/ restate-sdk-test-suite.jar + +# macOS +.DS_Store diff --git a/python/restate/agents/.DS_Store b/python/restate/agents/.DS_Store deleted file mode 100644 index 9aca61ea22f85cc5e7db46bcca792f4d5bb07410..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 6148 zcmeHK%Sr=55Ukc50=eYqael!+7()Dld_WWf5|C9yh1~Y!cll|lejtX82M=DP8oFzy zr)!q2!}c}+TfObBfhB-B-4P!irsn7F6T7I45$U{Rz$^B+#};>!>hlTbF7S*4cI15L z@835s!{gKNkhhC5X)v#9lZiGDEYZ5|e5*Vp!6dZ&lX|$Hb(=YWT2vvekrQ@pRtbq8!!} z6{Ua_I91>>w@dH;H}oIo|5K87Qa}p)D+O$}-mF)ArRuGdm-AlR=(lvQ`J}sX9TbLW l$HZvI+;}^_ilVG*zUFx^920}ie9(#d8E{=>QsA!@I0A<&73Kf{ From 9d5c0fee4061ed159e1b1376573f6d1e99941ae7 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 18:15:19 +0100 Subject: [PATCH 37/45] restate.agents: the default model comes from RESTATE_AGENT_MODEL Namespaced, so it cannot collide with another tool's MODEL; MODEL is no longer read. The name is internal.MODEL_ENV, and error messages name it. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- examples-agents/basic/langgraph_agent.py | 2 +- examples-agents/basic/pydantic_ai_agent.py | 2 +- python/restate/agents/_common.py | 11 +++++++---- python/restate/agents/internal.py | 2 ++ python/restate/agents/langgraph/_agent.py | 7 ++++--- python/restate/agents/pydantic_ai/_agent.py | 4 ++-- tests/agents_definitions.py | 11 ++++++++++- 7 files changed, 27 insertions(+), 12 deletions(-) diff --git a/examples-agents/basic/langgraph_agent.py b/examples-agents/basic/langgraph_agent.py index a3758fae..1c57dea8 100644 --- a/examples-agents/basic/langgraph_agent.py +++ b/examples-agents/basic/langgraph_agent.py @@ -9,7 +9,7 @@ curl localhost:8080/assistant/alice/chat -d 'What is the weather in Lisbon?' curl localhost:8080/assistant/alice/messages --json '{}' # the conversation -The model is ``openai/gpt-6-sol`` unless MODEL names another: +The model is ``openai/gpt-6-sol`` unless RESTATE_AGENT_MODEL names another: ``anthropic/`` takes ``langchain-anthropic`` (and ANTHROPIC_API_KEY) instead of ``langchain-openai``. """ diff --git a/examples-agents/basic/pydantic_ai_agent.py b/examples-agents/basic/pydantic_ai_agent.py index f6051154..309dfeea 100644 --- a/examples-agents/basic/pydantic_ai_agent.py +++ b/examples-agents/basic/pydantic_ai_agent.py @@ -9,7 +9,7 @@ curl localhost:8080/assistant/alice/chat -d 'What is the weather in Lisbon?' curl localhost:8080/assistant/alice/messages --json '{}' # the conversation -The model is ``openai/gpt-6-sol`` unless MODEL names another, of any of +The model is ``openai/gpt-6-sol`` unless RESTATE_AGENT_MODEL names another, of any of Pydantic AI's providers: ``anthropic/`` takes ``pydantic-ai-slim[anthropic]`` (and ANTHROPIC_API_KEY) instead. """ diff --git a/python/restate/agents/_common.py b/python/restate/agents/_common.py index 2bfd2f98..0fe65593 100644 --- a/python/restate/agents/_common.py +++ b/python/restate/agents/_common.py @@ -63,6 +63,9 @@ def next_prompt_id(ctx: SessionContext) -> str: return ctx.request.id if n == 1 else f"{ctx.request.id}:{n}" +MODEL_ENV = "RESTATE_AGENT_MODEL" +"""The environment variable of the model that agents use when given none.""" + DEFAULT_MODEL = "openai/gpt-6-sol" """The default model of the frameworks that take any provider's, with OPENAI_API_KEY.""" @@ -71,8 +74,8 @@ def next_prompt_id(ctx: SessionContext) -> str: def model_of(model: str | None, fallback: str = DEFAULT_MODEL) -> str: - """The model, as ``provider/model-id``: ``model``, or else MODEL, or else the framework's default.""" - return model or os.environ.get("MODEL") or fallback + """Return ``model``, else ``$RESTATE_AGENT_MODEL``, else the framework's default (``fallback``).""" + return model or os.environ.get(MODEL_ENV) or fallback def provider_of(model: str) -> tuple[str, str]: @@ -95,7 +98,7 @@ def provider_of(model: str) -> tuple[str, str]: API key, an unknown model): Failed Dependency, as its HTTP status. """ -CHECK_MODEL = "check the agent's model (its `model` option, or MODEL) and its provider's API key" +CHECK_MODEL = f"check the agent's model (its `model` option, or {MODEL_ENV}) and its provider's API key" """What to check when a model's provider refuses a request.""" @@ -120,7 +123,7 @@ def model_id_of(model: str, provider: str) -> str: """The model id of ``model``, for a framework with only ``provider``'s models: ``provider/model-id``, or the id alone.""" named, model_id = provider_of(model) if "/" in model or ":" in model else (provider, model) if named != provider or not model_id: - from_env = " (from MODEL)" if os.environ.get("MODEL") == model else "" + from_env = f" (from {MODEL_ENV})" if os.environ.get(MODEL_ENV) == model else "" raise ModelError(f'the model "{model}"{from_env} is not {provider}\'s: this agent takes {provider}/') return model_id diff --git a/python/restate/agents/internal.py b/python/restate/agents/internal.py index a4976332..e383a962 100644 --- a/python/restate/agents/internal.py +++ b/python/restate/agents/internal.py @@ -20,6 +20,7 @@ CHECK_MODEL, CLAUDE_MODEL, DEFAULT_MODEL, + MODEL_ENV, REFUSED, UNANSWERED, MessageSerde, @@ -55,6 +56,7 @@ "CHECK_MODEL", "CLAUDE_MODEL", "DEFAULT_MODEL", + "MODEL_ENV", "REFUSED", "UNANSWERED", "AttemptEnded", diff --git a/python/restate/agents/langgraph/_agent.py b/python/restate/agents/langgraph/_agent.py index c3041ea8..5c67c939 100644 --- a/python/restate/agents/langgraph/_agent.py +++ b/python/restate/agents/langgraph/_agent.py @@ -51,6 +51,7 @@ from restate.agents._common import ( CANCELLED, + MODEL_ENV, UNANSWERED, ModelError, PerSession, @@ -107,7 +108,7 @@ def _provider_of_model(named: str) -> tuple[str, str]: """``provider/model-id`` as its provider and model id, for a provider the usual graph takes.""" provider, model_id = provider_of(named) if provider not in _PROVIDERS or not model_id: - from_env = " (from MODEL)" if os.environ.get("MODEL") == named else "" + from_env = f" (from {MODEL_ENV})" if os.environ.get(MODEL_ENV) == named else "" raise ModelError(f'the model "{named}"{from_env} is neither openai/ nor anthropic/') return provider, model_id @@ -320,7 +321,7 @@ async def summarize(ctx: sessions.SessionContext, topic: str) -> str: ctx: The handler's context. text: The prompt. model: ``"openai/"`` or ``"anthropic/"``, or a - LangChain chat model that calls tools. Defaults to ``$MODEL``, + LangChain chat model that calls tools. Defaults to ``$RESTATE_AGENT_MODEL``, else ``"openai/gpt-6-sol"``. system: The system prompt. tools: ``restate.agents`` tools, or LangChain tools. @@ -370,7 +371,7 @@ def agent( name: The object's name, as Restate serves it. model: ``"openai/"`` (needs ``langchain-openai``) or ``"anthropic/"`` (needs ``langchain-anthropic``), or a - LangChain chat model that calls tools. Defaults to ``$MODEL``, + LangChain chat model that calls tools. Defaults to ``$RESTATE_AGENT_MODEL``, else ``"openai/gpt-6-sol"``. system: The system prompt. tools: ``restate.agents`` tools, or LangChain tools. diff --git a/python/restate/agents/pydantic_ai/_agent.py b/python/restate/agents/pydantic_ai/_agent.py index 31750069..42aa62ce 100644 --- a/python/restate/agents/pydantic_ai/_agent.py +++ b/python/restate/agents/pydantic_ai/_agent.py @@ -375,7 +375,7 @@ async def summarize(ctx: sessions.SessionContext, topic: str) -> str: Defaults to one with ``system`` as its instructions. model: ``"/"`` of any of Pydantic AI's providers, or a Pydantic AI model. Defaults to ``agent``'s own, else - ``$MODEL``, else ``"openai/gpt-6-sol"``. + ``$RESTATE_AGENT_MODEL``, else ``"openai/gpt-6-sol"``. system: The instructions, when there is no ``agent``. tools: ``restate.agents`` tools, Pydantic AI tools or tool functions, beside ``agent``'s own. @@ -438,7 +438,7 @@ def agent( model: ``"/"`` of any of Pydantic AI's providers (each needs its extra, such as ``pydantic-ai-slim[openai]``), or a Pydantic AI model. Defaults to ``pydantic_agent``'s own, else - ``$MODEL``, else ``"openai/gpt-6-sol"``. + ``$RESTATE_AGENT_MODEL``, else ``"openai/gpt-6-sol"``. system: The instructions, when there is no ``pydantic_agent``. tools: ``restate.agents`` tools, Pydantic AI tools or tool functions. description: The object's description, shown in Restate's UI. diff --git a/tests/agents_definitions.py b/tests/agents_definitions.py index 6db95aaf..67b0c907 100644 --- a/tests/agents_definitions.py +++ b/tests/agents_definitions.py @@ -16,7 +16,7 @@ from restate import sessions from restate.agents import AgentObject -from restate.agents.internal import provider_of +from restate.agents.internal import DEFAULT_MODEL, model_of, provider_of def test_migrations_given_as_one_string_are_refused(): @@ -88,3 +88,12 @@ def test_a_model_names_its_provider_with_a_slash_or_a_colon(): assert provider_of("anthropic:claude-sonnet-5-5") == ("anthropic", "claude-sonnet-5-5") assert provider_of("openai/ft:gpt-4o:acme") == ("openai", "ft:gpt-4o:acme") assert provider_of("gpt") == ("", "gpt") + + +def test_the_default_model_comes_from_restate_agent_model(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("MODEL", "openai/ignored") + monkeypatch.delenv("RESTATE_AGENT_MODEL", raising=False) + assert model_of(None) == DEFAULT_MODEL + monkeypatch.setenv("RESTATE_AGENT_MODEL", "anthropic:claude-sonnet-5-5") + assert model_of(None) == "anthropic:claude-sonnet-5-5" + assert model_of("openai/gpt-6-sol") == "openai/gpt-6-sol" From f38d0e4b78c436fc24c6f59fee1cd2b2efd44fb1 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 18:22:33 +0100 Subject: [PATCH 38/45] Docstrings in plain English, PEP 257 and Google style; README for newcomers restate.sessions, restate.agents and the LangGraph and Pydantic AI integrations: an imperative one-line summary for each public name, then Args/Returns/Raises where they help; plain sentences instead of the compressed phrases of the port; no references to the TypeScript library. SessionObject and AgentObject list what they offer (an agent's built-in, reserved views); the integrations say how they differ from restate.ext's langchain and pydantic ones. Docstrings and comments only: the code is unchanged (checked on the ASTs). examples-agents/README.md starts from pip install "restate-sdk[serve]", says sessions.serve needs Python 3.11+, and explains why chat takes plain text while other handlers take JSON. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- examples-agents/README.md | 50 ++- python/restate/agents/__init__.py | 27 +- python/restate/agents/_api.py | 104 ++--- python/restate/agents/_calls.py | 72 ++-- python/restate/agents/_common.py | 107 ++--- python/restate/agents/_log.py | 80 ++-- python/restate/agents/_object.py | 222 ++++++----- python/restate/agents/_signals.py | 167 ++++---- python/restate/agents/_tools.py | 261 ++++++++---- python/restate/agents/_transcript.py | 43 +- python/restate/agents/_turn.py | 144 ++++--- python/restate/agents/internal.py | 18 +- python/restate/agents/langgraph/__init__.py | 28 +- python/restate/agents/langgraph/_agent.py | 105 ++--- .../restate/agents/langgraph/_checkpointer.py | 99 ++--- python/restate/agents/pydantic_ai/__init__.py | 26 +- python/restate/agents/pydantic_ai/_agent.py | 115 +++--- python/restate/sessions/_calls.py | 40 +- python/restate/sessions/_database.py | 263 +++++++------ python/restate/sessions/_link.py | 26 +- python/restate/sessions/_object.py | 272 ++++++++----- python/restate/sessions/_schema.py | 129 +++--- python/restate/sessions/_scope.py | 45 ++- python/restate/sessions/_selector.py | 112 +++--- python/restate/sessions/_serve.py | 107 +++-- python/restate/sessions/_session.py | 371 ++++++++++-------- python/restate/sessions/internal.py | 31 +- 27 files changed, 1779 insertions(+), 1285 deletions(-) diff --git a/examples-agents/README.md b/examples-agents/README.md index 7fe93f8d..6bf2f6d8 100644 --- a/examples-agents/README.md +++ b/examples-agents/README.md @@ -1,17 +1,26 @@ # Session objects and agents: examples Examples of [`restate.sessions`](../python/restate/sessions/__init__.py), -virtual objects with a SQL database per key, and -[`restate.agents`](../python/restate/agents/__init__.py), AI agents in them. -They are the Python counterparts of -[sdk-sessions](https://github.com/restatedev/sdk-sessions)' tutorial and -examples. Each is a file that runs on its own. +which gives each key of a virtual object its own SQL database, and of +[`restate.agents`](../python/restate/agents/__init__.py), AI agents built on +top of it. Each example is a single file that runs on its own. ## Run an example -From the repository's root, once (it builds the SDK, with its native -extension, and installs the `serve` extra the examples run on: uvicorn's -HTTP/2, on Python 3.11 or later): +The examples serve themselves with `sessions.serve`, which needs Python 3.11 +or later. (The SDK itself supports Python 3.10; there, serve +`sessions.app(...)` with any ASGI server, such as hypercorn.) + +Install the SDK with the `serve` extra, which adds the HTTP/2 server that +`sessions.serve` runs on: + +```sh +pip install "restate-sdk[serve]" +``` + +Or, to run the examples against the SDK from this repository, sync it from +the repository's root once. This builds the SDK's native extension, so it +needs a Rust toolchain: ```sh uv sync --all-extras --all-packages @@ -20,17 +29,25 @@ uv sync --all-extras --all-packages Then run an example, start Restate, and tell Restate where the example is: ```sh -uv run examples-agents/tutorial/01_object.py # the example, on port 9080 +python examples-agents/tutorial/01_object.py # the example, on port 9080 (or: uv run ...) npx @restatedev/restate-server # Restate, in another terminal npx @restatedev/restate -y deployments register --force http://localhost:9080 # tell Restate where it is curl localhost:8080/wallet/alice/hello ``` -Session objects need Restate 1.8 or later. An agent's framework comes with its extra, and its model's provider with -its package: the comment at the top of each agent's file says which (`uv run ---extra langgraph --with langchain-openai ...`, say). Every example serves on port 9080 (or `PORT`): stop one before you start the -next, and register again. The comment at the top of each file lists the -commands to try it. +Session objects need Restate 1.8 or later. Every example serves on port 9080 +(or `PORT`), so stop one before you start the next, and register again. The +comment at the top of each file lists the commands to try it. + +The agent examples also need their framework's extra and the package of their +model's provider; the comment at the top of each agent's file says which. For +example: `pip install "restate-sdk[serve,langgraph]" langchain-openai`, or +`uv run --extra langgraph --with langchain-openai ...` from source. + +Most handlers take and return JSON, so their requests use `curl --json`. An +agent's `chat` handler is the exception: it takes the message as plain text +(`curl -d 'Hi'`; a JSON string such as `"Hi"` works too), so that chatting +from a terminal needs no quoting. It answers with a JSON string. ## Tutorial: a wallet, from a database to agents @@ -58,8 +75,9 @@ real application keeps its data across releases. | [`07_calls.py`](tutorial/07_calls.py) | Calling another object, with an idempotency key | | [`08_sends.py`](tutorial/08_sends.py) | Transfers: a send that is delivered only if the debit commits | -Part 3, agents in the wallet (a prompt, an assistant, its tools, a person in -the loop, a subagent), comes with the agents' framework integrations. +A third part, agents in the wallet (a prompt, an assistant and its tools, a +person in the loop, a subagent), is not written yet. Until then, the agent +examples under [Basic](#basic) show the agents' framework integrations. ## Basic diff --git a/python/restate/agents/__init__.py b/python/restate/agents/__init__.py index f0c069ea..d2d4a34a 100644 --- a/python/restate/agents/__init__.py +++ b/python/restate/agents/__init__.py @@ -8,10 +8,10 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -"""AI agents in session objects, each conversation kept in its key's database. +"""AI agents built on session objects, with each conversation kept in its key's database. -A crash, a restart, or a person who takes a day to approve changes nothing: -each message is answered once, and each tool call is made once:: +A crash, a restart, or a person who takes a day to approve something changes +nothing: each message is answered once, and each tool call is made once:: from restate import sessions from restate.agents import langgraph, tool @@ -28,13 +28,20 @@ async def get_weather(city: str) -> str: if __name__ == "__main__": sessions.serve(assistant) -This package has what every framework shares: ``AgentObject``, session -objects set up for agents; tools for any framework (``tool``, and the tools -that are Restate itself: ``call``, ``send``, ``workflow``, ``subagent``, -``wait_for_signal``); and the types of the conversation that clients read. -Each framework is a subpackage with an extra of its own: -``restate.agents.langgraph`` and ``restate.agents.pydantic_ai``. What -integrations need beside this is in ``restate.agents.internal``. +This package holds what all frameworks share: + +- ``AgentObject``, a session object set up for agents. +- Tools that work with any framework: ``tool`` for your own functions, and + ``call``, ``send``, ``workflow``, ``subagent`` and ``wait_for_signal``, + which are built on Restate. +- ``PerSession``, for agent options that differ per invocation. +- The types of the conversation that clients read and send. + +Each framework has a subpackage and an extra of its own: +``restate.agents.langgraph`` (``pip install "restate-sdk[langgraph]"``) and +``restate.agents.pydantic_ai`` (``pip install "restate-sdk[pydantic_ai]"``). +``restate.agents.internal`` has the building blocks for writing another +integration. """ from restate.agents._api import ( diff --git a/python/restate/agents/_api.py b/python/restate/agents/_api.py index f980aa15..0a525707 100644 --- a/python/restate/agents/_api.py +++ b/python/restate/agents/_api.py @@ -8,11 +8,12 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -What clients of an agent read and send, as JSON: its conversation's messages -and pages of them, and the message it answers. The field names are -sdk-sessions' (``invocationId``), so that a client follows an agent of either -SDK. +"""The JSON types that an agent's clients send and receive. + +These cover the conversation's messages and pages of them, the agent's +status, and signals. Some keys are camelCase (``invocationId``, +``timeoutSeconds``): they are wire-format names, shared with Restate's other +SDKs so that the same client can follow an agent written in any of them. """ from __future__ import annotations @@ -21,11 +22,11 @@ from typing import Any, Literal, TypedDict Role = Literal["user", "assistant"] -"""Who said a message.""" +"""Who wrote a message.""" class TextPart(TypedDict): - """Text a message says.""" + """A message's text.""" type: Literal["text"] text: str @@ -39,10 +40,10 @@ class _ToolCall(TypedDict): class ToolPart(_ToolCall, total=False): - """ - A tool call a message makes, with its output once it ran, and ``error: - true`` if the output is the tool's error (``Error: `` of a - ``restate.TerminalError``, say). + """A tool call in a message, with its output once the tool has run. + + ``error`` is ``True`` if the output is an error, such as the + ``"Error: "`` of a ``restate.TerminalError``. """ output: str @@ -50,36 +51,36 @@ class ToolPart(_ToolCall, total=False): Part = TextPart | ToolPart -"""What a message says: text, or a tool call.""" +"""A part of a message: text, or a tool call.""" class Said(TypedDict): - """A message as a framework's transcript gives it: its role and parts, without an id.""" + """A message as a framework's transcript produces it: its role and parts, without an id.""" role: Role parts: list[Part] class ChatMessage(Said): - """A message of a conversation, as an agent's ``messages`` view gives it.""" + """A message of a conversation, as an agent's ``messages`` view returns it.""" id: int - """Its id, from 0, in the order of the conversation; the same however it changes. ``before`` takes it.""" + """The message's id, which increases through the conversation and stays the same when the message changes.""" class MessagesRequest(TypedDict, total=False): - """ - Which messages to read: the latest ``limit``, or those before a message, - for a page of the conversation; or those new or changed since a - ``cursor``, to follow it. + """The request of the ``messages`` view. + + For a page of the conversation, pass ``limit`` and optionally ``before``. + To follow the conversation, pass the ``cursor`` of the previous response. """ limit: int - """At most this many messages (default 50).""" + """The maximum number of messages to return (default 50).""" before: int - """The messages before this one (its ``id``): older ones, such as for scrolling back.""" + """Return the messages before the one with this ``id``, for scrolling back.""" cursor: int - """The messages new or changed since this cursor (a page's ``cursor``), by when they changed.""" + """Return the messages that are new or changed since this cursor, in the order they changed.""" class _Messages(TypedDict): @@ -89,34 +90,37 @@ class _Messages(TypedDict): class Messages(_Messages, total=False): - """ - A page of messages: oldest first; or, with a ``cursor``, those new or - changed since it, by when they changed. Keep them by ``id``: a message - comes again once it changed (its tool call's output came, say). ``more`` - says whether there are more (older ones, or with a ``cursor``, newer - ones); ``cursor`` is the one to read on from. + """The response of the ``messages`` view: a page of messages. + + Without a ``cursor`` in the request, the messages are oldest first. With + one, they are the messages that are new or changed since it, in the order + they changed. Store messages by ``id``: a message comes again whenever it + changes, for example when its tool call's output arrives. + + ``more`` says whether there are more messages (older ones, or with a + ``cursor``, newer ones), and ``cursor`` is the cursor to read on from. """ first: int - """The id of the conversation's first message, once it was cleared: a client drops the messages it kept from before.""" + """The id of the first message after the conversation was cleared; a client drops the messages it kept from before.""" class Turn(TypedDict): - """The message a ``chat`` handler is answering: its invocation, and when it started.""" + """The message a ``chat`` handler is answering: its invocation and when it started.""" invocationId: str - """The invocation's id: ``/restate/invocation//attach`` waits for its answer.""" + """The invocation's id; ``/restate/invocation//attach`` waits for its answer.""" since: str - """When it started, as an ISO date.""" + """When it started, as an ISO 8601 timestamp.""" def text_part(text: str) -> TextPart: - """A text part.""" + """Return a text part with ``text``.""" return {"type": "text", "text": text} def is_tool_part(part: Part) -> typing.TypeGuard[ToolPart]: - """Whether ``part`` is a tool call's.""" + """Return whether ``part`` is a tool call.""" return part["type"] == "tool" @@ -127,38 +131,38 @@ class _Waiting(TypedDict): class Waiting(_Waiting, total=False): - """ - A wait for a signal, as ``pending`` lists it: its ``id``, which ``signal`` - takes to resolve this wait and no other; the ``name`` of the signal it - waits for (such as ``approval``); and ``since``, when it started, as an - ISO date. + """A wait for a signal, as the ``pending`` view lists it. + + ``id`` identifies the wait, so that ``signal`` can resolve this wait and + no other; ``name`` is the signal it waits for (such as ``"approval"``); + and ``since`` is when it started, as an ISO 8601 timestamp. """ input: Any - """What it waits for, as the model said it (such as the refund to approve), if the tool takes an input.""" + """What it waits for, as the model described it (such as the refund to approve), if the tool takes an input.""" class Signal(TypedDict, total=False): - """Which wait a signal resolves, and with what.""" + """The request of the ``signal`` view: which wait to resolve, and with what value.""" id: str - """The wait of this id (from ``pending``).""" + """Resolve the wait with this id, as ``pending`` lists it.""" name: str - """Else the oldest wait for the signal of this name.""" + """Without an ``id``, resolve the oldest wait for the signal of this name.""" value: Any - """Its value: what the wait answers the model with.""" + """The value, which the tool returns to the model.""" keep: bool - """With a ``name`` and no one waiting for it, keep the signal for the next wait for it (by default it is dropped).""" + """If nothing waits for ``name`` yet, keep the signal for the next wait instead of dropping it.""" class Resolved(TypedDict): - """What ``signal`` did: the id of the wait it resolved, or None if none waited (and it was dropped, or kept).""" + """The response of the ``signal`` view: the id of the wait it resolved, or ``None`` if nothing was waiting.""" resolved: str | None class Notifications(TypedDict): - """The revision of an agent's conversation: each change of what a client reads increases it.""" + """The revision of an agent's conversation, which increases with every change a client can read.""" revision: int @@ -168,19 +172,19 @@ class _Watch(TypedDict): class Watch(_Watch, total=False): - """What ``watch`` waits for: a revision newer than ``revision``, for at most ``timeoutSeconds`` (default 50, at most 60).""" + """The request of the ``watch`` view: wait for a revision newer than ``revision``, for at most ``timeoutSeconds`` (default 50, at most 60).""" timeoutSeconds: float class Status(TypedDict): - """What an agent is doing: the message it answers, if any, and what its tools wait for, oldest first.""" + """What an agent is doing: the message it is answering, if any, and what its tools wait for, oldest first.""" answering: Turn | None waiting: list[Waiting] class Stopped(TypedDict): - """What ``stop`` did: the invocation of the message it stopped, or None if none was answered.""" + """The response of the ``stop`` view: the invocation it cancelled, or ``None`` if no message was being answered.""" stopped: str | None diff --git a/python/restate/agents/_calls.py b/python/restate/agents/_calls.py index b3e6952c..b3cff1e9 100644 --- a/python/restate/agents/_calls.py +++ b/python/restate/agents/_calls.py @@ -8,22 +8,24 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -The Restate calls of an invocation's tools (``call``, ``signal``, -``subagent``...), kept in the object's database (``restate_agent_calls``) so -that a call is made once, however often its tool call runs. As sdk-sessions' -``calls.ts``. - -A crash or a suspension restarts the handler. A framework then runs a tool call -that was cut off again, with its id; or tells the model it was interrupted, and -the model makes it again with a new id: made with the same tool and input, it -is the call an earlier attempt made, so it takes that call's id, and so the -same idempotency key: the same Restate call, which it waits for again, or -whose answer it gets. - -A call is kept until the framework's transcript has its result under the last -id it was made with (``settle``, at each publish, and before each tool call). -An error there is not its result, unless the call answered with that error. +"""Track the Restate calls that an invocation's tools make, so each is made only once. + +The tools built on Restate (``call``, ``send``, ``workflow``, ``subagent``, +``wait_for_signal``) record their calls in the object's database, in +``restate_agent_calls``, so that a call is made once however often its tool +call runs. + +A crash or a suspension restarts the handler. The framework then either runs +the interrupted tool call again with the same id, or tells the model it was +interrupted, and the model makes it again under a new id. Either way, a call +with the same tool and input as one an earlier attempt made reuses that +call's id, and so its idempotency key: it is the same Restate call, which it +waits for again or whose result it gets. + +A call is kept until the framework's transcript has its result under the +latest id it was made with (``settle`` checks this at each publish and before +each tool call). An error in the transcript does not count as the result, +unless the call itself returned that error. """ from __future__ import annotations @@ -50,7 +52,7 @@ class _Made(typing.TypedDict): - """A tool's call: the tool, its input as JSON, the ids it was made with (the first is its own), and its attempt.""" + """A tool's call: the tool, its input as JSON, the tool call ids it was made under (the first is its own), and its attempt.""" tool: str input: str @@ -59,24 +61,24 @@ class _Made(typing.TypedDict): class _Call(_Made, total=False): - """A tool's call, and whether it answered.""" + """A tool's call, and whether it has returned.""" done: bool - """Answered, until the transcript has its result.""" + """Whether the call has returned; kept until the transcript has its result.""" failed: str - """It answered with this error (``Error: ``), a terminal one: its result, as the transcript has it.""" + """The terminal error the call returned (``Error: ``), which is its result in the transcript.""" @dataclass class _Calls: - """The calls an invocation's tools made that the transcript has no result of yet.""" + """The calls an invocation's tools made whose results are not yet in the transcript.""" invocation: str calls: list[_Call] = field(default_factory=list[_Call]) async def _store(ctx: SessionContext, state: _Calls) -> None: - """Store ``state``, committed before the call it records is made.""" + """Store ``state``; it is committed before the call it records is made.""" await ctx.db.execute( "INSERT OR REPLACE INTO restate_agent_calls (row, invocation, calls) VALUES (0, ?, ?)", [state.invocation, json.dumps(state.calls)], @@ -84,7 +86,7 @@ async def _store(ctx: SessionContext, state: _Calls) -> None: async def _load(ctx: SessionContext, invocation: str) -> _Calls: - """The calls of ``invocation``, as stored: none if the row is another invocation's.""" + """Load the stored calls of ``invocation``; none if the row belongs to another invocation.""" await SCHEMA.ensure(ctx.db) row = await ctx.db.get("SELECT invocation, calls FROM restate_agent_calls") if row is None or row["invocation"] != invocation: @@ -102,7 +104,7 @@ class _Attempt: async def _calls_of(ctx: SessionContext) -> tuple[_Attempt, _Calls]: - """The attempt of ``ctx``, and the calls of its invocation.""" + """Return the current attempt and the calls of its invocation.""" attempt = _ATTEMPT.of(ctx, lambda: _Attempt(id=str(uuid.uuid4()))) if attempt.calls is None: attempt.calls = await _load(ctx, ctx.request.id) @@ -111,7 +113,7 @@ async def _calls_of(ctx: SessionContext) -> tuple[_Attempt, _Calls]: @dataclass class Call: - """The call a tool call makes: its id (an earlier attempt's, maybe), and how to say it answered.""" + """The call a tool call makes: its id (possibly from an earlier attempt), and a way to record that it returned.""" id: str _ctx: SessionContext @@ -119,7 +121,7 @@ class Call: _call: _Call async def answered(self, failed: str | None = None) -> None: - """It answered, with ``failed`` if a terminal error answered it: kept until the transcript has its result.""" + """Record that the call returned, with ``failed`` if it returned a terminal error.""" self._call["done"] = True if failed is not None: self._call["failed"] = failed @@ -127,10 +129,11 @@ async def answered(self, failed: str | None = None) -> None: async def call_of(ctx: SessionContext, tool: str, input: object, call_id: str) -> Call: - """ - The call of tool ``tool`` with ``input`` that tool call ``call_id`` makes: - the one it made before (a run again), or the one an earlier attempt made - with the same tool and input (made again under a new id), or else a new one. + """Return the call that tool call ``call_id`` makes to ``tool`` with ``input``. + + That is the call it made before (if it runs again), or else the call an + earlier attempt made with the same tool and input (if the model made it + again under a new id), or else a new call. """ attempt, state = await _calls_of(ctx) encoded = json.dumps(input, sort_keys=True) @@ -150,10 +153,11 @@ async def call_of(ctx: SessionContext, tool: str, input: object, call_id: str) - async def settle(ctx: SessionContext, stored: Callable[[str, str | None], bool]) -> None: - """ - Forget the calls that answered whose result the framework stored: those of - which ``stored(id, failed)`` is true for the last id they were made with, - ``failed`` the error that answered the call, if one did. + """Forget the calls that returned and whose results the framework has stored. + + A call's result is stored if ``stored(id, failed)`` is true, where ``id`` + is the latest id the call was made under and ``failed`` is the error it + returned, if any. """ _, state = await _calls_of(ctx) kept = [c for c in state.calls if not c.get("done") or not stored(c["ids"][-1], c.get("failed"))] diff --git a/python/restate/agents/_common.py b/python/restate/agents/_common.py index 0fe65593..fbb1a973 100644 --- a/python/restate/agents/_common.py +++ b/python/restate/agents/_common.py @@ -8,10 +8,11 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -What the frameworks share: their options (``PerSession``, the model), the ids of -an invocation's prompts, a model provider's refusal, and the ``chat`` message. -As sdk-sessions' ``common.ts``. +"""Helpers shared by the framework integrations. + +This covers per-invocation options (``PerSession``) and model names, the ids +of an invocation's prompts, recognising a model provider's refusal, and how +the ``chat`` handler reads its message. """ from __future__ import annotations @@ -32,9 +33,10 @@ T = typing.TypeVar("T") PerSession: typing.TypeAlias = T | Callable[[SessionContext], T | Awaitable[T]] -"""A value, or a function of the invocation's context (sync or async) that returns it. +"""A value, or a function (sync or async) that takes the invocation's context and returns it. -For options that are the same for every message, or that differ per user:: +Use a plain value for options that are the same for every message, and a +function for options that differ per user:: langgraph.agent("assistant", model="anthropic/claude-sonnet-5-5") langgraph.agent("assistant", system=lambda ctx: f"You help {ctx.key}.") @@ -42,7 +44,7 @@ async def configured(value: PerSession[T], ctx: SessionContext) -> T: - """``value`` for the invocation of ``ctx``.""" + """Return ``value`` for the invocation of ``ctx``, calling and awaiting it if it is a function.""" if not callable(value): return value made = typing.cast(Callable[[SessionContext], T | Awaitable[T]], value)(ctx) @@ -53,10 +55,11 @@ async def configured(value: PerSession[T], ctx: SessionContext) -> T: def next_prompt_id(ctx: SessionContext) -> str: - """ - The id of the next prompt of the invocation of ``ctx``: the invocation's - id, then ``:2``, ``:3``... A retried attempt makes its prompts in - the same order, so each gets the same id again, and its answer. + """Return the id of the invocation's next prompt. + + The first prompt's id is the invocation id, the next ones are ``:2``, + ``:3``, and so on. A retried attempt makes its prompts in the same + order, so each prompt gets the same id again, and finds its stored answer. """ n = (_PROMPTS.get(ctx) or 0) + 1 _PROMPTS.set(ctx, n) @@ -64,13 +67,13 @@ def next_prompt_id(ctx: SessionContext) -> str: MODEL_ENV = "RESTATE_AGENT_MODEL" -"""The environment variable of the model that agents use when given none.""" +"""The environment variable that names the model agents use when none is given.""" DEFAULT_MODEL = "openai/gpt-6-sol" -"""The default model of the frameworks that take any provider's, with OPENAI_API_KEY.""" +"""The default model of the frameworks that accept any provider; it needs ``OPENAI_API_KEY``.""" CLAUDE_MODEL = "anthropic/claude-sonnet-5-5" -"""An Anthropic model, as a framework that takes only Anthropic's defaults to (with ANTHROPIC_API_KEY).""" +"""The default model of the frameworks that accept only Anthropic models; it needs ``ANTHROPIC_API_KEY``.""" def model_of(model: str | None, fallback: str = DEFAULT_MODEL) -> str: @@ -79,11 +82,11 @@ def model_of(model: str | None, fallback: str = DEFAULT_MODEL) -> str: def provider_of(model: str) -> tuple[str, str]: - """Split ``"provider/model-id"`` (or the frameworks' ``"provider:model-id"``) into its provider and model id. + """Split ``"provider/model-id"`` (or ``"provider:model-id"``) into its provider and model id. - It splits at the first ``/``, else at the first ``:``, so model ids that - contain a colon (``"openai/ft:gpt-4o:acme"``) keep it. A name with neither - has no provider: ``("", model)``. + It splits at the first ``/``, else at the first ``:``, so a model id that + contains a colon (``"openai/ft:gpt-4o:acme"``) keeps it. A name with + neither has no provider, and returns ``("", model)``. """ for separator in ("/", ":"): provider, found, model_id = model.partition(separator) @@ -93,20 +96,21 @@ def provider_of(model: str) -> tuple[str, str]: REFUSED = 424 -""" -The error code of a message the model's provider refused (a missing or wrong -API key, an unknown model): Failed Dependency, as its HTTP status. +"""The error code of a message that the model's provider refused, such as for a wrong API key or an unknown model. + +It is the HTTP status 424 Failed Dependency. """ CHECK_MODEL = f"check the agent's model (its `model` option, or {MODEL_ENV}) and its provider's API key" -"""What to check when a model's provider refuses a request.""" +"""The advice in the error message when a model's provider refuses a request.""" class ModelError(TerminalError): - """ - An agent's model that it cannot take: not ``provider/model-id``, another - provider's, or one its framework does not know. Terminal (``REFUSED``), - wherever it is raised: a retry takes the same model. + """An error for a model name that an agent cannot use. + + The name is not ``provider/model-id``, belongs to another provider, or is + unknown to the framework. It is a terminal error with status + ``REFUSED``, since a retry would use the same model. """ def __init__(self, message: str) -> None: @@ -114,13 +118,19 @@ def __init__(self, message: str) -> None: def check_model_name(framework: str, model: str | None) -> None: - """Raise ``ModelError`` if ``model``, given explicitly to ``framework``'s agent, names no provider.""" + """Raise ``ModelError`` if ``model``, passed explicitly to a ``framework`` agent, does not name a provider.""" if model is not None and not re.fullmatch(r"[^/:\s]+[/:]\S+", model): raise ModelError(f'{framework}: the model "{model}" is not provider/model-id, such as {CLAUDE_MODEL}') def model_id_of(model: str, provider: str) -> str: - """The model id of ``model``, for a framework with only ``provider``'s models: ``provider/model-id``, or the id alone.""" + """Return the model id of ``model`` for a framework that only takes ``provider``'s models. + + ``model`` is either ``provider/model-id`` or the bare model id. + + Raises: + ModelError: If ``model`` names another provider, or no model id. + """ named, model_id = provider_of(model) if "/" in model or ":" in model else (provider, model) if named != provider or not model_id: from_env = f" (from {MODEL_ENV})" if os.environ.get(MODEL_ENV) == model else "" @@ -129,7 +139,7 @@ def model_id_of(model: str, provider: str) -> str: class MessageSerde(Serde[str]): - """A chat message: the request's body as text, or the string it is as JSON (``"Hi"``).""" + """The serde of a chat message: the request body as plain text, or a JSON string (``"Hi"``).""" def deserialize(self, buf: bytes) -> str | None: text = buf.decode() @@ -144,18 +154,18 @@ def serialize(self, obj: str | None) -> bytes: CANCELLED = "Cancelled: the turn was stopped before it finished." -"""What a cancelled turn answers, and each tool call it cut off: the model sees it in the conversation that goes on.""" +"""The answer of a cancelled turn, and the result of each tool call it cut off, which the model sees later in the conversation.""" UNANSWERED = "Unanswered: the model's provider refused the request." -""" -What a prompt answers that the model's provider refused for good: the prompt is -over, and the model sees it in the conversation that goes on, rather than -answer it with the next message. +"""The stored answer of a prompt that the model's provider refused for good. + +It ends the prompt, so that the model sees it in the rest of the conversation +instead of answering the refused prompt along with the next message. """ class ToolError(Exception): - """A tool's error, which fails the attempt (Restate retries it): the error of tool ``tool``.""" + """An error raised by tool ``tool``, which fails the attempt so that Restate retries it.""" def __init__(self, tool: str, cause: BaseException) -> None: super().__init__(f"the tool {tool} failed: {cause}") @@ -164,20 +174,21 @@ def __init__(self, tool: str, cause: BaseException) -> None: _TOOL_ERRORS: ScopeKey[weakref.WeakSet[BaseException]] = ScopeKey("from_tool") -"""The errors the attempt's framework tools (not ``tool()``'s) failed with: see ``from_tool``.""" +"""The errors that framework tools (not those made by ``tool()``) raised in the attempt; see ``from_tool``.""" def from_tool(ctx: SessionContext, error: BaseException) -> None: - """ - ``error`` is a tool's, in the attempt of ``ctx``: a framework's own tool - that failed, whose error is never taken for the model's provider refusing - (a tool that calls a model API may fail as one does). + """Record that a framework's own tool raised ``error`` in the attempt of ``ctx``. + + Such an error is never mistaken for a refusal by the model's provider, + even if it looks like one (a tool that calls a model API can fail the same + way). """ _TOOL_ERRORS.of(ctx, weakref.WeakSet[BaseException]).add(error) _FOR_GOOD = frozenset({"AuthenticationError", "PermissionDeniedError", "NotFoundError", "BadRequestError"}) -"""The model SDKs' errors (Anthropic's, OpenAI's) of a request that fails again however often it is made.""" +"""The names of the model SDKs' errors (Anthropic's, OpenAI's) for requests that would fail again on retry.""" _REFUSED_TEXT = re.compile( r"API key not found|no API key|Could not resolve authentication method|Provider is not configured" @@ -187,11 +198,11 @@ def from_tool(ctx: SessionContext, error: BaseException) -> None: _FOR_GOOD_STATUS = frozenset({400, 401, 403, 404}) -"""The HTTP statuses of a model request that fails again however often it is made (Pydantic AI's ``ModelHTTPError``'s).""" +"""The HTTP statuses of Pydantic AI's ``ModelHTTPError`` for requests that would fail again on retry.""" def _refused(error: BaseException) -> bool: - """Whether ``error`` is a model provider's refusal that a retry does not change: a missing key, an unknown model...""" + """Return whether ``error`` is a model provider's refusal that a retry would not change, such as a missing key.""" if isinstance(error, ToolError): return False if type(error).__name__ in _FOR_GOOD: @@ -202,7 +213,11 @@ def _refused(error: BaseException) -> bool: def refusal(error: BaseException) -> TerminalError | None: - """``error`` as a terminal error (``REFUSED``), if a retry does not change it: a model the agent cannot take, or its provider's refusal.""" + """Return ``error`` as a terminal error with status ``REFUSED`` if retrying would not help, else ``None``. + + That is the case for a model the agent cannot use, and for a refusal by + the model's provider. + """ if isinstance(error, ModelError): return error if isinstance(error, TerminalError) or not _refused(error): @@ -213,7 +228,7 @@ def refusal(error: BaseException) -> TerminalError | None: def refusal_of(ctx: SessionContext, error: BaseException) -> TerminalError | None: - """``refusal``, but not of an error a framework's tool failed with in the attempt of ``ctx`` (``from_tool``).""" + """Return ``refusal(error)``, except for errors that a framework's tool raised in the attempt (see ``from_tool``).""" failed_tool = _TOOL_ERRORS.get(ctx) if failed_tool is not None and error in failed_tool: return None @@ -221,7 +236,7 @@ def refusal_of(ctx: SessionContext, error: BaseException) -> TerminalError | Non def stored_answer(answer: str) -> str: - """A prompt's stored answer, as the prompt returns it: ``UNANSWERED`` fails as the prompt did then (``REFUSED``).""" + """Return a prompt's stored answer, or raise the refusal again if it is ``UNANSWERED``.""" if answer != UNANSWERED: return answer raise TerminalError(f"the model's provider refused the request: {CHECK_MODEL}", status_code=REFUSED) diff --git a/python/restate/agents/_log.py b/python/restate/agents/_log.py index 2dc92f72..fc91adb4 100644 --- a/python/restate/agents/_log.py +++ b/python/restate/agents/_log.py @@ -8,23 +8,24 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -An agent's conversation as clients read it: a log in the object's database, -kept in step with the framework's transcript each time the agent publishes a -change, and read by ``messages`` and ``sync``. As sdk-sessions' ``log.ts``, -whose tables it shares. +"""The conversation log that an agent's clients read. + +The log lives in the object's database. It is brought in step with the +framework's transcript each time the agent publishes a change, and the +``messages`` and ``sync`` views read it. Each message has an ``id`` that never changes, and a ``seq``: the sequence -number of its last change, which is the cursor. A message that grows in place -(a tool call's output comes) keeps its id and takes a new ``seq``; a -transcript replaced (compacted, reset) leaves its messages in the log, and the -messages after them are new ones. So a client reads from a cursor, and gets -each message that is new or changed since, once. - -``clear`` forgets the log with the rest of the conversation, but not where it -got to: the next conversation's ids and cursors go on from there (its start, a -row of ``restate_agent_log_start``), so a client following it reads on, and -drops the messages before ``first``. +number of its latest change, which serves as the cursor. A message that grows +in place (when a tool call's output arrives, for example) keeps its id and +gets a new ``seq``. When a transcript is replaced (compacted or reset), its +old messages stay in the log, and the replacement's messages are new ones. So +a client that reads from a cursor gets each new or changed message exactly +once. + +``clear`` deletes the log with the rest of the conversation, but remembers +where it got to, in a row of ``restate_agent_log_start``. The next +conversation's ids and cursors carry on from there, so a client that follows +the conversation keeps reading, and drops the messages before ``first``. """ from __future__ import annotations @@ -48,7 +49,7 @@ position INTEGER UNIQUE, message TEXT NOT NULL )""", - # Where the log starts, which `clear` sets: no row before. + # Where the log starts; there is no row until the first `clear`. f"""CREATE TABLE {TABLE}_start ( row INTEGER PRIMARY KEY CHECK (row = 0), seq INTEGER NOT NULL, @@ -60,32 +61,31 @@ @dataclass(frozen=True) class Start: - """Where the log starts: past the ids and cursors of the conversations ``clear`` forgot.""" + """Where the log starts: past the ids and cursors of the conversations that ``clear`` deleted.""" seq: int - """The cursor it starts at.""" + """The cursor the log starts at.""" id: int - """Its first message's id.""" + """The id of the log's first message.""" async def _start_of(db: DatabaseReader) -> Start | None: - """Where the log of ``db`` starts, if ``clear`` set it.""" + """Return where the log of ``db`` starts, or ``None`` if ``clear`` never set it.""" row = await db.get(f"SELECT seq, id FROM {TABLE}_start") return Start(seq=int(typing.cast(int, row["seq"])), id=int(typing.cast(int, row["id"]))) if row else None async def restart(tx: DatabaseWriter, start: Start) -> None: - """ - Start the log at ``start``, in the transaction ``tx`` of ``clear`` that - emptied the database: the next conversation's ids and cursors go on from - there. + """Restart the log at ``start``, inside the transaction of ``clear`` that emptied the database. + + The next conversation's ids and cursors carry on from ``start``. """ await tx.execute(f"INSERT OR REPLACE INTO {TABLE}_start (row, seq, id) VALUES (0, ?, ?)", [start.seq, start.id]) @dataclass class _Row: - """A message of the log's head, the transcript as last stored: its id, and the message, also as JSON.""" + """A message of the log's head (the transcript as last stored): its id, the message, and its JSON.""" id: int message: Said @@ -94,7 +94,7 @@ class _Row: @dataclass class _Head: - """The log's head as a session stored it last, with the last cursor and the next id.""" + """The log's head as last stored, with the latest cursor and the next id.""" rows: list[_Row] seq: int @@ -102,7 +102,7 @@ class _Head: _HEAD: ScopeKey[_Head] = ScopeKey("the log") -"""The log's head, as the attempt's last ``sync`` stored it: the next compares the transcript with it.""" +"""The log's head as the attempt's latest ``sync`` stored it; the next ``sync`` compares the transcript with it.""" def _json(message: Said) -> str: @@ -110,7 +110,7 @@ def _json(message: Said) -> str: async def _head_of(ctx: SessionContext) -> _Head: - """The log's head in the attempt of ``ctx``: as it stored it last, or else as the log has it.""" + """Return the log's head: as the attempt last stored it, or else as read from the database.""" cached = _HEAD.get(ctx) if cached is not None: return cached @@ -131,11 +131,11 @@ async def _head_of(ctx: SessionContext) -> _Head: async def sync(ctx: SessionContext, messages: list[Said]) -> None: - """Store the transcript's ``messages``, in order, in the log: those that changed since the last time, in one transaction.""" + """Store the transcript's ``messages`` in the log, writing only those that changed, in one transaction.""" db = ctx.db await SCHEMA.ensure(db) head = await _head_of(ctx) - # A message the transcript did not change is the one stored (see Transcript): only the others are compared, as JSON. + # A message the transcript did not change is the very object stored (see Transcript): compare only the others, as JSON. jsons: dict[int, str] = {} def json_of(position: int) -> str: @@ -165,7 +165,7 @@ def unchanged(position: int) -> bool: seq += 1 await tx.execute(f"UPDATE {TABLE} SET seq = ?, message = ? WHERE id = ?", [seq, row.json, row.id]) rows.append(row) - # The rest of the transcript as it was is replaced: its messages stay in the log, out of it. + # The rest of the old transcript is replaced: its messages stay in the log, but leave the transcript. await tx.execute(f"UPDATE {TABLE} SET position = NULL WHERE position >= ?", [grown]) for position in range(grown, len(messages)): row = _Row(id=next_id, message=messages[position], json=json_of(position)) @@ -174,18 +174,20 @@ def unchanged(position: int) -> bool: await tx.execute(f"INSERT INTO {TABLE} VALUES (?, ?, ?, ?)", [row.id, seq, position, row.json]) rows.append(row) except BaseException: - # Rolled back: the log is as it was, which the next sync reads again. + # Rolled back: the log is unchanged, and the next sync reads it again. _HEAD.delete(ctx) raise _HEAD.set(ctx, _Head(rows=rows, seq=seq, next=next_id)) async def page(db: DatabaseReader, request: MessagesRequest | None = None) -> Messages: - """ - A page of the log. With ``cursor``: the messages new or changed since it, - by when they changed, and the cursor to read on from. Without: the latest - ``limit`` messages (or those before ``before``, an id), oldest first, and - the cursor to follow the conversation from. + """Return a page of the log. + + With a ``cursor``, the page holds the messages that are new or changed + since it, in the order they changed, and the cursor to read on from. + Without one, it holds the latest ``limit`` messages (or the latest before + message id ``before``), oldest first, and the cursor to follow the + conversation from. """ request = request or {} limit = request.get("limit", 50) @@ -220,13 +222,13 @@ async def page(db: DatabaseReader, request: MessagesRequest | None = None) -> Me def _message(row: list[typing.Any]) -> ChatMessage: - """A row of the log, as ``messages`` gives it.""" + """Convert a row of the log to a message, as ``messages`` returns it.""" said: Said = json.loads(typing.cast(str, row[1])) return {"id": int(row[0]), "role": said["role"], "parts": said["parts"]} def _grown_from(now: Said, was: Said) -> bool: - """Whether ``now`` is ``was`` grown in place: the same role and parts, with more (a tool call's output, say).""" + """Return whether ``now`` is ``was`` grown in place: the same role and parts, with more added (such as a tool output).""" if was["role"] != now["role"] or len(now["parts"]) < len(was["parts"]): return False for part, other in zip(was["parts"], now["parts"], strict=False): diff --git a/python/restate/agents/_object.py b/python/restate/agents/_object.py index a13f175e..ab4c4a28 100644 --- a/python/restate/agents/_object.py +++ b/python/restate/agents/_object.py @@ -8,10 +8,10 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -Session objects set up for agents: ``AgentObject``, whose ``chat`` makes it an -agent. As sdk-sessions' ``sessionObject`` of ``@restatedev/agents`` -(``object.ts``). +"""Session objects set up for agents. + +``AgentObject`` adds the views every agent has; its ``chat()`` method turns it +into an agent that answers messages. """ from __future__ import annotations @@ -42,7 +42,7 @@ T = typing.TypeVar("T") A = typing.TypeVar("A", bound=Callable[[SessionContext, str], Awaitable[Any]]) -"""A chat turn's answer function: a message in, the answer out.""" +"""The function that answers a chat message: it takes the message and returns the answer.""" RETRIES = InvocationRetryPolicy( initial_interval=timedelta(milliseconds=50), @@ -50,25 +50,29 @@ max_interval=timedelta(seconds=10), max_attempts=100_000, ) -""" -An agent's retries, for weeks, since each attempt resumes from the database: -soon after a crash, then twice as long each time, up to 10 seconds apart (a -provider's rate limit, say, is not asked again twice a second). +"""The default retry policy of an agent's handlers. + +Each attempt resumes from the database, so an agent can keep retrying for +weeks. The first retry comes soon after a crash, and each wait doubles, up to +10 seconds, so that a provider that rate-limits is not asked again twice a +second. """ WAITS = timedelta(minutes=5) -""" -The inactivity timeout of the views that wait (``watch``, ``sync``, ``clear``): -Restate suspends one that waits past it, which restarts it, so it is longer -than the wait. +"""The inactivity timeout of the views that wait (``watch``, ``sync`` and ``clear``). + +Restate suspends a view that waits longer than this, which restarts it, so the +timeout is longer than any of their waits. """ class SyncRequest(TypedDict, total=False): - """ - What ``sync`` reads: what is new or changed since ``cursor`` (default 0, - all of it), at most ``limit`` messages (default 100); once the revision is - past ``revision``, if given, or at ``timeoutSeconds`` (default 50, at most + """The request of the ``sync`` view. + + ``sync`` returns the messages that are new or changed since ``cursor`` + (default 0, which means the whole conversation), at most ``limit`` of them + (default 100). If ``revision`` is given, it first waits until the revision + is newer than that, or until ``timeoutSeconds`` pass (default 50, at most 60). """ @@ -79,14 +83,14 @@ class SyncRequest(TypedDict, total=False): class Synced(Messages): - """What ``sync`` answers: the messages since the cursor, as ``messages`` gives them, the status and the revision.""" + """The response of the ``sync`` view: a page of messages, as ``messages`` returns it, plus the status and the revision.""" revision: int status: Status def _start_after(page: Messages) -> _log.Start: - """Where the log starts after the latest message of ``page``: past its ids and cursor.""" + """Return where the log restarts after the latest message of ``page``, past its ids and cursor.""" latest = page["messages"][0] if page["messages"] else None return _log.Start(seq=page["cursor"], id=latest["id"] + 1 if latest else page.get("first", 0)) @@ -94,26 +98,57 @@ def _start_after(page: Messages) -> _log.Start: BUILT_IN = frozenset( {"signal", "pending", "status", "stop", "notifications", "watch", "chat", "messages", "clear", "sync", "erase"} ) -"""The names of the handlers and views every agent has (``chat`` and the rest once it chats).""" +"""The names of the built-in handlers and views of every agent, which are reserved.""" class AgentObject(SessionObject): - """ - A Restate virtual object with a SQL database per key, ``ctx.db``, set up - for agents. Besides its own handlers and views, it has the views - ``status``, what it is doing (the message it answers, and what its tools - wait for), ``stop``, which stops that message, ``pending``, the waits of - its tools, ``signal({id, value})``, which resolves one, and - ``notifications`` and ``watch({revision})``, to follow its conversation. - - ``chat(answer)`` makes it an agent: a ``chat`` handler, a message in and - the answer out, the ``messages`` view (its conversation), ``clear``, which - forgets it (its database too), and ``sync``, which a client following it - reads. A framework's ``agent()`` makes one. - - ``migrations``: tables of your own, beside the agent's (named ``restate_...``). - ``background``: it works unwatched (a subagent, say): its turn publishes at - its end only, not at each step. + """A session object set up for agents. + + Like any session object, it is a Restate virtual object with a SQL + database per key (``ctx.db``), and it can have handlers and views of your + own. Every agent also has these built-in views: + + - ``status``: what the agent is doing: the message it is answering, if + any, and the signals its tools are waiting for. + - ``stop``: cancel the message being answered. + - ``pending``: list the signals its tools are waiting for, oldest first. + - ``signal``: resolve one of those waits with a ``value``, by its ``id`` + or by signal ``name``. + - ``notifications``: return the conversation's revision, which increases + with each change. + - ``watch``: wait for a revision newer than a given one. + + ``chat()`` turns the object into an agent, and adds: + + - ``chat``: a handler that takes a message and returns the answer. + - ``messages``: read the conversation, one page at a time. + - ``clear``: stop the message being answered and forget the + conversation, emptying the key's database. + - ``sync``: read what changed since a cursor, together with the status + and the revision, for clients that follow the conversation. + + These names, and ``erase`` (which ``clear`` uses), are reserved: + registering a handler or view of your own under one of them raises + ``ValueError``. You rarely create an ``AgentObject`` yourself: a + framework's ``agent()`` creates one and calls ``chat()`` for you. + + Args: + name: The object's name, as Restate serves it. + description: The object's description, shown in Restate's UI. + migrations: Tables of your own, beside the agent's (whose names + start with ``restate_``). + background: Whether the agent works unwatched (a subagent, for + example). A background agent publishes its progress only when a + turn ends, not after each step. + metadata: The object's metadata. + inactivity_timeout: Restate's inactivity timeout for its handlers. + abort_timeout: Restate's abort timeout for its handlers. + journal_retention: How long Restate keeps an invocation's journal. + idempotency_retention: How long Restate keeps the result of an + invocation that has an idempotency key. + ingress_private: Whether only other services may call it. + invocation_retry_policy: How Restate retries a failed handler. + Defaults to ``RETRIES``. """ def __init__( @@ -158,11 +193,11 @@ def _taken(self, name: str) -> str: return super()._taken(name) def _attempt_fn(self, fn: Callable[..., Awaitable[Any]], view: bool) -> Callable[..., Awaitable[Any]]: - """A handler's attempt: a model provider's refusal it fails with ends its invocation, as ``chat``'s does.""" + """Wrap a handler so that a model provider's refusal fails its invocation for good, as in ``chat``.""" if view: return fn - # With fn's signature, from which the SDK takes whether it has an input. + # Keep fn's signature: the SDK reads from it whether the handler takes an input. @functools.wraps(fn) async def answering(ctx: SessionContext, *args: Any) -> Any: try: @@ -176,33 +211,33 @@ async def answering(ctx: SessionContext, *args: Any) -> Any: return answering def _built_in(self) -> None: - """The views every agent has: each asks the companion, which knows what the object does.""" + """Register the views every agent has; each one asks the agent's companion ``signals`` object.""" async def signal(ctx: ViewContext, signal: Signal) -> Resolved: - """Resolve a wait of the object's tools: by its ``id``, or the oldest of a ``name``, with ``value``.""" + """Resolve a wait of the agent's tools with ``value``: the wait with ``id``, or the oldest one for ``name``.""" return await _companion(ctx, _signals.resolve, "signal")(signal) async def pending(ctx: ViewContext) -> list[Waiting]: - """The waits of the object's tools, oldest first.""" + """List the signals the agent's tools are waiting for, oldest first.""" return await _companion(ctx, _signals.pending)() async def status(ctx: ViewContext) -> Status: - """What the object is doing: the message it answers, if any, and what its tools wait for.""" + """Return what the agent is doing: the message it is answering, if any, and what its tools wait for.""" return await _companion(ctx, _signals.status)() async def stop(ctx: ViewContext) -> Stopped: - """Cancel the invocation of the message being answered (``status``'s), as Restate's admin API cancels one.""" + """Cancel the invocation that is answering the current message, if there is one.""" return await _stopped(ctx, "stop") async def notifications(ctx: ViewContext) -> Notifications: - """The revision of the conversation: each change of what ``messages`` and ``status`` give increases it.""" + """Return the conversation's revision, which increases whenever ``messages`` or ``status`` change.""" return await _companion(ctx, _signals.notifications)() async def watch(ctx: ViewContext, request: Watch) -> Notifications: - """ - The revision, once there is one newer than ``revision``, or at the - timeout (50 s by default, at most 60 s): a client then reads what - changed. It waits in Restate, holding no process. + """Wait for a revision newer than ``revision``, then return the current one. + + It returns anyway after ``timeoutSeconds`` (50 by default, at most + 60). The wait happens in Restate and holds no process. """ return await _companion(ctx, _signals.watch, "watch")(request) @@ -234,18 +269,30 @@ def chat( *, clear: Callable[[SessionContext], Awaitable[None]] | None = None, ) -> Any: - """ - Make this object an agent, which answers each message with ``answer``: - a framework's turn, which keeps the conversation in the log - (``transcribing``) and stops with ``stop_signal_of(ctx)``. ``clear``, - what the agent keeps beside its database, which ``clear`` empties too. + """Turn this object into an agent that answers each message with ``answer``. + + Call it directly or use it as a decorator. Besides the ``chat`` + handler, it adds the ``messages``, ``clear`` and ``sync`` views. - A message comes as the request's body, text or a JSON string:: + The message is the request's body, as plain text or as a JSON + string:: curl localhost:8080/assistant/alice/chat -d 'Hi' - A key answers one message at a time; one its model's provider refused - (a missing API key, an unknown model) answers with an error (424). + Each key answers one message at a time. A message that the model's + provider refuses (a missing API key, an unknown model) fails with + status 424 instead of being retried. + + Args: + answer: A framework's turn: it takes the handler's context and + the message, and returns the answer. It should keep the + conversation in the log with ``transcribing``, and stop when + ``stop_signal_of(ctx)`` stops. + clear: Clears whatever the agent keeps outside its database; the + ``clear`` view calls it after emptying the database. + + Returns: + ``answer`` itself. """ if answer is None: @@ -269,10 +316,10 @@ async def chat(ctx: SessionContext, text: str) -> Any: def _companion( ctx: ViewContext, handler: Callable[..., Awaitable[T]], step: str | None = None ) -> Callable[..., Awaitable[T]]: - """ - A call of the companion's of ``ctx``; with ``step``, a change, keyed by the - request and ``step``: a restarted view joins the same call. A read has no - key: asked again, it answers again. + """Return a function that calls ``handler`` on the companion of the agent of ``ctx``. + + With ``step``, the call gets an idempotency key made of the request id and + ``step``, so a restarted view joins the same call. Reads have no key. """ key = None if step is None else f"{ctx.request.id}:{step}" companion = _signals.signals_of(ctx) @@ -285,11 +332,10 @@ def call(*args: Any) -> Awaitable[T]: async def _stopped(ctx: ViewContext, step: str) -> Stopped: - """ - Cancel the invocation of the message being answered, if any, and say - which. The companion's ``status`` says which, with a key of the request and - ``step``: a restarted view cancels the same invocation again (which does - nothing), not the message that started once it was cancelled. + """Cancel the invocation answering the current message, if any, and return its id. + + The status is read with an idempotency key, so a restarted view cancels the + same invocation again (a no-op) rather than the next message's. """ turn = (await _companion(ctx, _signals.status, step)())["answering"] if turn is not None: @@ -298,24 +344,19 @@ async def _stopped(ctx: ViewContext, step: str) -> Stopped: def _followed(agent: AgentObject) -> None: - """ - An agent's views beside those of every object: ``messages``, its - conversation; ``clear``, which forgets it; and ``sync``, the conversation - since ``cursor``, what the agent is doing, and the revision, in one - request. And ``erase``, ``clear``'s work, which only its views call. - """ + """Register the views that ``chat()`` adds (``messages``, ``clear``, ``sync``), and the private ``erase``.""" async def messages(ctx: ViewContext, request: MessagesRequest | None = None) -> Messages: - """ - The conversation as clients read it, from its log: a page of it - (``{limit, before}``), or what is new or changed since a cursor - (``{cursor}``). It answers while a message is answered, with what was - stored so far. + """Read the conversation: a page of it, or what changed since a cursor. + + Send ``{"limit": n, "before": id}`` for a page, or ``{"cursor": c}`` + for the messages that are new or changed since ``c``. While a message + is being answered, it returns what has been stored so far. """ return await _log.page(ctx.db, request) async def erase(ctx: SessionContext) -> None: - """``clear``'s work, once the handlers before it are done: the key's database empties, and its signals'.""" + """Empty the key's database and its signals, once the handlers queued before it are done.""" start = _start_after(await _log.page(ctx.db, {"limit": 1})) async with ctx.db.transaction() as tx: await reset(tx) @@ -324,31 +365,32 @@ async def erase(ctx: SessionContext) -> None: await ctx.call(_signals.reset, key=_signals.signals_of(ctx), idempotency_key=f"{ctx.request.id}:reset") async def clear(ctx: ViewContext) -> Stopped: - """ - Forget the conversation: stop the message being answered (as - ``stop``), then empty the key's database, and what its ``signals`` - hold. The log's ids and cursors go on, so a follower reads on, and - drops the messages before ``first``. A message sent before it is - answered first, then forgotten too. + """Forget the conversation, and empty the key's database. + + It stops the message being answered, as ``stop`` does. Messages sent + before the ``clear`` are still answered first, then forgotten too. + Message ids and cursors carry on from where they were, so a client + that follows the conversation keeps reading, and drops the messages + before ``first``. """ stopped = await _stopped(ctx, "clear") await ctx.call(erase, key=ctx.key, idempotency_key=f"{ctx.request.id}:erase") return stopped async def sync(ctx: ViewContext, request: SyncRequest | None = None) -> Synced: - """ - The conversation since ``cursor`` (by default, all of it), what the - agent is doing, and the revision; with ``revision``, once there is a - newer one, or at the timeout. A client follows the conversation by - sending each answer's ``revision`` and ``cursor`` back, at once while - it has ``more``. + """Return the conversation since ``cursor``, what the agent is doing, and the revision. + + With ``revision``, it first waits for a newer revision, or for the + timeout. A client follows the conversation by sending back the + ``revision`` and ``cursor`` of each response, right away while + ``more`` is true. """ request = request or {} revision = request.get("revision") if revision is not None: watch: Watch = {"revision": revision, "timeoutSeconds": request.get("timeoutSeconds", 50)} await _companion(ctx, _signals.watch, "watch")(watch) - # The revision first: a change after it shows in the next sync. + # Read the revision first, so that any change after it shows up in the next sync. now = await _companion(ctx, _signals.notifications)() status = await _companion(ctx, _signals.status)() page = await _log.page(ctx.db, {"cursor": request.get("cursor", 0), "limit": request.get("limit", 100)}) diff --git a/python/restate/agents/_signals.py b/python/restate/agents/_signals.py index 0fc3014b..f7d11881 100644 --- a/python/restate/agents/_signals.py +++ b/python/restate/agents/_signals.py @@ -8,16 +8,18 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -Signals to agents: a virtual object per agent, keyed ``:`` (such as -``assistant:alice``), whose ``wait`` returns once ``resolve`` gives the wait its -value; the message the agent answers; and the revision of the agent's -conversation, which its ``watch`` waits on. As sdk-sessions' ``signals.ts``, -with the same handlers. - -An agent's own handler cannot wait on an awakeable (it has no journal to replay -one from), so it calls ``wait``, which can: a journaled handler of this object, -which waits without holding a process. +"""The ``signals`` virtual object: each agent key's companion. + +The companion of an agent's key is keyed ``:`` (such as +``assistant:alice``). It holds the signals the agent's tools wait for (``wait`` +returns once ``resolve`` gives the wait its value), the message the agent is +answering, and the conversation's revision, which ``watch`` waits on. Its +handler names, including the camelCase ``awaitTurn``, are wire names shared +with Restate's other SDKs. + +An agent's own handler cannot wait on an awakeable, since it has no journal to +replay one from. Instead it calls ``wait``, a regular journaled handler of +this object, which waits without holding a process. """ from __future__ import annotations @@ -39,14 +41,14 @@ D = typing.TypeVar("D") signals = restate.VirtualObject("signals") -""" -The signals of every agent, keyed ``:`` (``signals_of``). -``sessions.app`` serves it beside the agents; its ingress-private handlers are -called by the agents' own handlers. +"""The companion object of every agent, keyed ``:`` (see ``signals_of``). + +``sessions.app`` serves it alongside the agents. Its ingress-private handlers +are called only by the agents' own handlers. """ WATCH = timedelta(hours=1) -"""How long one ``awaitTurn`` waits for its turn to end, before a new one takes over.""" +"""How long one ``awaitTurn`` invocation waits for its turn to end before handing over to a new one.""" class _WaitRequest(TypedDict): @@ -55,19 +57,19 @@ class _WaitRequest(TypedDict): class WaitRequest(_WaitRequest, total=False): - """A tool's wait: its id, the signal it waits for, and what for (the model's input), if anything.""" + """A tool's wait: its id, the signal it waits for, and optionally the model's input describing what for.""" input: Any class _WaitingOn(Waiting): - """A wait, with its awakeable.""" + """A wait, together with its awakeable.""" awakeable: str class _Watching(TypedDict): - """A ``watch`` waiting for a revision newer than ``revision``, by its awakeable.""" + """A ``watch`` waiting for a revision newer than ``revision``, identified by its awakeable.""" revision: int awakeable: str @@ -77,13 +79,13 @@ class _Watching(TypedDict): async def _get(ctx: Context, key: str, default: T) -> T: - """The state ``key``, or ``default`` while it has none.""" + """Return the state ``key``, or ``default`` if it is unset.""" value = await ctx.get(key, type_hint=object) return typing.cast(T, value) if value is not None else default def _public(wait: _WaitingOn) -> Waiting: - """A wait, as ``pending`` lists it: without its awakeable.""" + """Return a wait as ``pending`` lists it, without its awakeable.""" shown: Waiting = {"id": wait["id"], "name": wait["name"], "since": wait["since"]} if "input" in wait: shown["input"] = wait["input"] @@ -91,13 +93,13 @@ def _public(wait: _WaitingOn) -> Waiting: def iso(seconds: float) -> str: - """``seconds`` since the epoch as an ISO date, as JavaScript's ``toISOString`` gives it.""" + """Format ``seconds`` since the epoch as an ISO 8601 UTC timestamp with milliseconds, like JavaScript's ``toISOString``.""" moment = datetime.fromtimestamp(seconds, tz=timezone.utc) return moment.isoformat(timespec="milliseconds").replace("+00:00", "Z") async def _publish(ctx: ObjectContext) -> Notifications: - """Increase the revision, and answer each ``watch`` waiting for it: the new revision.""" + """Increase the revision, wake up every waiting ``watch``, and return the new revision.""" revision = await _get(ctx, "revision", 0) + 1 ctx.set("revision", revision) notified: Notifications = {"revision": revision} @@ -108,12 +110,11 @@ async def _publish(ctx: ObjectContext) -> Notifications: async def _closed(ctx: ObjectContext, turn: str) -> None: - """ - ``turn`` (an invocation id) is over: forget it, if the agent answers it, - and its tools' waits, which no one waits for any more, ending their - ``wait``s (a turn that ends as it should has none left, but a killed one - leaves them). It publishes the change; a close again changes nothing, and - publishes nothing. + """Close ``turn`` (an invocation id): forget it and its tools' waits, and publish the change. + + The waits are rejected, which ends their ``wait`` calls. A turn that ends + normally has no waits left, but a killed one can leave some behind. + Closing a turn twice changes nothing and publishes nothing. """ answering = await _get(ctx, "turn", typing.cast(Turn | None, None)) is_answered = answering is not None and answering["invocationId"] == turn @@ -132,7 +133,7 @@ async def _closed(ctx: ObjectContext, turn: str) -> None: @signals.handler(kind="shared") async def wait(ctx: ObjectSharedContext, request: WaitRequest) -> Any: - """Wait for the signal ``name``, as wait ``id`` (for ``input``), and return its value.""" + """Wait for the signal ``name`` as wait ``id``, and return its value.""" awakeable, value = ctx.awakeable(type_hint=object) since = iso(await ctx.time()) waiting: _WaitingOn = {"id": request["id"], "name": request["name"], "since": since, "awakeable": awakeable} @@ -142,7 +143,7 @@ async def wait(ctx: ObjectSharedContext, request: WaitRequest) -> Any: try: return await value except TerminalError as error: - # The agent's invocation was cancelled: no one waits any more. + # The agent's invocation was cancelled, so nothing waits any more. if error.status_code == 409: ctx.object_send(forget, ctx.key(), request["id"]) raise @@ -150,7 +151,7 @@ async def wait(ctx: ObjectSharedContext, request: WaitRequest) -> Any: @signals.handler() async def resolve(ctx: ObjectContext, signal: Signal) -> Resolved: - """Resolve the wait ``id``, or else the oldest wait for ``name``, with ``value``.""" + """Resolve the wait ``id``, or else the oldest wait for signal ``name``, with ``value``.""" wait_id, name = signal.get("id"), signal.get("name") if wait_id is None and name is None: raise TerminalError("a signal needs an id or a name") @@ -174,13 +175,13 @@ async def resolve(ctx: ObjectContext, signal: Signal) -> Resolved: @signals.handler(kind="shared") async def pending(ctx: ObjectSharedContext) -> list[Waiting]: - """The waits, oldest first.""" + """List the waits, oldest first.""" return [_public(wait) for wait in await _get(ctx, "waiting", typing.cast(list[_WaitingOn], []))] @signals.handler(ingress_private=True) async def listen(ctx: ObjectContext, waiting: _WaitingOn) -> None: - """A wait: resolved now with a signal kept for it, or else when one comes.""" + """Register a wait, resolving it right away if a signal was kept for it.""" kept = await _get(ctx, "kept", typing.cast(dict[str, list[Any]], {})) values = kept.get(waiting["name"], []) if values: @@ -193,11 +194,12 @@ async def listen(ctx: ObjectContext, waiting: _WaitingOn) -> None: @signals.handler(ingress_private=True) async def start(ctx: ObjectContext, turn: Turn) -> None: - """ - The agent's ``chat`` started answering ``turn``: the same invocation again - (a retry) keeps when it started. A turn before it that never closed itself - (its invocation was killed) closes now. ``awaitTurn`` watches the new turn, - so that it closes even if no turn comes after it. + """Record that the agent's ``chat`` started answering ``turn``. + + A retry of the same invocation keeps the original start time. A previous + turn that never closed itself (because its invocation was killed) is + closed now. ``awaitTurn`` then watches the new turn, so that it is closed + even if no other turn follows it. """ before = await _get(ctx, "turn", typing.cast(Turn | None, None)) if before is not None and before["invocationId"] == turn["invocationId"]: @@ -212,23 +214,23 @@ async def start(ctx: ObjectContext, turn: Turn) -> None: name="awaitTurn", kind="shared", ingress_private=True, - # It leaves nothing behind once done: a turn has one at a time. + # Keep nothing once it is done: there is one at a time per turn. journal_retention=timedelta(0), idempotency_retention=timedelta(0), ) async def await_turn(ctx: ObjectSharedContext, turn: str) -> None: - """ - Wait for ``turn`` to end, without holding a process, and close it if it - failed: its waits end, and ``status`` no longer says the agent answers it. - A turn that never closed itself failed (killed, out of retries, lost); one - that ended as it should, or was cancelled, closed itself already. - - It waits ``WATCH`` at most: a turn still answered then gets a new watch, - which runs on the deployment of the day. A cancelled watch (an - operator's) closes nothing. + """Wait for ``turn`` to end, without holding a process, and close it if it failed. + + Closing ends the turn's waits, and ``status`` stops reporting it. A turn + that ended normally or was cancelled has closed itself already; one that + did not (killed, out of retries, lost) is closed here. + + It waits at most ``WATCH``: if the turn is still running then, a new + watch takes over, which runs on the current deployment. Cancelling the + watch itself (an operator might) closes nothing. """ attached = ctx.attach_invocation(turn, type_hint=object) - # This watch's own cancellation raises here, and closes nothing; the turn's failure does not. + # Cancelling this watch raises here, and closes nothing; the turn failing does not raise here. completed, _ = await restate.wait_completed(attached, ctx.sleep(WATCH)) if attached in completed: try: @@ -243,7 +245,7 @@ async def await_turn(ctx: ObjectSharedContext, turn: str) -> None: @signals.handler(kind="shared") async def status(ctx: ObjectSharedContext) -> Status: - """What the agent does: the message it answers, if any, and the waits of its tools.""" + """Return what the agent is doing: the message it is answering, if any, and its tools' waits.""" return { "answering": await _get(ctx, "turn", typing.cast(Turn | None, None)), "waiting": [_public(wait) for wait in await _get(ctx, "waiting", typing.cast(list[_WaitingOn], []))], @@ -252,11 +254,11 @@ async def status(ctx: ObjectSharedContext) -> Status: @signals.handler(kind="shared") async def watch(ctx: ObjectSharedContext, request: Watch) -> Notifications: - """ - The revision, once there is one newer than ``revision``, or at the timeout - (50 s by default, at most 60 s), whichever comes first. It waits in - Restate, holding no process; a caller that retries it with the same - idempotency key joins the same wait. + """Wait for a revision newer than ``revision``, or the timeout, and return the current revision. + + The timeout is ``timeoutSeconds`` (50 by default, at most 60). The wait + happens in Restate and holds no process; a caller that retries with the + same idempotency key joins the same wait. """ awakeable, changed = ctx.awakeable(type_hint=object) watching: _Watching = {"revision": request["revision"], "awakeable": awakeable} @@ -271,13 +273,13 @@ async def watch(ctx: ObjectSharedContext, request: Watch) -> Notifications: @signals.handler(ingress_private=True) async def close(ctx: ObjectContext, turn: str) -> None: - """The message ``turn`` (an invocation id) is answered, or its turn ended otherwise: see ``_closed``.""" + """Close ``turn`` (an invocation id), which has been answered or ended otherwise.""" await _closed(ctx, turn) @signals.handler(ingress_private=True) async def reset(ctx: ObjectContext) -> None: - """The agent forgot its conversation (``clear``): forget its waits and kept signals, and publish the change.""" + """Forget all waits and kept signals after the agent's ``clear``, and publish the change.""" ctx.clear("waiting") ctx.clear("kept") await _publish(ctx) @@ -285,7 +287,7 @@ async def reset(ctx: ObjectContext) -> None: @signals.handler(ingress_private=True) async def forget(ctx: ObjectContext, wait_id: str) -> None: - """Forget wait ``wait_id``, whose agent no longer waits.""" + """Forget wait ``wait_id``, which the agent no longer waits for.""" waiting = await _get(ctx, "waiting", typing.cast(list[_WaitingOn], [])) ctx.set("waiting", [wait for wait in waiting if wait["id"] != wait_id]) await _publish(ctx) @@ -293,19 +295,19 @@ async def forget(ctx: ObjectContext, wait_id: str) -> None: @signals.handler(ingress_private=True) async def publish(ctx: ObjectContext) -> Notifications: - """The conversation changed: increase its revision, and answer the ``watch``es.""" + """Increase the conversation's revision after a change, and wake up every waiting ``watch``.""" return await _publish(ctx) @signals.handler(kind="shared") async def notifications(ctx: ObjectSharedContext) -> Notifications: - """The conversation's revision.""" + """Return the conversation's revision.""" return {"revision": await _get(ctx, "revision", 0)} @signals.handler(ingress_private=True) async def subscribe(ctx: ObjectContext, watching: _Watching) -> None: - """Answer ``awakeable`` once there is a revision newer than ``revision``: at once, if there is.""" + """Resolve the watch's awakeable once there is a revision newer than its ``revision``, right away if there is one.""" revision = await _get(ctx, "revision", 0) if revision > watching["revision"]: ctx.resolve_awakeable(watching["awakeable"], {"revision": revision}) @@ -315,7 +317,7 @@ async def subscribe(ctx: ObjectContext, watching: _Watching) -> None: @signals.handler(ingress_private=True) async def unsubscribe(ctx: ObjectContext, awakeable: str) -> None: - """Forget the ``watch`` of ``awakeable``, which no longer waits.""" + """Forget the ``watch`` with ``awakeable``, which no longer waits.""" watching = await _get(ctx, "watching", typing.cast(list[_Watching], [])) ctx.set("watching", [one for one in watching if one["awakeable"] != awakeable]) @@ -324,12 +326,12 @@ async def unsubscribe(ctx: ObjectContext, awakeable: str) -> None: def key_of(service: str, key: str) -> str: - """The key of the signals of agent ``service``'s key ``key``: ``:``.""" + """Return the companion's key for key ``key`` of agent ``service``: ``:``.""" return f"{service}:{key}" def signals_of(ctx: SessionContext | ViewContext) -> str: - """The key of the signals of the agent of ``ctx``.""" + """Return the companion's key for the agent key that ``ctx`` runs on.""" definition = definition_of(ctx) if definition is None: raise TypeError("signals_of takes the context of an agent's handler or view") @@ -338,14 +340,14 @@ def signals_of(ctx: SessionContext | ViewContext) -> str: @dataclass(frozen=True) class Companion: - """ - An agent's companion, for one of its keys: its ``signals`` object, which - knows what the agent is doing (the message it answers, what its tools wait - for), and how to reach it. + """The address of an agent key's companion in the ``signals`` object. + + The companion knows what the agent is doing: the message it is answering + and what its tools wait for. """ key: str - """The key of the agent's key there: ``:``.""" + """The companion's key in ``signals``: ``:``.""" agent: str """The agent's name.""" agent_key: str @@ -353,18 +355,22 @@ class Companion: _COMPANION: Link[Callable[[str], Companion]] = Link("restate-agents.companion") -"""How an agent reaches its companion, for one of its keys.""" +"""How an agent finds the companion of each of its keys.""" def with_companion(definition: D) -> D: - """``definition`` (an agent), pointing to its companion; ``sessions.app`` serves its ``signals`` object beside it.""" + """Attach the companion to ``definition`` (an agent), so that ``sessions.app`` serves ``signals`` alongside it.""" name = typing.cast(str, getattr(definition, "name")) # noqa: B009 (a definition of the SDK's) _COMPANION.on(definition, lambda key: Companion(key=key_of(name, key), agent=name, agent_key=key)) return needing(definition, lambda: [signals]) def companion_of(agent: object, key: str) -> Companion: - """The companion of ``agent``'s key ``key``: for anyone holding an agent to reach what it is doing.""" + """Return the companion of ``agent``'s key ``key``, to find out what that agent is doing. + + Raises: + TypeError: If ``agent`` is not a ``restate.agents`` agent. + """ companion = _COMPANION.of(agent) if companion is None: raise TypeError(f"{agent!r} is not an agent of restate.agents") @@ -372,18 +378,21 @@ def companion_of(agent: object, key: str) -> Companion: def has_companion(agent: object) -> bool: - """Whether ``agent`` is an object of restate.agents, with a companion.""" + """Return whether ``agent`` is a ``restate.agents`` agent, with a companion.""" return _COMPANION.of(agent) is not None def deliver(ctx: SessionContext, agent: object, key: str, text: str, *, idempotency_key: str | None = None) -> None: - """ - Deliver ``text`` to ``agent``'s key ``key``, from a session handler: a - message of its own (as ``chat``), which it answers after the messages - before it. It is sent, not awaited: ``idempotency_key`` makes it - delivered once:: + """Send ``text`` to ``agent``'s key ``key`` as a chat message, from a session object's handler. + + The agent answers it after any messages queued before it. The message is + sent, not awaited; pass ``idempotency_key`` so that a retried attempt + delivers it only once:: + + deliver(ctx, triage, "ENG-1", "The issue was closed", idempotency_key=f"{ctx.request.id}:closed") - deliver(ctx, triage, "ENG-1", "The issue was closed", idempotency_key=delivery) + Raises: + TypeError: If ``agent`` is not a ``restate.agents`` agent. """ companion = companion_of(agent, key) # an agent of restate.agents, with a `chat` handler ctx.generic_send(companion.agent, "chat", text, key=key, idempotency_key=idempotency_key) diff --git a/python/restate/agents/_tools.py b/python/restate/agents/_tools.py index 8680c3cb..998189d3 100644 --- a/python/restate/agents/_tools.py +++ b/python/restate/agents/_tools.py @@ -8,19 +8,21 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -Tools for any of the frameworks: defined once with ``tool``, given to any -agent's ``tools``, which turns them into the framework's own. And the tools -that are Restate itself: calling a handler, sending to one, running a workflow, -handing a task to another agent, and waiting for a signal. As sdk-sessions' -``tools.ts``. - -A tool runs with the agent's context, and with the id of the model's tool -call: Restate calls carry an idempotency key made of the invocation id and that -id, so a tool call run again after a crash makes the same call, once. - -A tool's input schema comes from type hints, by pydantic, imported when a tool -is made: an agent framework's extra brings it. +"""Tools that work with any of the agent frameworks. + +Define a tool once with ``tool``, and pass it in any agent's ``tools``: the +integration converts it to the framework's own tool type. The other tools +here are built on Restate itself: ``call`` calls a handler, ``send`` sends to +one, ``workflow`` runs a workflow, ``subagent`` hands a task to another agent, +and ``wait_for_signal`` waits for a signal. + +A tool runs with the agent's context and the id of the model's tool call. The +Restate calls these tools make use an idempotency key built from the +invocation id and that id, so a tool call that runs again after a crash makes +the same call only once. + +A tool's input schema is derived from its type hints by pydantic, which is +imported when the tool is created. Each framework extra installs it. """ from __future__ import annotations @@ -44,11 +46,11 @@ T = typing.TypeVar("T") INJECTED = ("ctx", "call_id", "idempotency_key") -"""The parameters a tool's function gets from where it runs, rather than from the model.""" +"""The parameters a tool function gets from the runtime rather than from the model.""" def _pydantic() -> Any: - """pydantic, which tools take their inputs' schemas from.""" + """Import and return pydantic, or raise ``ImportError`` explaining how to install it.""" try: import pydantic except ImportError as error: @@ -61,29 +63,36 @@ def _pydantic() -> Any: @dataclass(frozen=True) class ToolCall: - """A tool call: the agent's context, and the id of the model's call (the same when a crash runs it again).""" + """A running tool call: the agent's context and the id of the model's tool call. + + The id stays the same when a crash makes the tool call run again. + """ ctx: SessionContext call_id: str @property def idempotency_key(self) -> str: - """ - The tool call's idempotency key, ``:``: for a - call the tool makes once, however often it runs (a Restate call's, or a - payment API's). + """The tool call's idempotency key, ``:``. + + Pass it to anything the tool must do only once however often it runs, + such as a Restate call or a payment API request. """ return f"{self.ctx.request.id}:{self.call_id}" @dataclass(frozen=True) class Tool: - """ - A tool of any framework: its ``name`` and ``description`` for the model, - its input's JSON schema, which a framework offers the model, and how it - runs: ``run(args, call)``, its result as text for the model, of the - model's arguments as they came (arguments that do not fit the input answer - with what is wrong, for the model to call it again). + """A tool that any framework's agent can use. + + Create one with ``tool``, ``call``, ``send``, ``workflow``, ``subagent`` + or ``wait_for_signal`` rather than directly. + + Attributes: + name: The tool's name, as the model sees it. + description: What the tool does, for the model. + input_schema: The JSON schema of the tool's input, which the + framework passes to the model. """ name: str @@ -93,22 +102,39 @@ class Tool: _validate: Callable[[dict[str, Any]], Any] async def run(self, args: object, call: ToolCall) -> str: - """The tool's result for the model, as text: see ``tool`` for how it fails.""" + """Run the tool with the model's raw arguments, and return its result as text for the model. + + Arguments that do not match the input schema are not an error: the + result then tells the model what is wrong, so it can call the tool + again. See ``tool`` for how errors are handled. + + Args: + args: The arguments, as the model gave them. + call: The tool call this run is for. + + Returns: + The result, or ``"Error: "`` for the model to act on. + + Raises: + ToolError: If the tool fails with an error other than + ``TerminalError``; Restate then retries the attempt. + """ try: arguments = self._validate(typing.cast(dict[str, Any], args if isinstance(args, dict) else {})) - except Exception as error: # pydantic's ValidationError: the model's to fix + except Exception as error: # pydantic's ValidationError, for the model to fix return f"Error: the arguments do not fit the tool's input: {error}" try: return _text(await self._run(arguments, call)) except BaseException as error: # The turn stopped (its invocation was cancelled, or its attempt - # ended): it ends, rather than the model going on with the tool's - # error as its answer (such as a call the cancellation cancelled). + # ended). End it, rather than letting the model carry on with the + # tool's error as the result (such as a call that the cancellation + # cancelled). stop = stop_signal_of(call.ctx) if stop is not None and stop.reason is not None: raise stop.reason from error if isinstance(error, TerminalError): - # A terminal error is the tool's answer, which the model can act on. + # A terminal error is the tool's result, which the model can act on. return f"Error: {error.message}" if not isinstance(error, Exception): raise @@ -117,10 +143,19 @@ async def run(self, args: object, call: ToolCall) -> str: async def run_tool(tool: Tool, args: object, call: ToolCall) -> str: - """ - Run ``tool`` for tool call ``call``, as a framework's integration runs one: - the conversation publishes its start and its end, for the clients that - follow it, and the calls whose results are stored are settled first. + """Run ``tool`` for a tool call, the way a framework integration should. + + Unlike ``Tool.run``, it publishes the tool call's start and end to the + clients that follow the conversation, and first settles the earlier calls + whose results the framework has stored. + + Args: + tool: The tool to run. + args: The arguments, as the model gave them. + call: The tool call this run is for. + + Returns: + The tool's result as text, as ``Tool.run`` returns it. """ await tool_call_started(call.ctx) try: @@ -130,7 +165,7 @@ async def run_tool(tool: Tool, args: object, call: ToolCall) -> str: def _text(value: object) -> str: - """A result as text: a string as it is, nothing as "done", anything else as JSON.""" + """Convert a result to text: a string as is, ``None`` as ``"done"``, anything else as JSON.""" if isinstance(value, str): return value if value is None: @@ -155,22 +190,40 @@ def tool(*, name: str | None = None, description: str | None = None) -> Callable def tool( fn: Callable[..., Any] | None = None, *, name: str | None = None, description: str | None = None ) -> Tool | Callable[[Callable[..., Any]], Tool]: - """ - A tool, from a function: the model gives its parameters, described by - their types and the docstring:: + """Create a tool from a function, sync or async. + + The model supplies the function's arguments; their type hints and the + docstring describe the tool to the model:: @tool async def get_weather(city: str) -> str: '''The current weather in a city.''' return f"Sunny in {city}" - Three parameters are not the model's: ``ctx``, the agent's context - (``ctx.db``, calls, sends); ``call_id``, the tool call's id; and - ``idempotency_key``, made of both, for what the tool must do once. A crash - runs an unfinished tool call again with the same id. A - ``restate.TerminalError`` the function raises (or a Restate call it makes - fails with) is its answer, ``Error: ``, for the model to act on; - any other error fails the attempt, which Restate retries. + Three parameter names are reserved, and filled in by the runtime instead + of the model: ``ctx``, the agent's ``SessionContext`` (for ``ctx.db``, + calls and sends); ``call_id``, the tool call's id; and + ``idempotency_key``, built from the invocation id and ``call_id``, for + whatever the tool must do only once. A crash runs an unfinished tool call + again with the same id. + + If the function raises ``restate.TerminalError`` (or a Restate call it + makes fails with one), the model gets ``"Error: "`` as the + result, and can act on it. Any other exception fails the attempt, and + Restate retries it. + + Args: + fn: The function. Omit it to use ``tool(name=..., description=...)`` + as a decorator. + name: The tool's name. Defaults to the function's name. + description: What the tool does, for the model. Defaults to the + function's docstring. + + Returns: + The tool, or a decorator that creates it. + + Raises: + ImportError: If pydantic is not installed. """ if fn is None: @@ -210,16 +263,16 @@ def validate(args: dict[str, Any]) -> dict[str, Any]: async def _once(name: str, input: object, call: ToolCall, body: Callable[[ToolCall], Awaitable[T]]) -> T: - """ - ``body``, a Restate call of tool ``name`` with ``input``, made once however - often the tool call runs (see _calls.py): it gets the id of the call, which - an earlier attempt may have made, and so its idempotency key. + """Run ``body``, a Restate call of tool ``name``, so that it is made once however often the tool call runs. + + ``body`` gets the id under which the call was first made, possibly by an + earlier attempt, and so the same idempotency key (see ``_calls.py``). """ made = await call_of(call.ctx, name, input, call.call_id) try: result = await body(ToolCall(call.ctx, made.id)) except TerminalError as error: - # A terminal error answers the call; any other fails the attempt, and the retry makes it again. + # A terminal error is the call's result; any other error fails the attempt, and the retry calls again. await made.answered(f"Error: {error.message}") raise await made.answered() @@ -227,12 +280,12 @@ async def _once(name: str, input: object, call: ToolCall, body: Callable[[ToolCa def _untyped(ctx: SessionContext) -> Any: - """``ctx``, to call and send to a handler whose input a tool knows only at run time.""" + """Return ``ctx`` untyped, to call a handler whose input type is known only at run time.""" return ctx def _inputs(handler: Handler[Any, Any], value: Any) -> tuple[Any, ...]: - """The positional input of a call of ``handler``: ``value``, if it takes an input.""" + """Return the positional arguments for calling ``handler``: ``(value,)`` if it takes an input.""" return (value,) if handler.arity >= 2 else () @@ -244,7 +297,7 @@ def _handler(target: Callable[..., Any]) -> Handler[Any, Any]: def _input_type(handler: Handler[Any, Any]) -> Any: - """The type of a handler's input, or None if it takes none.""" + """Return the type of a handler's input, or ``None`` if it takes none.""" if handler.arity < 2: return None parameters = list(inspect.signature(handler.fn, eval_str=True).parameters.values()) @@ -254,17 +307,18 @@ def _input_type(handler: Handler[Any, Any]) -> Any: @dataclass(frozen=True) class _Input: - """A Restate tool's input: the handler's (flattened if an object), and the key if the model chooses it.""" + """The input of a Restate tool: the handler's input (flattened if it is an object), and the key if the model picks it.""" schema: dict[str, Any] validate: Callable[[dict[str, Any]], tuple[str | None, Any]] def _input_of(input_type: Any, model_keys: str | None) -> _Input: - """ - The tool's input: ``input_type``'s, with ``key`` beside it if the model - chooses the key of ``model_keys`` (an object's, a workflow's). An object's - fields are the tool's own; any other input is the tool's field ``input``. + """Build a Restate tool's input from a handler's ``input_type``. + + If ``model_keys`` is given, the model also picks the key, in a ``key`` + field. An object type's fields become the tool's own fields; any other + type becomes a single ``input`` field. """ pydantic = _pydantic() adapter = pydantic.TypeAdapter(input_type) if input_type is not None else None @@ -310,13 +364,28 @@ def call( name: str | None = None, key: PerSession[str] | None = None, ) -> Tool: - """ - A tool that calls handler ``target`` and answers with its result:: + """Create a tool that calls a handler and returns its result. + + The model fills in the handler's input, and for a virtual object or a + workflow it also picks the key, unless ``key`` is given:: call(order_status, description="The status of an order.") - The model fills in the handler's input; of a virtual object, it chooses the - key too, unless ``key`` gives it. + The call is made once per tool call, even if the tool call runs again + after a crash. + + Args: + target: A handler of a service, a virtual object or a workflow. + description: What the tool does, for the model. + name: The tool's name. Defaults to ``"_"``. + key: The key to call, or a function of the context that returns it. + If omitted, the model picks it. + + Returns: + The tool, for any framework's agent. + + Raises: + TypeError: If ``target`` is not a handler. """ handler = _handler(target) service = handler.service_tag.name @@ -352,7 +421,25 @@ def send( key: PerSession[str] | None = None, delay: timedelta | None = None, ) -> Tool: - """A tool that sends to handler ``target``, without waiting for it: a one-way call, made once, after ``delay``.""" + """Create a tool that sends to a handler without waiting for it to run. + + Like ``call``, but a one-way call: the model gets a confirmation instead + of the handler's result. The send is made once per tool call. + + Args: + target: A handler of a service, a virtual object or a workflow. + description: What the tool does, for the model. + name: The tool's name. Defaults to ``"_"``. + key: The key to send to, or a function of the context that returns + it. If omitted, the model picks it. + delay: How long to wait before the handler runs. + + Returns: + The tool, for any framework's agent. + + Raises: + TypeError: If ``target`` is not a handler. + """ handler = _handler(target) service = handler.service_tag.name model_keys = service if _keyed(handler) and key is None else None @@ -390,11 +477,22 @@ def workflow( name: str | None = None, key: PerSession[str] | None = None, ) -> Tool: - """ - A tool that runs the workflow of ``target`` (its main handler) to the end, - and answers with its result. Each tool call is a workflow of its own, with - an id made of the agent's key and the tool call's id, unless ``key`` gives - one. + """Create a tool that runs a workflow to completion and returns its result. + + Each tool call starts a workflow of its own, whose id is built from the + agent's key and the tool call's id, unless ``key`` is given. + + Args: + target: The workflow's main handler. + description: What the tool does, for the model. + name: The tool's name. Defaults to the workflow's name. + key: The workflow id, or a function of the context that returns it. + + Returns: + The tool, for any framework's agent. + + Raises: + TypeError: If ``target`` is not a handler. """ handler = _handler(target) service = handler.service_tag.name @@ -421,7 +519,7 @@ async def made(c: ToolCall) -> Any: def _split(shape: _Input, args: dict[str, Any]) -> dict[str, Any]: - """The tool call's key and the handler's input, validated, and the raw arguments (what makes the call the same).""" + """Validate the arguments into the key and the handler's input, keeping the raw arguments to identify the call.""" key, value = shape.validate(args) return {"key": key, "value": value, "raw": args} @@ -436,16 +534,25 @@ def subagent( name: str | None = None, handler: str = "chat", ) -> Tool: - """ - A tool, named after ``agent`` (an agent, or any virtual object) unless - ``name`` names it, that hands a task to it: to the object keyed - ``.``, through its ``handler`` (``chat``), and answers - with its answer. Tasks of different names run at once; the same name - continues the same task. The model reads what the agent is from its - description:: + """Create a tool that hands a task to another agent and returns its answer. + + The model names each task, and the task goes to the agent's key + ``.``. Tasks with different names run concurrently; + reusing a name continues the same task. The model learns what the agent + does from its description:: coder = langgraph.agent("coder", description="Builds and tests code in Python") lead = langgraph.agent("lead", tools=[subagent(coder)]) + + Args: + agent: An agent, or any virtual object with a handler that takes text. + description: What the agent does, for the model. Defaults to the + agent's own description. + name: The tool's name. Defaults to the agent's name. + handler: The handler that takes the task. + + Returns: + The tool, for any framework's agent. """ agent_name = typing.cast(str, getattr(agent, "name")) # noqa: B009 (an SDK's definition) tag = getattr(agent, "service_tag", None) diff --git a/python/restate/agents/_transcript.py b/python/restate/agents/_transcript.py index 7a88699a..c2679c98 100644 --- a/python/restate/agents/_transcript.py +++ b/python/restate/agents/_transcript.py @@ -8,10 +8,7 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -A conversation as ``messages`` gives it, built from a framework's transcript. -As sdk-sessions' ``Transcript`` (``common.ts``). -""" +"""Build the conversation that clients read from a framework's transcript.""" from __future__ import annotations @@ -22,29 +19,29 @@ class Transcript: - """ - A conversation as ``messages`` gives it, built from a framework's - transcript: each message is started by ``say``, a tool call goes to the - assistant's message, and its result to its call. Messages without parts + """A conversation in the form the ``messages`` view returns, built from a framework's transcript. + + ``say`` starts each message, ``call`` adds a tool call to the assistant's + message, and ``result`` adds the call's output. Messages without parts are left out. - A message, once given out by ``messages()``, is never changed: a change - makes a new one in its place. So a framework can keep its transcript and - add to it as its conversation grows, and the log stores only the messages - that are not the ones it stored last (see ``_log.py``). + A message returned by ``messages()`` is never mutated: a change replaces + it with a new one. So a framework can keep its ``Transcript`` and add to + it as the conversation grows, and the log stores only the messages that + are not the same objects it stored last time (see ``_log.py``). """ def __init__(self) -> None: self._messages: list[Said] = [] - # The message of each tool call, by the call's id: its index. + # The index of each tool call's message, by the call's id. self._calls: dict[str, int] = {} def say(self, role: Role, text: str = "") -> None: - """A new message of ``role``, with ``text`` if there is any.""" + """Start a new message of ``role``, with ``text`` if it is not blank.""" self._messages.append({"role": role, "parts": [text_part(text)] if text.strip() else []}) def call(self, call_id: str, name: str, input: Any) -> None: - """Tool call ``call_id`` of the assistant's message.""" + """Add tool call ``call_id`` to the assistant's message, starting one if needed.""" if not self._messages or self._messages[-1]["role"] != "assistant": self.say("assistant") at = len(self._messages) - 1 @@ -53,7 +50,11 @@ def call(self, call_id: str, name: str, input: Any) -> None: self._calls[call_id] = at def result(self, call_id: str, output: str, error: bool = False) -> None: - """The output of tool call ``call_id``; an error's if the framework says so, or it is a tool's ``Error: ``.""" + """Set the output of tool call ``call_id``. + + The output counts as an error if ``error`` is true or it starts with + ``"Error: "``. + """ at = self._calls.get(call_id) if at is None: return @@ -74,7 +75,11 @@ def with_output(parts: list[Part]) -> list[Part]: self._change(at, with_output) def answered(self, call_id: str, failed: str | None = None) -> bool: - """Whether tool call ``call_id`` has its output: not an error's, unless it is ``failed``, the error the call answered with.""" + """Return whether tool call ``call_id`` has its output. + + An error output does not count, unless it is ``failed``: the error + the call itself returned. + """ at = self._calls.get(call_id) if at is None: return False @@ -87,10 +92,10 @@ def answered(self, call_id: str, failed: str | None = None) -> bool: ) def messages(self) -> list[Said]: - """Its messages, in order: those with parts.""" + """Return the messages in order, leaving out those without parts.""" return [message for message in self._messages if message["parts"]] def _change(self, at: int, change: Callable[[list[Part]], list[Part]]) -> None: - """Message ``at`` with its parts as ``change`` makes them: a new message, in its place.""" + """Replace message ``at`` with a new message whose parts ``change`` computes.""" message = self._messages[at] self._messages[at] = {"role": message["role"], "parts": change(message["parts"])} diff --git a/python/restate/agents/_turn.py b/python/restate/agents/_turn.py index 62f2decd..4642f841 100644 --- a/python/restate/agents/_turn.py +++ b/python/restate/agents/_turn.py @@ -8,11 +8,11 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -The turn of a ``chat`` handler: its status, its cancellation, and publishing -for the clients that follow the conversation, with the framework's transcript -kept in the log as it goes. As sdk-sessions' ``chatHandler`` and -``transcribing`` (``common.ts``). +"""The turn of a ``chat`` handler. + +A turn reports its status, handles its cancellation, and publishes each change +to the clients that follow the conversation, copying the framework's +transcript into the log as it goes (``transcribing``). """ from __future__ import annotations @@ -37,15 +37,16 @@ T = typing.TypeVar("T") Read = Callable[[], Awaitable[Transcript]] -"""A framework's transcript as it stands, while it answers: a new one, or the one it keeps and adds to.""" +"""A function that returns the framework's transcript as it stands, either new or one the framework keeps adding to.""" class AttemptEnded(Exception): - """ - Why a turn stopped when its attempt ended before it did: Restate suspended - the attempt, gave up on it, or lost its connection. Nothing the attempt does - after is stored, and the next attempt goes on with the turn, so it stops - without answering: unlike a cancellation, nothing is cancelled. + """Raised to stop a turn whose attempt ended before the turn did. + + The attempt ends when Restate suspends it, gives up on it, or loses its + connection. Nothing the attempt does afterwards is stored, and the next + attempt carries on with the turn, so the turn stops without answering. + Unlike a cancellation, nothing is cancelled. """ def __init__(self) -> None: @@ -53,46 +54,47 @@ def __init__(self) -> None: def cancellation() -> TerminalError: - """What a cancelled turn fails with: Restate's cancellation, 409.""" + """Return the error a cancelled turn fails with: Restate's cancellation, status 409.""" return TerminalError("cancelled", status_code=409) class Stop: - """ - What stops a turn: a framework stops its work with it (a model's request, - say). It stops with Restate's cancellation (``409``) once Restate cancels - the invocation, the turn then ending as cancelled; or with an - ``AttemptEnded`` once the attempt ends first, the work stopping and - answering nothing. + """The stop signal of a turn, which a framework uses to cut its work short (a model request, for example). + + It stops with Restate's cancellation (status 409) when Restate cancels the + invocation, and the turn then ends as cancelled. It stops with + ``AttemptEnded`` when the attempt ends first, and the work then stops + without answering. """ def __init__(self) -> None: self._event = asyncio.Event() self.reason: BaseException | None = None - """Why it stopped, once it has.""" + """Why it stopped, or ``None`` while it has not.""" @property def stopped(self) -> bool: - """Whether it stopped.""" + """Whether it has stopped.""" return self.reason is not None def stop(self, reason: BaseException) -> None: - """Stop, with ``reason``; once.""" + """Stop with ``reason``, unless already stopped.""" if self.reason is None: self.reason = reason self._event.set() async def wait(self) -> BaseException: - """Return once it stopped: why.""" + """Wait until it stops, and return the reason.""" await self._event.wait() assert self.reason is not None return self.reason async def guard(self, work: Awaitable[T]) -> T: - """ - ``work``, unless it stops first: then ``work`` is cancelled, and this - raises why it stopped. For a framework's request that the turn's stop - must cut short. + """Await ``work`` and return its result, unless the turn stops first. + + If it stops first, ``work`` is cancelled and the stop reason is + raised. Use it for a framework request that stopping the turn must cut + short. """ task = asyncio.ensure_future(work) stopping = asyncio.ensure_future(self._event.wait()) @@ -113,20 +115,20 @@ async def guard(self, work: Awaitable[T]) -> T: def stop_signal_of(ctx: SessionContext) -> Stop | None: - """The ``Stop`` of the turn of ``ctx``, a ``chat`` handler's: see ``Stop``.""" + """Return the ``Stop`` of the ``chat`` turn running on ``ctx``, or ``None`` outside one.""" return _STOP.get(ctx) @dataclass(frozen=True) class AgentInfo: - """An agent's own: whether it works unwatched (``background``), and what its ``clear`` also empties.""" + """An agent's settings: whether it works unwatched (``background``), and what its ``clear`` also empties.""" background: bool = False clear: Callable[[SessionContext], Awaitable[None]] | None = None AGENT: Link[AgentInfo] = Link("an agent") -"""An agent's own, on its definition.""" +"""The agent's settings, attached to its definition.""" def _agent_of(ctx: SessionContext) -> AgentInfo | None: @@ -135,27 +137,27 @@ def _agent_of(ctx: SessionContext) -> AgentInfo | None: async def cleared(ctx: SessionContext) -> None: - """Clear what the agent of ``ctx`` keeps beside its database.""" + """Clear whatever the agent of ``ctx`` keeps outside its database.""" agent = _agent_of(ctx) if agent is not None and agent.clear is not None: await agent.clear(ctx) def _unfollowed(ctx: SessionContext) -> bool: - """Whether the agent of ``ctx`` runs in the background: it publishes the turn's end only.""" + """Return whether the agent of ``ctx`` runs in the background, publishing only when a turn ends.""" agent = _agent_of(ctx) return agent is not None and agent.background _READING: ScopeKey[Read] = ScopeKey("transcribing") -"""The transcript of the attempt's turn, while it answers.""" +"""The transcript reader of the attempt's turn, while it answers.""" _WRITING: ScopeKey[asyncio.Future[None]] = ScopeKey("the log's writes") -"""The attempt's last log write: they run one after another (tool calls in parallel publish at once).""" +"""The attempt's latest log write; writes run one after another, as parallel tool calls publish at the same time.""" def _logged(ctx: SessionContext, read: Read) -> Awaitable[None]: - """The log of the conversation of ``ctx`` takes its transcript as ``read`` gives it.""" + """Copy the transcript that ``read`` returns into the conversation's log, after any earlier write.""" before = _WRITING.get(ctx) async def write() -> None: @@ -164,7 +166,7 @@ async def write() -> None: await before transcript = await read() await _log.sync(ctx, transcript.messages()) - # The calls whose results the framework stored are made: see _calls.py. + # Forget the calls whose results the framework has stored: see _calls.py. await settle(ctx, transcript.answered) done = asyncio.ensure_future(write()) @@ -173,10 +175,11 @@ async def write() -> None: async def publish(ctx: SessionContext) -> None: - """ - The conversation of ``ctx`` changed: its log takes the change (while a - framework answers: ``transcribing``), its revision increases, and the - clients that follow it read it again. + """Publish a change to the conversation of ``ctx``. + + While a framework is answering (inside ``transcribing``), the change is + first copied into the log. Then the revision increases, and the clients + that follow the conversation read it again. """ read = _READING.get(ctx) if read is not None: @@ -185,16 +188,17 @@ async def publish(ctx: SessionContext) -> None: async def stepped(ctx: SessionContext) -> None: - """A step of a turn of ``ctx``: ``publish`` it, unless the agent runs in the background.""" + """Publish a step of the turn, unless the agent runs in the background.""" if not _unfollowed(ctx): await publish(ctx) async def tool_call_started(ctx: SessionContext) -> None: - """ - A tool call of the turn of ``ctx`` starts: ``publish`` it, which forgets - the calls whose results are stored before the tool makes its own. In the - background, it forgets them alone. + """Publish that a tool call of the turn is starting. + + Publishing also forgets the earlier calls whose results are stored, before + the tool makes calls of its own. A background agent only forgets them, + without publishing. """ if not _unfollowed(ctx): await publish(ctx) @@ -206,16 +210,26 @@ async def tool_call_started(ctx: SessionContext) -> None: async def tool_call_ended(ctx: SessionContext) -> None: - """A tool call of the turn of ``ctx`` ended: ``publish`` it, unless the agent runs in the background.""" + """Publish that a tool call of the turn has ended, unless the agent runs in the background.""" await stepped(ctx) async def transcribing(ctx: SessionContext, read: Read, work: Callable[[], Awaitable[T]]) -> T: - """ - Run ``work``, a framework's turn (or a compaction, a reset), with ``read`` - giving the framework's transcript as it stands: each publish meanwhile (a - tool call's start and end) takes it into the conversation's log, and so - does ``work``'s end. ``messages`` and ``sync`` read the log. + """Run ``work`` while keeping the conversation's log in step with the framework's transcript. + + ``work`` is a framework's turn (or a compaction, or a reset), and ``read`` + returns the framework's transcript as it stands. Every publish while + ``work`` runs (such as a tool call's start and end) copies the transcript + into the log, and so does the end of ``work``. The ``messages`` and + ``sync`` views read the log. + + Args: + ctx: The handler's context. + read: Returns the framework's current transcript. + work: The work to run. + + Returns: + What ``work`` returns. """ outer = _READING.get(ctx) _READING.set(ctx, read) @@ -227,7 +241,7 @@ async def transcribing(ctx: SessionContext, read: Read, work: Callable[[], Await raise finally: try: - # A failed attempt stores nothing more: its retry runs again. + # A failed attempt stores nothing more, as its retry runs again. if not failed: await _logged(ctx, read) finally: @@ -242,26 +256,26 @@ def _now() -> str: async def chat_turn(ctx: SessionContext, text: str, answer: Callable[[SessionContext, str], Awaitable[T]]) -> T: - """ - ``answer`` as a chat handler's turn. While it answers, the object's - ``status`` says so (its ``Turn``). Its invocation's cancellation (Restate's, - explicit) stops the turn (``stop_signal_of``) whenever it comes, and cancels - the calls its tools made: the turn ends, and the handler fails with the - cancellation. An attempt that ends first stops it too, with an - ``AttemptEnded``: the work stops, and the next attempt goes on with it. It - publishes when it starts and ends, for the clients that follow it. + """Run ``answer`` as the turn of a ``chat`` handler. + + While it answers, the object's ``status`` reports the turn. Cancelling the + invocation stops the turn (through ``stop_signal_of``) whenever it comes, + and cancels the calls its tools made; the handler then fails with the + cancellation. If the attempt ends first, the turn stops with + ``AttemptEnded`` and the next attempt carries on. The turn publishes when + it starts and when it ends. """ invocation_id = ctx.request.id companion = _signals.signals_of(ctx) turn: Turn = {"invocationId": invocation_id, "since": _now()} - # The companion knows the turn from here to its close: `status` and `stop` ask it. + # From here until the turn closes, the companion knows about it: `status` and `stop` ask it. await ctx.call(_signals.start, turn, key=companion) await stepped(ctx) async def end() -> None: - # The log takes the turn's last change, then the companion closes the - # turn, which publishes it: one revision. The close is a call the turn - # waits for, so that it is done before the turn ends. + # Copy the turn's last change into the log, then have the companion + # close the turn, which publishes both as one revision. The close is a + # call, awaited so that it is done before the turn ends. read = _READING.get(ctx) if read is not None: await _logged(ctx, read) @@ -283,12 +297,12 @@ async def on_attempt_end() -> None: try: try: answered = await answer(ctx, text) - # Cancelled once answered: the answer stands. + # A cancellation after the answer does not change it. except TerminalError: await end() raise except Exception as error: - # Any other error fails the attempt, and Restate retries the turn, but a refusal. + # Any other error fails the attempt and Restate retries the turn, unless it is a refusal. terminal = refusal_of(ctx, error) if terminal is None: raise diff --git a/python/restate/agents/internal.py b/python/restate/agents/internal.py index e383a962..57bd83b3 100644 --- a/python/restate/agents/internal.py +++ b/python/restate/agents/internal.py @@ -8,11 +8,19 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -What a framework's integration needs beside ``restate.agents``: its turn's -transcript kept in the log (``transcribing``, ``Transcript``), its stop -(``stop_signal_of``), its tool calls (``run_tool``), its model (``model_of`` -and the rest), and a model provider's refusal (``refusal_of``, ``from_tool``). +"""Building blocks for writing a framework integration of ``restate.agents``. + +An integration needs, beyond the public API: + +- ``transcribing`` and ``Transcript``, to keep the conversation's log in step + with the framework's transcript. +- ``stop_signal_of``, to stop a turn when it is cancelled or its attempt ends. +- ``run_tool``, to run a ``restate.agents`` tool for a tool call. +- ``model_of``, ``provider_of`` and the related helpers, to resolve the model. +- ``refusal_of`` and ``from_tool``, to tell a model provider's refusal from a + tool's error. +- ``companion_of`` and ``deliver``, to see what an agent is doing and to send + it a message from another session object's handler. """ from restate.agents._common import ( diff --git a/python/restate/agents/langgraph/__init__.py b/python/restate/agents/langgraph/__init__.py index 04504f9e..c09d5adc 100644 --- a/python/restate/agents/langgraph/__init__.py +++ b/python/restate/agents/langgraph/__init__.py @@ -8,28 +8,38 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -LangGraph agents in session objects: a tool-calling graph whose checkpoints -are in the object's database. As sdk-sessions' ``@restatedev/agents/langgraph``:: +"""LangGraph agents as session objects, with checkpoints stored in the object's database. + +For example:: from restate import sessions from restate.agents import langgraph, tool @tool async def get_weather(city: str) -> str: - '''The current weather in a city.''' + '''Return the current weather in a city.''' return f"Sunny in {city}" assistant = langgraph.agent("assistant", system="Answer briefly.", tools=[get_weather]) sessions.serve(assistant) -``agent`` creates an agent; ``prompt`` answers a prompt in a handler of your -own; ``SqlCheckpointSaver`` keeps a graph of your own in ``ctx.db``. +``agent`` creates an agent, ``prompt`` answers a prompt from a handler of your +own, and ``SqlCheckpointSaver`` stores the checkpoints of a graph of your own +in ``ctx.db``. + +This differs from the older ``restate.ext.langchain`` integration, whose +middleware makes a LangChain ``create_agent`` agent durable inside any regular +Restate handler by journaling each model response, and leaves tool side +effects for you to wrap in ``ctx.run_typed``. Pick ``restate.ext.langchain`` +to add durability to an agent you run in your own services; pick this module +for a ready-made agent per key, whose conversation lives in a database and +which comes with built-in handlers to chat, read the conversation, watch its +progress and stop a turn. Install the ``langgraph`` extra (``pip install restate-sdk[langgraph]``). A -model named ``openai/...`` also needs ``langchain-openai``, and -``anthropic/...`` needs ``langchain-anthropic``; or pass any LangChain chat -model as ``model``. +model named ``openai/...`` also needs ``langchain-openai``, and one named +``anthropic/...`` needs ``langchain-anthropic``. Alternatively, pass any +LangChain chat model as ``model``. """ try: diff --git a/python/restate/agents/langgraph/_agent.py b/python/restate/agents/langgraph/_agent.py index 5c67c939..7e07d81a 100644 --- a/python/restate/agents/langgraph/_agent.py +++ b/python/restate/agents/langgraph/_agent.py @@ -10,24 +10,25 @@ # # LangGraph's graph API (langgraph.graph, its StateGraph and compiled graphs) is untyped: no py.typed. # pyright: reportMissingTypeStubs=false, reportUnknownMemberType=false -""" -LangGraph agents: a tool-calling graph whose checkpoints are in the object's -database (``_checkpointer.py``). As sdk-sessions' ``langgraph/index.ts``. - -The graph is the usual tool-calling loop: a model node, a tools node, and back -to the model until it answers. The handler is not replayed: a retried attempt -(a crash, a lost connection) runs again from the start, and LangGraph carries -on from the last step it stored: - -- Each step is stored before the next starts (``durability="sync"``), in order - with everything else the handler does. So the tool calls a model made are - stored before any tool runs, and a crash during the tools resumes them with - the same tool call ids, which a tool passes on to stay idempotent. -- A tool's error fails the attempt, which Restate retries (LangGraph's - ToolNode would hand it to the model as the tool's result); a call of a tool - the agent lacks, or with arguments that do not fit, is the model's to fix. -- A message is answered once per invocation: the graph records each answer in - its state, by prompt id, and a retried attempt returns it. +"""LangGraph agents whose checkpoints are stored in the object's database. + +The graph is the usual tool-calling loop: a model node and a tools node, going +back to the model until it answers. Checkpoints are stored by +``SqlCheckpointSaver`` (``_checkpointer.py``). The handler is not replayed: a +retried attempt (after a crash or a lost connection) runs again from the +start, and LangGraph resumes from the last step it stored: + +- Each step is stored before the next one starts (``durability="sync"``), in + order with everything else the handler does. The tool calls a model made are + therefore stored before any tool runs, and a crash while tools are running + resumes them with the same tool call ids, which a tool can pass on to stay + idempotent. +- A tool's error fails the attempt, and Restate retries it. (LangGraph's + ``ToolNode`` would otherwise hand the error to the model as the tool's + result.) A call to a tool the agent does not have, or with arguments that do + not validate, is reported back to the model to fix. +- Each message is answered once per invocation: the graph records each answer + in its state under the prompt's id, and a retried attempt returns it. """ from __future__ import annotations @@ -73,7 +74,7 @@ from restate.sessions import Migration, SessionContext _CANCELLATION = 409 -"""The status of Restate's cancellation, as a turn that was stopped fails with it.""" +"""The status code of Restate's cancellation error, which a stopped turn fails with.""" def _merged(old: dict[str, str], new: dict[str, str]) -> dict[str, str]: @@ -81,7 +82,7 @@ def _merged(old: dict[str, str], new: dict[str, str]) -> dict[str, str]: class State(TypedDict, total=False): - """The agent's graph's state: its conversation, the prompt it answers, and each prompt's answer.""" + """The state of the agent's graph: the conversation, the prompt being answered, and each prompt's answer.""" messages: Annotated[list[AnyMessage], add_messages] prompt: str @@ -94,18 +95,18 @@ class State(TypedDict, total=False): """A model: ``"openai/"`` or ``"anthropic/"``, or a LangChain chat model that calls tools.""" AnyTool: typing.TypeAlias = Tool | BaseTool -"""A tool: one of ``restate.agents`` (``tool``, ``call``, ``signal``...), or a LangChain tool.""" +"""A tool: one made with ``restate.agents`` (``tool``, ``call``, ``wait_for_signal``, ...) or a LangChain tool.""" _PROVIDERS = { "openai": ("langchain_openai", "ChatOpenAI", "langchain-openai"), "anthropic": ("langchain_anthropic", "ChatAnthropic", "langchain-anthropic"), } -"""The providers the usual graph takes by name: LangChain's integration of each, its class, and its package.""" +"""The providers that can be named in a model string, each with its LangChain module, class and package.""" def _provider_of_model(named: str) -> tuple[str, str]: - """``provider/model-id`` as its provider and model id, for a provider the usual graph takes.""" + """Split ``provider/model-id`` into provider and model id, or raise ``ModelError`` for an unsupported provider.""" provider, model_id = provider_of(named) if provider not in _PROVIDERS or not model_id: from_env = f" (from {MODEL_ENV})" if os.environ.get(MODEL_ENV) == named else "" @@ -114,10 +115,11 @@ def _provider_of_model(named: str) -> tuple[str, str]: def _chat_model(named: str) -> BaseChatModel: - """ - The chat model ``provider/model-id`` names, through LangChain's - integration of the provider, imported when a model of it is used: only - that one need be installed, beside the ``langgraph`` extra. + """Create the LangChain chat model that ``provider/model-id`` names. + + The provider's LangChain package is imported only when one of its models is + used, so only that package needs to be installed next to the ``langgraph`` + extra. """ provider, model_id = _provider_of_model(named) module, class_name, package = _PROVIDERS[provider] @@ -131,7 +133,7 @@ def _chat_model(named: str) -> BaseChatModel: class _RestateTool(BaseTool): - """A ``restate.agents`` tool as a LangChain tool: it gets its tool call's id, and checks the model's arguments itself.""" + """A ``restate.agents`` tool as a LangChain tool; it receives its tool call id and validates its own arguments.""" _tool: Tool = PrivateAttr() _ctx: SessionContext = PrivateAttr() @@ -156,12 +158,12 @@ def _native(tool: Tool, ctx: SessionContext) -> BaseTool: def _tools_node(ctx: SessionContext, tools: list[BaseTool]) -> Callable[[State, RunnableConfig], Awaitable[Any]]: - """ - The tools node: each tool call of the model's last message, run by a - ToolNode, which answers the model with what is wrong when a call names a - tool the agent lacks or arguments that do not fit; any other error of a - tool fails the attempt. The round is published as it starts and ends, for - the clients that follow the conversation. + """Return the graph's tools node, which runs the tool calls of the model's last message. + + A ``ToolNode`` runs them, and tells the model what went wrong when a call + names an unknown tool or has invalid arguments; any other tool error fails + the attempt. The start and end of each round are published for clients that + follow the conversation. """ node = ToolNode(tools) @@ -170,7 +172,7 @@ async def run_tools(state: State, config: RunnableConfig) -> Any: try: ran = await node.ainvoke(state, config) except Exception as error: - # A tool's error, never taken for the model's provider refusing. + # Mark the error as the tool's, so that it is never mistaken for the model provider refusing. from_tool(ctx, error) raise await tool_call_ended(ctx) @@ -184,7 +186,7 @@ def _text(message: BaseMessage) -> str: def _graph(ctx: SessionContext, model: BaseChatModel, tools: list[BaseTool], system: str | None) -> Any: - """The agent's graph: ``model``, calling ``tools``, until it answers the prompt.""" + """Build the agent's graph: ``model`` calls ``tools`` until it answers the prompt.""" bound = model.bind_tools(tools) if tools else model async def call_model(state: State) -> dict[str, Any]: @@ -215,7 +217,8 @@ async def _answer_once( ) -> str: chat_model = model if isinstance(model, BaseChatModel) else _chat_model(model_of(model)) tools = [_native(t, ctx) if isinstance(t, Tool) else t for t in given] - # The latest checkpoint is all an agent resumes from: earlier ones would store the conversation again at each step. + # An agent resumes only from the latest checkpoint; keeping earlier ones would store the conversation again + # at each step. app: Compiled = _graph(ctx, chat_model, tools, system).compile( checkpointer=SqlCheckpointSaver(ctx.db, history="latest") ) @@ -226,7 +229,7 @@ async def _answer_once( answers = values.get("answers", {}) if prompt_id in answers: return stored_answer(answers[prompt_id]) - # Finish an interrupted run first (this prompt's, after a crash), then start this prompt's. + # Finish an interrupted run first (this prompt's own, after a crash), then start this prompt's run. # A run is unfinished while its prompt has no answer. last = values.get("prompt") interrupted = last is not None and last not in answers @@ -236,21 +239,21 @@ async def _answer_once( try: await (stop.guard(run) if stop is not None else run) except AttemptEnded: - # The attempt ended: it stops, and the next attempt goes on with the run. + # The attempt ended: stop here, and let the next attempt continue the run. raise except TerminalError as error: - # Cancelled: the run ends, and the next prompt starts afresh. + # Cancelled: close the run, so that the next prompt starts afresh. if error.status_code == _CANCELLATION: await _close(app, config, CANCELLED) raise - # Failed for good (a LangChain tool's terminal error): the run ends with it, rather than the next - # prompt meeting it again. An earlier prompt's, this one goes on. + # Failed for good (a LangChain tool's terminal error): close the run with the error, so that the next + # prompt does not run into it again. If the run was an earlier prompt's, carry on with this one. await _close(app, config, f"Error: {error.message}") if not interrupted or last == prompt_id: raise except Exception as error: - # Refused for good (a wrong API key): the run ends unanswered, and fails as terminal. - # An earlier prompt's, this one goes on. + # Refused for good (a wrong API key, say): close the run unanswered and fail with a terminal error. + # If the run was an earlier prompt's, carry on with this one. refused = refusal_of(ctx, error) if refused is None: raise @@ -260,12 +263,12 @@ async def _answer_once( async def _close(app: Compiled, config: RunnableConfig, text: str) -> None: - """ - End the run of ``app`` that was cancelled, refused or failed, as its ``answer`` - node: a result of ``text`` for each tool call of the last reply that has - none, then ``text`` as the reply and the prompt's answer. A run that - stored no input yet (cancelled at once) has nothing to end: the prompt the - state has is an earlier one, which has its answer. + """Close a cancelled, refused or failed run of ``app`` with ``text``, as its ``answer`` node would. + + Each tool call of the last reply that has no result gets ``text`` as its + result, and ``text`` becomes the reply and the prompt's answer. A run that + stored no input yet (cancelled right away) has nothing to close: the state's + prompt is an earlier one, which already has its answer. """ values = typing.cast(State, (await app.aget_state(config)).values) prompt = values.get("prompt") @@ -283,7 +286,7 @@ async def _close(app: Compiled, config: RunnableConfig, text: str) -> None: async def _transcript(ctx: SessionContext) -> Transcript: - """The thread (the object's key) from its last checkpoint, as ``messages`` gives it.""" + """Return the thread (the object's key) from its last checkpoint, in the form ``messages`` returns.""" transcript = Transcript() saved = await SqlCheckpointSaver(ctx.db).aget_tuple({"configurable": {"thread_id": ctx.key}}) messages: list[BaseMessage] = saved.checkpoint["channel_values"].get("messages", []) if saved else [] diff --git a/python/restate/agents/langgraph/_checkpointer.py b/python/restate/agents/langgraph/_checkpointer.py index 4d74c63f..f99ad404 100644 --- a/python/restate/agents/langgraph/_checkpointer.py +++ b/python/restate/agents/langgraph/_checkpointer.py @@ -8,33 +8,36 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -A LangGraph checkpointer on a session object's database: ``SqlCheckpointSaver``, -over a database session ``db`` (in a handler, ``ctx.db``). As sdk-sessions' -``langgraph/checkpointer.ts``. - -Three tables, with values as LangGraph's serializer writes them: - -- ``restate_langgraph_checkpoints``: each checkpoint, without its channel values. -- ``restate_langgraph_channel_values``: each channel's value at each version, - stored when a step changes it. A checkpoint's values are those at its - channel versions, so a step stores only what it changed, not the whole - state again (an agent's message list, say). Versions are unique (a count - and a random part): a thread forked from an earlier checkpoint takes the - same counts again, and its values must not replace those of the branch it - forked from. -- ``restate_langgraph_writes``: the pending writes of each checkpoint's tasks. - -Each ``aput`` and ``aput_writes`` is one SQL transaction, so a step is stored -whole, in order with everything else the handler does. Operations run one at -a time, in call order: LangGraph writes the results of parallel tasks at once. -It is asynchronous only, as ``ctx.db`` is. - -By default a thread keeps every checkpoint (LangGraph's history: ``alist``, -time travel, forks). With ``history="latest"``, each ``aput`` drops the -thread's earlier checkpoints, their writes and the values only they had: a -thread then takes the space of its state once, not once per step (a message -list stored again at each step grows with the square of its length). +"""A LangGraph checkpointer that stores checkpoints in a session object's database. + +``SqlCheckpointSaver`` works on any database session; in a handler, that is +``ctx.db``. It uses three tables, and stores values as LangGraph's serializer +writes them: + +- ``restate_langgraph_checkpoints`` holds each checkpoint, without its + channel values. +- ``restate_langgraph_channel_values`` holds each channel's value at each + version, written when a step changes it. A checkpoint's values are the ones + at its channel versions, so a step stores only what it changed rather than + the whole state (an agent's message list, for example). Versions are unique, + made of a counter and a random part: a thread forked from an earlier + checkpoint reuses the same counter values, and its values must not + overwrite those of the branch it forked from. +- ``restate_langgraph_writes`` holds the pending writes of each checkpoint's + tasks. + +Each ``aput`` and ``aput_writes`` runs in one SQL transaction, so a step is +stored whole, in order with everything else the handler does. Operations run +one at a time, in the order they are called, even when LangGraph writes the +results of parallel tasks at once. The saver is async-only, because ``ctx.db`` +is. + +By default a thread keeps every checkpoint, which LangGraph needs for +``alist``, time travel and forks. With ``history="latest"``, each ``aput`` +deletes the thread's earlier checkpoints, their writes, and the channel values +that only they referenced. A thread then stores its state once instead of once +per step; otherwise a message list stored again at each step grows with the +square of its length. """ from __future__ import annotations @@ -97,7 +100,7 @@ )""", ], ) -"""The checkpointer's tables, a schema of their own, applied where it is used.""" +"""The checkpointer's tables, a schema of their own that is applied on first use.""" _TABLES = ("restate_langgraph_checkpoints", "restate_langgraph_channel_values", "restate_langgraph_writes") @@ -108,10 +111,21 @@ class SqlCheckpointSaver(BaseCheckpointSaver[str]): - """ - LangGraph's checkpoints in a database session, such as ``ctx.db``: see - the module. ``history``: the checkpoints a thread keeps, ``"all"`` (the - default) or the ``"latest"`` only. + """A LangGraph checkpoint saver that stores checkpoints in a database session. + + Pass it to ``compile`` to keep a graph of your own in ``ctx.db``:: + + app = graph.compile(checkpointer=SqlCheckpointSaver(ctx.db)) + await app.ainvoke(inputs, {"configurable": {"thread_id": ctx.key}}) + + See the module docstring for how checkpoints are stored. + + Args: + db: The database to store checkpoints in, usually ``ctx.db``. + history: Which checkpoints a thread keeps: ``"all"`` (the default), + or only the ``"latest"``. + serde: LangGraph's serializer for checkpoints and values; LangGraph's + default if omitted. """ def __init__( @@ -128,10 +142,11 @@ def __init__( self._ready = False def get_next_version(self, current: str | None, channel: None = None) -> str: - """ - The version after ``current``: its count plus one, zero-padded so that - versions compare as their counts do, and a random part, so that - branches of a thread never share one. + """Return the channel version that follows ``current``. + + A version is the counter of ``current`` plus one, zero-padded so that + versions sort in counter order, followed by a random part so that two + branches of a thread never share a version. """ if current is None: count = 0 @@ -150,7 +165,7 @@ def _load(self, type_: SqlValue, value: SqlValue) -> Any: return self.serde.loads_typed((str(type_), _bytes(value))) async def _channel_values(self, thread_id: str, ns: str, checkpoint: Checkpoint) -> dict[str, Any]: - """The values of ``checkpoint``'s channels, at its versions.""" + """Return the values of ``checkpoint``'s channels at its channel versions.""" versions = list(checkpoint.get("channel_versions", {}).items()) if not versions: return {} @@ -336,7 +351,7 @@ async def adelete_thread(self, thread_id: str) -> None: def _thread(config: RunnableConfig) -> tuple[str, str]: - """The thread of ``config``, and its namespace.""" + """Return the thread id and checkpoint namespace of ``config``.""" configurable = _configurable(config) return str(configurable["thread_id"]), str(configurable.get("checkpoint_ns", "")) @@ -350,7 +365,7 @@ def _bytes(value: SqlValue) -> bytes: def _matches(actual: object, expected: object) -> bool: - """Whether ``actual`` has every field of ``expected``, recursively (a metadata filter).""" + """Return whether ``actual`` has every field of ``expected``, recursively, for metadata filters.""" if isinstance(expected, dict): if not isinstance(actual, dict): return False @@ -360,11 +375,7 @@ def _matches(actual: object, expected: object) -> bool: async def _drop_earlier(tx: DatabaseWriter, thread_id: str, ns: str, checkpoint: Checkpoint) -> None: - """ - Drop what the thread keeps from before ``checkpoint``: its other - checkpoints, their pending writes, and the channel values at versions - other than ``checkpoint``'s. - """ + """Delete the thread's checkpoints, writes and channel values that ``checkpoint`` does not use.""" for table in ("restate_langgraph_checkpoints", "restate_langgraph_writes"): await tx.execute( f"DELETE FROM {table} WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id <> ?", diff --git a/python/restate/agents/pydantic_ai/__init__.py b/python/restate/agents/pydantic_ai/__init__.py index d08ac0e9..f6c896d4 100644 --- a/python/restate/agents/pydantic_ai/__init__.py +++ b/python/restate/agents/pydantic_ai/__init__.py @@ -8,27 +8,37 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -Pydantic AI agents in session objects: each conversation in the object's -database, stored a step at a time:: +"""Pydantic AI agents as session objects, with conversations stored in the object's database. + +For example:: from restate import sessions from restate.agents import pydantic_ai, tool @tool async def get_weather(city: str) -> str: - '''The current weather in a city.''' + '''Return the current weather in a city.''' return f"Sunny in {city}" assistant = pydantic_ai.agent("assistant", system="Answer briefly.", tools=[get_weather]) sessions.serve(assistant) -``agent`` creates an agent; ``prompt`` answers a prompt in a handler of your own. +``agent`` creates an agent, and ``prompt`` answers a prompt from a handler of +your own. The conversation is stored one step at a time. + +This differs from the older ``restate.ext.pydantic`` integration, whose +``RestateAgent`` wraps a Pydantic AI ``Agent`` so that it runs durably inside +any regular Restate handler, journaling model calls and MCP tool calls (and, +optionally, function tools) with ``ctx.run_typed``. Pick +``restate.ext.pydantic`` to add durability to an agent you run in your own +services; pick this module for a ready-made agent per key, whose conversation +lives in a database and which comes with built-in handlers to chat, read the +conversation, watch its progress and stop a turn. Install the ``pydantic_ai`` extra (``pip install restate-sdk[pydantic_ai]``). -A model's provider also needs Pydantic AI's extra for it, such as -``pydantic-ai-slim[openai]`` for ``openai/...`` models; or pass any Pydantic -AI model as ``model``. +Each model provider also needs Pydantic AI's extra for it, such as +``pydantic-ai-slim[openai]`` for ``openai/...`` models. Alternatively, pass +any Pydantic AI model as ``model``. """ try: diff --git a/python/restate/agents/pydantic_ai/_agent.py b/python/restate/agents/pydantic_ai/_agent.py index 42aa62ce..183bb614 100644 --- a/python/restate/agents/pydantic_ai/_agent.py +++ b/python/restate/agents/pydantic_ai/_agent.py @@ -8,25 +8,24 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -Pydantic AI agents: each conversation in the object's database, stored a step -at a time (sqlexperiment's ``restate_agents.pydantic_ai``, on this package). - -The handler is not replayed: a retried attempt (a crash, a lost connection) -runs again from the start, and the run carries on from the last step stored: - -- The conversation is stored before each model request, and with each model - response, in order with everything else the handler does. So the tool calls - a model made are stored before any tool runs, and a crash during the tools - runs them again with the same tool call ids, which a tool passes on to stay - idempotent. -- A tool's error fails the attempt, which Restate retries; a - ``restate.TerminalError``, which a retry would repeat, is the tool's answer - to the model; and so are arguments that do not fit (Pydantic AI asks the - model again). -- A message is answered once per invocation: its request is marked with the - prompt's id, its answer too, and a retried attempt that finds them returns - the answer, or finishes the run first. +"""Pydantic AI agents whose conversations are stored in the object's database, one step at a time. + +The handler is not replayed: a retried attempt (after a crash or a lost +connection) runs again from the start, and the run resumes from the last step +stored: + +- The conversation is stored before each model request and with each model + response, in order with everything else the handler does. The tool calls a + model made are therefore stored before any tool runs, and a crash while tools + are running runs them again with the same tool call ids, which a tool can + pass on to stay idempotent. +- A tool's error fails the attempt, and Restate retries it. A + ``restate.TerminalError``, which a retry would only repeat, is instead + returned to the model as the tool's answer, and so are arguments that do not + validate (Pydantic AI asks the model again). +- Each message is answered once per invocation: its request and its answer are + both marked with the prompt's id, and a retried attempt that finds them + returns the answer, or finishes the run first. """ from __future__ import annotations @@ -78,19 +77,19 @@ from restate.sessions.internal import Schema _CANCELLATION = 409 -"""The status of Restate's cancellation, as a turn that was stopped fails with it.""" +"""The status code of Restate's cancellation error, which a stopped turn fails with.""" _PROMPT = "restate_agents_prompt" -"""The metadata of a prompt's request (which models do not read): the prompt's id.""" +"""The metadata key that marks a prompt's request with the prompt's id; models do not read it.""" _ANSWERS = "restate_agents_answers" -"""The metadata of the response that answers a prompt: the prompt's id.""" +"""The metadata key that marks the response answering a prompt with the prompt's id.""" AnyModel: typing.TypeAlias = str | Model """A model: ``"/"`` (or Pydantic AI's ``":"``), or a Pydantic AI model.""" AnyTool: typing.TypeAlias = Tool | PydanticTool[Any] | Callable[..., Any] -"""A tool: one of ``restate.agents`` (``tool``, ``call``, ``signal``...), a Pydantic AI tool, or a tool function.""" +"""A tool: one made with ``restate.agents`` (``tool``, ``call``, ...), a Pydantic AI tool, or a plain function.""" def _check(own: Agent[Any, str] | None, system: object) -> None: @@ -99,10 +98,11 @@ def _check(own: Agent[Any, str] | None, system: object) -> None: def _toolset(given: Sequence[AnyTool], ctx: SessionContext) -> FunctionToolset[Any]: - """ - ``given`` as a Pydantic AI toolset for a run of the turn of ``ctx``: - ``restate.agents`` tools, each with its tool call's id (it checks the - model's arguments itself); Pydantic AI's tools and tool functions as they are. + """Return ``given`` as a Pydantic AI toolset for a run in ``ctx``. + + ``restate.agents`` tools are wrapped so that they receive their tool call id + (they validate their own arguments); Pydantic AI tools and tool functions are + passed through unchanged. """ def native(tool: Tool) -> PydanticTool[Any]: @@ -131,10 +131,10 @@ async def run(run_ctx: RunContext[Any], **args: Any) -> str: class _Conversation: - """ - A conversation in the object's database, a message per row. Saving it - writes the messages that changed since it was loaded or saved, in one - transaction. + """A conversation in the object's database, stored one message per row. + + ``save`` writes the messages that changed since the last load or save, in + one transaction. """ def __init__(self, db: Database, thread: str) -> None: @@ -170,11 +170,12 @@ async def save(self, messages: Sequence[ModelMessage]) -> None: class _Restate(AbstractCapability[Any]): - """ - A run in a Restate handler: its conversation stored before each model - request and with each response; each tool call published as it starts and - ends, for the clients that follow the conversation; a tool's - ``restate.TerminalError`` its answer, and any other error the attempt's. + """Make a Pydantic AI run durable in a Restate handler. + + It stores the conversation before each model request and with each + response, and publishes each tool call as it starts and ends for clients + that follow the conversation. A tool's ``restate.TerminalError`` becomes its + answer to the model; any other error fails the attempt. """ def __init__(self, ctx: SessionContext, conversation: _Conversation) -> None: @@ -214,16 +215,16 @@ async def on_tool_execute_error( if isinstance(error, TerminalError) and error.status_code != _CANCELLATION: await tool_call_ended(self.ctx) return f"Error: {error.message}" - # A tool's error, never taken for the model's provider refusing. + # Mark the error as the tool's, so that it is never mistaken for the model provider refusing. from_tool(self.ctx, error) raise error def _model(model: str | Model | None) -> Model: - """ - The model, as Pydantic AI resolves it: ``provider/model-id`` (or - ``provider:model-id``). A provider whose package is missing, or without - its API key, refuses every message alike: a ``ModelError``. + """Resolve ``provider/model-id`` (or ``provider:model-id``) to a Pydantic AI model. + + Raise ``ModelError`` if the provider's package is missing or its API key is + not set, since every message would then fail the same way. """ if isinstance(model, Model): return model @@ -236,17 +237,17 @@ def _model(model: str | Model | None) -> Model: def _prompt_of(message: ModelMessage) -> str | None: - """The id of the prompt that ``message`` is, if it is one.""" + """Return the prompt id ``message`` is marked with, or ``None`` if it is not a prompt.""" return (message.metadata or {}).get(_PROMPT) if isinstance(message, ModelRequest) else None def _answer_of(history: list[ModelMessage], asked: int) -> int | None: - """The index of the response that answers the prompt at ``asked``, if it was answered.""" + """Return the index of the response that answers the prompt at index ``asked``, or ``None``.""" return next((i for i in range(asked, len(history)) if (history[i].metadata or {}).get(_ANSWERS)), None) def _result(part: ToolReturnPart | RetryPromptPart) -> str: - """A tool call's result, as the model reads it: its return, or what to fix.""" + """Return a tool call's result as the model sees it: the return value, or the retry prompt.""" return part.model_response() if isinstance(part, RetryPromptPart) else part.model_response_str() @@ -255,11 +256,11 @@ def _text(response: ModelMessage) -> str: async def _close(conversation: _Conversation, prompt_id: str, text: str) -> None: - """ - End the run of prompt ``prompt_id``, which was cancelled, refused or - failed: a result of ``text`` for each tool call of its last response that - has none, then ``text`` as its answer. A run that stored no prompt yet has - nothing to end. + """Close the cancelled, refused or failed run of prompt ``prompt_id`` with ``text``. + + Each tool call of the last response gets ``text`` as its result, and ``text`` + becomes the prompt's answer. A run that stored no prompt yet has nothing to + close. """ messages = await conversation.load() if not any(_prompt_of(m) == prompt_id for m in messages): @@ -280,7 +281,7 @@ async def _answer_once( system: str | None, given: Sequence[AnyTool], ) -> str: - # An Agent of your own's model too, resolved here: refused as any other (a missing key or package). + # Resolve the model of an Agent of your own here as well, so that a missing key or package fails the same way. model = _model(named if named is not None else own.model if own is not None else None) agent: Agent[Any, str] = own or Agent(instructions=system) toolsets = [_toolset(given, ctx)] if given else [] @@ -291,7 +292,7 @@ async def _answer_once( asked = next((i for i, m in enumerate(history) if _prompt_of(m) == prompt_id), None) if asked is not None and (answer := _answer_of(history, asked)) is not None: return stored_answer(_text(history[answer])) - # Finish an interrupted run first (this prompt's, after a crash), then start this prompt's. + # Finish an interrupted run first (this prompt's own, after a crash), then start this prompt's run. last = next((i for i in reversed(range(len(history))) if _prompt_of(history[i]) is not None), None) interrupted = last is not None and _answer_of(history, last) is None running = _prompt_of(history[last]) if interrupted and last is not None else prompt_id @@ -305,21 +306,21 @@ async def _answer_once( try: result = await (stop.guard(run) if stop is not None else run) except AttemptEnded: - # The attempt ended: it stops, and the next attempt goes on with the run. + # The attempt ended: stop here, and let the next attempt continue the run. raise except TerminalError as error: if error.status_code == _CANCELLATION: - # Cancelled: the run ends, and the next prompt starts afresh. + # Cancelled: close the run, so that the next prompt starts afresh. await _close(conversation, running, CANCELLED) raise - # Failed for good: the run ends with it, rather than the next prompt meeting it again. + # Failed for good: close the run with the error, so that the next prompt does not run into it again. await _close(conversation, running, f"Error: {error.message}") if running == prompt_id: raise continue except Exception as error: - # Refused for good (a wrong API key): the run ends unanswered, and fails as terminal. - # An earlier prompt's, this one goes on. + # Refused for good (a wrong API key, say): close the run unanswered and fail with a terminal error. + # If the run was an earlier prompt's, carry on with this one. refused = refusal_of(ctx, error) if refused is None: raise @@ -334,7 +335,7 @@ async def _answer_once( async def _transcript(ctx: SessionContext) -> Transcript: - """The conversation (the object's key), as ``messages`` gives it.""" + """Return the conversation (the object's key) in the form ``messages`` returns.""" transcript = Transcript() for message in await _Conversation(ctx.db, ctx.key).load(): for part in message.parts: diff --git a/python/restate/sessions/_calls.py b/python/restate/sessions/_calls.py index 8c05a1fb..4ecdb9cd 100644 --- a/python/restate/sessions/_calls.py +++ b/python/restate/sessions/_calls.py @@ -8,11 +8,12 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -Calls and sends from a session object's handler: ``ctx.call`` and -``ctx.send``, and a transaction's ``tx.send``. A target is a handler (with -its key, for an object or a workflow), typed by its signature; a target known -only at run time is named, with ``generic_call`` and ``generic_send``:: +"""Calls and sends from session objects: ``ctx.call``, ``ctx.send`` and ``tx.send``. + +The target of a call or send is a handler function, passed together with its +key for a virtual object or a workflow, and the call is typed by the +handler's signature. A target known only at run time is addressed by name, +with ``generic_call`` and ``generic_send``:: count = await ctx.call(increment, 5, key="a", idempotency_key=f"{ctx.request.id}:inc") ctx.send(record, entry, key="audit", delay=timedelta(hours=1)) @@ -38,7 +39,7 @@ @dataclass(frozen=True) class GenericSend: - """A send, its input serialized: as ``ctx.send`` makes it, or a transaction's outbox keeps it.""" + """A send with its input already serialized, as ``ctx.send`` makes it or a transaction's outbox stores it.""" service: str handler: str @@ -50,7 +51,7 @@ class GenericSend: class CallHost(Protocol): - """What calls and sends need of the SDK's context.""" + """The part of the SDK's context that calls and sends use.""" def generic_call( self, @@ -75,10 +76,11 @@ def generic_send( def _target(handler: Callable[..., Any], key: str | None) -> tuple[str, str, Serde[Any], Serde[Any]]: - """ - The service, handler name and serdes of ``handler``, called with ``key``. An - object's or a workflow's handler needs its key, and a service's has none: - a ``TerminalError`` otherwise, since a retry would fail the same way. + """Return the service name, handler name and serdes of ``handler``, checking ``key``. + + Raises ``TerminalError`` (since a retry would fail the same way) if + ``handler`` is not a handler, if an object's or workflow's handler has no + key, or if a service's handler has one. """ try: target: Handler[Any, Any] = handler_from_callable(handler) @@ -95,7 +97,11 @@ def _target(handler: Callable[..., Any], key: str | None) -> tuple[str, str, Ser class CallFuture(RestateDurableCallFuture[O]): - """The result of a call, as its output serde reads it: awaitable, with the call's invocation id.""" + """A future for the result of a call, decoded with the handler's output serde. + + Await it for the result. ``invocation_id()`` returns the id of the call's + invocation, and ``cancel_invocation()`` cancels it. + """ def __init__(self, call: RestateDurableCallFuture[bytes], serde: Serde[O]) -> None: super().__init__() @@ -116,7 +122,7 @@ async def cancel_invocation(self) -> None: class Calls: - """``call`` and ``generic_call`` (see ``Caller`` in ``_api.py``), over ``_call_host``, the SDK's context.""" + """A mixin that implements ``call`` and ``generic_call`` (see ``Caller`` in ``_api.py``) on top of ``_call_host``.""" _call_host: CallHost @@ -175,7 +181,11 @@ def generic_call( class Sends: - """``send`` and ``generic_send`` (see ``Sender`` in ``_api.py``), into ``_send_sink``: made at once, or at a commit.""" + """A mixin that implements ``send`` and ``generic_send`` (see ``Sender`` in ``_api.py``) by passing each send to ``_send_sink``. + + The sink either sends right away (``ctx.send``) or when a transaction + commits (``tx.send``). + """ _send_sink: Callable[[GenericSend], None] @@ -236,7 +246,7 @@ def generic_send( def send_now(host: CallHost) -> Callable[[GenericSend], None]: - """The sink of ``ctx.send``: each send is made at once.""" + """Return the sink of ``ctx.send``, which sends each send right away.""" def sink(send: GenericSend) -> None: host.generic_send( diff --git a/python/restate/sessions/_database.py b/python/restate/sessions/_database.py index 924ca26e..a1240251 100644 --- a/python/restate/sessions/_database.py +++ b/python/restate/sessions/_database.py @@ -8,16 +8,17 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -A session's database, through the native extension (``src/sessions.rs``): -SQLite-compatible SQL, kept in the object's state. The extension reads and -writes the state through a store, in the order it makes its requests; this -module orders the calls and runs transactions:: +"""The SQL database of a session, backed by the native extension (``src/sessions.rs``). + +The database speaks SQLite-compatible SQL and is stored in the object's +state. The extension reads and writes that state through a ``Store``, in the +order it issues its requests; this module serializes calls and runs +transactions:: db = await open("Ledger/a", store) async with db.transaction() as tx: await tx.execute("INSERT INTO transfers VALUES(?, ?, ?, ?)", [id, src, dst, amount]) - await db.save() # ends the session, leaving its engine to the next one + await db.save() # end the session, handing its engine to the next one """ from __future__ import annotations @@ -47,16 +48,16 @@ from restate.exceptions import TerminalError SqlValue = None | int | float | str | bytes -"""An SQL value: ``None``, ``int`` (64 bits), ``float``, ``str``, or ``bytes`` (BLOB).""" +"""A value read from the database: ``None``, ``int`` (64-bit), ``float``, ``str`` or ``bytes`` (a BLOB).""" SqlParam = SqlValue | bool | bytearray | memoryview -"""An SQL parameter: an ``SqlValue``, a ``bool`` (stored as 0 or 1), or a BLOB as any bytes-like object.""" +"""A query parameter: any ``SqlValue``, a ``bool`` (stored as 0 or 1), or any bytes-like object (a BLOB).""" Row = dict[str, SqlValue] -"""A row, by column name.""" +"""A row, as a dict that maps column names to values.""" Rows = list[list[SqlValue]] -"""Rows as lists, in column order.""" +"""Rows as lists of values, in column order.""" class SqlError(TerminalError): @@ -79,27 +80,29 @@ class SqlError(TerminalError): class Store(Protocol): - """ - A key/value store, the object's state. - - ``set`` and ``clear`` return at once, and ``sync`` completes once every - write before it is stored. Writes are stored in the order they were made: - a crash keeps a prefix of the writes since the last sync and loses the - rest. ``get`` issues its read when it is called, and returns the stored - value or this request's own latest write (``None`` when absent), or an - awaitable of it; so a plain ``async def get`` does not do, as its read - would be issued only when it first runs, after writes made later. - - A store may also have ``keys()``: every key, as ``ctx.state_keys()``, or an - awaitable of them. Then a commit that checkpoints deletes the packs no - checkpoint needs any more; a store without it keeps them. - - A store over a Restate context may give the invocation's VM, as ``vm`` - (``restate._internal.PyVM``): then the extension writes to it, and reads - from it, directly, without the store's ``get``, ``set`` and ``clear``; it - calls ``wait(handle)`` (an awaitable, done once the VM has the read whose - handle it is) only for a read it must wait for. ``database_name`` names the - key's database, apart from the object's own (see ``_session.database``). + """The key/value store that holds a database: the object's state. + + ``set`` and ``clear`` return immediately, and ``sync`` completes once + every earlier write is stored. Writes are stored in the order they were + made, so a crash keeps some prefix of the writes since the last sync and + loses the rest. ``get`` must issue its read when it is called, and return + the stored value (or this request's own latest write, or ``None`` if + absent), either directly or as an awaitable. A plain ``async def get`` + does not qualify: its read would only be issued when the coroutine first + runs, after any writes made in the meantime. + + Optional members: + + - ``keys()`` returns every key, as ``ctx.state_keys()`` does, or an + awaitable of them. With it, a commit that writes a checkpoint deletes + the packs no checkpoint needs any more; without it, they are kept. + - ``vm`` is the invocation's VM (``restate._internal.PyVM``), for a store + over a Restate context. The extension then reads and writes the VM + directly instead of calling ``get``, ``set`` and ``clear``, and calls + ``wait(handle)`` (an awaitable that completes once the VM has the read + for ``handle``) only when it must wait for a read. + - ``database_name`` names a database of the key other than the object's + own (see ``_session.database``). """ def get(self, key: str, /) -> bytes | None | Awaitable[bytes | None]: ... @@ -109,11 +112,12 @@ def sync(self) -> None | Awaitable[None]: ... async def open(namespace: str, store: Store, *, reader: bool = False) -> Connection: # pylint: disable=redefined-builtin - """ - Open the database of ``namespace`` over ``store``. One process-wide page - cache serves every namespace. A ``reader`` reads beside the namespace's - owner (a shared handler): the state as of open, whatever the owner commits - meanwhile, until the owner's second checkpoint after it; it writes nothing. + """Open the database of ``namespace``, stored in ``store``. + + All namespaces share one process-wide page cache. A ``reader`` (for a + shared handler) reads alongside the namespace's owner and never writes: + it sees the state as of when it was opened, whatever the owner commits + meanwhile, until the owner's second checkpoint after that. """ boundary = _Boundary(store, asyncio.get_running_loop()) vm = getattr(store, "vm", None) @@ -131,37 +135,37 @@ async def open(namespace: str, store: Store, *, reader: bool = False) -> Connect async def _settled(answer: T | Awaitable[T]) -> T: - """A call's answer: its value, as it gave it, or once the asyncio future it gave is done.""" + """Return ``answer``, awaiting it first if it is awaitable.""" if inspect.isawaitable(answer): return await answer return answer class _Boundary: - """ - The store as the extension calls it, through ``dispatch``, on the loop's - thread, in the order of the extension's requests. An - error of the store's crosses the extension as text, which ends the - session: the first is kept here, and the calls that fail on the session it - ended raise it as it was raised (the SDK's suspension, say). A failed write - cannot fail its caller, which has moved on: the next ``sync`` fails. And a - commit's sync first runs what its transaction gave - ``ConnectionTransaction.with_commit``: after the commit's writes, before - anything is awaited. + """An adapter that serves the extension's requests, which arrive through ``dispatch``, from a ``Store``. + + Requests are handled on the event loop's thread, in the order the + extension makes them. A store error reaches the extension only as text, + and ends the session, so the first one is kept here: the calls that then + fail re-raise the original exception (the SDK's suspension, for example). + A failed write cannot fail its caller, which has moved on, so the next + ``sync`` fails instead. Before a commit's sync, the callbacks the + transaction registered with ``ConnectionTransaction.with_commit`` run, + after the commit's writes and before anything is awaited. """ def __init__(self, store: Store, loop: asyncio.AbstractEventLoop) -> None: self.store: Store = store self.loop: asyncio.AbstractEventLoop = loop - # Its optional methods (see `Store`). + # The store's optional methods (see `Store`). self._keys: Callable[[], Sequence[str] | Awaitable[Sequence[str]]] | None = getattr(store, "keys", None) self._wait: Callable[[int], Awaitable[None]] | None = getattr(store, "wait", None) self.lists_keys: bool = callable(self._keys) - # The first error of the store's. + # The store's first error. self.failed: BaseException | None = None - # What the next sync runs first: the committing transaction's. + # What the next sync runs first: the committing transaction's callbacks. self.with_sync: Callable[[], None] | None = None - # The first write that failed: every sync after it fails. + # The first failed write; every sync after it fails. self._write_failed: BaseException | None = None def _keep(self, error: BaseException) -> None: @@ -169,7 +173,7 @@ def _keep(self, error: BaseException) -> None: self.failed = error def dispatch(self, op: int, key: Any, value: bytes | None, reply: SqlReply | None) -> None: - """One request of the extension's: ``key`` is a key of the store, or for ``SQL_WAIT`` a handle of the VM.""" + """Handle one request of the extension; ``key`` is a store key, or a VM handle for ``SQL_WAIT``.""" if op in (SQL_SET, SQL_CLEAR): try: if op == SQL_SET: @@ -181,7 +185,7 @@ def dispatch(self, op: int, key: Any, value: bytes | None, reply: SqlReply | Non if self._write_failed is None: self._write_failed = error return - # Every request but a write has a reply (src/sessions.rs). + # Every request except a write has a reply (src/sessions.rs). assert reply is not None try: if op == SQL_GET: @@ -225,7 +229,7 @@ def done(future: asyncio.Future[object]) -> None: def _reply_value(op: int, result: object) -> object: - """A request's answer as the extension reads it: a value, the keys as a list, or nothing for a sync.""" + """Convert a request's result to what the extension expects: a value, a list of keys, or ``None``.""" if op in (SQL_SYNC, SQL_WAIT): return None if op == SQL_KEYS: @@ -234,7 +238,7 @@ def _reply_value(op: int, result: object) -> object: def objects_of(columns: list[str], rows: Rows) -> list[Row]: - """The rows as dicts by column name; a column name twice is an ``SqlError``.""" + """Return ``rows`` as dicts keyed by column name, raising ``SqlError`` if a column name repeats.""" if len(set(columns)) != len(columns): twice = next(name for i, name in enumerate(columns) if name in columns[:i]) raise SqlError(f"the column {twice} is named twice: alias one (AS), or read the rows with values()") @@ -242,7 +246,7 @@ def objects_of(columns: list[str], rows: Rows) -> list[Row]: class _Sql(abc.ABC): - """SQL on a session of the extension: ``query``, ``get``, ``values`` and ``execute`` by way of ``query_columns``.""" + """Base class for SQL access: ``query``, ``get``, ``values`` and ``execute``, built on ``query_columns``.""" def __init__(self, native: PySqlDatabase, boundary: _Boundary) -> None: self._native: PySqlDatabase = native @@ -250,36 +254,36 @@ def __init__(self, native: PySqlDatabase, boundary: _Boundary) -> None: @abc.abstractmethod async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: - """The rows of one statement, with its column names: ``(columns, rows)``.""" + """Run one statement and return its column names and rows: ``(columns, rows)``.""" async def query(self, sql: str, params: Sequence[SqlParam] | None = None) -> list[Row]: - """ - The rows of one statement, each a dict by column name; parameters bind - to a single statement. Columns are named as SQL names them (an - expression needs ``AS``), and two of the same name are an - ``SqlError``: alias them, or use ``values``. + """Run one statement and return its rows, each a dict keyed by column name. + + Columns are named as SQL names them, so an expression needs ``AS``; + two columns of the same name raise ``SqlError`` (alias one, or use + ``values``). """ return objects_of(*await self.query_columns(sql, params)) async def get(self, sql: str, params: Sequence[SqlParam] | None = None) -> Row | None: - """The first row of ``query``, or ``None`` when there is none.""" + """Return the first row of ``query``, or ``None`` if there is none.""" rows = await self.query(sql, params) return rows[0] if rows else None async def values(self, sql: str, params: Sequence[SqlParam] | None = None) -> Rows: - """The rows of one statement as lists, in column order: for adapters, and for columns that share a name.""" + """Run one statement and return its rows as lists, in column order.""" return (await self.query_columns(sql, params))[1] async def execute(self, sql: str, params: Sequence[SqlParam] | None = None) -> None: - """Run SQL (one or more statements), discarding rows; parameters only with a single statement.""" + """Run SQL and discard any rows; ``sql`` may hold several statements, but parameters bind only to one.""" await self.query_columns(sql, params) async def _request(self, call: Callable[[], T | Awaitable[T]]) -> T: - """ - ``call``'s result, or its ``_failure``: ``call`` makes the request, as - the extension raises as it is called (arguments it cannot convert, and - the error of a call that answers at once), and returns a value or an - awaitable of it. + """Make a request with ``call`` and return its result, translating errors with ``_failure``. + + ``call`` makes the request itself, because the extension may raise as + soon as it is called (for arguments it cannot convert, or a request + that fails right away); it returns a value or an awaitable of one. """ try: return await _settled(call()) @@ -287,9 +291,10 @@ async def _request(self, call: Callable[[], T | Awaitable[T]]) -> T: raise self._failure(error) from None def _failure(self, error: Exception) -> BaseException: - """ - A request's ``error``: an ``SqlError`` if a retry would repeat it; the - store's own error if that ended the session. + """Return the exception to raise for ``error``. + + That is an ``SqlError`` if a retry would fail the same way, or the + store's own error if that is what ended the session. """ if isinstance(error, SqlError): return error @@ -304,19 +309,20 @@ def _failure(self, error: Exception) -> BaseException: return error -# The transaction whose body is running, in that body's context (and the tasks -# it starts): the innermost, which knows the ones it runs inside. +# The innermost transaction whose body is running in this context (or in a task +# it started). It refers to the transactions it runs inside. _running: contextvars.ContextVar[ConnectionTransaction | None] = contextvars.ContextVar( "restate_sessions_running_transaction", default=None ) class Connection(_Sql): - """ - A session of a database: calls run one at a time, in call order, and a - transaction holds the session until it commits or rolls back. Calls made - meanwhile wait, but a call from inside the transaction's body is refused, - as it would wait for that body: the body uses its ``tx``. + """A session of a database. + + Calls run one at a time, in the order they were made. A transaction holds + the session until it commits or rolls back, and calls made meanwhile wait. + A call made from inside the transaction's own body is refused instead, as + it would wait forever; the body must use its ``tx``. """ def __init__(self, native: PySqlDatabase, boundary: _Boundary) -> None: @@ -326,7 +332,7 @@ def __init__(self, native: PySqlDatabase, boundary: _Boundary) -> None: @asynccontextmanager async def _serial(self) -> AsyncGenerator[None, None]: - """Run the body once the calls before it have, unless it comes from the body of a transaction it would wait for.""" + """Wait for earlier calls to finish, refusing calls from inside a transaction that holds this session.""" running = _running.get() if running is not None and running.holds(self): raise SqlError("inside a transaction, use the transaction handle its body gets") @@ -339,17 +345,16 @@ async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None @asynccontextmanager async def transaction(self) -> AsyncGenerator[ConnectionTransaction, None]: - """ - One SQL transaction, as the body of an ``async with``:: + """Run the body of an ``async with`` in one SQL transaction:: async with db.transaction() as tx: await tx.execute("UPDATE stock SET n = n - 1 WHERE id = ?", [7]) - It commits when the body ends, with the records it set, and rolls back - when the body raises, with the body's own error, once the calls it made - have completed. A statement in it that ends the transaction - (``COMMIT``, ``ROLLBACK``) ends the session. A failing COMMIT rolls - back. + The transaction commits, together with the records it set, when the + body ends. If the body raises, the transaction rolls back once the + calls it made have completed, and the body's exception propagates. A + statement that ends the transaction itself (``COMMIT``, ``ROLLBACK``) + ends the session. If the commit fails, the transaction rolls back. """ async with self._serial(): await self._request(self._native.begin) @@ -364,15 +369,15 @@ async def transaction(self) -> AsyncGenerator[ConnectionTransaction, None]: finally: _running.reset(token) records = await tx.end() - # Its commit's sync runs what the body gave `with_commit`; one that - # wrote nothing synced nothing, and it runs after. + # The commit's sync runs the callbacks given to `with_commit`. A commit + # that wrote nothing does not sync, so they run afterwards instead. self._boundary.with_sync = tx.commit_runs() try: await _settled(self._native.commit(records)) except Exception as error: - # A COMMIT that fails (a deferred constraint) leaves the - # transaction open: without the rollback, the calls after it - # would run inside it. + # A failed COMMIT (a deferred constraint, say) leaves the + # transaction open; without a rollback, later calls would run + # inside it. await self._rollback() raise self._failure(error) from None finally: @@ -381,39 +386,42 @@ async def transaction(self) -> AsyncGenerator[ConnectionTransaction, None]: unsynced() async def record(self, name: str) -> bytes: - """The record ``name``, as the last commit that set it left it (see ``ConnectionTransaction.set_record``); empty if none did.""" + """Return the committed value of the record ``name``, or empty bytes if it was never set. + + See ``ConnectionTransaction.set_record``. + """ async with self._serial(): return await self._request(lambda: self._native.record(name)) def is_usable(self) -> bool: - """Whether SQL can still run: an SQL error leaves the session usable, a storage error ends it.""" + """Return whether SQL can still run; an SQL error leaves the session usable, a storage error does not.""" return self._native.is_usable() async def save(self) -> None: - """ - End the session after the calls before it, leaving its engine to the - next one. It fails if they left a transaction open, or a failure ended - the session. + """End the session once earlier calls finish, handing its engine to the next session. + + Fails if those calls left a transaction open, or if a failure already + ended the session. """ async with self._serial(): await self._request(self._native.save) async def close(self) -> None: - """End the session now, whatever its state: queued calls fail.""" + """End the session immediately, whatever its state; queued calls fail.""" await _settled(self._native.close()) async def _rollback(self) -> None: - """Roll back, unless a storage error ended the session: the caller raises its own error, which says why.""" + """Roll back, ignoring errors; the caller raises its own error, which explains the failure.""" with contextlib.suppress(Exception): await _settled(self._native.rollback()) class ConnectionTransaction(_Sql): - """ - SQL in one transaction, while its body runs, and the records it commits - with. Its calls run one at a time, in call order, also when the body - awaits several together, or none: each call made while it is open runs - before its commit or rollback, so none outlives it. + """SQL inside one transaction, available while its body runs, and the records it commits. + + Calls run one at a time, in call order, even when the body awaits several + of them together or not at all. Every call made while the transaction is + open runs before it commits or rolls back, so none outlives it. """ def __init__( @@ -421,8 +429,8 @@ def __init__( ) -> None: super().__init__(native, boundary) self._connection: Connection = connection - # The transaction whose body this one's runs in, if any: of another - # connection, since a connection's own are refused. + # The transaction whose body this one runs in, if any. It belongs to + # another connection, since nesting on the same connection is refused. self._outer: ConnectionTransaction | None = outer self._records: dict[str, bytes] = {} self._with_commits: list[Callable[[], None]] = [] @@ -431,11 +439,11 @@ def __init__( @property def open(self) -> bool: - """Whether it is open: its body runs, and holds its connection.""" + """Whether the transaction is still open, holding its connection while its body runs.""" return not self._ended def holds(self, connection: Connection) -> bool: - """Whether this transaction, or one it runs inside, is open and holds ``connection``.""" + """Return whether this transaction, or one it runs inside, is open and holds ``connection``.""" if not self._ended and self._connection is connection: return True return self._outer is not None and self._outer.holds(connection) @@ -445,19 +453,20 @@ async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None return await self._request(lambda: self._native.query_columns(sql, params)) def set_record(self, name: str, value: bytes) -> None: - """ - Set the record ``name`` to ``value`` with this transaction, in the same - write: a whole value kept beside the SQL, in no table, stored exactly - when the transaction is, even if it changes no row. It stays until a - later transaction sets it again; an empty value reads as never set. A - name is at most 255 bytes, and not ``db``; a value at most 16 MiB. + """Set the record ``name`` to ``value``, committed atomically with this transaction. + + A record is a whole value stored next to the SQL tables, not in one. + It is written exactly when the transaction commits, even if no row + changed, and keeps its value until a later transaction sets it again. + An empty value reads as never set. Names are at most 255 bytes and may + not be ``db``; values are at most 16 MiB. """ self._refuse_ended() self._records[name] = bytes(value) async def record(self, name: str) -> bytes: - """The record ``name`` as this transaction would commit it: as it set it, or as committed.""" - # As set when it is called, in call order with `set_record`. + """Return the record ``name`` as this transaction would commit it.""" + # Read what was set at call time, to stay in call order with `set_record`. value = self._records.get(name) async with self._serial(): if value is not None: @@ -465,25 +474,25 @@ async def record(self, name: str) -> bytes: return await self._request(lambda: self._native.record(name)) def with_commit(self, run: Callable[[], None]) -> None: - """ - Run ``run`` with the commit, right after its writes and before it waits - for them to be stored: the store's host gets what ``run`` makes after - the writes, and before anything that fails the commit's wait (Restate's - sends, which it stores in order with the writes). Never if the - transaction rolls back. + """Call ``run`` during the commit, after its writes and before waiting for them to be stored. + + Whatever ``run`` hands the store's host (Restate sends, which are + stored in order with the writes) therefore follows the writes, and + comes before anything that could fail the wait. ``run`` is not called + if the transaction rolls back. """ self._refuse_ended() self._with_commits.append(run) async def end(self) -> dict[str, bytes]: - """End it, refusing calls from now on, once the calls made before have completed; then the records it set.""" + """Refuse further calls, wait for pending ones, and return the records this transaction set.""" self._ended = True async with self._lock: pass return dict(self._records) def commit_runs(self) -> Callable[[], None] | None: - """What runs with its commit, once it ended: see ``with_commit``.""" + """Return a function that runs the ``with_commit`` callbacks, or ``None`` if there are none.""" runs = self._with_commits if not runs: return None @@ -496,7 +505,7 @@ def run_all() -> None: @asynccontextmanager async def _serial(self) -> AsyncGenerator[None, None]: - """Run the body, a call made while the transaction is open, once the calls before it have.""" + """Wait for earlier calls of this transaction to finish, refusing calls once it has ended.""" self._refuse_ended() async with self._lock: yield diff --git a/python/restate/sessions/_link.py b/python/restate/sessions/_link.py index db69f3ae..9d404581 100644 --- a/python/restate/sessions/_link.py +++ b/python/restate/sessions/_link.py @@ -8,11 +8,12 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -A private link on the objects this package makes: set once, as it makes one, -and read by its code alone (a context's attempt, a database's attempt, a -transaction's reset, an object's entries). The object reads as its fields -alone, and its link goes with it. As sdk-sessions' ``link.ts``. +"""Private data attached to objects that this package creates. + +A ``Link`` attaches a value to an object (a context's attempt, a database's +attempt, a transaction's reset, an object's entries) without adding a field +to it. The value is set once, when the object is created, is read only by +this package, and is dropped together with the object. """ from __future__ import annotations @@ -25,11 +26,12 @@ class Link(typing.Generic[T]): - """ - A link of one kind, on the objects that have one: compared by identity, - kept as long as the object is. A value that refers to its owner (an - attempt, to its database) would keep both forever, as a link holds its - value: ``weak`` holds it weakly, for the owner to keep it alive. + """A mapping from objects, by identity, to one kind of attached value. + + Entries live as long as their owner. A value that refers back to its + owner (an attempt that refers to its database, say) would keep both alive + forever; pass ``weak=True`` to hold values weakly, so that the owner must + keep the value alive itself. """ def __init__(self, what: str, *, weak: bool = False) -> None: @@ -38,12 +40,12 @@ def __init__(self, what: str, *, weak: bool = False) -> None: self._links: weakref.WeakKeyDictionary[object, T | weakref.ref[typing.Any]] = weakref.WeakKeyDictionary() def on(self, owner: O, value: T) -> O: - """``owner``, linked to ``value``: once, as it is made.""" + """Attach ``value`` to ``owner`` and return ``owner``; call it once, when ``owner`` is created.""" self._links[owner] = weakref.ref(value) if self._weak else value return owner def of(self, owner: object) -> T | None: - """What ``owner`` is linked to, if it has this link (and None for anything that cannot have one).""" + """Return the value attached to ``owner``, or ``None`` if it has none (or cannot have one).""" try: value = self._links.get(owner) except TypeError: diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index 5a166559..0b6bac23 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -8,20 +8,22 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -Session objects: Restate virtual objects whose state is a SQL database per -key. Its handlers run one at a time per key, its views beside them; each gets a -context of its own, made of the key, the request, the database and calls, and -nothing else of the SDK's context:: +"""Session objects: Restate virtual objects whose state is a SQL database per key. + +Handlers run one at a time per key, and views run alongside them. Each gets a +context of its own that holds the key, the request, the database, and calls +and sends, and nothing else from the SDK's context:: notes = sessions.SessionObject("notes", migrations=["CREATE TABLE notes (text TEXT NOT NULL)"]) + @notes.handler async def add(ctx: sessions.SessionContext, text: str) -> None: await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [text]) - @notes.view - async def list(ctx: sessions.ViewContext) -> list[str]: + + @notes.view(name="list") + async def list_notes(ctx: sessions.ViewContext) -> list[str]: return [row["text"] for row in await ctx.db.query("SELECT text FROM notes ORDER BY rowid")] """ @@ -57,22 +59,23 @@ async def list(ctx: sessions.ViewContext) -> list[str]: F = typing.TypeVar("F", bound=Callable[..., Awaitable[Any]]) HEAD_KEY = "v22/head" -"""The key of the object database's head, its first read as it opens (shared-core's state format).""" +"""The state key of the database's head, the first thing read when it opens.""" SQL_METADATA = { "dev.restate.tooling.journal.hidden.names": "sync", "dev.restate.tooling.journal.hidden.name.prefix": "v22/", } -""" -The metadata every session object has: Restate's UI hides the database from -the journal view (its head, tail, manifests and packs in the object's state, -``v22/...``; and the ``sync`` run after its writes). +"""Metadata added to every session object. + +It tells Restate's UI to hide the database from the journal view: its head, +tail, manifests and packs (state keys under ``v22/``), and the ``sync`` run +after its writes. """ @dataclass(frozen=True, eq=False) class AttemptContext(Calls, Sends): - """``SessionContext`` (see ``_api.py``): one attempt's, made as it starts. Compared by identity.""" + """The ``SessionContext`` of one attempt, created when the attempt starts; compared by identity.""" key: str request: Request @@ -87,7 +90,7 @@ def cancel(self, invocation_id: str) -> None: @dataclass(frozen=True, eq=False) class AttemptViewContext(Calls): - """``ViewContext`` (see ``_api.py``): one attempt's, made as it starts. Compared by identity.""" + """The ``ViewContext`` of one attempt, created when the attempt starts; compared by identity.""" key: str request: Request @@ -100,10 +103,10 @@ def cancel(self, invocation_id: str) -> None: class _ContextHost: - """ - An attempt's session's host (``SqlHost``) over the SDK's context: its waits - go through the selector, and its state commands are named after their keys, - so that tooling can hide the database by name. + """A ``SqlHost`` over the SDK's context, for one attempt. + + Its waits go through the selector, and its state commands are named after + their keys, so that tooling can hide the database by name. """ def __init__(self, context: ServerInvocationContext, shared: bool) -> None: @@ -123,15 +126,15 @@ def on_attempt_finished() -> None: @property def vm(self) -> Any: - """ - The invocation's VM, which the database's store writes and reads - directly (see ``_database.Store``); none with a journal codec, whose - values only the context's state commands encode. + """The invocation's VM, which the database reads and writes directly (see ``_database.Store``). + + It is ``None`` when a journal codec is configured, since only the + context's state commands apply the codec. """ return None if self.context.journal_codec is not None else self.context.vm.vm def wait(self, handle: int) -> Awaitable[None]: - """Done once the VM has the read whose handle is ``handle``: through the selector, as every wait.""" + """Wait, through the selector, until the VM has the result of the read ``handle``.""" return self.selector.wait(SingleUnresolvedFuture(handle)) def cancel_invocation(self, invocation_id: str) -> None: @@ -147,7 +150,7 @@ def clear(self, key: str) -> None: self.context.clear(key, entry_name=key) def state_keys(self) -> Awaitable[list[str]]: - # Named under the format's prefix, which tooling hides. + # Named under the state format's prefix, which tooling hides. return self.context.state_keys(entry_name="v22/keys") def sync(self) -> Awaitable[Any]: @@ -203,14 +206,14 @@ def _notify_attempt_end(self) -> None: async def _nothing() -> None: - """The action of an empty run.""" + """Do nothing; the action of an empty run.""" def _cancel_of(host: SqlHost) -> Callable[[str], None]: - """``cancel`` over the SDK's context.""" + """Return the context's ``cancel``, implemented on ``host``.""" def cancel(invocation_id: str) -> None: - # Not an id: a retry would fail the same way. + # Terminal: a retry would fail the same way. if not invocation_id.startswith("inv"): raise TerminalError(f"not an invocation id: {invocation_id}") host.cancel_invocation(invocation_id) @@ -220,22 +223,22 @@ def cancel(invocation_id: str) -> None: @dataclass(frozen=True) class _Entry: - """A handler or view of a session object: its function, kind, and object's config.""" + """A registered handler or view of a session object.""" fn: Callable[..., Awaitable[Any]] view: bool config: ObjectConfig definition: weakref.ref[SessionObject] - """Its session object, weakly: the object's entries, which it holds, hold this.""" + """The session object, held weakly because the object holds its entries.""" explicit_cancellation: bool = False ENTRIES: Link[dict[str, _Entry]] = Link("entries") -"""A session object's handlers and views, by name: for ``invoke``.""" +"""A session object's handlers and views, by name, for ``invoke``.""" def _context_of(entry: _Entry, host: SqlHost, request: Request, attempt: Attempt) -> SessionContext | ViewContext: - """The context of an attempt of ``entry``.""" + """Create the context for an attempt of ``entry``.""" cancel = _cancel_of(host) sdk = typing.cast("ServerInvocationContext | None", getattr(host, "context", None)) scope = Scope(sdk=sdk, explicit_cancellation=entry.explicit_cancellation, definition=entry.definition()) @@ -256,7 +259,7 @@ def _context_of(entry: _Entry, host: SqlHost, request: Request, attempt: Attempt async def _attempt(entry: _Entry, host: SqlHost, request: Request, args: Sequence[Any]) -> Any: - """One attempt of ``entry`` over ``host``: its own context and session, which ends with it.""" + """Run one attempt of ``entry`` over ``host``, with a context and database session of its own.""" attempt = Attempt(host, entry.config) attempt.open_now() context = _context_of(entry, host, request, attempt) @@ -264,7 +267,7 @@ async def _attempt(entry: _Entry, host: SqlHost, request: Request, args: Sequenc def _sdk_handler(entry: _Entry) -> Callable[..., Awaitable[Any]]: - """``entry`` as the SDK's handler: one loop drives each attempt's awaits (the selector), until it ends.""" + """Wrap ``entry`` as an SDK handler whose attempts drive their awaits through a selector.""" @functools.wraps(entry.fn) async def handler(ctx: ServerInvocationContext, *args: Any) -> Any: @@ -275,7 +278,7 @@ async def handler(ctx: ServerInvocationContext, *args: Any) -> Any: def _explicit_of(ctx: SessionContext, operation: str) -> ServerInvocationContext: - """The SDK's context behind ``ctx``, a handler's with ``explicit_cancellation=True``.""" + """Return the SDK's context behind ``ctx``, which must be a handler's with ``explicit_cancellation=True``.""" scope = SCOPE.of(ctx) if scope is None or scope.sdk is None: raise TerminalError(f"{operation} takes a session object's ctx") @@ -285,20 +288,28 @@ def _explicit_of(ctx: SessionContext, operation: str) -> ServerInvocationContext async def cancelled(ctx: SessionContext) -> None: - """ - Returns when Restate cancels the invocation of ``ctx``, a handler's with - ``explicit_cancellation=True``, whose awaits are left as they are: for - code that stops its own work then (a model's request, say):: + """Wait until Restate cancels the invocation of ``ctx``. + + For handlers declared with ``explicit_cancellation=True``, whose awaits + are not interrupted by a cancellation, so that they can stop their own + work (a model request, for example):: async def stop_when_cancelled(): await internal.cancelled(ctx) await internal.cancel_calls(ctx) + asyncio.ensure_future(stop_when_cancelled()) - If the attempt ends first (the handler returned, or Restate suspended - it), it raises ``asyncio.CancelledError``: a task that waits for it ends - cancelled, quietly, as a task whose work is over does. + Args: + ctx: The context of a handler with ``explicit_cancellation=True``. + + Raises: + asyncio.CancelledError: If the attempt ends first (the handler + returned, or Restate suspended it), so that a task waiting on this + ends quietly, as cancelled. + TerminalError: If ``ctx`` is not the context of a handler with + ``explicit_cancellation=True``. """ try: await _preview.cancellation(_explicit_of(ctx, "cancelled")) @@ -307,18 +318,29 @@ async def stop_when_cancelled(): async def cancel_calls(ctx: SessionContext) -> list[str]: - """ - Cancel the calls the attempt of ``ctx`` made and that have not completed, - as a cancelled invocation's would be without explicit cancellation: a - handler's with ``explicit_cancellation=True``. Returns their invocation ids. + """Cancel the calls made by the attempt of ``ctx`` that have not completed yet. + + This is what a cancellation does automatically for handlers without + explicit cancellation. + + Args: + ctx: The context of a handler with ``explicit_cancellation=True``. + + Returns: + The invocation ids of the cancelled calls. + + Raises: + TerminalError: If ``ctx`` is not the context of a handler with + ``explicit_cancellation=True``. """ return await _preview.cancel_previous_calls(_explicit_of(ctx, "cancel_calls")) def definition_of(ctx: SessionContext | ViewContext) -> SessionObject | None: - """ - The session object whose handler or view ``ctx`` is the context of: what - a library keeps per object (an agent's options, say) lives on it. + """Return the session object whose handler or view ``ctx`` belongs to, or ``None``. + + Libraries use it to find what they keep per object, such as an agent's + options. """ scope = SCOPE.of(ctx) definition = scope.definition if scope is not None else None @@ -326,7 +348,23 @@ def definition_of(ctx: SessionContext | ViewContext) -> SessionObject | None: async def invoke(obj: SessionObject, name: str, host: SqlHost, request: Request, *args: Any) -> Any: - """Run handler or view ``name`` of ``obj`` for one attempt over ``host``, a context of a test's own, without Restate.""" + """Run one attempt of the handler or view ``name`` of ``obj`` on ``host``, without Restate. + + This is for tests, which provide a ``SqlHost`` of their own. + + Args: + obj: The session object. + name: The name of the handler or view. + host: The host that stores state and makes calls. + request: The request the context reports. + *args: The handler's input, if it takes one. + + Returns: + The handler's result. + + Raises: + ValueError: If ``obj`` has no handler or view named ``name``. + """ entry = (ENTRIES.of(obj) or {}).get(name) if entry is None: raise ValueError(f"{obj.name}.{name} is no handler or view of a session object") @@ -334,20 +372,42 @@ async def invoke(obj: SessionObject, name: str, host: SqlHost, request: Request, class SessionObject(VirtualObject): - """ - A virtual object whose state is a SQL database per key. - - - ``migrations``: the steps that make each key's database, in order, applied - once per key by the first handler that uses it. A step that ran never - changes: a change is a step of its own. - - ``@obj.handler``: handlers run one at a time per key, each with a - ``SessionContext``. - - ``@obj.view``: views run beside the handlers, each with a - ``ViewContext``: they read, and write nothing. - - A name is a handler or a view, not both. The other options are the SDK's - virtual object's, but for its state options: a session object's state is - its database's. + """A virtual object whose state is a SQL database per key. + + A session object offers: + + - A SQL database per key, ``ctx.db``, built by ``migrations``. + - Handlers (``@obj.handler``), which run one at a time per key and may + write to the database. + - Views (``@obj.view``), which run alongside the handlers and only read. + - Calls and sends to other handlers (``ctx.call``, ``ctx.send``), and + sends that commit with a transaction (``tx.send``). + + :: + + notes = sessions.SessionObject("notes", migrations=["CREATE TABLE notes (text TEXT NOT NULL)"]) + + Each name is either a handler or a view, not both. The remaining options + are those of ``restate.VirtualObject``, minus its state options, since a + session object's state is its database. + + Args: + name: The object's name, as Restate serves it. + migrations: The migrations that build each key's database, in order. + They are applied once per key, by the first handler that uses the + database. Never change a migration that has run; add a new one. + description: The object's description, shown in Restate's UI. + metadata: The object's metadata. + inactivity_timeout: Restate's inactivity timeout for its handlers. + abort_timeout: Restate's abort timeout for its handlers. + journal_retention: How long Restate keeps an invocation's journal. + idempotency_retention: How long Restate keeps the result of an + invocation with an idempotency key. + ingress_private: Whether only other services may call it. + invocation_retry_policy: How Restate retries a failed invocation. + + Raises: + TypeError: If ``migrations`` is a string rather than a list. """ def __init__( @@ -372,9 +432,9 @@ def __init__( abort_timeout=abort_timeout, journal_retention=journal_retention, idempotency_retention=idempotency_retention, - # An invocation starts with the database's head only, its first read as it opens: the rest it - # reads as it needs it (its tail entries and packs change keys with each flush: only the head - # can be named). As sdk-sessions' sqlOptions; restate-server 1.8 or later. + # Preload only the database's head, the first thing read when it opens; the rest is read on + # demand (the tail entries and packs change keys with each flush, so only the head has a fixed + # name). Needs restate-server 1.8 or later. state_preload=[HEAD_KEY], ingress_private=ingress_private, invocation_retry_policy=invocation_retry_policy, @@ -386,7 +446,7 @@ def __repr__(self) -> str: return f"{type(self).__name__}({self.name!r})" def _taken(self, name: str) -> str: - """Why ``name`` is taken already, for the error of a second handler of that name.""" + """Return the error message for registering ``name`` a second time.""" return f"{self.name} has a handler or view named {name!r} already: give this one another name" @typing.overload # type: ignore[override] @@ -428,19 +488,36 @@ def handler( # pyright: ignore[reportIncompatibleMethodOverride] invocation_retry_policy: InvocationRetryPolicy | None = None, explicit_cancellation: bool = False, ) -> F | Callable[[F], F]: - """ - A handler: it runs one at a time per key, and is the only one that - writes. Its function takes a ``SessionContext``, and its input, if any:: + """Register a handler, used bare (``@obj.handler``) or with options (``@obj.handler(...)``). + + Handlers run one at a time per key, and only handlers can write to the + database. The function takes a ``SessionContext`` and its input, if + any, and stays a plain function, so a test can call it directly:: @notes.handler async def add(ctx: sessions.SessionContext, text: str) -> None: await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [text]) - With options, ``@notes.handler(name=..., ...)``: the SDK's handler's. - With ``explicit_cancellation=True``, Restate's cancellation of its - invocation fails no await, nor cancels its calls: the handler learns of - it with ``internal.cancelled(ctx)``, and decides what to stop (its calls, - with ``internal.cancel_calls(ctx)``), of ``restate.sessions.internal``. + The options are those of ``restate.VirtualObject.handler``, plus + ``explicit_cancellation``. + + Args: + name: The handler's name. Defaults to the function's name. + explicit_cancellation: If true, cancelling the invocation neither + fails its awaits nor cancels its calls. The handler learns of + the cancellation with ``internal.cancelled(ctx)`` and decides + what to stop, for example its calls with + ``internal.cancel_calls(ctx)`` (from + ``restate.sessions.internal``). + + Returns: + The function itself, or a decorator if called with options. + + Raises: + TypeError: If the function is not an ``async def``, or is already + a handler. + ValueError: If the object already has a handler or view of that + name. """ return self._decorate( False, @@ -498,19 +575,32 @@ def view( ingress_private: bool | None = None, invocation_retry_policy: InvocationRetryPolicy | None = None, ) -> F | Callable[[F], F]: - """ - A view: it runs beside the handlers, and reads the database as - committed when its invocation started. Its function takes a - ``ViewContext``, and its input, if any:: + """Register a view, used bare (``@obj.view``) or with options (``@obj.view(...)``). + + Views run alongside the handlers and read the database as it was + committed when the view started. The function takes a + ``ViewContext`` and its input, if any:: @notes.view async def count(ctx: sessions.ViewContext) -> int: row = await ctx.db.get("SELECT count(*) AS n FROM notes") return row["n"] if row else 0 - A view that is still running after 16 commits of the handlers fails, - and Restate retries it. With options, ``@notes.view(name=..., ...)``: - the SDK's handler's. + A view still running after the handlers have committed 16 times + fails, and Restate retries it. The options are those of + ``restate.VirtualObject.handler``. + + Args: + name: The view's name. Defaults to the function's name. + + Returns: + The function itself, or a decorator if called with options. + + Raises: + TypeError: If the function is not an ``async def``, or is already + a handler. + ValueError: If the object already has a handler or view of that + name. """ return self._decorate( True, @@ -531,17 +621,13 @@ async def count(ctx: sessions.ViewContext) -> int: ) def _decorate(self, view: bool, name: str | F | None, options: dict[str, Any]) -> F | Callable[[F], F]: - """The decorator of ``@obj.handler(...)``, or, bare (``@obj.handler``), the function it registers.""" + """Return the decorator for ``@obj.handler(...)``, or register the function directly for a bare ``@obj.handler``.""" if callable(name): return self._register(view, None, options)(name) return self._register(view, name, options) def _attempt_fn(self, fn: Callable[..., Awaitable[Any]], view: bool) -> Callable[..., Awaitable[Any]]: - """ - What an attempt of handler (or view, with ``view``) ``fn`` runs: ``fn`` - itself; a subclass's may wrap it (an agent's turns a model provider's - refusal into a terminal error). - """ + """Return the function an attempt of ``fn`` runs; subclasses may wrap it (agents do).""" del view return fn @@ -573,16 +659,16 @@ def register(fn: F) -> F: _sdk_handler(entry) ) sdk: Handler[Any, Any] = handler_from_callable(registered) - # A later attempt makes different reads than a replayed journal - # would have: each runs from the start, against the state as it is. + # Attempts are not replayed: each runs from the start, against the + # current state, so it may read differently from earlier attempts. sdk.explicit_recovery = True sdk.explicit_cancellation = explicit entries = ENTRIES.of(self) assert entries is not None entries[sdk.name] = entry - # The function itself, marked as the handler it is: `ctx.call(fn, ...)` - # and the client find it, and a test calls it with a context of - # its own (a fake of the protocols). + # Return the function itself, marked as a handler: `ctx.call(fn, ...)` + # and the client can find it, and a test can call it with a fake + # context. vars(fn)[RESTATE_UNIQUE_HANDLER_SYMBOL] = sdk return fn diff --git a/python/restate/sessions/_schema.py b/python/restate/sessions/_schema.py index 7e034e28..6286fd68 100644 --- a/python/restate/sessions/_schema.py +++ b/python/restate/sessions/_schema.py @@ -8,28 +8,32 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -The schema of a session object's databases: each key's database is made by -the same steps, in order, applied once per key, as its first exclusive attempt -opens it. A database keeps how far each schema got in its record ``schema`` -(beside its tables, in none of them), committed with the steps, so a step -added later (a column, an index) runs on each key's database once, as that key -is next opened:: - - notes = sessions.SessionObject("notes", migrations=[ - "CREATE TABLE notes (text TEXT NOT NULL)", - "ALTER TABLE notes ADD COLUMN at TEXT", - ]) - -A step is a statement, or an async function of the transaction (to move data, -say). Steps are never changed or removed once they ran: a change is a step of -its own. - -A view applies the steps its key's database lacks in its own session, which -stores nothing, then reads: a key never written reads as empty tables. -Libraries (``restate.agents``) keep their tables in a schema of their own -name, which ``schema.ensure(ctx.db)`` applies where they use them, in any -object. +"""Migrations for the databases of session objects. + +Every key of an object has its own database, built by the same list of +migrations, applied in order. Each key's database is migrated the first time +a handler opens it, and records how many migrations it has applied (in a +record named ``schema``, stored outside its tables, and committed together +with the migrations). A migration added later, such as a new column or index, +therefore runs once on each key's database, the next time that key is used:: + + notes = sessions.SessionObject( + "notes", + migrations=[ + "CREATE TABLE notes (text TEXT NOT NULL)", + "ALTER TABLE notes ADD COLUMN at TEXT", + ], + ) + +A migration is a SQL statement, or an async function that takes the +transaction (to move data, for example). Never change or remove a migration +once it has run; add a new one instead. + +A view applies any missing migrations in a session of its own, which stores +nothing, before it reads, so a key that was never written reads as empty +tables. Libraries such as ``restate.agents`` keep their tables in a named +``Schema`` of their own and apply it with ``await schema.ensure(ctx.db)`` +wherever they use them, in any object. """ from __future__ import annotations @@ -42,21 +46,21 @@ from restate.sessions._database import Connection, ConnectionTransaction RECORD = "schema" -""" -The database's record of its schemas: how many steps of each it has, by name. -One record, so that ``reset`` clears all of them, a library's this process -never loaded too. +"""The name of the record that stores how many migrations of each schema were applied. + +It is a single record for all schemas, so that ``reset`` clears all of them, +including those of libraries that this process never loaded. """ Migration = str | Callable[[Any], Awaitable[None]] -"""One step of a database's migrations: a SQL statement, or an async function of the transaction.""" +"""One migration: a SQL statement, or an async function that takes the transaction.""" _ENSURE = "_restate_sessions_ensure" -"""The attribute of a database of this package's: how it applies schemas (see ``Schema.ensure``).""" +"""The attribute through which this package's databases apply schemas (see ``Schema.ensure``).""" def steps_of(steps: Sequence[Migration]) -> tuple[Migration, ...]: - """Return ``steps`` as a tuple; raise ``TypeError`` for a lone string, which is a sequence of characters.""" + """Return ``steps`` as a tuple, rejecting a lone string (which is a sequence of characters).""" if isinstance(steps, str): shown = f"{steps[:40]!r}{'...' if len(steps) > 40 else ''}" raise TypeError(f"migrations takes a list of steps, not a string: migrations=[{shown}]") @@ -64,59 +68,76 @@ def steps_of(steps: Sequence[Migration]) -> tuple[Migration, ...]: class Schema: - """A part of a database's schema: its ``name``, and the ``steps`` that make it, in order.""" + """A named part of a database's schema, built by its own list of migrations. + + Libraries use this for the tables they own, so that their migrations are + tracked separately from the object's own ``migrations``:: + + AUDIT = Schema("mylib.audit", ["CREATE TABLE mylib_audit (at TEXT NOT NULL, what TEXT NOT NULL)"]) + + Args: + name: The schema's name, unique within a database. + steps: Its migrations, in order. + + Raises: + TypeError: If ``steps`` is a string rather than a list of migrations. + """ def __init__(self, name: str, steps: Sequence[Migration]) -> None: self.name: str = name self.steps: tuple[Migration, ...] = steps_of(steps) async def ensure(self, db: Any) -> None: - """ - Apply the steps ``db`` lacks: a library's tables, where it uses them, in - any object. On a handler's ``ctx.db``, once per attempt; a view's are - its own, vanishing with its session. + """Apply the migrations of this schema that ``db`` has not applied yet. + + On a handler's ``ctx.db``, the migrations are committed and checked at + most once per attempt. On a view's ``ctx.db``, they are applied to the + view's own session and discarded with it. + + Args: + db: A session object's ``ctx.db``, or a database from ``database``. """ ensure = ensure_of(db) await (ensure([self]) if ensure is not None else apply(db, [self])) Ensure = Callable[[Sequence[Schema]], Awaitable[None]] -"""How a database applies schemas: a handler's ``ctx.db``, its reader, or one of ``as_database``.""" +"""A function that applies schemas to a database (a handler's ``ctx.db``, a view's, or one from ``as_database``).""" def ensuring(db: Any, ensure: Ensure) -> Any: - """``db``, which applies schemas with ``ensure`` (see ``Schema.ensure``).""" + """Make ``db`` apply schemas with ``ensure``, and return it.""" object.__setattr__(db, _ENSURE, ensure) return db def ensure_of(db: Any) -> Ensure | None: - """How ``db`` applies schemas, if it is a database of this package's.""" + """Return the function ``db`` applies schemas with, or ``None`` if it is not one of this package's databases.""" return getattr(db, _ENSURE, None) async def _applied_in(db: Connection | ConnectionTransaction) -> dict[str, int]: - """How many steps of each schema the database has, by name.""" + """Return how many migrations of each schema the database has applied, by schema name.""" record = await db.record(RECORD) return json.loads(record.decode()) if record else {} def _behind(applied: dict[str, int], schemas: Sequence[Schema]) -> list[Schema]: - """The schemas of ``schemas`` that ``applied`` lacks steps of.""" + """Return the schemas that have migrations missing from ``applied``.""" return [schema for schema in schemas if applied.get(schema.name, 0) < len(schema.steps)] async def apply(db: Connection, schemas: Sequence[Schema], *, reader: bool = False) -> None: - """Apply the steps of ``schemas`` that ``db`` lacks, in one transaction, with how far each got.""" + """Apply the missing migrations of ``schemas`` to ``db`` in one transaction, and record them.""" if not _behind(await _applied_in(db), schemas): return - # Again in the transaction, which holds the session: a call before it may - # have applied them. + # Check again inside the transaction, which holds the session: another call + # may have applied them in the meantime. async with db.transaction() as tx: if not reader: await apply_in(tx, schemas) return - # A reader's session is query-only, but for its schemas. + # A reader's session is query-only, except while it applies schemas. await tx.execute("PRAGMA query_only = 0") try: await apply_in(tx, schemas) @@ -125,7 +146,7 @@ async def apply(db: Connection, schemas: Sequence[Schema], *, reader: bool = Fal async def apply_in(tx: ConnectionTransaction, schemas: Sequence[Schema]) -> None: - """Apply the steps of ``schemas`` that the transaction ``tx`` lacks, and record how far each got.""" + """Apply the missing migrations of ``schemas`` inside the transaction ``tx``, and record them.""" applied = await _applied_in(tx) lacking = _behind(applied, schemas) if not lacking: @@ -138,14 +159,14 @@ async def apply_in(tx: ConnectionTransaction, schemas: Sequence[Schema]) -> None async def drop_all(tx: ConnectionTransaction, schemas: Sequence[Schema]) -> None: - """ - Empty the database in the transaction ``tx``: drop every table, view, index - and trigger, and forget every schema, then apply ``schemas`` again. Its - other records (the outbox) stay. It commits with ``tx``: what ``tx`` writes - after it are the new database's first rows. + """Empty the database inside the transaction ``tx``, then apply ``schemas`` again. + + Drops every table, view, index and trigger and forgets every schema; other + records, such as the outbox, are kept. The change commits with ``tx``, so + whatever ``tx`` writes afterwards becomes the new database's first rows. """ made = await tx.values("SELECT type, name FROM sqlite_schema WHERE substr(name, 1, 7) != 'sqlite_'") - # A trigger or view may name a table; an index goes with its table. + # Triggers and views may refer to tables, so drop them first; indexes go with their tables. for kind in ("trigger", "view"): for made_kind, name in made: if made_kind == kind: @@ -163,10 +184,10 @@ async def _drop(tx: ConnectionTransaction, kind: str, name: str) -> None: async def _children_first(tx: ConnectionTransaction, tables: list[str]) -> list[str]: - """ - ``tables``, each after every table whose foreign keys name it: with foreign - keys on, dropping a table deletes its rows, which a child's rows may not - outlive. Tables in a cycle stay in their order. + """Order ``tables`` so that each comes after every table whose foreign keys refer to it. + + With foreign keys enforced, dropping a parent table deletes its rows, which + its children's rows may not outlive. Tables in a cycle keep their order. """ parents: dict[str, set[str]] = {} for name in tables: diff --git a/python/restate/sessions/_scope.py b/python/restate/sessions/_scope.py index b2dafac1..58b1cb57 100644 --- a/python/restate/sessions/_scope.py +++ b/python/restate/sessions/_scope.py @@ -8,11 +8,12 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -An attempt's scope: what its code keeps for the attempt (an agent's turn, the -calls of its tools...), which its context holds and which ends with it. A -retry gets a context, and so a scope, of its own. A module holds only the key -of what it keeps (``ScopeKey``), never the state. As sdk-sessions' ``scope.ts``. +"""Per-attempt state, kept on the attempt's context. + +Code running in an attempt can keep state for the duration of that attempt +(an agent's turn, the calls its tools made, ...). The state lives on the +context and ends with it; a retry gets a new context and starts empty. +Modules hold only a ``ScopeKey``, never the state itself. """ from __future__ import annotations @@ -31,50 +32,58 @@ @dataclass class Scope: - """An attempt's own, as this package made it, on its context (``SCOPE``), and what its code keeps for it.""" + """The per-attempt data this package attaches to a context (see ``SCOPE``).""" sdk: ServerInvocationContext | None - """The SDK's context of the attempt: none over a test's own host.""" + """The SDK's context of the attempt, or ``None`` when running on a test's host.""" explicit_cancellation: bool = False """Whether its handler has ``explicit_cancellation=True``.""" definition: object | None = None - """The session object whose handler or view it runs (``definition_of``).""" + """The session object whose handler or view the attempt runs (see ``definition_of``).""" values: dict[ScopeKey[Any], Any] = field(default_factory=dict[Any, Any]) - """What its code keeps for it, by key: only ``ScopeKey`` reaches it.""" + """The state kept for the attempt, by key; only ``ScopeKey`` accesses it.""" SCOPE: Link[Scope] = Link("scope") -"""The attempt's scope, on its context.""" +"""The attempt's scope, attached to its context.""" class ScopeKey(typing.Generic[T]): - """ - A kind of state an attempt's code keeps, by its context: its key in the - attempt's scope. Kept per attempt, so a retry starts without it:: + """A key for one kind of state kept per attempt, on the attempt's context. + + The state is not carried over to a retry, which starts without it:: TURN: ScopeKey[Turn] = ScopeKey("an agent's turn") turn = TURN.of(ctx, lambda: Turn(...)) + + Args: + what: A description of the state, used in error messages. """ def __init__(self, what: str) -> None: self.what: str = what - """What it is, for an error.""" + """A description of the state, used in error messages.""" def get(self, ctx: SessionContext | ViewContext) -> T | None: - """What the attempt of ``ctx`` keeps under this key, if anything.""" + """Return the value kept under this key for the attempt of ``ctx``, or ``None``.""" return typing.cast("T | None", self._values(ctx).get(self)) def set(self, ctx: SessionContext | ViewContext, value: T) -> None: - """Keep ``value`` under this key, for the attempt of ``ctx``.""" + """Keep ``value`` under this key for the attempt of ``ctx``.""" self._values(ctx)[self] = value def delete(self, ctx: SessionContext | ViewContext) -> None: - """Keep nothing under this key, for the attempt of ``ctx``.""" + """Remove the value kept under this key for the attempt of ``ctx``, if any.""" self._values(ctx).pop(self, None) def of(self, ctx: SessionContext | ViewContext, make: Callable[[], T]) -> T: - """What the attempt of ``ctx`` keeps under this key, made by ``make`` the first time.""" + """Return the value kept under this key for the attempt of ``ctx``, creating it with ``make`` if absent. + + Raises: + TypeError: If ``ctx`` is not the context of a session object's + handler or view. + """ values = self._values(ctx) if self not in values: values[self] = make() diff --git a/python/restate/sessions/_selector.py b/python/restate/sessions/_selector.py index 5b1f7777..8fb1185c 100644 --- a/python/restate/sessions/_selector.py +++ b/python/restate/sessions/_selector.py @@ -8,16 +8,15 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -One loop that drives every pending wait of an invocation without journal replay. - -Without journal replay, the handler's futures are plain awaitables that user -code awaits and combines as it likes: ``asyncio.gather``, ``asyncio.wait``, -its own tasks, or a library's. The protocol has one await point per -invocation, though: a suspension, the awaiting-on hint and a cancellation each -take one unresolved future. So every wait is registered with the selector, -which asks the VM for progress on the first of all of them to complete, and is -the only one reading the invocation's input. +"""A single loop that drives every pending wait of an invocation without journal replay. + +Without journal replay, a handler's futures are plain awaitables that user +code awaits and combines however it likes: with ``asyncio.gather``, +``asyncio.wait``, its own tasks, or a library's. The protocol, however, has +one await point per invocation: a suspension, the awaiting-on hint and a +cancellation each take a single unresolved future. So every wait registers +with the selector, which asks the VM for progress on whichever of them +completes first, and which is the only reader of the invocation's input. """ from __future__ import annotations @@ -44,47 +43,47 @@ class SelectorDriver(Protocol): - """What the selector needs of the invocation.""" + """The operations of the invocation that the selector uses.""" def do_progress(self, future: UnresolvedFuture) -> Any: - """The VM's progress on ``future``: a DoProgress* response, Suspended, or an exception.""" + """Return the VM's progress on ``future``: a DoProgress* response, ``Suspended``, or an exception.""" ... def is_completed(self, handle: int) -> bool: - """Whether the VM has the result of ``handle``.""" + """Return whether the VM has the result of ``handle``.""" ... async def flush_output(self) -> None: - """Writes out what the VM has buffered.""" + """Write out whatever the VM has buffered.""" ... async def receive(self) -> Mapping[str, Any]: - """The next input event: a chunk of the request body, a disconnect, or a restate.* event.""" + """Return the next input event: a chunk of the request body, a disconnect, or a ``restate.*`` event.""" ... def wake(self) -> None: - """Makes ``receive`` return a restate.* event.""" + """Make ``receive`` return a ``restate.*`` event.""" ... def execute_run(self, handle: int) -> None: - """Starts the run of ``handle``, which signals its completion through ``receive``.""" + """Start the run of ``handle``, which signals its completion through ``receive``.""" ... def notify_input(self, chunk: Mapping[str, Any]) -> None: - """Passes a chunk of the request body to the VM.""" + """Pass a chunk of the request body to the VM.""" ... def track(self, task: asyncio.Future[Any]) -> None: - """Cancels ``task`` at the end of the attempt.""" + """Cancel ``task`` when the attempt ends.""" ... def on_end(self) -> None: - """The attempt cannot go on: the selector has failed every wait.""" + """Handle the end of the attempt, after the selector has failed every wait.""" ... def is_ready(future: UnresolvedFuture, completed: Callable[[int], bool]) -> bool: - """Whether ``future`` needs no more progress.""" + """Return whether ``future`` needs no more progress.""" if isinstance(future, SingleUnresolvedFuture): return completed(future.handle) if isinstance(future, FirstCompletedUnresolvedFuture): @@ -93,15 +92,15 @@ def is_ready(future: UnresolvedFuture, completed: Callable[[int], bool]) -> bool class FutureSelector: - """ - The one loop of an invocation without journal replay. - - A wait stays registered until it is ready. The loop settles whatever is - ready, then asks the VM for progress on the first of all the pending waits - to complete: so a suspension waits for any of them, a pending run starts - whichever wait reaches it first, and a cancellation reaches all of them. A - wait registered while the loop reads the input wakes it, so it is part of - the next check. The loop runs while something is pending. + """The single loop that drives the waits of an invocation without journal replay. + + A wait stays registered until it is ready. Each pass, the loop resolves + the waits that are ready, then asks the VM for progress on whichever + pending wait completes first. So a suspension waits for any of them, a + pending run starts whichever wait reaches it first, and a cancellation + reaches all of them. Registering a wait while the loop reads input wakes + the loop, so the new wait is included in the next check. The loop runs as + long as something is pending. """ def __init__(self, driver: SelectorDriver, disconnected: Callable[[], BaseException]) -> None: @@ -109,23 +108,23 @@ def __init__(self, driver: SelectorDriver, disconnected: Callable[[], BaseExcept self._disconnected: Callable[[], BaseException] = disconnected self._waits: dict[asyncio.Future[None], UnresolvedFuture] = {} self._task: asyncio.Task[None] | None = None - # Between the VM's WaitExternalProgress and the next input. + # True between the VM's WaitExternalProgress and the next input. self._waiting: bool = False - # Why the attempt ended, once it has: each wait then raises a new one. + # Why the attempt ended, once it has; each later wait raises a fresh instance. self._end: Callable[[], BaseException] | None = None @property def ended(self) -> BaseException | None: - """Why the attempt ended (a suspension, an error, a disconnect), or None while it goes on.""" + """Why the attempt ended (a suspension, an error, a disconnect), or ``None`` while it continues.""" return None if self._end is None else self._end() @property def size(self) -> int: - """The waits still pending, for tests.""" + """The number of waits still pending, for tests.""" return len(self._waits) async def wait(self, future: UnresolvedFuture) -> None: - """Returns once ``future`` is ready; raises if the attempt ends first, or is cancelled.""" + """Wait until ``future`` is ready; raise if the attempt ends first or is cancelled.""" if self._end is not None: raise self._end() if is_ready(future, self._driver.is_completed): @@ -143,8 +142,8 @@ def _start(self) -> None: self._task = asyncio.create_task(self._loop()) self._driver.track(self._task) elif self._waiting: - # Only a loop waiting for input needs waking, once: any other picks - # up the new wait on its next pass. + # Only a loop that waits for input needs waking, and only once; otherwise + # the loop picks up the new wait on its next pass. self._waiting = False self._driver.wake() @@ -159,7 +158,7 @@ def _fail_all(self, error: Callable[[], BaseException]) -> None: waiter.set_exception(error()) def _finish(self, end: Callable[[], BaseException]) -> None: - """The attempt cannot go on: every wait, now and later, raises ``end()``.""" + """End the attempt: every wait, current or future, raises ``end()``.""" self._end = end self._fail_all(end) self._driver.on_end() @@ -169,9 +168,9 @@ async def _loop(self) -> None: try: while True: self._settle() - # Let what that settled run, and register what it waits for - # next, so that one check covers them. An optimization: what - # registers later wakes the loop anyway. + # Let the code that was just resolved run and register its next + # waits, so that one check covers them. This is only an + # optimization: a later registration wakes the loop anyway. await asyncio.sleep(0) pending = [future for waiter, future in self._waits.items() if not waiter.done()] if not pending: @@ -190,7 +189,7 @@ async def _loop(self) -> None: if isinstance(response, DoProgressAnyCompleted): continue if isinstance(response, DoProgressCancelSignalReceived): - # Reaches every pending wait; what waits after it is not cancelled. + # Cancel every pending wait; waits registered afterwards are not cancelled. self._fail_all(lambda: TerminalError("cancelled", 409)) continue if isinstance(response, DoProgressExecuteRun): @@ -220,7 +219,7 @@ async def _loop(self) -> None: def _internal(cause: BaseException) -> Callable[[], BaseException]: - """New SdkInternalExceptions, caused by ``cause``.""" + """Return a factory of ``SdkInternalException`` instances caused by ``cause``.""" def internal() -> BaseException: error = SdkInternalException() @@ -230,16 +229,17 @@ def internal() -> BaseException: return internal -# The selector takes over an invocation of the SDK's (``hook``): the waits of -# its futures, which the SDK makes through ``create_poll_or_cancel_coroutine``, -# go to the selector, and each future starts when it is made. It reaches into -# ``ServerInvocationContext``, of this same package, whose tests cover it. +# ``hook`` makes the selector take over an SDK invocation: the waits of its +# futures, which the SDK makes through ``create_poll_or_cancel_coroutine``, go +# to the selector, and each future starts as soon as it is created. This +# reaches into ``ServerInvocationContext``, from the same package, whose tests +# cover it. # TODO: once sessions run on it, the context could drive its waits through -# the selector itself, rather than be hooked. +# the selector itself, instead of being hooked. class _ContextDriver: - """What the selector needs of the SDK's context.""" + """A ``SelectorDriver`` over the SDK's invocation context.""" def __init__(self, context: ServerInvocationContext, on_end: Callable[[], None]) -> None: self.context: ServerInvocationContext = context @@ -258,7 +258,7 @@ async def receive(self) -> Mapping[str, Any]: return await self.context.receive() def wake(self) -> None: - # Any restate.* event: the loop takes it as a wake-up. + # Any restate.* event will do: the loop treats it as a wake-up. event = self.context.receive.enqueue_restate_event({"type": "restate.run_completed", "data": None}) self.track(asyncio.ensure_future(event)) @@ -287,11 +287,11 @@ def track(self, task: asyncio.Future[Any]) -> None: def hook(context: ServerInvocationContext, on_end: Callable[[], None]) -> FutureSelector: - """ - Make the waits of ``context``, an invocation without journal replay, go - through one selector, which calls ``on_end`` if it fails them all; each - future starts when it is made, whether or not anything awaits it. The - hook lasts as long as ``context``, one attempt. + """Route the waits of ``context``, an invocation without journal replay, through one selector. + + Each future starts as soon as it is created, whether or not anything + awaits it. The selector calls ``on_end`` if it has to fail every wait. + The hook lasts as long as ``context``, which is one attempt. """ selector = FutureSelector(_ContextDriver(context, on_end), DisconnectedException) started: set[asyncio.Task[Any]] = set() @@ -302,7 +302,7 @@ def make_started(*args: Any, **kwargs: Any) -> Any: lazy = future.future if lazy.task is None: lazy.task = asyncio.create_task(lazy.async_def()) - # Its error, if nothing awaits it, is dropped. + # If nothing awaits it, its error is dropped. lazy.task.add_done_callback(_retrieve) started.add(lazy.task) lazy.task.add_done_callback(started.discard) diff --git a/python/restate/sessions/_serve.py b/python/restate/sessions/_serve.py index 80288094..baa2e6ac 100644 --- a/python/restate/sessions/_serve.py +++ b/python/restate/sessions/_serve.py @@ -8,11 +8,12 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -``app``: the ASGI app of a program's services, with what they need beside them -(a library's objects, such as an agent's ``signals``). As sdk-sessions' -``serve.ts``; serving it is the program's (hypercorn, uvicorn...), as for -``restate.app``, or ``serve``'s. +"""Serving session objects. + +``app`` builds the ASGI app of a program's services, together with any +services they depend on (such as the ``signals`` object that agents use). Run +it with an ASGI server of your choice (hypercorn, uvicorn, ...), as you would +``restate.app``, or call ``serve`` to run it in one line. """ from __future__ import annotations @@ -34,23 +35,44 @@ D = typing.TypeVar("D") _NEEDS: Link[Callable[[], Sequence[Service]]] = Link("needs") -"""The services a definition needs served beside it: see ``needing``.""" +"""The services a definition needs served next to it (see ``needing``).""" def needing(definition: D, services: Callable[[], Sequence[Service]]) -> D: - """ - ``definition``, which needs ``services`` served beside it: ``app`` serves - them too. For a library whose objects call one of their own (an agent's - ``signals``, say). ``services`` is called once ``app`` asks. + """Declare that ``definition`` needs ``services`` served next to it, and return it. + + ``app`` then serves those services too. This is for libraries whose + objects call a service of their own, such as the ``signals`` object of + agents. + + Args: + definition: The service, virtual object or workflow that depends on + ``services``. + services: A function returning the services needed. It is called + lazily, when ``app`` collects the services to serve. + + Returns: + ``definition`` itself. """ return _NEEDS.on(definition, services) def served(services: Iterable[Service], names: str | None = None) -> list[Service]: - """ - What ``app`` serves of ``services``: each with what it needs beside it, - once, of those ``names`` names (comma separated; all of them without). A - name of none is an error. + """Return the services that ``app`` would serve for ``services``. + + Each service is followed by the services it needs (see ``needing``), + transitively, and each appears once. + + Args: + services: The services to serve. + names: A comma-separated list of service names to keep; if empty or + ``None``, all of them are kept. + + Returns: + The services to serve, in order. + + Raises: + ValueError: If ``names`` contains a name that matches no service. """ every: dict[str, Service] = {} @@ -79,13 +101,30 @@ def app( identity_keys: list[str] | None = None, journal_value_codec: JournalValueCodec | JournalValueCodecProvider | None = None, ) -> RestateAppT: - """ - The ASGI app of ``services``, with the services they need beside them, - each once: ``restate.app``'s, whose options it takes. ``RESTATE_SERVE_ONLY`` - (names separated by commas, such as ``triage,planner,signals``) serves only - those of them: one program (or image), deployed in parts:: + """Return the ASGI app that serves ``services`` and the services they need. + + This is ``restate.app`` with the dependencies declared by ``needing`` + added, and it takes the same options. Run it with any ASGI server:: + + app = sessions.app(notes, orders) # then: hypercorn module:app + + If the ``RESTATE_SERVE_ONLY`` environment variable is set to a + comma-separated list of names (such as ``triage,planner,signals``), only + those services are served. This lets one program, or one container image, + be deployed in parts. + + Args: + *services: The services, virtual objects (session objects included) + and workflows to serve. + protocol: Passed to ``restate.app``. + identity_keys: Passed to ``restate.app``. + journal_value_codec: Passed to ``restate.app``. - app = sessions.app(notes, orders) # hypercorn module:app + Returns: + The ASGI app. + + Raises: + ValueError: If ``RESTATE_SERVE_ONLY`` names a service that is not served. """ return restate.app( served(services, os.environ.get("RESTATE_SERVE_ONLY")), @@ -101,18 +140,28 @@ def serve( identity_keys: list[str] | None = None, journal_value_codec: JournalValueCodec | JournalValueCodecProvider | None = None, ) -> None: - """ - Serve ``services`` (``app``'s: with what they need beside them, of those - ``RESTATE_SERVE_ONLY`` names) on ``port``, or PORT, or 9080, until the - process stops: a program of session objects and agents in one line:: + """Serve ``services`` over HTTP until the process is stopped. + + This runs ``app(*services)`` on uvicorn with cleartext HTTP/2, which is + how Restate connects to services:: if __name__ == "__main__": sessions.serve(wallet, handle) - It runs them on uvicorn's HTTP/2 (cleartext, as Restate calls them), - imported here: install the ``serve`` extra (uvicorn 0.53 or later, and - zttp; Python 3.11 or later), or serve ``app``'s ASGI app with a server of - your own. It stops on SIGINT or SIGTERM, within a few seconds. + It needs the ``serve`` extra (``pip install "restate-sdk[serve]"``, which + installs uvicorn 0.53 or later and zttp) and Python 3.11 or later. On + other setups, serve ``app(...)`` with an ASGI server of your own. It stops + within a few seconds of SIGINT or SIGTERM. + + Args: + *services: The services, virtual objects (session objects included) + and workflows to serve, as for ``app``. + port: The port to listen on. Defaults to ``$PORT``, else 9080. + identity_keys: Passed to ``restate.app``. + journal_value_codec: Passed to ``restate.app``. + + Raises: + ImportError: If uvicorn 0.53 or later, or zttp, is not installed. """ try: uvicorn = importlib.import_module("uvicorn") @@ -141,4 +190,4 @@ def serve( ) _GRACEFUL_SHUTDOWN = 3 -"""Seconds a stopping ``serve`` waits for the requests in flight, which Restate retries elsewhere.""" +"""How many seconds a stopping ``serve`` waits for requests in flight; Restate retries the rest.""" diff --git a/python/restate/sessions/_session.py b/python/restate/sessions/_session.py index 336ee2b2..c9f50757 100644 --- a/python/restate/sessions/_session.py +++ b/python/restate/sessions/_session.py @@ -8,11 +8,12 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -An attempt's database: the session's database, kept in the object's state -(``_database.py``), opened as an exclusive attempt starts (a view's on its -first read) and saved when the handler ends; and its transactions, with the -sends each commits with (``tx.send``). +"""The database of one attempt, and its transactions. + +The database itself is stored in the object's state (see ``_database.py``). +An exclusive handler's attempt opens it when it starts, a view's on its +first read, and the attempt saves it when the handler ends. Transactions +commit together with the sends made through ``tx.send``. """ from __future__ import annotations @@ -56,11 +57,14 @@ class SqlHost(CallHost, Protocol): - """What an attempt's session needs of the handler's context.""" + """The host that an attempt's database runs on: the handler's context, or a test's fake. + + Implement it to run handlers without Restate, with ``internal.invoke``. + """ @property def shared(self) -> bool: - """A view's: it reads beside the exclusive handler, and stores nothing.""" + """Whether this is a view's attempt, which reads alongside the exclusive handler and stores nothing.""" ... @property @@ -74,74 +78,75 @@ def key(self) -> str: ... def get(self, key: str) -> Awaitable[bytes | None]: - """A state read, issued now.""" + """Read the state ``key``; the read is issued immediately.""" ... def set(self, key: str, value: bytes) -> None: - """A state write.""" + """Write the state ``key``.""" ... def clear(self, key: str) -> None: - """A state clear.""" + """Clear the state ``key``.""" ... def state_keys(self) -> Awaitable[list[str]]: - """The object's state keys, issued now.""" + """List the object's state keys; the request is issued immediately.""" ... def sync(self) -> Awaitable[Any]: - """An empty run: done once every command before it is stored.""" + """Make an empty run, which completes once every earlier command is stored.""" ... def attempt_end(self) -> BaseException | None: - """Why the attempt ended (a suspension, say), or None while it goes on.""" + """Return why the attempt ended (a suspension, say), or ``None`` while it continues.""" ... def on_attempt_end(self, listener: Callable[[], None]) -> Callable[[], None]: - """Call ``listener`` when the attempt ends; the returned function cancels it.""" + """Call ``listener`` when the attempt ends, and return a function that unregisters it.""" ... @property def vm(self) -> PyVM | None: - """The invocation's VM, which the store writes and reads directly; none without one (a test's host).""" + """The invocation's VM, which the database reads and writes directly, or ``None`` (as on a test's host).""" ... def wait(self, handle: int) -> Awaitable[None]: - """Done once the VM has the read whose handle is ``handle``.""" + """Wait until the VM has the result of the read ``handle``.""" ... def cancel_invocation(self, invocation_id: str) -> None: - """Cancel the invocation ``invocation_id``, as Restate's admin API does.""" + """Cancel the invocation ``invocation_id``.""" ... @dataclass class ObjectConfig: - """What a session object gives its attempts: its schema, and a wrapper of its store (fault injection, in tests).""" + """Per-object settings shared by its attempts. + + Attributes: + schema: The object's migrations, as a ``Schema``. + wrap_store: A function that wraps the store, for fault injection in tests. + """ schema: Schema | None = None wrap_store: Callable[[Store], Store] | None = field(default=None) -# --- Errors --------------------------------------------------------------------------------------- - - # --- The object's state as a store ---------------------------------------------------------------- def _beside(name: str, key: str) -> str: - """ - The key of the database ``name``, apart from the object's own: each of the - format's keys (``v22/head``, ``v22/tail/...``) as ``v22/db//head``, - under the format's prefix still, which tooling hides, and under none of the - object's database's. + """Return the state key of ``key`` for the separate database ``name``. + + ``v22/head`` becomes ``v22/db//head``: still under the format's + prefix, which tooling hides, and clear of the object's own database's keys. """ at = key.index("/") + 1 return f"{key[:at]}db/{name}/{key[at:]}" def _unbeside(name: str, key: str) -> str | None: - """The format's key of a key of the database ``name`` (undoes ``_beside``); None for any other key of the object.""" + """Undo ``_beside``: return the format's key, or ``None`` if ``key`` is not in the database ``name``.""" at = key.find("/") + 1 prefix = f"db/{name}/" if at > 0 and key.startswith(prefix, at): @@ -150,14 +155,15 @@ def _unbeside(name: str, key: str) -> str | None: class _ObjectStore: - """ - The object's state as an ordered store: ``get`` is a state get, ``set`` and - ``clear`` are state writes, which return at once and which the runtime - stores in order; ``sync`` is an empty run, done once the commands before it - are stored. A ``get``, ``wait`` or ``sync`` the attempt's end leaves - waiting raises, so the session that made it ends instead of holding its - lock. With the invocation's VM (``vm``), the extension writes and reads it - directly, and waits for a read with ``wait``. + """The object's state as a ``Store``. + + ``get`` reads the state; ``set`` and ``clear`` write it, returning at once, + and Restate stores the writes in order; ``sync`` is an empty run, which + completes once the commands before it are stored. A ``get``, ``wait`` or + ``sync`` still pending when the attempt ends raises, so that the session + ends instead of holding its lock. With the invocation's VM (``vm``), the + extension reads and writes the VM directly, and waits for reads with + ``wait``. """ def __init__(self, host: SqlHost, name: str | None) -> None: @@ -167,16 +173,16 @@ def __init__(self, host: SqlHost, name: str | None) -> None: self.database_name: str | None = name def wait(self, handle: int) -> Awaitable[None]: - """Done once the VM has the read whose handle is ``handle``.""" + """Wait until the VM has the result of the read ``handle``.""" return _until_attempt_end(self._host, self._host.wait(handle)) def _key(self, key: str) -> str: return key if self._name is None else _beside(self._name, key) def get(self, key: str) -> Awaitable[bytes | None]: - # TODO: read with protocol v8's ephemeral get, left out of the journal, - # once a Restate release ships v8 and the SDK can make one. Restate 1.8 - # negotiates v7 at most: a journaled get. + # TODO: use protocol v8's ephemeral get, which stays out of the journal, + # once a Restate release ships v8 and the SDK supports it. Restate 1.8 + # negotiates v7 at most, so this is a journaled get. return _until_attempt_end(self._host, self._host.get(self._key(key))) def set(self, key: str, value: bytes) -> None: @@ -186,7 +192,7 @@ def clear(self, key: str) -> None: self._host.clear(self._key(key)) def keys(self) -> Awaitable[list[str]]: - """The database's keys, for a garbage collection: listed now.""" + """List the database's state keys, for garbage collection; the request is issued immediately.""" listing = _until_attempt_end(self._host, self._host.state_keys()) name = self._name if name is None: @@ -198,12 +204,12 @@ async def own() -> list[str]: return own() def sync(self) -> Awaitable[None]: - """Done once every write before it is stored: the run is made now, in order with them.""" + """Wait until every earlier write is stored; the run is issued now, in order with the writes.""" return _until_attempt_end(self._host, self._host.sync()) class _ReaderStore: - """The object's state as a view's store: it reads beside the exclusive handler, and writes nothing.""" + """The object's state as a view's ``Store``, which reads alongside the exclusive handler and never writes.""" def __init__(self, host: SqlHost, name: str | None) -> None: self._host: SqlHost = host @@ -212,7 +218,7 @@ def __init__(self, host: SqlHost, name: str | None) -> None: self.database_name: str | None = name def wait(self, handle: int) -> Awaitable[None]: - """Done once the VM has the read whose handle is ``handle``.""" + """Wait until the VM has the result of the read ``handle``.""" return _until_attempt_end(self._host, self._host.wait(handle)) def get(self, key: str) -> Awaitable[bytes | None]: @@ -230,9 +236,9 @@ async def sync(self) -> None: async def _until_attempt_end(host: SqlHost, request: Awaitable[T]) -> T: - """ - ``request``'s result, or why the attempt ended, once it has: an attempt that - is suspended or aborted never completes what it was waiting for. + """Return the result of ``request``, or raise why the attempt ended if it ends first. + + A suspended or aborted attempt never receives what it was waiting for. """ ended = host.attempt_end() if ended is not None: @@ -251,33 +257,34 @@ def ended_now() -> None: cancel() if task.done(): return task.result() - # Its error, if it ever has one, is dropped. + # Drop the request's error, if it ever has one. task.add_done_callback(lambda t: t.cancelled() or t.exception()) raise host.attempt_end() or SuspendedException() # --- The outbox ----------------------------------------------------------------------------------- -# A transaction's sends commit with it, in the database's record `outbox` -# (beside its tables, in none of them), which holds each committed -# transaction's sends until they are stored. They are made with its commit -# (see `Attempt.transaction`); a later open makes again those still in the -# outbox, with the same idempotency keys (`transaction//`, with a -# random id of the transaction), so making them again is dropped. Restate -# keeps an attempt's commands in the order they were made, so a send made -# before a commit is stored whenever the commit is: the next commit leaves the -# sends this attempt made out of the outbox, and so does the end of the -# attempt. The format is sdk-sessions' (TypeScript). +# A transaction's sends commit with it, in the database's `outbox` record +# (stored next to its tables, not in one), which holds each committed +# transaction's sends until they are stored by Restate. They are sent during +# the commit (see `Attempt.transaction`). A later open sends again whatever is +# still in the outbox, with the same idempotency keys +# (`transaction//`, where the id is random per transaction), so +# Restate drops the duplicates. Restate keeps an attempt's commands in the +# order they were made, so a send made before a commit is stored whenever the +# commit is: the next commit removes the sends this attempt made from the +# outbox, and so does the end of the attempt. The record's format is a wire +# format shared with other SDKs: keep it compatible. OUTBOX = "outbox" -"""The database's record of the outbox.""" +"""The name of the record that holds the outbox.""" Outgoing = dict[str, Any] -"""The sends of a committed transaction, in the outbox: ``{"id": ..., "sends": [...]}``.""" +"""The sends of one committed transaction, as stored in the outbox: ``{"id": ..., "sends": [...]}``.""" def _encode_outbox(outbox: list[Outgoing]) -> bytes: - """The outbox as its record: none is empty, as a record never set reads.""" + """Encode the outbox as its record; an empty outbox is stored as empty bytes, like a record never set.""" return json.dumps(outbox).encode() if outbox else b"" @@ -290,7 +297,7 @@ def _now_ms() -> int: class _Sends: - """The sends of one transaction, as its body makes them (``tx.send``), until it ends.""" + """The sends that one transaction's body makes with ``tx.send``, collected until it ends.""" def __init__(self) -> None: self._sends: list[dict[str, Any]] = [] @@ -315,7 +322,7 @@ def record(self, send: GenericSend) -> None: self._sends.append(stored) def close(self) -> Outgoing | None: - """The body ended: the transaction's sends, if it made any, their delays counted from now.""" + """Close the collection and return the transaction's sends, if any, with delays counted from now.""" self._open = False if not self._sends: return None @@ -327,7 +334,7 @@ def close(self) -> Outgoing | None: def _refused_send(_: GenericSend) -> None: - """The sends of a transaction that has no outbox: a view's, or another database's.""" + """Refuse a send from a transaction without an outbox (a view's, or another database's).""" raise TerminalError("a send of a transaction is an exclusive handler's, in the object's own database") @@ -335,42 +342,42 @@ def _refused_send(_: GenericSend) -> None: class Reads(abc.ABC): - """``DatabaseReader`` (see ``_api.py``) by way of ``query_columns``: an SQL error that escapes is terminal.""" + """A ``DatabaseReader`` (see ``_api.py``) built on ``query_columns``.""" @abc.abstractmethod async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None) -> tuple[list[str], Rows]: - """The rows of one statement as lists, with the column names of the statement: ``(columns, rows)``.""" + """Run one statement and return its column names and its rows as lists: ``(columns, rows)``.""" async def query(self, sql: str, params: Sequence[SqlParam] | None = None) -> list[Row]: - """ - The rows of one statement, each a dict by column name; parameters bind - to a single statement. Columns are named as SQL names them (an - expression needs ``AS``), and two of the same name are an error: alias - them, or use ``values``. + """Run one statement and return its rows, each a dict keyed by column name. + + Columns are named as SQL names them, so an expression needs ``AS``; + two columns of the same name raise ``SqlError`` (alias one, or use + ``values``). """ columns, rows = await self.query_columns(sql, params) return objects_of(columns, rows) async def get(self, sql: str, params: Sequence[SqlParam] | None = None) -> Row | None: - """The first row of ``query``, or ``None`` when there is none.""" + """Return the first row of ``query``, or ``None`` if there is none.""" rows = await self.query(sql, params) return rows[0] if rows else None async def values(self, sql: str, params: Sequence[SqlParam] | None = None) -> Rows: - """The rows of one statement as lists, in column order: for adapters, and for columns that share a name.""" + """Run one statement and return its rows as lists, in column order.""" return (await self.query_columns(sql, params))[1] class Writes(Reads, abc.ABC): - """``DatabaseWriter`` (see ``_api.py``) by way of ``query_columns``.""" + """A ``DatabaseWriter`` (see ``_api.py``) built on ``query_columns``.""" async def execute(self, sql: str, params: Sequence[SqlParam] | None = None) -> None: - """Run SQL (one or more statements), discarding rows; parameters only with a single statement.""" + """Run SQL and discard any rows; ``sql`` may hold several statements, but parameters bind only to one.""" await self.query_columns(sql, params) class AttemptTransaction(Writes, Sends): - """``Transaction`` (see ``_api.py``): the SQL of a connection's transaction, and the sends it commits with.""" + """The ``Transaction`` (see ``_api.py``) of an attempt: SQL in one transaction, and the sends it commits.""" def __init__(self, handle: ConnectionTransaction, sink: Callable[[GenericSend], None]) -> None: self._handle: ConnectionTransaction = handle @@ -381,7 +388,7 @@ async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None class AttemptDatabase(Writes): - """``Database`` (see ``_api.py``): an attempt's session, opened on its first use.""" + """The ``Database`` (see ``_api.py``) of an attempt, opened on first use.""" def __init__(self, attempt: Attempt) -> None: self._attempt: Attempt = attempt @@ -395,7 +402,7 @@ def transaction(self) -> typing.AsyncContextManager[Transaction]: class _ReaderOf(Reads): - """The reads of a database: a view's ``ctx.db``.""" + """A read-only view of an ``AttemptDatabase``: a view's ``ctx.db``.""" def __init__(self, db: AttemptDatabase) -> None: self._db: AttemptDatabase = db @@ -405,7 +412,7 @@ async def query_columns(self, sql: str, params: Sequence[SqlParam] | None = None def reader_of(db: AttemptDatabase) -> DatabaseReader: - """The reads of ``db``, which apply a library's schema as ``db`` does (``Schema.ensure``).""" + """Return a read-only wrapper of ``db`` that still applies schemas (see ``Schema.ensure``).""" reader = _ReaderOf(db) ensure = ensure_of(db) return ensuring(reader, ensure) if ensure is not None else reader @@ -414,19 +421,32 @@ def reader_of(db: AttemptDatabase) -> DatabaseReader: _DATABASE_NAME = re.compile(r"^[\w.-]{1,64}$") _ATTEMPT: Link[Attempt] = Link("database", weak=True) -"""The attempt of a ``ctx.db``, which keeps it alive (and refers to it): for ``database``.""" +"""The attempt that owns a ``ctx.db``, held weakly as it refers back to it; used by ``database``.""" _RESET: Link[Callable[[], Awaitable[None]]] = Link("reset") -"""How a transaction of an exclusive handler's ``ctx.db`` empties the database: for ``reset``.""" +"""How a transaction of an exclusive handler's ``ctx.db`` empties the database; used by ``reset``.""" def database(db: Database, name: str) -> Database: - """ - The database ``name`` of the key, apart from the object's own (``ctx.db``, - of the attempt whose database ``db`` is): for a library whose tables must - be alone in their database. It opens on its first use, and ends with the - attempt, as ``ctx.db`` does. A name is 1 to 64 letters, digits, ``_``, - ``-`` or ``.``. + """Return another database of the same key, separate from the object's own ``ctx.db``. + + This is for libraries whose tables must live in a database of their own. + Like ``ctx.db``, it opens on first use and closes when the attempt ends. + It has neither the object's migrations nor an outbox, so ``tx.send`` is + not available in its transactions:: + + checkpoints = internal.database(ctx.db, "checkpoints") + + Args: + db: A handler's ``ctx.db``. + name: The database's name: 1 to 64 letters, digits, ``_``, ``-`` or ``.``. + + Returns: + The database ``name`` of the key. + + Raises: + TerminalError: If ``name`` is not a valid name, or ``db`` is not a + handler's ``ctx.db``. """ if not _DATABASE_NAME.match(name): raise TerminalError(f"database: {name!r} is not a database name") @@ -437,12 +457,22 @@ def database(db: Database, name: str) -> Database: async def reset(tx: Transaction) -> None: - """ - Empty the key's database in ``tx``, a transaction of an exclusive handler's - ``ctx.db``: every table, view, index and trigger is dropped, the schemas are - applied again (the object's, and the libraries' the attempt applied), and - the outbox stays. It commits with ``tx``: what ``tx`` writes after it are - the new database's first rows. + """Empty the key's database inside the transaction ``tx``. + + Every table, view, index and trigger is dropped, and then the schemas are + applied again: the object's migrations, and those of any library this + attempt used. The outbox is kept. The reset commits with ``tx``, so + whatever ``tx`` writes afterwards becomes the new database's first rows:: + + async with ctx.db.transaction() as tx: + await internal.reset(tx) + + Args: + tx: A transaction of an exclusive handler's ``ctx.db``. + + Raises: + TerminalError: If ``tx`` is not a transaction of an exclusive + handler's ``ctx.db``. """ empty = _RESET.of(tx) if empty is None: @@ -459,15 +489,15 @@ def _close_soon(db: Connection) -> None: class _Session: - """ - One session of an attempt's database, from its open on. A failure ends it: - its open's, or a request's since. It closes as the attempt ends, unless the - attempt ends it first (``end``). + """One session of an attempt's database, from the moment it starts opening. + + A failure, while opening or afterwards, ends it. It closes when the + attempt ends, unless the attempt ends it first with ``end``. """ def __init__(self, ready: asyncio.Task[Connection], host: SqlHost) -> None: self.ready: asyncio.Task[Connection] = ready - # A call failed on it once a failure ended it: the next call opens another. + # Set when a call fails on a session that a failure ended; the next call opens a new one. self.retired: bool = False self._release: Callable[[], None] = lambda: None ready.add_done_callback(lambda task: self._track(task, host)) @@ -477,13 +507,13 @@ def _track(self, task: asyncio.Task[Connection], host: SqlHost) -> None: return db = task.result() # A handler still waiting for something else when the attempt ends (a - # model's request, say) keeps the session no longer: it closes once no - # statement runs, and the handler's later calls fail. + # model request, say) no longer holds the session: it closes once no + # statement is running, and the handler's later calls fail. self._release = host.on_attempt_end(lambda: _close_soon(db)) @property def ended(self) -> bool: - """Whether a failure ended it.""" + """Whether a failure ended the session.""" if not self.ready.done(): return False if self.ready.cancelled() or self.ready.exception() is not None: @@ -491,11 +521,11 @@ def ended(self) -> bool: return not self.ready.result().is_usable() async def end(self) -> Connection | None: - """Its connection, or none if the open failed; the attempt's end no longer closes it.""" + """Detach the session from the attempt's end and return its connection, or ``None`` if opening failed.""" try: db: Connection | None = await asyncio.shield(self.ready) except asyncio.CancelledError: - # Ours, as the task awaiting is cancelled, unless the open itself was. + # The awaiting task was cancelled, unless the open itself was. if not self.ready.cancelled(): raise db = None @@ -505,7 +535,7 @@ async def end(self) -> Connection | None: return db def close(self) -> None: - """Close it now.""" + """Close the session now.""" async def close() -> None: db = await self.end() @@ -517,26 +547,26 @@ async def close() -> None: class Attempt: - """ - An attempt's SQL session: opened on its first use, or as an exclusive - attempt starts (``open_now``), shared by everything the handler does, and - saved when the handler ends (``end``). With ``name``, of the key's database - ``name`` (``database``), which has neither the object's schema nor an - outbox. + """The SQL session of one attempt. + + It opens on first use, or when an exclusive attempt starts (``open_now``), + is shared by everything the handler does, and is saved when the handler + ends (``end``). With ``name``, it is the key's separate database ``name`` + (see ``database``), which has neither the object's schema nor an outbox. """ def __init__(self, host: SqlHost, config: ObjectConfig | None = None, name: str | None = None) -> None: self._host: SqlHost = host self._config: ObjectConfig = config or ObjectConfig() self._name: str | None = name - # The object's own database of an exclusive handler: the one with an outbox. + # Only the object's own database, in an exclusive handler, has an outbox. self._has_outbox: bool = not host.shared and name is None # The key's other databases this attempt opened, by name. self._databases: dict[str, Attempt] = {} self._session: _Session | None = None - # The ids of the outbox's transactions this attempt made the sends of. + # The ids of the outbox transactions whose sends this attempt has made. self._made: set[str] = set() - # The schemas this attempt's database has all the steps of, by name. + # The schemas fully applied to this attempt's database, by name. self._ensured: dict[str, Schema] = {} self._ended: bool = False self.db: AttemptDatabase = AttemptDatabase(self) @@ -547,7 +577,7 @@ async def _ensure(self, schemas: Sequence[Schema]) -> None: await self.call(lambda db: self._ensure_schemas(db, schemas)) def database(self, name: str) -> Database: - """The key's database ``name``, in this attempt: see ``database``.""" + """Return the key's separate database ``name`` for this attempt (see ``database``).""" if self._ended: raise RuntimeError("the handler has ended: its SQL session is closed") other = self._databases.get(name) @@ -556,19 +586,20 @@ def database(self, name: str) -> Database: return other.db def open_now(self) -> None: - """ - Open the session now, as the attempt starts, rather than on its first - call: the open makes the sends a crashed attempt committed and left - unmade, which must not wait for a handler that uses the database. + """Open the session now, as the attempt starts, rather than on first use. + + Opening makes any sends that a crashed attempt committed but did not + make, and those should not wait until the handler uses the database. """ if self._has_outbox and self._session is None: - # Its failure goes to the first call, as it would had that call opened it. + # A failure is reported to the first call, as if that call had opened the session. self._current().ready.add_done_callback(lambda t: t.cancelled() or t.exception()) async def end(self, error: BaseException | None) -> None: - """ - The handler ended, with ``error`` or not: the key's other databases end, - then this one. It raises what the handler failed with, or a failure to end the session. + """End the attempt's databases after the handler ends, the key's other databases first. + + Re-raises ``error``, the handler's exception, if there is one; + otherwise raises any failure to end a session. """ failures: list[BaseException] = [] for other in self._databases.values(): @@ -581,10 +612,11 @@ async def end(self, error: BaseException | None) -> None: raise failures[0] async def call(self, call: Callable[[Connection], Awaitable[T]]) -> T: - """ - ``call`` on the session, opening it on the first; its failure as the - handler fails with it. The call that fails on a session that a failure - ended is its last: the next opens another, which reads what is stored. + """Run ``call`` on the session, opening it if needed, and return its result. + + Errors propagate to the handler. If a call fails on a session that a + failure ended, the next call opens a new session, which reads what was + stored. """ session = self._current() try: @@ -596,16 +628,18 @@ async def call(self, call: Callable[[Connection], Awaitable[T]]) -> T: @asynccontextmanager async def transaction(self) -> AsyncGenerator[AttemptTransaction, None]: - """ - One transaction. In the object's own database of an exclusive handler, - its sends commit with it, in the outbox, and are made with its commit, - right after its writes and before it waits for anything - (``with_commit``): Restate stores an attempt's commands in order, so - they are stored only if the commit is, and whatever fails the commit's - wait (a cancellation, say) comes after them. A second commit empties - the outbox before the transaction returns: left until the handler's - end, they would be made again by an attempt that resumes a long wait, - past the idempotency retention that drops a repeat. + """Run the body of an ``async with`` in one transaction. + + In the object's own database of an exclusive handler, the transaction's + sends are committed to the outbox with it, and are made during the + commit, right after its writes and before it waits for anything (see + ``with_commit``). Since Restate stores an attempt's commands in order, + the sends are stored only if the commit is, and anything that fails the + commit's wait (a cancellation, say) comes after them. A second commit + empties the outbox before the transaction returns. If the sends stayed + until the handler's end, an attempt that resumes after a long wait could + make them again after the idempotency retention has expired, so that + the duplicates would no longer be dropped. """ outgoing: Outgoing | None = None session = self._current() @@ -613,7 +647,7 @@ async def transaction(self) -> AsyncGenerator[AttemptTransaction, None]: db = await session.ready async with db.transaction() as handle: sends = _Sends() if self._has_outbox else None - # A view's writes stay in its session: it has nothing to reset. + # A view's writes stay in its own session, so there is nothing to reset. tx = AttemptTransaction(handle, sends.record if sends is not None else _refused_send) if not self._host.shared: _RESET.on(tx, lambda: self._reset(handle)) @@ -631,7 +665,7 @@ async def transaction(self) -> AsyncGenerator[AttemptTransaction, None]: await self.call(self._flush) async def _commit_outbox(self, handle: ConnectionTransaction, outgoing: Outgoing | None) -> None: - """Commit with ``handle`` the outbox as committed, but the sends stored since, and with ``outgoing``, which the commit makes.""" + """Write the outbox in ``handle``: the unstored committed sends, plus ``outgoing``, which the commit makes.""" outbox = await self._unstored(handle) if outgoing is not None: outbox.append(outgoing) @@ -640,7 +674,7 @@ async def _commit_outbox(self, handle: ConnectionTransaction, outgoing: Outgoing handle.set_record(OUTBOX, _encode_outbox(outbox)) def _current(self) -> _Session: - """The session, opening one if there is none, or the one there is was retired.""" + """Return the current session, opening a new one if there is none or it was retired.""" if self._ended: raise RuntimeError("the handler has ended: its SQL session is closed") if self._session is not None and self._session.retired: @@ -653,13 +687,13 @@ def _open(self) -> _Session: return _Session(asyncio.ensure_future(self._start()), self._host) def _own_schema(self) -> Schema | None: - """The object's schema, which the key's other databases have none of.""" + """Return the object's schema, or ``None`` for the key's separate databases, which do not have it.""" return self._config.schema if self._name is None else None async def _start(self) -> Connection: host = self._host - # The page cache's: the key escaped, so that no key reads as another's - # database (a name has no `/`). + # The page cache's namespace. The key is escaped so that no key can be + # mistaken for another key's separate database (names have no `/`). key = f"{host.service}/{quote(host.key, safe='')}" namespace = key if self._name is None else f"{key}/{self._name}" store: Store = _ReaderStore(host, self._name) if host.shared else _ObjectStore(host, self._name) @@ -667,20 +701,20 @@ async def _start(self) -> Connection: store = self._config.wrap_store(store) db = await open_connection(namespace, store, reader=host.shared) try: - # A view writes nothing: a write fails, rather than vanish with the attempt. + # Views write nothing: make a write fail, rather than vanish with the attempt. if host.shared: await db.execute("PRAGMA query_only = 1") if self._has_outbox: - # What a crash left unmade; an open again in this attempt finds - # what this attempt made, stored before it, as made. + # Make the sends a crash left unmade. When this attempt reopens the + # database, the sends it already made were stored first, so skip them. left = _decode_outbox(await db.record(OUTBOX)) for outgoing in left: if outgoing["id"] not in self._made: self._make(outgoing, again=True) - # Emptied at once, as a transaction's own sends are. + # Empty the outbox right away, as for a transaction's own sends. if left: await self._flush(db) - # A view's too, whose writes are never stored: a key never written reads as empty tables. + # Views too, although their writes are never stored: a key never written reads as empty tables. own = self._own_schema() if own is not None: await self._ensure_schemas(db, [own]) @@ -690,11 +724,12 @@ async def _start(self) -> Connection: return db async def _end_session(self, error: BaseException | None) -> None: - """ - End the session, then raise as the handler did. The handler's error stands: what it - committed is stored, and an outbox left full is made by the next open. - A failure to end the session fails a handler that succeeded, which a - retry then ends (a transaction it left open, say). + """End the session, then re-raise the handler's error, if any. + + The handler's error takes precedence: what it committed is stored, and + the next open makes any sends left in the outbox. A failure to end the + session (a transaction left open, say) fails a handler that succeeded, + and a retry then ends it. """ self._ended = True failed: BaseException | None = None @@ -708,12 +743,14 @@ async def _end_session(self, error: BaseException | None) -> None: raise failed async def _close(self) -> None: - """ - Close the session as the attempt ends. A view's closes: it read the - snapshot it opened on. Another is saved, its outbox emptied. One that a - failure ended closes, and the object's own database is opened again and - saved: it may hold sends unmade (a crash left them, and this attempt's - open failed) or made and not emptied (the commit's wait failed). + """Close the session as the attempt ends. + + A view's session just closes, as it only read its snapshot. Any other + session is saved, after emptying its outbox. If a failure ended it, it + is closed, and the object's own database is reopened and saved, since + its outbox may hold sends not yet made (left by a crash, when this + attempt's open failed) or made but not removed (when the commit's wait + failed). """ if self._session is None: return @@ -737,7 +774,7 @@ async def _close(self) -> None: await again.end() async def _save(self, db: Connection) -> None: - """Empty the outbox and save the session, which closes if that fails.""" + """Empty the outbox and save the session, closing it if that fails.""" try: if self._has_outbox: await self._flush(db) @@ -748,13 +785,13 @@ async def _save(self, db: Connection) -> None: raise async def _reset(self, handle: ConnectionTransaction) -> None: - """Empty the database in the transaction ``handle`` (``reset``), and apply the schemas this attempt applied again.""" + """Empty the database in the transaction ``handle`` and reapply the schemas this attempt applied (see ``reset``).""" own = self._own_schema() schemas = ([own] if own is not None else []) + [s for s in self._ensured.values() if s is not own] await drop_all(handle, schemas) async def _ensure_schemas(self, db: Connection, schemas: Sequence[Schema]) -> None: - """Apply the steps of ``schemas`` that the database lacks, each schema once per attempt.""" + """Apply the missing migrations of ``schemas``, checking each schema at most once per attempt.""" lacking = [s for s in schemas if s.name not in self._ensured] if not lacking: return @@ -763,12 +800,12 @@ async def _ensure_schemas(self, db: Connection, schemas: Sequence[Schema]) -> No self._ensured[schema.name] = schema def _make(self, outgoing: Outgoing, again: bool = False) -> None: - """Make the sends of a committed transaction: each with an idempotency key of the transaction (or its own), so that making them again is dropped.""" + """Make the sends of a committed transaction, with idempotency keys so that repeats are dropped.""" for index, send in enumerate(outgoing["sends"]): delay = send.get("delay") due = send.get("due") if delay is not None and again and due is not None: - # Made again, a delayed send waits only what remains of its delay. + # When made again, a delayed send waits only for the rest of its delay. delay = max(0, due - _now_ms()) self._host.generic_send( send["service"], @@ -782,24 +819,24 @@ def _make(self, outgoing: Outgoing, again: bool = False) -> None: self._made.add(outgoing["id"]) async def _unstored(self, db: Connection | ConnectionTransaction) -> list[Outgoing]: - """The outbox as committed, but the transactions this attempt made the sends of: a commit made after them stores them first.""" + """Return the committed outbox without the transactions whose sends this attempt made (a later commit stores those).""" return [outgoing for outgoing in _decode_outbox(await db.record(OUTBOX)) if outgoing["id"] not in self._made] async def _flush(self, db: Connection) -> None: - """Empty the outbox, so that a later open does not make its sends again: a commit after them, which stores them first.""" + """Remove the sends this attempt made from the outbox, so that a later open does not make them again.""" if not await db.record(OUTBOX): return - # Read in the transaction that writes it: one that committed sends since keeps them. + # Read the outbox in the transaction that writes it, to keep sends committed in the meantime. async with db.transaction() as tx: tx.set_record(OUTBOX, _encode_outbox(await self._unstored(tx))) async def with_attempt(handler: Callable[[], Awaitable[T]], attempt: Attempt) -> T: - """Run ``handler``, then end ``attempt``'s SQL session, if it opened one (see ``Attempt.end``).""" + """Run ``handler``, then end the SQL session of ``attempt``, if it opened one (see ``Attempt.end``).""" try: value = await handler() except BaseException as error: await attempt.end(error) - raise # `end` raises: this is for the type checker + raise # `end` already raises; this is for the type checker await attempt.end(None) return value diff --git a/python/restate/sessions/internal.py b/python/restate/sessions/internal.py index 0209366c..4b4ad689 100644 --- a/python/restate/sessions/internal.py +++ b/python/restate/sessions/internal.py @@ -8,22 +8,23 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -""" -What libraries built on session objects (``restate.agents``) and tests need, -beside the API handlers code against: as sdk-sessions' -``@restatedev/sessions/internal``. +"""Building blocks for libraries built on session objects, and for tests. + +Application code should not need this module; ``restate.agents`` is built on +it. It provides: -- ``Schema``: tables a library owns, applied where it uses them - (``schema.ensure(ctx.db)``); ``database``, another database of the same - key; ``reset``, the key's database emptied in a transaction. -- ``cancelled`` and ``cancel_calls``: explicit cancellation, for a handler - with ``explicit_cancellation=True``. -- ``ScopeKey`` and ``definition_of``: what a library keeps for an attempt, or - for an object. -- ``needing`` and ``served``: the services a definition needs served beside - it, which ``restate.sessions.app`` serves. -- ``invoke``: a handler's attempt over a host of a test's own, without - Restate (see ``SqlHost``). +- ``Schema``: tables that a library owns and creates where it uses them + (``await schema.ensure(ctx.db)``). +- ``database``: another named database under the same key. +- ``reset``: empty the key's database inside a transaction. +- ``cancelled`` and ``cancel_calls``: explicit cancellation, for handlers + declared with ``explicit_cancellation=True``. +- ``ScopeKey`` and ``definition_of``: state a library keeps per attempt, and + the object an attempt belongs to. +- ``needing`` and ``served``: declare the services an object needs served + next to it, which ``restate.sessions.app`` then serves too. +- ``invoke``, ``SqlHost`` and ``ObjectConfig``: run one attempt of a handler + against a host of your own, without Restate, for tests. """ from restate.sessions._object import cancel_calls, cancelled, definition_of, invoke From c1766c04160410148b5a9a28fb0f717d2c4ad7cf Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 18:26:18 +0100 Subject: [PATCH 39/45] restate.agents: refuse a built-in name at once; send's text names the key it used An agent's handler or view named chat, messages, clear or sync was refused only once chat() registered the built-ins; every built-in name is now refused when it is registered. The send tool's confirmation named a key even for a service, which takes none. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/agents/_object.py | 30 +++++++++++++++++++----------- python/restate/agents/_tools.py | 8 ++------ python/restate/sessions/_object.py | 8 ++++++++ tests/agents_definitions.py | 10 ++++++++++ 4 files changed, 39 insertions(+), 17 deletions(-) diff --git a/python/restate/agents/_object.py b/python/restate/agents/_object.py index ab4c4a28..a5fda5dd 100644 --- a/python/restate/agents/_object.py +++ b/python/restate/agents/_object.py @@ -182,15 +182,19 @@ def __init__( self._clear: Callable[[SessionContext], Awaitable[None]] | None = None AGENT.on(self, AgentInfo(background=background)) _signals.with_companion(self) - self._built_in() - - def _taken(self, name: str) -> str: - if name in BUILT_IN: - return ( - f"{name!r} is a built-in handler of every agent ({', '.join(sorted(BUILT_IN))}): " - f"give yours another name, such as @{self.name}.handler(name='my_{name}')" - ) - return super()._taken(name) + self._building = True + try: + self._built_in() + finally: + self._building = False + + def _reserved(self, name: str) -> str | None: + if name not in BUILT_IN or getattr(self, "_building", False): + return None + return ( + f"{name!r} is a built-in handler of every agent ({', '.join(sorted(BUILT_IN))}): " + f"give yours another name, such as @{self.name}.handler(name='my_{name}')" + ) def _attempt_fn(self, fn: Callable[..., Awaitable[Any]], view: bool) -> Callable[..., Awaitable[Any]]: """Wrap a handler so that a model provider's refusal fails its invocation for good, as in ``chat``.""" @@ -308,8 +312,12 @@ def decorate( async def chat(ctx: SessionContext, text: str) -> Any: return await chat_turn(ctx, text, answering) - self.handler(name="chat", accept="*/*", input_serde=MessageSerde(), explicit_cancellation=True)(chat) - _followed(self) + self._building = True + try: + self.handler(name="chat", accept="*/*", input_serde=MessageSerde(), explicit_cancellation=True)(chat) + _followed(self) + finally: + self._building = False return answer diff --git a/python/restate/agents/_tools.py b/python/restate/agents/_tools.py index 998189d3..0c409e72 100644 --- a/python/restate/agents/_tools.py +++ b/python/restate/agents/_tools.py @@ -449,13 +449,9 @@ async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: chosen, value = arguments["key"], arguments["value"] async def made(c: ToolCall) -> str: - target_key = chosen if key is None else await configured(key, c.ctx) + target_key = (chosen if key is None else await configured(key, c.ctx)) if _keyed(handler) else None _untyped(c.ctx).send( - target, - *_inputs(handler, value), - key=target_key if _keyed(handler) else None, - delay=delay, - idempotency_key=c.idempotency_key, + target, *_inputs(handler, value), key=target_key, delay=delay, idempotency_key=c.idempotency_key ) return f"sent to {service}{'' if target_key is None else f'/{target_key}'}/{handler.name}" diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index 0b6bac23..45231397 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -445,6 +445,11 @@ def __init__( def __repr__(self) -> str: return f"{type(self).__name__}({self.name!r})" + def _reserved(self, name: str) -> str | None: + """Return why ``name`` is reserved for a built-in handler, if it is (none are, for a plain session object).""" + del name + return None + def _taken(self, name: str) -> str: """Return the error message for registering ``name`` a second time.""" return f"{self.name} has a handler or view named {name!r} already: give this one another name" @@ -644,6 +649,9 @@ def register(fn: F) -> F: kind = "view" if view else "handler" raise TypeError(f"{self.name}.{fn.__name__}: a {kind} must be a coroutine function (async def)") chosen = name or fn.__name__ + reserved = self._reserved(chosen) + if reserved is not None: + raise ValueError(reserved) if chosen in self.handlers: raise ValueError(self._taken(chosen)) explicit = bool(options.get("explicit_cancellation")) diff --git a/tests/agents_definitions.py b/tests/agents_definitions.py index 67b0c907..4187f1d5 100644 --- a/tests/agents_definitions.py +++ b/tests/agents_definitions.py @@ -57,6 +57,16 @@ async def status(ctx: sessions.SessionContext) -> str: agent.handler(status) +def test_a_name_chat_will_take_is_refused_before_chat_registers_it(): + agent = AgentObject("early") + + async def messages(ctx: sessions.SessionContext) -> str: + return "mine" + + with pytest.raises(ValueError, match="'messages' is a built-in handler of every agent"): + agent.handler(messages) + + def test_chat_is_a_decorator_too_and_returns_the_function(): agent = AgentObject("decorated") From 98b1383d62db3563a948626d53d25953c6b84a20 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 18:38:48 +0100 Subject: [PATCH 40/45] Take the SIGTERM chaining out of this branch, for a PR of its own server.py is main's again (but for passing explicit_cancellation to the VM). The SDK-wide fix, where install_shutdown_hook(chain=True) keeps the ASGI server's handler, is f14507d's and goes to main separately. sessions.serve no longer needs it: it registers uvicorn's exit through install_shutdown_hook itself, before serving, and uvicorn captures no signals. The endpoint's own SIGTERM hook then joins the same dispatcher on its first request, instead of replacing the server's handler. It stops on SIGINT and SIGTERM after serving requests (tests/examples_agents.py). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/server.py | 40 +++++--------------------- python/restate/sessions/_serve.py | 28 +++++++++++++++++- tests/server_sigterm.py | 48 ------------------------------- 3 files changed, 34 insertions(+), 82 deletions(-) delete mode 100644 tests/server_sigterm.py diff --git a/python/restate/server.py b/python/restate/server.py index f4208bdf..77cc1830 100644 --- a/python/restate/server.py +++ b/python/restate/server.py @@ -15,7 +15,7 @@ import logging import signal import weakref -from typing import Any, Callable, Dict, List, Optional, Set, TypedDict, Literal +from typing import Callable, Dict, List, Optional, Set, TypedDict, Literal from restate.discovery import compute_discovery_json from restate.endpoint import Endpoint @@ -39,18 +39,13 @@ def install_shutdown_hook( - loop: asyncio.AbstractEventLoop, hook: Callable[[], None], signum: int = signal.SIGTERM, *, chain: bool = False + loop: asyncio.AbstractEventLoop, hook: Callable[[], None], signum: int = signal.SIGTERM ) -> None: """ Run ``hook`` when ``signum`` (default SIGTERM) is received while ``loop`` runs. - With ``chain``, the handler the process already has for ``signum`` still runs too: the ASGI - server's (hypercorn's, uvicorn's), which stops it. ``loop.add_signal_handler`` keeps a single - handler per signal, and would replace it. Without, ``hook``'s dispatcher replaces it, for a - program that owns its signals (``restate.connect``). The first hook of a loop and signal decides. - - Raises what ``loop.add_signal_handler`` (or ``signal.signal``, with ``chain``) raises when - signals cannot be handled here (Windows, or a non-main thread). + Raises what ``loop.add_signal_handler`` raises when signals cannot be handled here + (Windows, or a non-main thread). """ hooks = _shutdown_hooks.setdefault(loop, {}) if signum not in hooks: @@ -63,33 +58,11 @@ def dispatch() -> None: except Exception: # pylint: disable=broad-except logger.exception("Shutdown hook failed") - if chain: - _chain_signal(loop, signum, dispatch) - else: - loop.add_signal_handler(signum, dispatch) + loop.add_signal_handler(signum, dispatch) hooks[signum] = registered hooks[signum].append(hook) -def _chain_signal(loop: asyncio.AbstractEventLoop, signum: int, on_signal: Callable[[], None]) -> None: - """ - Call ``on_signal`` on ``loop`` when ``signum`` arrives, then the signal's previous handler. - - A signal-module handler leaves the loop's own handlers working too (they are woken by its wakeup - fd). Python's default handlers are not chained: the default SIGINT handler would raise - KeyboardInterrupt into whatever runs. - """ - previous = signal.getsignal(signum) - - def handler(received: int, frame: Any) -> None: - if not loop.is_closed(): - loop.call_soon_threadsafe(on_signal) - if callable(previous) and previous is not signal.default_int_handler: - previous(received, frame) - - signal.signal(signum, handler) - - X_RESTATE_SERVER = header_to_binary([("x-restate-server", f"restate-sdk-python/{SDK_VERSION}")]) @@ -322,8 +295,9 @@ def _on_sigterm() -> None: async def app(scope: Scope, receive: Receive, send: Send): nonlocal sigterm_installed if not sigterm_installed: + loop = asyncio.get_running_loop() try: - install_shutdown_hook(asyncio.get_running_loop(), _on_sigterm, signal.SIGTERM, chain=True) + install_shutdown_hook(loop, _on_sigterm, signal.SIGTERM) except (NotImplementedError, RuntimeError, ValueError): pass # Windows or non-main thread sigterm_installed = True diff --git a/python/restate/sessions/_serve.py b/python/restate/sessions/_serve.py index baa2e6ac..6d9178e3 100644 --- a/python/restate/sessions/_serve.py +++ b/python/restate/sessions/_serve.py @@ -18,14 +18,19 @@ from __future__ import annotations +import asyncio +import contextlib import importlib import os import re +import signal import typing from collections.abc import Callable, Iterable, Sequence +from typing import Any import restate from restate.entry_codec import JournalValueCodec, JournalValueCodecProvider +from restate.server import install_shutdown_hook from restate.server_types import RestateAppT from restate.sessions._link import Link @@ -181,7 +186,28 @@ def serve( lifespan="off", timeout_graceful_shutdown=_GRACEFUL_SHUTDOWN, ) - uvicorn.Server(config).run() + asyncio.run(_until_stopped(uvicorn.Server(config))) + + +async def _until_stopped(server: Any) -> None: + """Run the uvicorn ``server`` until SIGINT or SIGTERM. + + The SDK's endpoint registers its own SIGTERM hook (to wind down the + invocations in flight) through ``install_shutdown_hook``, which keeps one + asyncio handler per signal. Uvicorn's handlers would be replaced by it on + the first request, so the server's exit is registered through the same + hook instead, and uvicorn captures no signals itself: both then run. + """ + loop = asyncio.get_running_loop() + + def exit_on(signum: int) -> Callable[[], None]: + return lambda: server.handle_exit(signum, None) + + for signum in (signal.SIGINT, signal.SIGTERM): + with contextlib.suppress(NotImplementedError, RuntimeError, ValueError): # Windows, or not the main thread + install_shutdown_hook(loop, exit_on(signum), signum) + server.capture_signals = contextlib.nullcontext + await server.serve() _NO_SERVER = ( diff --git a/tests/server_sigterm.py b/tests/server_sigterm.py deleted file mode 100644 index 55c03062..00000000 --- a/tests/server_sigterm.py +++ /dev/null @@ -1,48 +0,0 @@ -# -# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH -# -# This file is part of the Restate SDK for Python, -# which is released under the MIT license. -# -# You can find a copy of the license in file LICENSE in the root -# directory of this repository or package, or at -# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE -# -"""The endpoint's SIGTERM handler leaves the ASGI server's own in place.""" - -import asyncio -import os -import signal -import sys - -import pytest - -from restate.server import install_shutdown_hook - -pytestmark = pytest.mark.skipif(sys.platform == "win32", reason="POSIX signals") - - -def test_sigterm_reaches_the_server_and_the_endpoint(): - async def main() -> list[str]: - loop = asyncio.get_running_loop() - heard: list[str] = [] - both = asyncio.Event() - - def hear(who: str) -> None: - heard.append(who) - if len(heard) == 2: - both.set() - - # The server's handler, as hypercorn installs it, and then the endpoint's, on its first request. - loop.add_signal_handler(signal.SIGTERM, hear, "server") - install_shutdown_hook(loop, lambda: hear("endpoint"), signal.SIGTERM, chain=True) - os.kill(os.getpid(), signal.SIGTERM) - await asyncio.wait_for(both.wait(), 5) - loop.remove_signal_handler(signal.SIGTERM) - return heard - - previous = signal.getsignal(signal.SIGTERM) - try: - assert sorted(asyncio.run(main())) == ["endpoint", "server"] - finally: - signal.signal(signal.SIGTERM, previous) From b5d1ff9c0c815eab9b899843db3d4319267f3e4b Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 19:25:28 +0100 Subject: [PATCH 41/45] restate.sessions: pass explicit_recovery and explicit_cancellation as keys Now that both are options of the SDK's handler decorators, a session object's handler passes them through instead of setting them on the registered handler afterwards. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkTUvGceAzHBUPULRvy --- python/restate/sessions/_object.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index 45231397..8ce7e9ae 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -654,23 +654,20 @@ def register(fn: F) -> F: raise ValueError(reserved) if chosen in self.handlers: raise ValueError(self._taken(chosen)) - explicit = bool(options.get("explicit_cancellation")) - sdk_options = {key: value for key, value in options.items() if key != "explicit_cancellation"} entry = _Entry( fn=self._attempt_fn(fn, view), view=view, config=self._config, definition=weakref.ref(self), - explicit_cancellation=explicit, - ) - registered = VirtualObject.handler(self, name, kind="shared" if view else "exclusive", **sdk_options)( - _sdk_handler(entry) + explicit_cancellation=bool(options.get("explicit_cancellation")), ) + # Attempts are not replayed (explicit_recovery): each runs from the + # start, against the current state, so it may read differently from + # earlier attempts. + registered = VirtualObject.handler( + self, name, kind="shared" if view else "exclusive", explicit_recovery=True, **options + )(_sdk_handler(entry)) sdk: Handler[Any, Any] = handler_from_callable(registered) - # Attempts are not replayed: each runs from the start, against the - # current state, so it may read differently from earlier attempts. - sdk.explicit_recovery = True - sdk.explicit_cancellation = explicit entries = ENTRIES.of(self) assert entries is not None entries[sdk.name] = entry From 3e18050fa8d7eb91462d43f7c2d576bea19a1dc4 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 19:50:37 +0100 Subject: [PATCH 42/45] CI: check the SDK as an ordinary user installs it, without extras A new job builds the wheel, installs it alone into a fresh environment (Python 3.10 and 3.14) and runs scripts/check_no_extras.py from outside the checkout. It fails if the package requires anything beyond its extras, if anything is installed beside it, or if import restate loads restate.sessions or restate.agents. It checks that an ordinary restate.app, and session objects and agents, serve their discovery; and that pydantic tools, the langgraph and pydantic_ai integrations and sessions.serve each fail only when used, naming what to install. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/test.yml | 25 ++++++++ pyproject.toml | 1 + scripts/check_no_extras.py | 123 +++++++++++++++++++++++++++++++++++++ 3 files changed, 149 insertions(+) create mode 100644 scripts/check_no_extras.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index df18a1d7..d70d8b51 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -32,3 +32,28 @@ jobs: run: just build - name: Verify run: just verify + + no-extras: + # An ordinary user's install, `pip install restate-sdk`: no new dependencies, and the optional parts + # (restate.sessions' server, restate.agents' integrations) fail only when used, naming their extra. + name: "Install without extras (Python ${{ matrix.python }})" + runs-on: warp-ubuntu-latest-x64-4x + strategy: + matrix: + python: [ "3.10", "3.14" ] + env: + UV_PYTHON: ${{ matrix.python }} + steps: + - uses: actions/checkout@v4.4.0 + - uses: astral-sh/setup-uv@v5.4.2 + with: + enable-cache: true + - name: Build the wheel + run: uv build --wheel --out-dir dist + - name: Install it alone into a fresh environment + run: | + uv venv /tmp/plain + uv pip install --python /tmp/plain/bin/python dist/restate_sdk-*.whl + - name: Check it as an ordinary user gets it + working-directory: /tmp + run: /tmp/plain/bin/python -I "$GITHUB_WORKSPACE/scripts/check_no_extras.py" diff --git a/pyproject.toml b/pyproject.toml index 7c13b70a..65407d98 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -86,6 +86,7 @@ include = [ "python/**/*.py", "examples/**/*.py", "examples-agents/**/*.py", + "scripts/**/*.py", "tests/**/*.py", "test-services//**/*.py", ] diff --git a/scripts/check_no_extras.py b/scripts/check_no_extras.py new file mode 100644 index 00000000..79b1c521 --- /dev/null +++ b/scripts/check_no_extras.py @@ -0,0 +1,123 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""Check the SDK as an ordinary user installs it: ``pip install restate-sdk``, without extras. + +Run it with the interpreter of a fresh environment where only the built wheel +is installed, from outside the checkout, so that it imports the installed +package rather than the sources:: + + uv venv /tmp/plain && uv pip install --python /tmp/plain/bin/python dist/restate_sdk-*.whl + cd /tmp && /tmp/plain/bin/python -I "$OLDPWD/scripts/check_no_extras.py" + +It checks that: + +- the package requires nothing beyond its extras, and nothing else is installed; +- ``import restate`` loads neither ``restate.sessions`` nor ``restate.agents``; +- an ordinary ``restate.app`` serves its discovery; +- ``restate.sessions`` and ``restate.agents`` import, and their apps serve too; +- each optional part fails only when used, with an ``ImportError`` that names + what to install. +""" + +import asyncio +import importlib.metadata +import json +import sys +from collections.abc import Callable +from typing import Any + +ALLOWED = {"restate-sdk", "pip", "setuptools", "wheel"} +"""What a fresh environment may hold beside the SDK.""" + + +def check(what: str, condition: bool, detail: object = "") -> None: + if not condition: + sys.exit(f"FAILED {what}: {detail}") + print(f"ok {what}") + + +async def discover(app: Any) -> dict[str, Any]: + """Ask ``app`` for its endpoint manifest, as Restate does on registration.""" + sent: list[dict[str, Any]] = [] + scope = { + "type": "http", + "http_version": "2", + "method": "GET", + "path": "/discover", + "headers": [(b"accept", b"application/vnd.restate.endpointmanifest.v5+json")], + "query_string": b"", + } + + async def receive() -> dict[str, Any]: + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(message: dict[str, Any]) -> None: + sent.append(message) + + await app(scope, receive, send) + check("discovery answers 200", sent[0]["status"] == 200, sent[:1]) + return json.loads(b"".join(m.get("body", b"") for m in sent if m["type"] == "http.response.body")) + + +def refused(what: str, action: Callable[[], object], names: str) -> None: + """Check that ``action`` raises an ``ImportError`` that names ``names``.""" + try: + action() + except ImportError as error: + check(f"{what} needs {names}", names in str(error), error) + return + sys.exit(f"FAILED {what} worked without {names}") + + +def main() -> None: + requirements = importlib.metadata.requires("restate-sdk") or [] + unconditional = [req for req in requirements if "extra ==" not in req] + check("no requirement beyond the extras", not unconditional, unconditional) + installed = {dist.metadata["Name"].lower() for dist in importlib.metadata.distributions()} + check("nothing installed beside restate-sdk", installed <= ALLOWED, sorted(installed - ALLOWED)) + + import restate + + loaded = [name for name in ("restate.sessions", "restate.agents") if name in sys.modules] + check("import restate loads neither sessions nor agents", not loaded, loaded) + + greeter = restate.Service("greeter") + + @greeter.handler() + async def greet(ctx: restate.Context, name: str) -> str: + return f"Hello {name}" + + manifest = asyncio.run(discover(restate.app([greeter]))) + check("an ordinary restate.app serves", [s["name"] for s in manifest["services"]] == ["greeter"], manifest) + + import restate.agents + import restate.agents.internal + import restate.sessions.internal + from restate import sessions + + notes = sessions.SessionObject("notes", migrations=["CREATE TABLE notes (text TEXT)"]) + + @notes.handler + async def add(ctx: sessions.SessionContext, text: str) -> None: + await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [text]) + + manifest = asyncio.run(discover(sessions.app(notes, restate.agents.AgentObject("assistant")))) + served = sorted(s["name"] for s in manifest["services"]) + check("sessions and agents serve", served == ["assistant", "notes", "signals"], served) + + refused("restate.agents.tool", lambda: restate.agents.tool(lambda city: city), "pydantic") + refused("restate.agents.langgraph", lambda: __import__("restate.agents.langgraph"), "restate-sdk[langgraph]") + refused("restate.agents.pydantic_ai", lambda: __import__("restate.agents.pydantic_ai"), "restate-sdk[pydantic_ai]") + refused("sessions.serve", lambda: sessions.serve(notes), "restate-sdk[serve]") + + +if __name__ == "__main__": + main() From ebfe66fc5bde42e5d061bc397e8c64c2337bce10 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 20:02:02 +0100 Subject: [PATCH 43/45] restate.serve replaces sessions.serve Session objects and agents are served as any Restate app: asyncio.run(restate.serve(sessions.app(...))). sessions.serve is gone, and the examples, docs and the no-extras check use restate.serve. Co-Authored-By: Claude Opus 5.5 --- examples-agents/README.md | 8 +- examples-agents/basic/langgraph_agent.py | 5 +- examples-agents/basic/notes.py | 5 +- examples-agents/basic/pydantic_ai_agent.py | 5 +- examples-agents/tutorial/01_object.py | 5 +- examples-agents/tutorial/02_database.py | 5 +- examples-agents/tutorial/03_views.py | 5 +- examples-agents/tutorial/04_migrations.py | 5 +- examples-agents/tutorial/05_transaction.py | 4 +- examples-agents/tutorial/06_retries.py | 4 +- examples-agents/tutorial/07_calls.py | 4 +- examples-agents/tutorial/08_sends.py | 4 +- python/restate/agents/__init__.py | 5 +- python/restate/agents/langgraph/__init__.py | 5 +- python/restate/agents/langgraph/_agent.py | 2 +- python/restate/agents/pydantic_ai/__init__.py | 5 +- python/restate/agents/pydantic_ai/_agent.py | 2 +- python/restate/sessions/__init__.py | 8 +- python/restate/sessions/_serve.py | 94 +------------------ scripts/check_no_extras.py | 4 +- tests/examples_agents.py | 4 +- 21 files changed, 71 insertions(+), 117 deletions(-) diff --git a/examples-agents/README.md b/examples-agents/README.md index 6bf2f6d8..92d11494 100644 --- a/examples-agents/README.md +++ b/examples-agents/README.md @@ -7,12 +7,12 @@ top of it. Each example is a single file that runs on its own. ## Run an example -The examples serve themselves with `sessions.serve`, which needs Python 3.11 -or later. (The SDK itself supports Python 3.10; there, serve -`sessions.app(...)` with any ASGI server, such as hypercorn.) +The examples serve themselves with `restate.serve(sessions.app(...))`, +which needs Python 3.11 or later. (The SDK itself supports Python 3.10; there, +serve `sessions.app(...)` with any ASGI server, such as hypercorn.) Install the SDK with the `serve` extra, which adds the HTTP/2 server that -`sessions.serve` runs on: +`restate.serve` runs on: ```sh pip install "restate-sdk[serve]" diff --git a/examples-agents/basic/langgraph_agent.py b/examples-agents/basic/langgraph_agent.py index 1c57dea8..e99eb52a 100644 --- a/examples-agents/basic/langgraph_agent.py +++ b/examples-agents/basic/langgraph_agent.py @@ -14,6 +14,9 @@ ANTHROPIC_API_KEY) instead of ``langchain-openai``. """ +import asyncio + +import restate from restate import sessions from restate.agents import langgraph, tool @@ -28,4 +31,4 @@ async def get_weather(city: str) -> str: if __name__ == "__main__": - sessions.serve(assistant) + asyncio.run(restate.serve(sessions.app(assistant))) diff --git a/examples-agents/basic/notes.py b/examples-agents/basic/notes.py index 06ce4712..214ca5b9 100644 --- a/examples-agents/basic/notes.py +++ b/examples-agents/basic/notes.py @@ -12,6 +12,9 @@ curl localhost:8080/notes/alice/list # ["Buy milk"] """ +import asyncio + +import restate from restate import sessions # The schema of each key's database: its steps, applied once per key (a step added later runs then). @@ -30,4 +33,4 @@ async def list_notes(ctx: sessions.ViewContext) -> list[str]: if __name__ == "__main__": - sessions.serve(notes) + asyncio.run(restate.serve(sessions.app(notes))) diff --git a/examples-agents/basic/pydantic_ai_agent.py b/examples-agents/basic/pydantic_ai_agent.py index 309dfeea..d890238a 100644 --- a/examples-agents/basic/pydantic_ai_agent.py +++ b/examples-agents/basic/pydantic_ai_agent.py @@ -14,6 +14,9 @@ ``pydantic-ai-slim[anthropic]`` (and ANTHROPIC_API_KEY) instead. """ +import asyncio + +import restate from restate import sessions from restate.agents import pydantic_ai, tool @@ -28,4 +31,4 @@ async def get_weather(city: str) -> str: if __name__ == "__main__": - sessions.serve(assistant) + asyncio.run(restate.serve(sessions.app(assistant))) diff --git a/examples-agents/tutorial/01_object.py b/examples-agents/tutorial/01_object.py index fc4b8d2c..c59d0874 100644 --- a/examples-agents/tutorial/01_object.py +++ b/examples-agents/tutorial/01_object.py @@ -18,6 +18,9 @@ and register again. """ +import asyncio + +import restate from restate import sessions # A session object: a virtual object that gets a database in the next lesson. @@ -31,4 +34,4 @@ async def hello(ctx: sessions.SessionContext) -> str: if __name__ == "__main__": - sessions.serve(wallet) + asyncio.run(restate.serve(sessions.app(wallet))) diff --git a/examples-agents/tutorial/02_database.py b/examples-agents/tutorial/02_database.py index f329d20f..aa24fa1b 100644 --- a/examples-agents/tutorial/02_database.py +++ b/examples-agents/tutorial/02_database.py @@ -15,6 +15,9 @@ curl localhost:8080/wallet/bob/deposit --json 7 # 7 """ +import asyncio + +import restate from restate import sessions wallet = sessions.SessionObject("wallet", migrations=["CREATE TABLE entries (amount INTEGER NOT NULL)"]) @@ -29,4 +32,4 @@ async def deposit(ctx: sessions.SessionContext, amount: int) -> int: if __name__ == "__main__": - sessions.serve(wallet) + asyncio.run(restate.serve(sessions.app(wallet))) diff --git a/examples-agents/tutorial/03_views.py b/examples-agents/tutorial/03_views.py index 1ece278c..771cd280 100644 --- a/examples-agents/tutorial/03_views.py +++ b/examples-agents/tutorial/03_views.py @@ -13,6 +13,9 @@ curl localhost:8080/wallet/alice/history # every deposit, oldest first """ +import asyncio + +import restate from restate import sessions wallet = sessions.SessionObject("wallet", migrations=["CREATE TABLE entries (amount INTEGER NOT NULL)"]) @@ -41,4 +44,4 @@ async def history(ctx: sessions.ViewContext) -> list[int]: if __name__ == "__main__": - sessions.serve(wallet) + asyncio.run(restate.serve(sessions.app(wallet))) diff --git a/examples-agents/tutorial/04_migrations.py b/examples-agents/tutorial/04_migrations.py index 5b83185a..5542bb8e 100644 --- a/examples-agents/tutorial/04_migrations.py +++ b/examples-agents/tutorial/04_migrations.py @@ -14,8 +14,11 @@ curl localhost:8080/wallet/alice/history # the earlier deposits have no note """ +import asyncio + from pydantic import BaseModel +import restate from restate import sessions wallet = sessions.SessionObject( @@ -54,4 +57,4 @@ async def history(ctx: sessions.ViewContext) -> list[sessions.Row]: if __name__ == "__main__": - sessions.serve(wallet) + asyncio.run(restate.serve(sessions.app(wallet))) diff --git a/examples-agents/tutorial/05_transaction.py b/examples-agents/tutorial/05_transaction.py index a79720e7..0a81d92a 100644 --- a/examples-agents/tutorial/05_transaction.py +++ b/examples-agents/tutorial/05_transaction.py @@ -20,6 +20,8 @@ curl localhost:8080/wallet/alice/history # lunch, but no boat """ +import asyncio + from pydantic import BaseModel import restate @@ -69,4 +71,4 @@ async def history(ctx: sessions.ViewContext) -> list[sessions.Row]: if __name__ == "__main__": - sessions.serve(wallet) + asyncio.run(restate.serve(sessions.app(wallet))) diff --git a/examples-agents/tutorial/06_retries.py b/examples-agents/tutorial/06_retries.py index 37c25aca..769a4608 100644 --- a/examples-agents/tutorial/06_retries.py +++ b/examples-agents/tutorial/06_retries.py @@ -20,6 +20,8 @@ curl localhost:8080/wallet/alice/deposit --json '{"id": "d1", "amount": 10, "note": "refund"}' # sent again: the balance stays the same """ +import asyncio + from pydantic import BaseModel import restate @@ -97,4 +99,4 @@ async def history(ctx: sessions.ViewContext) -> list[sessions.Row]: if __name__ == "__main__": - sessions.serve(wallet) + asyncio.run(restate.serve(sessions.app(wallet))) diff --git a/examples-agents/tutorial/07_calls.py b/examples-agents/tutorial/07_calls.py index 685ebfd4..8e5ecea0 100644 --- a/examples-agents/tutorial/07_calls.py +++ b/examples-agents/tutorial/07_calls.py @@ -18,6 +18,8 @@ curl localhost:8080/handle/ace/owner # "alice" """ +import asyncio + from pydantic import BaseModel import restate @@ -85,4 +87,4 @@ async def balance(ctx: sessions.ViewContext) -> int: if __name__ == "__main__": - sessions.serve(wallet, handle) + asyncio.run(restate.serve(sessions.app(wallet, handle))) diff --git a/examples-agents/tutorial/08_sends.py b/examples-agents/tutorial/08_sends.py index eaa71bfe..316fc1d6 100644 --- a/examples-agents/tutorial/08_sends.py +++ b/examples-agents/tutorial/08_sends.py @@ -18,6 +18,8 @@ curl localhost:8080/wallet/alice/transfer --json '{"to": "bob", "amount": 5000}' # fails, and bob gets nothing """ +import asyncio + from pydantic import BaseModel import restate @@ -87,4 +89,4 @@ async def history(ctx: sessions.ViewContext) -> list[sessions.Row]: if __name__ == "__main__": - sessions.serve(wallet) + asyncio.run(restate.serve(sessions.app(wallet))) diff --git a/python/restate/agents/__init__.py b/python/restate/agents/__init__.py index d2d4a34a..f7eadc5b 100644 --- a/python/restate/agents/__init__.py +++ b/python/restate/agents/__init__.py @@ -13,6 +13,9 @@ A crash, a restart, or a person who takes a day to approve something changes nothing: each message is answered once, and each tool call is made once:: + import asyncio + + import restate from restate import sessions from restate.agents import langgraph, tool @@ -26,7 +29,7 @@ async def get_weather(city: str) -> str: assistant = langgraph.agent("assistant", system="Answer briefly.", tools=[get_weather]) if __name__ == "__main__": - sessions.serve(assistant) + asyncio.run(restate.serve(sessions.app(assistant))) This package holds what all frameworks share: diff --git a/python/restate/agents/langgraph/__init__.py b/python/restate/agents/langgraph/__init__.py index c09d5adc..d16f3545 100644 --- a/python/restate/agents/langgraph/__init__.py +++ b/python/restate/agents/langgraph/__init__.py @@ -12,6 +12,9 @@ For example:: + import asyncio + + import restate from restate import sessions from restate.agents import langgraph, tool @@ -21,7 +24,7 @@ async def get_weather(city: str) -> str: return f"Sunny in {city}" assistant = langgraph.agent("assistant", system="Answer briefly.", tools=[get_weather]) - sessions.serve(assistant) + asyncio.run(restate.serve(sessions.app(assistant))) ``agent`` creates an agent, ``prompt`` answers a prompt from a handler of your own, and ``SqlCheckpointSaver`` stores the checkpoints of a graph of your own diff --git a/python/restate/agents/langgraph/_agent.py b/python/restate/agents/langgraph/_agent.py index 7e07d81a..0e7c760f 100644 --- a/python/restate/agents/langgraph/_agent.py +++ b/python/restate/agents/langgraph/_agent.py @@ -389,7 +389,7 @@ def agent( invocation_retry_policy: How Restate retries a failed turn. Returns: - The agent, to serve with ``sessions.serve`` or ``sessions.app``. + The agent, to serve with ``restate.serve(sessions.app(agent))``, or any ASGI server. Raises: TerminalError: If ``model`` names a provider other than OpenAI or Anthropic. diff --git a/python/restate/agents/pydantic_ai/__init__.py b/python/restate/agents/pydantic_ai/__init__.py index f6c896d4..76f1ffc4 100644 --- a/python/restate/agents/pydantic_ai/__init__.py +++ b/python/restate/agents/pydantic_ai/__init__.py @@ -12,6 +12,9 @@ For example:: + import asyncio + + import restate from restate import sessions from restate.agents import pydantic_ai, tool @@ -21,7 +24,7 @@ async def get_weather(city: str) -> str: return f"Sunny in {city}" assistant = pydantic_ai.agent("assistant", system="Answer briefly.", tools=[get_weather]) - sessions.serve(assistant) + asyncio.run(restate.serve(sessions.app(assistant))) ``agent`` creates an agent, and ``prompt`` answers a prompt from a handler of your own. The conversation is stored one step at a time. diff --git a/python/restate/agents/pydantic_ai/_agent.py b/python/restate/agents/pydantic_ai/_agent.py index 183bb614..537e2443 100644 --- a/python/restate/agents/pydantic_ai/_agent.py +++ b/python/restate/agents/pydantic_ai/_agent.py @@ -453,7 +453,7 @@ def agent( invocation_retry_policy: How Restate retries a failed turn. Returns: - The agent, to serve with ``sessions.serve`` or ``sessions.app``. + The agent, to serve with ``restate.serve(sessions.app(agent))``, or any ASGI server. Raises: TypeError: If both ``pydantic_agent`` and ``system`` are given. diff --git a/python/restate/sessions/__init__.py b/python/restate/sessions/__init__.py index d78e0f0f..c1ee0745 100644 --- a/python/restate/sessions/__init__.py +++ b/python/restate/sessions/__init__.py @@ -13,6 +13,9 @@ Each key's database (``ctx.db``) is kept in the object's state in Restate, so there is no database to run:: + import asyncio + + import restate from restate import sessions notes = sessions.SessionObject("notes", migrations=["CREATE TABLE notes (text TEXT NOT NULL)"]) @@ -29,7 +32,7 @@ async def list_notes(ctx: sessions.ViewContext) -> list[str]: if __name__ == "__main__": - sessions.serve(notes) + asyncio.run(restate.serve(sessions.app(notes))) Handlers run one at a time per key; views run beside them and read what was committed. A handler is not replayed from a journal: a retried attempt runs @@ -56,7 +59,7 @@ async def list_notes(ctx: sessions.ViewContext) -> list[str]: from restate.sessions._database import Row, Rows, SqlError, SqlParam, SqlValue from restate.sessions._object import SessionObject from restate.sessions._schema import Migration -from restate.sessions._serve import app, serve +from restate.sessions._serve import app __all__ = [ "CallFuture", @@ -76,5 +79,4 @@ async def list_notes(ctx: sessions.ViewContext) -> list[str]: "Transaction", "ViewContext", "app", - "serve", ] diff --git a/python/restate/sessions/_serve.py b/python/restate/sessions/_serve.py index 6d9178e3..a3107744 100644 --- a/python/restate/sessions/_serve.py +++ b/python/restate/sessions/_serve.py @@ -11,26 +11,20 @@ """Serving session objects. ``app`` builds the ASGI app of a program's services, together with any -services they depend on (such as the ``signals`` object that agents use). Run -it with an ASGI server of your choice (hypercorn, uvicorn, ...), as you would -``restate.app``, or call ``serve`` to run it in one line. +services they depend on (such as the ``signals`` object that agents use). Serve +it as you would ``restate.app``: with ``restate.serve``, or any ASGI server:: + + asyncio.run(restate.serve(sessions.app(notes))) """ from __future__ import annotations -import asyncio -import contextlib -import importlib import os -import re -import signal import typing from collections.abc import Callable, Iterable, Sequence -from typing import Any import restate from restate.entry_codec import JournalValueCodec, JournalValueCodecProvider -from restate.server import install_shutdown_hook from restate.server_types import RestateAppT from restate.sessions._link import Link @@ -137,83 +131,3 @@ def app( identity_keys=identity_keys, journal_value_codec=journal_value_codec, ) - - -def serve( - *services: Service, - port: int | None = None, - identity_keys: list[str] | None = None, - journal_value_codec: JournalValueCodec | JournalValueCodecProvider | None = None, -) -> None: - """Serve ``services`` over HTTP until the process is stopped. - - This runs ``app(*services)`` on uvicorn with cleartext HTTP/2, which is - how Restate connects to services:: - - if __name__ == "__main__": - sessions.serve(wallet, handle) - - It needs the ``serve`` extra (``pip install "restate-sdk[serve]"``, which - installs uvicorn 0.53 or later and zttp) and Python 3.11 or later. On - other setups, serve ``app(...)`` with an ASGI server of your own. It stops - within a few seconds of SIGINT or SIGTERM. - - Args: - *services: The services, virtual objects (session objects included) - and workflows to serve, as for ``app``. - port: The port to listen on. Defaults to ``$PORT``, else 9080. - identity_keys: Passed to ``restate.app``. - journal_value_codec: Passed to ``restate.app``. - - Raises: - ImportError: If uvicorn 0.53 or later, or zttp, is not installed. - """ - try: - uvicorn = importlib.import_module("uvicorn") - importlib.import_module("zttp") - except ImportError as error: - raise ImportError(_NO_SERVER) from error - version = re.match(r"(\d+)\.(\d+)", str(getattr(uvicorn, "__version__", ""))) - if version is None or (int(version[1]), int(version[2])) < (0, 53): - raise ImportError(_NO_SERVER) - config = uvicorn.Config( - app(*services, identity_keys=identity_keys, journal_value_codec=journal_value_codec), - host="0.0.0.0", - port=port or int(os.environ.get("PORT", "9080")), - http="zttp", - http2=True, - # The endpoint has no lifespan; and a graceful shutdown would wait for Restate's open connections. - lifespan="off", - timeout_graceful_shutdown=_GRACEFUL_SHUTDOWN, - ) - asyncio.run(_until_stopped(uvicorn.Server(config))) - - -async def _until_stopped(server: Any) -> None: - """Run the uvicorn ``server`` until SIGINT or SIGTERM. - - The SDK's endpoint registers its own SIGTERM hook (to wind down the - invocations in flight) through ``install_shutdown_hook``, which keeps one - asyncio handler per signal. Uvicorn's handlers would be replaced by it on - the first request, so the server's exit is registered through the same - hook instead, and uvicorn captures no signals itself: both then run. - """ - loop = asyncio.get_running_loop() - - def exit_on(signum: int) -> Callable[[], None]: - return lambda: server.handle_exit(signum, None) - - for signum in (signal.SIGINT, signal.SIGTERM): - with contextlib.suppress(NotImplementedError, RuntimeError, ValueError): # Windows, or not the main thread - install_shutdown_hook(loop, exit_on(signum), signum) - server.capture_signals = contextlib.nullcontext - await server.serve() - - -_NO_SERVER = ( - "sessions.serve runs on uvicorn's HTTP/2: install restate-sdk[serve] (uvicorn 0.53 or later, and zttp; " - "Python 3.11 or later), or serve sessions.app(...) with an ASGI server of your own" -) - -_GRACEFUL_SHUTDOWN = 3 -"""How many seconds a stopping ``serve`` waits for requests in flight; Restate retries the rest.""" diff --git a/scripts/check_no_extras.py b/scripts/check_no_extras.py index 79b1c521..6560824c 100644 --- a/scripts/check_no_extras.py +++ b/scripts/check_no_extras.py @@ -23,7 +23,7 @@ - ``import restate`` loads neither ``restate.sessions`` nor ``restate.agents``; - an ordinary ``restate.app`` serves its discovery; - ``restate.sessions`` and ``restate.agents`` import, and their apps serve too; -- each optional part fails only when used, with an ``ImportError`` that names +- each optional part (``restate.serve`` included) fails only when used, with an ``ImportError`` that names what to install. """ @@ -116,7 +116,7 @@ async def add(ctx: sessions.SessionContext, text: str) -> None: refused("restate.agents.tool", lambda: restate.agents.tool(lambda city: city), "pydantic") refused("restate.agents.langgraph", lambda: __import__("restate.agents.langgraph"), "restate-sdk[langgraph]") refused("restate.agents.pydantic_ai", lambda: __import__("restate.agents.pydantic_ai"), "restate-sdk[pydantic_ai]") - refused("sessions.serve", lambda: sessions.serve(notes), "restate-sdk[serve]") + refused("restate.serve", lambda: asyncio.run(restate.serve(sessions.app(notes))), "restate-sdk[serve]") if __name__ == "__main__": diff --git a/tests/examples_agents.py b/tests/examples_agents.py index f7ec05a1..3f44ecd4 100644 --- a/tests/examples_agents.py +++ b/tests/examples_agents.py @@ -164,11 +164,11 @@ async def test_basic_pydantic_ai_agent(): def serves_http2() -> bool: - """Whether ``sessions.serve`` can run here: uvicorn 0.53+ with zttp.""" + """Whether ``restate.serve`` can run here: uvicorn 0.53+ with zttp.""" return importlib.util.find_spec("uvicorn") is not None and importlib.util.find_spec("zttp") is not None -@pytest.mark.skipif(not serves_http2(), reason="sessions.serve needs the serve extra (Python 3.11+)") +@pytest.mark.skipif(not serves_http2(), reason="restate.serve needs the serve extra (Python 3.11+)") async def test_an_example_as_a_program(): """``uv run`` an example: it serves on PORT, Restate calls it over HTTP/2, and it stops on SIGTERM.""" port = find_free_port() From af790b341ea62717c825fa44ca3a5e0ea453d103 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 20:10:31 +0100 Subject: [PATCH 44/45] Review fixes: a selector stall, a stale sync page, a re-asked answer, and three smaller ones - The selector could miss a wait registered while it flushed output before checking the VM: not in that pass's snapshot, and not waking it (it was not waiting yet). If the VM then waited for input, so did the loop, for a run (a commit's sync, say) that only starts once awaited. Such a wait now makes the loop take another pass. - An agent's sync returned a fresh revision with a page from its view's snapshot, taken before (by the whole wait, when it waited). It now reads the page in an invocation of its own (the messages view), after the revision. (A view's snapshot was confirmed against a real Restate.) - Pydantic AI: a final answer stored just before a crash, not yet marked as the prompt's, made the retry ask the model again. It is now marked. - Pydantic AI: messages store bytes as base64, as pydantic-ai does; a tool result of non-UTF-8 bytes failed every save, and so every attempt. - A stopped turn no longer reports a tool call whose Restate call returned as cancelled (the model would repeat it): it gets the error it returned, or a note that it happened. - Transaction's docstring no longer claims a retried body's sends are dropped: only a committed transaction's are made once. Co-Authored-By: Claude Opus 5.5 --- python/restate/agents/_calls.py | 13 ++++- python/restate/agents/_common.py | 19 ++++++++ python/restate/agents/_object.py | 7 ++- python/restate/agents/langgraph/_agent.py | 16 +++++-- python/restate/agents/pydantic_ai/_agent.py | 30 +++++++++--- python/restate/sessions/_api.py | 11 ++++- python/restate/sessions/_selector.py | 12 ++++- tests/agents_object.py | 24 ++++++++++ tests/agents_pydantic_ai.py | 53 +++++++++++++++++++++ tests/agents_tools.py | 31 ++++++++++++ tests/sessions_selector.py | 22 +++++++++ 11 files changed, 219 insertions(+), 19 deletions(-) diff --git a/python/restate/agents/_calls.py b/python/restate/agents/_calls.py index b3cff1e9..88e33188 100644 --- a/python/restate/agents/_calls.py +++ b/python/restate/agents/_calls.py @@ -33,7 +33,7 @@ import json import typing import uuid -from collections.abc import Callable +from collections.abc import Callable, Iterable from dataclasses import dataclass, field from restate.sessions import SessionContext @@ -165,3 +165,14 @@ async def settle(ctx: SessionContext, stored: Callable[[str, str | None], bool]) return state.calls = kept await _store(ctx, state) + + +async def outcomes(ctx: SessionContext, call_ids: Iterable[str]) -> dict[str, str | None]: + """Return which of ``call_ids`` made a Restate call that returned in this invocation. + + Each such id maps to the terminal error the call returned (``Error: ``), or to ``None`` if + it succeeded. A tool call whose Restate call has not returned, or that made none, is left out. + """ + _, state = await _calls_of(ctx) + wanted = set(call_ids) + return {i: call.get("failed") for call in state.calls if call.get("done") for i in call["ids"] if i in wanted} diff --git a/python/restate/agents/_common.py b/python/restate/agents/_common.py index fbb1a973..bd7dbe3c 100644 --- a/python/restate/agents/_common.py +++ b/python/restate/agents/_common.py @@ -156,6 +156,25 @@ def serialize(self, obj: str | None) -> bytes: CANCELLED = "Cancelled: the turn was stopped before it finished." """The answer of a cancelled turn, and the result of each tool call it cut off, which the model sees later in the conversation.""" +RETURNED_UNRECORDED = ( + "Done, but the turn was stopped before this call's result was recorded: it did happen, so check its " + "effect before doing it again." +) +"""The result of a tool call cut off by a stopped turn after its Restate call had returned successfully.""" + + +def cut_off_result(outcomes: dict[str, str | None], call_id: str, text: str) -> str: + """Return what a stopped turn records as the result of the tool call ``call_id``, which has none. + + A call whose Restate call returned is not reported as ``text`` (cancelled, say): the model would take + it for undone, and repeat its effect. It gets the error it returned, or ``RETURNED_UNRECORDED``. + """ + if call_id not in outcomes: + return text + failed = outcomes[call_id] + return failed if failed is not None else RETURNED_UNRECORDED + + UNANSWERED = "Unanswered: the model's provider refused the request." """The stored answer of a prompt that the model's provider refused for good. diff --git a/python/restate/agents/_object.py b/python/restate/agents/_object.py index a5fda5dd..3dd48a71 100644 --- a/python/restate/agents/_object.py +++ b/python/restate/agents/_object.py @@ -398,10 +398,13 @@ async def sync(ctx: ViewContext, request: SyncRequest | None = None) -> Synced: if revision is not None: watch: Watch = {"revision": revision, "timeoutSeconds": request.get("timeoutSeconds", 50)} await _companion(ctx, _signals.watch, "watch")(watch) - # Read the revision first, so that any change after it shows up in the next sync. + # Read the revision first, then the page in an invocation of its own, which starts after it: a view + # reads the database as committed when it started, and this one started before the revision (by + # the whole wait, when it waited). So the page has every change the revision counts. now = await _companion(ctx, _signals.notifications)() status = await _companion(ctx, _signals.status)() - page = await _log.page(ctx.db, {"cursor": request.get("cursor", 0), "limit": request.get("limit", 100)}) + wanted: MessagesRequest = {"cursor": request.get("cursor", 0), "limit": request.get("limit", 100)} + page = await ctx.call(messages, wanted, key=ctx.key) synced: Synced = {**page, "revision": now["revision"], "status": status} return synced diff --git a/python/restate/agents/langgraph/_agent.py b/python/restate/agents/langgraph/_agent.py index 0e7c760f..e1cd1792 100644 --- a/python/restate/agents/langgraph/_agent.py +++ b/python/restate/agents/langgraph/_agent.py @@ -50,6 +50,7 @@ from langgraph.prebuilt import ToolNode, tools_condition from pydantic import PrivateAttr +from restate.agents._calls import outcomes from restate.agents._common import ( CANCELLED, MODEL_ENV, @@ -57,6 +58,7 @@ ModelError, PerSession, configured, + cut_off_result, from_tool, model_of, next_prompt_id, @@ -244,11 +246,11 @@ async def _answer_once( except TerminalError as error: # Cancelled: close the run, so that the next prompt starts afresh. if error.status_code == _CANCELLATION: - await _close(app, config, CANCELLED) + await _close(ctx, app, config, CANCELLED) raise # Failed for good (a LangChain tool's terminal error): close the run with the error, so that the next # prompt does not run into it again. If the run was an earlier prompt's, carry on with this one. - await _close(app, config, f"Error: {error.message}") + await _close(ctx, app, config, f"Error: {error.message}") if not interrupted or last == prompt_id: raise except Exception as error: @@ -257,12 +259,12 @@ async def _answer_once( refused = refusal_of(ctx, error) if refused is None: raise - await _close(app, config, UNANSWERED) + await _close(ctx, app, config, UNANSWERED) if not interrupted or last == prompt_id: raise refused from error -async def _close(app: Compiled, config: RunnableConfig, text: str) -> None: +async def _close(ctx: SessionContext, app: Compiled, config: RunnableConfig, text: str) -> None: """Close a cancelled, refused or failed run of ``app`` with ``text``, as its ``answer`` node would. Each tool call of the last reply that has no result gets ``text`` as its @@ -278,8 +280,12 @@ async def _close(app: Compiled, config: RunnableConfig, text: str) -> None: answered = {m.tool_call_id for m in messages if isinstance(m, ToolMessage)} last = next((m for m in reversed(messages) if isinstance(m, AIMessage)), None) calls = [c for c in (last.tool_calls if last else []) if c.get("id") and c.get("id") not in answered] + returned = await outcomes(ctx, [str(c.get("id")) for c in calls]) closing: list[BaseMessage] = [ - *(ToolMessage(text, tool_call_id=str(c.get("id")), status="error") for c in calls), + *( + ToolMessage(cut_off_result(returned, str(c.get("id")), text), tool_call_id=str(c.get("id")), status="error") + for c in calls + ), AIMessage(text), ] await app.aupdate_state(config, {"messages": closing, "answers": {prompt: text}}, as_node="answer") diff --git a/python/restate/agents/pydantic_ai/_agent.py b/python/restate/agents/pydantic_ai/_agent.py index 537e2443..917d519d 100644 --- a/python/restate/agents/pydantic_ai/_agent.py +++ b/python/restate/agents/pydantic_ai/_agent.py @@ -35,7 +35,7 @@ from datetime import timedelta from typing import Any -from pydantic import TypeAdapter +from pydantic import ConfigDict, TypeAdapter from pydantic_ai.capabilities import AbstractCapability from pydantic_ai.exceptions import UserError from pydantic_ai.messages import ( @@ -54,12 +54,14 @@ from pydantic_ai import Agent, RunContext from pydantic_ai import Tool as PydanticTool +from restate.agents._calls import outcomes from restate.agents._common import ( CANCELLED, UNANSWERED, ModelError, PerSession, configured, + cut_off_result, from_tool, model_of, next_prompt_id, @@ -127,7 +129,10 @@ async def run(run_ctx: RunContext[Any], **args: Any) -> str: ) # ModelMessage is an Annotated union (discriminated), which TypeAdapter takes though its type says otherwise. -_MESSAGE: TypeAdapter[ModelMessage] = TypeAdapter(typing.cast(Any, ModelMessage)) +# Bytes (a tool's binary result, say) as base64, as pydantic-ai's own ModelMessagesTypeAdapter stores them. +_MESSAGE: TypeAdapter[ModelMessage] = TypeAdapter( + typing.cast(Any, ModelMessage), config=ConfigDict(ser_json_bytes="base64", val_json_bytes="base64") +) class _Conversation: @@ -255,7 +260,7 @@ def _text(response: ModelMessage) -> str: return "".join(part.content for part in response.parts if isinstance(part, TextPart)) -async def _close(conversation: _Conversation, prompt_id: str, text: str) -> None: +async def _close(ctx: SessionContext, conversation: _Conversation, prompt_id: str, text: str) -> None: """Close the cancelled, refused or failed run of prompt ``prompt_id`` with ``text``. Each tool call of the last response gets ``text`` as its result, and ``text`` @@ -267,7 +272,11 @@ async def _close(conversation: _Conversation, prompt_id: str, text: str) -> None return last = messages[-1] if isinstance(last, ModelResponse) and last.tool_calls: - returns = [ToolReturnPart(call.tool_name, text, call.tool_call_id) for call in last.tool_calls] + returned = await outcomes(ctx, [call.tool_call_id for call in last.tool_calls]) + returns = [ + ToolReturnPart(call.tool_name, cut_off_result(returned, call.tool_call_id, text), call.tool_call_id) + for call in last.tool_calls + ] messages.append(ModelRequest(parts=returns)) messages.append(ModelResponse(parts=[TextPart(text)], metadata={_ANSWERS: prompt_id})) await conversation.save(messages) @@ -297,6 +306,13 @@ async def _answer_once( interrupted = last is not None and _answer_of(history, last) is None running = _prompt_of(history[last]) if interrupted and last is not None else prompt_id assert running is not None + final = history[-1] if history else None + if interrupted and isinstance(final, ModelResponse) and not final.tool_calls: + # The run had stored its final answer, but not yet marked it as the prompt's: mark it, rather + # than run on, which would ask the model again. + final.metadata = {**(final.metadata or {}), _ANSWERS: running} + await conversation.save(history) + continue if not interrupted: history = [*history, ModelRequest(parts=[UserPromptPart(text)], metadata={_PROMPT: prompt_id})] stop = stop_signal_of(ctx) @@ -311,10 +327,10 @@ async def _answer_once( except TerminalError as error: if error.status_code == _CANCELLATION: # Cancelled: close the run, so that the next prompt starts afresh. - await _close(conversation, running, CANCELLED) + await _close(ctx, conversation, running, CANCELLED) raise # Failed for good: close the run with the error, so that the next prompt does not run into it again. - await _close(conversation, running, f"Error: {error.message}") + await _close(ctx, conversation, running, f"Error: {error.message}") if running == prompt_id: raise continue @@ -324,7 +340,7 @@ async def _answer_once( refused = refusal_of(ctx, error) if refused is None: raise - await _close(conversation, running, UNANSWERED) + await _close(ctx, conversation, running, UNANSWERED) if running == prompt_id: raise refused from error continue diff --git a/python/restate/sessions/_api.py b/python/restate/sessions/_api.py index 87742120..e9c12b35 100644 --- a/python/restate/sessions/_api.py +++ b/python/restate/sessions/_api.py @@ -237,8 +237,15 @@ class Transaction(DatabaseWriter, Sender, Protocol): """One transaction of a ``Database``: its SQL, and the sends it commits with. ``tx.send`` sends right after the transaction commits, and never if it - rolls back. A send without an ``idempotency_key`` gets one from the - transaction, so a send that a retried attempt makes again is dropped. + rolls back. Each committed send goes out once: if the process fails after + the commit but before the send, the next attempt makes it, under the same + idempotency key, and Restate drops any duplicate. + + A retried attempt that runs the transaction's body again commits a new + transaction, with sends of its own. To send once across attempts, give the + send an ``idempotency_key`` that is the same on every attempt (one made + from ``ctx.request.id``, say), or skip the body when its rows show that it + already ran. """ diff --git a/python/restate/sessions/_selector.py b/python/restate/sessions/_selector.py index 8fb1185c..2f5184c3 100644 --- a/python/restate/sessions/_selector.py +++ b/python/restate/sessions/_selector.py @@ -110,6 +110,9 @@ def __init__(self, driver: SelectorDriver, disconnected: Callable[[], BaseExcept self._task: asyncio.Task[None] | None = None # True between the VM's WaitExternalProgress and the next input. self._waiting: bool = False + # True when a wait was registered after the loop took its snapshot of the pending waits: the loop + # then takes another pass instead of waiting for input, which might never come for that wait. + self._missed: bool = False # Why the attempt ended, once it has; each later wait raises a fresh instance. self._end: Callable[[], BaseException] | None = None @@ -142,10 +145,12 @@ def _start(self) -> None: self._task = asyncio.create_task(self._loop()) self._driver.track(self._task) elif self._waiting: - # Only a loop that waits for input needs waking, and only once; otherwise - # the loop picks up the new wait on its next pass. + # A loop that waits for input needs waking, once. self._waiting = False self._driver.wake() + else: + # A loop between passes (flushing output, say) may have taken its snapshot already. + self._missed = True def _settle(self) -> None: for waiter, future in self._waits.items(): @@ -172,6 +177,7 @@ async def _loop(self) -> None: # waits, so that one check covers them. This is only an # optimization: a later registration wakes the loop anyway. await asyncio.sleep(0) + self._missed = False pending = [future for waiter, future in self._waits.items() if not waiter.done()] if not pending: return @@ -196,6 +202,8 @@ async def _loop(self) -> None: driver.execute_run(response.handle) continue if isinstance(response, DoProgressWaitExternalProgress): + if self._missed: + continue self._waiting = True await driver.flush_output() chunk = await driver.receive() diff --git a/tests/agents_object.py b/tests/agents_object.py index 6ed621c0..ad39351e 100644 --- a/tests/agents_object.py +++ b/tests/agents_object.py @@ -152,6 +152,30 @@ async def test_a_message_is_answered_and_the_conversation_read_and_followed(harn assert [m["id"] for m in newer["messages"]] == [2, 3] +async def test_a_sync_that_waited_returns_every_change_its_revision_counts(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + await call(harness, "chat", key, "Hi") + synced = (await call(harness, "sync", key, {})).json() + timeout = 20 + started = asyncio.get_running_loop().time() + + async def follow() -> list[int]: + # As a client follows: each sync waits for a revision newer than the last one it got. + state, seen = synced, [] + while 3 not in seen: + request = {"revision": state["revision"], "cursor": state["cursor"], "timeoutSeconds": timeout} + state = (await call(harness, "sync", key, request)).json() + seen += [message["id"] for message in state["messages"]] + return seen + + following = asyncio.ensure_future(follow()) + await asyncio.sleep(0.5) # the first sync waits + await call(harness, "chat", key, "Again") + assert 2 in await following and 3 in following.result() + # Every message came with the revision that counted it: no sync had to time out to bring it. + assert asyncio.get_running_loop().time() - started < timeout + + async def test_status_says_which_message_is_answered_and_stop_cancels_it(harness: HarnessEnvironment): key = str(uuid.uuid4()) gate.clear() diff --git a/tests/agents_pydantic_ai.py b/tests/agents_pydantic_ai.py index e7c6bd60..6b951244 100644 --- a/tests/agents_pydantic_ai.py +++ b/tests/agents_pydantic_ai.py @@ -328,3 +328,56 @@ async def test_options_may_be_functions_of_the_context(harness: HarnessEnvironme url = f"{harness.ingress_url}/pai_per_user/{k}/chat" assert (await ingress.post(url, content=f"system {k}", headers=TEXT)).json() == f"system: You help {k}." assert (await ingress.post(url, content=f"weather {k}", headers=TEXT)).json() == f"tool said: sunny in {k}" + + +async def test_an_answer_stored_before_a_crash_is_not_asked_for_again(monkeypatch: pytest.MonkeyPatch): + from restate.agents.pydantic_ai import _agent # pylint: disable=import-outside-toplevel + from sessions_database import Backend # pylint: disable=import-error,import-outside-toplevel + from sessions_object import Host, request # pylint: disable=import-error,import-outside-toplevel + + from restate.sessions.internal import invoke # pylint: disable=import-outside-toplevel + + replies: list[str] = [] + + def answer(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + replies.append("asked") + return ModelResponse(parts=[TextPart(f"answer {len(replies)}")]) + + save = _agent._Conversation.save # pyright: ignore[reportPrivateUsage] + crashed: list[bool] = [] + + async def save_then_crash_once(self: Any, messages: Any) -> None: + await save(self, messages) + last = messages[-1] + if isinstance(last, ModelResponse) and not (last.metadata or {}) and not crashed: + crashed.append(True) + raise RuntimeError("the process crashed after storing the answer") + + monkeypatch.setattr(_agent._Conversation, "save", save_then_crash_once) # pyright: ignore[reportPrivateUsage] + asker = sessions.SessionObject("asker") + + @asker.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + return await pydantic_ai.prompt(ctx, text, model=FunctionModel(answer), system="Answer briefly.") + + backend, invocation = Backend(), request() + with pytest.raises(RuntimeError, match="crashed after storing"): + await invoke(asker, "ask", Host(backend, service="asker"), invocation, "hi") + # The retry (the same invocation) returns the answer that was stored, without asking the model again. + assert await invoke(asker, "ask", Host(backend, service="asker"), invocation, "hi") == "answer 1" + assert replies == ["asked"] + + +def test_a_tool_result_of_bytes_is_stored_as_pydantic_ai_stores_it(): + from pydantic_ai.messages import ModelMessagesTypeAdapter # pylint: disable=import-outside-toplevel + + from restate.agents.pydantic_ai import _agent # pylint: disable=import-outside-toplevel + + message = ModelRequest(parts=[ToolReturnPart("picture", b"\xff\x00\xfe", tool_call_id="call-1")]) + # Storing it does not fail (a failure would fail every attempt, and run the tool again each time) ... + stored = _agent._MESSAGE.dump_json(message) # pyright: ignore[reportPrivateUsage] + read = _agent._MESSAGE.validate_json(stored) # pyright: ignore[reportPrivateUsage] + # ... and it reads back as pydantic-ai's own adapter reads it back. + [expected] = ModelMessagesTypeAdapter.validate_json(ModelMessagesTypeAdapter.dump_json([message])) + assert isinstance(read, ModelRequest) and isinstance(expected, ModelRequest) + assert read.parts[0].content == expected.parts[0].content # type: ignore[union-attr] diff --git a/tests/agents_tools.py b/tests/agents_tools.py index a93983e4..eb75f3ce 100644 --- a/tests/agents_tools.py +++ b/tests/agents_tools.py @@ -151,3 +151,34 @@ async def test_subagent_hands_a_task_to_the_agent_keyed_by_this_key_and_the_task assert (service, handler, key, json.loads(arg)) == ("helper", "chat", "alice.lookup", "Find it") bad = await run_tool(host, "helper", {"name": "Not A Name", "task": "x"}, "s2") assert bad.startswith("Error: the arguments do not fit") + + +async def test_a_stopped_turn_reports_a_call_that_returned_as_done_not_cancelled(): + from restate.agents._calls import call_of, outcomes # pylint: disable=import-outside-toplevel + from restate.agents._common import ( # pylint: disable=import-outside-toplevel + CANCELLED, + RETURNED_UNRECORDED, + cut_off_result, + ) + + recorded = sessions.SessionObject("stopped") + seen: dict[str, str] = {} + ids = ("call-refund", "call-status", "call-charge", "call-none") + + @recorded.handler + async def turn(ctx: sessions.SessionContext) -> None: + refund = await call_of(ctx, "refund", {"order": "o1"}, "call-refund") + await refund.answered() + await call_of(ctx, "status", {"order": "o1"}, "call-status") # made, but not returned + charge = await call_of(ctx, "charge", {"order": "o1"}, "call-charge") + await charge.answered(failed="Error: card declined") + returned = await outcomes(ctx, ids) + seen.update({call_id: cut_off_result(returned, call_id, CANCELLED) for call_id in ids}) + + await internal.invoke(recorded, "turn", Host(Backend(), service="stopped"), request()) + assert seen == { + "call-refund": RETURNED_UNRECORDED, # it happened: the model must not take it for undone + "call-status": CANCELLED, + "call-charge": "Error: card declined", + "call-none": CANCELLED, + } diff --git a/tests/sessions_selector.py b/tests/sessions_selector.py index 8fd95639..02781e02 100644 --- a/tests/sessions_selector.py +++ b/tests/sessions_selector.py @@ -151,6 +151,28 @@ async def test_a_wait_made_while_the_loop_waits_for_input_joins_the_next_check() await asyncio.gather(first, second, third) +async def test_a_wait_made_while_the_loop_flushes_is_checked_before_it_waits_for_input(): + vm = FakeVm() + selector = FutureSelector(vm, Disconnected) + late: list[asyncio.Future] = [] + + async def flush_letting_a_wait_register() -> None: + if not late: + late.append(asyncio.ensure_future(selector.wait(single(2)))) + await asyncio.sleep(0) + await asyncio.sleep(0) + + vm.flush_output = flush_letting_a_wait_register # type: ignore[method-assign] + first = asyncio.ensure_future(selector.wait(single(1))) + await settle() + await settle() + # Not left out of the check the loop was making while it registered, nor waiting for input or a wake. + assert vm.checks[-1] == FirstCompletedUnresolvedFuture([single(1), single(2)]) + vm.complete(2) + vm.complete(1) + await asyncio.gather(first, *late) + + async def test_combinators_wait_for_their_children(): vm = FakeVm() selector = FutureSelector(vm, Disconnected) From 226a649de1ca0dbcebf57bfa055f8e4f035c7b56 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Wed, 4 Nov 2026 20:27:35 +0100 Subject: [PATCH 45/45] tests/examples_agents.py: start Restate before choosing the example's port, and wait until it listens As in tests/serve.py (#266): Docker could map one of the Restate container's ports onto the port chosen for the example before the example bound it. The container now starts first, the test waits until the example accepts connections, and SIGTERM comes while Restate still holds its connections open. Co-Authored-By: Claude Opus 5.5 --- tests/examples_agents.py | 50 +++++++++++++++++++++++++--------------- 1 file changed, 32 insertions(+), 18 deletions(-) diff --git a/tests/examples_agents.py b/tests/examples_agents.py index 3f44ecd4..f2ca94d8 100644 --- a/tests/examples_agents.py +++ b/tests/examples_agents.py @@ -169,36 +169,50 @@ def serves_http2() -> bool: @pytest.mark.skipif(not serves_http2(), reason="restate.serve needs the serve extra (Python 3.11+)") -async def test_an_example_as_a_program(): +async def test_an_example_as_a_program(tmp_path: pathlib.Path): """``uv run`` an example: it serves on PORT, Restate calls it over HTTP/2, and it stops on SIGTERM.""" - port = find_free_port() - program = subprocess.Popen( - [sys.executable, str(EXAMPLES / "basic/notes.py")], - env={**os.environ, "PORT": str(port)}, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) + # Restate first: Docker maps its ports to free host ports, which could otherwise take the program's. container = RestateContainer("ghcr.io/restatedev/restate:latest", always_replay=False, disable_retries=False) + container.start() + port = find_free_port() + output = tmp_path / "program.log" + with output.open("wb") as log: + program = subprocess.Popen( + [sys.executable, str(EXAMPLES / "basic/notes.py")], + env={**os.environ, "PORT": str(port)}, + stdout=log, + stderr=log, + ) try: - container.start() + await listening(program, port, output) async with httpx.AsyncClient(base_url=container.admin_url(), timeout=30) as admin: uri = f"http://host.docker.internal:{port}" registered = await admin.post("/deployments", json={"uri": uri, "force": True}) - for _ in range(100): # until the program serves - if registered.status_code < 300: - break - await asyncio.sleep(0.1) - registered = await admin.post("/deployments", json={"uri": uri, "force": True}) assert registered.status_code < 300, registered.text deployment = (await admin.get(f"/deployments/{registered.json()['id']}")).json() assert deployment.get("http_version") == "HTTP/2.0", deployment ingress = Ingress(container.ingress_url()) await ingress.post("notes/alice/add", "Buy milk") assert (await ingress.post("notes/alice/list")).json() == ["Buy milk"] + + # Stopped while Restate still holds its connections open: it stops within its graceful timeout. + program.send_signal(signal.SIGTERM) + program.wait(timeout=10) finally: + program.kill() container.stop() - program.send_signal(signal.SIGTERM) + + +async def listening(program: subprocess.Popen[bytes], port: int, output: pathlib.Path) -> None: + """Wait until ``program`` accepts connections on ``port``, or fail with what it printed.""" + for _ in range(200): + assert program.poll() is None, f"the program exited:\n{output.read_text()}" try: - program.wait(timeout=10) - finally: - program.kill() + _, writer = await asyncio.open_connection("127.0.0.1", port) + except OSError: + await asyncio.sleep(0.05) + continue + writer.close() + await writer.wait_closed() + return + raise AssertionError(f"the program never listened on {port}:\n{output.read_text()}")