[Fix] Adding all of the files in the project directory to each image.

examples/portfolio/agents/metrics_agent.py

@@ -8,6 +8,7 @@
 # downstream RiskAgent can build the portfolio covariance.
 #
 # Resource profile: cheap CPU, high fan-out — one compute() call per holding.
+import os
 import sys
 import os
 

examples/portfolio/workflow/portfolio_workflow.py

@@ -46,7 +46,7 @@ def main(
     advisor = AdvisorAgent()
 
     # Stage 0: parse the free-text request into structured holdings + window.
-    intent = intent_agent.parse(query=query).value()
+    intent = intent_agent.parse(query=query)
     # parse() returns a dict, but a Future's .value() only ever gives back the
     # raw string ventis stored in Redis -- same deserialization requirement as
     # every other dict-returning agent call below.

ventis/cli.py

@@ -226,6 +226,23 @@ def cmd_build(args):
     if not yaml_files:
         logger.warning("No agent YAML files found in %s", agents_dir)
 
+    import yaml
+
+    # Looks up a config entry's YAML and to map stubs to entrypoints.
+    yaml_by_name = {}
+    for yaml_path in yaml_files:
+        with open(yaml_path) as f:
+            name = yaml.safe_load(f).get("agent", {}).get("name")
+        if name:
+            yaml_by_name[name] = yaml_path
+
+    entrypoints_by_name = {a["name"]: a.get("entrypoint") for a in agents}
+    stub_entrypoints = {
+        f"{os.path.splitext(os.path.basename(p))[0]}.py": entrypoints_by_name[n]
+        for n, p in yaml_by_name.items()
+        if entrypoints_by_name.get(n)
+    }
+
     stub_paths = []
     for yaml_path in yaml_files:
         base_name = os.path.splitext(os.path.basename(yaml_path))[0]
@@ -288,6 +305,8 @@ def cmd_build(args):
                 output_dir=docker_context,
                 grpc_stubs_dir=grpc_stubs_dir,
                 api_port=agent_cfg.get("api_port", 8080),
+                project_dir=project_dir,
+                stub_entrypoints=stub_entrypoints,
                 requirements=_normalize_requirements(agent_cfg),
             )
 
@@ -306,16 +325,7 @@ def cmd_build(args):
                 continue
 
             # Find matching YAML by agent name
-            matching_yaml = None
-            for yaml_path in yaml_files:
-                import yaml
-
-                with open(yaml_path) as f:
-                    ydata = yaml.safe_load(f)
-                if ydata.get("agent", {}).get("name") == agent_name:
-                    matching_yaml = yaml_path
-                    break
-
+            matching_yaml = yaml_by_name.get(agent_name)
             if not matching_yaml:
                 logger.warning(
                     "No YAML definition found for agent '%s', skipping Docker",
@@ -331,6 +341,8 @@ def cmd_build(args):
                 output_dir=docker_context,
                 grpc_stubs_dir=grpc_stubs_dir,
                 stub_files=stub_paths,
+                project_dir=project_dir,
+                stub_entrypoints=stub_entrypoints,
                 requirements=_normalize_requirements(agent_cfg),
             )
 

ventis/stub_generator.py

@@ -271,12 +271,66 @@ def _format_source(source):
     return "\n".join(formatted) + "\n"
 
 
+# Directories ventis build itself generates inside a project -- never swept.
+_GENERATED_DIRS = {"docker_container", "stubs", "grpc_stubs"}
+
+
+def _sweep_py_files(project_dir):
+    """Recursively collect (abs_src, rel_dst) for every .py file under project_dir, preserving its directory structure."""
+    swept = []
+    for root, dirs, files in os.walk(project_dir):
+        dirs[:] = [
+            d
+            for d in dirs
+            if not d.startswith(".")
+            and not (root == project_dir and d in _GENERATED_DIRS)
+        ]
+        for fname in files:
+            abs_src = os.path.join(root, fname)
+            if fname.endswith(".py") and not os.path.islink(abs_src):
+                rel_dst = os.path.relpath(abs_src, project_dir)
+                swept.append((abs_src, rel_dst))
+    return swept
+
+
+def _stub_destination(stub_file, stub_entrypoints):
+    """Where to copy a stub so it overwrites the real file it replaces, falling back to flat if that's unsafe."""
+    basename = os.path.basename(stub_file)
+    entrypoint = stub_entrypoints.get(basename)
+    if entrypoint:
+        normalized = entrypoint.replace("\\", "/")
+        if not normalized.startswith("/") and ".." not in normalized.split("/"):
+            return normalized
+        print(f"  Warning: unsafe entrypoint '{entrypoint}' for stub {basename}, placing flat instead")
+    elif stub_entrypoints:
+        print(f"  Warning: no entrypoint mapping for stub {basename}, placing flat instead")
+    return basename
+
+
+def _copy_files(output_dir, files_to_copy):
+    """Copy each (src, dst) pair into output_dir, refusing to write outside it (e.g. via a symlinked destination parent)."""
+    real_output_dir = os.path.realpath(output_dir)
+    for src, dst in files_to_copy:
+        if not os.path.isfile(src):
+            print(f"  Warning: source file not found, skipping: {src}")
+            continue
+        dest_path = os.path.join(output_dir, dst)
+        real_dest = os.path.realpath(dest_path)
+        if os.path.commonpath([real_output_dir, real_dest]) != real_output_dir:
+            print(f"  Warning: destination escapes build context, skipping: {dst}")
+            continue
+        os.makedirs(os.path.dirname(dest_path), exist_ok=True)
+        shutil.copy2(src, dest_path)
+
+
 def generate_docker(
     yaml_path,
     agent_file,
     output_dir=None,
     grpc_stubs_dir=None,
     stub_files=None,
+    project_dir=None,
+    stub_entrypoints=None,
     requirements=None,
 ):
     """
@@ -286,11 +340,13 @@ def generate_docker(
     source files needed to run the agent with its own local controller.
 
     Args:
-        yaml_path:      Path to the YAML agent definition.
-        agent_file:     Path to the original Python agent implementation.
-        output_dir:     Optional output directory (default: docker_container/<AgentName>/).
-        grpc_stubs_dir: Optional path to compiled gRPC stubs (default: <repo_root>/grpc_stubs).
-        stub_files:     Optional list of agent stub files to copy into the context.
+        yaml_path:         Path to the YAML agent definition.
+        agent_file:        Path to the original Python agent implementation.
+        output_dir:        Optional output directory (default: docker_container/<AgentName>/).
+        grpc_stubs_dir:    Optional path to compiled gRPC stubs (default: <repo_root>/grpc_stubs).
+        stub_files:        Optional list of agent stub files to copy into the context.
+        project_dir:       Optional project root to sweep for extra .py helper files.
+        stub_entrypoints:  Optional {stub_basename: entrypoint} map for exact stub placement.
         requirements:   Optional list of extra pip packages this agent needs.
     """
     with open(yaml_path, "r") as f:
@@ -314,8 +370,13 @@ def generate_docker(
     with open(os.path.join(output_dir, "requirements.txt"), "w") as f:
         f.write(requirements_txt)
 
+    # Sweep the project for extra .py helper files not on the explicit list below.
+    files_to_copy = []
+    if project_dir:
+        files_to_copy += _sweep_py_files(project_dir)
+
     # Copy general agent files
-    files_to_copy = [
+    files_to_copy += [
         # (source_path, destination_filename)
         (os.path.join(script_dir, "future.py"), "future.py"),
         (os.path.join(script_dir, "ventis_context.py"), "ventis_context.py"),
@@ -336,13 +397,16 @@ def generate_docker(
         (os.path.join(script_dir, "llm", "bedrock.py"), "bedrock.py"),
     ]
 
-    # Copy provided agent stubs
+    # Copy provided agent stubs, overwriting the swept real file at the same path
     if stub_files:
         for stub_file in stub_files:
             files_to_copy.append(
-                (os.path.abspath(stub_file), os.path.basename(stub_file))
+                (
+                    os.path.abspath(stub_file),
+                    _stub_destination(stub_file, stub_entrypoints or {}),
+                )
             )
-          
+
     files_to_copy.append((os.path.abspath(agent_file), os.path.basename(agent_file)))
 
     # Copy gRPC generated stubs if they exist
@@ -351,11 +415,7 @@ def generate_docker(
             if fname.endswith(".py"):
                 files_to_copy.append((os.path.join(grpc_stubs_dir, fname), fname))
 
-    for src, dst in files_to_copy:
-        if os.path.isfile(src):
-            shutil.copy2(src, os.path.join(output_dir, dst))
-        else:
-            print(f"  Warning: source file not found, skipping: {src}")
+    _copy_files(output_dir, files_to_copy)
 
     # Copy the YAML definition too
     shutil.copy2(
@@ -396,6 +456,8 @@ def generate_workflow_docker(
     output_dir=None,
     grpc_stubs_dir=None,
     api_port=8080,
+    project_dir=None,
+    stub_entrypoints=None,
     requirements=None,
 ):
     """
@@ -406,10 +468,12 @@ def generate_workflow_docker(
     with its own local controller.
 
     Args:
-        workflow_file:  Path to the workflow Python file.
-        stub_files:     List of stub file paths to include.
-        output_dir:     Optional output directory (default: docker_container/Workflow/).
-        grpc_stubs_dir: Optional path to compiled gRPC stubs (default: <repo_root>/grpc_stubs).
+        workflow_file:     Path to the workflow Python file.
+        stub_files:        List of stub file paths to include.
+        output_dir:        Optional output directory (default: docker_container/Workflow/).
+        grpc_stubs_dir:    Optional path to compiled gRPC stubs (default: <repo_root>/grpc_stubs).
+        project_dir:       Optional project root to sweep for extra .py helper files.
+        stub_entrypoints:  Optional {stub_basename: entrypoint} map for exact stub placement.
         requirements:   Optional list of extra pip packages this workflow needs.
     """
     script_dir = os.path.dirname(os.path.abspath(__file__))
@@ -432,7 +496,10 @@ def generate_workflow_docker(
     # ---- Copy source files into the build context ------------------------
     workflow_basename = os.path.basename(workflow_file)
 
-    files_to_copy = [
+    # Sweep the project for extra .py helper files not on the explicit list below.
+    files_to_copy = _sweep_py_files(project_dir) if project_dir else []
+
+    files_to_copy += [
         (os.path.abspath(workflow_file), workflow_basename),
         (os.path.join(script_dir, "future.py"), "future.py"),
         (os.path.join(script_dir, "ventis_context.py"), "ventis_context.py"),
@@ -452,22 +519,23 @@ def generate_workflow_docker(
             for name in ("gpu_metrics.py", "session_logging.py")
         ],
     ]
-
-    # Copy stub files
+          
+    # Copy stub files, overwriting the swept real file at the same path
     for stub_file in stub_files:
-        files_to_copy.append((os.path.abspath(stub_file), os.path.basename(stub_file)))
+        files_to_copy.append(
+            (
+                os.path.abspath(stub_file),
+                _stub_destination(stub_file, stub_entrypoints or {}),
+            )
+        )
 
     # Copy gRPC generated stubs if they exist
     if os.path.isdir(grpc_stubs_dir):
         for fname in os.listdir(grpc_stubs_dir):
             if fname.endswith(".py"):
                 files_to_copy.append((os.path.join(grpc_stubs_dir, fname), fname))
 
-    for src, dst in files_to_copy:
-        if os.path.isfile(src):
-            shutil.copy2(src, os.path.join(output_dir, dst))
-        else:
-            print(f"  Warning: source file not found, skipping: {src}")
+    _copy_files(output_dir, files_to_copy)
 
     # ---- workflow_launcher.py --------------------------------------------
     launcher = f"""import threading