-
Notifications
You must be signed in to change notification settings - Fork 162
feat(shell): one persistent shell session per task so cwd and env survive - #3400
feat(shell): one persistent shell session per task so cwd and env survive #3400kovtcharov-amd wants to merge 4 commits into
Conversation
Ports the C++ toolbelt's ShellSession (cpp/include/gaia/process.h, #2810) to Python. What persists is the session state, not a resident child: each command runs a generated script that restores the session's directory and variables, sources the command from its own file, then reports the resulting pwd and environment to a side file the session absorbs. Not wired into the shell tool yet.
...and a reset A directory change or a variable set in one call is now still in effect for the next. Adds get_shell_state so the agent can read where it is rather than infer it, set_shell_variable because the shell's own export is not on the read-only whitelist, and reset_shell_session to recover a session that is wedged or in the wrong place. Command validation is unchanged and still applies per segment. Two seams needed care: a skill-granted binary still runs as argv, never as a string handed to a shell, and a cd into a directory the path policy refuses is not absorbed - otherwise persistence would be a way to read files by bare name from anywhere.
Also closes the chat agent's shell session in its teardown path, so the temp directory goes with the agent rather than waiting on the OS.
Also registers the session's bandit suppressions and drops the Popen kwargs dict that hid the call's overload from mypy.
Request changes
This makes shell commands share one session per agent, so a cd or a variable set in one call is still in effect for the next — a real improvement, carefully built, and unusually well tested for a change this close to the security boundary. Two things need fixing before merge, both caught by running it.
Setting a variable and then using it in a command doesn't work on Linux or macOS. The new set_shell_variable tool reports success, and get_shell_state shows the variable, but a later echo $VAR prints the literal name instead of the value — because every command's arguments are re-quoted before the shell sees them (which is the right call for security, and the reason ~, * and $VAR no longer expand). Only programs that read the environment directly get the value. The PR's own new test asserts the opposite and fails on Linux, and a second test in the same file asserts the quoting that makes it fail — so the two contradict each other. Decide which behaviour you want and make the tests and the tool's description agree.
The session invents environment changes nobody made. On Linux, every single command comes back reporting AWKPATH and AWKLIBPATH as session changes, on a host where nothing set them — an artifact of how the environment is captured. They are then carried forward and re-exported to every later command, and shown to the model as if the agent had changed them. One-line fix.
Also worth a look before merge: when a cd lands outside the agent's allowed directories, the session asks the path policy in its "and prompt the user if not" mode. On an interactive CLI that pops a security prompt after the command already ran, and answering "always" permanently widens what the agent may read — including for a bare cd, which walks the session to the user's home directory. Asking without the prompt is almost certainly what was meant.
Real-world evidence
An evidence run exercised the tools the way the agent's tool-call layer does — real hosts, real subprocesses — and both findings above come from it, not from reading the diff. What it showed working: cd persisting across calls, relative arguments validated against the session's directory (cat ../inside.txt after a cd build), the whitelist and operator refusals unchanged, reset_shell_session restoring the start state, a one-shot working_directory leaving the session where it was, and a 3-second timeout leaving no find process behind. The granted-CLI path was confirmed to still receive argv: x|whoami, $HOME and ; rm -rf / all arrived as data. Per-command overhead was 2–4 ms.
🔍 Technical details
Evidence, verbatim from the run:
$ run_shell_command(command='cd build') → session_cwd .../build
$ run_shell_command(command='cat notes.txt')
{ "stdout": "planted fact: violet-otter-92\n", "return_code": 0, "status": "success" }
$ run_shell_command(command='cd /etc') # host exposing only _is_path_allowed
{ "status": "success", "session_cwd": "/tmp/gaia_evidence_q26di2p2",
"warning": "Directory change to '/etc' was not applied: ..." }
$ session.run_argv(['echo', 'x|whoami', '$HOME', '; rm -rf /'])
{ "return_code": 0, "stdout": "x|whoami $HOME ; rm -rf /\n" }
pgrep -x find BEFORE: [] → timeout=3, timed_out: true → pgrep AFTER: []
Not exercised on this lane: the Windows cmd.exe batch path (_batch_quote, set-parsing, taskkill), the Agent UI rendering of the three tools, and the ShellSessionBusy/ShellSessionClosed tool responses. gaia chat --list-tools — the user-visible surface where the new tools appear — needs Lemonade and is pending the strix-halo lane. So the verdict on the Windows branch and the UI rests on static review alone.
I reproduced both findings independently against shell_session.py:
raw : 'kept\n' # session.run('echo $GAIA_X')
printenv: 'kept\n'
divergence after `echo hi`: {'AWKLIBPATH': '.../gawk', 'AWKPATH': '.:/usr/share/awk', 'GAIA_X': 'kept'}
🔴 set_shell_variable is unusable from a command on POSIX, and two tests disagree about it
_requote_for_posix (src/gaia/agents/tools/shell_tools.py:282) shlex.quotes every validated token, so echo $GAIA_TOOL_VAR reaches sh as echo '$GAIA_TOOL_VAR'. That is deliberate and correct — test_the_posix_shell_gets_argv_back_not_a_string_to_expand pins exactly this (("echo $HOME", "echo '$HOME'")). But test_a_variable_set_through_the_tool_reaches_later_commands (tests/unit/test_shell_tools_session.py:75) asserts the value comes back, which can only hold on the Windows branch where the original string goes to cmd.exe. On Linux it fails deterministically:
E AssertionError: assert 'kept' in '$GAIA_TOOL_VAR\n'
1 failed, 17 passed
The variable is delivered to the child — printenv GAIA_X returns it, and so would git, make, or anything else reading its own environment. Only shell-level interpolation is suppressed. Two ways out:
- Keep the quoting (recommended). Make the test assert what actually holds — e.g.
printenv GAIA_TOOL_VAR— and say so inset_shell_variable's docstring, since the model will otherwise reach forecho $VARand conclude the tool is broken. Something like: "Later commands receive the variable in their environment. Shell-level$VARexpansion is off by design, so read it withprintenv NAMErather thanecho $NAME." - Allow expansion for session-known names only. More surface, and it re-opens the question the quoting closes; I would not.
Either way the spec/shell-tools-mixin.mdx "Persistent Session" section should say expansion is off, next to where it already explains the re-quoting.
🟡 gawk's own variables become permanent session overrides (src/gaia/agents/tools/shell_session.py:64)
The POSIX state dump reads awk's ENVIRON, and gawk injects AWKPATH / AWKLIBPATH into it. _absorb_state sees two names absent from _baseline_env, records them as overrides, and _build_script then exports them on every subsequent command — so the session drifts a little further from the parent with each call, and get_shell_state reports changes the agent never made. Reproduced on a clean session with a single echo hi.
_VOLATILE_ENV_NAMES = frozenset(
{
"_",
"PWD",
"OLDPWD",
"SHLVL",
"PS1",
"PS2",
"RANDOM",
"SECONDS",
"LINENO",
"PROMPT",
"CD",
"ERRORLEVEL",
"CMDCMDLINE",
"CMDEXTVERSION",
# gawk puts these into its own ENVIRON, so the state dump reports them
# as session changes on a host where nothing set them.
"AWKPATH",
"AWKLIBPATH",
"__GAIA_RC",
}
)
Worth a regression test alongside test_inherited_variables_are_not_reported_as_changes: assert session.environment() == {} after a plain echo, rather than the current < len(os.environ) / 4 bound, which passes with the pollution present.
🟡 The cwd guard prompts the user, and approval widens the path policy (src/gaia/agents/tools/shell_tools.py:353)
validator.is_path_allowed defaults to prompt_user=True (src/gaia/security.py:375), so _absorb_state calling the guard can reach _prompt_user_for_access — which prints a security warning and, on y/always, adds the directory to allowed_paths (persisting it for always). That fires from state absorption after the command has run, for a decision the user is being asked to make about bookkeeping.
The reachable case is a bare cd: no arguments means nothing for the per-argument path check to see, the shell moves to $HOME, and absorption then asks the user whether the agent may have their home directory. Non-interactive contexts auto-deny, so this only bites the CLI — but that is the interactive surface. A guard should decide, not negotiate:
validator = getattr(self, "path_validator", None)
if validator is not None:
# prompt_user=False: this is bookkeeping after the command already
# ran, not an access request the user asked to arbitrate.
def guard(path: str) -> bool:
return validator.is_path_allowed(path, prompt_user=False)
return guard
🟢 Nits
ShellToolsMixin's "Tools provided" list is missingset_shell_variable(src/gaia/agents/tools/shell_tools.py:317); the mdx spec lists all three.close_shell_session()is wired only intoChatAgent.__del__. Any other host composing the mixin leaves itsgaia_shell_*temp directory behind — cheap to also hook wherever the mixin's own teardown lives.
Strengths
- The security reasoning at the two seams is the best part of this PR, and it holds up under testing. Re-quoting the validated tokens so
shgets back exactly the argv it had before is the right instinct — a session needs a shell, and this pays for it without letting the shell re-read anything the checks read literally. Same for keeping the granted-CLI path onrun_argv. - Not keeping a resident child is the correct call, and the docstring explains why (a timed-out command inside a shared shell is precisely the corruption at issue) rather than just what.
- The tests run real subprocesses where a mock would have proved nothing, and
_env_key's Windows case-folding comment names a bug that would have been genuinely hard to find later. - The
# nosecentries are specific and justified, and the now-unreachableB602suppression was removed with theshell=Truecall it covered.
An agent that ran
cd buildand then a build got the wrong directory, because every shell command started a fresh subprocess from the same cold state — the working directory, exported variables and any activated virtualenv were discarded the moment the command returned. Commands now share one shell session per agent, so a directory change or a variable survives to the next call. The agent can read that state back withget_shell_stateinstead of inferring it, set a variable withset_shell_variable(the read-only whitelist has noexport), and recover a session that is wedged or in the wrong place withreset_shell_session.This is a port of the C++ toolbelt's
ShellSession(#2810), so what persists is the session state, not a resident child process — there is nothing to orphan, on any platform. A timeout now kills the command's whole process tree rather than just the shell it started.Command validation is unchanged and still applies per pipeline segment. Two seams needed care and are worth a reviewer's eye: a skill-granted binary still runs as argv, never as a string handed to a shell; and a
cdinto a directory the path policy refuses is not absorbed, because otherwise persistence would be a way to reach a forbidden directory and then read files by bare name, which the per-argument check never sees.Closes #3380.
🔍 Technical details
Needs a rebase once #3394 lands — that PR reworks
shell_tools.pyheavily (bypass mode, +192 lines). This branch is cut frommainand does not build on it. The two are orthogonal: #3394 decides what may run, this decides what state carries over. Env persistence is largely unreachable under the default read-only whitelist (noexport, nosource) —set_shell_variableis what makes it usable today, and bypass mode makes the shell's ownexportwork through the same machinery for free.POSIX execution changed. A session needs a shell —
cdandexportexist nowhere else — but POSIX commands previously ran as argv with no shell at all. Handingshthe raw string would newly let it expand what every check above reads literally:cat ~/../secretis one token to the path validator and a different file to the shell. The validated tokens are re-quoted withshlex.quoteand the pipeline rebuilt, soshgets back exactly the argv it had before. Windows keeps the original string, as today — cmd's quoting rules are not shlex's, and re-quoting a PowerShell-Commandbody breaks it.Serialisation guards #2600, not parallel tool_calls. The loop runs tool calls sequentially. The hazard is a call that outlives its timeout, leaving its worker thread inside the subprocess when the next one starts. One lock per session, taken for the whole command, with a bounded wait: a call that cannot take it returns a "session busy" error naming
reset_shell_sessionrather than blocking forever.Windows folds environment variable names and
os.environupper-cases them whilesetreports original casing. Comparing raw made every inherited variable look both changed and removed — the session would have replayed the entire environment and unset it at the same time.PATHis deliberately absorbed (it is what venv activation changes) but refused byset_shell_variable, along withPYTHONPATH,LD_*and the rest of the names that decide which binary or which code runs next.get_shell_state/set_shell_variable/reset_shell_sessionare added to theshellbundle in both profiles; the flagship's is now at the 6-member ceiling.Test plan
python -m pytest tests/unit/test_shell_session.py tests/unit/test_shell_tools_session.py -q— 32 pass. Covers cwd and an exported variable surviving to the next call, serialised access, the busy error, session reset, and a timed-out command leaving no child still running.GAIA_SHELL=/path/to/sh python -m pytest tests/unit/test_shell_session.py -q— the same suite against the POSIX script generator. On Windows this is the only way to exercise it; the test helpers key off the session's interpreter, notos.name.python -m pytest tests/unit/test_shell_guardrails.py tests/unit/test_skill_binary_grants.py tests/unit/test_shell_output_encoding.py -q— 347 pass. The guardrails are untouched; the granted-CLI argv guards moved to the new seam and still assert the binary is run directly.python -m pytest tests/unit/test_chat_tool_bundles.py -q— the drift gate, with the three new tools bundled.python util/lint.py --black --isort --bandit --security --imports --agents— pass. The three new# nosecentries are justified in.security-suppressions.json; theshell_tools.pyB602 entry is removed with theshell=Truecall it covered.tests/unit/(email agent excluded — it fails on this box for an unrelated editable-install reason): 626 failed / 488 errors, identical to this box's baseline. The failures I sampled reproduce on a cleanorigin/mainworktree.The
rag_quality + context_retention + tool_selectioncheck is red on every PR right now — #3341 tracks it failing identically onmain.