Repository navigation
feat(core): add checkpoint namespaces - #148
kaihao-zhao wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
7 issues found.
About Unblocked
Unblocked has been set up to automatically review your team's pull requests to identify genuine bugs and issues.
📖 Documentation — Learn more in our docs.
💬 Ask questions — Mention @unblocked-local-kaihao to request a review or summary, or ask follow-up questions.
👍 Give feedback — React to comments with 👍 or 👎 to help us improve.
⚙️ Customize — Adjust settings in your preferences.
| */ | ||
| public static class Builder extends HasMetadata.Builder<Builder> { | ||
| private String threadId; | ||
| private String checkpointNamespace; |
There was a problem hiding this comment.
The new checkpointNamespace field was added to Builder, but the copy constructor Builder(RunnableConfig config) at lines 285-291 was not updated. It copies threadId, checkPointId, nextNode and streamMode, but never copies checkpointNamespace. Every RunnableConfig.builder(config) call therefore silently resets the namespace to null, which means the default namespace. That includes withStreamMode, withCheckPointId, updateMetadata, removeMetadata, AbstractCheckpointSaver.put (line 88), CompiledGraph.updateState (line 269), CompiledGraph.AsyncNodeGenerator (line 674), SubCompiledGraphNodeAction and StateSnapshot.
Effect: once a graph is invoked with .checkpointNamespace("tenant-acme"), the execution config is rebuilt right away and the namespace is lost. Checkpoints from different tenants that share a thread ID are then written to and read from the same $default slot. Tenants can read and overwrite each other's state. The tests don't catch this because they call saver.put/saver.get directly with the original config, and withCheckpointNamespace re-sets the field after copying.
Fix: add this.checkpointNamespace = config.checkpointNamespace; to the copy constructor. Also add a test that round-trips through withStreamMode/withCheckPointId and through the config returned by saver.put.
| if (Objects.equals(this.threadId, checkpointNamespace)) { | ||
| return this; | ||
| } |
There was a problem hiding this comment.
The early-return guard compares this.threadId with the requested namespace. It looks copied from a sibling method. Two things go wrong:
- If the new namespace happens to equal the thread ID, the method returns
thisunchanged. The namespace is not applied and checkpoints go to the old or default namespace. - If the namespace already equals the requested value, a new config is still built every time.
| if (Objects.equals(this.threadId, checkpointNamespace)) { | |
| return this; | |
| } | |
| if (Objects.equals(this.checkpointNamespace, checkpointNamespace)) { | |
| return this; | |
| } |
| /** | ||
| * Returns the namespace used to partition checkpoints. | ||
| * | ||
| * @return the configured namespace, or {@code $default} when no namespace is set |
There was a problem hiding this comment.
checkpointNamespace() returns ofNullable(checkpointNamespace). When no namespace is set, that is Optional.empty(), not $default. The $default fallback only happens in BaseCheckpointSaver.checkpointNamespace(config). Callers who trust this Javadoc will be misled. Change the doc to say: "the configured namespace, or an empty Optional if none is set".
| | Attribute | Type | Description | | ||
| | --------- | ---- | ----------- | | ||
| | **threadId** | `String` | A unique identifier for the execution thread/session. Essential for checkpoint-based persistence, as it groups related executions together. Allows resuming interrupted graphs or maintaining conversation history. | | ||
| | **checkpointNamespace** | `String` | Optional partition for checkpoint data. Namespace values are trimmed and must contain only letters, digits, hyphens, or underscores. | |
There was a problem hiding this comment.
The table says namespace values "are trimmed and must contain only letters, digits, hyphens, or underscores". RunnableConfig.Builder.checkpointNamespace stores the raw string, and neither saver trims or validates it. So " tenant-a" and "tenant-a" end up as different namespaces, and values containing : or / are accepted. Either add the validation or correct the documentation.
| default String checkpointKey(RunnableConfig config) { | ||
| return "%s:%s".formatted(checkpointNamespace(config), threadId(config)); | ||
| } |
There was a problem hiding this comment.
checkpointKey joins the namespace and thread ID with :. Neither value is restricted, so (namespace="a:b", thread="c") and (namespace="a", thread="b:c") both map to "a:b:c". MemorySaver would then share one checkpoint list between two tenants, which defeats the isolation this feature exists for.
Fix: use an unambiguous composite key. Options are a nested map keyed by namespace and then thread, a record key such as record Key(String ns, String thread), or escaping or rejecting the separator in namespace values.
| return config.checkpointNamespace() | ||
| .map(targetFolder::resolve) | ||
| .orElse(targetFolder); |
There was a problem hiding this comment.
BaseCheckpointSaver defines CHECKPOINT_NAMESPACE_DEFAULT = "$default", and MemorySaver treats an absent namespace and an explicit "$default" as the same key. FileSystemSaver, however, uses config.checkpointNamespace() directly: an absent namespace maps to targetFolder, while an explicit "$default" maps to targetFolder/$default. The two savers behave differently, and a user who explicitly sets the documented default namespace in FileSystemSaver will not see their existing checkpoints.
Fix: either use checkpointNamespace(config) here and treat CHECKPOINT_NAMESPACE_DEFAULT as the root folder, or document the difference.
| private Path getNamespaceFolder(RunnableConfig config) { | ||
| return config.checkpointNamespace() | ||
| .map(targetFolder::resolve) | ||
| .orElse(targetFolder); | ||
| } |
There was a problem hiding this comment.
checkpointNamespace is attacker-controlled (set freely through RunnableConfig.builder().checkpointNamespace(...) / withCheckpointNamespace(...)) and is passed unsanitized into targetFolder::resolve. Path.resolve returns the argument verbatim when it is absolute, so a namespace like /etc/... writes outside targetFolder, and a relative namespace like ../../... escapes the store as well. This getNamespaceFolder result flows into getPath/getFile, which are used by put (serialize → Files.createDirectories + Files.writeString/ObjectOutputStream), get/loadCheckpoints (deserialize → arbitrary file read), releaseCheckpoints (Files.copy/Files.delete), and deleteFile (File.delete). Net effect: arbitrary file write, read, and delete relative to (or, with absolute paths, anywhere under) the process's filesystem permissions.
The documentation added in src/site/mkdocs/core/core-library.md (R145) states namespace values "are trimmed and must contain only letters, digits, hyphens, or underscores," but no trimming or charset validation exists anywhere in RunnableConfig or the savers — the contract is unenforced. Validate the namespace before resolving it (and trim it, as documented):
| private Path getNamespaceFolder(RunnableConfig config) { | |
| return config.checkpointNamespace() | |
| .map(targetFolder::resolve) | |
| .orElse(targetFolder); | |
| } | |
| private Path getNamespaceFolder(RunnableConfig config) { | |
| return config.checkpointNamespace() | |
| .map(namespace -> { | |
| var trimmed = namespace.trim(); | |
| if (!trimmed.matches("[A-Za-z0-9_-]+")) { | |
| throw new IllegalArgumentException("Invalid checkpoint namespace: " + namespace); | |
| } | |
| return targetFolder.resolve(trimmed); | |
| }) | |
| .orElse(targetFolder); | |
| } |
Risk AssessmentFileSystemSaver resolves the caller-provided checkpoint namespace directly beneath its storage folder without enforcing the documented character restrictions; values containing path traversal components can direct checkpoint reads and writes outside that folder. This should be fixed before merging. |
There was a problem hiding this comment.
10 issues found.
About Unblocked
Unblocked has been set up to automatically review your team's pull requests to identify genuine bugs and issues.
📖 Documentation — Learn more in our docs.
💬 Ask questions — Mention @dev-unblocked to request a review or summary, or ask follow-up questions.
👍 Give feedback — React to comments with 👍 or 👎 to help us improve.
⚙️ Customize — Adjust settings in your preferences.
| */ | ||
| public static class Builder extends HasMetadata.Builder<Builder> { | ||
| private String threadId; | ||
| private String checkpointNamespace; |
There was a problem hiding this comment.
The new checkpointNamespace field was added to Builder, but the copy constructor Builder(RunnableConfig config) (RunnableConfig.java:285-291) was not updated — it copies threadId, checkPointId, nextNode, streamMode, and metadata, but not checkpointNamespace. Every config round-trip through RunnableConfig.builder(config) therefore silently loses the namespace:
AbstractCheckpointSaver.put(AbstractCheckpointSaver.java:88) returnsRunnableConfig.builder(config).checkPointId(...).build()— the returned config has no namespace.CompiledGraph'sAsyncNodeGeneratorconstructor (CompiledGraph.java:674) rebuilds the config for everyinvoke/streamcall, sothis.configloses the namespace at the start of each run and all checkpoints written during the run (addCheckpoint→saver.put) go to the default namespace instead of the configured one.updateState(CompiledGraph.java:267-274),withCheckPointId,withStreamMode,updateMetadata, andremoveMetadataare affected the same way.
This breaks the feature end-to-end and silently writes tenant data into the shared default partition (cross-tenant contamination). The new tests pass only because they call saver.put/saver.get directly with the original config and discard put's return value.
Fix: add the missing copy in the copy constructor:
Builder( RunnableConfig config ) {
super( requireNonNull(config, "config cannot be null!").metadata );
this.threadId = config.threadId;
this.checkpointNamespace = config.checkpointNamespace;
this.checkPointId = config.checkPointId;
this.nextNode = config.nextNode;
this.streamMode = config.streamMode;
}| * @return a configuration using the requested namespace | ||
| */ | ||
| public RunnableConfig withCheckpointNamespace(String checkpointNamespace) { | ||
| if (Objects.equals(this.threadId, checkpointNamespace)) { |
There was a problem hiding this comment.
The short-circuit guard compares this.threadId against the new namespace instead of this.checkpointNamespace. If the thread ID happens to equal the requested namespace (e.g. threadId("tenant-acme") then withCheckpointNamespace("tenant-acme")), the method returns this and silently ignores the requested change; conversely, when the namespace is genuinely unchanged it needlessly rebuilds. The added test uses threadId("conversation-42") with namespaces "tenant-acme"/"tenant-globex", so it cannot catch this.
| if (Objects.equals(this.threadId, checkpointNamespace)) { | |
| if (Objects.equals(this.checkpointNamespace, checkpointNamespace)) { |
| private Path getNamespaceFolder(RunnableConfig config) { | ||
| return config.checkpointNamespace() | ||
| .map(targetFolder::resolve) | ||
| .orElse(targetFolder); | ||
| } |
There was a problem hiding this comment.
getNamespaceFolder passes the user-supplied checkpointNamespace straight into targetFolder::resolve with no sanitization. A namespace such as "../../etc" (or any value containing path separators) resolves outside targetFolder, and serialize/loadCheckpoints/releaseCheckpoints/deleteFile will then read, write, back up, and delete files outside the intended checkpoint root. Since the namespace comes from RunnableConfig (caller/user input), this is a path-traversal write primitive. Validate the namespace before resolving it — e.g. reject null/empty values and anything not matching [A-Za-z0-9_-]+ (which is exactly what the documentation already promises, see the core-library.md comment) — or at minimum reject values containing /, \, or ...
| default String checkpointNamespace(RunnableConfig config) { | ||
| return config.checkpointNamespace().orElse(CHECKPOINT_NAMESPACE_DEFAULT); | ||
| } |
There was a problem hiding this comment.
The interface now advertises namespace-partitioned checkpoint storage, but only MemorySaver and FileSystemSaver consume it. Every other built-in saver still keys storage by bare thread ID: PostgresSaver.java:262, RedisSaver.java:347, DynamoDBSaver.java:284, HazelcastSaver.java:133, OracleSaver.java:230, AbstractMysqlServer.java:155, CockroachDBSaver.java:284. A user who sets checkpointNamespace with any of those savers gets silently zero isolation — two tenants sharing a thread ID read and overwrite each other's persisted state, with no error or warning. At minimum, those savers should reject a config that carries a non-default namespace (fail loudly rather than silently share), or the documentation must state plainly that namespace isolation is currently limited to MemorySaver and FileSystemSaver.
| final var checkpointKey = checkpointKey(config); | ||
| return new Tag( threadId(config), _checkpointsByThread.remove( checkpointKey ) ); |
There was a problem hiding this comment.
MemorySaver.releaseCheckpoints now removes entries by the namespace-qualified checkpointKey, but its only caller with history, VersionedMemorySaver.release (VersionedMemorySaver.java:176), stores the returned Tag into _checkpointsHistoryByThread keyed by the raw threadId. Two tenants using the same thread ID under different namespaces therefore merge their released checkpoint histories into one map, and versionsByThreadId/lastVersionByThreadId return one tenant's versions for the other. VersionedMemorySaver should key its history by checkpointKey(config) as well.
| * @return namespace-qualified thread key | ||
| */ | ||
| default String checkpointKey(RunnableConfig config) { | ||
| return "%s:%s".formatted(checkpointNamespace(config), threadId(config)); |
There was a problem hiding this comment.
"%s:%s".formatted(namespace, threadId) is not injective: namespace "a:b" with threadId "c" and namespace "a" with threadId "b:c" both produce the key "a:b:c", so two distinct tenants' checkpoints share one MemorySaver entry. Nothing validates that either component is free of ":" (see the documentation comment on core-library.md — the promised validation does not exist). Use a separator that cannot appear in either component, or validate both components when building the key.
| | Attribute | Type | Description | | ||
| | --------- | ---- | ----------- | | ||
| | **threadId** | `String` | A unique identifier for the execution thread/session. Essential for checkpoint-based persistence, as it groups related executions together. Allows resuming interrupted graphs or maintaining conversation history. | | ||
| | **checkpointNamespace** | `String` | Optional partition for checkpoint data. Namespace values are trimmed and must contain only letters, digits, hyphens, or underscores. | |
There was a problem hiding this comment.
The attribute table states "Namespace values are trimmed and must contain only letters, digits, hyphens, or underscores", but no code in this change trims or validates the namespace: RunnableConfig.Builder.checkpointNamespace (RunnableConfig.java:314-316) stores the raw string, and neither BaseCheckpointSaver.checkpointKey nor FileSystemSaver.getNamespaceFolder checks it. The documented contract and the code contradict each other — and the absent validation is what makes the FileSystemSaver path traversal and the MemorySaver key-collision issues reachable. Either implement the documented normalization (trim + [A-Za-z0-9_-]+ check, rejecting invalid values) or correct the documentation.
| /** | ||
| * Returns the namespace used to partition checkpoints. | ||
| * | ||
| * @return the configured namespace, or {@code $default} when no namespace is set |
There was a problem hiding this comment.
The @return says the method yields "$default" when no namespace is set, but the implementation returns ofNullable(checkpointNamespace) — an empty Optional. The $default fallback exists only in BaseCheckpointSaver.checkpointNamespace(RunnableConfig). A caller following this javadoc would never see "$default" from this accessor.
| * @return the configured namespace, or {@code $default} when no namespace is set | |
| * @return the configured namespace, or an empty {@code Optional} when no namespace is set |
| return config.checkpointNamespace() | ||
| .map(targetFolder::resolve) | ||
| .orElse(targetFolder); |
There was a problem hiding this comment.
getNamespaceFolder branches on the raw Optional instead of the normalizing checkpointNamespace(config) default method: an absent namespace maps to the targetFolder root, while an explicitly configured "$default" (the value CHECKPOINT_NAMESPACE_DEFAULT advertises as the default) maps to a targetFolder/$default subfolder. The same logical namespace therefore has two storage locations, and a config that explicitly sets the documented default silently loses sight of existing root-level checkpoints. MemorySaver does not have this divergence (both cases produce the "$default:threadId" key). Consider mapping CHECKPOINT_NAMESPACE_DEFAULT to the root folder, or always resolving through checkpointNamespace(config).
| @Override | ||
| protected Tag releaseCheckpoints(RunnableConfig config, LinkedList<Checkpoint> checkpoints) throws Exception { | ||
| final var currentPath = getPath(config); | ||
| final var namesapceFolder = getNamespaceFolder(config); |
There was a problem hiding this comment.
The new local variable is misspelled (namesapceFolder); it is used again at lines 157, 166, and 172. Rename it to namespaceFolder in all four places (a single-line suggestion would break compilation, so no auto-fix is attached).
Risk Assessment
|
Adds checkpoint namespace support to RunnableConfig and applies namespace isolation to the built-in memory and filesystem checkpoint savers. Includes API documentation, backward-compatible default behavior, and focused tests for tenant isolation. Test command: ./mvnw -pl langgraph4j-core test.