diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 02ec88f..92a49ad 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -1,3 +1,30 @@ +## v1.48.0 — 2026-10-24 + +**An inline run binds its store per context, not in the process env.** +`_run_inline_workers` set `PUPPETMASTER_STATE_DIR` in `os.environ` for the +whole run. Two inline runs on different stores in one process (Marionette +runs swarms inline in its backend) then saw each other's state dir. A worker +of one run wrote its deltas, stdout sidecars and tool offloads into the +other run's store, and the restore order could leave the wrong value. + +- `state.bind_state_dir` sets a context variable. `state.state_dir_from_env` + reads it first, then the env. Every in-process reader uses it. +- The inline role pool and the agentic parallel tool pool copy the context + into their threads. +- `stamp_worker_env` and the agentic shell env give child processes the + bound dir, so children see the same value as before. +- A subprocess worker still sets the env in its own process. +- Repro: two concurrent inline runs, held inside their workers with a + barrier. Before, one run's worker resolved the other store. Now each run + resolves its own, and `os.environ` does not change. + +**Opt-in: map items as threads of the walker.** Set +`"defaults": {"inprocess": true}` or `PUPPETMASTER_FLOW_INPROCESS=1` to walk +map items as threads of the parent walker, not as new walker processes. Run +locks, markers, stop, cut, resume and reuse do not change. A thread walk that +raises leaves its run interrupted, and the map starts it again. See +`docs/FLOWS.md`. + ## v1.47.0 — 2026-10-21 **Opt-in: one shared Codex host for a flow walk.** A Codex worker starts its diff --git a/docs/FLOWS.md b/docs/FLOWS.md index e208c6e..3bb3726 100644 --- a/docs/FLOWS.md +++ b/docs/FLOWS.md @@ -444,6 +444,19 @@ a resumed walker kills the command the dead walker left running (when the pid still names that command) and runs it again. A pid that now names another process is left alone. +### Map items in the walker process + +By default a map starts one walker process for each item, and a nested map +does the same again for its own items. Set `"defaults": {"inprocess": true}`, +or `PUPPETMASTER_FLOW_INPROCESS=1` for the walker, to walk the items as +threads of the parent walker instead. A graph value wins. + +- Run locks, markers, stop, cut, resume and reuse work as with processes. +- A thread walk that raises leaves its run interrupted, and the map starts + it again (at most 4 times), as for a walker process that exits. +- It saves one process start for each item. That is about 0.2 s on a plain + machine, but more where each process start is checked or guarded. + ### One shared Codex host A Codex worker usually starts its own `codex exec` process: the launcher, the diff --git a/puppetmaster/__init__.py b/puppetmaster/__init__.py index fcef00a..983e5a9 100644 --- a/puppetmaster/__init__.py +++ b/puppetmaster/__init__.py @@ -1,4 +1,4 @@ """Puppetmaster: distributed agent workers with stitched shared memory.""" __all__ = ["__version__"] -__version__ = "1.47.0" +__version__ = "1.48.0" diff --git a/puppetmaster/adapters/_delta_stream.py b/puppetmaster/adapters/_delta_stream.py index c481a23..568613a 100644 --- a/puppetmaster/adapters/_delta_stream.py +++ b/puppetmaster/adapters/_delta_stream.py @@ -28,16 +28,17 @@ from puppetmaster.fs_permissions import mkdir_private, open_private from puppetmaster.models import Task from puppetmaster.redaction import redact_secrets +from puppetmaster.state import state_dir_from_env _DELTA_FILE = "agentic_deltas.ndjson" def _resolve_state_dir() -> Optional[Path]: - """The active state dir from ``PUPPETMASTER_STATE_DIR`` (exported by - ``worker_runtime`` after resolving --state-dir), or ``None`` when no state - dir is in scope -- mirroring :func:`_streaming._resolve_sidecar_state_dir`. + """The active state dir (bound by an inline run, or ``PUPPETMASTER_STATE_DIR`` + exported by ``worker_runtime``), or ``None`` when no state dir is in scope -- + mirroring :func:`_streaming._resolve_sidecar_state_dir`. """ - raw = os.environ.get("PUPPETMASTER_STATE_DIR") + raw = state_dir_from_env() if not raw: return None try: diff --git a/puppetmaster/adapters/_dir_manifest.py b/puppetmaster/adapters/_dir_manifest.py index 2a3d528..bb74144 100644 --- a/puppetmaster/adapters/_dir_manifest.py +++ b/puppetmaster/adapters/_dir_manifest.py @@ -11,6 +11,8 @@ from pathlib import Path from typing import Optional +from puppetmaster.state import state_dir_from_env + # Caches and tool state that no task owns. The gate also skips bytecode. SKIP_DIRS = frozenset({ ".git", ".hg", ".svn", "node_modules", "__pycache__", ".venv", "venv", @@ -23,7 +25,7 @@ def _excluded_roots(root: Path) -> list[Path]: """The live Puppetmaster state dir changes during every run; it is not the worker's.""" - raw = os.environ.get("PUPPETMASTER_STATE_DIR") + raw = state_dir_from_env() if not raw: return [] try: diff --git a/puppetmaster/adapters/_prompts.py b/puppetmaster/adapters/_prompts.py index 00d81da..fccd138 100644 --- a/puppetmaster/adapters/_prompts.py +++ b/puppetmaster/adapters/_prompts.py @@ -609,19 +609,17 @@ def _open_store_for_task(task: Task, *, discover: bool = True): so worker subprocesses see the same store the orchestrator wrote into. Returns ``None`` on any miss or error. """ - import os - job_id = getattr(task, "job_id", None) or "" if not job_id: return None from puppetmaster.adapters._streaming import _resolve_sidecar_state_dir - from puppetmaster.state import STATE_DIR_ENV, find_state_dir_for_job, resolve_state_dir + from puppetmaster.state import find_state_dir_for_job, resolve_state_dir, state_dir_from_env from puppetmaster.store_factory import create_store state_dir = _resolve_sidecar_state_dir() if state_dir is None and discover: state_dir = find_state_dir_for_job(job_id) - if state_dir is None and os.environ.get(STATE_DIR_ENV): + if state_dir is None and state_dir_from_env(): try: state_dir = resolve_state_dir() except Exception: diff --git a/puppetmaster/adapters/_streaming.py b/puppetmaster/adapters/_streaming.py index f6c3527..6798759 100644 --- a/puppetmaster/adapters/_streaming.py +++ b/puppetmaster/adapters/_streaming.py @@ -12,6 +12,7 @@ from puppetmaster.invocation import current_attempt from puppetmaster.models import Task, new_id from puppetmaster.redaction import redact_secrets +from puppetmaster.state import state_dir_from_env from puppetmaster.worker_fence import stamp_worker_env _STDOUT_HEAD_CHARS = 1000 @@ -46,10 +47,11 @@ def _resolve_sidecar_state_dir() -> Optional[Path]: Returns ``None`` if no state dir is in scope (e.g. direct adapter unit tests). Falling back to the default state dir would write logs into a workspace-hashed path that may not own the job, so we only honor an - explicit ``PUPPETMASTER_STATE_DIR`` env var (which ``worker_runtime`` - exports after resolving its --state-dir flag). + explicit state dir: the one an inline run bound, or the + ``PUPPETMASTER_STATE_DIR`` env var (which ``worker_runtime`` exports after + resolving its --state-dir flag). """ - raw = os.environ.get("PUPPETMASTER_STATE_DIR") + raw = state_dir_from_env() if not raw: return None try: diff --git a/puppetmaster/adapters/agentic.py b/puppetmaster/adapters/agentic.py index 1937fad..8304d3f 100644 --- a/puppetmaster/adapters/agentic.py +++ b/puppetmaster/adapters/agentic.py @@ -48,6 +48,7 @@ import subprocess import tempfile import time +from contextvars import copy_context from pathlib import Path from typing import Any, Callable, Optional @@ -78,7 +79,7 @@ from puppetmaster.redaction import redact_secrets from puppetmaster.usage import selected_token_usage from puppetmaster.swarm_reasoning import DEFAULT_SWARM_REASONING_EFFORT -from puppetmaster.state import resolve_state_dir +from puppetmaster.state import export_state_dir, resolve_state_dir from puppetmaster.hashline import ( SnapshotStore, apply_patch, @@ -2260,7 +2261,8 @@ def _execute_one(call: dict) -> tuple[dict, str]: # Segmentation / barrier rules are unchanged; only pool size is bounded. worker_count = parallel_executor_max_workers(len(calls)) with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as executor: - futures = {executor.submit(_execute_one, call): idx for idx, call in enumerate(calls)} + futures = {executor.submit(copy_context().run, _execute_one, call): idx + for idx, call in enumerate(calls)} results = [None] * len(calls) for future in concurrent.futures.as_completed(futures): idx = futures[future] @@ -2731,7 +2733,7 @@ def _run_owned_shell( owner = new_id("process") env = { - **os.environ, + **export_state_dir(dict(os.environ)), "GIT_PAGER": "cat", "PAGER": "cat", "GIT_TERMINAL_PROMPT": "0", diff --git a/puppetmaster/flow.py b/puppetmaster/flow.py index 228d49d..cdc1806 100644 --- a/puppetmaster/flow.py +++ b/puppetmaster/flow.py @@ -50,7 +50,9 @@ import shutil import subprocess import sys +import threading import time +import traceback import unicodedata import uuid from dataclasses import dataclass, field @@ -137,6 +139,8 @@ def validate_graph(graph: Any, *, child: bool = False, depth: int = 0, problems.extend(_effort_problems(own_defaults, "defaults")) if "codex_host" in own_defaults and not isinstance(own_defaults["codex_host"], bool): problems.append("defaults.codex_host must be true or false") + if "inprocess" in own_defaults and not isinstance(own_defaults["inprocess"], bool): + problems.append("defaults.inprocess must be true or false") lanes = own_defaults.get("lanes") if lanes is not None: if not isinstance(lanes, dict) or not set(lanes) <= set(LANES): @@ -1558,6 +1562,42 @@ def spawn_background_walk(state_dir: Path, run_id: str, *, answer: Optional[str] return process.pid +INPROCESS_ENV = "PUPPETMASTER_FLOW_INPROCESS" + + +def inprocess_requested(graph: dict, env: Optional[dict] = None) -> bool: + """Whether map items walk as threads of the parent walker, not as new processes.""" + env = os.environ if env is None else env + defaults = graph.get("defaults") if isinstance(graph.get("defaults"), dict) else {} + if "inprocess" in defaults: + return defaults.get("inprocess") is True + return str(env.get(INPROCESS_ENV, "")).strip().lower() in ("1", "true", "on", "yes") + + +def walk_in_thread(state_dir: Path, run_id: str) -> threading.Thread: + """Walk ``run_id`` on a thread of this process. + + The run lock, markers, stop and resume work as for a walker process. Each + walker process pays its own start (interpreter, imports, guards); a thread + does not. A walk that raises leaves the run ``running`` with a free lock, + so the map sees it as interrupted and starts it again. + """ + def target() -> None: + try: + walk(state_dir, run_id) + except FlowError: + pass # another walker holds the run + except Exception: + try: + with open(run_path(state_dir, run_id).with_name("walker.log"), "a", encoding="utf-8") as log: + log.write(traceback.format_exc()) + except OSError: + pass + thread = threading.Thread(target=target, name=f"flow-walk-{run_id}", daemon=True) + thread.start() + return thread + + def _fail_unwalked(state_dir: Path, run_id: str, reason: str) -> None: """Fail a run only if no walker holds it; never race a slow-starting walker's writes.""" lock = _lock(state_dir, run_id, timeout=0) @@ -1614,6 +1654,7 @@ def __init__(self, state_dir: Path, *, backend: str = "sqlite", worker_mode: str self.backend = backend self.worker_mode = worker_mode self.popen = popen or subprocess.Popen + self.custom_spawn = spawn is not None self.spawn = spawn or (lambda state_dir, run_id: spawn_background_walk( state_dir, run_id, popen=self.popen)) self.poll_seconds = poll_seconds @@ -2041,6 +2082,7 @@ def _run_map(self, node: dict, run: FlowRun, prev: str) -> NodeOutcome: def _drive_children(self, run: FlowRun, node: dict, children: dict[str, str]) -> None: concurrency = int(node.get("concurrency", DEFAULT_MAP_CONCURRENCY)) + spawn = self.spawn if self.custom_spawn or not inprocess_requested(run.graph) else walk_in_thread started: dict[str, float] = {} spawns: dict[str, int] = {} finished: dict[str, str] = {} @@ -2073,7 +2115,7 @@ def _drive_children(self, run: FlowRun, node: dict, children: dict[str, str]) -> continue spawns[child_id] = spawns.get(child_id, 0) + 1 started[child_id] = now - self.spawn(self.state_dir, child_id) + spawn(self.state_dir, child_id) live.add(child_id) time.sleep(self.poll_seconds) diff --git a/puppetmaster/job_brief.py b/puppetmaster/job_brief.py index ecee70c..4d7c57c 100644 --- a/puppetmaster/job_brief.py +++ b/puppetmaster/job_brief.py @@ -171,12 +171,12 @@ def resolve_job_brief_for_task(task) -> str: if isinstance(inline, str) and inline.strip(): return inline.strip() + ("\n" if not inline.endswith("\n") else "") from puppetmaster.adapters._streaming import _resolve_sidecar_state_dir - from puppetmaster.state import STATE_DIR_ENV, find_state_dir_for_job, resolve_state_dir + from puppetmaster.state import find_state_dir_for_job, resolve_state_dir, state_dir_from_env state_dir = _resolve_sidecar_state_dir() if state_dir is None: state_dir = find_state_dir_for_job(job_id) - if state_dir is None and os.environ.get(STATE_DIR_ENV): + if state_dir is None and state_dir_from_env(): # Explicit env pin — resolve_state_dir short-circuits before git. try: state_dir = resolve_state_dir() diff --git a/puppetmaster/orchestrator.py b/puppetmaster/orchestrator.py index c04a786..70c5190 100644 --- a/puppetmaster/orchestrator.py +++ b/puppetmaster/orchestrator.py @@ -1,11 +1,11 @@ from __future__ import annotations -import contextlib import os import subprocess import sys import time from concurrent.futures import ThreadPoolExecutor +from contextvars import copy_context from dataclasses import dataclass, replace from datetime import datetime, timezone from pathlib import Path @@ -25,6 +25,7 @@ now_iso, ) from puppetmaster.stitcher import Stitcher +from puppetmaster.state import bind_state_dir from puppetmaster.store import SwarmStore from puppetmaster.swarm_reasoning import NO_EFFORT_ADAPTERS, apply_swarm_reasoning from puppetmaster.worker_fence import stamp_worker_env @@ -184,19 +185,6 @@ def _env_int(name: str, default: int) -> int: return default -@contextlib.contextmanager -def _temporary_env_var(name: str, value: str): - previous = os.environ.get(name) - os.environ[name] = value - try: - yield - finally: - if previous is None: - os.environ.pop(name, None) - else: - os.environ[name] = previous - - def merge_routing_payload( payload: dict, decision, extra_fields: Optional[dict] = None, *, registry=None, previous_adapter: Optional[str] = None, @@ -2818,7 +2806,7 @@ def _run_inline_workers( allowed_task_ids: Optional[set[str]] = None, ) -> None: allowed_task_ids = allowed_task_ids or {task.id for task in tasks} - with _temporary_env_var("PUPPETMASTER_STATE_DIR", str(self.store.root)): + with bind_state_dir(self.store.root): while True: record_orchestrator_heartbeat(self.store, job.id) self.store.recover_stale_tasks(job.id) @@ -2882,7 +2870,7 @@ def _run_role(role: str) -> int: pool_size = min(len(roles), _INLINE_ROLE_WORKER_POOL_CAP) with ThreadPoolExecutor(max_workers=pool_size) as executor: - futures = [executor.submit(_run_role, role) for role in roles] + futures = [executor.submit(copy_context().run, _run_role, role) for role in roles] completed = 0 errors: list[BaseException] = [] for future in futures: diff --git a/puppetmaster/state.py b/puppetmaster/state.py index b1a8c0c..2321133 100644 --- a/puppetmaster/state.py +++ b/puppetmaster/state.py @@ -2,6 +2,8 @@ import hashlib from collections import OrderedDict +from contextlib import contextmanager +from contextvars import ContextVar import threading import os import re @@ -9,13 +11,40 @@ import subprocess import sys from pathlib import Path -from typing import Optional, Union +from typing import Iterator, Optional, Union from puppetmaster.readonly import ReadUnavailable STATE_DIR_ENV = "PUPPETMASTER_STATE_DIR" +# An in-process run binds its store here, not in os.environ: two inline runs +# on different stores in one process must not see each other's state dir. +_BOUND_STATE_DIR: ContextVar[Optional[str]] = ContextVar("puppetmaster_state_dir", default=None) + + +@contextmanager +def bind_state_dir(path: Union[Path, str]) -> Iterator[None]: + """Make ``path`` the state dir of this context and of the threads it copies to.""" + token = _BOUND_STATE_DIR.set(str(path)) + try: + yield + finally: + _BOUND_STATE_DIR.reset(token) + + +def state_dir_from_env() -> Optional[str]: + """The state dir bound by an in-process run, else ``PUPPETMASTER_STATE_DIR``.""" + return _BOUND_STATE_DIR.get() or os.environ.get(STATE_DIR_ENV) or None + + +def export_state_dir(env: dict) -> dict: + """Give a child process the bound state dir. Mutates and returns ``env``.""" + bound = _BOUND_STATE_DIR.get() + if bound: + env[STATE_DIR_ENV] = bound + return env + def state_identity(path: Union[Path, str]) -> str: """Return a stable opaque identity for a resolved state directory.""" @@ -28,7 +57,7 @@ def resolve_state_dir(value: Optional[Union[Path, str]] = None, cwd: Optional[Pa root = cwd or Path.cwd() if value: return _resolve_user_path(value, root) - env_value = os.environ.get(STATE_DIR_ENV) + env_value = state_dir_from_env() if env_value: return _resolve_user_path(env_value, root) return default_state_dir(root) diff --git a/puppetmaster/worker_fence.py b/puppetmaster/worker_fence.py index 75bc235..08b89dd 100644 --- a/puppetmaster/worker_fence.py +++ b/puppetmaster/worker_fence.py @@ -11,6 +11,8 @@ import os from typing import Mapping, Optional +from puppetmaster.state import export_state_dir + WORKER_ENV = "PUPPETMASTER_WORKER" WORKER_VALUE = "1" ALLOW_NESTED_ENV = "PUPPETMASTER_ALLOW_NESTED" @@ -71,6 +73,7 @@ def stamp_worker_env( role: str = "", ) -> dict: """Mark ``env`` as a Puppetmaster worker. Mutates and returns ``env``.""" + export_state_dir(env) env[WORKER_ENV] = WORKER_VALUE if job_id: env[JOB_ID_ENV] = str(job_id) diff --git a/pyproject.toml b/pyproject.toml index c8dd9e4..86d25d1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "puppetmaster-ai" -version = "1.47.0" +version = "1.48.0" description = "Gunicorn-style agent swarms with Redis-like shared memory and artifact stitching. (Imported as `puppetmaster`; published as `puppetmaster-ai` because the bare PyPI name is held by an abandoned 2019 project, name-reassignment pending.)" readme = "README.md" requires-python = ">=3.9" diff --git a/tests/test_flow_inprocess.py b/tests/test_flow_inprocess.py new file mode 100644 index 0000000..97f1f7e --- /dev/null +++ b/tests/test_flow_inprocess.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import os +import sys + +sys.path.insert(0, os.path.dirname(__file__)) +import hermetic_env # noqa: F401 # process-wide host-env isolation + +import json +import shutil +import tempfile +import threading +import unittest +from pathlib import Path +from unittest.mock import patch + +from puppetmaster import flow + + +def nested_graph(inprocess=True): + check = f'"{sys.executable}" -c "print(\'VERDICT: PASS - {{{{item}}}} ok\')"' + region = {"entry": "check", "nodes": [{"id": "check", "kind": "shell", "command": check}, + {"id": "done", "kind": "end", "status": "pass"}], + "edges": [{"from": "check", "to": "done", "when": "ok"}]} + group = {"entry": "regions", "nodes": [{"id": "regions", "kind": "map", "items": "{{item.regions}}", + "concurrency": 2, "graph": region}, + {"id": "done", "kind": "end", "status": "pass"}], + "edges": [{"from": "regions", "to": "done", "when": "ok"}]} + return {"id": "inproc", "entry": "groups", "defaults": {"inprocess": inprocess}, + "nodes": [{"id": "groups", "kind": "map", "concurrency": 2, "graph": group, + "items": [{"id": "g0", "regions": ["r0", "r1"]}, {"id": "g1", "regions": ["r2", "r3"]}]}, + {"id": "done", "kind": "end", "status": "pass"}], + "edges": [{"from": "groups", "to": "done", "when": "ok"}]} + + +class InprocessWalkTests(unittest.TestCase): + def setUp(self): + self.folder = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, self.folder, True) + self.state = self.folder / "state" + + def run_graph(self, graph): + run = flow.new_run(self.state, graph, cwd=str(self.folder)) + return flow.walk(self.state, run.run_id, execute=flow.JobNodeExecutor(self.state, poll_seconds=0.05)) + + def test_nested_map_walks_without_a_new_process(self): + def no_process(*_a, **_k): + raise AssertionError("an in-process walk started a walker process") + threads = [] + real = flow.walk_in_thread + with patch.object(flow, "spawn_background_walk", no_process), \ + patch.object(flow, "walk_in_thread", lambda s, r: threads.append(real(s, r)) or threads[-1]): + run = self.run_graph(nested_graph()) + self.assertEqual((run.status, run.reason), ("done", "pass"), run.steps) + self.assertEqual(len(threads), 6) # 2 groups and 4 regions + for thread in threads: + thread.join(5) + self.assertFalse(thread.is_alive()) + + def test_graph_false_wins_over_the_operator(self): + self.assertFalse(flow.inprocess_requested({"defaults": {"inprocess": False}}, + env={flow.INPROCESS_ENV: "1"})) + self.assertTrue(flow.inprocess_requested({}, env={flow.INPROCESS_ENV: "1"})) + self.assertFalse(flow.inprocess_requested({}, env={})) + graph = nested_graph() + graph["defaults"]["inprocess"] = "yes" + self.assertIn("defaults.inprocess must be true or false", flow.validate_graph(graph)) + + def test_a_walk_that_raises_is_started_again(self): + calls = {"n": 0} + real_walk = flow.walk + lock = threading.Lock() + + def flaky(state_dir, run_id, **kwargs): + run = flow.load_run(state_dir, run_id) + with lock: + first = run.parent is not None and calls["n"] == 0 + if first: + calls["n"] += 1 + if first: + raise RuntimeError("walker crashed") + return real_walk(state_dir, run_id, **kwargs) + with patch.object(flow, "walk", flaky): + run = flow.new_run(self.state, nested_graph(), cwd=str(self.folder)) + run = real_walk(self.state, run.run_id, execute=flow.JobNodeExecutor(self.state, poll_seconds=0.05)) + self.assertEqual((run.status, run.reason), ("done", "pass"), run.steps) + self.assertEqual(calls["n"], 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_puppetmaster.py b/tests/test_puppetmaster.py index e6ff881..0cf53a0 100644 --- a/tests/test_puppetmaster.py +++ b/tests/test_puppetmaster.py @@ -28159,7 +28159,9 @@ def test_inline_workers_scope_state_dir_env_to_store_root(self) -> None: seen = [] def complete_task(runtime): - seen.append(os.environ.get("PUPPETMASTER_STATE_DIR")) + from puppetmaster.state import state_dir_from_env + + seen.append((state_dir_from_env(), os.environ.get("PUPPETMASTER_STATE_DIR"))) stored = store.get_task_by_id(task.id) store.update_task_status(stored, TaskStatus.COMPLETE) return 1 @@ -28173,7 +28175,8 @@ def complete_task(runtime): Orchestrator(store)._run_inline_workers(job, [task]) self.assertEqual(os.environ["PUPPETMASTER_STATE_DIR"], "parent-state") - self.assertEqual(seen, [str(store.root)]) + # The run binds its store; the process env keeps the parent value. + self.assertEqual(seen, [(str(store.root), "parent-state")]) def test_redact_secrets_scrubs_argument_supplied_keys(self) -> None: from puppetmaster.redaction import ( diff --git a/tests/test_state_dir_binding.py b/tests/test_state_dir_binding.py new file mode 100644 index 0000000..fa29bf9 --- /dev/null +++ b/tests/test_state_dir_binding.py @@ -0,0 +1,150 @@ +"""An inline run binds its store per context, never in os.environ. + +Two inline runs on different stores in one process (Marionette runs swarms +inline in its backend) must each resolve their own state dir. Before the +binding, ``_run_inline_workers`` set ``PUPPETMASTER_STATE_DIR`` process-wide, +so worker output of one run could land in the other run's store. +""" +from __future__ import annotations + +import os +import threading +import unittest +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import patch + +from puppetmaster.adapters import _delta_stream, _streaming +from puppetmaster.adapters.agentic import AgenticAdapter +from puppetmaster.models import Task +from puppetmaster.orchestrator import Orchestrator +from puppetmaster.state import STATE_DIR_ENV, bind_state_dir, resolve_state_dir, state_dir_from_env +from puppetmaster.store import SwarmStore +from puppetmaster.worker_fence import stamp_worker_env +from puppetmaster.workers import WorkerSpec + + +def _seen() -> dict: + return { + "bound": state_dir_from_env(), + "sidecar": str(_streaming._resolve_sidecar_state_dir()), + "delta": str(_delta_stream._resolve_state_dir()), + "resolved": str(resolve_state_dir()), + "env": os.environ.get(STATE_DIR_ENV), + } + + +class StateDirBindingTests(unittest.TestCase): + def setUp(self) -> None: + patcher = patch.dict(os.environ) + patcher.start() + self.addCleanup(patcher.stop) + os.environ.pop(STATE_DIR_ENV, None) + + def test_binding_is_local_to_each_thread_context(self) -> None: + inside = threading.Barrier(2, timeout=5) + seen: dict = {} + + def run(name: str, path: str) -> None: + with bind_state_dir(path): + inside.wait() + seen[name] = _seen() + inside.wait() + seen[name + "_after"] = state_dir_from_env() + + with TemporaryDirectory() as a, TemporaryDirectory() as b: + threads = [threading.Thread(target=run, args=("a", a)), threading.Thread(target=run, args=("b", b))] + for t in threads: + t.start() + for t in threads: + t.join() + for name, path in (("a", a), ("b", b)): + self.assertEqual(seen[name]["bound"], path) + self.assertEqual(seen[name]["sidecar"], path) + self.assertEqual(seen[name]["delta"], path) + self.assertEqual(Path(seen[name]["resolved"]).resolve(), Path(path).resolve()) + self.assertIsNone(seen[name]["env"]) + self.assertIsNone(seen[name + "_after"]) + + def test_binding_wins_over_the_process_env(self) -> None: + os.environ[STATE_DIR_ENV] = "/env/state" + self.assertEqual(state_dir_from_env(), "/env/state") + with bind_state_dir("/bound/state"): + self.assertEqual(state_dir_from_env(), "/bound/state") + self.assertEqual(state_dir_from_env(), "/env/state") + + def test_child_env_gets_the_bound_state_dir(self) -> None: + self.assertNotIn(STATE_DIR_ENV, stamp_worker_env({})) + self.assertEqual(stamp_worker_env({STATE_DIR_ENV: "/kept"})[STATE_DIR_ENV], "/kept") + with bind_state_dir("/bound/state"): + self.assertEqual(stamp_worker_env({STATE_DIR_ENV: "/old"})[STATE_DIR_ENV], "/bound/state") + + def test_concurrent_inline_runs_each_see_their_own_store(self) -> None: + from puppetmaster.worker_runtime import WorkerRuntime + + original = WorkerRuntime.run_until_idle + inside = threading.Barrier(2, timeout=10) + seen: dict = {} + lock = threading.Lock() + + def run_until_idle(runtime: WorkerRuntime) -> int: + first = False + with lock: + if runtime.store.root not in seen: + seen[runtime.store.root] = None + first = True + if first: + inside.wait() + observed = _seen() + inside.wait() + with lock: + seen[runtime.store.root] = observed + return original(runtime) + + with TemporaryDirectory() as tmp: + stores = [SwarmStore(Path(tmp) / name) for name in ("plain", "sealed")] + errors: list = [] + + def run(store: SwarmStore) -> None: + try: + Orchestrator(store).run( + "binding probe", specs=[WorkerSpec(role="explore", instruction="look")], + worker_mode="inline", + ) + except BaseException as exc: # surfaced below + errors.append(exc) + + with patch("puppetmaster.orchestrator.WorkerRuntime.run_until_idle", run_until_idle): + threads = [threading.Thread(target=run, args=(s,)) for s in stores] + for t in threads: + t.start() + for t in threads: + t.join() + self.assertEqual(errors, []) + for store in stores: + observed = seen[store.root] + self.assertEqual(observed["bound"], str(store.root)) + self.assertEqual(observed["sidecar"], str(store.root)) + self.assertEqual(observed["delta"], str(store.root)) + self.assertIsNone(observed["env"]) + self.assertNotIn(STATE_DIR_ENV, os.environ) + + def test_agentic_parallel_tool_calls_see_the_binding(self) -> None: + adapter = AgenticAdapter() + seen: list = [] + + def execute_tool(name, args, cwd, implement, task): + seen.append(state_dir_from_env()) + return "ok" + + adapter._execute_tool = execute_tool + task = Task(job_id="job_x", role="explore", instruction="look") + calls = [{"name": "read_file", "arguments": {"path": str(i)}} for i in range(4)] + with bind_state_dir("/bound/state"): + results = adapter._execute_tool_segment_parallel(calls, Path("."), False, task) + self.assertEqual([out for _, out in results], ["ok"] * 4) + self.assertEqual(seen, ["/bound/state"] * 4) + + +if __name__ == "__main__": + unittest.main()