fix(coding-agent): pickle callable Python skill modules by reference - #1278
Open
15297839035 wants to merge 1 commit into
Open
fix(coding-agent): pickle callable Python skill modules by reference#127815297839035 wants to merge 1 commit into
15297839035 wants to merge 1 commit into
Conversation
The kernel bootstrap wraps every Python skill that exposes run() in a _PrimeAgentCallableSkillModule, a types.ModuleType subclass. dill saves modules by reference, but pickle's dispatch table is keyed on the exact type, so the subclass fell through to the generic reduce path and raised TypeError: cannot pickle '_PrimeAgentCallableSkillModule' object. Every such skill was therefore dropped from the kernel state snapshot, and so was any user variable holding a reference to one - the skill globals are re-created by the bootstrap on resume, but the user variable was silently lost. Give the wrapper a __reduce__ that restores through importlib.import_module. The reducer names a stdlib function so it pickles by reference and does not drag the wrapper class into the payload. fixes PrimeIntellect-ai#1211
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #1211.
Root cause
The kernel bootstrap wraps every Python skill that exposes
run()in_PrimeAgentCallableSkillModule, atypes.ModuleTypesubclass (packages/coding-agent/src/core/tools/ipython.ts). dill saves modules by reference, but pickle's dispatch table is keyed on the exact type, so the subclass fell through to the generic reduce path and raisedTypeError: cannot pickle '_PrimeAgentCallableSkillModule' object.This predicts the reported split exactly: the five skills listed as skipped (
attach-image,compact,edit,refine,websearch) are the ones definingrun; the four listed as saved (agent-message,agent-observe,goal,rlm-heartbeat) do not, so they stayed plain modules and pickled by reference.The part that actually loses data
The skill globals themselves are rebuilt on resume regardless —
restoreState()runs before the bootstrap specifically so the bootstrap overwrites them with fresh live handles.The real loss is that an unpicklable wrapper also poisons any user variable holding a reference to it. Reproduced in a real kernel:
toolswas reported inskippedand dropped, with no bootstrap to rebuild it.Fix
Give the wrapper a
__reduce__that restores throughimportlib.import_module. The reducer names a stdlib function, so it pickles by reference rather than dragging the wrapper class into the payload — which would reintroduce the same failure.Tests
test/kernel-state-roundtrip.test.ts: real-kernel round-trip covering both the skill module and a user variable referencing it. Written first and confirmed failing against the unfixed code, with bothdemo_skillandtoolsinskipped.test/suite/regressions/1211-callable-skill-module-pickle.test.ts: contract test on the generated bootstrap, so the guarantee also holds in the default shards wherekernel-heavyis filtered out.Verification
npm run test:kernelscope: 7/7 inkernel-state-roundtrip.test.ts.npm run checkpasses.Four failures in
kernel-goal-skill.test.tsandacp-kernel-features.test.tsare pre-existing in my environment, not from this change: the same four fail atHEADwith the patch reverted. My kernel venv'srlmlacks the host bridge.Known limitation
A variable restored before the bootstrap holds the unwrapped module, so
tools['doubler'].run(21)works but callingtools['doubler'](21)directly does not. Closing that gap means switching the bootstrap from copy-wrapping to an in-place__class__upgrade, which makes every existing reference callable and drops the__dict__copy. I verified that works but left it out — it changes module-wrapping semantics for all users and is beyond this issue. Happy to send it separately if you want it.Note
Fix pickling of callable Python skill modules in kernel state snapshots
Callable Python skill modules (those exposing
run()) were being dropped from kernel state snapshots with aTypeErrorduring pickling. A__reduce__method is injected into theModuleTypesubclass generated inipython.ts, causing the wrapper to pickle by reference viaimportlib.import_moduleinstead of triggering generic reduction. Regression tests verify both the generated bootstrap code and full kernel snapshot round-trip behavior.Macroscope summarized f5c7aa5.