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/AGENTS.md b/AGENTS.md index 5efd936..2b1c37b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -75,7 +75,7 @@ src/cdwagent/ 1. **pip/uvx package** (`src/cdwagent/`): `cli.py` reads env vars → `server.main()` → `create_cdw_server(config)` → `mcp.run()` 2. **BioRouter extension**: install `extensions/cdwagent.brxt`; BioRouter creates its environment and configures the declared credentials. See `docs/BIOROUTER.md`. -3. **Job CLI**: `cdwagent submit/status/watch/list/cancel/purge` uses the same durable job runtime as the three MCP job tools. `cdwagent auth` configures a separate standalone OS-keyring profile. +3. **Job CLI**: `cdwagent submit/status/watch/list/cancel/purge` uses the same durable job runtime as the six MCP job tools. `cdwagent auth` configures a separate standalone OS-keyring profile. ### Schema Reference @@ -102,7 +102,7 @@ All config via environment variables (see `.env.example`): - `CDW_SCHEMA` (database schema for table qualification, default "deid_uf") - `CDW_LOG_LEVEL` -### 25 MCP Tools (namespace-prefixed with `CDW-`) +### 28 MCP Tools (namespace-prefixed with `CDW-`) | Module | Tools | |---|---| @@ -113,7 +113,7 @@ All config via environment variables (see `.env.example`): | concepts.py | `search_diagnoses_by_code`, `search_medications_by_code`, `search_labs_by_code`, `search_procedures_by_code` | | stats.py | `summarize_table`, `cohort_summary` | | cohort.py | `build_cohort` | -| jobs.py | `submit_query_job`, `query_job_status`, `cancel_query_job` | +| jobs.py | `submit_query_job`, `query_job_capacity`, `query_job_status`, `wait_for_query_job`, `read_query_job_result`, `cancel_query_job` | `search_note_concepts` and `search_note_sdoh` were added in v0.4.0 to expose the NLP-extracted concept layer (`note_concepts`, `note_concepts_sdoh`, populated by @@ -181,7 +181,7 @@ The `crossmap_patient` tool resolves an OMOP `person_id` to a CDW `PatientDurabl BioRouter does NOT use a classifier to pick agents; it flattens all enabled extensions' tools into one list and lets the LLM choose. Consequences for this repo: - Tool names and descriptions ARE the routing signal — keep the `CDW-` prefix to disambiguate from OMOPAgent's `query_ucsf_omop` etc. -- Use `submit_query_job` for expensive queries/exports and poll `query_job_status` at its recommended interval. Increasing an MCP timeout alone does not make an export durable. +- Use `submit_query_job` for expensive queries/exports and use native `wait_for_query_job`, or poll `query_job_status` at its recommended interval. Increasing an MCP timeout alone does not make an export durable. - Publish tagged releases so BioRouter's `extension_malware_check` pins a stable git ref. ### Context strategy (LLM dispatch optimization) diff --git a/README.md b/README.md index cca3f7b..956b4ee 100644 --- a/README.md +++ b/README.md @@ -44,9 +44,15 @@ job-routing instructions; the skill folders can also be installed separately. ## Long queries, CLI and progress Use `CDW-submit_query_job` for an export or a query that could exceed the -interactive timeout. It returns a job ID immediately. Poll -`CDW-query_job_status` at its recommended interval; use -`CDW-cancel_query_job` to stop. Only `completed` means the file is complete. +interactive timeout. Call `CDW-query_job_capacity` first; submission returns a +job ID immediately and atomically checks the two-job limit. Use +`CDW-wait_for_query_job` for native waits of up to 30 seconds, or +`CDW-query_job_status` for an immediate snapshot. Use `CDW-cancel_query_job` +to stop. Only `completed` means the file is complete. Read bounded UTF-8 chunks +with `CDW-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. @@ -63,7 +69,8 @@ uv run cdwagent watch JOB_ID No-argument `uv run cdwagent` 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. See [architecture, storage, security and release details](docs/QUERY_JOBS.md). @@ -84,7 +91,7 @@ Architecture is based on the [MedCP](https://github.com/BaranziniLab/MedCP) temp ## Features -- 25 MCP tools: 22 domain tools across 7 modules and 3 durable query-job tools +- 28 MCP tools: 22 domain tools across 7 modules and 6 durable query-job tools - **One-call multimodal cohort building** (`build_cohort`) across 8 modalities: diagnosis, medication, procedure, lab, **imaging/radiology**, **immunization**, **allergy**, and **vitals/flowsheets** — resolves a term/code, builds the diff --git a/docs/QUERY_JOBS.md b/docs/QUERY_JOBS.md index 41c7bba..9a8c648 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/cdwagent.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/cdwagent/` (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 cdwagent cancel JOB_ID uv run cdwagent 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":"cdwagent: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 `CDW-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 f8f5c5e..7aa8168 100644 --- a/extensions/SHA256SUMS +++ b/extensions/SHA256SUMS @@ -1 +1 @@ -91372a3a9e69ae99bd76d849c882432ddfcd79b6bb6d8fd45f101286d911d0e2 cdwagent.brxt +1e0923a6616a47719ed4acc7f2684bde8304f5a14db616959b8ddd851cff3bcd cdwagent.brxt diff --git a/extensions/cdwagent.brxt b/extensions/cdwagent.brxt index 1aefff6..4d60c02 100644 Binary files a/extensions/cdwagent.brxt and b/extensions/cdwagent.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/cdwagent-query-jobs/SKILL.md b/skills/cdwagent-query-jobs/SKILL.md index 546f3ce..7914839 100644 --- a/skills/cdwagent-query-jobs/SKILL.md +++ b/skills/cdwagent-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 `CDW-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 `CDW-cancel_query_job` for unwanted work, then poll to terminal state. +3. Call `CDW-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 `CDW-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 `CDW-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 + `CDW-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 `CDW-cancel_query_job` for + unwanted work, then wait for a terminal state before replacing it. CLI status needs no credentials: from the installed extension directory, run `uv run cdwagent status JOB_ID` or `uv run cdwagent watch JOB_ID`. diff --git a/src/cdwagent/_private_storage_windows.py b/src/cdwagent/_private_storage_windows.py new file mode 100644 index 0000000..645e3ba --- /dev/null +++ b/src/cdwagent/_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/cdwagent/job_worker.py b/src/cdwagent/job_worker.py index 6602812..8836b72 100644 --- a/src/cdwagent/job_worker.py +++ b/src/cdwagent/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/cdwagent/jobs.py b/src/cdwagent/jobs.py index 6b18080..58b93c9 100644 --- a/src/cdwagent/jobs.py +++ b/src/cdwagent/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/cdwagent/private_storage.py b/src/cdwagent/private_storage.py new file mode 100644 index 0000000..cdc9dfc --- /dev/null +++ b/src/cdwagent/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/cdwagent/query_backend.py b/src/cdwagent/query_backend.py index 9ee5bfd..3c24c67 100644 --- a/src/cdwagent/query_backend.py +++ b/src/cdwagent/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/cdwagent/query_cli.py b/src/cdwagent/query_cli.py index f7b7af8..95094a1 100644 --- a/src/cdwagent/query_cli.py +++ b/src/cdwagent/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/cdwagent/server.py b/src/cdwagent/server.py index 48838a9..9acfd17 100644 --- a/src/cdwagent/server.py +++ b/src/cdwagent/server.py @@ -226,7 +226,7 @@ def clinical_data_exploration() -> str: "2. Search the schema for tables related to my topic of interest\n" "3. Describe the relevant tables to understand their columns\n" "4. Run a bounded preview; use submit_query_job for exports or uncertain-cost queries\n" - "5. Poll query_job_status at recommended_poll_seconds until completed; report failures and incomplete results\n\n" + "5. Use wait_for_query_job until completed, or query_job_status at recommended_poll_seconds; report failures and incomplete results\n\n" f"All tables are in the {schema} schema (e.g., {schema}.PatientDim).\n\n" "CRITICAL — PATIENT IDENTIFIERS:\n" "- PatientDurableKey is the STABLE patient identifier. Always use it for cohort queries.\n" @@ -272,7 +272,7 @@ def cohort_building() -> str: "3. Use cohort_summary with the patient_key_query to get counts and demographics\n" f"4. Retrieve demographics: SELECT ... FROM {schema}.PatientDim WHERE IsCurrent = 1 AND PatientDurableKey IN (subquery)\n" "5. For clinical details (labs, meds, encounters): filter fact tables WHERE PatientDurableKey IN (subquery)\n" - "6. Export with submit_query_job(format='csv'); monitor query_job_status until completed\n" + "6. Export with submit_query_job(format='csv'); use wait_for_query_job until completed\n" " Do not run a costly count only for an ETA. Do not duplicate active jobs.\n\n" "KEY COLUMN NAMES:\n" "- PatientDim: PatientKey, PatientDurableKey, Sex, BirthDate, DeathDate, FirstRace, Ethnicity\n" @@ -298,7 +298,7 @@ def notes_analysis() -> str: "3. Review note metadata (note_type, encounter_type, enc_dept_specialty, deid_service_date)\n" "4. Read full note text for relevant findings using get_note\n" "5. Summarize patterns, disclose NLP exclusions and population sampling notices\n" - "6. For expensive population queries or exports, submit_query_job and poll query_job_status\n\n" + "6. For expensive population queries or exports, submit_query_job and wait_for_query_job\n\n" "What patient or keyword should we start searching for?" ) diff --git a/tests/test_driver_timeout.py b/tests/test_driver_timeout.py new file mode 100644 index 0000000..ff9a2c7 --- /dev/null +++ b/tests/test_driver_timeout.py @@ -0,0 +1,31 @@ +import importlib +import json +import unittest + +import pymssql + +PACKAGE = 'cdwagent' +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..9fc0ffb --- /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 = 'cdwagent' +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..aebd452 --- /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 = 'cdwagent' +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 c9d4970..e8c2014 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..5052cfb --- /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 = 'cdwagent' +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('