297 lines
9.3 KiB
Python
297 lines
9.3 KiB
Python
"""Unit tests for :mod:`omnigent.inner.os_env` helper-env construction."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import shutil
|
|
import tracemalloc
|
|
from pathlib import Path
|
|
|
|
from omnigent.inner.os_env import _read_impl, _shell_impl, build_helper_env
|
|
from omnigent.inner.sandbox import SandboxPolicy
|
|
from omnigent.runner.identity import (
|
|
OMNIGENT_SESSION_ENV_VALUE,
|
|
OMNIGENT_SESSION_ENV_VAR,
|
|
RUNNER_TUNNEL_BINDING_TOKEN_ENV_VAR,
|
|
)
|
|
|
|
|
|
def _inactive_policy() -> SandboxPolicy:
|
|
"""A ``sandbox.type: none`` policy (user opted out of sandboxing).
|
|
|
|
:returns: An inactive :class:`SandboxPolicy` whose ``build_helper_env``
|
|
branch mirrors the parent environment.
|
|
"""
|
|
return SandboxPolicy(
|
|
backend_type="none",
|
|
active=False,
|
|
read_roots=None,
|
|
write_roots=[],
|
|
write_files=[],
|
|
allow_network=True,
|
|
)
|
|
|
|
|
|
def _active_policy() -> SandboxPolicy:
|
|
"""An active policy that drives ``build_helper_env``'s allowlist branch.
|
|
|
|
``build_helper_env`` only consults ``active`` and ``env_passthrough``;
|
|
the ``backend_type`` is never activated here, so ``"none"`` is fine.
|
|
|
|
:returns: An active :class:`SandboxPolicy`.
|
|
"""
|
|
return SandboxPolicy(
|
|
backend_type="none",
|
|
active=True,
|
|
read_roots=None,
|
|
write_roots=[],
|
|
write_files=[],
|
|
allow_network=True,
|
|
)
|
|
|
|
|
|
def test_build_helper_env_inactive_strips_binding_token() -> None:
|
|
"""``sandbox.type: none`` mirrors parent env MINUS the binding token.
|
|
|
|
Opting out of sandboxing grants the agent broad
|
|
file/network access, but it must NOT additionally leak the runner's
|
|
control-plane auth secret. Asserts ``PATH`` survives (the opt-out
|
|
still mirrors the parent env) while the token is dropped.
|
|
|
|
:returns: None.
|
|
"""
|
|
parent = {
|
|
"PATH": "/usr/bin",
|
|
RUNNER_TUNNEL_BINDING_TOKEN_ENV_VAR: "bug-binding-token-secret",
|
|
}
|
|
|
|
env = build_helper_env(parent, _inactive_policy())
|
|
|
|
assert RUNNER_TUNNEL_BINDING_TOKEN_ENV_VAR not in env
|
|
assert "bug-binding-token-secret" not in env.values()
|
|
assert env["PATH"] == "/usr/bin"
|
|
|
|
|
|
def test_build_helper_env_active_drops_binding_token() -> None:
|
|
"""The active allowlist branch never admits the binding token.
|
|
|
|
The deny-by-default allowlist excludes the token's name, so even if
|
|
it is present in the parent env it does not reach the helper.
|
|
|
|
:returns: None.
|
|
"""
|
|
parent = {
|
|
"PATH": "/usr/bin",
|
|
RUNNER_TUNNEL_BINDING_TOKEN_ENV_VAR: "bug-binding-token-secret",
|
|
}
|
|
|
|
env = build_helper_env(parent, _active_policy())
|
|
|
|
assert RUNNER_TUNNEL_BINDING_TOKEN_ENV_VAR not in env
|
|
assert "bug-binding-token-secret" not in env.values()
|
|
assert env["PATH"] == "/usr/bin" # PATH is in the default allowlist
|
|
|
|
|
|
def test_build_helper_env_active_passes_omnigent_session_marker() -> None:
|
|
"""The ``OMNIGENT`` session marker survives the active allowlist.
|
|
|
|
The marker (set once on the runner process) must reach an agent's
|
|
sandboxed shell so code running there can detect it is inside an
|
|
Omnigent session, the way ``CLAUDE_CODE`` / ``CODEX`` are visible in
|
|
their own agents' shells.
|
|
|
|
:returns: None.
|
|
"""
|
|
parent = {
|
|
"PATH": "/usr/bin",
|
|
OMNIGENT_SESSION_ENV_VAR: OMNIGENT_SESSION_ENV_VALUE,
|
|
}
|
|
|
|
env = build_helper_env(parent, _active_policy())
|
|
|
|
assert env[OMNIGENT_SESSION_ENV_VAR] == OMNIGENT_SESSION_ENV_VALUE
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _shell_impl — timeout result shape
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_shell_impl_timeout_includes_exit_code(tmp_path: Path) -> None:
|
|
"""Timed-out shell commands still return the documented result fields."""
|
|
shell_path = shutil.which("bash") or shutil.which("sh")
|
|
assert shell_path is not None
|
|
|
|
result = _shell_impl(
|
|
command="sleep 2",
|
|
timeout=1,
|
|
shell_path=shell_path,
|
|
cwd=tmp_path,
|
|
)
|
|
|
|
assert result["stdout"] == ""
|
|
assert result["stderr"] == ""
|
|
assert result["exit_code"] is None
|
|
assert result["timed_out"] is True
|
|
assert result["error"] == "Command timed out after 1 seconds"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _read_impl — binary file handling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_BINARY = b"\x89PNG\r\n\x1a\n\x00\x01\x02\xff"
|
|
|
|
|
|
def test_read_impl_binary_descriptor_for_agent(tmp_path: Path) -> None:
|
|
"""With no byte cap (agent ``sys_os_read`` path) binary is not inlined.
|
|
|
|
The base64 payload would be useless to the model and could saturate the
|
|
context window, so only a descriptor is returned.
|
|
|
|
:returns: None.
|
|
"""
|
|
f = tmp_path / "logo.png"
|
|
f.write_bytes(_BINARY)
|
|
|
|
result = _read_impl(f, offset=1, limit=2_000)
|
|
|
|
assert result["encoding"] == "base64"
|
|
assert result["content"] == ""
|
|
assert result["total_bytes"] == len(_BINARY)
|
|
# Not truncated — the payload was deliberately omitted, not cut short.
|
|
assert result["truncated"] is False
|
|
assert "note" in result
|
|
|
|
|
|
def test_read_impl_binary_inlined_within_cap(tmp_path: Path) -> None:
|
|
"""A byte cap larger than the file inlines the whole payload, untruncated.
|
|
|
|
:returns: None.
|
|
"""
|
|
f = tmp_path / "logo.png"
|
|
f.write_bytes(_BINARY)
|
|
|
|
result = _read_impl(f, offset=1, limit=2_000, max_binary_bytes=10 * 1024 * 1024)
|
|
|
|
assert result["encoding"] == "base64"
|
|
assert base64.b64decode(result["content"]) == _BINARY
|
|
assert result["total_bytes"] == len(_BINARY)
|
|
assert result["truncated"] is False
|
|
|
|
|
|
def test_read_impl_binary_truncated_at_cap(tmp_path: Path) -> None:
|
|
"""A byte cap smaller than the file truncates and flags it.
|
|
|
|
:returns: None.
|
|
"""
|
|
f = tmp_path / "logo.png"
|
|
f.write_bytes(_BINARY)
|
|
|
|
result = _read_impl(f, offset=1, limit=2_000, max_binary_bytes=4)
|
|
|
|
assert base64.b64decode(result["content"]) == _BINARY[:4]
|
|
assert result["returned_bytes"] == 4
|
|
assert result["total_bytes"] == len(_BINARY)
|
|
assert result["truncated"] is True
|
|
|
|
|
|
def _make_large_binary(path: Path, size: int) -> None:
|
|
"""Write a sparse file with a binary prefix and a logical size of *size*.
|
|
|
|
The 8 KB binary prefix forces the prefix-sniff to classify it binary; the
|
|
``truncate`` extends the (sparse) file to *size* without writing the bytes,
|
|
so the test stays cheap while exercising a large logical file.
|
|
|
|
:returns: None.
|
|
"""
|
|
with path.open("wb") as fh:
|
|
fh.write(b"\xff\xfe\x00\x01" * 2_048) # 8 KB of non-UTF-8 bytes
|
|
fh.truncate(size)
|
|
|
|
|
|
def test_read_impl_binary_descriptor_does_not_read_whole_file(tmp_path: Path) -> None:
|
|
"""The descriptor path is O(1): it stats the size, never reading content.
|
|
|
|
Regression guard for inlining the whole file (``path.read_bytes()``) just
|
|
to compute ``total_bytes`` — which would OOM on large workspace blobs.
|
|
|
|
:returns: None.
|
|
"""
|
|
size = 256 * 1024 * 1024 # 256 MB logical, only ~8 KB on disk
|
|
f = tmp_path / "big.bin"
|
|
_make_large_binary(f, size)
|
|
|
|
tracemalloc.start()
|
|
try:
|
|
result = _read_impl(f, offset=1, limit=2_000)
|
|
_, peak = tracemalloc.get_traced_memory()
|
|
finally:
|
|
tracemalloc.stop()
|
|
|
|
assert result["total_bytes"] == size
|
|
assert result["content"] == ""
|
|
# A full read would have allocated ~256 MB; bounded reads stay tiny.
|
|
assert peak < 10 * 1024 * 1024
|
|
|
|
|
|
def test_read_impl_binary_cap_reads_only_the_cap(tmp_path: Path) -> None:
|
|
"""The byte-capped path reads at most ``max_binary_bytes``, not the file.
|
|
|
|
:returns: None.
|
|
"""
|
|
size = 256 * 1024 * 1024
|
|
f = tmp_path / "big.bin"
|
|
_make_large_binary(f, size)
|
|
|
|
tracemalloc.start()
|
|
try:
|
|
result = _read_impl(f, offset=1, limit=2_000, max_binary_bytes=16)
|
|
_, peak = tracemalloc.get_traced_memory()
|
|
finally:
|
|
tracemalloc.stop()
|
|
|
|
assert result["returned_bytes"] == 16
|
|
assert result["total_bytes"] == size
|
|
assert result["truncated"] is True
|
|
assert peak < 10 * 1024 * 1024
|
|
|
|
|
|
def test_read_impl_multibyte_char_straddling_sniff_boundary_is_text(tmp_path: Path) -> None:
|
|
"""A multi-byte char split across the 8 KB sniff boundary stays text.
|
|
|
|
The incremental decoder must treat the truncated trailing sequence as
|
|
*incomplete*, not invalid — otherwise valid UTF-8 would be misread as
|
|
binary purely because of where the prefix happened to be cut.
|
|
|
|
:returns: None.
|
|
"""
|
|
# 8 KB sniff window cuts the 3-byte '€' (0xE2 0x82 0xAC) at byte 8191.
|
|
text = "a" * 8_190 + "€" + "tail\n"
|
|
f = tmp_path / "wide.txt"
|
|
f.write_text(text, encoding="utf-8")
|
|
|
|
result = _read_impl(f, offset=1, limit=2_000)
|
|
|
|
assert result["encoding"] == "utf-8"
|
|
assert result["content"] == text
|
|
|
|
|
|
def test_read_impl_nul_byte_file_classified_binary(tmp_path: Path) -> None:
|
|
"""A NUL byte marks a file binary even though ``\\x00`` is valid UTF-8.
|
|
|
|
UTF-16/NUL-laden files decode cleanly as UTF-8, so without an explicit NUL
|
|
check they'd be misread as text and line-windowed into garbage.
|
|
|
|
:returns: None.
|
|
"""
|
|
# UTF-16-LE-style ASCII: every byte is valid UTF-8, but the interleaved
|
|
# NULs make this binary.
|
|
f = tmp_path / "utf16.bin"
|
|
f.write_bytes(b"H\x00e\x00l\x00l\x00o\x00")
|
|
|
|
result = _read_impl(f, offset=1, limit=2_000)
|
|
|
|
assert result["encoding"] == "base64"
|
|
assert result["total_bytes"] == 10
|