PsycheFoundation · arilotter · Feb 2, 2026 · Feb 3, 2026 · Feb 3, 2026
diff --git a/architectures/centralized/client/packages.nix b/architectures/centralized/client/packages.nix
@@ -1,6 +1,6 @@
-{ psycheLib, ... }:
+{ buildRustPackage, ... }:
 
-psycheLib.buildRustPackage {
+buildRustPackage {
   needsPython = "optional";
   needsGpu = true;
   cratePath = ./.;

diff --git a/architectures/centralized/local-testnet/packages.nix b/architectures/centralized/local-testnet/packages.nix
@@ -1,5 +1,5 @@
-{ psycheLib, ... }:
+{ buildRustPackage, ... }:
 
-psycheLib.buildRustPackage {
+buildRustPackage {
   cratePath = ./.;
 }
diff --git a/architectures/centralized/server/packages.nix b/architectures/centralized/server/packages.nix
@@ -1,5 +1,5 @@
-{ psycheLib, ... }:
+{ buildRustPackage, ... }:
 
-psycheLib.buildRustPackage {
+buildRustPackage {
   cratePath = ./.;
 }
diff --git a/architectures/decentralized/solana-client/packages.nix b/architectures/decentralized/solana-client/packages.nix
@@ -1,6 +1,6 @@
-{ psycheLib, ... }:
+{ buildRustPackage, ... }:
 
-psycheLib.buildRustPackage {
+buildRustPackage {
   needsPython = "optional";
   needsGpu = true;
   cratePath = ./.;

diff --git a/architectures/decentralized/testing/packages.nix b/architectures/decentralized/testing/packages.nix
@@ -1,5 +1,5 @@
 {
-  psycheLib,
+  buildRustPackage,
   pkgs,
   inputs,
   ...
@@ -8,7 +8,7 @@
 let
   system = pkgs.stdenv.hostPlatform.system;
 in
-psycheLib.buildRustPackage {
+buildRustPackage {
   cratePath = ./.;
   # all tests need solana CLI and just
   buildInputs.test = [

diff --git a/architectures/inference-only/inference-node/packages.nix b/architectures/inference-only/inference-node/packages.nix
@@ -1,12 +1,16 @@
-{ psycheLib, ... }:
+{ buildRustPackage, ... }:
 
-psycheLib.buildRustPackage {
+buildRustPackage {
   needsPython = true;
   needsGpu = true;
   cratePath = ./.;
+
   # vllm doesn't build on macos
   supportedSystems = [
     "x86_64-linux"
     "aarch64-linux"
   ];
+  pythonBuildInputs.main = [
+    "vllm"
+  ];
 }
diff --git a/nix/checks.nix b/nix/checks.nix
@@ -10,9 +10,19 @@
       inherit (pkgs.psycheLib)
         craneLib
         rustWorkspaceArgs
-        rustWorkspaceArgsWithPython
+        mkRustWorkspaceArgsWithPython
+        mkPythonVenv
+        psychePythonExtension
+        allPythonDeps
         cargoArtifacts
         ;
+
+      rustWorkspaceArgsWithPython = mkRustWorkspaceArgsWithPython (mkPythonVenv {
+        extraPackages = {
+          psyche = psychePythonExtension;
+        };
+        additionalNixPackages = allPythonDeps;
+      });
     in
     {
       checks =

diff --git a/nix/devShell.nix b/nix/devShell.nix
@@ -13,9 +13,18 @@
         rustWorkspaceArgs
         craneLib
         env
-        psychePythonVenv
-        psychePythonVenvWithExtension
+        mkPythonVenv
+        psychePythonExtension
+        allPythonDeps
         ;
+
+      psychePythonVenv = mkPythonVenv { };
+      psychePythonVenvWithExtension = mkPythonVenv {
+        extraPackages = {
+          psyche = psychePythonExtension;
+        };
+        additionalNixPackages = allPythonDeps;
+      };
     in
     {
       # fmt as precommit hook

diff --git a/nix/lib.nix b/nix/lib.nix
@@ -66,12 +66,25 @@ let
     cargoExtraArgs = "--features python" + lib.optionalString (pkgs.config.cudaSupport) ",parallelism";
   };
 
-  rustWorkspaceArgsWithPython = rustWorkspaceArgs // {
-    buildInputs = rustWorkspaceArgs.buildInputs ++ [
-      psychePythonVenvWithExtension
-    ];
-    NIX_LDFLAGS = "-L${psychePythonVenvWithExtension}/lib -lpython3.12";
-  };
+  psychePythonExtension = pkgs.callPackage ../python { };
+
+  mkPythonVenv =
+    {
+      extraPackages ? { }, # actual packages from pkgs.callPackage
+      additionalNixPackages ? [ ], # like "vllm", etc
+    }:
+    pkgs.callPackage ./python.nix {
+      inherit (inputs) uv2nix pyproject-nix pyproject-build-systems;
+      inherit extraPackages additionalNixPackages;
+    };
+
+  mkRustWorkspaceArgsWithPython =
+    venv:
+    rustWorkspaceArgs
+    // {
+      buildInputs = rustWorkspaceArgs.buildInputs ++ [ venv ];
+      NIX_LDFLAGS = "-L${venv}/lib -lpython3.12";
+    };
 
   rustWorkspaceArgsNoPython = rustWorkspaceDeps // {
     inherit env src;
@@ -82,28 +95,6 @@ let
 
   cargoArtifacts = craneLib.buildDepsOnly rustWorkspaceArgs;
   cargoArtifactsNoPython = craneLib.buildDepsOnly rustWorkspaceArgsNoPython;
-
-  psychePythonExtension = pkgs.callPackage ../python { };
-
-  # python venv without the psyche extension (vllm, etc)
-
-  pythonDeps = { inherit (inputs) uv2nix pyproject-nix pyproject-build-systems; };
-  psychePythonVenv = pkgs.callPackage ./python.nix (
-    {
-      extraPackages = { };
-    }
-    // pythonDeps
-  );
-
-  # python venv with the psyche extension
-  psychePythonVenvWithExtension = pkgs.callPackage ./python.nix (
-    {
-      extraPackages = {
-        psyche = psychePythonExtension;
-      };
-    }
-    // pythonDeps
-  );
   # builds a rust package
   # Returns an attrset of packages: { packageName = ...; packageName-nopython = ...; }
   # Automatically discovers and builds examples from the crate's examples/ directory
@@ -116,21 +107,23 @@ let
   #   - buildInputs.main = [ deps ] applies to src/main.rs
   #   - buildInputs.<type> = [ deps ] applies to all binaries of type (bin/test/example)
   #   - buildInputs.<type>.<name> = [ deps ] applies to specific binary
+  # pythonBuildInputs: attrset of runtime python deps for each built binary, same format as buildInputs
   buildRustPackage =
     {
       needsPython ? false,
       needsGpu ? false,
       cratePath, # path to the crate dir
       supportedSystems ? null,
       buildInputs ? { },
+      pythonBuildInputs ? { },
     }:
     let
       # type: "main" | "bin" | "test" | "example"
       # name: the binary / test name
       getRuntimeDepsForArtifact =
-        type: name:
+        inputsAttrset: inputsAttrsetName: type: name:
         let
-          typeConfig = buildInputs.${type} or null;
+          typeConfig = inputsAttrset.${type} or null;
 
           # there's only one "main"
           mainDeps =
@@ -139,7 +132,7 @@ let
                 if lib.isList typeConfig then
                   typeConfig
                 else if typeConfig != null then
-                  throw "buildInputs.main must be a list, got ${builtins.typeOf typeConfig}"
+                  throw "${inputsAttrsetName}.main must be a list, got ${builtins.typeOf typeConfig}"
                 else
                   [ ]
               )
@@ -168,7 +161,32 @@ let
           "example"
         ]) "type must be 'bin', 'test', or 'example', got: ${type}";
         let
-          workspaceArgs = if withPython then rustWorkspaceArgsWithPython else rustWorkspaceArgsNoPython;
+          # custom venv for this package with its specific python dependencies
+          pythonRuntimeDeps =
+            getRuntimeDepsForArtifact pythonBuildInputs "pythonBuildInputs" type
+              originalName;
+          customPythonVenv =
+            if withPython then
+              mkPythonVenv {
+                extraPackages = {
+                  psyche = psychePythonExtension;
+                };
+                additionalNixPackages = pythonRuntimeDeps;
+              }
+            else
+              null;
+
+          workspaceArgs =
+            if withPython then
+              (
+                rustWorkspaceArgs
+                // {
+                  buildInputs = rustWorkspaceArgs.buildInputs ++ [ customPythonVenv ];
+                  NIX_LDFLAGS = "-L${customPythonVenv}/lib -lpython3.12";
+                }
+              )
+            else
+              rustWorkspaceArgsNoPython;
           artifacts = if withPython then cargoArtifacts else cargoArtifactsNoPython;
 
           # delete conflicting bins from other crates to prevent ambiguous --bin/--example/--test
@@ -277,8 +295,8 @@ let
             }
           );
 
-          runtimeDeps = getRuntimeDepsForArtifact type originalName;
-          allRuntimeDeps = (lib.optionals withPython [ psychePythonVenvWithExtension ]) ++ runtimeDeps;
+          runtimeDeps = getRuntimeDepsForArtifact buildInputs "buildInputs" type originalName;
+          allRuntimeDeps = (lib.optionals withPython [ customPythonVenv ]) ++ runtimeDeps;
 
           wrappedRustPackage =
             pkgs.runCommand "${name}"
@@ -458,18 +476,37 @@ let
         else
           assert (builtins.all (c: c == null) typeChecks);
           x;
+      flatPythonDeps =
+        let
+          getDepsFromType =
+            typeValue:
+            if lib.isList typeValue then
+              typeValue
+            else if lib.isAttrs typeValue then
+              lib.attrValues typeValue
+            else
+              [ ];
+        in
+        lib.pipe pythonBuildInputs [
+          lib.attrValues
+          (map getDepsFromType)
+          lib.flatten
+          lib.unique
+        ];
     in
-    validateBuildInputs (
-      lib.optionalAttrs shouldBuildForThisSystem (
-        util.mergeAttrsetsNoConflicts "can't merge binary package sets" [
-          mainRsPackage
-          binDirPackages
-          examplePackages
-          testPackages
-        ]
-      )
-    );
-
+    {
+      inherit flatPythonDeps;
+      packages = validateBuildInputs (
+        lib.optionalAttrs shouldBuildForThisSystem (
+          util.mergeAttrsetsNoConflicts "can't merge binary package sets" [
+            mainRsPackage
+            binDirPackages
+            examplePackages
+            testPackages
+          ]
+        )
+      );
+    };
   # TODO: i can't set the rust build target to WASM for the build deps for wasm-pack, since *some* of them don't build.
   # really, i want like a wasm-only set of deps to build... can I do that?
   # like do the buildDepsOnly for not the workspace, but my specific package that *happens* to be in a workspace.
@@ -612,23 +649,95 @@ let
         doInstallCargoArtifacts = false;
       }
     );
+
+  workspaceCargoToml = builtins.fromTOML (builtins.readFile ../Cargo.toml);
+
+  # expand globs in workspace members from cargo.toml
+  expandWorkspaceMembers =
+    members:
+    lib.flatten (
+      lib.map (
+        memberPattern:
+        if lib.hasSuffix "/*" memberPattern then
+          let
+            dir = lib.removeSuffix "/*" memberPattern;
+            dirPath = ../${dir};
+            entries = builtins.readDir dirPath;
+            subdirs = lib.filterAttrs (n: v: v == "directory") entries;
+          in
+          lib.mapAttrsToList (name: _: "${dir}/${name}") subdirs
+        else
+          [ memberPattern ]
+      ) members
+    );
+
+  expandedMembers = expandWorkspaceMembers workspaceCargoToml.workspace.members;
+
+  # find all crates with packages.nix
+  discoverCratesWithPackagesNix =
+    members:
+    lib.filter (pkg: pkg != null) (
+      lib.map (
+        memberPath:
+        let
+          fullPath = ../${memberPath};
+          packagesNixPath = fullPath + "/packages.nix";
+          cargoTomlPath = fullPath + "/Cargo.toml";
+
+          isExcluded = builtins.elem memberPath [
+            "python/" # python venv with special dependencies
+          ];
+
+          hasCargoToml = builtins.pathExists cargoTomlPath;
+          hasPackagesNix = builtins.pathExists packagesNixPath;
+        in
+        if hasCargoToml && hasPackagesNix && !isExcluded then
+          let
+            cargoToml = builtins.fromTOML (builtins.readFile cargoTomlPath);
+            packageName = cargoToml.package.name or (baseNameOf memberPath);
+          in
+          {
+            name = packageName;
+            path = fullPath;
+          }
+        else
+          null
+      ) members
+    );
+
+  rustPackageSets = lib.map (
+    pkg: import (pkg.path + "/packages.nix") { inherit buildRustPackage pkgs inputs; }
+  ) (discoverCratesWithPackagesNix expandedMembers);
+  # a packages.nix returns an attrset of packages (including examples)
+  rustPackages = util.mergeAttrsetsNoConflicts "can't merge rust package sets." (
+    lib.map (pkg: pkg.packages) rustPackageSets
+  );
+
+  allPythonDeps = lib.pipe rustPackageSets [
+    (map (pkg: pkg.flatPythonDeps or [ ]))
+    lib.flatten
+    lib.unique
+  ];
+
 in
 {
   inherit
     rustToolchain
     craneLib
     buildSolanaIdl
     rustWorkspaceArgs
-    rustWorkspaceArgsWithPython
+    mkRustWorkspaceArgsWithPython
     cargoArtifacts
     buildRustPackage
     buildRustWasmTsPackage
     useHostGpuDrivers
     env
     src
     gitcommit
-    psychePythonVenv
-    psychePythonVenvWithExtension
+    mkPythonVenv
+    psychePythonExtension
+    allPythonDeps
+    rustPackages
     ;
 
   mkWebsitePackage = pkgs.callPackage ../website/common.nix { };