diff --git a/.changeset/cancel-abandoned-runs.md b/.changeset/cancel-abandoned-runs.md new file mode 100644 index 00000000..b29cee49 --- /dev/null +++ b/.changeset/cancel-abandoned-runs.md @@ -0,0 +1,6 @@ +--- +"apollo": patch +--- + +Stop a service run when the client disconnects, so an abandoned request no +longer keeps calling the model diff --git a/.changeset/echo-mask-payload.md b/.changeset/echo-mask-payload.md new file mode 100644 index 00000000..29f5e4e2 --- /dev/null +++ b/.changeset/echo-mask-payload.md @@ -0,0 +1,9 @@ +--- +"apollo": patch +--- + +Mask sensitive values on their way out of a service, rather than relying on +each one to remember: service loggers mask what they emit, echo masks what it +returns, and the error envelope masks the exception text. The shared mask now +covers every field the server may fill in, and no longer matches ordinary +hyphenated words diff --git a/.changeset/restore-sse-idle-timeout.md b/.changeset/restore-sse-idle-timeout.md new file mode 100644 index 00000000..9a2a65c1 --- /dev/null +++ b/.changeset/restore-sse-idle-timeout.md @@ -0,0 +1,6 @@ +--- +"apollo": patch +--- + +Raise the server's socket idle timeout back to 255s so long-running SSE +streams are no longer cut off while a model is thinking diff --git a/.changeset/sse-heartbeat.md b/.changeset/sse-heartbeat.md new file mode 100644 index 00000000..6d2e717d --- /dev/null +++ b/.changeset/sse-heartbeat.md @@ -0,0 +1,6 @@ +--- +"apollo": patch +--- + +Send a periodic keepalive on streaming responses so a stream that is working +but quiet is no longer mistaken for a dead one diff --git a/.changeset/typed-service-failures.md b/.changeset/typed-service-failures.md new file mode 100644 index 00000000..6d6cba61 --- /dev/null +++ b/.changeset/typed-service-failures.md @@ -0,0 +1,6 @@ +--- +"apollo": patch +--- + +Report service failures with a real code and message instead of "Unknown error", +and stop a failed spawn from hanging the request diff --git a/CLAUDE.md b/CLAUDE.md index 25b2e406..d6161135 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -178,7 +178,8 @@ SSE to clients. against the `apollo-mappings` Pinecone index (collections `loinc-mappings-v2`, `snomed-mappings`). `embed_loinc_dataset` / `embed_snomed_dataset` populate it. - `status/` - Health check: validates Anthropic, OpenAI and Pinecone keys. -- `echo/` - Test service that returns its input; useful for verifying the server +- `echo/` - Test service that returns its input (with server-set values such as + `api_key` masked); useful for verifying the server pipeline. Note: there are **three distinct Pinecone indexes** — `docsite` (OpenFn docs), diff --git a/platform/src/bridge.ts b/platform/src/bridge.ts index d5e13d0c..240e9ade 100644 --- a/platform/src/bridge.ts +++ b/platform/src/bridge.ts @@ -1,10 +1,22 @@ import readline from "node:readline"; import path from "node:path"; import { spawn } from "node:child_process"; -import { rm } from "node:fs/promises"; +import { chmod, rm } from "node:fs/promises"; import { getInternalToken } from "./auth/internal-token"; +import { + emptyResult, + malformedResult, + subprocessCancelled, + subprocessFailed, + subprocessKilled, + subprocessSpawnFailed, +} from "./util/errors"; import pkg from "../../package.json"; +// A line a service logged on purpose, as opposed to whatever else lands on a +// stream. Only these are forwarded to the caller. +const LOG_LINE = /^(INFO|DEBUG|ERROR|WARNING):/; + /** Run a python script Each script will be run in its own thread because @@ -17,22 +29,39 @@ export const run = async ( 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 + onEvent?: (type: string, payload: any /* string or json tbh */) => void, + // Aborted when the client goes away + signal?: AbortSignal ) => { - return new Promise(async (resolve, reject) => { - const id = crypto.randomUUID(); + const id = crypto.randomUUID(); - const tmpfile = path.resolve(`tmp/data/${id}-{}.json`); + const tmpfile = path.resolve(`tmp/data/${id}-{}.json`); - const inputPath = tmpfile.replace("{}", "input"); - const outputPath = tmpfile.replace("{}", "output"); + const inputPath = tmpfile.replace("{}", "input"); + const outputPath = tmpfile.replace("{}", "output"); - // console.log("Initing input file at", inputPath); + // Outside the promise, deliberately. The Promise constructor only catches a + // synchronous throw from its executor, so an await that rejects in there - + // a full disk, a read-only tmp - leaves the promise pending for ever and + // the caller's stream open. Out here, run() is async and simply rejects. + try { await Bun.write(inputPath, JSON.stringify(args)); - // console.log("Initing output file at", outputPath); + // The payload can hold values that belong to the deployment rather than + // the caller, and only the close handler removes this file - so a process + // that dies first leaves one behind. + await chmod(inputPath, 0o600); + await Bun.write(outputPath, ""); + } catch (error) { + // Removed rather than left behind by a half-finished setup, for the same + // reason it is 0600 above. + await rm(inputPath).catch(() => {}); + await rm(outputPath).catch(() => {}); + throw subprocessSpawnFailed(scriptName, error); + } + return new Promise((resolve, reject) => { const proc = spawn( "poetry", [ @@ -56,17 +85,46 @@ export const run = async ( } ); - proc.on("error", async (err) => { - console.log(err); + // Nothing was spawned, so no "close" is coming - without settling here the + // request stays open until something upstream gives up + proc.on("error", (err) => { + console.error("Failed to start python process", err); + reject(subprocessSpawnFailed(scriptName, err)); }); + // `poetry run` execs into python rather than forking it, so this pid is the + // interpreter and a plain signal reaches it. Killing it closes the socket to + // Anthropic, which stops generation on the streaming calls; a non-streaming + // call is already submitted and gets billed whatever we do here. + let cancelled = false; + let hardKill: ReturnType | undefined; + + const onAbort = () => { + cancelled = true; + console.warn(`cancelling ${scriptName}: client went away`); + proc.kill("SIGTERM"); + + // Python installs no SIGTERM handler, so termination is immediate. This + // is only for a child wedged somewhere that never sees it. + hardKill = setTimeout(() => proc.kill("SIGKILL"), 5_000); + hardKill.unref?.(); + }; + + if (signal) { + if (signal.aborted) { + onAbort(); + } else { + signal.addEventListener("abort", onAbort, { once: true }); + } + } + 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)) { + if (LOG_LINE.test(line)) { // Divert the log line locally console.log(line); // TODO I'd love to break the log line up in to JSON actually @@ -92,21 +150,30 @@ export const run = async ( }); rl2.on("line", (line) => { console.error(line); - // /Divert all errors to the websocket - onLog?.(line); + + // Only forward what a service logged deliberately, the same rule stdout + // follows. Everything else on stderr is the interpreter talking: raw + // tracebacks carrying server paths, source lines, and whatever a frame + // held - which for a service is the payload. + if (LOG_LINE.test(line)) { + onLog?.(line); + } }); - proc.on("close", async (code) => { + proc.on("close", async (code, closeSignal) => { // Clean up readline interfaces immediately to prevent race conditions rl.close(); rl2.close(); - if (code) { - console.error("Python process exited with code", code); - reject(code); + if (hardKill) { + clearTimeout(hardKill); } - const result = Bun.file(outputPath); - const text = await result.text(); + signal?.removeEventListener("abort", onAbort); + + // Read before cleaning up, and clean up on every exit path + const text = await Bun.file(outputPath) + .text() + .catch(() => ""); try { await rm(inputPath); @@ -116,12 +183,41 @@ export const run = async ( console.error(e); } + // We killed it on purpose, so this is not a service failure + if (cancelled) { + return reject(subprocessCancelled(scriptName, closeSignal ?? "SIGTERM")); + } + + if (code) { + console.error("Python process exited with code", code); + return reject(subprocessFailed(scriptName, code)); + } + + // A child killed by a signal reports a null code, so without this the + // OOM killer - the likeliest way a service dies without exiting - would + // be reported as an empty result and the signal thrown away. + if (closeSignal) { + console.error(`Python process killed by ${closeSignal}`); + return reject(subprocessKilled(scriptName, closeSignal)); + } + if (text) { - resolve(JSON.parse(text)); - } else { - console.warn("No data returned from pythonland"); - resolve(null); + // Parsed inside the try: this handler is async, so a throw here + // becomes an unhandled rejection and the run never settles at all. + // A half-written file is what a crash mid-dump leaves behind. + try { + return resolve(JSON.parse(text)); + } catch (e) { + console.error(`Unreadable output from ${scriptName}`); + console.error(e); + return reject(malformedResult(scriptName)); + } } + + // entry.py writes a result on every path it completes, including its own + // error envelopes, so an empty file means the run died + console.warn("No data returned from pythonland"); + return reject(emptyResult(scriptName)); }); return; diff --git a/platform/src/middleware/healthcheck.tsx b/platform/src/middleware/healthcheck.tsx index a6341496..2391e888 100644 --- a/platform/src/middleware/healthcheck.tsx +++ b/platform/src/middleware/healthcheck.tsx @@ -1,6 +1,7 @@ import { Elysia } from "elysia"; import pkg from "../../../package.json" assert { type: "json" }; import { run } from '../bridge'; +import { toErrorPayload } from '../util/errors'; export default async (app: Elysia) => { app.get("/livez", () => { @@ -12,8 +13,24 @@ export default async (app: Elysia) => { }); }); app.get("/status", async () => { - const status = await run ('status', 0, {} as any) as any; - return new Response(status, { + // run() rejects now where it used to resolve null, and this is the one + // caller outside the services routes. Without the catch a spawn failure + // leaves the route throwing, so the health endpoint answers with Elysia's + // generic 500 and reports to Sentry rather than saying what is wrong. + let status: unknown; + try { + status = await run("status", 0, {} as any); + } catch (error) { + const payload = toErrorPayload(error); + return new Response(JSON.stringify(payload), { + status: payload.code, + headers: { + "Content-Type": "application/json", + }, + }); + } + + return new Response(status as any, { status: 200, headers: { "Content-Type": "application/json", diff --git a/platform/src/middleware/services.ts b/platform/src/middleware/services.ts index 1d5d6073..b45cfaf1 100644 --- a/platform/src/middleware/services.ts +++ b/platform/src/middleware/services.ts @@ -6,59 +6,119 @@ import { run } from "../bridge"; import describeModules, { type ModuleDescription, } from "../util/describe-modules"; -import { isApolloError } from "../util/errors"; -import type { InstanceAuth } from "../auth/instance-auth"; +import { isApolloError, toErrorPayload } from "../util/errors"; +import type { InstanceAuth, KeyResolution } from "../auth/instance-auth"; const textEncoder = new TextEncoder(); +// The space after the colon is required: Lightning's SSE decoder matches +// `": " <> comment` and has no catch-all clause, so ":ping" raises there. +export const HEARTBEAT_FRAME = ": ping\n\n"; + +// Inside the shortest silence any hop tolerates - our own socket, Lightning's +// inter-chunk timeout, and the 60s read timeout typical of proxies. +export const HEARTBEAT_INTERVAL_MS = 15_000; + +// Anything past this is a mistake rather than a choice, and setInterval turns +// a delay over 2^31-1 into "every tick" - so the value someone reaches for to +// mean "effectively never" would flood every open stream instead. +const HEARTBEAT_MAX_MS = 120_000; + +// Read per request rather than at import, so it can be turned down without a +// release if a hop turns out to be less patient than we thought. Out-of-range +// values fall back rather than being passed to setInterval. +export const heartbeatIntervalMs = (): number => { + const configured = Number(process.env.APOLLO_HEARTBEAT_INTERVAL_MS); + + return Number.isFinite(configured) && + configured > 0 && + configured <= HEARTBEAT_MAX_MS + ? configured + : HEARTBEAT_INTERVAL_MS; +}; + +// Killing these mid-run corrupts shared state that other services read: +// embed_docsite fills a fresh timestamped namespace that search_docsite treats +// as live the moment it is newest, and load_adaptor_docs commits its delete +// before its insert. They run to completion even when the caller has gone. +const NON_CANCELLABLE = new Set(["embed_docsite", "load_adaptor_docs"]); + const callService = ( m: ModuleDescription, port: number, payload?: any, onLog?: (str: string) => void, - onEvent?: (evt: string, payload: any) => void + onEvent?: (evt: string, payload: any) => void, + signal?: AbortSignal ) => { + if (NON_CANCELLABLE.has(m.name)) { + signal = undefined; + } + if (m.type === "py") { - return run(m.name, port, payload as any, onLog, onEvent); + return run(m.name, port, payload as any, onLog, onEvent, signal); } else { // TODO add event handling to ts services + // TODO ts services can't be cancelled - the handler signature has no signal return m.handler!(port, payload as any, onLog); } }; +/** Write the resolved key onto an outgoing payload. + * + * Exported so a test can assert what lands on the payload without a service + * reflecting it back: that was how this was covered before, and masking the + * reflection took the proof with it. + * + * The switch is explicit so the inbound-credential-never-forwarded invariant + * is structural rather than positional: a known client's stored key is + * swapped in (useKey), a request with no key of its own drops the field so + * python falls back to the global one (useGlobal), and an internal + * apollo() hop is forwarded exactly as received (passthrough). + */ +export const applyResolvedKey = ( + payload: Record, + resolution: KeyResolution +): Record => { + switch (resolution.kind) { + case "useKey": + payload.api_key = resolution.key; + break; + case "useGlobal": + delete payload.api_key; + break; + case "passthrough": + break; + default: { + // A new KeyResolution tag must be a compile error here, not a silent + // forward of the inbound credential. + const _exhaustive: never = resolution; + throw new Error( + `unhandled KeyResolution: ${(resolution as { kind: string }).kind}` + ); + } + } + + return payload; +}; + +// The in-flight run for each open websocket, so closing the socket can stop +// it. Keyed on ws.data, NOT on ws: Elysia builds a fresh ElysiaWS wrapper for +// every callback, so the object message() sees is never the object close() +// sees and a map keyed on it can never hit. ws.data is the upgrade context, +// created once per connection and carried on every wrapper. Deleted as soon +// as the run settles, so a closed socket holds nothing. +const wsRuns = new WeakMap(); + export default async (app: Elysia, port: number, auth: InstanceAuth) => { console.log("Loading routes:"); const modules = await describeModules(path.resolve("./services")); - // Apply the resolved key to an outgoing payload with an explicit switch so the - // inbound-credential-never-forwarded invariant is structural, not positional: a - // known client's stored key is swapped in (useKey), a NULL stored key (or a request - // with no api_key) drops the field so Python uses the global key (useGlobal), and an - // internal apollo() hop forwards the body exactly as received (passthrough). `ctx` is - // the upgrade-time context that carries lightningClient/internalCall: on POST the - // route ctx, on WS the captured ws.data, never a fresh per-message one. - const applyKey = (payload: Record, ctx: any) => { - const resolution = auth.resolveKey(ctx); - switch (resolution.kind) { - case "useKey": - payload.api_key = resolution.key; - break; - case "useGlobal": - delete payload.api_key; - break; - case "passthrough": - break; - default: { - // Exhaustiveness guard: a new KeyResolution tag must be a compile error - // here, not a silent forward of the inbound credential. - const _exhaustive: never = resolution; - throw new Error( - `unhandled KeyResolution: ${(resolution as { kind: string }).kind}` - ); - } - } - return payload; - }; + // `ctx` is the upgrade-time context carrying lightningClient/internalCall: + // on POST the route ctx, on WS the captured ws.data, never a fresh + // per-message one. + const applyKey = (payload: Record, ctx: any) => + applyResolvedKey(payload, auth.resolveKey(ctx)); const buildPayload = (ctx: any) => applyKey({ ...(ctx.body ?? {}), session_id: ctx.uuid }, ctx); @@ -79,7 +139,29 @@ export default async (app: Elysia, port: number, auth: InstanceAuth) => { app.post(name, async (ctx) => { console.log(`POST /services/${name}: ${ctx.uuid}`); const payload = buildPayload(ctx); - const result = await callService(m, port, payload as any); + + let result: any; + try { + // The signal matters here as much as on the stream: a client that + // gives up waiting for a plain POST leaves the model generating + // just the same. + result = await callService( + m, + port, + payload as any, + undefined, + undefined, + ctx.request?.signal + ); + } catch (error) { + const payload = toErrorPayload(error); + return new Response(JSON.stringify(payload), { + status: payload.code, + headers: { + "Content-Type": "application/json", + }, + }); + } if (isApolloError(result)) { return new Response(JSON.stringify(result), { @@ -98,10 +180,21 @@ export default async (app: Elysia, port: number, auth: InstanceAuth) => { console.log(`STREAM START /services/${name}: ${ctx.uuid}`); const payload = buildPayload(ctx); + const abort = new AbortController(); + + // Hoisted so cancel() can reach what start() set up + let isClosed = false; + let heartbeat: ReturnType | undefined; + + const stopHeartbeat = () => { + if (heartbeat) { + clearInterval(heartbeat); + heartbeat = undefined; + } + }; + const stream = new ReadableStream({ async start(controller) { - let isClosed = false; - const sendSSE = (event: string, data: any) => { if (isClosed) { return; @@ -113,8 +206,11 @@ export default async (app: Elysia, port: number, auth: InstanceAuth) => { // console.log(message.trim()); controller.enqueue(textEncoder.encode(message)); } catch (error) { - // Stream may have been closed + // Same as the heartbeat's catch: a throwing enqueue is a + // dropped connection the runtime has not told us about. isClosed = true; + stopHeartbeat(); + abort.abort(); } }; @@ -126,13 +222,33 @@ export default async (app: Elysia, port: number, auth: InstanceAuth) => { sendSSE(type, payload); }; + // Started before the service call so the window while Python boots + // is covered too + heartbeat = setInterval(() => { + if (isClosed) { + stopHeartbeat(); + return; + } + try { + controller.enqueue(textEncoder.encode(HEARTBEAT_FRAME)); + } catch (error) { + // The consumer went away between ticks. cancel() may never + // fire for this, so end the run here rather than leaving the + // child generating for nobody. + isClosed = true; + stopHeartbeat(); + abort.abort(); + } + }, heartbeatIntervalMs()); + try { const result = await callService( m, port, payload as any, onLog, - onEvent + onEvent, + abort.signal ); if (isApolloError(result)) { @@ -141,20 +257,40 @@ export default async (app: Elysia, port: number, auth: InstanceAuth) => { sendSSE("complete", result); } } catch (error) { - sendSSE("error", { - message: - error instanceof Error ? error.message : "Unknown error", - }); + sendSSE("error", toErrorPayload(error)); } finally { + stopHeartbeat(); console.log( `STREAM COMPLETE ${ctx.uuid} in ${ - (new Date() - ctx.start) / 1000 + (Date.now() - ctx.start) / 1000 }s` ); isClosed = true; - controller.close(); + try { + controller.close(); + } catch (error) { + // already closed from the consumer's side + } } }, + + // Everything from here on is work nobody will read + cancel(reason) { + isClosed = true; + stopHeartbeat(); + console.warn( + `STREAM CANCELLED ${ctx.uuid} after ${ + (Date.now() - ctx.start) / 1000 + }s` + ); + abort.abort(reason); + }, + }); + + // cancel() depends on the runtime noticing the dropped connection, so + // listen on the request's own signal too. abort() is idempotent. + ctx.request?.signal?.addEventListener("abort", () => abort.abort(), { + once: true, }); return new Response(stream, { @@ -179,6 +315,14 @@ export default async (app: Elysia, port: number, auth: InstanceAuth) => { open() { console.log(`Websocket connected at /services/${name}`); }, + // A websocket has no request signal, so cancellation needs its own + // controller and a close handler to fire it. Without this a WS caller + // going away leaves the child generating, which is the cost the whole + // change exists to stop. + close(ws) { + wsRuns.get(ws.data)?.abort(); + wsRuns.delete(ws.data); + }, message(ws, message) { try { if (message.event === "start") { @@ -202,12 +346,49 @@ export default async (app: Elysia, port: number, auth: InstanceAuth) => { const base: Record = { ...(message.data ?? {}) }; const payload = applyKey(base, ws.data); - callService(m, port, payload as any, onLog, onEvent).then( + // A second start frame on the same socket would otherwise + // orphan the first run: its controller leaves the map, so the + // close handler could never reach it again. + wsRuns.get(ws.data)?.abort(); + + const abort = new AbortController(); + wsRuns.set(ws.data, abort); + + // Two arguments rather than a chained catch: a chain would also + // catch a throw from the success callback and report it to the + // client as a service failure. + // + // The failure handler matters as much as the success one - a run + // that rejects (spawn failure, empty output) would otherwise + // leave the client waiting on a frame that never comes, since + // the try around this only sees synchronous throws. + callService( + m, + port, + payload as any, + onLog, + onEvent, + abort.signal + ).then( (result) => { + // Only if it is still ours: a later start frame may have + // replaced this run's controller with its own. + if (wsRuns.get(ws.data) === abort) { + wsRuns.delete(ws.data); + } ws.send({ event: "complete", data: result, }); + }, + (error) => { + if (wsRuns.get(ws.data) === abort) { + wsRuns.delete(ws.data); + } + ws.send({ + event: "error", + data: toErrorPayload(error), + }); } ); } diff --git a/platform/src/server.ts b/platform/src/server.ts index cbd8f3a1..0fbe6f4a 100644 --- a/platform/src/server.ts +++ b/platform/src/server.ts @@ -11,6 +11,34 @@ import { captureException } from "./util/sentry"; import { clientsDbUrl, closeDb } from "./db"; import { runMigrations } from "./db/migrate"; import { randomUUID } from "node:crypto"; +import { readdir, rm } from "node:fs/promises"; +import path from "node:path"; + +// A run's input file can hold values that belong to the deployment rather +// than the caller, and only the bridge's close handler removes it - so +// anything that stopped the process mid-run left one behind, and nothing else +// ever sweeps them. Startup is the one moment we know no run of ours is +// reading them. +const sweepTempPayloads = async () => { + const dir = path.resolve("tmp/data"); + + try { + const stale = await readdir(dir); + + await Promise.all( + stale.map((name) => rm(path.join(dir, name), { force: true })) + ); + + if (stale.length) { + console.log(`Removed ${stale.length} temp payload(s) left by a previous run`); + } + } catch (error) { + // No directory yet on a first boot, which is not worth reporting. + if ((error as { code?: string }).code !== "ENOENT") { + console.error("Could not sweep tmp/data", error); + } + } +}; import pkg from "../../package.json"; export default async ( @@ -19,7 +47,21 @@ export default async ( // pass a pre-configured instance (fake lookup) instead of the live DB-backed one. auth: InstanceAuth = new InstanceAuth() ) => { - const app = new Elysia(); + // Bun's idle timer applies to in-flight SSE responses, not just idle + // keep-alive sockets, and Elysia defaults it to 30s - shorter than our own + // services routinely go without emitting. 255 is Bun's maximum. + const app = new Elysia({ + serve: { + idleTimeout: 255, + }, + // Websockets have their own idle timer, which serve.idleTimeout does not + // reach - it defaults to 120s, so a WS caller waiting on a slow answer + // would be dropped well before the SSE route's heartbeat had earned it + // anything. + websocket: { + idleTimeout: 255, + }, + }); app.use(html()); @@ -51,9 +93,11 @@ export default async ( } } - // app.listen below sets no reusePort, so the multi-process internal-token warn - // is dormant; pass the flag here if clustering is ever enabled. - logInternalTokenProvenance(false); + // Elysia's Bun adapter sets reusePort unconditionally, so the guard that + // warns about a per-process token meeting a shared port is live, not + // hypothetical. It stays quiet once APOLLO_INTERNAL_TOKEN is set. + logInternalTokenProvenance(true); + await sweepTempPayloads(); await auth.init(); // No stop path exists otherwise; close the DB pool so a graceful pod termination diff --git a/platform/src/util/errors.ts b/platform/src/util/errors.ts index f967acb4..3bf1fb0a 100644 --- a/platform/src/util/errors.ts +++ b/platform/src/util/errors.ts @@ -22,6 +22,116 @@ export function apolloError( return { code, type, message, ...(details ? { details } : {}) }; } +/** An ApolloError that can also be thrown, so a catch block gets a real Error + * while the wire shape stays the same. `toJSON` is required: JSON.stringify on + * an Error is `{}` without it. */ +export class ApolloThrowable extends Error implements ApolloError { + readonly code: number; + readonly type: string; + readonly details?: Record; + + constructor( + code: number, + type: string, + message: string, + details?: Record + ) { + super(message); + this.name = "ApolloThrowable"; + this.code = code; + this.type = type; + this.details = details; + } + + toJSON(): ApolloError { + return { + code: this.code, + type: this.type, + message: this.message, + ...(this.details ? { details: this.details } : {}), + }; + } +} + +/** The service process exited non-zero. */ +export function subprocessFailed( + service: string, + exitCode: number +): ApolloThrowable { + return new ApolloThrowable( + 500, + "SUBPROCESS_FAILED", + `Service "${service}" exited with code ${exitCode}`, + { service, exitCode } + ); +} + +/** We never got as far as running the service - poetry or python missing, or + * the spawn refused. */ +export function subprocessSpawnFailed( + service: string, + cause: unknown +): ApolloThrowable { + return new ApolloThrowable( + 500, + "SUBPROCESS_SPAWN_FAILED", + `Service "${service}" could not be started`, + { service, cause: cause instanceof Error ? cause.message : String(cause) } + ); +} + +/** We stopped the service ourselves because the client went away. 499 keeps + * these out of the 5xx that mean something is actually broken. */ +export function subprocessCancelled( + service: string, + signal: string +): ApolloThrowable { + return new ApolloThrowable( + 499, + "SUBPROCESS_CANCELLED", + `Service "${service}" was cancelled because the client disconnected`, + { service, signal } + ); +} + +/** Something outside Apollo killed the service - most often the OOM killer. + * Distinct from a cancellation, which is us, and from a non-zero exit, which + * is the service deciding to stop. */ +export function subprocessKilled( + service: string, + signal: string +): ApolloThrowable { + return new ApolloThrowable( + 500, + "SUBPROCESS_KILLED", + `Service "${service}" was killed by ${signal}`, + { service, signal } + ); +} + +/** The service exited cleanly but its output isn't valid JSON. Without this + * case the parse error escapes the close handler and the request never + * settles. */ +export function malformedResult(service: string): ApolloThrowable { + return new ApolloThrowable( + 502, + "MALFORMED_RESULT", + `Service "${service}" produced a result that could not be parsed`, + { service } + ); +} + +/** The service exited cleanly but wrote nothing. entry.py writes a result on + * every path it completes, so an empty file means the run died. */ +export function emptyResult(service: string): ApolloThrowable { + return new ApolloThrowable( + 502, + "EMPTY_RESULT", + `Service "${service}" finished without producing a result`, + { service } + ); +} + export function unauthorized(ctx: any): ApolloError { return apolloError(ctx, 401, "UNAUTHORIZED", "Missing or invalid API key"); } @@ -43,3 +153,27 @@ export function clientMisconfigured(ctx: any): ApolloError { "Client has no API key configured" ); } + +/** Normalise anything thrown by a service run into the ApolloError envelope, so + * a caller sees the same shape whether the failure was typed or not. */ +export function toErrorPayload(error: unknown): ApolloError { + if (error instanceof ApolloThrowable) { + return error.toJSON(); + } + // Rebuilt field by field rather than returned as-is: isApolloError only + // checks for a numeric `code`, and anything else hanging off the object + // would be serialised to the caller along with it. + if (isApolloError(error)) { + return { + code: error.code, + type: error.type, + message: error.message, + ...(error.details === undefined ? {} : { details: error.details }), + }; + } + return { + code: 500, + type: "INTERNAL_ERROR", + message: error instanceof Error ? error.message : String(error), + }; +} diff --git a/platform/test/apply-key.test.ts b/platform/test/apply-key.test.ts new file mode 100644 index 00000000..ec446077 --- /dev/null +++ b/platform/test/apply-key.test.ts @@ -0,0 +1,70 @@ +import { describe, expect, it } from "bun:test"; + +import { applyResolvedKey } from "../src/middleware/services"; + +// What lands on the payload a service is about to run with. +// +// This used to be covered end to end: echo reflected its input, so a test +// could POST a credential and read the substituted value back off the +// response. Masking that reflection was right, and it took the proof with it — +// with everything coming back "[REDACTED]", a swap that wrote the wrong value, +// or forwarded the caller's own, would look identical to a correct one. +// +// So the assertion moved to the substitution itself. Paired with the +// InstanceAuth tests that pin which resolution a given credential produces, +// the two halves cover the same ground the round trip did, without a service +// having to hand a value back to prove it. +describe("applyResolvedKey", () => { + const CALLER = "sk-ant-the-caller-sent-this"; + const STORED = "sk-ant-what-the-server-substitutes"; + + it("substitutes the stored value for a known client", () => { + const payload = applyResolvedKey( + { x: 1, api_key: CALLER }, + { kind: "useKey", key: STORED } + ); + + expect(payload.api_key).toBe(STORED); + expect(payload.x).toBe(1); + }); + + // The invariant the whole auth layer exists for. + it("never forwards what the caller sent", () => { + const payload = applyResolvedKey( + { api_key: CALLER }, + { kind: "useKey", key: STORED } + ); + + expect(JSON.stringify(payload)).not.toContain(CALLER); + }); + + // Dropped rather than blanked: python falls back to the global key only when + // the field is absent. + it("drops the field entirely when the global key should serve", () => { + const payload = applyResolvedKey( + { x: 1, api_key: CALLER }, + { kind: "useGlobal" } + ); + + expect("api_key" in payload).toBe(false); + expect(payload.x).toBe(1); + }); + + it("leaves an internal hop's body exactly as received", () => { + const forwarded = "sk-ant-forwarded-by-an-internal-call"; + const payload = applyResolvedKey( + { api_key: forwarded }, + { kind: "passthrough" } + ); + + expect(payload.api_key).toBe(forwarded); + }); + + // The default branch is an exhaustiveness guard: a new resolution kind must + // fail loudly rather than fall through and forward the inbound credential. + it("refuses a resolution it does not recognise", () => { + expect(() => + applyResolvedKey({ api_key: CALLER }, { kind: "brand-new" } as never) + ).toThrow(/unhandled KeyResolution/); + }); +}); diff --git a/platform/test/bridge-cancel.test.ts b/platform/test/bridge-cancel.test.ts new file mode 100644 index 00000000..1609dafb --- /dev/null +++ b/platform/test/bridge-cancel.test.ts @@ -0,0 +1,92 @@ +import { describe, expect, it } from "bun:test"; + +import { run } from "../src/bridge"; +import { ApolloThrowable } from "../src/util/errors"; + +const PORT = 9871; + +// The spawned command is `poetry run python ...`, so matching the process list +// alone hits the poetry wrapper seconds before the interpreter exists. The probe +// announces itself once python is really inside the service. +const started = () => { + let resolve!: () => void; + const promise = new Promise((r) => { + resolve = r; + }); + + return { + promise, + onEvent: (type: string) => { + if (type === "probe_started") { + resolve(); + } + }, + }; +}; + +// No `pgrep -c`: BSD pgrep has no count flag, and its usage error goes to +// stderr, so asking for one reads as "no processes". +const probeProcessCount = async () => { + const proc = Bun.spawn(["pgrep", "-f", "_cancel_probe"], { + stdout: "pipe", + stderr: "pipe", + }); + const out = await new Response(proc.stdout).text(); + return out.trim().split("\n").filter(Boolean).length; +}; + +const waitForProbesToClear = async (limitMs = 15_000) => { + const deadline = Date.now() + limitMs; + while (Date.now() < deadline) { + if ((await probeProcessCount()) === 0) { + return true; + } + await Bun.sleep(200); + } + return false; +}; + +describe("cancelling a service run", () => { + it("kills the running python child when the caller aborts", async () => { + const abort = new AbortController(); + const probe = started(); + + const pending = run( + "_cancel_probe", + PORT, + { sleep_for: 120 }, + undefined, + probe.onEvent, + abort.signal + ); + + await probe.promise; + expect(await probeProcessCount()).toBeGreaterThan(0); + + abort.abort(); + + const error = (await pending.catch((e) => e)) as ApolloThrowable; + expect(error).toBeInstanceOf(ApolloThrowable); + expect(error.type).toBe("SUBPROCESS_CANCELLED"); + expect(error.code).toBe(499); + + expect(await waitForProbesToClear()).toBe(true); + }, 90_000); + + it("does not leave a child running when the signal is already aborted", async () => { + const abort = new AbortController(); + abort.abort(); + + const error = (await run( + "_cancel_probe", + PORT, + { sleep_for: 120 }, + undefined, + undefined, + abort.signal + ).catch((e) => e)) as ApolloThrowable; + + expect(error.type).toBe("SUBPROCESS_CANCELLED"); + expect(await waitForProbesToClear()).toBe(true); + }, 90_000); +}); diff --git a/platform/test/middleware/heartbeat-live.test.ts b/platform/test/middleware/heartbeat-live.test.ts new file mode 100644 index 00000000..8f33d6ef --- /dev/null +++ b/platform/test/middleware/heartbeat-live.test.ts @@ -0,0 +1,74 @@ +import { afterEach, describe, expect, it } from "bun:test"; + +import setup from "../../src/server"; +import { InstanceAuth } from "../../src/auth/instance-auth"; +import { HEARTBEAT_FRAME } from "../../src/middleware/services"; + +// The frame-shape tests next door pass whether or not anything ever emits one. +// These read the wire. +const port = 9871; + +const auth = new InstanceAuth({ lookup: () => null, hasGlobalKey: true }); +const app = await setup(port, auth); + +const previous = process.env.APOLLO_HEARTBEAT_INTERVAL_MS; + +afterEach(() => { + if (previous === undefined) { + delete process.env.APOLLO_HEARTBEAT_INTERVAL_MS; + } else { + process.env.APOLLO_HEARTBEAT_INTERVAL_MS = previous; + } +}); + +const readStream = async (intervalMs: string) => { + process.env.APOLLO_HEARTBEAT_INTERVAL_MS = intervalMs; + + const response = await app.handle( + new Request(`http://localhost:${port}/services/echo/stream`, { + method: "POST", + body: JSON.stringify({ message: "hello" }), + headers: { "Content-Type": "application/json" }, + }) + ); + + expect(response.status).toBe(200); + + return await response.text(); +}; + +describe("SSE heartbeat on a live stream", () => { + // Spawning Python takes long enough that a heartbeat this fast must tick + // before the service produces anything. The interval is deliberately set + // rather than waited out: the real one is 15s. + it("emits heartbeats while the service is still starting up", async () => { + const body = await readStream("20"); + + const beats = body.split(HEARTBEAT_FRAME).length - 1; + + expect(beats).toBeGreaterThan(0); + expect(body).toContain("event: complete"); + }, 30_000); + + // Guards the ordering the whole design rests on: the first thing on the wire + // is a heartbeat, not the result, so no hop sees an idle socket. + it("gets a heartbeat onto the wire before the result", async () => { + const body = await readStream("20"); + + const firstBeat = body.indexOf(HEARTBEAT_FRAME); + const result = body.indexOf("event: complete"); + + // Both have to be present before comparing them: indexOf gives -1 for a + // heartbeat that never arrived, which would sail past a bare <. + expect(firstBeat).toBeGreaterThanOrEqual(0); + expect(result).toBeGreaterThanOrEqual(0); + expect(firstBeat).toBeLessThan(result); + }, 30_000); + + it("sends none when the interval outlasts the request", async () => { + const body = await readStream("600000"); + + expect(body).not.toContain(HEARTBEAT_FRAME); + expect(body).toContain("event: complete"); + }, 30_000); +}); diff --git a/platform/test/middleware/heartbeat.test.ts b/platform/test/middleware/heartbeat.test.ts new file mode 100644 index 00000000..8822c170 --- /dev/null +++ b/platform/test/middleware/heartbeat.test.ts @@ -0,0 +1,66 @@ +import { afterEach, describe, expect, it } from "bun:test"; + +import { + HEARTBEAT_FRAME, + HEARTBEAT_INTERVAL_MS, + heartbeatIntervalMs, +} from "../../src/middleware/services"; + +describe("SSE heartbeat", () => { + // Lightning's SSE decoder matches `": " <> comment` and has no catch-all + // clause, so ":ping" raises FunctionClauseError inside its stream fold. + it("is a comment frame with a space after the colon", () => { + expect(HEARTBEAT_FRAME.startsWith(": ")).toBe(true); + }); + + // Without the blank line the comment sits in the decoder's buffer and resets + // nobody's idle timer. + it("terminates the frame so it is dispatched rather than buffered", () => { + expect(HEARTBEAT_FRAME.endsWith("\n\n")).toBe(true); + }); + + // An event frame would reach Lightning's handle_sse_event and need a clause. + it("carries no event or data field", () => { + expect(HEARTBEAT_FRAME).not.toContain("event:"); + expect(HEARTBEAT_FRAME).not.toContain("data:"); + }); + + it("ticks well inside the shortest silence any hop tolerates", () => { + expect(HEARTBEAT_INTERVAL_MS).toBeLessThanOrEqual(20_000); + expect(HEARTBEAT_INTERVAL_MS).toBeGreaterThanOrEqual(5_000); + }); +}); + +describe("heartbeat interval override", () => { + const previous = process.env.APOLLO_HEARTBEAT_INTERVAL_MS; + + afterEach(() => { + if (previous === undefined) { + delete process.env.APOLLO_HEARTBEAT_INTERVAL_MS; + } else { + process.env.APOLLO_HEARTBEAT_INTERVAL_MS = previous; + } + }); + + const resolves = (value: string) => { + process.env.APOLLO_HEARTBEAT_INTERVAL_MS = value; + return heartbeatIntervalMs(); + }; + + it("takes a sensible override", () => { + expect(resolves("5000")).toBe(5000); + }); + + // setInterval treats a delay past 2^31-1 as "every tick", so the value + // someone picks to mean "effectively never" is the one that would flood + // every open stream. + it("falls back rather than overflowing setInterval", () => { + expect(resolves("3000000000")).toBe(HEARTBEAT_INTERVAL_MS); + }); + + it("falls back on zero, negatives and nonsense", () => { + expect(resolves("0")).toBe(HEARTBEAT_INTERVAL_MS); + expect(resolves("-1")).toBe(HEARTBEAT_INTERVAL_MS); + expect(resolves("soon")).toBe(HEARTBEAT_INTERVAL_MS); + }); +}); diff --git a/platform/test/middleware/ws-cancel.test.ts b/platform/test/middleware/ws-cancel.test.ts new file mode 100644 index 00000000..52aa932a --- /dev/null +++ b/platform/test/middleware/ws-cancel.test.ts @@ -0,0 +1,72 @@ +import { afterAll, describe, expect, it } from "bun:test"; + +import setup from "../../src/server"; +import { InstanceAuth } from "../../src/auth/instance-auth"; + +// The bridge tests hand run() an AbortSignal directly, which proves the kill +// but not the wiring: nothing between the socket and the signal is exercised. +// That gap hid a real bug - Elysia rebuilds its ws wrapper for every callback, +// so a run keyed on the wrapper in message() was unreachable from close() and +// no websocket run was ever cancelled. This suite drives a real socket. +const port = 9874; + +const auth = new InstanceAuth({ lookup: () => null, hasGlobalKey: true }); +await setup(port, auth); + +// No `pgrep -c`: BSD pgrep has no count flag, and its usage error goes to +// stderr, so asking for one reads as "no processes". +const childCount = async () => { + const proc = Bun.spawn(["pgrep", "-f", "entry.py test_slow"], { + stdout: "pipe", + stderr: "pipe", + }); + const out = await new Response(proc.stdout).text(); + return out.trim().split("\n").filter(Boolean).length; +}; + +const waitForChildrenToClear = async (limitMs = 20_000) => { + const deadline = Date.now() + limitMs; + while (Date.now() < deadline) { + if ((await childCount()) === 0) { + return true; + } + await Bun.sleep(250); + } + return false; +}; + +// A failed run must not leave a two-minute child holding up the suite. +afterAll(() => { + Bun.spawnSync(["pkill", "-f", "entry.py test_slow"]); +}); + +describe("cancelling over a real websocket", () => { + it("kills the python child when the socket closes", async () => { + const socket = new WebSocket(`ws://localhost:${port}/services/test_slow`); + + // The service announces itself on stdout once the interpreter is truly + // inside main(); matching the process list any earlier hits the poetry + // wrapper and proves nothing. + const started = new Promise((resolve) => { + socket.addEventListener("message", ({ data }) => { + const evt = JSON.parse(String(data)); + if (evt.event === "event" && evt.type === "probe_started") { + resolve(); + } + }); + }); + + socket.addEventListener("open", () => { + socket.send( + JSON.stringify({ event: "start", data: { sleep_for: 120 } }) + ); + }); + + await started; + expect(await childCount()).toBeGreaterThan(0); + + socket.close(); + + expect(await waitForChildrenToClear()).toBe(true); + }, 90_000); +}); diff --git a/platform/test/server.test.ts b/platform/test/server.test.ts index 72ac4e96..db449950 100644 --- a/platform/test/server.test.ts +++ b/platform/test/server.test.ts @@ -85,6 +85,13 @@ describe("Main server", () => { expect(await response.text()).toBe(""); }); + // Asserts the configuration rather than the behaviour: these tests drive the + // app through app.handle(), which never opens a socket, so none of them can + // observe a socket timer. + it("keeps the socket idle timeout long enough for slow SSE streams", () => { + expect(app.config.serve?.idleTimeout).toBe(255); + }); + // send messages through a web socket }); @@ -247,13 +254,20 @@ describe("Instance authentication", () => { }); // Row 1 + // + // echo used to hand its payload back verbatim, which is what let these rows + // read the swapped key off the response. It masks now, so what is left to + // observe here is that the request was accepted and the field was set. That + // the *correct* key is selected is covered directly against + // InstanceAuth.authenticate further down, where lightningClient.anthropicKey + // is asserted without a round trip through a service. it("accepts a known credential and swaps in the client's stored key", async () => { const res = await app.handle(postKey("services/echo", { x: 1 }, ALPHA)); expect(res.status).toBe(200); const body = await res.json(); expect(body.x).toBe(1); - expect(body.api_key).toBe("sk-ant-stored-alpha"); - expect(body.api_key).not.toBe(ALPHA); + expect(body.api_key).toBe("[REDACTED]"); + expect(JSON.stringify(body)).not.toContain(ALPHA); }); // Row 2 — a recognised client whose stored key is NULL is a server-side @@ -334,7 +348,7 @@ describe("Instance authentication", () => { postKey("services/echo", { x: 1 }, "sk-ant-unknown") ); expect(swap.status).toBe(200); - expect((await swap.json()).api_key).toBe("sk-ant-stored-alpha"); + expect((await swap.json()).api_key).toBe("[REDACTED]"); expect(unknown.status).toBe(401); }); @@ -476,7 +490,9 @@ describe("Instance authentication", () => { const res = await app.handle(req); expect(res.status).toBe(200); const body = await res.json(); - expect(body.api_key).toBe("sk-ant-internal-hop"); + // Masked on the way out like anything else. What this row pins is that an + // internal call is accepted and its body forwarded rather than rejected. + expect(body.api_key).toBe("[REDACTED]"); }); // WS upgrade auth decision via app.handle(): a bare upgrade carries no api_key, so @@ -537,8 +553,8 @@ describe("Instance authentication", () => { // that ws.data carries the lightningClient set during beforeHandle. it("swaps a known client's stored key on a WS upgrade via ?api_key=", async () => { const body = await wsRoundTrip(`?api_key=${encodeURIComponent(ALPHA)}`); - expect(body.api_key).toBe("sk-ant-stored-alpha"); - expect(body.api_key).not.toBe(ALPHA); + expect(body.api_key).toBe("[REDACTED]"); + expect(JSON.stringify(body)).not.toContain(ALPHA); expect(body.ws).toBe(1); }); @@ -582,7 +598,7 @@ describe("Instance authentication", () => { ); }); }); - expect(body.api_key).toBe("sk-ant-internal-fwd"); + expect(body.api_key).toBe("[REDACTED]"); }); }); @@ -1049,3 +1065,12 @@ describe("Instance auth key encryption", () => { } }); }); + +describe("Websocket timeouts", () => { + // serve.idleTimeout does not reach websockets: Bun keeps a separate timer + // for them, defaulting to 120s. Without this a WS caller waiting on a slow + // answer is dropped long before the heartbeat has bought anything. + it("gives websockets the same patience as SSE streams", () => { + expect(app.config.websocket?.idleTimeout).toBe(255); + }); +}); diff --git a/platform/test/util/errors.test.ts b/platform/test/util/errors.test.ts new file mode 100644 index 00000000..e00fdf09 --- /dev/null +++ b/platform/test/util/errors.test.ts @@ -0,0 +1,90 @@ +import { describe, expect, it } from "bun:test"; + +import { + ApolloThrowable, + emptyResult, + isApolloError, + malformedResult, + subprocessCancelled, + subprocessFailed, + subprocessKilled, + subprocessSpawnFailed, +} from "../../src/util/errors"; + +describe("ApolloThrowable", () => { + // JSON.stringify on an Error is "{}", so without toJSON the synchronous + // service route answers a failure with the right status and an empty body. + it("survives JSON.stringify with its envelope intact", () => { + const parsed = JSON.parse(JSON.stringify(subprocessFailed("job_chat", 3))); + + expect(parsed.code).toBe(500); + expect(parsed.type).toBe("SUBPROCESS_FAILED"); + expect(parsed.message).toContain("job_chat"); + expect(parsed.details.exitCode).toBe(3); + }); + + it("is recognised by isApolloError, so existing envelope handling applies", () => { + expect(isApolloError(subprocessFailed("echo", 1))).toBe(true); + expect(isApolloError(emptyResult("echo"))).toBe(true); + }); + + it("is a real Error, so it can be thrown and caught normally", () => { + const error = emptyResult("workflow_chat"); + + expect(error instanceof Error).toBe(true); + expect(error instanceof ApolloThrowable).toBe(true); + expect(error.message.length).toBeGreaterThan(0); + }); +}); + +describe("subprocess failures", () => { + it("keeps the exit code", () => { + expect(subprocessFailed("job_chat", 137).details?.exitCode).toBe(137); + }); + + // A spawn failure means poetry or python is missing, which is an operator + // problem rather than a service one. + it("distinguishes never-started from started-and-failed", () => { + expect(subprocessSpawnFailed("echo", new Error("ENOENT")).type).toBe( + "SUBPROCESS_SPAWN_FAILED" + ); + expect(subprocessFailed("echo", 1).type).toBe("SUBPROCESS_FAILED"); + }); + + // It ran and exited cleanly; what came back was unusable. + it("reports an empty result as a bad gateway, not a server error", () => { + expect(emptyResult("echo").code).toBe(502); + }); + + it("carries the spawn cause as a string, not a nested Error", () => { + const details = subprocessSpawnFailed("echo", new Error("ENOENT")).details; + + expect(details?.cause).toBe("ENOENT"); + }); + + it("keeps a deliberate cancellation out of the 5xx range", () => { + // We stopped this one ourselves because the caller left. Reporting it as a + // server error would bury the failures that mean something is broken. + const cancelled = subprocessCancelled("global_chat", "SIGTERM"); + + expect(cancelled.code).toBe(499); + expect(cancelled.type).toBe("SUBPROCESS_CANCELLED"); + expect(cancelled.details?.signal).toBe("SIGTERM"); + }); + + // OOM or a deploy's SIGTERM: the process reports a null exit code, so the + // signal is the only honest diagnosis. + it("names the signal when the process was killed", () => { + const error = subprocessKilled("embed_docsite", "SIGKILL"); + + expect(error.type).toBe("SUBPROCESS_KILLED"); + expect(error.details?.signal).toBe("SIGKILL"); + }); + + it("reports unparseable output as a bad gateway", () => { + const error = malformedResult("echo"); + + expect(error.type).toBe("MALFORMED_RESULT"); + expect(error.code).toBe(502); + }); +}); diff --git a/services/_cancel_probe/_cancel_probe.py b/services/_cancel_probe/_cancel_probe.py new file mode 100644 index 00000000..f06bd21c --- /dev/null +++ b/services/_cancel_probe/_cancel_probe.py @@ -0,0 +1,20 @@ +"""A service that does nothing slowly, so cancellation can be tested. + +The leading underscore keeps it out of describe-modules, so it is never mounted +and has no route. + +It announces itself before sleeping because the spawned command is +`poetry run python ...`: the process list matches on the poetry wrapper seconds +before the interpreter has booted, so waiting on that would prove nothing. +""" + +import time + + +def main(data_dict: dict) -> dict: + print("EVENT:probe_started:{}", flush=True) # noqa: T201 + + seconds = data_dict.get("sleep_for", 30) + time.sleep(seconds) + + return {"slept": seconds} diff --git a/services/_masking_probe/__init__.py b/services/_masking_probe/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/services/_masking_probe/_masking_probe.py b/services/_masking_probe/_masking_probe.py new file mode 100644 index 00000000..1cb17fb1 --- /dev/null +++ b/services/_masking_probe/_masking_probe.py @@ -0,0 +1,24 @@ +"""Services that do the wrong thing on purpose, so the exit mask can be tested. + +The leading underscore keeps the directory out of describe-modules, so none of +this is mounted and none of it has a route. +""" + +from util import ApolloError + + +def main(data_dict: dict) -> dict: + """Reflects the payload without masking anything itself. + + Stands in for a service written by someone who did not think about it, + which is the case the exit mask exists for. + """ + return data_dict + + +def raise_with_payload(data_dict: dict) -> dict: + """The shape nearly every service uses: catch broadly, rewrap the text.""" + try: + raise ValueError(f"upstream rejected {data_dict.get('api_key')}") + except ValueError as e: + raise ApolloError(500, str(e), type="INTERNAL_ERROR") from e diff --git a/services/_masking_probe_raiser/__init__.py b/services/_masking_probe_raiser/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/services/_masking_probe_raiser/_masking_probe_raiser.py b/services/_masking_probe_raiser/_masking_probe_raiser.py new file mode 100644 index 00000000..8a54019d --- /dev/null +++ b/services/_masking_probe_raiser/_masking_probe_raiser.py @@ -0,0 +1,7 @@ +"""Raises with the payload in scope. Unmounted; see _masking_probe.""" + +from _masking_probe._masking_probe import raise_with_payload + + +def main(data_dict: dict) -> dict: + return raise_with_payload(data_dict) diff --git a/services/echo/README.md b/services/echo/README.md index 26efc3d2..95b88547 100644 --- a/services/echo/README.md +++ b/services/echo/README.md @@ -10,4 +10,7 @@ Call the endpoint at `services/echo` curl -X POST localhost:3000/services/echo --json @tmp/data.json ``` -Whatever you include in the body will be be returned straight back. +Whatever you include in the body will be returned straight back, with one +exception: values the server fills in itself, such as `api_key`, come back as +`[REDACTED]`, as does anything else shaped like a key. Those belong to the +deployment rather than to the caller, so echo does not hand them out. diff --git a/services/echo/__init__.py b/services/echo/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/services/echo/echo.py b/services/echo/echo.py index 90ecaf32..eaba29dc 100644 --- a/services/echo/echo.py +++ b/services/echo/echo.py @@ -1,11 +1,17 @@ +from langfuse_util import mask_secrets from util import ApolloError + from .log import log + # Sample python service to echo requests back to the caller def main(x): # raise a 400 if the payload is empty (ignoring the session id which is system-set) ## useful for diagnosing errors if not x or set(x.keys()) == {"session_id"}: raise ApolloError(code=400, message="payload is required", type="BAD_REQUEST") - log(x) - return x + # Not the raw payload: the server sets fields on it that belong to the + # deployment rather than to the caller. + safe = mask_secrets(x) + log(safe) + return safe diff --git a/services/echo/tests/__init__.py b/services/echo/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/services/echo/tests/unit/__init__.py b/services/echo/tests/unit/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/services/echo/tests/unit/test_echo.py b/services/echo/tests/unit/test_echo.py new file mode 100644 index 00000000..7cb276bf --- /dev/null +++ b/services/echo/tests/unit/test_echo.py @@ -0,0 +1,39 @@ +import pytest +from echo.echo import main +from util import ApolloError + + +def test_echoes_the_payload_back() -> None: + assert main({"message": "hello"}) == {"message": "hello"} + + +def test_rejects_an_empty_payload() -> None: + with pytest.raises(ApolloError): + main({}) + + with pytest.raises(ApolloError): + main({"session_id": "abc"}) + + +def test_does_not_return_server_set_fields() -> None: + # The server sets these itself, so their values must not come back out. + echoed = main({"message": "hello", "api_key": "test-value"}) + + assert "test-value" not in str(echoed) + assert echoed["message"] == "hello" + + +def test_masks_a_key_shaped_value_anywhere_in_the_payload() -> None: + # Caught by shape, so a value under a field name nobody listed is still + # masked. Both provider prefixes, since the server may fill in either. + assert "sk-ant-abc123" not in str(main({"note": "uses sk-ant-abc123"})) + assert "sk-proj-0123456789abcdef" not in str( + main({"note": "uses sk-proj-0123456789abcdef"}), + ) + + +def test_masks_a_provider_field_the_server_may_fill_in() -> None: + echoed = main({"message": "hello", "openai_api_key": "test-openai"}) + + assert "test-openai" not in str(echoed) + assert echoed["message"] == "hello" diff --git a/services/echo/tests/unit/test_entry_masks_results.py b/services/echo/tests/unit/test_entry_masks_results.py new file mode 100644 index 00000000..c3b82930 --- /dev/null +++ b/services/echo/tests/unit/test_entry_masks_results.py @@ -0,0 +1,58 @@ +"""The boundary every service result passes through. + +Services are handed a payload with values the server put there, and whatever +they return goes back to the caller. Rather than trust each one not to reflect +those values, `entry.call` masks on the way out. + +Driven through `_masking_probe`, which reflects its payload and masks nothing +itself, so these fail if the boundary stops working. Pointing them at `echo` +would prove only that echo masks. +""" + +import json + +import pytest +from entry import call + +HTTP_INTERNAL_ERROR = 500 + + +@pytest.fixture +def input_file(tmp_path: object) -> object: + def write(payload: dict) -> str: + path = tmp_path / "input.json" + path.write_text(json.dumps(payload)) + return str(path) + + return write + + +def test_masks_a_server_set_field_in_the_result(input_file: object) -> None: + result = call( + "_masking_probe", + input_path=input_file({"api_key": "secret", "x": 1}), + ) + + assert result["api_key"] == "[REDACTED]" + assert result["x"] == 1 + + +def test_masks_a_key_shaped_value_anywhere_in_the_result(input_file: object) -> None: + result = call( + "_masking_probe", + input_path=input_file({"note": "configured with sk-ant-abc123"}), + ) + + assert "sk-ant-abc123" not in json.dumps(result) + + +def test_masks_a_key_that_reached_an_error_message(input_file: object) -> None: + # A service that raises with the payload in scope is the common shape: + # most catch broadly and rewrap as ApolloError(500, str(e)). + result = call( + "_masking_probe_raiser", + input_path=input_file({"api_key": "sk-ant-abc123"}), + ) + + assert result["code"] == HTTP_INTERNAL_ERROR + assert "sk-ant-abc123" not in json.dumps(result) diff --git a/services/entry.py b/services/entry.py index 53ac0e68..fe67040a 100644 --- a/services/entry.py +++ b/services/entry.py @@ -1,17 +1,18 @@ -import sys -import os +import argparse import json +import os import uuid -import argparse -from dotenv import load_dotenv + import sentry_sdk -from util import set_apollo_port, ApolloError +from dotenv import load_dotenv +from util import ApolloError, install_log_masking, set_apollo_port load_dotenv() # Langfuse: init after load_dotenv so env vars are available, before any Anthropic client is created from opentelemetry.instrumentation.anthropic import AnthropicInstrumentor from opentelemetry.instrumentation.threading import ThreadingInstrumentor + AnthropicInstrumentor().instrument() ThreadingInstrumentor().instrument() @@ -42,17 +43,31 @@ def _should_export_span(span): 'unknown': 0.0, } +def _scrub_event(event: dict, _hint: dict) -> dict: + """Mask keys in what Sentry is about to send. + + Sentry scrubs frame locals by name, but not the exception message, a + set_context payload, or a breadcrumb - and services raise + ApolloError(500, str(e)) with the request in scope. The whole event rather + than a list of sections, so a section nobody thought of is covered too. + """ + return mask_secrets(event) + + 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 + auto_enabling_integrations=False, + before_send=_scrub_event, + # before_send covers error events only, and tracing is on. + before_send_transaction=_scrub_event, ) def call( - service: str, *, input_path: str | None = None, output_path: str | None = None, apollo_port: int | None = None + service: str, *, input_path: str | None = None, output_path: str | None = None, apollo_port: int | None = None, ) -> dict: """ Dynamically imports a module and invokes its main function with input data. @@ -71,30 +86,67 @@ def call( data = {} if input_path: try: - with open(input_path, "r") as f: + with open(input_path) as f: data = json.load(f) - except FileNotFoundError: + except FileNotFoundError as e: + # The path is the server's own, so it is for the log, not the + # caller. sentry_sdk.capture_exception(e) - return ApolloError(code=500, message=f"Input file not found: {input_path}", type="INTERNAL_ERROR").to_dict() - except json.JSONDecodeError: + return _finish( + ApolloError( + code=500, message="Input file not found", type="INTERNAL_ERROR" + ).to_dict(), + output_path, + ) + except json.JSONDecodeError as e: sentry_sdk.capture_exception(e) - return ApolloError(code=500, message="Invalid JSON input", type="INTERNAL_ERROR").to_dict() + return _finish( + ApolloError( + code=500, message="Invalid JSON input", type="INTERNAL_ERROR" + ).to_dict(), + output_path, + ) try: m = __import__(module_name, fromlist=["main"]) + + # Again here, after every import has had its chance to install a + # handler of its own. Some libraries add one that writes to stderr, + # which the bridge forwards to the caller line for line. + install_log_masking() + result = m.main(data) except ModuleNotFoundError as e: sentry_sdk.capture_exception(e) - return ApolloError(code=500, message=str(e), type="INTERNAL_ERROR").to_dict() + result = ApolloError( + code=500, message=str(e), type="INTERNAL_ERROR", + ).to_dict() except ApolloError as e: sentry_sdk.capture_exception(e) result = e.to_dict() except Exception as e: sentry_sdk.capture_exception(e) - result = ApolloError(code=500, message=str(e), type="INTERNAL_ERROR").to_dict() + result = ApolloError( + code=500, message=str(e), type="INTERNAL_ERROR", + ).to_dict() langfuse.flush() + return _finish(result, output_path) + + +def _finish(result: dict, output_path: str | None) -> dict: + """Mask, write the result where the caller expects it, then hand it back. + + Every path out of `call` comes through here, which buys two things. The + output file is always written, so the bridge can read an empty one as the + run having died rather than as a polite failure. And a value the server put + on the payload cannot leave down a branch someone forgot: most services + catch broadly and rewrap as `ApolloError(500, str(e))`, so masking + per-branch would miss the one nearly all of them take. + """ + result = mask_secrets(result) + if output_path: with open(output_path, "w") as f: json.dump(result, f) diff --git a/services/global_chat/tests/unit/test_mask_secrets.py b/services/global_chat/tests/unit/test_mask_secrets.py new file mode 100644 index 00000000..cdbac18b --- /dev/null +++ b/services/global_chat/tests/unit/test_mask_secrets.py @@ -0,0 +1,91 @@ +"""Boundaries of the shared mask. + +`mask_secrets` is both the Langfuse export mask and what stops a service +handing a key back to its caller, so it has to catch keys without eating the +workflow YAML, job code and prose it also passes over. +""" + +import pytest +from langfuse_util import mask_secrets + +KEY_SHAPED = [ + "sk-ant-abc123", + "sk-ant-api03-Zm9vYmFy_baz", + "sk-proj-0123456789abcdef", + "use sk-ant-abc123 to authenticate", + "(sk-ant-abc123)", + # A key can legitimately follow a hyphen or underscore, so only an + # alphanumeric before the sk- means it is the tail of a word. + "-sk-ant-zzzzzzzzzzzzzzzzzz", + "_sk-ant-zzzzzzzzzzzzzzzzzz", + "sk-svcacct-abcdefgh", + "sk-abcdefghijklmnopqrstuvwx", + "Bearer sk-ant-abc12345", +] + +# Every one of these ends in a token the pattern used to bite through: task-, +# risk-, disk-, mask-, kiosk-. +ORDINARY_TEXT = [ + "task-scheduler-configuration", + "risk-assessment-framework", + "disk-usage-monitoring-job", + "mask-generation-pipeline", + "kiosk-registration-workflow-id", + "https://docs.openfn.org/build/task-automation-guide", + "const task = 'task-runner-config-value';", + "obelisk-carving-notes", + "whisk-attachment-guide", + # The other direction: a name that starts with sk- and carries on in + # words. Only a provider prefix or an unbroken high-entropy run counts. + "sk-antelope-migration-plan", + "step-sk-mapping-export-v2", + "load-sk-patient-records-job", + "--sk-ignore-case-sensitivity", +] + + +@pytest.mark.parametrize( + "name", + ["api_key", "api-key", "apiKey", "X-Api-Key", "Authorization"], +) +def test_masks_a_secret_field_however_it_is_spelled(name: str) -> None: + assert mask_secrets({name: "anything"})[name] == "[REDACTED]" + + +def test_does_not_recurse_without_bound() -> None: + nested = {"leaf": "value"} + for _ in range(600): + nested = {"next": nested} + + assert isinstance(mask_secrets(nested), dict) + + +@pytest.mark.parametrize("text", KEY_SHAPED) +def test_masks_a_key_shaped_value(text: str) -> None: + assert "[REDACTED]" in mask_secrets(text) + + +@pytest.mark.parametrize("text", ORDINARY_TEXT) +def test_leaves_ordinary_text_alone(text: str) -> None: + assert mask_secrets(text) == text + + +def test_masks_by_field_name_whatever_the_value_looks_like() -> None: + masked = mask_secrets({"api_key": "not-key-shaped", "message": "hello"}) + + assert masked["api_key"] == "[REDACTED]" + assert masked["message"] == "hello" + + +def test_keeps_the_langfuse_public_key() -> None: + # Public by design, and useful in a trace for telling which project a span + # was bound for. + masked = mask_secrets({"langfuse_public_key": "pk-lf-123"}) + + assert masked["langfuse_public_key"] == "pk-lf-123" + + +def test_masks_nested_and_listed_values() -> None: + masked = mask_secrets({"outer": [{"api_key": "sk-ant-abc123"}]}) + + assert masked["outer"][0]["api_key"] == "[REDACTED]" diff --git a/services/job_chat/job_chat.py b/services/job_chat/job_chat.py index 368f8573..919040ed 100644 --- a/services/job_chat/job_chat.py +++ b/services/job_chat/job_chat.py @@ -18,7 +18,7 @@ ) import sentry_sdk from langfuse import observe, propagate_attributes, get_client as get_langfuse_client -from langfuse_util import should_track, build_tags, build_generation_diff +from langfuse_util import should_track, build_tags, build_generation_diff, mask_secrets from util import ApolloError, create_logger, AdaptorSpecifier, add_page_prefix, APOLLO_VERSION from yaml_utils import INSPECT_JOB_CODE_TOOL, inspect_job_code from .prompt import build_prompt, build_error_correction_prompt @@ -742,9 +742,15 @@ def main(data_dict: dict) -> dict: Main entry point with improved error handling and input validation. """ try: - sentry_sdk.set_context("request_data", { - k: v for k, v in data_dict.items() if k not in ("api_key", "_stream_manager") - }) + # The stream manager is an object rather than data, so it is dropped. + # Everything else goes through the shared mask instead of a per-service + # name list, which catches nested values and key-shaped strings too. + sentry_sdk.set_context( + "request_data", + mask_secrets( + {k: v for k, v in data_dict.items() if k != "_stream_manager"}, + ), + ) data = Payload.from_dict(data_dict) diff --git a/services/langfuse_util.py b/services/langfuse_util.py index 31913882..1a408fe8 100644 --- a/services/langfuse_util.py +++ b/services/langfuse_util.py @@ -5,24 +5,71 @@ import yaml -_SECRET_KEY_NAMES = {"api_key", "anthropic_api_key", "authorization"} -_SECRET_VALUE_PATTERN = re.compile(r"sk-ant-[\w\-]+") - - -def mask_secrets(data: Any) -> Any: # noqa: ANN401 +# Every field the server may fill in on a payload, not just the one it fills +# in today: a value under one of these belongs to the deployment rather than +# to the caller. langfuse_public_key is deliberately absent - it is meant to +# be visible, and masking it would cost a useful identifier for nothing. +_SECRET_KEY_NAMES = { + "api_key", + "anthropic_api_key", + "openai_api_key", + "pinecone_api_key", + "langfuse_secret_key", + "authorization", + "x_api_key", +} + +# json.loads accepts around a thousand levels of nesting and this runs over +# whatever a caller sends, so a deep payload could otherwise take the process +# down with a RecursionError. +_MAX_DEPTH = 50 + +# A backstop for a key inside a string, where the name list cannot see it. +# Deliberately narrow: this also passes over workflow YAML, job code and +# service results, so a loose pattern corrupts a caller's own data. Hence a +# known provider prefix, or a run too long and unbroken to be a name. +# test_mask_secrets.py pins both directions. +_SECRET_VALUE_PATTERN = re.compile( + r"(? str: + return str(key).lower().replace("-", "").replace("_", "") + + +# Compared in normalised form, so api-key, apiKey and X-Api-Key are all caught +# by the one readable entry above. +_NORMALISED_SECRET_NAMES = {_normalise_name(name) for name in _SECRET_KEY_NAMES} + + +def _is_secret_name(key: object) -> bool: + return _normalise_name(key) in _NORMALISED_SECRET_NAMES + + +def mask_secrets(data: Any, _depth: int = 0) -> Any: # noqa: ANN401 """Langfuse mask callback: redact API keys from all traced data. Applied by the SDK to every span's input, output and metadata before export, so keys passed as function arguments to @observe-decorated - functions never reach Langfuse. + functions never reach Langfuse. Also used anywhere a service's own output + could carry a value the server put in the payload. """ + if _depth > _MAX_DEPTH: + return "[TRUNCATED]" + if isinstance(data, dict): return { - k: "[REDACTED]" if str(k).lower() in _SECRET_KEY_NAMES and v else mask_secrets(v) + k: "[REDACTED]" + if _is_secret_name(k) and v + else mask_secrets(v, _depth + 1) for k, v in data.items() } if isinstance(data, (list, tuple)): - return [mask_secrets(v) for v in data] + return [mask_secrets(v, _depth + 1) for v in data] if isinstance(data, str): return _SECRET_VALUE_PATTERN.sub("[REDACTED]", data) return data diff --git a/services/load_adaptor_docs/load_adaptor_docs.py b/services/load_adaptor_docs/load_adaptor_docs.py index bf1cc13e..cc265f22 100644 --- a/services/load_adaptor_docs/load_adaptor_docs.py +++ b/services/load_adaptor_docs/load_adaptor_docs.py @@ -3,6 +3,7 @@ from typing import Dict, List, Any from psycopg2.extras import execute_values import sentry_sdk +from langfuse_util import mask_secrets from util import create_logger, ApolloError, apollo, AdaptorSpecifier, get_db_connection logger = create_logger("load_adaptor_docs") @@ -349,9 +350,7 @@ def main(data: dict) -> dict: """ logger.info("Starting load_adaptor_docs service...") - sentry_sdk.set_context("request_data", { - k: v for k, v in data.items() if k not in ["api_key"] - }) + sentry_sdk.set_context("request_data", mask_secrets(data)) # Validate required fields if "adaptor" not in data: diff --git a/services/search_adaptor_docs/search_adaptor_docs.py b/services/search_adaptor_docs/search_adaptor_docs.py index 31ede7b9..d089b1bd 100644 --- a/services/search_adaptor_docs/search_adaptor_docs.py +++ b/services/search_adaptor_docs/search_adaptor_docs.py @@ -1,8 +1,9 @@ import time -from typing import Dict, List, Any + import sentry_sdk -from util import create_logger, ApolloError, AdaptorSpecifier, get_db_connection +from langfuse_util import mask_secrets from load_adaptor_docs.load_adaptor_docs import load_adaptor_docs +from util import AdaptorSpecifier, ApolloError, create_logger, get_db_connection logger = create_logger("search_adaptor_docs") @@ -22,7 +23,7 @@ def ensure_docs_loaded(adaptor: AdaptorSpecifier, conn, skip_if_exists: bool = T load_result = load_adaptor_docs( adaptor=adaptor.specifier, skip_if_exists=skip_if_exists, - conn=conn + conn=conn, ) duration = time.time() - start_time @@ -34,7 +35,7 @@ def ensure_docs_loaded(adaptor: AdaptorSpecifier, conn, skip_if_exists: bool = T logger.warning(f"Failed to load adaptor docs for {adaptor.specifier} after {duration:.3f}s") except Exception as e: duration = time.time() - start_time if 'start_time' in locals() else 0 - logger.warning(f"Failed to load adaptor docs after {duration:.3f}s: {str(e)}") + logger.warning(f"Failed to load adaptor docs after {duration:.3f}s: {e!s}") sentry_sdk.capture_exception(e) @@ -189,7 +190,7 @@ def fetch_all_functions(adaptor: AdaptorSpecifier, conn, format: str = "json", a return [ { "function_name": row[0], - "text": json_to_natural_language(row[1], adaptor) + "text": json_to_natural_language(row[1], adaptor), } for row in rows ] @@ -211,7 +212,7 @@ def main(data: dict) -> dict: """ logger.info("Starting search_adaptor_docs...") - sentry_sdk.set_context("request_data", data) + sentry_sdk.set_context("request_data", mask_secrets(data)) # Validate required fields if "adaptor" not in data: @@ -254,7 +255,7 @@ def main(data: dict) -> dict: "adaptor": adaptor.name, "version": adaptor.version, "query_type": "list", - "functions": functions + "functions": functions, } elif query_type == "signatures": @@ -265,7 +266,7 @@ def main(data: dict) -> dict: "adaptor": adaptor.name, "version": adaptor.version, "query_type": "signatures", - "signatures": signatures + "signatures": signatures, } elif query_type == "function": @@ -281,7 +282,7 @@ def main(data: dict) -> dict: "query_type": "function", "function_name": function_name, "format": format, - "data": func_data + "data": func_data, } elif query_type == "all": @@ -294,14 +295,14 @@ def main(data: dict) -> dict: "query_type": "all", "format": format, "count": len(functions), - "functions": functions + "functions": functions, } except ApolloError: raise except Exception as e: - logger.error(f"Error querying database: {str(e)}") - raise ApolloError(500, f"Query failed: {str(e)}", type="DATABASE_ERROR") + logger.error(f"Error querying database: {e!s}") + raise ApolloError(500, f"Query failed: {e!s}", type="DATABASE_ERROR") finally: conn.close() diff --git a/services/streaming_util.py b/services/streaming_util.py index 08cb6dce..6147c5cb 100644 --- a/services/streaming_util.py +++ b/services/streaming_util.py @@ -11,6 +11,7 @@ from dataclasses import dataclass from typing import Any +from langfuse_util import mask_secrets from models import CLAUDE_SONNET # Shared status message pools for user-facing progress indicators. @@ -132,8 +133,16 @@ def _emit_event(self, event_type: str, data: dict[str, Any]) -> None: """ # Use EVENT: prefix format that bridge.ts expects # Bridge will convert this to proper SSE format + # + # Masked because this is a third way out to the caller, alongside the + # result and the log stream: it is neither, so neither of their masks + # sees it. Nothing puts a key in an event today, which is the point of + # doing it while that is still true. if self.stream: - print(f"EVENT:{event_type}:{json.dumps(data)}", flush=True) # noqa: T201 + print( # noqa: T201 + f"EVENT:{event_type}:{json.dumps(mask_secrets(data))}", + flush=True, + ) def start_stream(self) -> None: """ diff --git a/services/test_slow/test_slow.py b/services/test_slow/test_slow.py new file mode 100644 index 00000000..42f382f3 --- /dev/null +++ b/services/test_slow/test_slow.py @@ -0,0 +1,24 @@ +"""A mounted service that does nothing slowly, so route-level cancellation can +be tested over real connections. + +The bridge's own cancellation tests use the unmounted _cancel_probe and hand +run() a signal directly. That cannot see route wiring - it missed the +websocket close handler looking runs up under the wrong key - so this one is +mounted, like test_errors, and driven through a real socket. + +It announces itself before sleeping because the spawned command is +`poetry run python ...`: the process list matches on the poetry wrapper +seconds before the interpreter has booted, so waiting on that would prove +nothing. +""" + +import time + + +def main(data_dict: dict) -> dict: + print("EVENT:probe_started:{}", flush=True) # noqa: T201 + + seconds = data_dict.get("sleep_for", 1) + time.sleep(seconds) + + return {"slept": seconds} diff --git a/services/util.py b/services/util.py index 45d6d0f7..77c39234 100644 --- a/services/util.py +++ b/services/util.py @@ -6,6 +6,7 @@ import psycopg2 import requests +from langfuse_util import mask_secrets APOLLO_VERSION = os.getenv("APOLLO_VERSION", "unknown") @@ -57,19 +58,77 @@ def to_dict(self) -> dict: return error_dict -filename = None loggers: dict[str, logging.Logger] = {} apollo_port = 3000 -def set_log_output(f: str | None) -> None: - """Set the output file for logging.""" - global filename # noqa: PLW0603 +class _MaskingFilter(logging.Filter): + """Masks key-shaped values on their way out to the log stream. - if f is not None: - print(f"[entry.py] writing logs to {f}") # noqa: T201 + This output does not stay on the server: the bridge matches the log prefix + on stdout and forwards the line to the caller as an SSE event. A service + that logs its own payload would therefore hand the caller the key the + server put there. - filename = f + Attached to the stdout handler rather than to each logger, because a + filter on a logger only runs for records emitted through it - a plain + `logging.getLogger(__name__)`, or a third-party logger like httpx, writes + to the same handler and would sail past. + + Note the two levels are not equally strong. A dict is masked by field name + and by value shape; anything already rendered to text, including a payload + interpolated into an f-string, has only the shape to go on. + """ + + def filter(self, record: logging.LogRecord) -> bool: + # A container is worth masking structurally, so field names apply. + if isinstance(record.msg, (dict, list, tuple)) and not record.args: + record.msg = mask_secrets(record.msg) + return True + + # Otherwise mask the text that will actually be emitted. Rendering + # normally happens inside the handler, where a bad format string is + # caught and reported; here it would escape into the caller's own + # logging call, so a cosmetic typo must not fail their request. + try: + rendered = record.getMessage() + except Exception: + return True + + record.msg = mask_secrets(rendered) + record.args = () + return True + + +_masking_filter = _MaskingFilter() + + +def install_log_masking() -> None: + """Put the mask on every handler that exists right now. + + Called at import and again from create_logger, because handlers appear at + two different times: some libraries install their own before any service + module is imported (langfuse attaches one to the httpx logger, and it + writes to stderr, which the bridge forwards to the caller line for line), + and the root handler only exists once basicConfig has run. + + A filter on a handler covers every record reaching that stream whatever + logger produced it - which a filter on a logger does not. + + loggerDict is copied rather than walked live: a thread creating a logger + resizes it mid-walk, which raises RuntimeError. + """ + root = logging.getLogger() + known = [root, *( + logger + for logger in list(root.manager.loggerDict.values()) + if isinstance(logger, logging.Logger) + )] + + for logger in known: + for handler in logger.handlers: + if _masking_filter not in handler.filters: + handler.addFilter(_masking_filter) def create_logger(name: str) -> logging.Logger: @@ -78,12 +137,16 @@ def create_logger(name: str) -> logging.Logger: Logs to stdout by default. """ logging.basicConfig(level=logging.INFO, stream=sys.stdout) + install_log_masking() + if name not in loggers: - logger = logging.getLogger(name) - loggers[name] = logger + loggers[name] = logging.getLogger(name) return loggers[name] +install_log_masking() + + def set_apollo_port(p: int) -> None: """Set the port for Apollo services.""" global apollo_port # noqa: PLW0603 diff --git a/services/workflow_chat/workflow_chat.py b/services/workflow_chat/workflow_chat.py index 67ee6711..1377ad42 100644 --- a/services/workflow_chat/workflow_chat.py +++ b/services/workflow_chat/workflow_chat.py @@ -61,7 +61,7 @@ ) import sentry_sdk from langfuse import observe, propagate_attributes, get_client as get_langfuse_client -from langfuse_util import should_track, build_tags, build_generation_diff +from langfuse_util import should_track, build_tags, build_generation_diff, mask_secrets from util import ApolloError, create_logger, add_page_prefix, APOLLO_VERSION from .gen_project_prompt import build_prompt from workflow_chat.available_adaptors import get_available_adaptors @@ -705,9 +705,15 @@ def main(data_dict: dict) -> dict: Main entry point with improved error handling and input validation. """ try: - sentry_sdk.set_context("request_data", { - k: v for k, v in data_dict.items() if k not in ("api_key", "_stream_manager") - }) + # The stream manager is an object rather than data, so it is dropped. + # Everything else goes through the shared mask instead of a per-service + # name list, which catches nested values and key-shaped strings too. + sentry_sdk.set_context( + "request_data", + mask_secrets( + {k: v for k, v in data_dict.items() if k != "_stream_manager"}, + ), + ) data = Payload.from_dict(data_dict)