aios/aios
1
0
Fork 0
Code Issues Pull Requests 5 Actions Packages Projects Releases Wiki Activity

stage-b: secure writer module updates (tests/task_plane/workers/test_codex_worker.py)

Browse Source
This commit is contained in:
chatgpt-gitops-bot
2026-07-25 00:02:39 +08:00
parent d276a8ecb3
commit 063949afe5
1 changed files with 211 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/task_plane/workers/test_codex_worker.py
+211 -21
View File
@@ -1,13 +1,63 @@
import json
import tempfile
import threading
import unittest
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
import os
import sys
import tempfile
import unittest
from pathlib import Path
ROOT = os.getcwd()
sys.path.insert(0, str(Path(ROOT) / "active-runtime"))
from app.task_plane.workers.command_policy import CommandPolicy, CommandPolicyError, isolated_environment
from app.task_plane.workers.command_policy import CommandPolicy
from app.task_plane.workers.task_client import (
CredentialSource,
StageAClientConfig,
TaskClientAuthError,
TaskClientProtocolError,
TaskClientResponseTooLargeError,
TaskPlaneClient,
)
def _start_server(responses):
class Handler(BaseHTTPRequestHandler):
request_count = 0
seen = []
def do_POST(self):
index = Handler.request_count
Handler.request_count += 1
payload_length = int(self.headers.get("Content-Length", "0"))
payload = self.rfile.read(payload_length)
try:
payload_obj = json.loads(payload.decode("utf-8") or "{}")
except Exception:
payload_obj = {"raw": payload.decode("utf-8", errors="replace")}
Handler.seen.append((self.path, payload_obj))
status, response_obj = responses[min(index, len(responses) - 1)]
body = json.dumps(response_obj).encode("utf-8")
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def log_message(self, *_args, **_kwargs):
return
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
return server, thread, Handler
def _stop_server(server, thread):
server.shutdown()
server.server_close()
thread.join(timeout=2)
class CodexWorkerCommandPolicyTests(unittest.TestCase):
@@ -15,24 +65,164 @@ class CodexWorkerCommandPolicyTests(unittest.TestCase):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
policy = CommandPolicy(root, root / "results")
with self.assertRaises(CommandPolicyError):
with self.assertRaises(ValueError):
policy.build("git-show", oid="HEAD~1", paths=["active-runtime/app/task_plane/api.py"])
with self.assertRaises(CommandPolicyError):
with self.assertRaises(ValueError):
policy.build("pytest", oid="0" * 40, paths=["../outside"])
def test_scrubs_secret_environment_and_allows_scoped_pytest(self):
class CodexWorkerClientStageBTests(unittest.TestCase):
def test_client_retry_on_retryable_status(self):
server, thread, handler = _start_server([(503, {"error": "try again"}), (200, {"ok": True})])
try:
client = TaskPlaneClient(
StageAClientConfig(base_url=f"http://127.0.0.1:{server.server_port}"),
token="stage-b-token",
)
body = client._post("/tests/retry", {"probe": "x"}, idempotency_key="retry")
self.assertEqual(body, {"ok": True})
self.assertEqual(handler.request_count, 2)
self.assertTrue(handler.seen[0][0].startswith("/tests/retry"))
finally:
_stop_server(server, thread)
def test_client_stops_retry_on_auth_error(self):
server, thread, handler = _start_server([(403, {"error": "denied"})])
try:
client = TaskPlaneClient(
StageAClientConfig(base_url=f"http://127.0.0.1:{server.server_port}"),
token="stage-b-token",
)
with self.assertRaises(TaskClientAuthError):
client._post("/tests/auth", {"probe": "x"}, idempotency_key="auth")
self.assertEqual(handler.request_count, 1)
finally:
_stop_server(server, thread)
def test_client_rejects_oversize_request_and_response(self):
server, thread, _ = _start_server([(200, {"ok": True})])
try:
small = TaskPlaneClient(
StageAClientConfig(
base_url=f"http://127.0.0.1:{server.server_port}",
max_request_bytes=8,
),
token="stage-b-token",
)
with self.assertRaises(TaskClientResponseTooLargeError):
small._post("/tests/request", {"probe": "too-large"}, idempotency_key="req")
finally:
_stop_server(server, thread)
oversize = {"blob": "x" * 1_100_000}
server, thread, _ = _start_server([(200, oversize)])
try:
client = TaskPlaneClient(
StageAClientConfig(base_url=f"http://127.0.0.1:{server.server_port}"),
token="stage-b-token",
)
with self.assertRaises(TaskClientResponseTooLargeError):
client._post("/tests/response", {"probe": "ok"}, idempotency_key="resp")
finally:
_stop_server(server, thread)
def test_client_credential_source_symlink_and_mode(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
previous = os.environ.get("TEST_SENTINEL_TOKEN")
os.environ["TEST_SENTINEL_TOKEN"] = "not-visible"
try:
env = isolated_environment(root / "home", root / "tmp", "active-runtime")
finally:
if previous is None:
os.environ.pop("TEST_SENTINEL_TOKEN", None)
else:
os.environ["TEST_SENTINEL_TOKEN"] = previous
self.assertNotIn("TEST_SENTINEL_TOKEN", env)
policy = CommandPolicy(root, root / "results")
command = policy.build("pytest", oid="0" * 40, paths=["tests/task_plane/public_interface"])
self.assertEqual(command[:4], ["python3", "-m", "pytest", "-q"])
directory = Path(directory)
token_file = directory / "token"
token_file.write_text("value", encoding="utf-8")
token_file.chmod(0o600)
import grp
import pwd
owner = pwd.getpwuid(os.getuid()).pw_name
group = grp.getgrgid(os.getgid()).gr_name
source = CredentialSource(
token_file,
token_owner=owner,
token_group=group,
expected_mode=0o600,
)
self.assertEqual(source.read_token(), "value")
symlink = directory / "token-link"
symlink.symlink_to(token_file)
with self.assertRaises(TaskClientAuthError):
CredentialSource(
symlink,
token_owner=owner,
token_group=group,
expected_mode=0o600,
).read_token()
bad_mode = directory / "bad-mode"
bad_mode.write_text("value", encoding="utf-8")
bad_mode.chmod(0o644)
with self.assertRaises(TaskClientAuthError):
CredentialSource(
bad_mode,
token_owner=owner,
token_group=group,
expected_mode=0o600,
).read_token()
def test_client_rejects_missing_claim_id_and_reports_heartbeat_defaults(self):
client = TaskPlaneClient(StageAClientConfig(), token="stage-b-token")
with self.assertRaises(TaskClientProtocolError):
client.progress("task-id", claim_id="", progress=10, stage="RUN")
with self.assertRaises(TaskClientProtocolError):
client.result(
"task-id",
claim_id="",
status="COMPLETED",
outputs=[],
evidence=[],
started_at="",
completed_at="",
)
with self.assertRaises(TaskClientProtocolError):
client.fail("task-id", claim_id="")
def test_heartbeat_request_uses_default_lease(self):
captured = {}
def handler_factory(responses):
class Handler(BaseHTTPRequestHandler):
def do_POST(self):
payload_length = int(self.headers.get("Content-Length", "0"))
captured["path"] = self.path
captured["payload"] = json.loads(self.rfile.read(payload_length).decode("utf-8") or "{}")
status, response_obj = responses[min(responses_index[0], len(responses) - 1)]
responses_index[0] += 1
response = json.dumps(response_obj).encode("utf-8")
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(response)))
self.end_headers()
self.wfile.write(response)
def log_message(self, *_args, **_kwargs):
return
return Handler
responses_index = [0]
server = ThreadingHTTPServer(("127.0.0.1", 0), handler_factory([(200, {"claim_id": "x"}), (200, {"ok": True})]))
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
client = TaskPlaneClient(
StageAClientConfig(base_url=f"http://127.0.0.1:{server.server_port}"),
token="stage-b-token",
)
client.claim_heartbeat("task-id", claim_id="claim-id")
self.assertEqual(captured["payload"], {"claim_id": "claim-id", "lease_seconds": 120})
self.assertTrue(captured["path"].endswith("/tasks/task-id/heartbeat"))
captured["payload"] = None
client.worker_heartbeat("worker-1")
self.assertEqual(captured["payload"], {})
finally:
_stop_server(server, thread)