From de38798ac7423f158dec02f8eae30ec2e3e8acfc Mon Sep 17 00:00:00 2001 From: Joe Clark Date: Thu, 2 Jul 2026 15:41:38 +0100 Subject: [PATCH 1/2] long-running python process: partial commit --- .changeset/tidy-worms-cheer.md | 5 + .../2-service-tests.md | 501 ++++++++++++++++ platform/src/bridge.ts | 555 ++++++++++++++---- platform/src/server.ts | 14 +- services/streaming_util.py | 13 +- services/util.py | 65 +- services/worker.py | 257 ++++++++ 7 files changed, 1291 insertions(+), 119 deletions(-) create mode 100644 .changeset/tidy-worms-cheer.md create mode 100644 agent-team-architecture-plan/2-service-tests.md create mode 100644 services/worker.py diff --git a/.changeset/tidy-worms-cheer.md b/.changeset/tidy-worms-cheer.md new file mode 100644 index 00000000..db211fab --- /dev/null +++ b/.changeset/tidy-worms-cheer.md @@ -0,0 +1,5 @@ +--- +"apollo": minor +--- + +Run Python services on a single long-lived worker over a Bun-owned Unix socket instead of spawning a process per request diff --git a/agent-team-architecture-plan/2-service-tests.md b/agent-team-architecture-plan/2-service-tests.md new file mode 100644 index 00000000..e58598b5 --- /dev/null +++ b/agent-team-architecture-plan/2-service-tests.md @@ -0,0 +1,501 @@ +# Section 2 — Service Tests Architecture + +> Scope: `services/global_chat/`, `services/workflow_chat/`, +> `services/job_chat/`, and the tools / sub-agents they invoke. + +--- + +## 1. Naming and position + +**Service tests** — direct calls to `main()` with Anthropic HTTP calls mocked. +One rung above unit, one below integration. + +| Tier | Scope | LLM calls | HTTP layer | Cost | Runs on PR push | +| ----------- | ----------------------- | ---------- | ---------------------- | -------- | ---------------- | +| unit | single function | no | no | free | yes | +| **service** | **`main()` end-to-end** | **mocked** | **no (direct python)** | **free** | **yes** | +| integration | `main()` via server | real | yes | $$ | manual / label | +| acceptance | behaviour spec | real | yes | $$$ | nightly / manual | + +Service tests verify _logic and information flow_: payload validation, routing, +prompt assembly, tool-call orchestration, sub-agent invocation, history/usage +aggregation, error paths, headers, api_key passthrough. They do **not** verify +model quality. + +--- + +## 2. The `test_hooks` second argument + +### 2.1 Signature change + +```python +def main(data_dict: dict, test_hooks: Optional[dict] = None) -> dict: ... +``` + +`entry.py` keeps calling `m.main(data)` with one positional arg — the HTTP path +never sees `test_hooks`. `test_hooks` is a test-only affordance. + +### 2.2 The `test_hooks` dict — minimum viable shape + +Plain Python `dict`. No `TypedDict`, no pydantic model — just a dict with +documented keys. The recognised keys are documented as a docstring in +`testing/anthropic_mock.py`: + +```python +# testing/anthropic_mock.py +""" +The `test_hooks` dict accepts (all optional; all default to absent): + +- "anthropic_http_client": an httpx.Client backed by httpx.MockTransport. + When present, threaded into every Anthropic(...) constructor site. +- "tool_calls": a list[dict] the test allocates. Production code appends + breadcrumbs via record_tool_call(test_hooks, entry). +- "tool_stubs": dict[str, Callable] keyed by tool name. When the planner + dispatches a tool, if a stub exists for that name, the stub is called + with the tool input and its return value used as the tool result. Today + only used for "search_documentation" — see §5. +""" +``` + +Start with three keys. Add more only when a concrete test can't be written +without one. Things intentionally left out until needed: + +- Sub-agent stub registry. Default behaviour: the sub-agent's `main()` runs + under the same mock HTTP client — that's usually what a test wants. A stub + registry (`test_hooks["subagent_stubs"]`) can be added if a test needs to + bypass the sub-agent's logic entirely. +- `seed`, `disable_langfuse`, `scratch`. Add when a test fails without them. + +### 2.3 Threading `test_hooks` through + +Each chat service's `main()` passes `test_hooks` into the agent / client +constructors it creates. Everywhere that currently calls +`Anthropic(api_key=...)` swaps to `build_anthropic_client(api_key, test_hooks)` +(new factory — see §3). + +Sites that change: + +- `services/job_chat/job_chat.py` — `AnthropicClient.__init__`. +- `services/workflow_chat/workflow_chat.py` — `AnthropicClient.__init__`. +- `services/global_chat/router.py` — `RouterAgent.__init__`. +- `services/global_chat/planner.py` — `PlannerAgent.__init__`. +- `services/global_chat/subagent_caller.py` — accepts `test_hooks` and forwards + to sub-agent `main()` calls. + +Production behaviour when `test_hooks is None` is byte-identical to today. Every +new kwarg defaults to `None`. + +--- + +## 3. Mock Anthropic HTTP client + +### 3.1 Factory in `services/util.py` + +```python +def build_anthropic_client(api_key: str, test_hooks: Optional[dict] = None) -> Anthropic: + http_client = (test_hooks or {}).get("anthropic_http_client") + kwargs = {"api_key": api_key} + if http_client is not None: + kwargs["http_client"] = http_client + return Anthropic(**kwargs) +``` + +Every `AnthropicClient` / `RouterAgent` / `PlannerAgent` constructor swaps +`Anthropic(api_key=...)` for `build_anthropic_client(api_key, test_hooks)`. + +### 3.2 `testing/anthropic_mock.py` + +```py + +service_input = make_service_input(history=history, content=content, context=context, meta=meta, suggest_code=True) + +hooks = { + #anthropic: MockAnthropicClient() + http_client_anthropic: MockAnthropicHttpClient() +} + +## deep inside the anthropic client + +http_client.request('anthropic.org/chat', { + query: {}, + headers: {}, + body: {} +}) + +## + +## mock request implementation +def mockResponse(req): + if req.body.match(/cat poems/): + return { message: [{ role: "assistant", content: "sure, I'll write a haiku" }]} + + +http_client_anthropic.setResponse(/write me a cat haiku/, "sure, I'll write a haiku") + + +main(service_input, hooks) + +assert response["suggested_code"] is not None, "JSON parsing failed - suggested_code is None" + +``` + +Single file — `MockAnthropicClient` class, canned response-body builders, +docstring documenting recognised `test_hooks` keys, and the +`record_tool_call(test_hooks, entry)` helper. Split later if it grows unwieldy. + +The `anthropic` Python SDK accepts a custom `http_client` — we build ours from +`httpx.MockTransport`: + +```python +class MockAnthropicClient: + """Thin wrapper over httpx.Client + httpx.MockTransport. + + Usage: + mc = MockAnthropicClient.always(response=text_response("hello")) + mc = MockAnthropicClient.script([resp1, resp2, resp3]) # multi-turn + mc = MockAnthropicClient.streaming(events=[...]) # SSE + + After the call: + mc.requests # list[RecordedRequest] + mc.last_request.json["messages"] + mc.last_request.headers["x-api-key"] + """ + @classmethod + def always(cls, response) -> "MockAnthropicClient": ... + @classmethod + def script(cls, responses) -> "MockAnthropicClient": ... + @classmethod + def streaming(cls, events) -> "MockAnthropicClient": ... + + @property + def httpx_client(self) -> httpx.Client: ... + @property + def requests(self) -> list[RecordedRequest]: ... + + def setResponse(self, question_regex, response_string): + self.responses[question_regex] = response_string + def requests(self) -> + # extract the user question from the request body + # for each response in self.responses + # find the first one which matches user question + # return the response_string associated with that match + # + # if there is no match... idk return error? + + @property + def last_request(self) -> RecordedRequest: ... +``` + +Response-body builders in the same file: + +- `text_response(text, model=..., usage=...)` +- `tool_use_response(tool_name, tool_input, tool_use_id="toolu_01")` +- `mixed_response(text, tool_uses=[...])` +- `router_decision_response(destination, confidence=4, job_key=None)` +- `stream_events(text="", tool_uses=None)` +- `usage_block(input_tokens=100, output_tokens=50, cache_creation=0, cache_read=0)` + +No new runtime dep — `httpx.MockTransport` is built into httpx, which is already +in `poetry.lock`. + +### 3.3 Scripted multi-turn example + +```python +def test_planner_calls_workflow_then_job_agents(test_hooks_factory): + mock = MockAnthropicClient.script([ + router_decision_response("planner", confidence=5), + tool_use_response("call_workflow_agent", {"message": "create workflow"}), + tool_use_response("call_job_code_agent", {"message": "code for step", "job_key": "fetch"}), + text_response("All done."), + ]) + test_hooks = test_hooks_factory(anthropic=mock) + result = global_chat_main(make_global_chat_payload("create a workflow"), test_hooks) + + assert [c["tool"] for c in test_hooks["tool_calls"]] == [ + "router_decision", "call_workflow_agent", "call_job_code_agent", + ] +``` + +When a test wants to bypass a sub-agent's real code, it scripts responses for +the planner's `/v1/messages` calls and lets the sub-agent's own `main()` run +under the same mock client. Stub registries aren't needed for the common case. + +--- + +## 4. Tool-call breadcrumbs (`test_hooks["tool_calls"]`) + +`test_hooks["tool_calls"]` is a list the test allocates and production code +appends to. One helper in `testing/anthropic_mock.py`: + +```python +def record_tool_call(test_hooks: Optional[dict], entry: dict) -> None: + if test_hooks is None: + return + crumbs = test_hooks.get("tool_calls") + if crumbs is not None: + crumbs.append(entry) +``` + +Dispatch sites (`planner._execute_tool`, `router.route_and_execute`) call +`record_tool_call(test_hooks, {"tool": ..., "input": ...})`. Two dict lookups +per call when `test_hooks is None` — negligible. + +Tests read: + +```python +assert [c["tool"] for c in test_hooks["tool_calls"]] == ["router_decision", "call_workflow_agent"] +``` + +--- + +## 5. Tool stubs (`test_hooks["tool_stubs"]`) + +Most planner tools don't need stubbing. `call_workflow_agent` and +`call_job_code_agent` inherit `test_hooks` and run the sub-agent's own mocked +`main()`. `inspect_job_code` is pure local code with no network. The one tool +that does need stubbing is **`search_documentation`** — without a stub it would +hit Pinecone (vector store) and OpenAI (embeddings) on every service test. + +Production change in `services/global_chat/planner.py::_execute_tool` — one +if/else at the top of the dispatch: + +```python +stub = (self._test_hooks or {}).get("tool_stubs", {}).get(tool_use_block.name) +if stub is not None: + tool_result = stub(tool_use_block.input) +else: + # original dispatch by name follows + ... +``` + +Test usage: + +```python +test_hooks = { + "anthropic_http_client": mock.httpx_client, + "tool_calls": [], + "tool_stubs": { + "search_documentation": lambda tool_input: "Cron triggers run on a schedule...", + }, +} +result = main(payload, test_hooks) +``` + +The stub returns whatever shape the real tool returns (here a string — the +planner feeds it back into the next Anthropic call). + +A `build_search_documentation_stub(docs=[...])` helper in +`testing/anthropic_mock.py` can emerge once a second test reuses the same shape +— not preemptively. + +--- + +## 6. Directory layout + +``` +services//tests/ + __init__.py + conftest.py # re-exports shared fixtures; auto-marks by filename suffix + test__unit.py # tier 1 (unit-tests-architect) + test__service.py # tier 2 (this tier) + fixtures/ # per-service fixture data (optional) +``` + +Test filenames this tier will add (illustrative, not exhaustive): + +- `services/global_chat/tests/test_router_service.py` — router decisions by + intent. +- `services/global_chat/tests/test_planner_service.py` — tool dispatch order, + test_hooks propagation. +- `services/global_chat/tests/test_subagent_passthrough_service.py` — global → + workflow / job wiring. +- `services/workflow_chat/tests/test_workflow_chat_service.py` — YAML + extraction, retry loop, streaming events. +- `services/job_chat/tests/test_job_chat_service.py` — RAG injection (with + stubbed retriever), suggest-code response shape, page-prefix detection, + error-correction loop. + +Cross-service end-to-end flow tests (planner chain over mocks) also live under +`services/global_chat/tests/` since `global_chat` owns the planner. + +--- + +## 7. Shared helpers in `testing/` + +``` +testing/ + __init__.py + anthropic_mock.py # MockAnthropicClient, response builders, test_hooks-keys docstring, record_tool_call + fixtures.py # pytest fixtures + YAML assertion helpers + payload builders + loaders + fixtures/ + workflows/*.yaml + histories/*.json +``` + +`fixtures.py` is the flat home for: + +- `make_global_chat_payload`, `make_workflow_chat_payload`, + `make_job_chat_payload`. +- `get_workflow_yaml_attachment`, `get_suggested_code_attachment`, `get_usage`. +- `assert_yaml_has_ids`, `assert_yaml_jobs_have_body`, + `assert_yaml_equal_except`, `path_matches`, `assert_no_special_chars`. +- Pytest fixtures: `mock_anthropic`, `test_hooks_factory`, `fake_api_key`, + `sample_workflow_yaml`, `anthropic_client_no_network`. +- `set_unit_test_env` (dummy keys, disable langfuse/sentry). +- `load_fixture_json`, `load_fixture_yaml`. + +One file until it gets unwieldy (~500 lines). Split then, not pre-emptively. + +Key fixture: + +```python +@pytest.fixture +def test_hooks_factory(): + def _factory(*, anthropic=None, **overrides): + opts = {"tool_calls": []} + if anthropic is not None: + opts["anthropic_http_client"] = anthropic.httpx_client + opts.update(overrides) + return opts + return _factory +``` + +Per-service `conftest.py` just exposes `pytest_plugins = ["testing.fixtures"]` +(inherited from root) plus any per-service niche fixtures. + +--- + +## 8. Pytest configuration + +Owned initially by this tier in PR #1 (bootstrap). See overview §5 for the full +block. Relevant keys: + +```toml +[tool.pytest.ini_options] +pythonpath = ["services", "."] +testpaths = ["services"] +python_files = ["test_*.py"] +markers = [ + "unit: ...", + "service: main() tests with mocked LLM HTTP client", + "integration: ...", + "acceptance: ...", +] +addopts = ["-ra"] +``` + +Markers applied by filename suffix in the root `apollo/conftest.py` — authors +don't decorate manually. + +--- + +## 9. CI integration + +Service runs in the same `tests.yaml` workflow as unit, via +`pytest -m "unit or service"`. No secrets on this job (invariant). See overview +§6. + +--- + +## 10. Migration recipe for existing `pass_fail` tests + +1. **Classify the assertion.** Content-sensitive + (`"response mentions Salesforce"`) → integration or acceptance. Structural + (`"workflow_yaml has 2 jobs"`) → service (with a canned mock producing that + structure). +2. **Replace the call site.** Swap `subprocess.run([..., "entry.py", ...])` for + `from . import main; main(payload, test_hooks)`. +3. **Build the mock.** Hand-craft an Anthropic response fixture (or a script for + planner multi-turn) that produces the shape under test. +4. **Assert on structure + breadcrumbs.** Replace content asserts with routing / + shape asserts; keep content in acceptance. +5. **Delete the old test** once the new one is stable. + +Expect ~50–70% of `pass_fail` tests to become service tests; the rest stay in +integration. + +--- + +## 11. Production-code edits (summary) + +| File | Edit | +| ----------------------------------------- | ------------------------------------------------------------------------ | +| `services/global_chat/global_chat.py` | `main(data_dict, test_hooks=None)`; thread through | +| `services/workflow_chat/workflow_chat.py` | same | +| `services/job_chat/job_chat.py` | same | +| `services/global_chat/router.py` | accept `test_hooks` in `__init__`; pass into sub-agent calls | +| `services/global_chat/planner.py` | accept `test_hooks`; use in `_execute_tool`; thread into subagent_caller | +| `services/global_chat/subagent_caller.py` | accept `test_hooks`; pass to sub-agent `main()` | +| `services/util.py` | add `build_anthropic_client()` | + +Everywhere: backward-compatible defaults. `test_hooks is None` ⇒ existing +behaviour, byte-for-byte. + +--- + +## 12. Extensibility — new sub-agent or tool + +**New sub-agent:** + +1. `services/my_new_agent/my_new_agent.py` with + `def main(data, test_hooks=None)`. +2. Every `Anthropic(...)` site uses + `build_anthropic_client(api_key, test_hooks)`. +3. Thread `test_hooks` through internal calls. +4. Add `services/my_new_agent/tests/test_*_service.py`. Conftest auto-inherits. + +**New tool in the planner:** + +1. Append to `services/global_chat/tools/tool_definitions.py`. +2. Dispatch branch in `planner._execute_tool` calls + `record_tool_call(test_hooks, ...)`. +3. Write service tests. + +Pattern: **one arg, one call to `record_tool_call`, one test**. No framework +changes. + +--- + +## 13. What this tier deliberately does NOT do + +- **No sub-agent stub registry on day one.** Default behaviour (sub-agent runs + under shared mock client) is what tests usually want. +- **No `test_hooks["seed"]` / `["disable_langfuse"]` / `["scratch"]`.** Add when + a test fails without them. +- **No `pytest-asyncio`, `pytest-randomly`, or other dev deps.** Add when + needed. +- **No frozen public API contract between tiers.** Shared helpers live in one + package; rename when the signature improves. + +--- + +## 14. What else belongs in this tier + +Good service-test targets: + +- **API key threading** — payload `api_key` ends up in + `mock.last_request.headers["x-api-key"]`; absent → env var is used. +- **Cache-control regression** — planner system prompt has + `cache_control: {"type": "ephemeral"}`; assert on outbound request body. +- **Context-management beta** — planner sets `context-management-2025-06-27` + header and `context_management` field on every call. +- **History round-trip** — returned `history` equals input + this turn's + user/assistant messages. +- **`AdaptorSpecifier` propagation** — payload + `context.adaptor = "@openfn/language-http@3.1.11"` shows up in the prompt. +- **Retry loops** — `workflow_chat` retries once on YAML parse failure; script + invalid-then-valid and assert count. +- **Negative paths** — missing `content` → `ApolloError(400)`; malformed + tool-use response → graceful fallback. + +--- + +## Summary + +`test_hooks` second arg on `main()` + +`build_anthropic_client(api_key, test_hooks)` factory + `MockAnthropicClient` +with `always`/`script`/`streaming` constructors + three `test_hooks` keys +(`anthropic_http_client`, `tool_calls`, `tool_stubs` — the last only used for +`search_documentation` today). Three files in `testing/`, filename-suffix +markers, shared workflow with the unit tier. Add sub-agent stub infrastructure +the first time a test can't be written without it. diff --git a/platform/src/bridge.ts b/platform/src/bridge.ts index d5e13d0c..d735ed97 100644 --- a/platform/src/bridge.ts +++ b/platform/src/bridge.ts @@ -1,129 +1,468 @@ -import readline from "node:readline"; import path from "node:path"; -import { spawn } from "node:child_process"; -import { rm } from "node:fs/promises"; +import os from "node:os"; +import fs from "node:fs"; +import { spawn, type ChildProcess } from "node:child_process"; +import { timingSafeEqual } from "node:crypto"; import { getInternalToken } from "./auth/internal-token"; +import type { ApolloError } from "./util/errors"; import pkg from "../../package.json"; -/** - Run a python script - Each script will be run in its own thread because - 1) It saves script writers having to worry about long writing process - 2) Removes any risk of stale credentials and ensures a pristine environment - 3) it makes capturing logs a bit easier +/* + Long-lived Python worker manager. + + Replaces the per-request `poetry run python entry.py` spawn with a single + worker process connected to a Bun-owned Unix domain socket. Bun multiplexes + every job over one newline-delimited-JSON stream, demuxing by job_id. `run()` + keeps its original signature so the HTTP surface (services.ts) is unchanged; + both END (success) and ERROR (failure) resolve into an in-band value so + isApolloError() keeps driving the HTTP status — nothing is rejected. */ -export const run = async ( - scriptName: string, - port: number, // needed for self-calling services in pythonland - args: any = {}, - onLog?: (str: string) => void, - onEvent?: (type: string, payload: any /* string or json tbh */) => void -) => { - return new Promise(async (resolve, reject) => { - const id = crypto.randomUUID(); - - const tmpfile = path.resolve(`tmp/data/${id}-{}.json`); - - const inputPath = tmpfile.replace("{}", "input"); - const outputPath = tmpfile.replace("{}", "output"); - - // console.log("Initing input file at", inputPath); - await Bun.write(inputPath, JSON.stringify(args)); - - // console.log("Initing output file at", outputPath); - await Bun.write(outputPath, ""); - - const proc = spawn( - "poetry", - [ - "run", - "python", - "services/entry.py", - scriptName, - ...(inputPath ? ["--input", inputPath] : []), - ...(outputPath ? ["--output", outputPath] : []), - ...(port ? ["--port", `${port}`] : []), - ], - // Hand the internal token to the child explicitly so its apollo() self-calls - // are recognised by the auth hook. Spawned from here (the honest owner) rather than - // written back onto this process's env. - { - env: { - ...process.env, - APOLLO_INTERNAL_TOKEN: getInternalToken(), - APOLLO_VERSION: pkg.version, - }, + +const JOB_TIMEOUT_MS = Number(process.env.APOLLO_JOB_TIMEOUT_MS ?? 300_000); +const READY_TIMEOUT_MS = Number(process.env.APOLLO_READY_TIMEOUT_MS ?? 10_000); +const MAX_QUEUE = Number(process.env.APOLLO_MAX_QUEUE ?? 100); +const RESTART_BACKOFF_BASE_MS = 1_000; +const RESTART_BACKOFF_CAP_MS = 30_000; +const CIRCUIT_BREAKER_THRESHOLD = 5; +const SHUTDOWN_GRACE_MS = 5_000; + +const NEWLINE = 0x0a; + +type JobHandler = { + onLog?: (str: string) => void; + onEvent?: (type: string, payload: any) => void; + resolve: (value: any) => void; + timer: ReturnType; +}; + +// Module singleton state — one worker per Bun process. +let socketPath = ""; +let server: ReturnType | null = null; +let worker: ChildProcess | null = null; +let conn: any = null; +let connBuf = Buffer.alloc(0); +let outBuf = Buffer.alloc(0); // unflushed START bytes (backpressure) +const pendingJobs = new Map(); + +let workerReady: Promise = Promise.resolve(); +let resolveReady: () => void = () => {}; +let rejectReady: (e: any) => void = () => {}; +let isReady = false; +let waiters = 0; + +let restartCount = 0; // consecutive crash-near-startup events +let lastSpawnTime = 0; +let circuitOpen = false; +let shuttingDown = false; + +const apolloErrorValue = ( + code: number, + type: string, + message: string +): ApolloError => ({ code, type, message }); + +function resetReady() { + isReady = false; + workerReady = new Promise((res, rej) => { + resolveReady = res; + rejectReady = rej; + }); +} +resetReady(); + +function withTimeout(p: Promise, ms: number): Promise { + return new Promise((res, rej) => { + const t = setTimeout(() => rej(new Error("timeout")), ms); + p.then( + (v) => { + clearTimeout(t); + res(v); + }, + (e) => { + clearTimeout(t); + rej(e); } ); + }); +} - proc.on("error", async (err) => { - console.log(err); - }); +function resolveSocketPath(): string { + if (process.env.APOLLO_SOCKET_PATH) return process.env.APOLLO_SOCKET_PATH; + const xdg = process.env.XDG_RUNTIME_DIR; + const dir = xdg + ? path.join(xdg, "apollo") + : path.join(os.homedir(), ".apollo"); + return path.join(dir, "apollo.sock"); +} - const rl = readline.createInterface({ - input: proc.stdout, - crlfDelay: Infinity, - }); - rl.on("line", (line) => { - // Then divert any logs from a logger object to the websocket - if (/^(INFO|DEBUG|ERROR|WARNING)\:/.test(line)) { - // Divert the log line locally - console.log(line); - // TODO I'd love to break the log line up in to JSON actually - // { source, level, message } - onLog?.(line); - } else if (/^(EVENT)\:/.test(line)) { - // TODO does the event encoding need to be any more complex than this? - // Nice that it stays human readable - const [_prefix, type, ...payload] = line.split(":"); - let processedPayload = payload.join(":"); - try { - processedPayload = JSON.parse(processedPayload); - } catch (e) { - // No json, no problem - } - onEvent?.(type, processedPayload); - } - }); +function ensureSocketDir(p: string) { + const dir = path.dirname(p); + fs.mkdirSync(dir, { recursive: true, mode: 0o700 }); + try { + fs.chmodSync(dir, 0o700); + } catch { + // best effort; a pre-existing dir owned by us is the common case + } +} - const rl2 = readline.createInterface({ - input: proc.stderr, - crlfDelay: Infinity, - }); - rl2.on("line", (line) => { - console.error(line); - // /Divert all errors to the websocket - onLog?.(line); +// Before binding: if the path exists, probe it. A live answer means another +// instance owns it -> fail loudly. Refused means it's stale -> unlink. The +// 0700 parent dir (owned by us) closes the unlink->bind symlink TOCTOU. +async function clearStaleSocket(p: string): Promise { + if (!fs.existsSync(p)) return; + let live = false; + try { + const probe = await Bun.connect({ + unix: p, + socket: { data() {}, open(s) { s.end(); } }, }); + live = true; + try { + probe.end(); + } catch {} + } catch { + live = false; + } + if (live) { + throw new Error( + `Apollo socket ${p} is already in use by another instance; set a distinct APOLLO_SOCKET_PATH` + ); + } + fs.unlinkSync(p); +} - proc.on("close", async (code) => { - // Clean up readline interfaces immediately to prevent race conditions - rl.close(); - rl2.close(); +const socketHandlers = { + open(socket: any) { + if (conn) { + // Exactly one worker connection is accepted; extras are hung up. + console.warn("Apollo worker socket: rejecting an extra connection"); + try { + socket.end(); + } catch {} + return; + } + conn = socket; + connBuf = Buffer.alloc(0); + outBuf = Buffer.alloc(0); + }, + data(socket: any, chunk: Buffer) { + if (socket !== conn) return; + connBuf = connBuf.length ? Buffer.concat([connBuf, chunk]) : chunk; + let idx: number; + while ((idx = connBuf.indexOf(NEWLINE)) !== -1) { + const line = connBuf.subarray(0, idx); + connBuf = connBuf.subarray(idx + 1); + if (line.length) handleLine(line); + } + }, + drain(socket: any) { + if (socket !== conn || outBuf.length === 0) return; + const written = socket.write(outBuf); + outBuf = outBuf.subarray(written); + }, + close(socket: any) { + if (socket === conn) conn = null; + }, + error(_socket: any, err: any) { + console.error("Apollo worker socket error", err); + }, +}; + +function handleLine(line: Buffer) { + let msg: any; + try { + msg = JSON.parse(line.toString("utf-8")); + } catch { + const preview = line.toString("utf-8").slice(0, 200); + console.warn("Apollo worker: skipping malformed line:", preview); + return; + } + try { + dispatchMessage(msg); + } catch (e) { + console.warn("Apollo worker: error handling message", e); + } +} + +function validateToken(token: any): boolean { + if (typeof token !== "string") return false; + const a = Buffer.from(token); + const b = Buffer.from(getInternalToken()); + return a.length === b.length && timingSafeEqual(a, b); +} - if (code) { - console.error("Python process exited with code", code); - reject(code); +const formatLog = (msg: any): string => + `${msg.level}:${msg.source}:${msg.message}`; + +function dispatchMessage(msg: any) { + // Control channel (job_id null): the ready handshake. + if (msg.job_id == null) { + if (msg.type === "STATUS" && msg.data?.ready) { + if (isReady) return; + if (!validateToken(msg.token)) { + console.warn("Apollo worker ready handshake: invalid token, ignoring"); + return; } - const result = Bun.file(outputPath); - const text = await result.text(); + isReady = true; + restartCount = 0; // a clean startup resets the breaker + resolveReady(); + console.log("Apollo worker ready"); + } + return; + } + + const handler = pendingJobs.get(msg.job_id); + if (!handler) return; // unknown / already completed — drop + + switch (msg.type) { + case "LOG": + handler.onLog?.(formatLog(msg)); + break; + case "EVENT": + handler.onEvent?.(msg.event, msg.data); + break; + case "STATUS": + handler.onEvent?.("status", msg.data); + break; + case "ATTACHMENT": + handler.onEvent?.("attachment", { name: msg.name, data: msg.data }); + break; + case "END": + finishJob(msg.job_id, msg.result); + break; + case "ERROR": + finishJob(msg.job_id, { + code: msg.code ?? 500, + type: msg.error_type ?? "INTERNAL_ERROR", + message: msg.message ?? "Unknown error", + ...(msg.details ? { details: msg.details } : {}), + }); + break; + default: + break; + } +} +function finishJob(jobId: string, value: any) { + const handler = pendingJobs.get(jobId); + if (!handler) return; + clearTimeout(handler.timer); + pendingJobs.delete(jobId); + handler.resolve(value); +} + +function failAllPending(code: number, type: string, message: string) { + // Resolve (never reject) with a scrubbed error shape; handlers hold no payload, + // so nothing sensitive is logged here. + for (const handler of pendingJobs.values()) { + clearTimeout(handler.timer); + handler.resolve(apolloErrorValue(code, type, message)); + } + pendingJobs.clear(); +} + +// Serialize START writes; buffer whatever the socket can't take now and flush on +// drain. The worker's own single writer thread guarantees inbound framing. +function writeToWorker(line: string): boolean { + if (!conn) return false; + const data = Buffer.from(line, "utf-8"); + if (outBuf.length) { + outBuf = Buffer.concat([outBuf, data]); + return true; + } + try { + const written = conn.write(data); + if (written < data.length) outBuf = data.subarray(written); + return true; + } catch { + return false; + } +} + +function spawnWorker() { + if (shuttingDown || circuitOpen) return; + lastSpawnTime = Date.now(); + worker = spawn("poetry", ["run", "python", "services/worker.py"], { + // Same env the per-request spawn injected, so apollo() self-calls keep + // authenticating; plus the socket path. getInternalToken() is per-process + // stable, so restarts re-inject the same token (handshake + self-call auth). + env: { + ...process.env, + APOLLO_INTERNAL_TOKEN: getInternalToken(), + APOLLO_VERSION: pkg.version, + APOLLO_SOCKET_PATH: socketPath, + }, + stdio: ["ignore", "inherit", "inherit"], + }); + worker.on("error", (err) => { + console.error("Apollo worker spawn error", err); + }); + worker.on("exit", onWorkerExit); +} + +function onWorkerExit(code: number | null, signal: NodeJS.Signals | null) { + if (shuttingDown) return; + console.error( + `Apollo worker exited (code=${code}, signal=${signal}); failing ${pendingJobs.size} in-flight job(s)` + ); + worker = null; + conn = null; + connBuf = Buffer.alloc(0); + outBuf = Buffer.alloc(0); + + failAllPending(500, "INTERNAL_ERROR", "Apollo worker crashed"); + resetReady(); + + restartCount++; + if (restartCount >= CIRCUIT_BREAKER_THRESHOLD) { + circuitOpen = true; + console.error( + `Apollo worker crash loop: circuit breaker tripped after ${restartCount} crashes; serving 503` + ); + rejectReady(new Error("worker circuit open")); + return; + } + + const backoff = Math.min( + RESTART_BACKOFF_BASE_MS * 2 ** (restartCount - 1), + RESTART_BACKOFF_CAP_MS + ); + console.log( + `Apollo worker restarting in ${backoff}ms (attempt ${restartCount})` + ); + setTimeout(() => { + if (!shuttingDown && !circuitOpen) spawnWorker(); + }, backoff); +} + +/** Boot the worker: create/own the socket, then spawn the child. Idempotent. */ +export async function startWorker(): Promise { + if (server || worker) return; + socketPath = resolveSocketPath(); + ensureSocketDir(socketPath); + await clearStaleSocket(socketPath); + + const prevUmask = process.umask(0o077); + try { + server = Bun.listen({ unix: socketPath, socket: socketHandlers }); + } finally { + process.umask(prevUmask); + } + try { + fs.chmodSync(socketPath, 0o600); + } catch (e) { + console.warn("Apollo worker socket: could not chmod 0600", e); + } + + // Best-effort orphan kill if Bun exits without running the async shutdown. + process.on("exit", () => { + if (worker) { try { - await rm(inputPath); - await rm(outputPath); - } catch (e) { - console.error("Error removing temporary files"); - console.error(e); - } + worker.kill("SIGKILL"); + } catch {} + } + }); - if (text) { - resolve(JSON.parse(text)); - } else { - console.warn("No data returned from pythonland"); - resolve(null); - } + spawnWorker(); +} + +/** Stop the worker and release the socket (called from the server shutdown). */ +export async function stopWorker(): Promise { + shuttingDown = true; + const w = worker; + if (w) { + w.kill("SIGTERM"); + await new Promise((res) => { + const t = setTimeout(() => { + try { + w.kill("SIGKILL"); + } catch {} + res(); + }, SHUTDOWN_GRACE_MS); + w.on("exit", () => { + clearTimeout(t); + res(); + }); }); + } + try { + server?.stop(); + } catch {} + try { + if (socketPath && fs.existsSync(socketPath)) fs.unlinkSync(socketPath); + } catch {} +} - return; +/** + Run a python service on the long-lived worker. + + Signature unchanged from the per-request model. Resolves (never rejects) with + either the service result (END) or an ApolloError-shaped value (ERROR / timeout + / worker-down), so services.ts's isApolloError() drives the HTTP status. +*/ +export const run = async ( + scriptName: string, + port: number, // needed for self-calling services in pythonland + args: any = {}, + onLog?: (str: string) => void, + onEvent?: (type: string, payload: any /* string or json tbh */) => void +): Promise => { + if (circuitOpen) { + return apolloErrorValue( + 503, + "SERVICE_UNAVAILABLE", + "Apollo worker is unavailable" + ); + } + + if (!isReady) { + if (waiters >= MAX_QUEUE) { + return apolloErrorValue( + 503, + "SERVICE_UNAVAILABLE", + "Apollo worker queue is full" + ); + } + waiters++; + try { + await withTimeout(workerReady, READY_TIMEOUT_MS); + } catch { + return apolloErrorValue( + 503, + "SERVICE_UNAVAILABLE", + "Apollo worker is not ready" + ); + } finally { + waiters--; + } + } + + const jobId = crypto.randomUUID(); + + return new Promise((resolve) => { + const timer = setTimeout(() => { + pendingJobs.delete(jobId); + console.warn(`Apollo job ${jobId} timed out after ${JOB_TIMEOUT_MS}ms`); + resolve(apolloErrorValue(504, "TIMEOUT", "Job timed out")); + }, JOB_TIMEOUT_MS); + + pendingJobs.set(jobId, { onLog, onEvent, resolve, timer }); + + // Never log the START payload: it carries the resolved api_key. + const start = { + type: "START", + job_id: jobId, + service: scriptName, + payload: args, + port, + }; + const ok = writeToWorker(JSON.stringify(start) + "\n"); + if (!ok) { + clearTimeout(timer); + pendingJobs.delete(jobId); + resolve( + apolloErrorValue(503, "SERVICE_UNAVAILABLE", "Apollo worker unavailable") + ); + } }); }; diff --git a/platform/src/server.ts b/platform/src/server.ts index cbd8f3a1..0571e998 100644 --- a/platform/src/server.ts +++ b/platform/src/server.ts @@ -11,6 +11,7 @@ import { captureException } from "./util/sentry"; import { clientsDbUrl, closeDb } from "./db"; import { runMigrations } from "./db/migrate"; import { randomUUID } from "node:crypto"; +import { startWorker, stopWorker } from "./bridge"; import pkg from "../../package.json"; export default async ( @@ -56,11 +57,16 @@ export default async ( logInternalTokenProvenance(false); await auth.init(); - // No stop path exists otherwise; close the DB pool so a graceful pod termination - // (or Ctrl-C in dev) exits cleanly without orphaned Postgres connections. In-flight - // requests, open SSE streams, and spawned Python children are intentionally not - // drained — termination drops them rather than waiting them out. + // Boot the single long-lived Python worker: Bun creates/owns the socket, then + // spawns the child that connects to it. + await startWorker(); + + // No stop path exists otherwise; kill the worker (no orphan), close the DB pool + // so a graceful pod termination (or Ctrl-C in dev) exits cleanly without orphaned + // Postgres connections. In-flight requests and open SSE streams are intentionally + // not drained — termination drops them rather than waiting them out. const shutdown = async () => { + await stopWorker(); await closeDb(); process.exit(0); }; diff --git a/services/streaming_util.py b/services/streaming_util.py index c95ec700..b866ca82 100644 --- a/services/streaming_util.py +++ b/services/streaming_util.py @@ -12,6 +12,7 @@ from typing import Any from models import CLAUDE_SONNET +from util import get_job_writer # Shared status message pools for user-facing progress indicators. # Services compose from these to build context-specific pools. @@ -130,9 +131,15 @@ def _emit_event(self, event_type: str, data: dict[str, Any]) -> None: event_type: SSE event type (e.g., 'message_start', 'content_block_delta') data: Event data dictionary """ - # Use EVENT: prefix format that bridge.ts expects - # Bridge will convert this to proper SSE format - if self.stream: + if not self.stream: + return + # In the worker, the context writer sends an EVENT message over the socket + # (it attaches this job's job_id). Standalone (bun py / CLI) there is no + # writer, so fall back to the EVENT: stdout protocol bridge.ts parses. + writer = get_job_writer() + if writer is not None: + writer({"type": "EVENT", "event": event_type, "data": data}) + else: print(f"EVENT:{event_type}:{json.dumps(data)}", flush=True) # noqa: T201 def start_stream(self) -> None: diff --git a/services/util.py b/services/util.py index 45d6d0f7..cefcdac6 100644 --- a/services/util.py +++ b/services/util.py @@ -1,6 +1,8 @@ +import contextvars import logging import os import sys +from collections.abc import Callable from dataclasses import dataclass from typing import Any @@ -61,6 +63,61 @@ def to_dict(self) -> dict: loggers: dict[str, logging.Logger] = {} apollo_port = 3000 +# The job writer, when set, routes this job's logs/events to the Bun worker +# socket. Absent (standalone entry.py / CLI) they fall back to stdout. It lives +# in a ContextVar so a per-job copy_context() in the worker isolates it per job. +_job_writer: contextvars.ContextVar[Callable[[dict], None] | None] = contextvars.ContextVar( + "apollo_job_writer", default=None, +) + + +def set_job_writer(writer: Callable[[dict], None] | None) -> None: + """Bind the current context's job writer (worker socket sink).""" + _job_writer.set(writer) + + +def get_job_writer() -> Callable[[dict], None] | None: + """The current context's job writer, or None when running standalone.""" + return _job_writer.get() + + +class _JobLogHandler(logging.Handler): + """Single root handler: routes each record to the job writer when one is + bound (worker context), else to stdout (standalone). Exactly one emission + per record — never both.""" + + def emit(self, record: logging.LogRecord) -> None: + try: + writer = get_job_writer() + if writer is not None: + writer({ + "type": "LOG", + "level": record.levelname, + "source": record.name, + "message": record.getMessage(), + }) + else: + sys.stdout.write(self.format(record) + "\n") + sys.stdout.flush() + except Exception: + self.handleError(record) + + +_handler_installed = False + + +def _ensure_root_handler() -> None: + """Install the single job/stdout handler on the root logger, once.""" + global _handler_installed # noqa: PLW0603 + if _handler_installed: + return + root = logging.getLogger() + root.setLevel(logging.INFO) + handler = _JobLogHandler() + handler.setFormatter(logging.Formatter(logging.BASIC_FORMAT)) + root.addHandler(handler) + _handler_installed = True + def set_log_output(f: str | None) -> None: """Set the output file for logging.""" @@ -75,12 +132,12 @@ def set_log_output(f: str | None) -> None: def create_logger(name: str) -> logging.Logger: """ Create or retrieve a logger with the given name. - Logs to stdout by default. + Records route to the job writer (worker socket) when one is bound for the + current job, else to stdout. """ - logging.basicConfig(level=logging.INFO, stream=sys.stdout) + _ensure_root_handler() if name not in loggers: - logger = logging.getLogger(name) - loggers[name] = logger + loggers[name] = logging.getLogger(name) return loggers[name] diff --git a/services/worker.py b/services/worker.py new file mode 100644 index 00000000..3ad256ab --- /dev/null +++ b/services/worker.py @@ -0,0 +1,257 @@ +""" +Long-lived Python worker for the Apollo Bun server. + +Connects to the Bun-owned Unix domain socket (APOLLO_SOCKET_PATH), receives +newline-delimited-JSON START messages, and runs each job on its own thread. +LOG/EVENT/STATUS/END/ERROR messages are streamed back over the same socket via +a single writer thread (so NDJSON framing can never interleave). Replaces the +per-request `poetry run python entry.py` spawn; `entry.py` remains the standalone +`bun py` entrypoint. +""" + +import contextlib +import contextvars +import ctypes +import json +import os +import queue +import signal +import socket +import sys +import threading +import time +from collections.abc import Callable +from typing import Any + +import sentry_sdk +from dotenv import load_dotenv +from langfuse import Langfuse +from langfuse.span_filter import is_default_export_span +from opentelemetry.instrumentation.anthropic import AnthropicInstrumentor +from opentelemetry.instrumentation.threading import ThreadingInstrumentor +from util import ApolloError, set_apollo_port, set_job_writer + +# Langfuse/OTel init: once per process, before any Anthropic client is created. +# Moved here from entry.py (the worker is now the long-lived process; entry.py +# keeps its own copy for the standalone path). load_dotenv first so the Langfuse +# client and Sentry read the environment. +load_dotenv() +AnthropicInstrumentor().instrument() +ThreadingInstrumentor().instrument() + + +def _should_export_span(span: Any) -> bool: # noqa: ANN401 + """Drop spans marked as tracing-disabled (user has not opted in).""" + attrs = getattr(span, "attributes", None) or {} + if attrs.get("langfuse.trace.metadata.tracing_disabled") == "true": + return False + return is_default_export_span(span) + + +langfuse = Langfuse(should_export_span=_should_export_span, release=os.getenv("APOLLO_VERSION", "unknown")) + +_env = os.getenv("ENVIRONMENT", "unknown") +_trace_rates = { + "development": 1, + "staging": 0.05, + "production": 0.03, + "unknown": 0.0, +} + +sentry_sdk.init( + dsn=os.getenv("SENTRY_DSN"), + environment=_env, + sample_rate=1.0, + traces_sample_rate=_trace_rates.get(_env, 0.0), + enable_tracing=True, + auto_enabling_integrations=False, +) + +# Terminal/control messages must never be dropped; noisy per-job LOG/STATUS may be +# under backpressure. The ready handshake is a STATUS with job_id null and is +# treated as non-droppable. +_DROPPABLE = {"LOG", "STATUS"} +_QUEUE_MAXSIZE = 10000 + +_sock: socket.socket | None = None +_send_queue: "queue.Queue[dict | None]" = queue.Queue(maxsize=_QUEUE_MAXSIZE) + + +def _enqueue(msg: dict) -> None: + """Queue a framed message for the writer thread, applying the drop policy.""" + if msg.get("type") in _DROPPABLE and msg.get("job_id") is not None: + with contextlib.suppress(queue.Full): # droppable under pressure + _send_queue.put_nowait(msg) + else: + _send_queue.put(msg) + + +def _writer_loop() -> None: + """Single owner of the socket write end: one whole line per send.""" + while True: + msg = _send_queue.get() + if msg is None: + return + line = (json.dumps(msg) + "\n").encode("utf-8") + try: + _sock.sendall(line) + except OSError: + # Socket gone (parent died). Nothing more to do; read loop exits too. + return + + +def _make_writer(job_id: str) -> Callable[[dict], None]: + """Writer closure bound to one job — attaches job_id so callers need not.""" + def writer(msg: dict) -> None: + _enqueue({**msg, "job_id": job_id}) + return writer + + +def _run_job(job_id: str, service: str, payload: dict, port: int | None) -> None: + """Run one service under the current (per-job) context. Guarantees exactly + one terminal END/ERROR message.""" + set_job_writer(_make_writer(job_id)) + if port is not None: + set_apollo_port(port) + + sentry_sdk.set_tag("service", service) + try: + m = __import__(f"{service}.{service}", fromlist=["main"]) + result = m.main(payload) + _enqueue({"type": "END", "job_id": job_id, "result": result}) + except ApolloError as e: + sentry_sdk.capture_exception(e) + d = e.to_dict() + err = { + "type": "ERROR", + "job_id": job_id, + "code": d["code"], + "message": d["message"], + "error_type": d.get("type", "APOLLO_ERROR"), + } + if d.get("details"): + err["details"] = d["details"] + _enqueue(err) + except Exception as e: # ModuleNotFoundError etc. -> 500 + sentry_sdk.capture_exception(e) + _enqueue({ + "type": "ERROR", + "job_id": job_id, + "code": 500, + "message": str(e), + "error_type": "INTERNAL_ERROR", + }) + finally: + with contextlib.suppress(Exception): + langfuse.flush() + + +def _dispatch_job(job_id: str, service: str, payload: dict, port: int | None) -> None: + """Run the job in a fresh context so its writer (and any copy_context pools it + spawns) are isolated per job.""" + ctx = contextvars.copy_context() + ctx.run(_run_job, job_id, service, payload, port) + + +def _handle_message(msg: dict) -> None: + if msg.get("type") != "START": + return + job_id = msg.get("job_id") + service = msg.get("service") + payload = msg.get("payload") or {} + port = msg.get("port") + threading.Thread( + target=_dispatch_job, + args=(job_id, service, payload, port), + daemon=True, + ).start() + + +def _read_loop() -> None: + """Read NDJSON from the socket, dispatch each START on its own thread. Exits + (and ends the process) when the socket closes — i.e. Bun went away.""" + buf = b"" + while True: + try: + chunk = _sock.recv(65536) + except OSError: + break + if not chunk: + break # Bun closed the connection + buf += chunk + while b"\n" in buf: + line, buf = buf.split(b"\n", 1) + if not line.strip(): + continue + try: + msg = json.loads(line.decode("utf-8")) + except Exception: # skip a malformed/spliced line + continue + _handle_message(msg) + + +def _connect_with_retry(path: str, timeout: float = 30.0) -> socket.socket: + """Bun creates+binds the socket; give it a short grace to appear on startup.""" + deadline = time.time() + timeout + while True: + s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + s.connect(path) + return s + except OSError: + s.close() + if time.time() > deadline: + raise + time.sleep(0.1) + + +def _install_parent_death_signal() -> None: + """On Linux, ask the kernel to signal us when our parent dies (Bun crash with + no clean SIGTERM). Best-effort; the ppid watcher backstops it.""" + if sys.platform != "linux": + return + with contextlib.suppress(Exception): + libc = ctypes.CDLL("libc.so.6", use_errno=True) + PR_SET_PDEATHSIG = 1 # noqa: N806 + libc.prctl(PR_SET_PDEATHSIG, signal.SIGTERM) + + +def _watch_parent() -> None: + """Portable backstop: exit if we get reparented (parent died).""" + initial = os.getppid() + while True: + time.sleep(2) + if os.getppid() != initial: + os._exit(0) + + +def main() -> None: + _install_parent_death_signal() + + socket_path = os.environ.get("APOLLO_SOCKET_PATH") + if not socket_path: + print("APOLLO_SOCKET_PATH not set; worker cannot start", file=sys.stderr) # noqa: T201 + sys.exit(1) + + global _sock # noqa: PLW0603 + _sock = _connect_with_retry(socket_path) + + threading.Thread(target=_writer_loop, daemon=True).start() + threading.Thread(target=_watch_parent, daemon=True).start() + + # Ready handshake: carries the shared internal token so Bun can authenticate + # "this is the child I spawned" (constant-time compare) before going live. + _send_queue.put({ + "type": "STATUS", + "job_id": None, + "data": {"ready": True}, + "token": os.environ.get("APOLLO_INTERNAL_TOKEN", ""), + }) + + _read_loop() + # Socket closed -> parent gone. Exit hard so we don't hold keys or bill LLM calls. + os._exit(0) + + +if __name__ == "__main__": + main() From 9bcfc3276f43fbf5979a4e974c38a0d15d6f47e8 Mon Sep 17 00:00:00 2001 From: Joe Clark Date: Thu, 9 Jul 2026 09:50:27 +0100 Subject: [PATCH 2/2] tmp commit --- platform/src/bridge.ts | 4 + platform/test/fixtures/fake-worker.ts | 180 ++++++++++++++ platform/test/worker-bridge.test.ts | 297 ++++++++++++++++++++++ pyproject.toml | 11 + services/tests/__init__.py | 0 services/tests/test_worker_ipc.py | 340 ++++++++++++++++++++++++++ 6 files changed, 832 insertions(+) create mode 100644 platform/test/fixtures/fake-worker.ts create mode 100644 platform/test/worker-bridge.test.ts create mode 100644 services/tests/__init__.py create mode 100644 services/tests/test_worker_ipc.py diff --git a/platform/src/bridge.ts b/platform/src/bridge.ts index d735ed97..34e5651a 100644 --- a/platform/src/bridge.ts +++ b/platform/src/bridge.ts @@ -67,6 +67,10 @@ function resetReady() { resolveReady = res; rejectReady = rej; }); + // The circuit breaker rejects this even when no request is awaiting it; a no-op + // catch keeps that from surfacing as an unhandled rejection (real awaiters + // attach their own handler via withTimeout) + workerReady.catch(() => {}); } resetReady(); diff --git a/platform/test/fixtures/fake-worker.ts b/platform/test/fixtures/fake-worker.ts new file mode 100644 index 00000000..b9e00c29 --- /dev/null +++ b/platform/test/fixtures/fake-worker.ts @@ -0,0 +1,180 @@ +/* + Fake Apollo worker for the bridge.ts unit tests. + + Spawned by bridge.ts exactly like the real worker (via a PATH-shimmed `poetry` + in worker-bridge.test.ts), so it exercises the real spawn -> connect -> ready + handshake -> START/reply machinery WITHOUT a real Python worker or any LLM call. + + It connects to APOLLO_SOCKET_PATH, sends the ready handshake carrying the + injected APOLLO_INTERNAL_TOKEN, then replies to each START with scripted NDJSON + frames driven by the START payload's `mode`. Raw-write modes let a test control + socket chunking to exercise bridge's NDJSON reassembly and malformed-line + handling. +*/ + +const socketPath = process.env.APOLLO_SOCKET_PATH ?? ""; +const token = process.env.APOLLO_INTERNAL_TOKEN ?? ""; + +// Crash before the handshake so bridge counts a crash-near-startup (circuit +// breaker test). Nothing connects; bridge sees the child exit non-zero. +if (process.env.FAKE_CRASH_ON_BOOT === "1") { + process.exit(1); +} + +let sock: any = null; + +const sendLine = (obj: unknown) => sock.write(JSON.stringify(obj) + "\n"); + +function handleStart(msg: any) { + if (msg.type !== "START") return; + const jobId = msg.job_id; + const p = msg.payload ?? {}; + switch (p.mode) { + case "end": + sendLine({ type: "END", job_id: jobId, result: p.result ?? {} }); + break; + + case "error": + sendLine({ + type: "ERROR", + job_id: jobId, + code: p.code, + error_type: p.error_type, + message: p.message, + ...(p.details ? { details: p.details } : {}), + }); + break; + + case "log_then_end": + // Two whole messages in ONE socket write: bridge must split them on the + // newline (LOG forwarded, then END resolves). + sock.write( + JSON.stringify({ + type: "LOG", + job_id: jobId, + level: p.level, + source: p.source, + message: p.message, + }) + + "\n" + + JSON.stringify({ type: "END", job_id: jobId, result: p.result ?? {} }) + + "\n" + ); + break; + + case "event": + sendLine({ type: "EVENT", job_id: jobId, event: p.event, data: p.data }); + sendLine({ type: "END", job_id: jobId, result: {} }); + break; + + case "status": + sendLine({ type: "STATUS", job_id: jobId, data: p.data }); + sendLine({ type: "END", job_id: jobId, result: {} }); + break; + + case "attachment": + sendLine({ + type: "ATTACHMENT", + job_id: jobId, + name: p.name, + data: p.data, + }); + sendLine({ type: "END", job_id: jobId, result: {} }); + break; + + case "split_end": { + // ONE END frame split across two socket writes: bridge must buffer the + // partial line and reassemble it. + const s = + JSON.stringify({ type: "END", job_id: jobId, result: p.result ?? {} }) + + "\n"; + const mid = Math.max(1, Math.floor(s.length / 2)); + sock.write(s.slice(0, mid)); + setTimeout(() => sock.write(s.slice(mid)), 40); + break; + } + + case "malformed_then_end": + // A non-JSON line must be skipped without killing the reader/hanging jobs. + sock.write("this is not json at all\n"); + sendLine({ type: "END", job_id: jobId, result: p.result ?? {} }); + break; + + case "unknown_then_end": + // A frame for an unknown/already-completed job_id must be dropped. + sendLine({ type: "END", job_id: "00000000-0000-0000-0000-000000000000", result: {} }); + sendLine({ type: "END", job_id: jobId, result: p.result ?? {} }); + break; + + case "noreply": + // Send nothing so the per-job timeout fires. + break; + + case "crash": + // Exit while the job is pending so bridge fails it (and all pending) 500. + process.exit(1); + break; + + default: + sendLine({ type: "END", job_id: jobId, result: { echo: p } }); + } +} + +let buf = ""; +function onData(chunk: Buffer) { + buf += chunk.toString("utf-8"); + let i: number; + while ((i = buf.indexOf("\n")) !== -1) { + const line = buf.slice(0, i); + buf = buf.slice(i + 1); + if (line.trim()) { + try { + handleStart(JSON.parse(line)); + } catch { + // ignore anything unparseable from bridge (there shouldn't be any) + } + } + } +} + +async function connectWithRetry(): Promise { + const deadline = Date.now() + 10_000; + for (;;) { + try { + await Bun.connect({ + unix: socketPath, + socket: { + open(s) { + sock = s; + // Ready handshake: control message (job_id null) carrying the token. + s.write( + JSON.stringify({ + type: "STATUS", + job_id: null, + data: { ready: true }, + token, + }) + "\n" + ); + }, + data(_s, chunk) { + onData(chunk as Buffer); + }, + close() { + // Bun went away (server stopped) -> exit like the real worker does. + process.exit(0); + }, + error() {}, + }, + }); + return; + } catch { + if (Date.now() > deadline) throw new Error("fake worker: connect timed out"); + await new Promise((r) => setTimeout(r, 50)); + } + } +} + +await connectWithRetry(); + +// Keep the process alive to service jobs. +setInterval(() => {}, 1 << 30); diff --git a/platform/test/worker-bridge.test.ts b/platform/test/worker-bridge.test.ts new file mode 100644 index 00000000..e69e3ff2 --- /dev/null +++ b/platform/test/worker-bridge.test.ts @@ -0,0 +1,297 @@ +import { + afterAll, + beforeAll, + describe, + expect, + it, +} from "bun:test"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +/* + Unit tests for the bridge.ts long-lived-worker manager, driven by a FAKE worker + (platform/test/fixtures/fake-worker.ts) — no real Python worker, no LLM calls. + + ISOLATION: server.test.ts boots a REAL worker through the canonical bridge + module singleton, and Bun shares module state across files in one process. To + avoid colliding with it, we import a FRESH, cache-busted bridge instance (same + trick auth.startup.test.ts uses) with its OWN socket server + singleton state, + pointed at a temp socket path. The fake worker is spawned by bridge's real + `spawn("poetry", ...)` via a PATH shim that runs `bun fixtures/fake-worker.ts`. + + Env (APOLLO_JOB_TIMEOUT_MS etc.) is read by bridge at import time, so we set it + BEFORE the dynamic import and restore it after, keeping the canonical module + (whenever it loads for server.test.ts) on its defaults. +*/ + +type Bridge = typeof import("../src/bridge"); + +const FIXTURE = path.resolve(import.meta.dir, "fixtures/fake-worker.ts"); + +// A fresh bridge instance with its own singleton state, its own temp socket, and +// a PATH-shimmed `poetry` that launches the fake worker. Returns a teardown fn. +async function makeBridge(opts: { + jobTimeoutMs?: number; + readyTimeoutMs?: number; + maxQueue?: number; + crashOnBoot?: boolean; +} = {}): Promise<{ bridge: Bridge; teardown: () => Promise }> { + const tmpRoot = fs.mkdtempSync(path.join(os.tmpdir(), "apollo-worker-test-")); + const shimDir = path.join(tmpRoot, "bin"); + fs.mkdirSync(shimDir, { recursive: true }); + const socket = path.join(tmpRoot, "apollo.sock"); + + // A `poetry` shim: bridge runs `spawn("poetry", ["run","python",...])`; the + // shim ignores those args and launches the fake worker under this Bun binary. + const shim = path.join(shimDir, "poetry"); + fs.writeFileSync(shim, `#!/bin/sh\nexec "${process.execPath}" "${FIXTURE}"\n`); + fs.chmodSync(shim, 0o755); + + const saved: Record = { + PATH: process.env.PATH, + APOLLO_SOCKET_PATH: process.env.APOLLO_SOCKET_PATH, + APOLLO_JOB_TIMEOUT_MS: process.env.APOLLO_JOB_TIMEOUT_MS, + APOLLO_READY_TIMEOUT_MS: process.env.APOLLO_READY_TIMEOUT_MS, + APOLLO_MAX_QUEUE: process.env.APOLLO_MAX_QUEUE, + FAKE_CRASH_ON_BOOT: process.env.FAKE_CRASH_ON_BOOT, + }; + + process.env.PATH = `${shimDir}:${process.env.PATH}`; + process.env.APOLLO_SOCKET_PATH = socket; + process.env.APOLLO_JOB_TIMEOUT_MS = String(opts.jobTimeoutMs ?? 1500); + process.env.APOLLO_READY_TIMEOUT_MS = String(opts.readyTimeoutMs ?? 4000); + process.env.APOLLO_MAX_QUEUE = String(opts.maxQueue ?? 100); + if (opts.crashOnBoot) process.env.FAKE_CRASH_ON_BOOT = "1"; + else delete process.env.FAKE_CRASH_ON_BOOT; + + // Fresh, isolated bridge module (own singleton state). + const bridge: Bridge = await import( + `../src/bridge?bust=${Date.now()}_${Math.random()}` + ); + await bridge.startWorker(); + + const teardown = async () => { + try { + await bridge.stopWorker(); + } catch {} + for (const [k, v] of Object.entries(saved)) { + if (v === undefined) delete process.env[k]; + else process.env[k] = v; + } + try { + fs.rmSync(tmpRoot, { recursive: true, force: true }); + } catch {} + }; + + return { bridge, teardown }; +} + +describe("bridge worker manager (fake worker)", () => { + let bridge: Bridge; + let teardown: () => Promise; + + beforeAll(async () => { + ({ bridge, teardown } = await makeBridge()); + }); + + afterAll(async () => { + await teardown(); + }); + + describe("handshake + terminal resolution", () => { + it("awaits the ready handshake and resolves END result as-is", async () => { + // run() gates on the ready handshake; the fake worker sends it on connect. + const result = await bridge.run("fake", 0, { + mode: "end", + result: { hello: "world", n: 42 }, + }); + expect(result).toEqual({ hello: "world", n: 42 }); + }); + + it("resolves an ERROR frame as an ApolloError-shaped value with code/type/message/details", async () => { + // Regression C fixed: details must survive the ERROR -> resolved-value hop. + const result = await bridge.run("fake", 0, { + mode: "error", + code: 429, + error_type: "RATE_LIMIT", + message: "slow down", + details: { retry_after: 60 }, + }); + expect(result).toEqual({ + code: 429, + type: "RATE_LIMIT", + message: "slow down", + details: { retry_after: 60 }, + }); + }); + + it("defaults an ERROR with missing fields to 500/INTERNAL_ERROR", async () => { + const result = await bridge.run("fake", 0, { mode: "error" }); + expect(result.code).toBe(500); + expect(result.type).toBe("INTERNAL_ERROR"); + expect(typeof result.message).toBe("string"); + expect(result.details).toBeUndefined(); + }); + }); + + describe("event routing", () => { + it("forwards LOG to onLog as the exact level:source:message string", async () => { + const logs: string[] = []; + const result = await bridge.run( + "fake", + 0, + { + mode: "log_then_end", + level: "INFO", + source: "echo", + message: "Echoing request", + result: { done: true }, + }, + (s) => logs.push(s) + ); + expect(logs).toEqual(["INFO:echo:Echoing request"]); + expect(result).toEqual({ done: true }); + }); + + it("forwards EVENT to onEvent(event, data)", async () => { + const events: Array<[string, any]> = []; + await bridge.run( + "fake", + 0, + { mode: "event", event: "content_block_delta", data: { index: 1 } }, + undefined, + (type, payload) => events.push([type, payload]) + ); + expect(events).toContainEqual(["content_block_delta", { index: 1 }]); + }); + + it("forwards STATUS to onEvent('status', data)", async () => { + const events: Array<[string, any]> = []; + await bridge.run( + "fake", + 0, + { mode: "status", data: { message: "Working on it..." } }, + undefined, + (type, payload) => events.push([type, payload]) + ); + expect(events).toContainEqual(["status", { message: "Working on it..." }]); + }); + + it("forwards ATTACHMENT to onEvent('attachment', {name, data})", async () => { + const events: Array<[string, any]> = []; + await bridge.run( + "fake", + 0, + { mode: "attachment", name: "chart.png", data: "base64==" }, + undefined, + (type, payload) => events.push([type, payload]) + ); + expect(events).toContainEqual([ + "attachment", + { name: "chart.png", data: "base64==" }, + ]); + }); + }); + + describe("reader / NDJSON framing / demux", () => { + it("reassembles a message split across two socket chunks", async () => { + const result = await bridge.run("fake", 0, { + mode: "split_end", + result: { split: true }, + }); + expect(result).toEqual({ split: true }); + }); + + it("splits two messages delivered in one chunk (LOG + END)", async () => { + const logs: string[] = []; + const result = await bridge.run( + "fake", + 0, + { + mode: "log_then_end", + level: "WARN", + source: "svc", + message: "one chunk", + result: { ok: 1 }, + }, + (s) => logs.push(s) + ); + expect(logs).toEqual(["WARN:svc:one chunk"]); + expect(result).toEqual({ ok: 1 }); + }); + + it("skips a malformed (non-JSON) line without hanging the job or others", async () => { + // The malformed job still completes, AND a concurrent normal job is not + // stalled by the bad line — proving the reader loop survived it. + const [bad, good] = await Promise.all([ + bridge.run("fake", 0, { mode: "malformed_then_end", result: { id: 1 } }), + bridge.run("fake", 0, { mode: "end", result: { id: 2 } }), + ]); + expect(bad).toEqual({ id: 1 }); + expect(good).toEqual({ id: 2 }); + }); + + it("drops a frame for an unknown/already-completed job_id", async () => { + // The worker emits an END for a bogus job_id before the real one; the + // bogus is dropped and the real job resolves normally. + const result = await bridge.run("fake", 0, { + mode: "unknown_then_end", + result: { real: true }, + }); + expect(result).toEqual({ real: true }); + }); + + it("multiplexes several concurrent jobs to distinct results", async () => { + const results = await Promise.all( + Array.from({ length: 8 }, (_, i) => + bridge.run("fake", 0, { mode: "end", result: { i } }) + ) + ); + expect(results).toEqual(Array.from({ length: 8 }, (_, i) => ({ i }))); + }); + }); + + describe("per-job timeout", () => { + it("resolves a 504 ApolloError when the worker never replies", async () => { + // jobTimeoutMs is 1500 for this bridge; the fake sends nothing. + const result = await bridge.run("fake", 0, { mode: "noreply" }); + expect(result).toEqual({ + code: 504, + type: "TIMEOUT", + message: "Job timed out", + }); + }, 6000); + }); + + describe("crash + restart", () => { + it("fails in-flight jobs with 500 on worker exit, then restarts and serves new jobs", async () => { + // Crash while a job is pending -> that job resolves a scrubbed 500. + const crashed = await bridge.run("fake", 0, { mode: "crash" }); + expect(crashed.code).toBe(500); + expect(crashed.type).toBe("INTERNAL_ERROR"); + + // The worker restarts (backoff) and re-handshakes; a subsequent job awaits + // readiness (the queue/await path) and then succeeds — proving the pending + // map was cleared and the restart path re-armed readiness. + const after = await bridge.run("fake", 0, { + mode: "end", + result: { recovered: true }, + }); + expect(after).toEqual({ recovered: true }); + }, 15000); + }); +}); + +// NOTE — circuit-breaker trip is deliberately NOT driven live here. Tripping it +// requires K=5 consecutive crash-near-startup exits (backoff is a fixed 1s..cap in +// bridge, ~15s wall clock), and at the trip bridge calls +// `rejectReady(new Error("worker circuit open"))` on a workerReady promise it +// created in the SAME synchronous tick — so nothing can ever be awaiting it +// (run() short-circuits on `circuitOpen` instead). That rejection is therefore +// always unattached, and Bun's test runner fails any test that observes an +// unhandled rejection regardless of a process listener. So a live trip test can't +// be made green without a production change. The restart/backoff path up to the +// breaker is covered by the crash+restart test above; the trip itself is verified +// by code-trace + Agent A/C. See agentE-report.md (flagged as a possible minor +// bridge robustness fix: attach a no-op .catch in resetReady()). diff --git a/pyproject.toml b/pyproject.toml index 896548e5..962f8a3c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,6 +51,7 @@ testpaths = [ "services/latest_adaptors/tests", "services/search_docsite/tests", "services/tools", + "services/tests", ] python_files = ["test_*.py"] @@ -114,3 +115,13 @@ ignore = [ ] line-length = 120 target-version = "py311" + +[tool.ruff.per-file-ignores] +# Test suites don't need full annotation ceremony or named constants for the +# literal values they assert against. +"services/tests/**" = [ + "ANN001", # missing type annotation for function argument (fixtures/helpers) + "ANN201", # missing return type on public test functions + "ANN202", # missing return type on private helpers + "PLR2004", # magic value used in comparison (expected values in asserts) +] diff --git a/services/tests/__init__.py b/services/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/services/tests/test_worker_ipc.py b/services/tests/test_worker_ipc.py new file mode 100644 index 00000000..3d96da3d --- /dev/null +++ b/services/tests/test_worker_ipc.py @@ -0,0 +1,340 @@ +"""Unit tests for the Python side of the Bun<->Python Unix-socket worker. + +Covers the shared logging/streaming plumbing (util.create_logger, the job-writer +ContextVar, StreamManager._emit_event) and worker.py's single-writer framing + +per-job terminal handling. All deterministic; no LLM/network calls. + +Run from repo root: + poetry run pytest services/tests/test_worker_ipc.py +""" +import contextvars +import json +import queue +import socket +import sys +import threading +import types +from collections.abc import Callable + +import pytest +import streaming_util +import util +import worker +from util import ApolloError + +pytestmark = pytest.mark.unit + + +@pytest.fixture(autouse=True) +def _reset_job_writer(): + """No writer leaks between tests (the writer is a process-wide ContextVar in + the test's own context).""" + util.set_job_writer(None) + yield + util.set_job_writer(None) + + +# --------------------------------------------------------------------------- # +# util.create_logger — single emission path (socket XOR stdout, never both) +# --------------------------------------------------------------------------- # + +class TestCreateLogger: + def test_writer_present_enqueues_log_and_skips_stdout(self, capsys): + captured: list[dict] = [] + util.set_job_writer(captured.append) + + log = util.create_logger("job_chat") + log.info("hello from a job") + + assert len(captured) == 1 + msg = captured[0] + assert msg["type"] == "LOG" + assert msg["level"] == "INFO" + assert msg["source"] == "job_chat" + assert msg["message"] == "hello from a job" + # Must NOT also hit stdout (no double-emit). + assert "hello from a job" not in capsys.readouterr().out + + def test_no_writer_writes_stdout_exactly_once(self, capsys): + util.set_job_writer(None) + + log = util.create_logger("standalone_svc") + log.info("unique-stdout-token-xyz") + + out = capsys.readouterr().out + assert out.count("unique-stdout-token-xyz") == 1 + + def test_emit_never_raises_when_writer_absent(self): + util.set_job_writer(None) + log = util.create_logger("noraise") + # Should not raise regardless of content. + log.info("plain") + log.warning("with %s formatting", "arg") + + +# --------------------------------------------------------------------------- # +# set_job_writer / get_job_writer — contextvar isolation +# --------------------------------------------------------------------------- # + +class TestJobWriterIsolation: + def test_writer_set_in_copied_context_does_not_leak_out(self): + util.set_job_writer(None) + seen: dict = {} + + def inner_a(): + util.set_job_writer(lambda _m: None) + seen["a"] = util.get_job_writer() + + def inner_b(): + seen["b"] = util.get_job_writer() + + contextvars.copy_context().run(inner_a) + seen["outer_after_a"] = util.get_job_writer() + contextvars.copy_context().run(inner_b) + + assert seen["a"] is not None # writer visible inside its context + assert seen["outer_after_a"] is None # did not leak back to the parent + assert seen["b"] is None # nor into a sibling context + + def test_two_contexts_carry_independent_writers(self): + util.set_job_writer(None) + got: dict = {} + + def make(tag): + # A distinct writer object per context so identity proves isolation. + writer = {"tag": tag} + def run(): + util.set_job_writer(writer.__setitem__) + got[tag] = util.get_job_writer() + return run + + ctx1 = contextvars.copy_context() + ctx2 = contextvars.copy_context() + ctx1.run(make("one")) + ctx2.run(make("two")) + + assert got["one"] is not got["two"] + assert util.get_job_writer() is None + + +# --------------------------------------------------------------------------- # +# streaming_util.StreamManager._emit_event +# --------------------------------------------------------------------------- # + +class TestEmitEvent: + def test_writer_present_sends_event_shape_with_job_id_from_closure(self, monkeypatch): + # The writer closure (worker._make_writer) is what attaches job_id; + # _emit_event itself only supplies type/event/data. Wire the real closure + # to a captured _enqueue to lock in that contract. + captured: list[dict] = [] + monkeypatch.setattr(worker, "_enqueue", captured.append) + util.set_job_writer(worker._make_writer("job-123")) + + sm = streaming_util.StreamManager(stream=True) + sm._emit_event("content_block_delta", {"index": 2}) + + assert captured == [ + { + "type": "EVENT", + "event": "content_block_delta", + "data": {"index": 2}, + "job_id": "job-123", + }, + ] + + def test_no_writer_falls_back_to_event_stdout_protocol(self, capsys): + util.set_job_writer(None) + sm = streaming_util.StreamManager(stream=True) + sm._emit_event("message_stop", {"type": "message_stop"}) + + out = capsys.readouterr().out + assert 'EVENT:message_stop:{"type": "message_stop"}' in out + + def test_no_writer_never_raises(self): + util.set_job_writer(None) + sm = streaming_util.StreamManager(stream=True) + sm._emit_event("evt", {"a": 1}) # must not raise + + def test_non_streaming_manager_emits_nothing(self, capsys, monkeypatch): + captured: list[dict] = [] + monkeypatch.setattr(worker, "_enqueue", captured.append) + util.set_job_writer(worker._make_writer("j")) + + sm = streaming_util.StreamManager(stream=False) + sm._emit_event("evt", {"a": 1}) + + assert captured == [] + assert capsys.readouterr().out == "" + + +# --------------------------------------------------------------------------- # +# worker._writer_loop — single writer serializes concurrent producers +# --------------------------------------------------------------------------- # + +class TestWriterThreadFraming: + def test_concurrent_enqueue_yields_whole_json_lines(self, monkeypatch): + # Drive the real _enqueue/_writer_loop over a socketpair from many threads + # and assert every received line is a complete, standalone JSON object — + # i.e. no interleaving across the single writer. + a, b = socket.socketpair() + fresh_q: queue.Queue = queue.Queue(maxsize=100_000) + monkeypatch.setattr(worker, "_sock", a) + monkeypatch.setattr(worker, "_send_queue", fresh_q) + + received = bytearray() + + def reader(): + while True: + chunk = b.recv(65536) + if not chunk: + break + received.extend(chunk) + + rt = threading.Thread(target=reader) + rt.start() + wt = threading.Thread(target=worker._writer_loop) + wt.start() + + n_threads, per_thread = 8, 250 + + def producer(tid: int): + for i in range(per_thread): + # EVENT is non-droppable, so nothing is lost under the drop policy. + worker._enqueue({ + "type": "EVENT", + "job_id": f"job-{tid}", + "event": "e", + "data": {"tid": tid, "i": i, "pad": "x" * 64}, + }) + + producers = [threading.Thread(target=producer, args=(t,)) for t in range(n_threads)] + for p in producers: + p.start() + for p in producers: + p.join() + + fresh_q.put(None) # stop the writer loop + wt.join() + a.close() # EOF for the reader + rt.join() + b.close() + + text = received.decode("utf-8") + assert text.endswith("\n") + lines = [ln for ln in text.split("\n") if ln] + assert len(lines) == n_threads * per_thread + # Every line parses on its own -> no framing interleave. + for ln in lines: + obj = json.loads(ln) + assert obj["type"] == "EVENT" + + +# --------------------------------------------------------------------------- # +# worker._run_job — exactly one terminal message per job +# --------------------------------------------------------------------------- # + +def _install_fake_service(name: str, main_fn: Callable[[dict], dict]) -> None: + """Register a fake `.` module so worker's + __import__(f"{name}.{name}", fromlist=["main"]) resolves to it.""" + pkg = types.ModuleType(name) + pkg.__path__ = [] # mark as a package + sub = types.ModuleType(f"{name}.{name}") + sub.main = main_fn + setattr(pkg, name, sub) + sys.modules[name] = pkg + sys.modules[f"{name}.{name}"] = sub + + +def _run_job_collecting(monkeypatch, service: str, payload: dict) -> list[dict]: + """Run _run_job in a fresh context (as production does) with _enqueue and + langfuse stubbed; return every enqueued message.""" + captured: list[dict] = [] + monkeypatch.setattr(worker, "_enqueue", captured.append) + monkeypatch.setattr(worker, "langfuse", types.SimpleNamespace(flush=lambda: None)) + contextvars.copy_context().run(worker._run_job, "job-1", service, payload, None) + return captured + + +class TestRunJobTerminal: + def test_success_sends_exactly_one_end(self, monkeypatch): + _install_fake_service("svc_ok", lambda payload: {"result": payload["x"] + 1}) + msgs = _run_job_collecting(monkeypatch, "svc_ok", {"x": 41}) + + terminals = [m for m in msgs if m["type"] in ("END", "ERROR")] + assert len(terminals) == 1 + end = terminals[0] + assert end["type"] == "END" + assert end["job_id"] == "job-1" + assert end["result"] == {"result": 42} + + def test_apollo_error_sends_one_error_from_to_dict_with_details(self, monkeypatch): + def boom(_payload): + raise ApolloError( + 429, "slow down", type="RATE_LIMIT", details={"retry_after": 60}, + ) + + _install_fake_service("svc_apollo", boom) + msgs = _run_job_collecting(monkeypatch, "svc_apollo", {}) + + terminals = [m for m in msgs if m["type"] in ("END", "ERROR")] + assert len(terminals) == 1 + err = terminals[0] + assert err["type"] == "ERROR" + assert err["code"] == 429 + assert err["error_type"] == "RATE_LIMIT" + assert err["message"] == "slow down" + assert err["details"] == {"retry_after": 60} + + def test_generic_exception_sends_one_500_internal_error(self, monkeypatch): + def boom(_payload): + raise ValueError("kaboom") + + _install_fake_service("svc_boom", boom) + msgs = _run_job_collecting(monkeypatch, "svc_boom", {}) + + terminals = [m for m in msgs if m["type"] in ("END", "ERROR")] + assert len(terminals) == 1 + err = terminals[0] + assert err["type"] == "ERROR" + assert err["code"] == 500 + assert err["error_type"] == "INTERNAL_ERROR" + assert err["message"] == "kaboom" + + def test_missing_service_module_sends_one_500(self, monkeypatch): + # ModuleNotFoundError from __import__ must map to a single 500 ERROR. + msgs = _run_job_collecting(monkeypatch, "no_such_service_xyz", {}) + terminals = [m for m in msgs if m["type"] in ("END", "ERROR")] + assert len(terminals) == 1 + assert terminals[0]["type"] == "ERROR" + assert terminals[0]["code"] == 500 + assert terminals[0]["error_type"] == "INTERNAL_ERROR" + + +class TestEnqueueDropPolicy: + def test_per_job_log_is_dropped_when_queue_full(self, monkeypatch): + full_q: queue.Queue = queue.Queue(maxsize=1) + full_q.put({"type": "seed"}) # occupy the only slot + monkeypatch.setattr(worker, "_send_queue", full_q) + + # A per-job LOG is droppable, so it is dropped rather than blocking. + worker._enqueue({"type": "LOG", "job_id": "j", "message": "noise"}) + + assert full_q.get_nowait() == {"type": "seed"} # LOG never landed + assert full_q.empty() + + def test_end_is_non_droppable_and_blocks_until_space(self, monkeypatch): + full_q: queue.Queue = queue.Queue(maxsize=1) + full_q.put({"type": "seed"}) + monkeypatch.setattr(worker, "_send_queue", full_q) + + # END must not be dropped: the enqueue blocks until a slot frees up. + def free_slot() -> None: + threading.Event().wait(0.05) + full_q.get() + + drainer = threading.Thread(target=free_slot) + drainer.start() + worker._enqueue({"type": "END", "job_id": "j", "result": {}}) # blocks, then lands + drainer.join() + + assert full_q.get_nowait() == {"type": "END", "job_id": "j", "result": {}}