739 lines
32 KiB
Python
739 lines
32 KiB
Python
"""Failure and byte-boundary regressions found in the backend review."""
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import os
|
|
import sys
|
|
import threading
|
|
import time
|
|
import zlib
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from backend.app import create_app
|
|
from backend.artifacts import ArtifactStore
|
|
from backend.containers import DockerContainerManager, FakeContainerManager
|
|
from backend.dev.fake_peer import ScriptedContainerPeer
|
|
from backend.errors import BusyError, CleanupFailedError, InvalidInputError, QueryTimeoutError, UpstreamFailedError, UpstreamResponseTooLargeError
|
|
from backend.history import HistoryManager, RESERVED_METADATA_BYTES_PER_ENTRY, WarningsManager
|
|
from backend.model import FakeModelAdapter, OpenAIModelAdapter
|
|
from backend.runner import QueryRunner, StderrDrainer
|
|
from backend.settings import Settings, canonicalize_url, validate_confluence_url
|
|
from backend.transport import BridgeState, BridgeTransport, NDJSONProtocolError
|
|
from backend.upstream import read_json
|
|
from tests.backend.conftest import make_test_confluence_client_factory
|
|
|
|
|
|
class Sink:
|
|
def write(self, data):
|
|
pass
|
|
async def drain(self):
|
|
pass
|
|
|
|
|
|
def stage_file(store, query, name="file.txt", content=b"data"):
|
|
staging = store.create_staging_session(query)
|
|
assert staging.handle_begin("t_1", name, len(content))[0] == "accept"
|
|
if content:
|
|
staging.handle_chunk("t_1", 0, base64.b64encode(content).decode())
|
|
staging.handle_end("t_1", len(content), int(bool(content)))
|
|
return staging
|
|
|
|
|
|
def runner_for(tmp_path, manager=None, timeout=1, cleanup_timeout=1, factory=None):
|
|
return QueryRunner(
|
|
Settings(query_timeout_seconds=timeout, cleanup_timeout_seconds=cleanup_timeout),
|
|
manager or FakeContainerManager(), ArtifactStore(tmp_path / "artifacts"),
|
|
FakeModelAdapter(), factory or make_test_confluence_client_factory(),
|
|
)
|
|
|
|
|
|
async def run(runner, disconnect=None):
|
|
return await runner.run("research", "https://approved.example.com", "valid-pat", "owner", disconnect)
|
|
|
|
|
|
def assert_empty(runner):
|
|
assert not list(runner.artifact_store.staging_dir.iterdir())
|
|
assert not list(runner.artifact_store.committed_dir.iterdir())
|
|
assert runner.artifact_store.global_reserved_bytes == 0
|
|
|
|
|
|
def test_partial_commit_keeps_other_query_reservation(tmp_path, monkeypatch):
|
|
store = ArtifactStore(tmp_path)
|
|
retained = stage_file(store, "old").commit("other")[0]
|
|
current = stage_file(store, "new")
|
|
current.handle_begin("t_2", "second.txt", 0)
|
|
current.handle_end("t_2", 0, 0)
|
|
import backend.artifacts as module
|
|
original = module.shutil.move
|
|
def fail_second(source, destination):
|
|
if source.endswith("t_2.tmp"):
|
|
raise OSError("disk failure")
|
|
return original(source, destination)
|
|
monkeypatch.setattr(module.shutil, "move", fail_second)
|
|
with pytest.raises(OSError):
|
|
current.commit("owner")
|
|
current.purge_committed_and_discard() # cleanup is idempotent
|
|
assert store.global_reserved_bytes == 4
|
|
assert len(list(store.committed_dir.iterdir())) == 1
|
|
assert store.get_artifact_for_download(retained["id"], "other") is not None
|
|
|
|
|
|
def test_purge_after_successful_commit(tmp_path):
|
|
store = ArtifactStore(tmp_path)
|
|
staging = stage_file(store, "q")
|
|
metadata = staging.commit("owner")
|
|
staging.purge_committed_and_discard()
|
|
assert store.global_reserved_bytes == 0
|
|
assert store.get_artifact_for_download(metadata[0]["id"], "owner") is None
|
|
assert not list(store.committed_dir.iterdir())
|
|
|
|
|
|
def test_active_download_survives_expiry(tmp_path):
|
|
store = ArtifactStore(tmp_path)
|
|
aid = stage_file(store, "q").commit("owner")[0]["id"]
|
|
path, _, _ = store.get_artifact_for_download(aid, "owner")
|
|
store._committed[aid].expires_at_ts = 0
|
|
store.expire_artifacts()
|
|
assert path.read_bytes() == b"data"
|
|
assert store.get_artifact_for_download(aid, "owner") is None
|
|
store.release_reader(aid)
|
|
store.expire_artifacts()
|
|
assert not path.exists()
|
|
assert store.global_reserved_bytes == 0
|
|
|
|
|
|
def test_duplicate_transfer_names_permissions_and_invalid_chunks(tmp_path):
|
|
store = ArtifactStore(tmp_path)
|
|
staging = stage_file(store, "q", "clé.txt")
|
|
with pytest.raises(InvalidInputError):
|
|
staging.handle_begin("t_1", "another.txt", 0)
|
|
assert staging.handle_begin("t_2", "cle\u0301.txt", 0)[0] == "skip"
|
|
assert staging.handle_begin("t_3", "extra.txt", 1)[0] == "accept"
|
|
assert os.stat(staging.open_transfer.staging_path).st_mode & 0o777 == 0o600
|
|
with pytest.raises(InvalidInputError):
|
|
staging.handle_chunk("t_3", 0, "!")
|
|
with pytest.raises(InvalidInputError):
|
|
staging.handle_end("t_3", 1, 1)
|
|
staging.discard()
|
|
assert store.global_reserved_bytes == 0
|
|
|
|
|
|
def test_reservations_are_atomic_between_threads(tmp_path, monkeypatch):
|
|
import backend.artifacts as module
|
|
monkeypatch.setattr(module, "GLOBAL_STORAGE_LIMIT", 5)
|
|
store = ArtifactStore(tmp_path)
|
|
sessions = [store.create_staging_session(f"q_{i}") for i in range(2)]
|
|
barrier = threading.Barrier(2)
|
|
def begin(session):
|
|
barrier.wait()
|
|
return session.handle_begin("t", "file.txt", 4)[0]
|
|
with ThreadPoolExecutor(2) as pool:
|
|
results = list(pool.map(begin, sessions))
|
|
assert sorted(results) == ["accept", "skip"]
|
|
assert store.global_reserved_bytes == 4
|
|
for session in sessions:
|
|
session.discard()
|
|
assert store.global_reserved_bytes == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_timeout_and_disconnect_while_waiting_for_frame(tmp_path):
|
|
manager = FakeContainerManager(lambda: ScriptedContainerPeer("timeout"))
|
|
runner = runner_for(tmp_path, manager, timeout=.03)
|
|
with pytest.raises(QueryTimeoutError):
|
|
await run(runner)
|
|
assert_empty(runner)
|
|
assert not manager.active_containers
|
|
disconnected = asyncio.Event()
|
|
runner.settings.query_timeout_seconds = 1
|
|
task = asyncio.create_task(run(runner, disconnected.is_set))
|
|
await asyncio.sleep(.02)
|
|
disconnected.set()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert_empty(runner)
|
|
assert not manager.active_containers
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("disconnect", [False, True])
|
|
async def test_deadline_disconnect_cancel_remote_dispatch(tmp_path, disconnect):
|
|
entered, cancelled, lost = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
|
class SlowClient:
|
|
def __init__(self, **kwargs):
|
|
self.closed = False
|
|
async def search(self, **kwargs):
|
|
entered.set()
|
|
try:
|
|
await asyncio.sleep(10)
|
|
finally:
|
|
cancelled.set()
|
|
async def close(self):
|
|
self.closed = True
|
|
client = SlowClient()
|
|
runner = runner_for(tmp_path, timeout=1 if disconnect else .05, factory=lambda **kw: client)
|
|
task = asyncio.create_task(run(runner, lost.is_set))
|
|
await entered.wait()
|
|
if disconnect:
|
|
lost.set()
|
|
with pytest.raises(asyncio.CancelledError if disconnect else QueryTimeoutError):
|
|
await task
|
|
assert cancelled.is_set() and client.closed
|
|
assert_empty(runner)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_repeated_cancellation_holds_gate_through_cleanup(tmp_path):
|
|
entered, finish = asyncio.Event(), asyncio.Event()
|
|
class SlowRemoval(FakeContainerManager):
|
|
async def kill_and_remove(self, container_id, timeout=10):
|
|
entered.set()
|
|
await finish.wait()
|
|
await super().kill_and_remove(container_id, timeout)
|
|
manager = SlowRemoval(lambda: ScriptedContainerPeer("agent_error"))
|
|
runner = runner_for(tmp_path, manager)
|
|
task = asyncio.create_task(run(runner))
|
|
await entered.wait()
|
|
task.cancel()
|
|
await asyncio.sleep(0)
|
|
task.cancel()
|
|
with pytest.raises(BusyError):
|
|
await run(runner)
|
|
assert not task.done()
|
|
finish.set()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert not runner._gate.locked()
|
|
assert not manager.active_containers
|
|
assert_empty(runner)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_uncertainty_blocks_queries_until_reconciliation(tmp_path):
|
|
class FailedRemoval(FakeContainerManager):
|
|
fail = True
|
|
async def kill_and_remove(self, container_id, timeout=10):
|
|
if self.fail:
|
|
raise CleanupFailedError()
|
|
await super().kill_and_remove(container_id, timeout)
|
|
manager = FailedRemoval(lambda: ScriptedContainerPeer("agent_error"))
|
|
runner = runner_for(tmp_path, manager)
|
|
with pytest.raises(CleanupFailedError):
|
|
await run(runner)
|
|
assert_empty(runner)
|
|
with pytest.raises(BusyError):
|
|
await run(runner)
|
|
manager.fail = False
|
|
await runner.reconcile()
|
|
assert runner.ready and not manager.active_containers
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_timeout_does_not_release_gate_early(tmp_path):
|
|
entered, finish = asyncio.Event(), asyncio.Event()
|
|
class DelayedRemoval(FakeContainerManager):
|
|
async def kill_and_remove(self, container_id, timeout=10):
|
|
entered.set()
|
|
await finish.wait()
|
|
await super().kill_and_remove(container_id, timeout)
|
|
runner = runner_for(tmp_path, DelayedRemoval(lambda: ScriptedContainerPeer("agent_error")), cleanup_timeout=.01)
|
|
task = asyncio.create_task(run(runner))
|
|
await entered.wait()
|
|
await asyncio.sleep(.03)
|
|
with pytest.raises(BusyError):
|
|
await run(runner)
|
|
finish.set()
|
|
with pytest.raises(CleanupFailedError):
|
|
await task
|
|
assert_empty(runner)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_failure_discards_staging(tmp_path):
|
|
class FailedStart(FakeContainerManager):
|
|
async def create_and_run(self, query_id):
|
|
await super().create_and_run(query_id)
|
|
raise UpstreamFailedError()
|
|
manager = FailedStart(lambda: ScriptedContainerPeer("timeout"))
|
|
runner = runner_for(tmp_path, manager)
|
|
with pytest.raises(UpstreamFailedError):
|
|
await run(runner)
|
|
assert not manager.active_containers
|
|
assert_empty(runner)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("scenario", ["no_artifacts", "skip_artifact"])
|
|
async def test_zero_stored_exports(tmp_path, scenario):
|
|
runner = runner_for(tmp_path, FakeContainerManager(lambda: ScriptedContainerPeer(scenario)))
|
|
result = await run(runner)
|
|
assert result["artifacts"] == []
|
|
if scenario == "skip_artifact":
|
|
assert result["warnings"][0]["code"] == "artifact_size_exceeded"
|
|
assert_empty(runner)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("failure", ["count", "error", "base64", "open"])
|
|
async def test_failed_transfer_protocol_clears_all_exports(tmp_path, failure):
|
|
class BrokenPeer(ScriptedContainerPeer):
|
|
async def _write_msg(self, writer, message):
|
|
if message["type"] == "complete":
|
|
message = {**message, "payload": {"accepted_transfer_count": 2}}
|
|
if failure == "error":
|
|
message = {"v": 1, "id": "a_terminal", "type": "error", "payload": {"code": "secret", "message": "valid-pat"}}
|
|
if failure == "base64" and message["type"] == "artifact_chunk":
|
|
message["payload"]["data_base64"] = "!"
|
|
if failure == "open" and message["type"] == "artifact_chunk":
|
|
message = {"v": 1, "id": "a_complete", "type": "complete", "payload": {"accepted_transfer_count": 0}}
|
|
await super()._write_msg(writer, message)
|
|
runner = runner_for(tmp_path, FakeContainerManager(BrokenPeer))
|
|
from backend.errors import AppError
|
|
with pytest.raises(AppError) as error:
|
|
await run(runner)
|
|
assert "valid-pat" not in error.value.message
|
|
assert_empty(runner)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multimib_frame_through_fake_manager(tmp_path):
|
|
markdown = "é" * (1024 * 1024)
|
|
async def steps(reader, writer, peer):
|
|
await peer._read_line(reader)
|
|
await peer._write_msg(writer, {"v": 1, "id": "a_1", "type": "collection_start", "payload": {"markdown": markdown, "warnings": []}})
|
|
await peer._read_line(reader)
|
|
await peer._write_msg(writer, {"v": 1, "id": "a_2", "type": "complete", "payload": {"accepted_transfer_count": 0}})
|
|
runner = runner_for(tmp_path, FakeContainerManager(lambda: ScriptedContainerPeer(custom_steps=steps)), timeout=5)
|
|
assert (await run(runner))["markdown"] == markdown
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_real_subprocess_large_stdout_and_chatty_stderr(monkeypatch):
|
|
original = asyncio.create_subprocess_exec
|
|
script = 'import sys,json; sys.stderr.buffer.write(b"x"*262144); sys.stderr.flush(); print(json.dumps({"v":1,"id":"a_1","type":"collection_start","payload":{"markdown":"z"*2097152,"warnings":[]}}),flush=True)'
|
|
async def spawn(*args, **kwargs):
|
|
assert kwargs["limit"] > 2 * 1024 * 1024
|
|
assert "--network" in args and "none" in args
|
|
return await original(sys.executable, "-c", script, **kwargs)
|
|
monkeypatch.setattr(asyncio, "create_subprocess_exec", spawn)
|
|
manager = DockerContainerManager(Settings())
|
|
handle = await manager.create_and_run("q")
|
|
drainer = StderrDrainer(handle.stderr)
|
|
drainer.start()
|
|
bridge = BridgeTransport(handle.stdout, handle.stdin)
|
|
await bridge.send_start("p", "s", 1000, {})
|
|
message = await asyncio.wait_for(bridge.read_message_bounded(), 3)
|
|
assert len(message["payload"]["markdown"]) == 2097152
|
|
await handle.process.wait()
|
|
assert len(drainer.buffer) == 64 * 1024
|
|
await drainer.stop()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fragmentation_multiple_frames_and_incremental_limit():
|
|
reader = asyncio.StreamReader()
|
|
bridge = BridgeTransport(reader, Sink(), max_frame_bytes=256)
|
|
await bridge.send_start("p", "s", 1000, {})
|
|
frames = [
|
|
{"v": 1, "type": "collection_start", "id": "a_1", "payload": {"markdown": "é", "warnings": []}},
|
|
{"v": 1, "type": "complete", "id": "a_2", "payload": {"accepted_transfer_count": 0}},
|
|
]
|
|
raw = b"".join(json.dumps(f, ensure_ascii=False).encode() + b"\n" for f in frames)
|
|
task = asyncio.create_task(bridge.read_message_bounded())
|
|
for byte in raw:
|
|
reader.feed_data(bytes([byte]))
|
|
await asyncio.sleep(0)
|
|
assert (await task)["payload"]["markdown"] == "é"
|
|
assert (await bridge.read_message_bounded())["type"] == "complete"
|
|
reader.feed_eof()
|
|
assert (await bridge.read_message_bounded())["type"] == "eof"
|
|
oversized = asyncio.StreamReader()
|
|
bridge = BridgeTransport(oversized, Sink(), max_frame_bytes=16)
|
|
oversized.feed_data(b"x" * 17) # no newline and no EOF: must fail immediately
|
|
with pytest.raises(NDJSONProtocolError):
|
|
await asyncio.wait_for(bridge.read_message_bounded(), .1)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("frame", [
|
|
{"v": 1, "id": "a_1", "type": "error", "payload": {"code": "x", "message": "x"}},
|
|
{"v": 1, "id": "a_1", "type": "tool_request", "payload": None},
|
|
{"v": 1, "id": "a_1", "type": "tool_request", "payload": {"tool": "x", "parameters": {}}, "reply_to": "b_1"},
|
|
{"v": 1, "id": "a_1", "type": "complete", "payload": {"accepted_transfer_count": True}},
|
|
])
|
|
async def test_malformed_envelopes_and_init_error(frame):
|
|
reader = asyncio.StreamReader()
|
|
reader.feed_data(json.dumps(frame).encode() + b"\n")
|
|
bridge = BridgeTransport(reader, Sink())
|
|
with pytest.raises(NDJSONProtocolError):
|
|
await bridge.read_message_bounded()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_eof_and_type_specific_payload_limit(monkeypatch):
|
|
reader = asyncio.StreamReader()
|
|
reader.feed_eof()
|
|
with pytest.raises(NDJSONProtocolError):
|
|
await BridgeTransport(reader, Sink()).read_message_bounded()
|
|
import backend.transport as module
|
|
monkeypatch.setattr(module, "MAX_TOOL_REQUEST_PAYLOAD", 1024)
|
|
reader = asyncio.StreamReader()
|
|
reader.feed_data(json.dumps({"v": 1, "id": "a_1", "type": "tool_request", "payload": {"tool": "x", "parameters": {"query": "é" * 1024}}}).encode() + b"\n")
|
|
bridge = BridgeTransport(reader, Sink())
|
|
await bridge.send_start("p", "s", 1000, {})
|
|
with pytest.raises(NDJSONProtocolError):
|
|
await bridge.read_message_bounded()
|
|
|
|
|
|
class RawStream(httpx.AsyncByteStream):
|
|
def __init__(self, chunks):
|
|
self.chunks, self.closed = chunks, False
|
|
async def __aiter__(self):
|
|
for chunk in self.chunks:
|
|
yield chunk
|
|
async def aclose(self):
|
|
self.closed = True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("compressed", [False, True])
|
|
async def test_upstream_wire_and_decompressed_limits(compressed):
|
|
raw = b'"' + b"x" * 2048 + b'"'
|
|
if compressed:
|
|
compressor = zlib.compressobj(wbits=31)
|
|
raw = compressor.compress(raw) + compressor.flush()
|
|
stream = RawStream([raw[:10], raw[10:]])
|
|
response = httpx.Response(200, stream=stream, headers={"content-encoding": "gzip"} if compressed else {})
|
|
with pytest.raises(UpstreamResponseTooLargeError):
|
|
await read_json(response, 1024)
|
|
assert stream.closed
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_valid_compressed_json_and_malformed_response():
|
|
compressor = zlib.compressobj(wbits=31)
|
|
raw = compressor.compress(b'{"results":[]}') + compressor.flush()
|
|
response = httpx.Response(200, stream=RawStream([raw]), headers={"content-encoding": "gzip"})
|
|
assert await read_json(response, 1024) == {"results": []}
|
|
response = httpx.Response(200, stream=RawStream([b"not JSON"]))
|
|
with pytest.raises(UpstreamFailedError):
|
|
await read_json(response, 1024)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("body,status,code", [
|
|
({"error": {"message": "valid-pat"}}, 400, "upstream_failed"),
|
|
({"error": {"code": "context_length_exceeded"}}, 400, "model_context_exceeded"),
|
|
({"choices": [{"message": {"content": "hi"}, "finish_reason": "stop"}], "usage": {"completion_tokens": 11}}, 200, "model_output_limit"),
|
|
([], 200, "upstream_failed"),
|
|
({"choices": [None]}, 200, "upstream_failed"),
|
|
({"choices": [{"message": {"tool_calls": [{"id": "tc", "function": {"name": "x", "arguments": "[]"}}]}, "finish_reason": "tool_calls"}]}, 200, "upstream_failed"),
|
|
])
|
|
async def test_provider_errors_limits_and_malformed_shapes(body, status, code):
|
|
adapter = OpenAIModelAdapter("key", max_output_tokens=10, transport=httpx.MockTransport(lambda req: httpx.Response(status, json=body)))
|
|
from backend.errors import AppError
|
|
try:
|
|
with pytest.raises(AppError) as error:
|
|
await adapter.complete([], [], "system")
|
|
assert error.value.code == code
|
|
assert "valid-pat" not in error.value.message
|
|
finally:
|
|
await adapter.close()
|
|
|
|
|
|
def test_history_byte_budgets_and_audit_before_truncation(monkeypatch):
|
|
import backend.history as module
|
|
monkeypatch.setattr(module, "MAX_HISTORY_BYTES", 128 * 1024)
|
|
monkeypatch.setattr(module, "MAX_HISTORY_ENTRIES", 2)
|
|
history = HistoryManager()
|
|
history.record_call("a_1", "confluence_view", {str(i): "v" for i in range(10000)}, "t", "t", "success", False, {"page_id": "1", "title": "Guide", "space": "OPS", "url": "https://approved.example.com/page", "markdown": "é" * 1024 * 1024, "truncated": False}, None)
|
|
history.record_call("a_2", "confluence_list_spaces", {}, "t", "t", "success", False, {"spaces": [{"key": "x", "name": "y" * 1024 * 1024}], "pagination": {"offset": 0, "limit": 1, "has_more": False}}, None)
|
|
assert len(json.dumps(history.get_tool_history()).encode()) <= 128 * 1024
|
|
first = history.get_tool_history()[0]
|
|
assert first["parameters_truncated"] and first["result_truncated"]
|
|
metadata = {**first, "result": None}
|
|
assert len(json.dumps(metadata).encode()) <= RESERVED_METADATA_BYTES_PER_ENTRY
|
|
assert history.get_pages_accessed()[0]["page_id"] == "1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_docker_uncertainty_budget_and_rootless_checks(monkeypatch):
|
|
manager = DockerContainerManager(Settings())
|
|
calls = []
|
|
async def daemon_down(args, timeout=15):
|
|
calls.append((args, timeout))
|
|
return 1, "", "unavailable"
|
|
monkeypatch.setattr(manager, "_exec_docker", daemon_down)
|
|
with pytest.raises(CleanupFailedError):
|
|
await manager.kill_and_remove("query", timeout=.02)
|
|
assert all(0 < timeout <= .02 for _, timeout in calls)
|
|
with pytest.raises(CleanupFailedError):
|
|
await manager.reconcile_orphans()
|
|
async def rootful(args, timeout=15):
|
|
return 0, json.dumps({"SecurityOptions": [], "DockerRootDir": "/home/docker"}), ""
|
|
monkeypatch.setattr(manager, "_exec_docker", rootful)
|
|
from backend.errors import ExecutionFailedError
|
|
with pytest.raises(ExecutionFailedError):
|
|
await manager.verify_rootless()
|
|
async def missing_limits(args, timeout=15):
|
|
return 0, json.dumps({"SecurityOptions": ["name=rootless"]}), ""
|
|
monkeypatch.setattr(manager, "_exec_docker", missing_limits)
|
|
with pytest.raises(ExecutionFailedError):
|
|
await manager.verify_rootless()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_startup_verifies_reconciles_purges_and_closes(tmp_path):
|
|
manager = DockerContainerManager(Settings())
|
|
events = []
|
|
async def verify():
|
|
events.append("verify")
|
|
async def reconcile(max_age_seconds=200):
|
|
assert max_age_seconds == 0
|
|
events.append("reconcile")
|
|
return 0
|
|
manager.verify_rootless, manager.reconcile_orphans = verify, reconcile
|
|
class ClosingAdapter(FakeModelAdapter):
|
|
async def close(self):
|
|
events.append("close")
|
|
store = ArtifactStore(tmp_path)
|
|
stage_file(store, "old").commit("owner")
|
|
app = create_app(Settings(), manager, store, ClosingAdapter())
|
|
async with app.router.lifespan_context(app):
|
|
assert events == ["verify", "reconcile"]
|
|
assert not list(store.committed_dir.iterdir())
|
|
assert events[-1] == "close"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_valid_large_prompt_origin_defaults_and_forged_cookie(tmp_path):
|
|
app = create_app(Settings(dev_mode=True, artifact_storage_dir=tmp_path), confluence_client_factory=make_test_confluence_client_factory())
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="https://testserver") as client:
|
|
payload = {"prompt": "é" * 65536, "credentials": {"url": "https://approved.example.com", "pat": "valid-pat"}}
|
|
for origin in ("http://testserver", "https://user@testserver", "https://testserver/path"):
|
|
response = await client.post("/api/v1/query", json=payload, headers={"origin": origin})
|
|
assert response.status_code == 403
|
|
assert response.headers["cache-control"] == "no-store"
|
|
client.cookies.set("cw_session", "a" * 64)
|
|
response = await client.post("/api/v1/query", json=payload, headers={"origin": "https://testserver:443"})
|
|
assert response.status_code == 200
|
|
assert response.cookies["cw_session"] != "a" * 64
|
|
assert "Secure" in response.headers["set-cookie"]
|
|
|
|
|
|
@pytest.mark.parametrize("url", ["https://example.com:abc", "https://example.com:99999", "http://[invalid", "https://example.com/\nsecret"])
|
|
def test_bad_urls_have_application_errors(url):
|
|
with pytest.raises(InvalidInputError):
|
|
canonicalize_url(url)
|
|
|
|
|
|
def test_ipv6_loopback_and_environment_errors(monkeypatch):
|
|
assert validate_confluence_url("http://[::1]:8080/wiki", ["http://[::1]:8080/wiki"]) == "http://[::1]:8080/wiki"
|
|
monkeypatch.setenv("CONFLUENCE_WEB_BIND_PORT", "valid-pat")
|
|
with pytest.raises(ValueError) as error:
|
|
Settings.from_env()
|
|
assert "valid-pat" not in str(error.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("operation", ["create", "commit"])
|
|
async def test_cancelled_filesystem_operation_finishes_before_discard(tmp_path, operation):
|
|
entered, finish = threading.Event(), threading.Event()
|
|
store = ArtifactStore(tmp_path)
|
|
original = store.create_staging_session
|
|
def create(query_id):
|
|
if operation == "create":
|
|
entered.set()
|
|
finish.wait(2)
|
|
session = original(query_id)
|
|
if operation == "commit":
|
|
commit = session.commit
|
|
def delayed_commit(*args, **kwargs):
|
|
entered.set()
|
|
finish.wait(2)
|
|
return commit(*args, **kwargs)
|
|
session.commit = delayed_commit
|
|
return session
|
|
store.create_staging_session = create
|
|
runner = runner_for(tmp_path)
|
|
runner.artifact_store = store
|
|
task = asyncio.create_task(run(runner))
|
|
while not entered.is_set():
|
|
await asyncio.sleep(.001)
|
|
task.cancel()
|
|
await asyncio.sleep(.01)
|
|
assert not task.done()
|
|
with pytest.raises(BusyError):
|
|
await run(runner)
|
|
finish.set()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert_empty(runner)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_network_free_dev_mode(tmp_path):
|
|
app = create_app(Settings(dev_mode=True, artifact_storage_dir=tmp_path))
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client:
|
|
verify = await client.post("/api/v1/auth/verify", json={"url": "https://approved.example.com", "pat": "dev-pat"}, headers={"origin": "http://testserver"})
|
|
assert verify.json() == {"valid": True}
|
|
result = await client.post("/api/v1/query", json={"prompt": "deploy", "credentials": {"url": "https://approved.example.com", "pat": "dev-pat"}}, headers={"origin": "http://testserver"})
|
|
assert result.status_code == 200
|
|
assert len(result.json()["pages_accessed"]) == 1
|
|
assert isinstance(app.state.runner.container_manager, FakeContainerManager)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_model_tool_result_roundtrip_and_state_association():
|
|
from backend.model import ModelDispatcher
|
|
calls = []
|
|
async def respond(request):
|
|
body = json.loads(request.content)
|
|
calls.append(body)
|
|
if len(calls) == 1:
|
|
return httpx.Response(200, json={"choices": [{"message": {"content": None, "tool_calls": [{"id": "tc_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}]}, "finish_reason": "tool_calls"}]})
|
|
assert body["messages"][2]["tool_calls"][0]["id"] == "tc_1"
|
|
assert body["messages"][3] == {"role": "tool", "tool_call_id": "tc_1", "content": "file contents"}
|
|
return httpx.Response(200, json={"choices": [{"message": {"content": "done"}, "finish_reason": "stop"}]})
|
|
adapter = OpenAIModelAdapter("key", transport=httpx.MockTransport(respond))
|
|
dispatcher = ModelDispatcher(adapter)
|
|
user = {"role": "user", "content": [{"type": "text", "text": "read file"}]}
|
|
try:
|
|
first, error = await dispatcher.dispatch({"messages": [user], "tools": []}, "trusted")
|
|
assert error is None
|
|
second, error = await dispatcher.dispatch({"messages": [user, {"role": "assistant", "content": first["content"]}, {"role": "tool", "tool_call_id": "tc_1", "name": "read", "content": "file contents", "is_error": False}], "tools": []}, "trusted")
|
|
assert error is None and second["content"][0]["text"] == "done"
|
|
finally:
|
|
await adapter.close()
|
|
fake = FakeModelAdapter([{"content": [{"type": "text", "text": "original"}], "stop_reason": "stop", "usage": {"input_tokens": 1, "output_tokens": 1}, "provider_state": "private"}])
|
|
dispatcher = ModelDispatcher(fake)
|
|
response, error = await dispatcher.dispatch({"messages": [], "tools": []}, "trusted")
|
|
assert error is None and response["provider_state"] != "private"
|
|
_, error = await dispatcher.dispatch({"messages": [{"role": "assistant", "content": [{"type": "text", "text": "changed"}], "provider_state": response["provider_state"]}], "tools": []}, "trusted")
|
|
assert error["code"] == "invalid_input"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("payload", [
|
|
{"messages": [None], "tools": []},
|
|
{"messages": [], "tools": [], "model": "untrusted"},
|
|
{"messages": [], "tools": [{"name": "x", "description": "x", "input_schema": []}]},
|
|
])
|
|
async def test_model_shapes_rejected_before_provider_call(payload):
|
|
from backend.model import ModelDispatcher
|
|
fake = FakeModelAdapter()
|
|
_, error = await ModelDispatcher(fake).dispatch(payload, "trusted")
|
|
assert error["code"] == "invalid_input" and fake.call_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_actual_chunked_gzip_upstream_bound(monkeypatch):
|
|
from backend.confluence import ConfluenceClient
|
|
import backend.confluence as module
|
|
monkeypatch.setattr(module, "MAX_RESPONSE_BYTES", 1024)
|
|
compressor = zlib.compressobj(wbits=31)
|
|
payload = compressor.compress(b'{"results":[],"padding":"' + b"x" * 2048 + b'"}') + compressor.flush()
|
|
async def serve(reader, writer):
|
|
await reader.readuntil(b"\r\n\r\n")
|
|
writer.write(b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Encoding: gzip\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n")
|
|
for offset in range(0, len(payload), 7):
|
|
chunk = payload[offset:offset + 7]
|
|
writer.write(f"{len(chunk):x}\r\n".encode() + chunk + b"\r\n")
|
|
writer.write(b"0\r\n\r\n")
|
|
await writer.drain()
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
server = await asyncio.start_server(serve, "127.0.0.1", 0)
|
|
url = f"http://127.0.0.1:{server.sockets[0].getsockname()[1]}"
|
|
try:
|
|
async with ConfluenceClient(url, "pat", [url]) as client:
|
|
with pytest.raises(UpstreamResponseTooLargeError):
|
|
await client.verify_auth()
|
|
finally:
|
|
server.close()
|
|
await server.wait_closed()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_conversion_worker_is_killed_on_cancel(monkeypatch):
|
|
from backend.confluence import ConfluenceClient
|
|
original = asyncio.create_subprocess_exec
|
|
created = asyncio.Event()
|
|
processes = []
|
|
async def spawn(*args, **kwargs):
|
|
assert "CONFLUENCE_WEB_MODEL_API_KEY" not in kwargs["env"]
|
|
process = await original(sys.executable, "-c", "import sys,time;sys.stdin.read();time.sleep(30)", **kwargs)
|
|
processes.append(process)
|
|
created.set()
|
|
return process
|
|
monkeypatch.setattr(asyncio, "create_subprocess_exec", spawn)
|
|
transport = httpx.MockTransport(lambda req: httpx.Response(200, json={"body": {"storage": {"value": "<p>text</p>"}}}))
|
|
async with ConfluenceClient("https://approved.example.com", "pat", ["https://approved.example.com"], transport=transport) as client:
|
|
task = asyncio.create_task(client.view("1"))
|
|
await created.wait()
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert processes[0].returncode is not None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("body", [{}, [], {"results": "login"}])
|
|
async def test_verify_auth_rejects_unrelated_json(body):
|
|
from backend.confluence import ConfluenceClient
|
|
async with ConfluenceClient("https://approved.example.com", "pat", ["https://approved.example.com"], transport=httpx.MockTransport(lambda req: httpx.Response(200, json=body))) as client:
|
|
with pytest.raises(UpstreamFailedError):
|
|
await client.verify_auth()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_startup_failure_closes_adapter(tmp_path):
|
|
from backend.errors import ExecutionFailedError
|
|
manager = DockerContainerManager(Settings())
|
|
async def fail():
|
|
raise ExecutionFailedError()
|
|
manager.verify_rootless = fail
|
|
class Adapter(FakeModelAdapter):
|
|
closed = False
|
|
async def close(self):
|
|
self.closed = True
|
|
adapter = Adapter()
|
|
app = create_app(Settings(artifact_storage_dir=tmp_path), container_manager=manager, model_adapter=adapter)
|
|
with pytest.raises(ExecutionFailedError):
|
|
async with app.router.lifespan_context(app):
|
|
pytest.fail("startup should fail")
|
|
assert adapter.closed
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_download_lookup_releases_reader(tmp_path):
|
|
app = create_app(Settings(dev_mode=True, artifact_storage_dir=tmp_path))
|
|
store = app.state.artifact_store
|
|
entered, finish = threading.Event(), threading.Event()
|
|
original = store.get_artifact_for_download
|
|
def delayed_lookup(*args, **kwargs):
|
|
entered.set()
|
|
finish.wait(2)
|
|
return original(*args, **kwargs)
|
|
store.get_artifact_for_download = delayed_lookup
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client:
|
|
session = (await client.get("/")).cookies["cw_session"]
|
|
aid = stage_file(store, "q").commit(session)[0]["id"]
|
|
task = asyncio.create_task(client.get(f"/api/v1/artifacts/{aid}"))
|
|
while not entered.is_set():
|
|
await asyncio.sleep(.001)
|
|
task.cancel()
|
|
await asyncio.sleep(.01)
|
|
assert not task.done()
|
|
finish.set()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert store._committed[aid].active_readers == 0
|
|
store._committed[aid].expires_at_ts = 0
|
|
store.expire_artifacts()
|
|
assert not list(store.committed_dir.iterdir())
|