Skip to content

Fix workspace temp-dir leaks and guard stateful agents without workspace #90

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dipeshbabu wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix workspace temp-dir leaks and guard stateful agents without workspace #90

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions src/nightshift/task.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,8 +56,15 @@ async def run_task(
# Workspace from agent config takes priority, then platform config.
# Empty workspace gets a minimal temp dir (cwd would be "/" under systemd).
workspace = agent.config.workspace or config.workspace
workspace_is_temp = False
if not workspace:
if agent.config.stateful:
raise RuntimeError(
"Stateful agents require a configured workspace "
"(set agent.workspace or NIGHTSHIFT_WORKSPACE)."
)
workspace = tempfile.mkdtemp(prefix="nightshift-empty-ws-")
workspace_is_temp = True

try:
# Package agent source code and manifest for VM injection
Expand Down Expand Up @@ -120,6 +127,8 @@ async def run_task(
cleanup_package(pkg_dir)
if staging_dir:
shutil.rmtree(staging_dir, ignore_errors=True)
if workspace_is_temp and workspace:
shutil.rmtree(workspace, ignore_errors=True)
await log.cleanup(run_id)


Expand Down
13 changes: 12 additions & 1 deletion src/nightshift/vm/pool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ class _PoolEntry:
pkg_dir: str
staging_dir: str
workspace_dest: str # original workspace path on host
workspace_is_temp: bool = False
stateful: bool
busy: bool = False
idle_task: asyncio.Task | None = field(default=None, repr=False)
Expand Down Expand Up @@ -91,16 +92,24 @@ async def checkout(
# 2. No idle entry but under the limit → create placeholder
if len(entries) < effective_max:
workspace = agent.config.workspace or config.workspace
workspace_is_temp = False
if not workspace:
if agent.config.stateful:
raise RuntimeError(
"Stateful agents require a configured workspace "
"(set agent.workspace or NIGHTSHIFT_WORKSPACE)."
)
# Empty workspace: create a minimal temp dir rather than
# falling back to cwd (which is "/" under systemd).
# falling back to cwd (which is '/' under systemd).
workspace = tempfile.mkdtemp(prefix="nightshift-empty-ws-")
workspace_is_temp = True
entry = _PoolEntry(
vm=None,
agent_id=agent_id,
pkg_dir="",
staging_dir="",
workspace_dest=workspace,
workspace_is_temp=workspace_is_temp,
stateful=agent.config.stateful,
busy=True,
)
Expand Down Expand Up @@ -306,6 +315,8 @@ async def _destroy_entry(self, entry: _PoolEntry) -> None:
cleanup_package(entry.pkg_dir)
if entry.staging_dir:
shutil.rmtree(entry.staging_dir, ignore_errors=True)
if entry.workspace_is_temp and entry.workspace_dest:
shutil.rmtree(entry.workspace_dest, ignore_errors=True)
logger.info("Pool entry cleanup complete for agent %s", entry.agent_id)


Expand Down