runkite-runner 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (63) hide show
  1. runkite_runner-0.1.0/LICENSE +113 -0
  2. runkite_runner-0.1.0/PKG-INFO +49 -0
  3. runkite_runner-0.1.0/README.md +15 -0
  4. runkite_runner-0.1.0/pyproject.toml +51 -0
  5. runkite_runner-0.1.0/runkite_runner/__init__.py +3 -0
  6. runkite_runner-0.1.0/runkite_runner/__main__.py +5 -0
  7. runkite_runner-0.1.0/runkite_runner/a2a.py +133 -0
  8. runkite_runner-0.1.0/runkite_runner/checkpoint.py +214 -0
  9. runkite_runner-0.1.0/runkite_runner/custom_app.py +62 -0
  10. runkite_runner-0.1.0/runkite_runner/custom_auth.py +82 -0
  11. runkite_runner-0.1.0/runkite_runner/custom_helpers.py +76 -0
  12. runkite_runner-0.1.0/runkite_runner/factory_graph.py +315 -0
  13. runkite_runner-0.1.0/runkite_runner/generic_worker.py +521 -0
  14. runkite_runner-0.1.0/runkite_runner/heartbeat.py +100 -0
  15. runkite_runner-0.1.0/runkite_runner/logging_config.py +64 -0
  16. runkite_runner-0.1.0/runkite_runner/otel_callbacks.py +176 -0
  17. runkite_runner-0.1.0/runkite_runner/pg_pool.py +28 -0
  18. runkite_runner-0.1.0/runkite_runner/py.typed +0 -0
  19. runkite_runner-0.1.0/runkite_runner/run_status.py +57 -0
  20. runkite_runner-0.1.0/runkite_runner/runner_loop.py +21 -0
  21. runkite_runner-0.1.0/runkite_runner/runner_pb2.py +57 -0
  22. runkite_runner-0.1.0/runkite_runner/runner_pb2_grpc.py +299 -0
  23. runkite_runner-0.1.0/runkite_runner/schema_introspect.py +113 -0
  24. runkite_runner-0.1.0/runkite_runner/store.py +505 -0
  25. runkite_runner-0.1.0/runkite_runner/tenant_ctx.py +57 -0
  26. runkite_runner-0.1.0/runkite_runner/tls_utils.py +101 -0
  27. runkite_runner-0.1.0/runkite_runner/tracing.py +169 -0
  28. runkite_runner-0.1.0/runkite_runner/vectorstore.py +380 -0
  29. runkite_runner-0.1.0/runkite_runner/worker.py +1021 -0
  30. runkite_runner-0.1.0/runkite_runner.egg-info/PKG-INFO +49 -0
  31. runkite_runner-0.1.0/runkite_runner.egg-info/SOURCES.txt +61 -0
  32. runkite_runner-0.1.0/runkite_runner.egg-info/dependency_links.txt +1 -0
  33. runkite_runner-0.1.0/runkite_runner.egg-info/entry_points.txt +2 -0
  34. runkite_runner-0.1.0/runkite_runner.egg-info/requires.txt +12 -0
  35. runkite_runner-0.1.0/runkite_runner.egg-info/top_level.txt +1 -0
  36. runkite_runner-0.1.0/setup.cfg +4 -0
  37. runkite_runner-0.1.0/tests/test_a2a.py +155 -0
  38. runkite_runner-0.1.0/tests/test_checkpoint_concurrent_setup.py +65 -0
  39. runkite_runner-0.1.0/tests/test_checkpoint_pool_recreate.py +63 -0
  40. runkite_runner-0.1.0/tests/test_custom_app.py +91 -0
  41. runkite_runner-0.1.0/tests/test_custom_auth.py +51 -0
  42. runkite_runner-0.1.0/tests/test_factory_graph.py +469 -0
  43. runkite_runner-0.1.0/tests/test_generic_worker.py +552 -0
  44. runkite_runner-0.1.0/tests/test_heartbeat.py +212 -0
  45. runkite_runner-0.1.0/tests/test_langchain_adapter.py +216 -0
  46. runkite_runner-0.1.0/tests/test_logging_config.py +111 -0
  47. runkite_runner-0.1.0/tests/test_otel_callbacks.py +146 -0
  48. runkite_runner-0.1.0/tests/test_protocol_execute_goldens.py +290 -0
  49. runkite_runner-0.1.0/tests/test_protocol_llm_structural.py +187 -0
  50. runkite_runner-0.1.0/tests/test_run_status.py +49 -0
  51. runkite_runner-0.1.0/tests/test_schema_introspect.py +172 -0
  52. runkite_runner-0.1.0/tests/test_store_batch_cross_loop.py +64 -0
  53. runkite_runner-0.1.0/tests/test_store_dual_mode.py +207 -0
  54. runkite_runner-0.1.0/tests/test_store_pool.py +62 -0
  55. runkite_runner-0.1.0/tests/test_store_pool_recreate.py +56 -0
  56. runkite_runner-0.1.0/tests/test_tenant_ctx.py +107 -0
  57. runkite_runner-0.1.0/tests/test_tls_utils.py +173 -0
  58. runkite_runner-0.1.0/tests/test_tool_call_hook.py +192 -0
  59. runkite_runner-0.1.0/tests/test_tracing.py +89 -0
  60. runkite_runner-0.1.0/tests/test_vectorstore_cross_loop.py +82 -0
  61. runkite_runner-0.1.0/tests/test_vectorstore_dual_mode.py +118 -0
  62. runkite_runner-0.1.0/tests/test_worker_cancel_race.py +85 -0
  63. runkite_runner-0.1.0/tests/test_worker_concurrency.py +315 -0
@@ -0,0 +1,113 @@
1
+ Business Source License 1.1
2
+
3
+ Parameters
4
+
5
+ Licensor: Sharan Harsoor
6
+ Licensed Work: Runkite
7
+ The Licensed Work is (c) 2026 Sharan Harsoor
8
+ Additional Use Grant: You may use, copy, modify, and self-host the Licensed
9
+ Work, including for internal production use within
10
+ your own organization, free of charge.
11
+
12
+ You may NOT offer the Licensed Work, or any
13
+ substantial part of its functionality, to third
14
+ parties as a hosted or managed service (e.g. a
15
+ commercial SaaS offering of Runkite itself) without a
16
+ separate commercial license from the Licensor.
17
+
18
+ Change Date: 2030-07-27
19
+
20
+ Change License: Apache License, Version 2.0
21
+
22
+ For information about alternative licensing arrangements for the Licensed
23
+ Work, please contact the Licensor.
24
+
25
+ Notice
26
+
27
+ The Business Source License (this document, or the "License") is not an
28
+ Open Source license. However, the Licensed Work will eventually be made
29
+ available under an Open Source License, as stated in this License.
30
+
31
+ License text copyright (c) 2017 MariaDB Corporation Ab, All Rights
32
+ Reserved. "Business Source License" is a trademark of MariaDB Corporation
33
+ Ab.
34
+
35
+ -----------------------------------------------------------------------------
36
+
37
+ Business Source License 1.1
38
+
39
+ Terms
40
+
41
+ The Licensor hereby grants you the right to copy, modify, create
42
+ derivative works, redistribute, and make non-production use of the
43
+ Licensed Work. The Licensor may make an Additional Use Grant, above,
44
+ permitting limited production use.
45
+
46
+ Effective on the Change Date, or the fourth anniversary of the first
47
+ publicly available distribution of a specific version of the Licensed
48
+ Work under this License, whichever comes first, the Licensor hereby
49
+ grants you rights under the terms of the Change License, and the rights
50
+ granted in the paragraph above terminate.
51
+
52
+ If your use of the Licensed Work does not comply with the requirements
53
+ currently in effect as described in this License, you must purchase a
54
+ commercial license from the Licensor, its affiliated entities, or
55
+ authorized resellers, or you must refrain from using the Licensed Work.
56
+
57
+ All copies of the original and modified Licensed Work, and derivative
58
+ works of the Licensed Work, are subject to this License. This License
59
+ applies separately for each version of the Licensed Work and the Change
60
+ Date may vary for each version of the Licensed Work released by
61
+ Licensor.
62
+
63
+ You must conspicuously display this License on each original or
64
+ modified copy of the Licensed Work. If you receive the Licensed Work in
65
+ original or modified form from a third party, the terms and conditions
66
+ set forth in this License apply to your use of that work.
67
+
68
+ Any use of the Licensed Work in violation of this License will
69
+ automatically terminate your rights under this License for the current
70
+ and all other versions of the Licensed Work.
71
+
72
+ This License does not grant you any right in any trademark or logo of
73
+ Licensor or its affiliates (provided that you may use a trademark or
74
+ logo of Licensor as expressly required by this License).
75
+
76
+ TO THE EXTENT PERMITTED BY APPLICABLE LAW, THE LICENSED WORK IS PROVIDED
77
+ ON AN "AS IS" BASIS. LICENSOR HEREBY DISCLAIMS ALL WARRANTIES AND
78
+ CONDITIONS, EXPRESS OR IMPLIED, INCLUDING (WITHOUT LIMITATION)
79
+ WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE,
80
+ NON-INFRINGEMENT, AND TITLE.
81
+
82
+ MariaDB hereby grants you permission to use this License's text to
83
+ license your works, and to refer to it using the trademark "Business
84
+ Source License", as long as you comply with the Covenants of Licensor
85
+ below.
86
+
87
+ Covenants of Licensor
88
+
89
+ In consideration of the right to use this License's text and the
90
+ "Business Source License" name and trademark, Licensor covenants to
91
+ MariaDB, and to all other recipients of the licensed work to be provided
92
+ by Licensor:
93
+
94
+ 1. To specify as the Change License the GPL Version 2.0 or any later
95
+ version, or a license that is compatible with GPL Version 2.0 or a
96
+ later version, where "compatible" means that software provided under
97
+ the Change License can be included in a program with software
98
+ provided under GPL Version 2.0 or a later version. Licensor may
99
+ specify additional Change Licenses without limitation.
100
+
101
+ 2. To either: (a) specify an additional grant of rights to use that does
102
+ not impose any additional restriction on the right granted in this
103
+ License, as the Additional Use Grant, or (b) insert the text "None".
104
+
105
+ 3. To specify a Change Date.
106
+
107
+ 4. Not to modify this License in any other way.
108
+
109
+ Notice
110
+
111
+ The Business Source License (this document, or the "License") is not an
112
+ Open Source license. However, the Licensed Work will eventually be made
113
+ available under an Open Source License, as stated in this License.
@@ -0,0 +1,49 @@
1
+ Metadata-Version: 2.4
2
+ Name: runkite-runner
3
+ Version: 0.1.0
4
+ Summary: Python runner SDK for the Runkite Agent Protocol control plane (LangGraph and adapters).
5
+ Author: Runkite contributors
6
+ License: BUSL-1.1
7
+ Project-URL: Homepage, https://getrunkite.github.io/runkite/
8
+ Project-URL: Repository, https://github.com/getrunkite/runkite
9
+ Project-URL: Documentation, https://github.com/getrunkite/runkite/tree/main/docs
10
+ Project-URL: Changelog, https://github.com/getrunkite/runkite/blob/main/CHANGELOG.md
11
+ Project-URL: Issues, https://github.com/getrunkite/runkite/issues
12
+ Keywords: runkite,agent-protocol,langgraph,grpc,agents
13
+ Classifier: Development Status :: 4 - Beta
14
+ Classifier: License :: Other/Proprietary License
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Programming Language :: Python :: 3 :: Only
17
+ Classifier: Topic :: Software Development :: Libraries
18
+ Requires-Python: >=3.11
19
+ Description-Content-Type: text/markdown
20
+ License-File: LICENSE
21
+ Requires-Dist: grpcio
22
+ Requires-Dist: protobuf
23
+ Requires-Dist: langgraph
24
+ Requires-Dist: langgraph-sdk
25
+ Requires-Dist: langchain-core
26
+ Requires-Dist: langgraph-checkpoint-postgres
27
+ Requires-Dist: psycopg[binary,pool]
28
+ Requires-Dist: httpx
29
+ Requires-Dist: uvicorn
30
+ Requires-Dist: opentelemetry-api
31
+ Requires-Dist: opentelemetry-sdk
32
+ Requires-Dist: opentelemetry-exporter-otlp
33
+ Dynamic: license-file
34
+
35
+ # runkite-runner (Python)
36
+
37
+ Python runner for the [Runkite](https://github.com/getrunkite/runkite) control plane.
38
+ Connects over gRPC, executes LangGraph (and adapter) agents, streams events back.
39
+
40
+ ```bash
41
+ pip install runkite-runner
42
+ runkite-runner --config path/to/langgraph.json \
43
+ --grpc-address 127.0.0.1:50051 \
44
+ --http-address http://127.0.0.1:2026
45
+ ```
46
+
47
+ Docs: [docs/runners.md](https://github.com/getrunkite/runkite/blob/main/docs/runners.md) · Site: https://getrunkite.github.io/runkite/
48
+
49
+ License: [BUSL-1.1](https://github.com/getrunkite/runkite/blob/main/LICENSE)
@@ -0,0 +1,15 @@
1
+ # runkite-runner (Python)
2
+
3
+ Python runner for the [Runkite](https://github.com/getrunkite/runkite) control plane.
4
+ Connects over gRPC, executes LangGraph (and adapter) agents, streams events back.
5
+
6
+ ```bash
7
+ pip install runkite-runner
8
+ runkite-runner --config path/to/langgraph.json \
9
+ --grpc-address 127.0.0.1:50051 \
10
+ --http-address http://127.0.0.1:2026
11
+ ```
12
+
13
+ Docs: [docs/runners.md](https://github.com/getrunkite/runkite/blob/main/docs/runners.md) · Site: https://getrunkite.github.io/runkite/
14
+
15
+ License: [BUSL-1.1](https://github.com/getrunkite/runkite/blob/main/LICENSE)
@@ -0,0 +1,51 @@
1
+ [build-system]
2
+ requires = ["setuptools>=68", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "runkite-runner"
7
+ version = "0.1.0"
8
+ description = "Python runner SDK for the Runkite Agent Protocol control plane (LangGraph and adapters)."
9
+ readme = "README.md"
10
+ license = { text = "BUSL-1.1" }
11
+ requires-python = ">=3.11"
12
+ authors = [{ name = "Runkite contributors" }]
13
+ keywords = ["runkite", "agent-protocol", "langgraph", "grpc", "agents"]
14
+ classifiers = [
15
+ "Development Status :: 4 - Beta",
16
+ "License :: Other/Proprietary License",
17
+ "Programming Language :: Python :: 3",
18
+ "Programming Language :: Python :: 3 :: Only",
19
+ "Topic :: Software Development :: Libraries",
20
+ ]
21
+ dependencies = [
22
+ "grpcio",
23
+ "protobuf",
24
+ "langgraph",
25
+ "langgraph-sdk",
26
+ "langchain-core",
27
+ "langgraph-checkpoint-postgres",
28
+ "psycopg[binary,pool]",
29
+ "httpx",
30
+ "uvicorn",
31
+ "opentelemetry-api",
32
+ "opentelemetry-sdk",
33
+ "opentelemetry-exporter-otlp",
34
+ ]
35
+
36
+ [project.urls]
37
+ Homepage = "https://getrunkite.github.io/runkite/"
38
+ Repository = "https://github.com/getrunkite/runkite"
39
+ Documentation = "https://github.com/getrunkite/runkite/tree/main/docs"
40
+ Changelog = "https://github.com/getrunkite/runkite/blob/main/CHANGELOG.md"
41
+ Issues = "https://github.com/getrunkite/runkite/issues"
42
+
43
+ [project.scripts]
44
+ runkite-runner = "runkite_runner.worker:main"
45
+
46
+ [tool.setuptools.packages.find]
47
+ where = ["."]
48
+ include = ["runkite_runner*"]
49
+
50
+ [tool.setuptools.package-data]
51
+ runkite_runner = ["py.typed"]
@@ -0,0 +1,3 @@
1
+ """Runkite Python Runner SDK - Execute LangGraph agents against the Runkite control plane."""
2
+
3
+ __version__ = "0.1.0"
@@ -0,0 +1,5 @@
1
+ """Allow running as: python -m runkite_runner --config langgraph.json"""
2
+
3
+ from .worker import main
4
+
5
+ main()
@@ -0,0 +1,133 @@
1
+ """Agent-to-Agent (A2A) delegation client: agent calls agent via the same
2
+ Agent Protocol API.
3
+
4
+ `call_agent` is what a running agent's own node code calls to invoke
5
+ another agent as a sub-task -- it POSTs to the control plane's
6
+ `/internal/a2a/runs` endpoint (see internal/api/a2a.go for the full
7
+ server-side design: auth propagation, recursion limits, cost
8
+ attribution via root_run_id).
9
+
10
+ Deliberately takes the graph node's own `config` dict, not separate
11
+ run_id/user parameters -- every value this needs (the current run_id to
12
+ set as parent_run_id, the authenticated user to forward as
13
+ on_behalf_of) is already there, set by `build_run_config` in worker.py
14
+ for every run. A node function just calls
15
+ ``await call_agent(config, "other_agent", {"messages": [...]})``.
16
+
17
+ Operational note: with ``wait=True`` (the default), the parent run's
18
+ worker slot stays occupied until the child finishes. A runner process
19
+ with ``concurrency=1`` therefore deadlocks on nested A2A -- the child
20
+ job cannot be dequeued until the parent frees its slot. Use
21
+ concurrency >= 2 (or ``wait=False`` + poll) for any graph that
22
+ delegates synchronously.
23
+ """
24
+
25
+ from __future__ import annotations
26
+
27
+ import os
28
+ from typing import Any
29
+
30
+ import httpx
31
+
32
+ from .tls_utils import httpx_tls_kwargs
33
+
34
+
35
+ class A2AError(Exception):
36
+ """Raised when the control plane rejects a delegation call --
37
+ e.g. HTTP 400 for a recursion depth limit exceeded (see
38
+ ErrA2ADepthExceeded in internal/api/server.go), or 404 for an
39
+ unknown parent_run_id."""
40
+
41
+
42
+ async def call_agent(
43
+ config: dict,
44
+ agent_id: str,
45
+ input: dict[str, Any],
46
+ *,
47
+ wait: bool = True,
48
+ thread_id: str | None = None,
49
+ run_config: dict[str, Any] | None = None,
50
+ control_plane_url: str | None = None,
51
+ timeout: float | None = None,
52
+ ) -> dict[str, Any]:
53
+ """Invoke another agent as a sub-task from within a running agent's
54
+ own node code.
55
+
56
+ Args:
57
+ config: the RunnableConfig LangGraph passes to every node --
58
+ must be the same `config` the calling node itself received
59
+ (or its unmodified `configurable` sub-dict), so run_id/user
60
+ can be forwarded correctly.
61
+ agent_id: which agent to delegate to.
62
+ input: the sub-agent's input, same shape as a normal run's input.
63
+ wait: block until the sub-agent's run reaches a terminal status
64
+ and return its result (default). False fires the sub-run
65
+ and returns immediately with just the created run's data --
66
+ the caller is then responsible for polling/streaming it
67
+ itself, same as any other async run.
68
+ thread_id: an existing thread to run on, e.g. to continue a
69
+ specific sub-conversation. Defaults to a fresh thread per
70
+ call -- most delegation calls are one-shot sub-tasks, not
71
+ multi-turn conversations with the sub-agent.
72
+ run_config: passed through as the sub-run's own `config`
73
+ (distinct from this function's own `config` argument, which
74
+ is the CALLER's LangGraph RunnableConfig, not the sub-run's).
75
+ control_plane_url: defaults to RUNKITE_HTTP_URL, same env var
76
+ convention as every other control-plane HTTP call in this
77
+ runner.
78
+ timeout: httpx request timeout in seconds. None (default) means
79
+ no timeout -- a `wait=True` call blocks for however long the
80
+ sub-agent actually takes, which the caller controls via its
81
+ own input, not an arbitrary client-side cutoff.
82
+
83
+ Returns:
84
+ When wait=True: {"run": {...}, "values": {...}} -- the sub-run's
85
+ final state and output, same shape as the client-facing
86
+ `/runs/{id}/wait` response.
87
+ When wait=False: just the created run object, status "pending".
88
+
89
+ Raises:
90
+ A2AError: the control plane rejected the call (e.g. recursion
91
+ depth exceeded, unknown parent_run_id).
92
+ RuntimeError: config has no run_id -- this wasn't called from
93
+ within an actual graph node's execution (or with a config
94
+ that build_run_config never touched).
95
+ """
96
+ configurable = config.get("configurable", {}) if config else {}
97
+ parent_run_id = configurable.get("run_id")
98
+ if not parent_run_id:
99
+ raise RuntimeError(
100
+ "call_agent: config has no configurable.run_id -- must be called "
101
+ "with the RunnableConfig a graph node itself received, not a "
102
+ "hand-built or empty one"
103
+ )
104
+
105
+ body: dict[str, Any] = {
106
+ "agent_id": agent_id,
107
+ "input": input,
108
+ "parent_run_id": parent_run_id,
109
+ "wait": wait,
110
+ }
111
+ if thread_id:
112
+ body["thread_id"] = thread_id
113
+ if run_config:
114
+ body["config"] = run_config
115
+
116
+ user = configurable.get("langgraph_auth_user")
117
+ if user is not None and hasattr(user, "to_dict"):
118
+ on_behalf_of = user.to_dict()
119
+ if on_behalf_of:
120
+ body["on_behalf_of"] = on_behalf_of
121
+
122
+ base_url = control_plane_url or os.environ.get("RUNKITE_HTTP_URL", "http://localhost:2026")
123
+ headers: dict[str, str] = {}
124
+ runner_token = os.environ.get("RUNNER_TOKEN")
125
+ if runner_token:
126
+ headers["X-Runner-Kind"] = "python-langgraph"
127
+ headers["X-Runner-Token"] = runner_token
128
+
129
+ async with httpx.AsyncClient(timeout=timeout, **httpx_tls_kwargs()) as client:
130
+ resp = await client.post(f"{base_url}/internal/a2a/runs", json=body, headers=headers)
131
+ if resp.status_code >= 400:
132
+ raise A2AError(f"call_agent({agent_id!r}) failed: {resp.status_code} {resp.text}")
133
+ return resp.json()
@@ -0,0 +1,214 @@
1
+ """Checkpoint dual mode.
2
+
3
+ Direct mode (default when POSTGRES_DSN is set -- production, shared DB with
4
+ the control plane): the runner holds its own connection and writes
5
+ checkpoints with LangGraph's native AsyncPostgresSaver. Zero added latency,
6
+ survives runner restarts. Correct only when the control plane also uses
7
+ POSTGRES_DSN against the same database; MySQL/Mongo/SQLite control planes
8
+ must unset POSTGRES_DSN on the runner (see README Checkpoint dual mode).
9
+
10
+ Local mode (no POSTGRES_DSN -- zero-dependency dev default): falls back to
11
+ LangGraph's in-memory MemorySaver. This is honestly ephemeral -- state does
12
+ NOT survive a runner restart. That's an accepted trade-off for the
13
+ zero-dependency default (same spirit as the control plane's own in-process
14
+ transport), not a hidden gap. Proxy mode (opaque-blob checkpoints via the
15
+ control plane's HTTP API, for non-Python runners or runners without DB
16
+ credentials) is not implemented by this Python runner -- it always has
17
+ direct DB access when Postgres is available, so proxy mode has no benefit
18
+ here; it exists in the protocol for other-language runners.
19
+
20
+ Connection pooling (runner-side concurrency): a runner process can now
21
+ process multiple jobs at once (see worker.py's --concurrency), so `start`
22
+ takes a `pool_size` and builds an `AsyncConnectionPool` instead of the
23
+ single connection `AsyncPostgresSaver.from_conn_string` opens -- otherwise
24
+ every concurrent job's checkpoint I/O would serialize on that one
25
+ connection's internal lock (correct, but not actually parallel).
26
+ `AsyncPostgresSaver.__init__`'s `conn` parameter accepts a pool directly
27
+ (`Conn = AsyncConnection | AsyncConnectionPool` in langgraph's own
28
+ `checkpoint/postgres/_ainternal.py`) -- confirmed it checks out a
29
+ connection per operation via that same module's `get_connection` helper,
30
+ so this is a supported usage, not a hack.
31
+
32
+ Concurrent-startup migration race (found live): `AsyncPostgresSaver.setup()` runs `CREATE TABLE IF NOT EXISTS` DDL
33
+ for its own `checkpoint_migrations` table, which is not actually race-free
34
+ on a truly fresh database -- the same class of bug this project's own
35
+ `internal/state/postgres/postgres.go` had and fixed with a session advisory
36
+ lock. When 2+ runner replicas start simultaneously against a fresh
37
+ Postgres, one `setup()` call can lose that race and crash with
38
+ `duplicate key value violates unique constraint "pg_type_typname_nsp_index"`.
39
+
40
+ Fixed here with an advisory lock too, but `pg_try_advisory_lock` polled in
41
+ a loop, NOT a blocking `pg_advisory_lock` -- tried the blocking version
42
+ first and hit a real deadlock, not just a slower path: `setup()` also runs
43
+ `CREATE INDEX CONCURRENTLY`, which must wait for every *other* backend's
44
+ in-flight statement in the whole database to finish before it can proceed
45
+ (a Postgres-wide barrier, unrelated to which table the other statement
46
+ touches). A losing replica blocked inside a single `SELECT
47
+ pg_advisory_lock(...)` call counts as an in-flight statement -- so the
48
+ winner's `CREATE INDEX CONCURRENTLY` waits on the losers, and the losers
49
+ wait on the winner to finish `setup()` and unlock. Circular wait,
50
+ confirmed live via `pg_stat_activity`/`pg_locks` (losers stuck on
51
+ `Lock/advisory`, winner stuck on `Lock/virtualxid` waiting on them).
52
+ Polling `pg_try_advisory_lock` avoids this because each poll is its own
53
+ complete, instantly-committed statement -- a losing replica is never
54
+ mid-statement between polls, so it never blocks the winner's
55
+ `CREATE INDEX CONCURRENTLY`.
56
+ """
57
+
58
+ import asyncio
59
+ import logging
60
+
61
+ logger = logging.getLogger("runkite.runner")
62
+
63
+ # Distinct from the Go control plane's own schema-init advisory lock key
64
+ # (894127001, internal/state/postgres/postgres.go) -- unrelated schemas
65
+ # (runner checkpoint tables vs. control-plane state tables), no reason for
66
+ # one to block the other.
67
+ _CHECKPOINT_SETUP_ADVISORY_LOCK_KEY = 894127002
68
+ _LOCK_POLL_INTERVAL_S = 0.2
69
+ _LOCK_POLL_TIMEOUT_S = 60.0
70
+
71
+
72
+ class CheckpointerManager:
73
+ """Owns the runner's single shared checkpointer for its whole lifetime.
74
+
75
+ One checkpointer instance is created at worker startup and attached to
76
+ every loaded graph (overriding whatever checkpointer, if any, the
77
+ graph module itself compiled with) -- so checkpoint mode is a runner
78
+ concern, not something agent authors need to configure in their own
79
+ graph.py. Agent authors get this for free: zero changes to graph code.
80
+ """
81
+
82
+ def __init__(self):
83
+ self._checkpointer = None
84
+ self._pool = None # AsyncConnectionPool, kept open for the runner's lifetime (postgres mode only)
85
+ self._dsn: str | None = None
86
+ self._pool_size: int = 4
87
+ self._attached: list = [] # graphs whose .checkpointer we own
88
+ self.mode = "none"
89
+
90
+ async def start(self, postgres_dsn: str | None, pool_size: int = 4):
91
+ if postgres_dsn:
92
+ import psycopg
93
+
94
+ self._dsn = postgres_dsn
95
+ self._pool_size = pool_size
96
+ await self._open_pool()
97
+
98
+ # Serialize setup() across concurrently-starting runner replicas so
99
+ # its CREATE TABLE IF NOT EXISTS / CREATE INDEX CONCURRENTLY DDL
100
+ # can't race on a fresh DB (see module docstring). Deliberately a
101
+ # standalone connection, NOT checked out from self._pool: setup()
102
+ # itself checks out a connection from that same pool internally,
103
+ # and with --concurrency 1 (pool max_size=1) holding the lock from
104
+ # inside the pool would starve setup()'s own checkout of the only
105
+ # connection available.
106
+ async with await psycopg.AsyncConnection.connect(postgres_dsn, autocommit=True) as lock_conn:
107
+ waited = 0.0
108
+ while True:
109
+ row = await (
110
+ await lock_conn.execute(
111
+ "SELECT pg_try_advisory_lock(%s)", (_CHECKPOINT_SETUP_ADVISORY_LOCK_KEY,)
112
+ )
113
+ ).fetchone()
114
+ if row and row[0]:
115
+ break
116
+ if waited >= _LOCK_POLL_TIMEOUT_S:
117
+ raise TimeoutError(
118
+ "timed out waiting for checkpoint setup() advisory lock "
119
+ f"(held by another runner replica for over {_LOCK_POLL_TIMEOUT_S}s)"
120
+ )
121
+ await asyncio.sleep(_LOCK_POLL_INTERVAL_S)
122
+ waited += _LOCK_POLL_INTERVAL_S
123
+ try:
124
+ await self._checkpointer.setup()
125
+ finally:
126
+ await lock_conn.execute("SELECT pg_advisory_unlock(%s)", (_CHECKPOINT_SETUP_ADVISORY_LOCK_KEY,))
127
+ self.mode = "direct-postgres"
128
+ logger.info(
129
+ "checkpoint mode: direct (postgres, pool_size=%s) -- LangGraph tables on "
130
+ "POSTGRES_DSN; requires the control plane to use the same Postgres database "
131
+ "(Supported profile). If the control plane is MySQL/Mongo/SQLite, unset "
132
+ "POSTGRES_DSN on this runner and set RUNKITE_HTTP_URL for store proxy mode.",
133
+ pool_size,
134
+ )
135
+ else:
136
+ from langgraph.checkpoint.memory import MemorySaver
137
+
138
+ self._checkpointer = MemorySaver()
139
+ self.mode = "memory"
140
+ logger.warning(
141
+ "checkpoint mode: in-memory (no POSTGRES_DSN set) -- "
142
+ "thread state will NOT survive a runner restart. "
143
+ "Set POSTGRES_DSN for production persistence."
144
+ )
145
+
146
+ async def _open_pool(self) -> None:
147
+ from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
148
+ from psycopg.rows import dict_row
149
+
150
+ from . import pg_pool
151
+
152
+ # Same connection kwargs from_conn_string used
153
+ # (autocommit/prepare_threshold/row_factory) -- AsyncPostgresSaver
154
+ # expects dict-row results, not psycopg's tuple default.
155
+ self._pool = pg_pool.make(
156
+ self._dsn,
157
+ max_size=self._pool_size,
158
+ conn_kwargs={"prepare_threshold": 0, "row_factory": dict_row},
159
+ )
160
+ await self._pool.open()
161
+ self._checkpointer = AsyncPostgresSaver(conn=self._pool)
162
+ for graph in self._attached:
163
+ graph.checkpointer = self._checkpointer
164
+
165
+ async def recreate_pool(self) -> None:
166
+ """Drop a wedged pool (idle overnight / laptop sleep) and open a new one.
167
+
168
+ Rebinds every previously attach()'d graph so LangGraph does not keep
169
+ calling through the closed pool.
170
+ """
171
+ if self._dsn is None:
172
+ return
173
+ old = self._pool
174
+ self._pool = None
175
+ if old is not None:
176
+ try:
177
+ await old.close()
178
+ except Exception:
179
+ logger.exception("error closing wedged checkpoint pool")
180
+ await self._open_pool()
181
+
182
+ async def recover_if_wedged(self) -> None:
183
+ """Cheap pre-job probe: if getconn hangs/fails, recreate once.
184
+
185
+ AsyncPostgresSaver owns the pool reference internally, so unlike
186
+ store.abatch we cannot catch PoolTimeout inside each checkpoint
187
+ op -- probe here before astream so an overnight-wedged pool does
188
+ not burn a full run on a 30s timeout mid-graph.
189
+ """
190
+ if self._pool is None:
191
+ return
192
+ from psycopg_pool import PoolTimeout
193
+
194
+ try:
195
+ async with self._pool.connection(timeout=5.0) as conn:
196
+ await conn.execute("SELECT 1")
197
+ except PoolTimeout:
198
+ logger.warning("checkpoint pool timed out on health probe; recreating pool")
199
+ await self.recreate_pool()
200
+ except Exception:
201
+ logger.warning("checkpoint pool health probe failed; recreating pool", exc_info=True)
202
+ await self.recreate_pool()
203
+
204
+ async def stop(self):
205
+ if self._pool is not None:
206
+ await self._pool.close()
207
+ self._pool = None
208
+
209
+ def attach(self, graph):
210
+ """Override a compiled graph's checkpointer with the shared one."""
211
+ if graph not in self._attached:
212
+ self._attached.append(graph)
213
+ graph.checkpointer = self._checkpointer
214
+ return graph
@@ -0,0 +1,62 @@
1
+ """In-runner mode for the Custom routes platform extension.
2
+
3
+ Loads a user-defined ASGI app (FastAPI, Starlette, or any other ASGI
4
+ framework) from langgraph.json's "custom_app" section and serves it via
5
+ uvicorn, running as a concurrent asyncio task alongside the runner's own
6
+ gRPC worker loop (see run_worker in worker.py) -- same process, same event
7
+ loop, similar to dropping a file into your project.
8
+
9
+ The control plane reverse-proxies /custom/* to wherever this ends up
10
+ listening (see internal/config.CustomRoutesEntry / cmd/serve.go's
11
+ initCustomRoutesProxy) -- from the control plane's side, in-runner mode and
12
+ a separately-run sidecar are the exact same mechanism, just different
13
+ processes hosting the target URL.
14
+
15
+ WSGI frameworks (e.g. Flask) aren't directly supported -- uvicorn only
16
+ serves ASGI. Wrap a WSGI app with an adapter (e.g. a2wsgi.WSGIMiddleware)
17
+ if that's the framework of choice; this loader doesn't care what kind of
18
+ object it gets back as long as it's ASGI-callable.
19
+ """
20
+
21
+ import importlib.util
22
+ import logging
23
+ import sys
24
+ from pathlib import Path
25
+ from typing import Any
26
+
27
+ logger = logging.getLogger("runkite.runner")
28
+
29
+
30
+ def load_asgi_app(config_dir: Path, module_ref: str) -> Any:
31
+ """Loads an ASGI app object from a "path/to/module.py:app_symbol" ref --
32
+ the same "path:symbol" convention langgraph.json's "graphs" section
33
+ already uses for agent graphs."""
34
+ file_path, export_name = module_ref.split(":", 1)
35
+ abs_path = (config_dir / file_path).resolve()
36
+
37
+ spec = importlib.util.spec_from_file_location("runkite_custom_app", str(abs_path))
38
+ if spec is None or spec.loader is None:
39
+ raise ValueError(f"Cannot load custom app module: {abs_path}")
40
+
41
+ module = importlib.util.module_from_spec(spec)
42
+ sys.modules[spec.name] = module
43
+ spec.loader.exec_module(module)
44
+
45
+ return getattr(module, export_name)
46
+
47
+
48
+ async def serve_custom_app(app: Any, host: str, port: int) -> None:
49
+ """Runs an ASGI app via uvicorn until cancelled. Meant to run as a
50
+ concurrent asyncio task (asyncio.create_task) alongside the gRPC
51
+ worker's own poll loop -- both share the process and event loop, so a
52
+ slow/blocking route handler in the user's app can, in principle, delay
53
+ the worker's own async work. That's an inherent trade-off of "in-runner,
54
+ same process" mode; use sidecar mode instead for routes that need
55
+ independent scaling or isolation.
56
+ """
57
+ import uvicorn
58
+
59
+ config = uvicorn.Config(app, host=host, port=port, log_level="info", access_log=True)
60
+ server = uvicorn.Server(config)
61
+ logger.info(f"Custom routes (in-runner mode): serving on http://{host}:{port}")
62
+ await server.serve()