diff --git a/.github/workflows/windows-private-storage.yml b/.github/workflows/windows-private-storage.yml new file mode 100644 index 0000000..a8cfa63 --- /dev/null +++ b/.github/workflows/windows-private-storage.yml @@ -0,0 +1,16 @@ +name: Windows private storage +on: + push: + pull_request: + workflow_dispatch: +jobs: + windows-private-storage: + runs-on: windows-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@v5 + - uses: astral-sh/setup-uv@v6 + with: + version: "0.12.13" + - run: uv sync --locked --python 3.13 + - run: uv run python -m unittest discover -s tests -p test_private_storage_windows.py -v diff --git a/README.md b/README.md index 58fd4a2..741b080 100644 --- a/README.md +++ b/README.md @@ -29,9 +29,15 @@ job-routing instructions; the skill folders can also be installed separately. ## Long queries, CLI and progress Use `omop-submit_query_job` for an export or a query that could exceed the -interactive timeout. It returns a job ID immediately. Poll -`omop-query_job_status` at its recommended interval; use -`omop-cancel_query_job` to stop. Only `completed` means the file is complete. +interactive timeout. Call `omop-query_job_capacity` first; submission returns a +job ID immediately and atomically checks the two-job limit. Use +`omop-wait_for_query_job` for native waits of up to 30 seconds, or +`omop-query_job_status` for an immediate snapshot. Use `omop-cancel_query_job` +to stop. Only `completed` means the file is complete. Read bounded UTF-8 chunks +with `omop-read_query_job_result` using the job ID and `next_offset`; do not pass +`result_path` to generic filesystem tools. The secure reader requires POSIX; +unsupported platforms retain the local CLI file. See [job boundaries and the +credential-free app fixture](docs/QUERY_JOBS.md#mcp-waiting-capacity-and-result-handoff). Status includes rows, bytes, elapsed time, phase and an advisory ETA when known. Set `mode="explain"` to obtain a plan without running the query. @@ -48,7 +54,8 @@ uv run ucsfomopagent watch JOB_ID No-argument `uv run ucsfomopagent` continues to start the MCP server. `status`, `watch`, `list`, `cancel` and `purge` need no database credentials. Results remain in the -local private job directory and are not sent to chat. Database/server limits can +local private job directory. The bounded result reader can return requested chunks +to chat; keep controlled data local and share only appropriate summaries. Database/server limits can still fail a query; jobs report those failures instead of silently truncating. ### Large fact-table studies diff --git a/docs/QUERY_JOBS.md b/docs/QUERY_JOBS.md index 33611bd..80136f8 100644 --- a/docs/QUERY_JOBS.md +++ b/docs/QUERY_JOBS.md @@ -13,8 +13,9 @@ Keep each connector independently installable and releasable. Database adapters live in query_backend.py, orchestration in jobs.py/job_worker.py, and command parsing in query_cli.py. Existing domain tools remain alongside them. This avoids a new separately published dependency or any requirement to clone sibling repos. -The three small job-runtime modules intentionally have identical source; when -changing the protocol, run the same contract tests in all three repositories. +OMOP and CDW keep jobs.py and job_worker.py identical. Their query adapters differ +only in connector defaults. Run the same contract tests in both repositories when +changing the protocol; the SPOKE connector has an independently released runtime. A future shared package is appropriate once its API and release ownership settle. `extensions/ucsfomopagent.brxt` is the current committed installable artifact. Build it @@ -46,8 +47,15 @@ malicious code running under the same OS account. ## Storage and progress Jobs live in `~/.local/share/ucsf-agents/ucsfomopagent/` (override the parent using -UCSF_AGENT_JOB_ROOT). Directories are private (0700) and DB/results are 0600 on -POSIX. On Windows use the account's private profile directory and its ACLs. +UCSF_AGENT_JOB_ROOT). The storage backend checks current-account ownership, +rejects symlinks, nonregular entries and hard-linked private files, and repairs +existing owned stores before reading them. On macOS/Linux it uses retained +no-follow descriptors, modes 0700/0600 and removal of inherited ACL grants. +New result files are protected before any result bytes are written. Windows uses +a protected current-user DACL through retained no-reparse handles. Unsupported +or unsafe storage returns an actionable `private_storage_unavailable` error from +the job tools; synchronous connector tools remain available. Normal job use +requires no additional prompts or changes to the six-tool interface. SQLite serializes admission (two active jobs per connector). Results stream in 1000-row batches to a `.partial` file and are renamed only after success and fsync. The default disk budget is 10 GB; exceeding it fails visibly rather than @@ -56,6 +64,30 @@ and JSON-encodes graph/list values. Use JSONL when empty string and null must be distinguished. SQL projections require unique column aliases. Exported text may contain spreadsheet formula strings: treat it as data when importing CSV. +These permissions protect against other ordinary local accounts where filesystem +ACLs are enforced; they do not revoke file handles another process opened before +an old store was repaired. They do not isolate arbitrary code running under the +same OS account, administrators or root. Choose storage beneath a private or +administrator-controlled parent directory. SQLite and process-lease/final-rename +operations still use pathnames after identity checks; this update does not claim +to eliminate every pathname replacement race in the complete job runtime. + +On Windows, automatic repair first checks every managed directory for mutation +rights granted to another ordinary account. A writable known child remains +accessible to that account even when its parent is private, so a protected parent +alone does not make legacy migration safe. Such stores return +`private_storage_requires_quiescent_repair` before changing existing ACLs. Unknown +ACL formats receive the same conservative refusal. Already-private stores and +legacy read-only inheritance continue through normal automatic repair. + +For this refusal, stop the connector and other writers and ask an administrator +to repair only the owned ordinary entries without following links, verify that +other ordinary accounts cannot mutate any managed directory, and retry. Another +option is a fresh private `UCSF_AGENT_JOB_ROOT`, retaining the old store untouched +for trusted recovery. Moving an old directory alone does not repair its existing +ACLs. The connector does not delete old data or silently migrate an actively +shared store. Ordinary job use needs no additional approval prompts. + Status records queued/running/terminal states and separate connecting/planning/ executing/streaming phases, rows written, bytes, heartbeat, elapsed time and advisory ETA. A heartbeat proves @@ -79,6 +111,78 @@ uv run ucsfomopagent cancel JOB_ID uv run ucsfomopagent purge JOB_ID ``` +## MCP waiting, capacity and result handoff + +Six namespace-prefixed job tools are available. `submit_query_job` starts work; +`query_job_status` is an immediate lookup; `cancel_query_job` requests cancellation. +`query_job_capacity` reports `active_jobs`, `max_active_jobs` (2), and +`available_slots`. Admission is checked atomically, so a later submission can +still return `capacity_exceeded` with current capacity. Waiting jobs hold their +slot until completion, failure or confirmed cancellation. + +`wait_for_query_job(job_id, timeout_seconds=30)` waits natively for 0–30 seconds. +It returns a current status plus `wait_timed_out`; true means the wait budget ended +while the job remains active. Cancelling the MCP wait does not cancel the job. +Use another bounded wait or respect `recommended_poll_seconds` for immediate +status polls. Status lookups run off the MCP event loop with a short database +lock budget. If the wait expires while storage is busy after an active state +was observed, it returns that last state with `wait_timed_out=true`, +`wait_status_stale=true` and an explicit warning. If no state could be read, +it returns a transient `status_lookup_busy` lookup error without claiming an +operation outcome. A busy lookup never marks the job failed. No shell sleep, +fabricated progress or automatic resubmission is needed. Unknown ETA is normal +during connection, planning and execution. + +`read_query_job_result(job_id, offset=0, max_bytes=8192)` reads only a completed +job's managed CSV/JSONL file. It accepts no path. It returns `text`, byte +`offset`, `bytes_read`, `total_bytes`, `eof` and `next_offset`. Limits are 4–65536 +bytes per call; follow `next_offset` until `eof=true`. Chunk boundaries preserve +UTF-8 characters but can split rows or quoted CSV records. Assemble the desired +bounded content before parsing it. A mid-character offset is rejected. + +The reader ignores saved `result_path` values and opens a fixed result name +relative to the managed job directory, using pinned-root identity, descriptor- +relative opens and no-follow flags. It rejects directory/file symlinks, hardlinks, +nonregular files and changed result sizes. POSIX secure-open primitives are +required. On unsupported platforms (including Windows) it fails closed with +`secure_reader_unavailable`; the existing local CLI result remains available. +This does not grant arbitrary filesystem access or isolate hostile code running +under the same OS account. Keep controlled data local and use bounded summaries +rather than placing complete patient-level exports in chat. + +Failed jobs retain `error_details` with stable `code`, `kind`, sanitized `message`, +original failure `phase`, and boolean `retryable`. Driver messages and query +literals are not exposed. Retryability is advisory; resubmission remains explicit. +New submissions, status/wait responses and cancellation responses also carry: + +```json +{"_meta":{"io.biorouter/operation":{"version":1,"id":"ucsfomopagent:JOB_ID","status":"failed","error":{"code":"query_rejected","kind":"invalid_args","message":"The database rejected the query. Check schema, projections and parameters.","phase":"executing","retryable":false}}}} +``` + +Allowed operation states are `queued`, `running`, `completed`, `failed` and +`cancelled`; job `timed_out` and `interrupted` map to operation `failed`. A status +lookup that finds a failed job still has MCP `isError=false`. Unknown-job, +admission and result-reader failures have `isError=true` and structured safe +errors, without changing the stored job outcome. Older failed records without +structured details receive a generic safe operation error. + +## Credential-free app smoke fixture + +From a source checkout, create a fresh empty absolute directory and run: + +```bash +uv run python scripts/synthetic_job_server.py --root /absolute/fresh-empty-dir --delay-seconds 2 +``` + +This stdio MCP server registers the actual six job tools with an isolated store. +It accepts only `SELECT 1 AS example` (three synthetic multilingual rows), +`SELECT 0 AS simulate_failure` (stable safe failure), and +`SELECT 2 AS simulate_slow` (bounded delayed success). No real database connection, +credential lookup or installed job store is used. An existing nonempty directory +is refused. Use `omop-submit_query_job`, wait/status, and the result reader to +exercise real MCP metadata and chunked handoff in a development app. Use separate +fresh roots for concurrent fixtures. + ## Read-only and failure boundaries Use database accounts with SELECT/read-only privileges. The shared query guard @@ -130,3 +234,17 @@ Security updates remove the former Intel-macOS cryptography<49 constraint. Patched cryptography 50+ is required. If a compatible Intel macOS wheel is not available, install the upstream Rust/OpenSSL build prerequisites or use a supported Python/platform combination; do not restore a vulnerable version cap. + +## Submission and CLI follow-up semantics + +Capacity responses and rejection context include `advisory: true`: another client +may consume an available slot before submission. Atomic admission remains the +authority. Once a job ID has been allocated, a local worker-start or pipe failure +returns that durable ID with failed operation metadata; it does not hide the job +behind a generic submission error. No active slot remains allocated to that failure. + +CLI `watch JOB_ID` uses the same native bounded wait as MCP instead of a separate +sleep/poll loop. The wait remains at most 30 seconds per call and can be repeated +until terminal. `status JOB_ID` remains an immediate lookup. Terminal payload +`phase` is the terminal state, while `error_details.phase` identifies where the +failure actually occurred (for example `executing` or `worker_start`). diff --git a/extensions/SHA256SUMS b/extensions/SHA256SUMS index 2e20340..8dc10a6 100644 --- a/extensions/SHA256SUMS +++ b/extensions/SHA256SUMS @@ -1 +1 @@ -62f53afd87a1378a65e89b04c0de453123c450e499543bbe2a25a81345438e6a ucsfomopagent.brxt +0eab7151e1005693205dcd76bded202a52dbeca33787e4c3ba99d3c790ad22cf ucsfomopagent.brxt diff --git a/extensions/ucsfomopagent.brxt b/extensions/ucsfomopagent.brxt index a7e9f5a..220f36d 100644 Binary files a/extensions/ucsfomopagent.brxt and b/extensions/ucsfomopagent.brxt differ diff --git a/scripts/synthetic_job_server.py b/scripts/synthetic_job_server.py new file mode 100644 index 0000000..ceebb10 --- /dev/null +++ b/scripts/synthetic_job_server.py @@ -0,0 +1,149 @@ +"""Credential-free MCP job fixture; only the three documented queries are accepted.""" +from __future__ import annotations + +import argparse +import importlib +import json +import math +from pathlib import Path +import sys +import threading +import time + +REPOSITORY = Path(__file__).resolve().parents[1] +PACKAGE = json.loads((REPOSITORY / 'manifest.json').read_text())['name'] +PREFIXES = {'ucsfomopagent': 'omop-', 'cdwagent': 'CDW-'} +if PACKAGE not in PREFIXES: + raise RuntimeError('This fixture supports only the OMOP and CDW connector packages.') +PREFIX = PREFIXES[PACKAGE] +sys.path.insert(0, str(REPOSITORY / 'src')) + +jobs = importlib.import_module(PACKAGE + '.jobs') +worker = importlib.import_module(PACKAGE + '.job_worker') + +EXAMPLE_QUERY = 'SELECT 1 AS example' +FAILURE_QUERY = 'SELECT 0 AS simulate_failure' +SLOW_QUERY = 'SELECT 2 AS simulate_slow' +ALLOWED_QUERIES = frozenset((EXAMPLE_QUERY, FAILURE_QUERY, SLOW_QUERY)) +EXAMPLE_ROWS = ((1, 'café'), (2, '漢字 😀'), (3, 'a, "quote"\nsecond line')) +SYNTHETIC_CREDENTIALS = dict( + server='synthetic.invalid', database='synthetic', + username='synthetic-only', password='synthetic-only', +) + + +def bounded_delay(value): + try: + delay = float(value) + except (TypeError, ValueError): + raise ValueError('delay_seconds must be a number from 0 to 30.') from None + if isinstance(value, bool) or not math.isfinite(delay) or not 0 <= delay <= 30: + raise ValueError('delay_seconds must be a number from 0 to 30.') + return delay + + +class SyntheticJobStore(jobs.JobStore): + def __init__(self, root, delay_seconds=2): + self.delay_seconds = bounded_delay(delay_seconds) + root = Path(root) + if not root.is_absolute() or root.is_symlink(): + raise ValueError('Supply an absolute, fresh temporary directory as root.') + if root.exists() and (not root.is_dir() or any(root.iterdir())): + raise ValueError('The synthetic job root must be empty; use a fresh temporary directory.') + root.mkdir(parents=True, exist_ok=True, mode=0o700) + super().__init__(PACKAGE, root) + self._threads = [] + self._threads_lock = threading.Lock() + self._slow_release = threading.Event() + self._closed = False + + def submit(self, query, parameters, credentials, **options): + if (not isinstance(query, str) or query.strip() not in ALLOWED_QUERIES + or parameters not in (None, {}) or options.get('mode', 'query') != 'query'): + raise jobs.JobError( + 'synthetic_query_rejected', 'invalid_args', + 'This fixture accepts only its three documented queries without parameters.', + 'validation', + ) + if self._closed: + raise jobs.JobError('synthetic_server_closed', 'internal', + 'The synthetic job server is closed.', 'admission') + return super().submit(query.strip(), {}, dict(SYNTHETIC_CREDENTIALS), **options) + + def _stream(self, request, credentials, progress): + if request['query'] == FAILURE_QUERY: + raise jobs.JobError('synthetic_failure', 'tool_failure', + 'Synthetic job failed as requested.', 'executing') + if request['query'] == SLOW_QUERY: + progress['phase'] = 'executing' + self._slow_release.wait(self.delay_seconds) + yield ['example', 'label'], EXAMPLE_ROWS + + def _start_worker(self, job_id, request, credentials): + def run(): + try: + with self.lease(job_id).acquire(timeout=0): + worker.execute(self, job_id, dict(request=request, credentials={}), self._stream) + except Exception: + data = self.status(job_id) + jobs.record_failure(data, dict( + code='synthetic_worker_failed', kind='internal', + message='The synthetic worker could not complete.', + phase='synthetic_worker', retryable=False, + )) + self.update(data) + + thread = threading.Thread(target=run, name='synthetic-query-job', daemon=True) + with self._threads_lock: + self._threads.append(thread) + thread.start() + + def release_slow_jobs(self): + self._slow_release.set() + + def close(self): + self._closed = True + self.release_slow_jobs() + deadline = time.monotonic() + 5 + with self._threads_lock: + threads = list(self._threads) + for thread in threads: + thread.join(max(0, deadline - time.monotonic())) + return not any(thread.is_alive() for thread in threads) + + +def create_server(root, delay_seconds=2): + from fastmcp import FastMCP + store = SyntheticJobStore(root, delay_seconds) + server = FastMCP( + PACKAGE + ' synthetic jobs', + instructions=( + 'This is a synthetic fixture with no database access or real records. ' + 'Allowed queries are exactly SELECT 1 AS example, SELECT 0 AS simulate_failure, ' + 'and SELECT 2 AS simulate_slow. Use query_job_capacity, submit_query_job, ' + 'wait_for_query_job and read_query_job_result to exercise the job contract.' + ), + ) + jobs.register_job_tools(server, PACKAGE, dict(SYNTHETIC_CREDENTIALS), PREFIX, store=store) + return server, store + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--root', type=Path, required=True, + help='Absolute path to a fresh, empty temporary directory.') + parser.add_argument('--delay-seconds', type=float, default=2, + help='Synthetic slow-job delay, from 0 to 30 seconds (default: 2).') + arguments = parser.parse_args(argv) + try: + server, store = create_server(arguments.root, arguments.delay_seconds) + except ValueError as error: + parser.error(str(error)) + try: + server.run(transport='stdio') + finally: + store.close() + + +if __name__ == '__main__': + main() diff --git a/skills/ucsfomopagent-query-jobs/SKILL.md b/skills/ucsfomopagent-query-jobs/SKILL.md index f5d4943..c4fba07 100644 --- a/skills/ucsfomopagent-query-jobs/SKILL.md +++ b/skills/ucsfomopagent-query-jobs/SKILL.md @@ -15,20 +15,36 @@ BioRouter. Never ask the user to paste credentials into chat or shell arguments. 2. If cost is uncertain, submit `mode="explain"` first and inspect the completed local plan. It does not execute the query. SQL Server SHOWPLAN may require additional permission; lack of a plan is not evidence of a cheap query. -3. Submit the query once. Set an explicit timeout budget (default 3600 seconds, - maximum 86400) and a disk budget (`max_bytes`, default 10 GB). Do not run more - than two simultaneous jobs per connector. Add `expected_rows` only if already - known or reasonably estimated, and tell the user it is an estimate. -4. Poll `omop-query_job_status` using its `recommended_poll_seconds` (5–60). - Status is immediate and contains phase, elapsed time, rows, bytes and heartbeat. - Report unknown ETA honestly during blocking execution. Historical-query and - streaming-rate estimates are advisory; changing load invalidates them. -5. Only `completed` supplies a complete `result_path`. Process CSV/JSONL locally; - send summaries to chat. Exports do not silently apply an MCP preview row cap. -6. On `failed`, `timed_out` or `interrupted`, inspect the cause and revise the plan - before a new submission. Partial files are not complete datasets. There is no - automatic retry, keyset checkpoint, or safe general-purpose query resume. - Use `omop-cancel_query_job` for unwanted work, then poll to terminal state. +3. Call `omop-query_job_capacity` before admission. It reports `active_jobs`, + `max_active_jobs` (2) and `available_slots`. Submit once with an explicit timeout + budget (default 3600 seconds, maximum 86400) and disk budget (`max_bytes`, + default 10 GB). Capacity is checked atomically again on submission; a + `capacity_exceeded` rejection includes current capacity. Wait for a slot or + cancel unwanted work, rather than retrying admission in a tight loop. Supply + `expected_rows` only when already known or estimated, and label it provisional. +4. Call `omop-wait_for_query_job` with the job ID and `timeout_seconds` from + 0 to 30. This is native bounded waiting; it does not require a shell sleep. + `wait_timed_out=true` means the wait ended with the job still active, not that + the job failed. Cancelling a wait leaves the job running. For an immediate + snapshot, use `omop-query_job_status`; if polling, respect its + `recommended_poll_seconds` (5–60). Report phase, elapsed time, rows and bytes; + unknown ETA during blocking execution is normal. Estimates are advisory. +5. Only `completed` means the export is complete. Use + `omop-read_query_job_result` with `job_id`, byte `offset` (initially 0) and + `max_bytes` (4–65536). Follow `next_offset` until `eof=true`; offsets are bytes, + not rows, and CSV records may span chunks. Do not send `result_path` to generic + filesystem tools or request a wider filesystem permission. Keep patient-level + exports local and return appropriate summaries. Large local analyses can use + the CLI file; the MCP reader requires POSIX secure-open support and refuses + unsupported platforms. Exports do not silently apply a preview row cap. +6. On `failed`, `timed_out` or `interrupted`, inspect sanitized `error_details` + (`code`, `kind`, `message`, `phase`, `retryable`) and revise the plan before a + new submission. A successful status lookup can report a failed job. The MCP + `_meta['io.biorouter/operation']` envelope describes that operation outcome + separately from lookup success. Partial files are incomplete; there is no + automatic retry, keyset checkpoint or general-purpose resume. `retryable` is + advisory and never triggers resubmission. Use `omop-cancel_query_job` for + unwanted work, then wait for a terminal state before replacing it. ## Staged extraction and monitoring diff --git a/src/ucsfomopagent/_private_storage_windows.py b/src/ucsfomopagent/_private_storage_windows.py new file mode 100644 index 0000000..645e3ba --- /dev/null +++ b/src/ucsfomopagent/_private_storage_windows.py @@ -0,0 +1,494 @@ +"""Handle-checked, current-owner-only storage for Windows durable jobs.""" +from __future__ import annotations + +from contextlib import contextmanager +import ctypes +from ctypes import wintypes +import errno +import os +from pathlib import Path + +if os.name != 'nt': + raise ImportError('The Windows private storage backend requires Windows.') + +import msvcrt + + +_kernel = ctypes.WinDLL('kernel32', use_last_error=True) +_security = ctypes.WinDLL('advapi32', use_last_error=True) +_POINTER = ctypes.c_void_p +_READ_CONTROL = 0x00020000 +_WRITE_DAC = 0x00040000 +_READ_ATTRIBUTES = 0x80 +_LIST_DIRECTORY = 1 +_SHARE_READ_WRITE = 3 +_SHARE_DELETE = 4 +_DIRECTORY = 0x10 +_REPARSE = 0x400 +_BACKUP_SEMANTICS = 0x02000000 +_OPEN_REPARSE_POINT = 0x00200000 +_INVALID_HANDLE = _POINTER(-1).value +_OWNER_INFO = 1 +_DACL_INFO = 4 +_PROTECTED_DACL_INFO = 0x80000000 +_DACL_PROTECTED = 0x1000 +_FILE_ALL_ACCESS = 0x001F01FF +_MUTATION_RIGHTS = 0x510D0156 +_KNOWN_FILE_RIGHTS = _FILE_ALL_ACCESS | 0xF1000000 + + +class _SecurityAttributes(ctypes.Structure): + _fields_ = [('nLength', wintypes.DWORD), ('lpSecurityDescriptor', _POINTER), + ('bInheritHandle', wintypes.BOOL)] + + +class _SidAttributes(ctypes.Structure): + _fields_ = [('Sid', _POINTER), ('Attributes', wintypes.DWORD)] + + +class _TokenUser(ctypes.Structure): + _fields_ = [('User', _SidAttributes)] + + +class _FileInformation(ctypes.Structure): + _fields_ = [('attributes', wintypes.DWORD), ('creation', wintypes.FILETIME), + ('access', wintypes.FILETIME), ('write', wintypes.FILETIME), + ('volume', wintypes.DWORD), ('size_high', wintypes.DWORD), + ('size_low', wintypes.DWORD), ('links', wintypes.DWORD), + ('index_high', wintypes.DWORD), ('index_low', wintypes.DWORD)] + + +class _AclHeader(ctypes.Structure): + _fields_ = [('revision', wintypes.BYTE), ('reserved', wintypes.BYTE), + ('size', wintypes.WORD), ('count', wintypes.WORD), + ('reserved2', wintypes.WORD)] + + +class _AceHeader(ctypes.Structure): + _fields_ = [('kind', wintypes.BYTE), ('flags', wintypes.BYTE), ('size', wintypes.WORD)] + + +_kernel.GetCurrentProcess.argtypes = [] +_kernel.GetCurrentProcess.restype = wintypes.HANDLE +_kernel.CloseHandle.argtypes = [wintypes.HANDLE] +_kernel.CloseHandle.restype = wintypes.BOOL +_kernel.LocalFree.argtypes = [_POINTER] +_kernel.LocalFree.restype = _POINTER +_kernel.CreateFileW.argtypes = [wintypes.LPCWSTR, wintypes.DWORD, wintypes.DWORD, + ctypes.POINTER(_SecurityAttributes), wintypes.DWORD, + wintypes.DWORD, wintypes.HANDLE] +_kernel.CreateFileW.restype = wintypes.HANDLE +_kernel.CreateDirectoryW.argtypes = [wintypes.LPCWSTR, ctypes.POINTER(_SecurityAttributes)] +_kernel.CreateDirectoryW.restype = wintypes.BOOL +_kernel.GetFileInformationByHandle.argtypes = [wintypes.HANDLE, ctypes.POINTER(_FileInformation)] +_kernel.GetFileInformationByHandle.restype = wintypes.BOOL +_kernel.GetFileType.argtypes = [wintypes.HANDLE] +_kernel.GetFileType.restype = wintypes.DWORD +_security.OpenProcessToken.argtypes = [wintypes.HANDLE, wintypes.DWORD, ctypes.POINTER(wintypes.HANDLE)] +_security.OpenProcessToken.restype = wintypes.BOOL +_security.GetTokenInformation.argtypes = [wintypes.HANDLE, ctypes.c_int, _POINTER, + wintypes.DWORD, ctypes.POINTER(wintypes.DWORD)] +_security.GetTokenInformation.restype = wintypes.BOOL +_security.ConvertSidToStringSidW.argtypes = [_POINTER, ctypes.POINTER(_POINTER)] +_security.ConvertSidToStringSidW.restype = wintypes.BOOL +_security.ConvertStringSidToSidW.argtypes = [wintypes.LPCWSTR, ctypes.POINTER(_POINTER)] +_security.ConvertStringSidToSidW.restype = wintypes.BOOL +_security.ConvertStringSecurityDescriptorToSecurityDescriptorW.argtypes = [ + wintypes.LPCWSTR, wintypes.DWORD, ctypes.POINTER(_POINTER), ctypes.POINTER(wintypes.DWORD)] +_security.ConvertStringSecurityDescriptorToSecurityDescriptorW.restype = wintypes.BOOL +_security.GetSecurityDescriptorDacl.argtypes = [ + _POINTER, ctypes.POINTER(wintypes.BOOL), ctypes.POINTER(_POINTER), ctypes.POINTER(wintypes.BOOL)] +_security.GetSecurityDescriptorDacl.restype = wintypes.BOOL +_security.GetSecurityDescriptorControl.argtypes = [ + _POINTER, ctypes.POINTER(wintypes.WORD), ctypes.POINTER(wintypes.DWORD)] +_security.GetSecurityDescriptorControl.restype = wintypes.BOOL +_security.GetSecurityInfo.argtypes = [wintypes.HANDLE, ctypes.c_int, wintypes.DWORD, + ctypes.POINTER(_POINTER), ctypes.POINTER(_POINTER), + ctypes.POINTER(_POINTER), ctypes.POINTER(_POINTER), + ctypes.POINTER(_POINTER)] +_security.GetSecurityInfo.restype = wintypes.DWORD +_security.SetSecurityInfo.argtypes = [wintypes.HANDLE, ctypes.c_int, wintypes.DWORD, + _POINTER, _POINTER, _POINTER, _POINTER] +_security.SetSecurityInfo.restype = wintypes.DWORD +_security.IsValidSid.argtypes = [_POINTER] +_security.IsValidSid.restype = wintypes.BOOL +_security.GetLengthSid.argtypes = [_POINTER] +_security.GetLengthSid.restype = wintypes.DWORD +_security.EqualSid.argtypes = [_POINTER, _POINTER] +_security.EqualSid.restype = wintypes.BOOL +_security.GetAce.argtypes = [_POINTER, wintypes.DWORD, ctypes.POINTER(_POINTER)] +_security.GetAce.restype = wintypes.BOOL + + +def _error(): + return ctypes.WinError(ctypes.get_last_error()) + + +def _name(name): + if (not isinstance(name, str) or name in ('', '.', '..') + or any(character in name for character in '/\\:\0') + or name.endswith((' ', '.'))): + raise ValueError('Private storage requires one ordinary Windows path component.') + stem = name.split('.')[0].upper() + if stem in {'CON', 'PRN', 'AUX', 'NUL', 'CONIN$', 'CONOUT$'} or ( + len(stem) == 4 and stem[:3] in ('COM', 'LPT') and stem[3] in '123456789¹²³'): + raise ValueError('Windows device names cannot be private storage entries.') + return name + + +def _absolute(path): + value = os.path.abspath(os.path.expanduser(os.fspath(path))) + if value.startswith(('\\\\?\\', '\\\\.\\')): + raise ValueError('Windows device namespace paths cannot be private job storage.') + result = Path(value) + if not result.anchor: + raise ValueError('Private storage requires an absolute directory.') + for component in result.parts[1:]: + _name(component) + if result == Path(result.anchor): + raise ValueError('A drive or share root cannot be private job storage.') + return result + + +def _close(handle): + if handle is not None and handle != _INVALID_HANDLE: + _kernel.CloseHandle(handle) + + +def _open(path, directory, *, repair=True): + access = _READ_CONTROL | _READ_ATTRIBUTES + if directory: + access |= _LIST_DIRECTORY + if repair: + access |= _WRITE_DAC + # Retained directory handles prevent pathname replacement. File metadata + # handles allow journal deletion while still binding repairs to the object. + share = _SHARE_READ_WRITE | (0 if directory else _SHARE_DELETE) + handle = _kernel.CreateFileW(str(path), access, share, None, 3, + _OPEN_REPARSE_POINT | _BACKUP_SEMANTICS, None) + if handle == _INVALID_HANDLE: + raise _error() + return handle + + +def _information(handle): + result = _FileInformation() + if not _kernel.GetFileInformationByHandle(handle, ctypes.byref(result)): + raise _error() + return result + + +def _identity(information): + return information.volume, (information.index_high << 32) | information.index_low + + +@contextmanager +def _security_info(handle): + owner, dacl, descriptor = _POINTER(), _POINTER(), _POINTER() + code = _security.GetSecurityInfo(handle, 1, _OWNER_INFO | _DACL_INFO, + ctypes.byref(owner), None, ctypes.byref(dacl), + None, ctypes.byref(descriptor)) + if code: + raise ctypes.WinError(code) + try: + yield owner, dacl, descriptor + finally: + _kernel.LocalFree(descriptor) + + +class _Security: + def __init__(self): + self._descriptors = [] + self._owned_sids = [] + token = wintypes.HANDLE() + if not _security.OpenProcessToken(_kernel.GetCurrentProcess(), 8, ctypes.byref(token)): + raise _error() + try: + size = wintypes.DWORD() + _security.GetTokenInformation(token, 1, None, 0, ctypes.byref(size)) + if not ctypes.sizeof(_TokenUser) <= size.value <= 65536: + raise OSError(errno.ENOTSUP, 'Cannot obtain the Windows process owner SID.') + self._user = ctypes.create_string_buffer(size.value) + if not _security.GetTokenInformation(token, 1, self._user, size, ctypes.byref(size)): + raise _error() + self.sid = ctypes.cast(self._user, ctypes.POINTER(_TokenUser)).contents.User.Sid + size = wintypes.DWORD() + _security.GetTokenInformation(token, 4, None, 0, ctypes.byref(size)) + if not ctypes.sizeof(_POINTER) <= size.value <= 65536: + raise OSError(errno.ENOTSUP, 'Cannot obtain the Windows process default owner SID.') + self._owner = ctypes.create_string_buffer(size.value) + if not _security.GetTokenInformation(token, 4, self._owner, size, ctypes.byref(size)): + raise _error() + self.owner_sid = ctypes.cast(self._owner, ctypes.POINTER(_POINTER)).contents.value + finally: + _close(token) + if not _security.IsValidSid(self.sid) or not _security.IsValidSid(self.owner_sid): + raise ValueError('A Windows process ownership SID is invalid.') + string = _POINTER() + if not _security.ConvertSidToStringSidW(self.sid, ctypes.byref(string)): + raise _error() + try: + self.sid_text = ctypes.wstring_at(string) + finally: + _kernel.LocalFree(string) + try: + self.system_sid = self._sid('S-1-5-18') + self.administrators_sid = self._sid('S-1-5-32-544') + self.directory = self._descriptor('OICI') + self.file = self._descriptor('') + except BaseException: + self.close() + raise + + def _sid(self, text): + sid = _POINTER() + if not _security.ConvertStringSidToSidW(text, ctypes.byref(sid)): + raise _error() + self._owned_sids.append(sid) + if not _security.IsValidSid(sid): + raise ValueError('A well-known Windows security SID is invalid.') + return sid + + def _owner_allowed(self, owner): + if _security.EqualSid(owner, self.sid): + return True + return bool(_security.EqualSid(owner, self.owner_sid) and any( + _security.EqualSid(self.owner_sid, privileged) + for privileged in (self.system_sid, self.administrators_sid) + )) + + def _descriptor(self, inheritance): + descriptor = _POINTER() + sddl = f'O:{self.sid_text}D:P(A;{inheritance};FA;;;{self.sid_text})' + if not _security.ConvertStringSecurityDescriptorToSecurityDescriptorW(sddl, 1, ctypes.byref(descriptor), None): + raise _error() + self._descriptors.append(descriptor) + present, defaulted, dacl = wintypes.BOOL(), wintypes.BOOL(), _POINTER() + if not _security.GetSecurityDescriptorDacl(descriptor, ctypes.byref(present), + ctypes.byref(dacl), ctypes.byref(defaulted)): + raise _error() + if not present or not dacl: + raise ValueError('An owner-only Windows DACL could not be constructed.') + attributes = _SecurityAttributes(ctypes.sizeof(_SecurityAttributes), descriptor, False) + return attributes, dacl + + def validate(self, handle, directory, *, owner=True): + information = _information(handle) + if (_kernel.GetFileType(handle) != 1 or information.attributes & _REPARSE + or bool(information.attributes & _DIRECTORY) != directory + or (not directory and information.links != 1)): + raise ValueError('Private job storage requires ordinary files and directories without reparse points or hard links.') + if owner: + with _security_info(handle) as (actual_owner, unused, descriptor): + # Elevated tokens give SQLite journals an administrator owner. + # A non-admin group owner would give other ordinary members + # implicit WRITE_DAC even after an owner-only DACL repair. + if (not actual_owner or not _security.IsValidSid(actual_owner) + or not self._owner_allowed(actual_owner)): + raise ValueError('Private job storage must be owned by the Windows account or its exact administrator process default owner.') + return _identity(information) + + @staticmethod + def _unsafe_legacy_directory(): + from .private_storage import PrivateStorageError + return PrivateStorageError( + 'Existing Windows job storage is writable by another ordinary account, ' + 'or its ACL cannot be safely inspected. Stop all processes using the store, ' + 'move it to an account-owned local directory and have its owner or administrator ' + 'remove other accounts\' write/delete/change-permission access throughout the tree ' + 'before retrying. Automatic ACL repair was refused; no existing ACL was changed.', + code='private_storage_requires_quiescent_repair', kind='permission_denied', + ) + + def preflight_directory(self, handle): + self.validate(handle, True) + with _security_info(handle) as (owner, dacl, descriptor): + if not dacl: + raise self._unsafe_legacy_directory() + header = ctypes.cast(dacl, ctypes.POINTER(_AclHeader)).contents + if header.revision not in (2, 4): + raise self._unsafe_legacy_directory() + for index in range(header.count): + ace = _POINTER() + if not _security.GetAce(dacl, index, ctypes.byref(ace)): + raise self._unsafe_legacy_directory() + entry = ctypes.cast(ace, ctypes.POINTER(_AceHeader)).contents + if entry.kind not in (0, 1) or entry.flags & ~0x1F or entry.size < 16: + raise self._unsafe_legacy_directory() + sid = _POINTER(ace.value + 8) + if not _security.IsValidSid(sid) or 8 + _security.GetLengthSid(sid) > entry.size: + raise self._unsafe_legacy_directory() + mask = ctypes.c_uint32.from_address(ace.value + 4).value + if mask & ~_KNOWN_FILE_RIGHTS: + raise self._unsafe_legacy_directory() + if entry.kind == 1 or entry.flags & 8: + continue + trusted = any(_security.EqualSid(sid, safe) for safe in ( + self.sid, self.system_sid, self.administrators_sid, + )) + # Conservatively refuse a foreign mutation allow ACE even if + # a separate deny ACE might narrow its effective rights. + if not trusted and mask & _MUTATION_RIGHTS: + raise self._unsafe_legacy_directory() + + def assert_private(self, handle, directory): + self.validate(handle, directory) + with _security_info(handle) as (owner, dacl, descriptor): + control, revision = wintypes.WORD(), wintypes.DWORD() + if not _security.GetSecurityDescriptorControl(descriptor, ctypes.byref(control), ctypes.byref(revision)): + raise _error() + if not dacl or not control.value & _DACL_PROTECTED or ctypes.cast(dacl, ctypes.POINTER(_AclHeader)).contents.count != 1: + raise OSError(errno.ENOTSUP, 'The filesystem did not enforce the protected owner-only Windows DACL.') + ace = _POINTER() + if not _security.GetAce(dacl, 0, ctypes.byref(ace)): + raise _error() + header = ctypes.cast(ace, ctypes.POINTER(_AceHeader)).contents + mask = ctypes.c_uint32.from_address(ace.value + 4).value + sid = _POINTER(ace.value + 8) + expected_flags = 3 if directory else 0 + if (header.kind != 0 or header.flags != expected_flags or mask != _FILE_ALL_ACCESS + or not _security.IsValidSid(sid) or not _security.EqualSid(sid, self.sid)): + raise OSError(errno.ENOTSUP, 'The filesystem did not retain the owner-only Windows access rule.') + + def protect(self, handle, directory): + self.validate(handle, directory) + dacl = (self.directory if directory else self.file)[1] + code = _security.SetSecurityInfo(handle, 1, _DACL_INFO | _PROTECTED_DACL_INFO, + None, None, dacl, None) + if code: + raise ctypes.WinError(code) + self.assert_private(handle, directory) + + def close(self): + for descriptor in self._descriptors: + _kernel.LocalFree(descriptor) + self._descriptors.clear() + for sid in self._owned_sids: + _kernel.LocalFree(sid) + self._owned_sids.clear() + + +class PrivateDirectory: + def __init__(self, path): + self._handles = [] + self._handle = None + self._security = None + self.path = _absolute(path) + try: + self._security = _Security() + current = Path(self.path.anchor) + anchor = _open(current, True, repair=False) + self._handles.append(anchor) + self._security.validate(anchor, True, owner=False) + for component in self.path.parts[1:]: + current /= component + created = False + try: + handle = _open(current, True, repair=current == self.path) + except FileNotFoundError: + attributes = self._security.directory[0] + if _kernel.CreateDirectoryW(str(current), ctypes.byref(attributes)): + created = True + else: + error = _error() + if error.winerror != 183: + raise error + handle = _open(current, True, repair=current == self.path) + self._handles.append(handle) + self._security.validate(handle, True, owner=created or current == self.path) + if created: + self._security.assert_private(handle, True) + self._handle = self._handles[-1] + self.identity = self._security.validate(self._handle, True) + self._repair_tree() + self.verify() + except BaseException: + self.close() + raise + + def _repair_tree(self): + entries = [(self._handle, True)] + retained = [] + try: + def inspect(path): + with os.scandir(path) as children: + for child in children: + try: + _name(child.name) + directory = child.is_dir(follow_symlinks=False) + handle = _open(child.path, directory) + except FileNotFoundError: + continue + retained.append(handle) + self._security.validate(handle, directory) + entries.append((handle, directory)) + if directory: + inspect(Path(child.path)) + inspect(self.path) + for handle, directory in entries: + if directory: + self._security.preflight_directory(handle) + # Validate the entire legacy tree before changing any existing ACL; + # protected children are repaired before inheritable parent rules. + for handle, directory in reversed(entries): + self._security.protect(handle, directory) + finally: + for handle in reversed(retained): + _close(handle) + + def verify(self): + if self._handle is None: + raise ValueError('The private job directory is closed.') + handle = _open(self.path, True, repair=False) + try: + if self._security.validate(handle, True) != self.identity: + raise ValueError('The private job directory changed; restore its original location before retrying.') + finally: + _close(handle) + self._security.assert_private(self._handle, True) + + def child(self, name): + self.verify() + return PrivateDirectory(self.path / _name(name)) + + def open_file(self, name, *, exclusive=False): + self.verify() + target = self.path / _name(name) + attributes = self._security.file[0] + handle = _kernel.CreateFileW(str(target), 0xC0000000 | _READ_CONTROL | _WRITE_DAC, + _SHARE_READ_WRITE, ctypes.byref(attributes), + 1 if exclusive else 4, _OPEN_REPARSE_POINT, None) + if handle == _INVALID_HANDLE: + raise _error() + try: + self._security.protect(handle, False) + descriptor = msvcrt.open_osfhandle(handle, os.O_RDWR | os.O_BINARY) + handle = None + try: + return os.fdopen(descriptor, 'r+b', buffering=0) + except BaseException: + os.close(descriptor) + raise + finally: + _close(handle) + + def duplicate_fd(self): + raise OSError(errno.ENOTSUP, 'Windows private directory handles are not POSIX directory descriptors.') + + def close(self): + for handle in reversed(getattr(self, '_handles', [])): + _close(handle) + self._handles = [] + self._handle = None + if getattr(self, '_security', None) is not None: + self._security.close() + self._security = None + + def __enter__(self): + return self + + def __exit__(self, *unused): + self.close() + + def __del__(self): + self.close() diff --git a/src/ucsfomopagent/job_worker.py b/src/ucsfomopagent/job_worker.py index 6602812..8836b72 100644 --- a/src/ucsfomopagent/job_worker.py +++ b/src/ucsfomopagent/job_worker.py @@ -8,7 +8,8 @@ import threading import time -from .jobs import JobStore, ACTIVE +from .jobs import JobStore, ACTIVE, failure_details, record_failure +from .private_storage import private_storage_call def execute(store, job_id, payload, stream): @@ -23,18 +24,20 @@ def execute(store, job_id, payload, stream): return data partial = directory / ('result.' + request['format'] + '.partial') final = directory / ('result.' + request['format']) - progress = {'phase': 'connecting', 'rows': 0, 'bytes': 0} + progress = {'phase': 'opening_output', 'rows': 0, 'bytes': 0} done = threading.Event() failure = [] def run(): try: - with partial.open('xb') as out: - os.chmod(partial, 0o600) + with private_storage_call(lambda: store._storage.child(job_id)) as private_output, private_storage_call( + lambda: private_output.open_file(partial.name, exclusive=True) + ) as out: csv_buffer = io.StringIO(newline='') writer = csv.writer(csv_buffer) if request['format'] == 'csv' else None header = False def write_row(row): + progress['phase'] = 'writing_output' if writer: csv_buffer.seek(0); csv_buffer.truncate(0) writer.writerow([json.dumps(v, default=str, ensure_ascii=False) if isinstance(v, (dict, list)) else v for v in row]) @@ -45,6 +48,7 @@ def write_row(row): raise OverflowError('output budget') out.write(encoded) progress['bytes'] += len(encoded) + progress['phase'] = 'connecting' for columns, batch in stream(request, payload['credentials'], progress): progress.setdefault('stream_started_at', time.time()) progress['phase'] = 'streaming' @@ -55,18 +59,11 @@ def write_row(row): write_row(row if writer else dict(zip(columns, row))) progress['rows'] += 1 out.flush() + progress['phase'] = 'finalizing' out.flush() os.fsync(out.fileno()) except Exception as error: - name = type(error).__name__ - if isinstance(error, OverflowError): - failure.append('Output byte budget exceeded; narrow the export or explicitly raise max_bytes.') - elif 'Auth' in name: - failure.append('Database authentication failed. Reconfigure credentials through the trusted UI or CLI.') - elif 'Timeout' in name: - failure.append('Database timeout. Inspect the plan and filters before explicitly resubmitting.') - else: - failure.append('Database or output operation failed (' + name + '). Check credentials, connectivity, schema, permissions and disk space; raw errors are withheld to protect data.') + failure.append(failure_details(error, progress['phase'])) finally: done.set() @@ -77,14 +74,17 @@ def write_row(row): if (directory / 'cancel').exists(): data.update(state='cancelled', phase='cancelled') elif time.monotonic() - started >= request['timeout_seconds']: - data.update(state='timed_out', phase='timed_out', error='Job wall-clock budget expired. Inspect the plan before explicitly resubmitting with a larger budget.') + record_failure(data, dict(code='job_deadline_exceeded', kind='timeout', message='Job wall-clock budget expired. Inspect the plan before explicitly resubmitting with a larger budget.', phase=progress['phase'], retryable=False), state='timed_out') elif done.is_set(): data.update(progress) if failure: - data.update(state='failed', phase='failed', error=failure[0]) + record_failure(data, failure[0]) else: - partial.replace(final) - data.update(state='completed', phase='completed', result_path=str(final), bytes=final.stat().st_size) + try: + partial.replace(final) + data.update(state='completed', phase='completed', result_path=str(final), bytes=final.stat().st_size) + except Exception as error: + record_failure(data, failure_details(error, 'finalizing')) if data['state'] != 'running': data['finished_at'] = time.time() if not store.update(data): @@ -107,7 +107,7 @@ def main(): execute(store, job_id, payload, backend.stream) except Exception: data = store.status(job_id) - data.update(state='failed', phase='failed', error='Worker initialization failed.', finished_at=time.time()) + record_failure(data, dict(code='worker_initialization_failed', kind='internal', message='Worker initialization failed.', phase='worker_initialization', retryable=False)) store.update(data) # Keep the OS lease until exit kills the driver thread too. Releasing it while # a blocked thread can still write would let purge race against an active file. diff --git a/src/ucsfomopagent/jobs.py b/src/ucsfomopagent/jobs.py index 6b18080..58b93c9 100644 --- a/src/ucsfomopagent/jobs.py +++ b/src/ucsfomopagent/jobs.py @@ -1,13 +1,17 @@ -"""Durable, bounded-memory query jobs. Kept identical in the three UCSF connectors.""" +"""Durable, bounded-memory query jobs shared by the OMOP and CDW connectors.""" from __future__ import annotations -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager, suppress +import asyncio +import codecs +import errno import hashlib import json import os from pathlib import Path import re import sqlite3 +import stat import statistics import subprocess import sys @@ -15,9 +19,88 @@ import time import uuid from filelock import FileLock, Timeout +from mcp.types import CallToolResult + +from .private_storage import PrivateStorageError, private_directory, private_storage_call ACTIVE = ('queued', 'running') TERMINAL = ('completed', 'failed', 'cancelled', 'timed_out', 'interrupted') +MAX_ACTIVE_JOBS = 2 +MAX_WAIT_SECONDS = 30 +WAIT_LOOKUP_SECONDS = 0.05 +MAX_RESULT_BYTES = 65536 +OPERATION_META_KEY = 'io.biorouter/operation' + + +class JobError(ValueError): + def __init__(self, code, kind, message, phase, retryable=False, **context): + super().__init__(message) + self.details = dict(code=code, kind=kind, message=message, phase=phase, retryable=retryable) + self.context = context + + +def failure_details(error, phase): + if isinstance(error, JobError): + return error.details + if isinstance(error, PrivateStorageError): + return dict(code=error.code, kind=error.kind, message=str(error), phase=phase, retryable=False) + code, kind, message, retryable = 'query_failed', 'tool_failure', 'Database execution or result processing failed. Check the query plan, schema and selected output format.', False + if isinstance(error, OverflowError): + code, message = 'output_limit_exceeded', 'Output byte budget exceeded; narrow the export or explicitly raise max_bytes.' + elif isinstance(error, TimeoutError): + code, kind, message = 'database_timeout', 'timeout', 'Database timeout. Inspect the plan and filters before explicitly resubmitting.' + elif isinstance(error, OSError): + if error.errno == errno.ENOSPC: + code, message = 'disk_full', 'Local result storage is full. Free space before explicitly resubmitting.' + elif error.errno in (errno.EACCES, errno.EPERM): + code, kind, message = 'storage_permission_denied', 'permission_denied', 'Local result storage access was denied.' + elif isinstance(error, ConnectionError): + code, kind, message, retryable = 'connection_interrupted', 'transient', 'The database connection was interrupted.', True + else: + code, message = 'output_io_failed', 'Local result file processing failed. Check available space and file permissions.' + elif type(error).__module__.split('.')[0] == 'pymssql': + number = error.args[0] if error.args and type(error.args[0]) is int else None + if number == 20003: + code, kind, message = 'database_timeout', 'timeout', 'Database timeout. Inspect the plan and filters before explicitly resubmitting.' + elif number == 18456: + code, kind, message = 'authentication_failed', 'permission_denied', 'Database authentication failed. Reconfigure credentials through the trusted UI or CLI.' + elif number == 229: + code, kind, message = 'database_permission_denied', 'permission_denied', 'The database denied this operation. Check read or SHOWPLAN permissions.' + elif number in (1205, 1222): + code, kind, message, retryable = 'database_busy', 'transient', 'The database could not obtain the required locks.', True + elif type(error).__name__ == 'ProgrammingError': + code, kind, message = 'query_rejected', 'invalid_args', 'The database rejected the query. Check schema, projections and parameters.' + elif phase == 'connecting': + code, message = 'connection_failed', 'Database connection failed. Check connectivity and the configured authentication route.' + return dict(code=code, kind=kind, message=message, phase=phase, retryable=retryable) + + +def record_failure(data, details, state='failed'): + data.update(state=state, phase=state, error=details['message'], error_details=details, finished_at=time.time()) + + +def operation_metadata(package, data): + state = data['state'] + status = 'failed' if state in ('failed', 'timed_out', 'interrupted') else state + operation = dict(version=1, id=package + ':' + data['job_id'], status=status) + if status == 'failed': + operation['error'] = data.get('error_details') or dict(code='job_' + state, kind='tool_failure', message='The job did not complete. Inspect its status before explicitly submitting a replacement.', phase=state, retryable=False) + return operation + + +def tool_result(data, package=None): + from mcp.types import TextContent + metadata = {OPERATION_META_KEY: operation_metadata(package, data)} if package else None + return CallToolResult(content=[TextContent(type='text', text=json.dumps(data))], structuredContent=data, isError=False, _meta=metadata) + + +def tool_failure(error, phase='request'): + from mcp.types import TextContent + details = failure_details(error, phase) + data = dict(error=details, **(error.context if isinstance(error, JobError) else {})) + return CallToolResult(content=[TextContent(type='text', text=details['message'])], structuredContent=data, isError=True) + + INSTRUCTIONS = """ For schema discovery and small previews use the existing synchronous tools. For exports, joins/scans with uncertain cost, or a query that timed out, use @@ -26,16 +109,21 @@ Resolve schema/concepts first; select only needed columns and use indexed filters. Use mode='explain' to inspect a plan without executing the query (SQL Server needs SHOWPLAN permission). Do not run an expensive COUNT(*) solely to estimate export size. -Poll query_job_status at its recommended_poll_seconds until a terminal state. +Check query_job_capacity before submission. Use wait_for_query_job for bounded native +waiting (up to 30 seconds per call); query_job_status remains an immediate lookup. Report phase, elapsed time, rows and bytes. An unknown ETA is not a stuck query; execution before the first row often has no measurable percentage. Estimates from prior identical queries or expected_rows are advisory, never guarantees. Only completed means the result file is complete; partial files are incomplete. Use cancel_query_job to stop unwanted work. Do not resubmit while a job is active. -Do not repeatedly retry failures: inspect schema/plan and narrow the query first. +Inspect error_details.code, phase and retryable before explicitly resubmitting. +Fix credentials, connectivity, permissions or disk conditions rather than blindly +rewriting SQL. Revise the query only when its schema, syntax or plan is the cause. Jobs survive chat/MCP disconnects, but not host reboot; interrupted queries require an explicit new submission. There is no unsafe automatic OFFSET resume or retry. -Process downloaded data locally and return summaries instead of full rows to chat. +Use read_query_job_result with the completed job ID for bounded UTF-8 chunks; +continue with next_offset until eof. Do not pass result_path to filesystem tools. +For large analyses process the complete local file and return summaries to chat. """ @@ -80,29 +168,25 @@ def strip(m): def private_dir(path: Path) -> Path: - if path.is_symlink(): - raise ValueError('Job directories must not be symbolic links.') - path.mkdir(parents=True, exist_ok=True, mode=0o700) - path.chmod(0o700) - return path + with private_directory(path) as directory: + return directory.path class JobStore: def __init__(self, package: str, root: str | Path | None = None): self.package = package - self.root = private_dir(Path(root or os.environ.get('UCSF_AGENT_JOB_ROOT', Path.home() / '.local/share/ucsf-agents')) / package) + self._storage = private_directory(Path(root or os.environ.get('UCSF_AGENT_JOB_ROOT', Path.home() / '.local/share/ucsf-agents')) / package) + self.root = self._storage.path + self._root_identity = self._storage.identity self.dbpath = self.root / 'jobs.sqlite3' - if self.dbpath.is_symlink(): - raise ValueError('Job database must not be a symbolic link.') - fd = os.open(self.dbpath, os.O_CREAT | os.O_RDWR, 0o600) - os.close(fd) with self.connect() as db: db.execute('CREATE TABLE IF NOT EXISTS jobs (id TEXT PRIMARY KEY, state TEXT, updated REAL, fingerprint TEXT, data TEXT)') - self.dbpath.chmod(0o600) @contextmanager - def connect(self): - db = sqlite3.connect(self.dbpath, timeout=10) + def connect(self, timeout=10): + with private_storage_call(lambda: self._storage.open_file('jobs.sqlite3')): + pass + db = sqlite3.connect(self.dbpath, timeout=timeout) try: with db: yield db @@ -111,7 +195,7 @@ def connect(self): def directory(self, job_id): if not re.fullmatch(r'[0-9a-f]{32}', job_id): - raise ValueError('Invalid job ID.') + raise JobError('invalid_job_id', 'invalid_args', 'Invalid job ID.', 'lookup') return self.root / job_id def lease(self, job_id): @@ -124,12 +208,21 @@ def update(self, data): (data['state'], data['updated_at'], json.dumps(data), data['job_id'])) return cursor.rowcount == 1 - def status(self, job_id): + def status(self, job_id, *, database_timeout=10): + deadline = time.monotonic() + database_timeout + try: + return self._status(job_id, deadline) + except sqlite3.OperationalError as error: + if (getattr(error, 'sqlite_errorcode', 0) & 255) not in (sqlite3.SQLITE_BUSY, sqlite3.SQLITE_LOCKED): + raise + raise JobError('status_lookup_busy', 'transient', 'Job status storage is busy. Retry the lookup after existing work releases it.', 'status_lookup', retryable=True) from None + + def _status(self, job_id, deadline): self.directory(job_id) - with self.connect() as db: + with self.connect(timeout=max(0, deadline - time.monotonic())) as db: row = db.execute('SELECT data FROM jobs WHERE id=?', (job_id,)).fetchone() if not row: - raise ValueError('Unknown job ID.') + raise JobError('job_not_found', 'not_found', 'Unknown job ID.', 'lookup') data = json.loads(row[0]) now = time.time() data['heartbeat_stale'] = data['state'] in ACTIVE and now - data['updated_at'] > 60 @@ -137,12 +230,12 @@ def status(self, job_id): try: with self.lease(job_id).acquire(timeout=0): previous_update = data['updated_at'] - data = dict(data, state='interrupted', phase='interrupted', error='Worker exited without completing. Inspect before explicitly submitting a new job.', finished_at=now) - with self.connect() as db: + record_failure(data, dict(code='worker_interrupted', kind='internal', message='Worker exited without completing. Inspect before explicitly submitting a new job.', phase=data.get('phase', 'unknown'), retryable=False), state='interrupted') + with self.connect(timeout=max(0, deadline - time.monotonic())) as db: cursor = db.execute("UPDATE jobs SET state=?, data=? WHERE id=? AND updated=? AND state IN ('queued','running')", ('interrupted', json.dumps(data), job_id, previous_update)) if not cursor.rowcount: - return self.status(job_id) + return self._status(job_id, deadline) except Timeout: data['heartbeat_warning'] = 'Worker still holds its process lease but heartbeat is stale; do not resubmit.' data['elapsed_seconds'] = round((data.get('finished_at') or now) - data['created_at'], 2) @@ -159,8 +252,87 @@ def status(self, job_id): elif data.get('historical_seconds') and data['elapsed_seconds'] < data['historical_seconds']: data['eta_seconds'] = round(data['historical_seconds'] - data['elapsed_seconds'], 1) data['estimate_basis'] = 'median of prior successful identical queries; load and data may differ' + data['capacity'] = self.capacity(database_timeout=max(0, deadline - time.monotonic())) return data + def capacity(self, *, database_timeout=10): + with self.connect(timeout=database_timeout) as db: + active = db.execute("SELECT COUNT(*) FROM jobs WHERE state IN ('queued','running')").fetchone()[0] + return dict(active_jobs=active, max_active_jobs=MAX_ACTIVE_JOBS, available_slots=max(0, MAX_ACTIVE_JOBS - active), advisory=True) + + async def wait(self, job_id, timeout_seconds=30): + if type(timeout_seconds) is not int or not 0 <= timeout_seconds <= MAX_WAIT_SECONDS: + raise JobError('invalid_wait', 'invalid_args', 'timeout_seconds must be an integer from 0 to 30.', 'validation') + deadline = time.monotonic() + timeout_seconds + data = None + while True: + remaining = max(0, deadline - time.monotonic()) + try: + data = await asyncio.to_thread(self.status, job_id, database_timeout=min(WAIT_LOOKUP_SECONDS, remaining)) + except JobError as error: + if error.details['code'] != 'status_lookup_busy': + raise + remaining = deadline - time.monotonic() + if remaining <= 0: + if data is None: + raise + data.update(wait_timed_out=True, wait_status_stale=True, wait_warning='The wait ended with status storage busy; this is the last successfully observed active state.') + return data + await asyncio.sleep(min(0.25, remaining)) + continue + remaining = deadline - time.monotonic() + if data['state'] in TERMINAL or remaining <= 0: + data['wait_timed_out'] = data['state'] not in TERMINAL + return data + await asyncio.sleep(min(0.25, remaining)) + + def read_result(self, job_id, offset=0, max_bytes=8192): + if type(offset) is not int or offset < 0 or type(max_bytes) is not int or not 4 <= max_bytes <= MAX_RESULT_BYTES: + raise JobError('invalid_result_chunk', 'invalid_args', 'offset must be a nonnegative byte offset; max_bytes must be 4..65536.', 'result_read') + try: + self._storage.verify() + except (OSError, ValueError): + raise JobError('result_scope_changed', 'permission_denied', 'The managed job directory changed or cannot be secured; result access was refused.', 'result_read') from None + data = self.status(job_id) + if data['state'] != 'completed': + raise JobError('result_not_complete', 'invalid_args', 'Only completed job results can be read.', 'result_read') + if data.get('format') not in ('csv', 'jsonl'): + raise JobError('invalid_result_format', 'internal', 'The saved result format is invalid.', 'result_read') + if not hasattr(os, 'O_NOFOLLOW') or not hasattr(os, 'O_DIRECTORY') or os.open not in os.supports_dir_fd: + raise JobError('secure_reader_unavailable', 'internal', 'This platform does not support the secure result reader; use the local CLI result file.', 'result_read') + try: + with ExitStack() as stack: + root_fd = self._storage.duplicate_fd() + stack.callback(os.close, root_fd) + root_stat = os.fstat(root_fd) + if (root_stat.st_dev, root_stat.st_ino) != self._root_identity: + raise JobError('result_scope_changed', 'permission_denied', 'The managed job directory changed; result access was refused.', 'result_read') + directory_fd = os.open(job_id, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=root_fd) + stack.callback(os.close, directory_fd) + result_fd = os.open('result.' + data['format'], os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=directory_fd) + stack.callback(os.close, result_fd) + result_stat = os.fstat(result_fd) + if not stat.S_ISREG(result_stat.st_mode) or result_stat.st_nlink != 1: + raise JobError('invalid_result_file', 'permission_denied', 'The result must be a regular managed file without links.', 'result_read') + size = result_stat.st_size + if size != data.get('bytes'): + raise JobError('result_size_changed', 'internal', 'The completed result size changed; access was refused.', 'result_read') + if offset > size: + raise JobError('invalid_result_offset', 'invalid_args', 'The byte offset exceeds the completed result size.', 'result_read') + os.lseek(result_fd, offset, os.SEEK_SET) + raw = os.read(result_fd, min(max_bytes, size - offset)) + if len(raw) != min(max_bytes, size - offset): + raise JobError('result_size_changed', 'internal', 'The completed result changed while it was being read.', 'result_read') + decoder = codecs.getincrementaldecoder('utf-8')() + text = decoder.decode(raw, final=offset + len(raw) == size) + consumed = len(raw) - len(decoder.getstate()[0]) + except UnicodeDecodeError: + raise JobError('invalid_result_encoding', 'invalid_args', 'The result or byte offset is not valid UTF-8.', 'result_read') from None + except OSError: + raise JobError('result_unavailable', 'permission_denied', 'The managed result file is missing, linked or inaccessible.', 'result_read') from None + next_offset = offset + consumed + return dict(job_id=job_id, format=data['format'], text=text, offset=offset, bytes_read=consumed, total_bytes=size, eof=next_offset == size, next_offset=None if next_offset == size else next_offset) + def list_jobs(self): with self.connect() as db: ids = db.execute('SELECT id FROM jobs ORDER BY updated DESC LIMIT 50').fetchall() @@ -186,7 +358,9 @@ def purge(self, job_id): def cancel(self, job_id): data = self.status(job_id) if data['state'] in ACTIVE: - (self.directory(job_id) / 'cancel').touch(mode=0o600) + with private_storage_call(lambda: self._storage.child(job_id)) as directory: + with private_storage_call(lambda: directory.open_file('cancel')): + pass data['cancel_requested'] = True return data @@ -214,45 +388,99 @@ def submit(self, query, parameters, credentials, *, mode='query', format='jsonl' with self.connect() as db: db.execute('BEGIN IMMEDIATE') # Stale workers retain their slot until the operator checks their status. - if db.execute("SELECT COUNT(*) FROM jobs WHERE state IN ('queued','running')").fetchone()[0] >= 2: - raise ValueError('Two jobs are already active. Monitor or cancel them before submitting another.') + active = db.execute("SELECT COUNT(*) FROM jobs WHERE state IN ('queued','running')").fetchone()[0] + if active >= MAX_ACTIVE_JOBS: + raise JobError('capacity_exceeded', 'transient', 'Two jobs are already active. Wait for completion or cancel unwanted work before submitting another.', 'admission', retryable=True, capacity=dict(active_jobs=active, max_active_jobs=MAX_ACTIVE_JOBS, available_slots=0, advisory=True)) history = db.execute("SELECT data FROM jobs WHERE fingerprint=? AND state='completed' ORDER BY updated DESC LIMIT 5", (fingerprint,)).fetchall() if history: data['historical_seconds'] = statistics.median(json.loads(r[0])['finished_at'] - json.loads(r[0])['created_at'] for r in history) - private_dir(self.directory(job_id)) + with private_storage_call(lambda: self._storage.child(job_id)): + pass db.execute('INSERT INTO jobs VALUES (?,?,?,?,?)', (job_id, 'queued', now, fingerprint, json.dumps(data))) - # Requests contain potentially sensitive query parameters, so only the status - # goes to disk. Credentials and query text reach the worker over a private pipe. + try: + self._start_worker(job_id, request, credentials) + except Exception: + details = dict(code='worker_start_failed', kind='internal', message='Could not start the local worker.', phase='worker_start', retryable=False) + record_failure(data, details) + self.update(data) + return self.status(job_id) + + def _start_worker(self, job_id, request, credentials): + # Keep query text and credentials off disk and out of child arguments. env = {k: v for k, v in os.environ.items() if k in ('PATH', 'HOME', 'USERPROFILE', 'SYSTEMROOT', 'WINDIR', 'TMPDIR', 'TEMP', 'TMP', 'LANG', 'LC_ALL', 'PYTHONPATH', 'SSL_CERT_FILE', 'SSL_CERT_DIR')} env['UCSF_AGENT_JOB_ROOT'] = str(self.root.parent) args = [sys.executable, '-m', f'{self.package}.job_worker', job_id] options = {'start_new_session': True} if os.name != 'nt' else {'creationflags': subprocess.CREATE_NEW_PROCESS_GROUP | subprocess.DETACHED_PROCESS} + wire = json.dumps(dict(request=request, credentials=credentials)).encode() + process = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, env=env, **options) try: - process = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, env=env, **options) - process.stdin.write(json.dumps(dict(request=request, credentials=credentials)).encode()) + process.stdin.write(wire) process.stdin.close() - threading.Thread(target=process.wait, daemon=True).start() except Exception: - data.update(state='failed', phase='failed', error='Could not start the local worker.', finished_at=time.time()) - self.update(data) - raise RuntimeError('Could not start the local worker.') from None - return self.status(job_id) + with suppress(OSError): + if process.stdin is not None: + process.stdin.close() + raise + finally: + with suppress(RuntimeError): + threading.Thread(target=process.wait, daemon=True).start() -def register_job_tools(mcp, package, credentials, prefix=''): +def register_job_tools(mcp, package, credentials, prefix='', *, store=None): from mcp.types import ToolAnnotations - store = JobStore(package) + storage_error = None + if store is None: + try: + store = JobStore(package) + except PrivateStorageError as error: + storage_error = error + + def call(action, operation=False): + try: + if storage_error is not None: + raise storage_error + data = action() + return tool_result(data, package if operation else None) + except ValueError as error: + if isinstance(error, PrivateStorageError): + return tool_failure(error, 'storage') + if not isinstance(error, JobError): + error = JobError('invalid_request', 'invalid_args', 'Invalid query or job arguments. Check the read-only query policy and declared argument bounds.', 'validation') + return tool_failure(error) + except Exception as error: + return tool_failure(error) + @mcp.tool(name=prefix + 'submit_query_job', annotations=ToolAnnotations(readOnlyHint=False, destructiveHint=False, idempotentHint=False, openWorldHint=True)) - def submit_query_job(query: str, parameters: dict | None = None, mode: str = 'query', format: str = 'jsonl', timeout_seconds: int = 3600, expected_rows: int | None = None, max_bytes: int = 10_000_000_000) -> dict: - """Start a detached read-only SQL/Cypher export, or mode=explain for a plan. Returns a job ID immediately; poll query_job_status. No result-row cap; max_bytes is a fail-visible disk budget. Credentials come from the connector, never tool arguments.""" - return store.submit(query, parameters, credentials, mode=mode, format=format, timeout_seconds=timeout_seconds, expected_rows=expected_rows, max_bytes=max_bytes) + def submit_query_job(query: str, parameters: dict | None = None, mode: str = 'query', format: str = 'jsonl', timeout_seconds: int = 3600, expected_rows: int | None = None, max_bytes: int = 10_000_000_000) -> CallToolResult: + """Start a detached read-only export or mode=explain plan. Check query_job_capacity first, then wait_for_query_job. Credentials come from the connector. No automatic retry or result-row cap.""" + return call(lambda: store.submit(query, parameters, credentials, mode=mode, format=format, timeout_seconds=timeout_seconds, expected_rows=expected_rows, max_bytes=max_bytes), operation=True) + + @mcp.tool(name=prefix + 'query_job_capacity', annotations=ToolAnnotations(readOnlyHint=True, destructiveHint=False, idempotentHint=True, openWorldHint=False)) + def query_job_capacity() -> CallToolResult: + """Read active_jobs, max_active_jobs and available_slots before submission. Capacity is advisory; admission is atomically checked again when submitting.""" + return call(lambda: store.capacity()) @mcp.tool(name=prefix + 'query_job_status', annotations=ToolAnnotations(readOnlyHint=True, destructiveHint=False, idempotentHint=True, openWorldHint=False)) - def query_job_status(job_id: str) -> dict: - """Read phase, heartbeat, rows, bytes, elapsed time, advisory ETA and completed file path. Poll at recommended_poll_seconds. No credentials needed; never reads result rows into chat.""" - return store.status(job_id) + def query_job_status(job_id: str) -> CallToolResult: + """Immediately read phase, heartbeat, rows, bytes, capacity and job outcome. Failed jobs include sanitized error_details. Use wait_for_query_job instead of rapid polling.""" + return call(lambda: store.status(job_id), operation=True) + + @mcp.tool(name=prefix + 'wait_for_query_job', annotations=ToolAnnotations(readOnlyHint=True, destructiveHint=False, idempotentHint=True, openWorldHint=False)) + async def wait_for_query_job(job_id: str, timeout_seconds: int = 30) -> CallToolResult: + """Wait natively for a terminal job state, for 0..30 seconds. wait_timed_out means this wait ended while the job remains active; it does not cancel or time out the job. Cancellation stops waiting only.""" + try: + if storage_error is not None: + raise storage_error + return tool_result(await store.wait(job_id, timeout_seconds), package) + except Exception as error: + return tool_failure(error) + + @mcp.tool(name=prefix + 'read_query_job_result', annotations=ToolAnnotations(readOnlyHint=True, destructiveHint=False, idempotentHint=True, openWorldHint=False)) + def read_query_job_result(job_id: str, offset: int = 0, max_bytes: int = 8192) -> CallToolResult: + """Read 4..65536 bytes of a completed managed result as UTF-8 text. Continue at next_offset until eof. Offsets are bytes, not rows; CSV rows may span chunks. Accepts job IDs, never arbitrary file paths.""" + return call(lambda: store.read_result(job_id, offset, max_bytes)) @mcp.tool(name=prefix + 'cancel_query_job', annotations=ToolAnnotations(readOnlyHint=False, destructiveHint=False, idempotentHint=True, openWorldHint=False)) - def cancel_query_job(job_id: str) -> dict: - """Request cancellation. Poll until cancelled/completed; closing the worker connection releases the query. Partial files are not complete results.""" - return store.cancel(job_id) + def cancel_query_job(job_id: str) -> CallToolResult: + """Request cancellation, then wait_for_query_job until terminal. Closing the worker releases its query connection; partial files remain incomplete.""" + return call(lambda: store.cancel(job_id), operation=True) diff --git a/src/ucsfomopagent/private_storage.py b/src/ucsfomopagent/private_storage.py new file mode 100644 index 0000000..cdc9dfc --- /dev/null +++ b/src/ucsfomopagent/private_storage.py @@ -0,0 +1,204 @@ +"""Owner-private storage for durable jobs; handles pin every permission repair.""" +from __future__ import annotations + +import ctypes +import errno +import os +from pathlib import Path +import stat +import sys + + +class PrivateStorageError(ValueError): + def __init__(self, message, code='private_storage_unavailable', kind='permission_denied'): + super().__init__(message) + self.code = code + self.kind = kind + + +def _name(name): + if not isinstance(name, str) or name in ('', '.', '..') or '/' in name or '\\' in name: + raise ValueError('Private storage requires one plain path component.') + return name + + +def _clear_acl(fd, directory): + if sys.platform == 'darwin': + library = ctypes.CDLL(None, use_errno=True) + library.acl_init.argtypes = [ctypes.c_int] + library.acl_init.restype = ctypes.c_void_p + library.acl_set_fd.argtypes = [ctypes.c_int, ctypes.c_void_p] + library.acl_set_fd.restype = ctypes.c_int + library.acl_free.argtypes = [ctypes.c_void_p] + library.acl_free.restype = ctypes.c_int + acl = library.acl_init(0) + if not acl: + raise OSError(ctypes.get_errno(), 'Cannot allocate an owner-private ACL.') + try: + if library.acl_set_fd(fd, acl) != 0: + raise OSError(ctypes.get_errno(), 'Cannot enforce an owner-private ACL.') + finally: + library.acl_free(acl) + elif sys.platform.startswith('linux'): + names = ['system.posix_acl_access'] + if directory: + names.append('system.posix_acl_default') + for name in names: + try: + os.removexattr(fd, name) + except OSError as error: + if error.errno not in (errno.ENODATA, errno.ENOTSUP): + raise + else: + raise OSError(errno.ENOTSUP, 'Owner-private ACL enforcement is unavailable on this platform.') + + +def _protect(fd, directory): + information = os.fstat(fd) + valid_type = stat.S_ISDIR(information.st_mode) if directory else stat.S_ISREG(information.st_mode) + if not valid_type or information.st_uid != os.geteuid() or (not directory and information.st_nlink != 1): + raise ValueError('Private job storage must contain only owned directories and regular files without links.') + mode = 0o700 if directory else 0o600 + os.fchmod(fd, mode) + _clear_acl(fd, directory) + if stat.S_IMODE(os.fstat(fd).st_mode) != mode: + raise OSError(errno.ENOTSUP, 'The filesystem does not enforce owner-private permissions.') + + +class PrivateDirectory: + def __init__(self, path): + if os.name != 'posix' or not hasattr(os, 'O_NOFOLLOW') or os.open not in os.supports_dir_fd: + raise OSError(errno.ENOTSUP, 'Secure private job storage is unavailable on this platform.') + path = Path(path).expanduser().absolute() + path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) + self.path = path.parent.resolve(strict=True) / _name(path.name) + self._fd = None + parent = os.open(self.path.parent, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW) + try: + try: + os.mkdir(self.path.name, mode=0o700, dir_fd=parent) + except FileExistsError: + pass + descriptor = os.open(self.path.name, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=parent) + finally: + os.close(parent) + self._initialize(descriptor) + + @classmethod + def _from_fd(cls, descriptor, path): + instance = cls.__new__(cls) + instance.path = path + instance._fd = None + instance._initialize(descriptor) + return instance + + def _initialize(self, descriptor): + self._fd = descriptor + try: + _protect(descriptor, True) + information = os.fstat(descriptor) + self.identity = (information.st_dev, information.st_ino) + self._repair_tree(descriptor) + self.verify() + except BaseException: + self.close() + raise + + def _repair_tree(self, descriptor): + # Enumeration and repairs remain beneath retained handles, including old stores. + with os.scandir(descriptor) as entries: + for entry in entries: + try: + information = entry.stat(follow_symlinks=False) + directory = stat.S_ISDIR(information.st_mode) + if not directory and not stat.S_ISREG(information.st_mode): + raise ValueError('Private job storage contains a symbolic link or nonregular entry; recover it before retrying.') + flags = os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK + if directory: + flags |= os.O_DIRECTORY + child = os.open(entry.name, flags, dir_fd=descriptor) + except FileNotFoundError: + continue + try: + _protect(child, directory) + if directory: + self._repair_tree(child) + finally: + os.close(child) + + def verify(self): + descriptor = os.open(self.path, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW) + try: + information = os.fstat(descriptor) + if (information.st_dev, information.st_ino) != self.identity: + raise ValueError('The private job directory changed; restore its original location before retrying.') + finally: + os.close(descriptor) + _protect(self._fd, True) + + def child(self, name): + self.verify() + name = _name(name) + try: + os.mkdir(name, mode=0o700, dir_fd=self._fd) + except FileExistsError: + pass + descriptor = os.open(name, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=self._fd) + return self._from_fd(descriptor, self.path / name) + + def open_file(self, name, *, exclusive=False): + self.verify() + flags = os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW | os.O_NONBLOCK + if exclusive: + flags |= os.O_EXCL + descriptor = os.open(_name(name), flags, 0o600, dir_fd=self._fd) + try: + _protect(descriptor, False) + return os.fdopen(descriptor, 'r+b') + except BaseException: + os.close(descriptor) + raise + + def duplicate_fd(self): + self.verify() + return os.dup(self._fd) + + def close(self): + if getattr(self, '_fd', None) is not None: + os.close(self._fd) + self._fd = None + + def __enter__(self): + return self + + def __exit__(self, *unused): + self.close() + + def __del__(self): + self.close() + + +if os.name == 'nt': + from ._private_storage_windows import PrivateDirectory + + +def private_storage_call(action): + try: + return action() + except (OSError, ValueError) as error: + if isinstance(error, PrivateStorageError): + raise + if isinstance(error, OSError) and error.errno == errno.ENOSPC: + raise PrivateStorageError( + 'Local job storage is full. Free space before explicitly resubmitting.', + code='disk_full', kind='tool_failure', + ) from error + raise PrivateStorageError( + 'Private job storage could not be secured. Use an owned local directory ' + 'with enforceable owner-only ACLs; recover symbolic/hard links or ' + 'foreign-owned entries before retrying. No new job was started.' + ) from error + + +def private_directory(path): + return private_storage_call(lambda: PrivateDirectory(path)) diff --git a/src/ucsfomopagent/query_backend.py b/src/ucsfomopagent/query_backend.py index e8d5f55..5449d48 100644 --- a/src/ucsfomopagent/query_backend.py +++ b/src/ucsfomopagent/query_backend.py @@ -1,6 +1,8 @@ """Streaming SQL Server adapter; each background job owns its connection.""" import os +from .jobs import JobError + DIALECT = 'sql' SECRET_KEYS = ('CLINICAL_RECORDS_USERNAME', 'CLINICAL_RECORDS_PASSWORD') ENV_KEYS = ('CLINICAL_RECORDS_USERNAME', 'CLINICAL_RECORDS_PASSWORD', 'CLINICAL_RECORDS_SERVER', 'CLINICAL_RECORDS_DATABASE') @@ -14,7 +16,7 @@ def from_environment(env=None): def check_credentials(values): if not values.get('username') or not values.get('password'): - raise ValueError('Configure CLINICAL_RECORDS_USERNAME and CLINICAL_RECORDS_PASSWORD through BioRouter or the interactive auth command.') + raise JobError('credentials_missing', 'permission_denied', 'Configure database credentials through BioRouter or the interactive auth command.', 'credentials') def stream(request, credentials, progress): @@ -22,17 +24,18 @@ def stream(request, credentials, progress): check_credentials(credentials) with pymssql.connect(server=credentials['server'], database=credentials['database'], user=credentials['username'], password=credentials['password'], timeout=request['timeout_seconds'], login_timeout=20, autocommit=True, appname=__package__ + '-query-job') as conn: with conn.cursor() as cursor: + progress['phase'] = 'planning' if request['mode'] == 'explain' else 'executing' if request['mode'] == 'explain': cursor.execute('SET SHOWPLAN_XML ON') - progress['phase'] = 'planning' if request['mode'] == 'explain' else 'executing' cursor.execute(request['query'], request['parameters'] or None) columns = [c[0] for c in cursor.description] if cursor.description else [] if len(columns) != len(set(columns)): - raise ValueError('Duplicate column names: alias each projection.') + raise JobError('duplicate_columns', 'invalid_args', 'Duplicate column names: alias each projection.', 'executing') while True: + progress['phase'] = 'streaming' rows = cursor.fetchmany(1000) yield columns, rows if not rows: break if cursor.nextset(): - raise ValueError("Multiple result sets are not supported; submit one SELECT per job.") + raise JobError('multiple_result_sets', 'invalid_args', 'Multiple result sets are not supported; submit one SELECT per job.', 'streaming') diff --git a/src/ucsfomopagent/query_cli.py b/src/ucsfomopagent/query_cli.py index f7b7af8..95094a1 100644 --- a/src/ucsfomopagent/query_cli.py +++ b/src/ucsfomopagent/query_cli.py @@ -1,12 +1,12 @@ """CLI front end for the same jobs exposed over MCP.""" import argparse +import asyncio import getpass import importlib import json import os from pathlib import Path import sys -import time from .jobs import JobStore, TERMINAL @@ -78,11 +78,10 @@ def main(argv=None): data = store.cancel(args.job_id) else: while True: - data = store.status(args.job_id) + data = asyncio.run(store.wait(args.job_id)) if args.command == 'watch' else store.status(args.job_id) if args.command != 'watch' or data['state'] in TERMINAL: break print(json.dumps(data), flush=True) - time.sleep(data['recommended_poll_seconds']) print(json.dumps(data), flush=True) return 1 if data['state'] in ('failed', 'timed_out', 'cancelled', 'interrupted') else 0 except (ValueError, RuntimeError) as error: diff --git a/src/ucsfomopagent/server.py b/src/ucsfomopagent/server.py index 7cf9af5..70f9fff 100644 --- a/src/ucsfomopagent/server.py +++ b/src/ucsfomopagent/server.py @@ -67,7 +67,7 @@ - query_ucsf_omop is for small previews and inexpensive queries, with a short timeout and capped output. A TOP bound limits returned rows, not database computation. For slow aggregates, large exports, uncertain costs or timeouts, use - omop-submit_query_job once, then omop-query_job_status at recommended intervals. + omop-submit_query_job once, then use omop-wait_for_query_job or status at recommended intervals. Use estimated plans when permitted, avoid full COUNT scans solely for estimates, and report unknown ETA honestly. Never claim every query is guaranteed to finish. @@ -246,7 +246,7 @@ def query_ucsf_omop( raise except Exception as e: msg = str(e) - hint = " Use omop-submit_query_job for slow queries or full exports; monitor omop-query_job_status." + hint = " Use omop-submit_query_job for slow queries or full exports; use omop-wait_for_query_job." if re.search(r"LIMIT", msg, re.IGNORECASE): hint = " HINT: use `TOP n` instead of `LIMIT` (SQL Server)." elif "Invalid object name" in msg: @@ -472,7 +472,7 @@ def find_measurement( 'specimens or assays. Samples are nonrandom and neither exhaustive nor ' 'population estimates. Filter each threshold by unit_concept_id; never mix ' 'units. Submit full_profile_query with omop-submit_query_job for exact ' - 'coverage, and monitor omop-query_job_status. A TOP limit does not ' + 'coverage, and use omop-wait_for_query_job. A TOP limit does not ' 'guarantee cheap execution; if sampling times out, use a monitored job.')} return ToolResult(content=[TextContent(type='text', text=json.dumps(payload, indent=2))]) except ToolError: diff --git a/tests/test_driver_timeout.py b/tests/test_driver_timeout.py new file mode 100644 index 0000000..233a661 --- /dev/null +++ b/tests/test_driver_timeout.py @@ -0,0 +1,31 @@ +import importlib +import json +import unittest + +import pymssql + +PACKAGE = 'ucsfomopagent' +jobs = importlib.import_module(PACKAGE + '.jobs') + + +class DriverTimeoutTests(unittest.TestCase): + def test_freetds_timeout_keeps_its_kind_and_phase_without_driver_text(self): + for phase in ('connecting', 'executing', 'streaming'): + with self.subTest(phase=phase): + error = pymssql.OperationalError(20003, b'private_marker query=private_marker') + details = jobs.failure_details(error, phase) + self.assertEqual(details['code'], 'database_timeout') + self.assertEqual(details['kind'], 'timeout') + self.assertEqual(details['phase'], phase) + self.assertFalse(details['retryable']) + self.assertNotIn('private_marker', json.dumps(details)) + failed_job = {'job_id': 'a' * 32, 'state': 'running'} + jobs.record_failure(failed_job, details) + operation = jobs.operation_metadata(PACKAGE, failed_job) + self.assertEqual(operation['status'], 'failed') + self.assertEqual(operation['error'], details) + self.assertNotIn('private_marker', json.dumps(operation)) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_job_contract.py b/tests/test_job_contract.py new file mode 100644 index 0000000..0c623b1 --- /dev/null +++ b/tests/test_job_contract.py @@ -0,0 +1,294 @@ +import asyncio +import errno +import importlib +import json +import os +from pathlib import Path +import shutil +import tempfile +import threading +import time +import unittest +from unittest.mock import patch + +PACKAGE = 'ucsfomopagent' +jobs = importlib.import_module(PACKAGE + '.jobs') +worker = importlib.import_module(PACKAGE + '.job_worker') +backend = importlib.import_module(PACKAGE + '.query_backend') +SECURE_READER = hasattr(os, 'O_NOFOLLOW') and hasattr(os, 'O_DIRECTORY') and os.open in os.supports_dir_fd + + +class ContractTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + self.store = jobs.JobStore(PACKAGE, self.temp.name) + + def queued(self, **options): + with patch.object(backend, 'check_credentials'), patch.object(jobs.subprocess, 'Popen') as process: + data = self.store.submit('SELECT 1 AS example', {}, {}, **options) + payload = json.loads(process.return_value.stdin.write.call_args.args[0]) + return data, payload + + def completed(self, value='example'): + data, payload = self.queued() + done = worker.execute(self.store, data['job_id'], payload, lambda *args: iter([(['example'], [[value]])])) + self.assertEqual(done['state'], 'completed') + return done + + def rewrite(self, data, **changes): + data = dict(data, **changes) + with self.store.connect() as db: + db.execute('UPDATE jobs SET state=?, data=? WHERE id=?', (data['state'], json.dumps(data), data['job_id'])) + return data + + def test_capacity_is_discoverable_and_rejection_is_structured(self): + self.assertEqual(self.store.capacity(), dict(active_jobs=0, max_active_jobs=2, available_slots=2, advisory=True)) + first, payload = self.queued() + self.queued() + self.assertEqual(self.store.status(first['job_id'])['capacity']['available_slots'], 0) + with self.assertRaises(jobs.JobError) as raised: + self.queued() + self.assertEqual(raised.exception.details['code'], 'capacity_exceeded') + self.assertEqual(raised.exception.context['capacity']['active_jobs'], 2) + worker.execute(self.store, first['job_id'], payload, lambda *args: iter([])) + self.assertEqual(self.store.capacity()['available_slots'], 1) + + def test_wait_returns_completion_without_resubmission(self): + data, payload = self.queued() + def finish(): + time.sleep(.05) + worker.execute(self.store, data['job_id'], payload, lambda *args: iter([])) + thread = threading.Thread(target=finish) + thread.start() + try: + result = asyncio.run(self.store.wait(data['job_id'], 2)) + finally: + thread.join(3) + self.assertEqual(result['state'], 'completed') + self.assertFalse(result['wait_timed_out']) + self.assertEqual(len(self.store.list_jobs()), 1) + + def test_wait_deadline_and_cancellation_do_not_stop_the_job(self): + data, _ = self.queued() + result = asyncio.run(self.store.wait(data['job_id'], 0)) + self.assertEqual(result['state'], 'queued') + self.assertTrue(result['wait_timed_out']) + async def cancel_wait(): + task = asyncio.create_task(self.store.wait(data['job_id'], 30)) + await asyncio.sleep(.02) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + asyncio.run(cancel_wait()) + self.assertEqual(self.store.status(data['job_id'])['state'], 'queued') + for timeout in [-1, 31, True, .1]: + with self.subTest(timeout=timeout), self.assertRaises(jobs.JobError): + asyncio.run(self.store.wait(data['job_id'], timeout)) + + def test_failure_keeps_safe_code_phase_and_retryability(self): + data, payload = self.queued() + def fail(request, credentials, progress): + progress['phase'] = 'executing' + raise ValueError('password=private_marker row=private_marker') + yield + done = worker.execute(self.store, data['job_id'], payload, fail) + error = done['error_details'] + self.assertEqual(error['code'], 'query_failed') + self.assertEqual(error['phase'], 'executing') + self.assertFalse(error['retryable']) + self.assertNotIn('private_marker', json.dumps(done)) + meta = jobs.operation_metadata(PACKAGE, done) + self.assertEqual(meta['status'], 'failed') + self.assertEqual(meta['error'], error) + + def test_error_categories_use_types_and_codes_without_raw_text(self): + import pymssql + cases = [(pymssql.OperationalError(18456, b'private_marker'), 'authentication_failed', False), + (pymssql.ProgrammingError(229, b'private_marker'), 'database_permission_denied', False), + (pymssql.ProgrammingError(208, b'private_marker'), 'query_rejected', False), + (pymssql.OperationalError(1205, b'private_marker'), 'database_busy', True), + (OSError(errno.ENOSPC, 'private_marker'), 'disk_full', False), + (PermissionError(errno.EACCES, 'private_marker'), 'storage_permission_denied', False), + (ConnectionResetError(errno.ECONNRESET, 'private_marker'), 'connection_interrupted', True), + (TimeoutError('private_marker'), 'database_timeout', False)] + for error, code, retryable in cases: + with self.subTest(code=code): + details = jobs.failure_details(error, 'connecting') + self.assertEqual(details['code'], code) + self.assertEqual(details['retryable'], retryable) + self.assertNotIn('private_marker', json.dumps(details)) + + def test_worker_start_error_is_private_and_releases_capacity(self): + with patch.object(backend, 'check_credentials'), patch.object(jobs.subprocess, 'Popen', side_effect=OSError('private_marker')): + data = self.store.submit('SELECT 1 AS example', {}, {}) + self.assertEqual(data['error_details']['code'], 'worker_start_failed') + self.assertEqual(data['state'], 'failed') + self.assertNotIn('private_marker', json.dumps(data)) + self.assertEqual(self.store.capacity()['active_jobs'], 0) + + def test_operation_outcomes_cover_legacy_failures_and_terminal_states(self): + data, _ = self.queued() + for state in ('failed', 'timed_out', 'interrupted'): + with self.subTest(state=state): + legacy = dict(data, state=state, error='private_marker') + operation = jobs.operation_metadata(PACKAGE, legacy) + self.assertEqual(operation['status'], 'failed') + self.assertEqual(operation['id'], PACKAGE + ':' + data['job_id']) + self.assertEqual(set(operation['error']), {'code', 'kind', 'message', 'phase', 'retryable'}) + self.assertNotIn('private_marker', json.dumps(operation)) + for state in ('queued', 'running', 'completed', 'cancelled'): + with self.subTest(state=state): + operation = jobs.operation_metadata(PACKAGE, dict(data, state=state)) + self.assertEqual(operation['status'], state) + self.assertNotIn('error', operation) + + def test_mcp_preserves_semantic_operation_metadata(self): + async def run(): + from fastmcp import Client, FastMCP + mcp = FastMCP('synthetic contract') + jobs.register_job_tools(mcp, PACKAGE, {}, prefix='fixture-', store=self.store) + async with Client(mcp) as client: + tools = {tool.name: tool for tool in await client.list_tools()} + self.assertEqual(set(tools['fixture-read_query_job_result'].input_schema['properties']), {'job_id', 'offset', 'max_bytes'}) + self.assertIn('timeout_seconds', tools['fixture-wait_for_query_job'].input_schema['properties']) + with patch.object(backend, 'check_credentials'), patch.object(jobs.subprocess, 'Popen'): + submitted = await client.call_tool('fixture-submit_query_job', {'query': 'SELECT 1 AS example'}) + data = json.loads(submitted.content[0].text) + self.assertEqual(submitted.meta[jobs.OPERATION_META_KEY]['status'], 'queued') + failed = self.rewrite(data, state='failed', error='Safe failure', error_details=dict(code='fixture_failed', kind='tool_failure', message='Safe failure', phase='executing', retryable=False)) + status = await client.call_tool('fixture-query_job_status', {'job_id': failed['job_id']}) + self.assertFalse(status.is_error) + self.assertEqual(status.meta[jobs.OPERATION_META_KEY]['status'], 'failed') + self.assertEqual(status.meta[jobs.OPERATION_META_KEY]['error']['code'], 'fixture_failed') + waited = await client.call_tool('fixture-wait_for_query_job', {'job_id': failed['job_id'], 'timeout_seconds': 0}) + self.assertEqual(waited.meta, status.meta) + unknown = await client.call_tool('fixture-query_job_status', {'job_id': '0' * 32}, raise_on_error=False) + self.assertTrue(unknown.is_error) + self.assertEqual(unknown.structured_content['error']['kind'], 'not_found') + capacity = await client.call_tool('fixture-query_job_capacity', {}) + self.assertEqual(capacity.structured_content['available_slots'], 2) + asyncio.run(run()) + + @unittest.skipUnless(SECURE_READER, 'secure POSIX file opens required') + def test_reader_reassembles_utf8_chunks_and_ignores_forged_path(self): + done = self.completed('α😀漢字' * 5) + original = Path(done['result_path']).read_text() + self.rewrite(done, result_path='/not/a/result') + for limit in range(4, 15): + offset, parts = 0, [] + while True: + chunk = self.store.read_result(done['job_id'], offset, limit) + parts.append(chunk['text']) + self.assertLessEqual(chunk['bytes_read'], limit) + if chunk['eof']: + self.assertIsNone(chunk['next_offset']) + break + self.assertGreater(chunk['next_offset'], offset) + offset = chunk['next_offset'] + self.assertEqual(''.join(parts), original) + end = self.store.read_result(done['job_id'], len(original.encode())) + self.assertEqual(end['text'], '') + self.assertTrue(end['eof']) + + @unittest.skipUnless(SECURE_READER, 'secure POSIX file opens required') + def test_reader_rejects_bad_offsets_limits_and_noncompleted_jobs(self): + done = self.completed('α') + for offset, limit in [(-1, 4), (True, 4), (0, 3), (0, 65537), (0, True), (9999, 4)]: + with self.subTest(offset=offset, limit=limit), self.assertRaises(jobs.JobError): + self.store.read_result(done['job_id'], offset, limit) + raw = Path(done['result_path']).read_bytes() + with self.assertRaises(jobs.JobError): + self.store.read_result(done['job_id'], raw.index('α'.encode()) + 1) + for state in ['queued', 'running', 'failed', 'cancelled', 'timed_out', 'interrupted']: + self.rewrite(done, state=state) + with self.subTest(state=state), self.assertRaises(jobs.JobError): + self.store.read_result(done['job_id']) + + @unittest.skipUnless(SECURE_READER, 'secure POSIX file opens required') + def test_reader_rejects_missing_changed_and_invalid_files(self): + done = self.completed() + path = Path(done['result_path']) + path.write_bytes(b'\xff' * done['bytes']) + with self.assertRaises(jobs.JobError) as raised: + self.store.read_result(done['job_id']) + self.assertEqual(raised.exception.details['code'], 'invalid_result_encoding') + path.write_text('changed') + with self.assertRaises(jobs.JobError): + self.store.read_result(done['job_id']) + path.unlink() + with self.assertRaises(jobs.JobError): + self.store.read_result(done['job_id']) + self.rewrite(done, format='../outside') + with self.assertRaises(jobs.JobError): + self.store.read_result(done['job_id']) + other = jobs.JobStore('other_connector', self.temp.name) + with self.assertRaises(jobs.JobError): + other.read_result(done['job_id']) + + @unittest.skipUnless(SECURE_READER, 'secure POSIX file opens required') + def test_reader_refuses_symlinks_hardlinks_and_fifo(self): + done = self.completed() + path = Path(done['result_path']) + outside = Path(self.temp.name) / 'outside' + outside.write_text('must not be read') + path.unlink() + path.symlink_to(outside) + with self.assertRaises(jobs.JobError): + self.store.read_result(done['job_id']) + path.unlink() + os.link(outside, path) + with self.assertRaises(jobs.JobError): + self.store.read_result(done['job_id']) + path.unlink() + os.mkfifo(path) + original_open = os.open + def checked_open(filename, flags, *args, **kwargs): + if filename == 'result.jsonl': + self.assertTrue(flags & os.O_NONBLOCK, 'opening an untrusted FIFO must not block') + return original_open(filename, flags, *args, **kwargs) + with patch.object(jobs.os, 'open', checked_open), patch.object(jobs.os, 'supports_dir_fd', {checked_open}): + with self.assertRaises(jobs.JobError): + self.store.read_result(done['job_id']) + + @unittest.skipUnless(SECURE_READER, 'secure POSIX file opens required') + def test_reader_refuses_job_directory_symlink(self): + done = self.completed() + directory = self.store.directory(done['job_id']) + moved = directory.with_name('moved') + directory.rename(moved) + directory.symlink_to(moved, target_is_directory=True) + with self.assertRaises(jobs.JobError): + self.store.read_result(done['job_id']) + + @unittest.skipUnless(SECURE_READER, 'secure POSIX file opens required') + def test_reader_refuses_replacement_of_the_pinned_root(self): + done = self.completed() + root = self.store.root + moved = root.with_name(root.name + '-moved') + root.rename(moved) + shutil.copytree(moved, root) + with self.assertRaises(jobs.JobError) as raised: + self.store.read_result(done['job_id']) + self.assertEqual(raised.exception.details['code'], 'result_scope_changed') + + @unittest.skipUnless(SECURE_READER, 'secure POSIX file opens required') + def test_reader_returns_eof_for_an_empty_complete_result(self): + data, payload = self.queued() + done = worker.execute(self.store, data['job_id'], payload, lambda *args: iter([])) + chunk = self.store.read_result(done['job_id']) + self.assertEqual(chunk['text'], '') + self.assertEqual(chunk['total_bytes'], 0) + self.assertEqual(chunk['bytes_read'], 0) + self.assertTrue(chunk['eof']) + self.assertIsNone(chunk['next_offset']) + + def test_reader_fails_closed_without_secure_open_support(self): + done = self.completed() + with patch.object(jobs.os, 'supports_dir_fd', set()), self.assertRaises(jobs.JobError) as raised: + self.store.read_result(done['job_id']) + self.assertEqual(raised.exception.details['code'], 'secure_reader_unavailable') + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_job_followup.py b/tests/test_job_followup.py new file mode 100644 index 0000000..c208755 --- /dev/null +++ b/tests/test_job_followup.py @@ -0,0 +1,99 @@ +"""Offline regressions for capacity discovery and worker-start handoff.""" +import asyncio +from contextlib import redirect_stdout +import importlib +import io +import json +import tempfile +import unittest +from unittest.mock import AsyncMock, Mock, patch + +PACKAGE = 'ucsfomopagent' +jobs = importlib.import_module(PACKAGE + '.jobs') +backend = importlib.import_module(PACKAGE + '.query_backend') + + +class FollowupTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + self.store = jobs.JobStore(PACKAGE, self.temp.name) + + def test_capacity_snapshots_explicitly_remain_advisory(self): + self.assertIs(self.store.capacity()['advisory'], True) + with patch.object(backend, 'check_credentials'), patch.object(jobs.subprocess, 'Popen'): + data = self.store.submit('SELECT 1 AS example', {}, {}) + self.store.submit('SELECT 1 AS example', {}, {}) + with self.assertRaises(jobs.JobError) as caught: + self.store.submit('SELECT 1 AS example', {}, {}) + self.assertIs(self.store.status(data['job_id'])['capacity']['advisory'], True) + self.assertIs(caught.exception.context['capacity']['advisory'], True) + + def test_cli_watch_reuses_native_wait(self): + cli = importlib.import_module(PACKAGE + '.query_cli') + queued = dict(job_id='a'*32, state='queued') + completed = dict(queued, state='completed') + with patch.object(cli, 'JobStore') as constructor, redirect_stdout(io.StringIO()): + store = constructor.return_value + store.status.return_value = completed + store.wait = AsyncMock(side_effect=[queued, completed]) + self.assertEqual(cli.main(['watch', queued['job_id']]), 0) + self.assertEqual(store.wait.await_count, 2) + store.wait.assert_awaited_with(queued['job_id']) + store.status.assert_not_called() + + def test_failed_start_returns_durable_id_and_failed_operation(self): + with patch.object(backend, 'check_credentials'), patch.object(jobs.subprocess, 'Popen', side_effect=OSError('private-start-detail')): + data = self.store.submit('SELECT 1 AS example', {}, {}) + self.assertEqual(self.store.status(data['job_id'])['state'], 'failed') + self.assertEqual(self.store.capacity()['active_jobs'], 0) + result = jobs.tool_result(data, PACKAGE) + self.assertFalse(result.is_error) + self.assertEqual(result.meta[jobs.OPERATION_META_KEY]['id'], PACKAGE + ':' + data['job_id']) + self.assertEqual(result.meta[jobs.OPERATION_META_KEY]['status'], 'failed') + self.assertNotIn('private-start-detail', result.model_dump_json()) + + def test_failed_pipe_cleanup_does_not_hide_id_or_leak_capacity(self): + process = Mock() + process.stdin.write.side_effect = BrokenPipeError('private-write-detail') + process.stdin.close.side_effect = BrokenPipeError('private-close-detail') + with patch.object(backend, 'check_credentials'), patch.object(jobs.subprocess, 'Popen', return_value=process): + data = self.store.submit('SELECT 1 AS example', {}, {}) + self.assertEqual(data['error_details']['code'], 'worker_start_failed') + self.assertEqual(self.store.status(data['job_id'])['state'], 'failed') + self.assertEqual(self.store.capacity()['available_slots'], 2) + process.stdin.close.assert_called_once() + self.assertNotIn('private-write-detail', json.dumps(data)) + self.assertNotIn('private-close-detail', json.dumps(data)) + + def test_missing_pipe_still_returns_a_failed_job_id(self): + process = Mock(stdin=None) + with patch.object(backend, 'check_credentials'), patch.object(jobs.subprocess, 'Popen', return_value=process): + data = self.store.submit('SELECT 1 AS example', {}, {}) + self.assertEqual(data['state'], 'failed') + self.assertEqual(self.store.capacity()['active_jobs'], 0) + + def test_reaper_thread_failure_does_not_fail_a_started_worker(self): + with patch.object(backend, 'check_credentials'), patch.object(jobs.subprocess, 'Popen'), patch.object(jobs.threading.Thread, 'start', side_effect=RuntimeError('private-thread-detail')): + data = self.store.submit('SELECT 1 AS example', {}, {}) + self.assertEqual(data['state'], 'queued') + self.assertEqual(self.store.capacity()['active_jobs'], 1) + self.assertNotIn('private-thread-detail', json.dumps(data)) + + def test_real_mcp_start_failure_observes_allocated_failed_operation(self): + async def run(): + from fastmcp import Client, FastMCP + mcp = FastMCP('startup failure') + jobs.register_job_tools(mcp, PACKAGE, {}, prefix='fixture-', store=self.store) + async with Client(mcp) as client: + with patch.object(backend, 'check_credentials'), patch.object(jobs.subprocess, 'Popen', side_effect=OSError('private-start-detail')): + result = await client.call_tool('fixture-submit_query_job', {'query': 'SELECT 1 AS example'}, raise_on_error=False) + self.assertFalse(result.is_error) + data = json.loads(result.content[0].text) + outcome = result.meta[jobs.OPERATION_META_KEY] + self.assertEqual(outcome['id'], PACKAGE + ':' + data['job_id']) + self.assertEqual(outcome['status'], 'failed') + self.assertEqual(outcome['error']['phase'], 'worker_start') + self.assertNotIn('capacity', outcome) + self.assertEqual(self.store.capacity()['active_jobs'], 0) + asyncio.run(run()) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 6a99f5f..6a89c25 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -17,11 +17,11 @@ async def run(): server = importlib.import_module(PACKAGE + '.server') if PACKAGE == 'ucsfomopagent': mcp = server.create_ucsf_omop_server(server.UCSFOMOPConfig(username='test',password='test')) - prefix, count = 'omop-', 8 + prefix, count = 'omop-', 11 elif PACKAGE == 'cdwagent': from cdwagent.config import CDWConfig, ClinicalDBConfig mcp = server.create_cdw_server(CDWConfig(clinical_db=ClinicalDBConfig(username='test',password='test'))) - prefix, count = 'CDW-', 25 + prefix, count = 'CDW-', 28 else: mcp = server.create_spoke_server(server.SPOKEConfig(uri='bolt://localhost:7687',username='test',password='test',database='neo4j')) prefix, count = 'spoke-', 8 @@ -29,7 +29,7 @@ async def run(): tools = await client.list_tools() names = {t.name for t in tools} self.assertEqual(len(names), count, names) - self.assertTrue({prefix+n for n in ['submit_query_job','query_job_status','cancel_query_job']} <= names) + self.assertTrue({prefix+n for n in ['submit_query_job','query_job_capacity','query_job_status','wait_for_query_job','read_query_job_result','cancel_query_job']} <= names) submit = next(t for t in tools if t.name == prefix+'submit_query_job') self.assertNotIn('credentials', submit.input_schema['properties']) query = 'SELECT 1 AS value' if PACKAGE != 'spokeagent' else 'RETURN 1 AS value' diff --git a/tests/test_private_storage.py b/tests/test_private_storage.py new file mode 100644 index 0000000..d941ebd --- /dev/null +++ b/tests/test_private_storage.py @@ -0,0 +1,213 @@ +"""Synthetic owner-private storage controls; no database credentials are needed.""" +import asyncio +import errno +import importlib +import os +from pathlib import Path +import stat +import struct +import subprocess +import sys +import tempfile +import unittest +from unittest.mock import patch + +PACKAGE = 'ucsfomopagent' +storage = importlib.import_module(PACKAGE + '.private_storage') +jobs = importlib.import_module(PACKAGE + '.jobs') + + +@unittest.skipUnless(os.name == 'posix' and (sys.platform == 'darwin' or sys.platform.startswith('linux')), 'POSIX private storage backend required') +class PrivateStorageTests(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory(prefix='connector-private-storage-') + self.addCleanup(self.temporary.cleanup) + self.parent = Path(self.temporary.name) + + def assert_mode(self, path, mode): + self.assertEqual(stat.S_IMODE(path.stat().st_mode), mode) + + def test_new_file_is_private_before_any_sensitive_bytes(self): + with storage.private_directory(self.parent / 'store') as directory: + with directory.open_file('new-result', exclusive=True) as result: + self.assertEqual(os.fstat(result.fileno()).st_size, 0) + self.assert_mode(directory.path, 0o700) + self.assert_mode(directory.path / 'new-result', 0o600) + result.write(b'synthetic-only') + + def test_existing_owned_tree_is_repaired_without_reading_contents(self): + root = self.parent / 'store' + child = root / 'old-job' + child.mkdir(parents=True) + result = child / 'result.csv' + result.write_text('synthetic-only') + root.chmod(0o755) + child.chmod(0o755) + result.chmod(0o644) + with storage.private_directory(root): + self.assert_mode(root, 0o700) + self.assert_mode(child, 0o700) + self.assert_mode(result, 0o600) + self.assertEqual(result.read_text(), 'synthetic-only') + + def test_symlink_root_and_descendant_never_change_outside_permissions(self): + outside = self.parent / 'outside' + outside.mkdir(mode=0o755) + root = self.parent / 'store' + root.symlink_to(outside, target_is_directory=True) + with self.assertRaises(storage.PrivateStorageError): + storage.private_directory(root) + self.assert_mode(outside, 0o755) + root.unlink() + root.mkdir() + (root / 'alias').symlink_to(outside, target_is_directory=True) + with self.assertRaises(storage.PrivateStorageError): + storage.private_directory(root) + self.assert_mode(outside, 0o755) + + def test_hardlinked_file_is_rejected_before_outside_inode_repair(self): + root = self.parent / 'store' + root.mkdir() + outside = self.parent / 'outside' + outside.write_text('synthetic-only') + outside.chmod(0o644) + os.link(outside, root / 'linked-result') + with self.assertRaises(storage.PrivateStorageError): + storage.private_directory(root) + self.assert_mode(outside, 0o644) + self.assertEqual(outside.read_text(), 'synthetic-only') + + def test_foreign_owner_refusal_precedes_permission_repair(self): + root = self.parent / 'store' + root.mkdir(mode=0o755) + other_uid = os.geteuid() + 1 + with patch.object(storage.os, 'geteuid', return_value=other_uid): + with self.assertRaises(storage.PrivateStorageError): + storage.private_directory(root) + self.assert_mode(root, 0o755) + + def test_replacement_directory_is_refused_without_writing_into_it(self): + root = self.parent / 'store' + with storage.private_directory(root) as directory: + root.rename(self.parent / 'original') + root.mkdir(mode=0o755) + with self.assertRaises(ValueError): + directory.open_file('result.csv', exclusive=True) + self.assertFalse((root / 'result.csv').exists()) + self.assert_mode(root, 0o755) + + @unittest.skipUnless(sys.platform == 'darwin', 'macOS inherited ACL control') + def test_macos_inherited_acl_removed_before_new_file_and_from_old_tree(self): + shared = self.parent / 'shared' + shared.mkdir() + subprocess.run(['/bin/chmod', '+a', 'everyone allow list,search,read,readattr,readextattr,readsecurity,file_inherit,directory_inherit', str(shared)], check=True) + root = shared / 'store' + old = root / 'old-job' + old.mkdir(parents=True) + result = old / 'result.csv' + result.write_text('synthetic-only') + before = subprocess.run(['/bin/ls', '-lde', str(result)], check=True, text=True, capture_output=True).stdout + self.assertIn('everyone', before) + with storage.private_directory(root) as directory: + for path in (root, old, result): + acl = subprocess.run(['/bin/ls', '-lde', str(path)], check=True, text=True, capture_output=True).stdout + self.assertNotIn('everyone', acl) + with directory.open_file('new-result', exclusive=True) as created: + self.assertEqual(os.fstat(created.fileno()).st_size, 0) + acl = subprocess.run(['/bin/ls', '-lde', str(root / 'new-result')], check=True, text=True, capture_output=True).stdout + self.assertNotIn('everyone', acl) + created.write(b'synthetic-only') + + @unittest.skipUnless(sys.platform.startswith('linux'), 'Linux default ACL control') + def test_linux_default_and_access_acls_are_removed(self): + shared = self.parent / 'shared' + shared.mkdir() + entries = [(1, 7, 0xffffffff), (2, 5, 65534), (4, 5, 0xffffffff), (16, 5, 0xffffffff), (32, 5, 0xffffffff)] + acl = struct.pack('