diff --git a/.agents/specs/REMOTE-2516-debug-archive-worker-logs.md b/.agents/specs/REMOTE-2516-debug-archive-worker-logs.md new file mode 100644 index 0000000..0dd8964 --- /dev/null +++ b/.agents/specs/REMOTE-2516-debug-archive-worker-logs.md @@ -0,0 +1,543 @@ +# REMOTE-2516: Self-hosted worker debug-archive logs + +## Document status +This is the `warpdotdev/oz-agent-worker` companion to the debug-archive product and technical contract in [`warpdotdev/warp-server#13839`](https://github.com/warpdotdev/warp-server/pull/13839). The server spec owns archive creation, authorization, storage, trigger policy, Temporal orchestration, GCP Pub/Sub fan-out, and the server half of the WebSocket protocol. This document owns the worker half: reporting the worker build on connection so the exact assigned version can be snapshotted, proving execution ownership, capturing Docker/Kubernetes/direct-backend logs, retaining ownership and retrievable logs for the existing cleanup grace, applying the requested content-transformer hook, uploading to a server-supplied target, and acknowledging the result. + +Both documents are required to implement REMOTE-2516. If their wire contracts differ, implementation must reconcile the specs before either PR is promoted; neither side should silently infer a different protocol. + +The requester additionally confirmed that the archive must preserve the client version, worker version, and server commit associated with run/execution creation, and that local client plus entrypoint-script logs are the sandbox stdout/stderr source. This companion spec implements the worker-version and self-hosted stdout/stderr portions of that decision. + +## Summary +Warp staff, run creators, and owning-team admins need a HAR-like debug archive for cloud-agent runs. For self-hosted executions, the logs live inside infrastructure controlled by `oz-agent-worker`. `warp-server` cannot query a customer's Docker daemon, Kubernetes cluster, or direct child process, and the current worker removes its task ownership record and often destroys the backend resource as soon as execution ends. + +V1 adds build provenance plus a request-driven worker protocol. Every authenticated worker connection reports the already build-stamped `main.Version`; `warp-server` snapshots the exact selected connection's value when an execution is claimed rather than reading a current version during archive collection. The server publishes a versioned worker-control request through its existing GCP Pub/Sub fan-out and forwards `debug_archive_logs_requested`, with one immutable 30-minute presigned upload target, to every connected process using the assigned worker ID. Each process checks an exact `(run_id, execution_id)` active/cleanup-grace ownership registry. Only the process that executed the assignment snapshots both stdout and stderr from the sandbox entrypoint/launcher and client process, applies the request's content-transformer descriptor while encoding message data, uploads bounded NDJSON directly to the target, and sends `debug_archive_logs_uploaded`. Other processes silently ignore the request. Missing logs, old workers, disconnected workers, upload failures, and expired cleanup grace make only the self-hosted source partial; they never fail or pause the cloud-agent run. + +V1 supports snapshots during an active run and retrieval after terminal reporting for the existing configured cleanup grace. It does not continuously append to one object. The design leaves room for a future mode that uploads immutable, monotonically numbered chunks while a run is active. + +## Product contract + +### Users and outcomes +The direct consumer is the debug-archive collector in `warp-server`; human authorization and archive download are server responsibilities. + +For a supported self-hosted execution: +- An active collection captures logs observed up to a consistent snapshot watermark without interrupting the agent. +- A collection triggered after an execution failure can still retrieve a bounded snapshot before the worker's existing cleanup grace expires; after that grace, partial/unavailable is expected. +- A later collection can replace the logical archive source with a newer snapshot. +- The archive manifest can distinguish uploaded, unavailable, unsupported, failed, and truncated worker-log outcomes; non-owning instances are silent and do not create a separate outcome. + +### Supported backends +V1 captures: +- Docker container stdout and stderr for the exact execution container, including all output emitted by the container entrypoint/launcher script and the client process it starts. +- Kubernetes stdout and stderr for every init container and regular container in every pod belonging to the exact execution Job, including entrypoint/launcher and client output plus current and best-effort previous logs after a restart. +- Direct-backend stdout and stderr for the entrypoint/launcher invocation, client process, and worker-managed setup/teardown phases. + +These are the sandbox's local stdout/stderr streams, not separate semantic transcript sources. The worker does not filter lines by presumed process. It preserves stdout/stderr and provider container/phase metadata when available; when a provider merges streams or does not identify whether a line came from the entrypoint or child client, the output is truthfully labeled `combined`/`unknown` rather than assigned a fabricated process identity. + +The command backend dispatches into an opaque operator-owned runtime and has no backend log API. A worker that owns a command-backend assignment responds `unavailable` with `backend_not_supported`; it must not claim that dispatch-command stdout is the remote agent's execution log. + +### Partial behavior +Log collection is best effort: +- A request is independent from task execution and never blocks assignment, process progress, terminal reporting, teardown, or worker reconnection. +- A failure to initialize or write direct-process capture never rejects or fails the assigned task. +- An inability to capture or upload returns a sanitized acknowledgement when possible. If the worker is disconnected or too old to understand the message, the server times out. +- The server publishes the remainder of the archive with a partial/stale self-hosted-log source according to the server spec. + +### Cleanup grace and bounds +V1 reuses the execution's existing cleanup grace instead of adding an archive-specific retention setting. The resolved duration follows the current precedence already emitted as `--idle-on-complete`: task `agent_config_snapshot.idle_timeout_minutes`, then worker `idle_on_complete`, then the Oz default (currently 45 minutes). The worker keeps the exact ownership/resource entry and any direct-process `TaskLogCapture` until that grace expires, then performs normal cleanup. A request never extends the configured grace. + +Self-hosted operators that want reliable `ANY_FAILURE` archive collection must configure a grace long enough for the terminal event to reach the server and for a request/upload that may wait up to 30 minutes. Kubernetes `ttl_seconds_after_finished` must not be shorter than the effective worker grace when provider logs are the retained source. A shorter grace is supported but may produce a partial archive. + +Every request is bounded to 64 MiB per execution, preserving first/last records with an explicit NDJSON truncation record. The server request may lower, but never raise, this ceiling. At most two upload handlers run concurrently per worker process, and direct captures plus request snapshots share a default 1-GiB process-local disk budget. Direct-process capture uses one secure bounded task-local temporary file through the existing grace; Docker and Kubernetes stay provider-backed and do not create a second post-terminal log spool. No V1 guarantee survives worker process/host replacement or loss of a Kubernetes worker pod. + +### Version skew +- An old worker ignores the unknown request; the server times out and records the source as unavailable. +- An old worker that does not report its build version remains eligible for ordinary task execution; the server records worker provenance as `not_reported` and makes a new archive partial rather than substituting a current image tag or connection's version. +- A new worker accepts protocol version 1 and ignores unknown optional JSON fields. +- A new worker receiving an unsupported protocol version sends `failed` with reason `unsupported_protocol_version`. +- A server that does not send archive requests does not change worker behavior. +- Protocol handling is disabled only by the absence of the new server message; no worker rollout may make ordinary task execution depend on archive availability. + +### Worker build provenance +`main.Version` is already stamped by release builds through `-ldflags="-X main.Version=..."`, defaults to `dev` locally, and is exported in `oz_worker_info`; it is not currently present on the WebSocket wire. + +Pass `main.Version` into `worker.Config` and send it as `X-Warp-Worker-Version` on every authenticated WebSocket dial and reconnect. The value is an opaque, non-secret build identifier capped at 128 UTF-8 bytes. The server validates and stores it on that exact `WorkerConnection`; when the connection is selected and successfully claims an execution, the server copies the value into the execution's set-once provenance snapshot in the claim transaction, before sandbox launch. Two processes sharing one worker ID may report different versions during a rolling deploy, so the selected connection—not worker ID, current presence, metrics, image tag, or collection-time state—is authoritative. + +An empty/missing header is accepted for backward compatibility and surfaced as `not_reported`; an overlong or control-character-bearing value is ignored as invalid provenance without exposing it in logs or preventing connection/task execution. Reconnects report again for future assignments but never rewrite a previously claimed execution's version. The version field is diagnostic only and is not used to negotiate the debug-log protocol. + +### Security and privacy +Worker logs may contain prompts, source code, identifiers, and secrets emitted by a process. V1 includes a versioned `ContentTransformer` hook in the request/upload path but implements only a no-op transformer, matching the server spec. Actual redaction rules are deferred. The worker: +- keeps direct-capture directories mode `0700` and files mode `0600`; +- never logs log content, presigned URLs, signed headers, multipart fields, or upload response bodies; +- sends bytes only to the exact authenticated-server-supplied target; +- permits HTTPS targets, with plain HTTP allowed only for loopback integration tests/local development; +- disables HTTP redirects so a signed target cannot redirect bytes elsewhere; +- deletes direct-capture bytes when the existing cleanup grace expires or earlier when normal task cleanup no longer needs them. + +The transformer is applied while NDJSON is first encoded/uploaded: only log-message `data` passes through it, while timestamps, sequence, stream/container identity, and source-error metadata remain structural. V1's no-op descriptor preserves bytes. Unsupported future descriptors return `failed/unsupported_content_transformer`; no worker silently uploads untransformed data. A future redaction implementation can therefore evolve the content rules without adding a server-side raw-object-to-transformed-object copy. + +## Current state +- `main.go:23` already exposes build-stamped `Version` and uses it for worker metrics, while `internal/worker/worker.go:187` sends only worker ID and authorization during WebSocket connection; the server cannot currently tie a worker build to an execution. +- `internal/types/messages.go` defines only assignment, lifecycle, and cancellation message DTOs. Unknown server messages are logged and ignored. +- `internal/worker/worker.go:323` decodes WebSocket messages synchronously on the read loop. Archive work cannot run inline there without blocking heartbeats, cancellation, or later assignments. +- `internal/worker/worker.go:62` tracks active tasks only by run/task ID. `executeTask` removes the record immediately after terminal reporting unless the command backend spawned the task. +- `internal/worker/backend.go:38` has execution, cancellation, shutdown, and shutdown-preservation methods, but no log snapshot contract. +- `internal/worker/docker.go:99` knows the exact container ID only inside `ExecuteTask`; a defer force-removes that container before the worker reports terminal state. Existing `getContainerLogs` reads the whole multiplexed stream into memory and does not preserve stdout/stderr identity. +- `internal/worker/kubernetes.go:406` deletes successful Jobs immediately. Existing diagnostic collection reads at most 1 MiB per container into memory, omits timestamps/stream metadata, and is not exposed outside failure logging. +- `internal/worker/direct.go:233` connects agent stdout/stderr directly to the worker process and removes the per-task workspace after execution. No task-scoped log bytes survive. +- `internal/common/task_utils.go:125` already resolves the task/worker/default `--idle-on-complete` cleanup grace. `internal/config/config.go:16` exposes the legacy worker-level override while task `idle_timeout_minutes` takes precedence. + +These constraints require changes in both the shared worker lifecycle and all three supported backends; adding only a WebSocket message would race terminal cleanup and fail the primary `ANY_FAILURE` use case. The implementation must reuse the resolved grace as a post-terminal ownership/provider-cleanup deadline rather than add a second retention configuration. + +## Technical design + +### Connection registration and build provenance +Extend `worker.Config` with `Version` and initialize it from `main.Version`. `Worker.connect` sends `X-Warp-Worker-Version` with every WebSocket upgrade request. This is connection metadata, not a new asynchronous message: it reaches the server before that connection is eligible for task selection. + +The server stores the validated value on its in-memory/presence connection record. The selected connection returns its version as part of the safe claim result; `warp-server` commits it to the execution's set-once `debug_provenance_snapshot` in the same transaction that claims the execution. The worker does not send a later “current version” response to archive collection, and the debug-log upload acknowledgement need not duplicate the persisted build value. + +Tests use two simultaneous connections with one worker ID and different `Version` values to prove the execution receives the version of the connection that actually accepts it. Missing/invalid version metadata does not fail connection, claim, or execution, but is observably absent for the archive's partial-provenance rules. + +### Wire protocol +Add the following v1 message types to `internal/types/messages.go`. + +Server to worker: + +```json +{ + "type": "debug_archive_logs_requested", + "data": { + "protocol_version": 1, + "request_id": "uuid", + "archive_id": "uuid", + "collection_id": "uuid", + "run_id": "uuid", + "execution_id": "uuid", + "requested_format": "application/x-ndjson", + "expires_at": "RFC3339 timestamp", + "max_bytes": 67108864, + "content_transformer": { + "kind": "noop", + "version": 1 + }, + "upload_target": { + "url": "https://...", + "method": "PUT or POST", + "headers": {"provider-header": "value"}, + "multipart_fields": {"provider-field": "value"} + } + } +} +``` + +Worker to server: + +```json +{ + "type": "debug_archive_logs_uploaded", + "data": { + "protocol_version": 1, + "request_id": "uuid", + "archive_id": "uuid", + "collection_id": "uuid", + "run_id": "uuid", + "execution_id": "uuid", + "outcome": "uploaded", + "backend_kind": "docker", + "bytes": 1234, + "crc32c": "base64", + "sha256": "hex", + "truncated": false, + "content_transformer_version": 1, + "capture_status": "complete", + "warning_codes": [], + "reason_code": "", + "message": "" + } +} +``` + +Allowed outcomes are `uploaded`, `unavailable`, and `failed`. Uploaded outcomes set `capture_status` to `complete` or `partial`; partial capture includes at most 16 deduplicated stable `warning_codes` such as `container_logs_unavailable`, `previous_logs_unavailable`, `output_dropped`, or `provider_snapshot_incomplete`. Non-upload outcomes omit byte/checksum/capture fields and use stable bounded reason codes: +- `unsupported_protocol_version` +- `unsupported_content_transformer` +- `invalid_request` +- `request_expired` +- `backend_not_supported` +- `resource_not_ready` +- `resource_not_found` +- `cleanup_grace_expired` +- `capture_unavailable` +- `snapshot_failed` +- `upload_rejected` +- `upload_expired` +- `upload_failed` +- `worker_shutting_down` +- `request_capacity_exhausted` + +Human-readable `message` is sanitized and capped at 256 UTF-8 bytes. It never includes provider output, a URL, request headers/fields, a local path, or an HTTP response body. + +Required request validation occurs before ownership/upload work: +- protocol version is 1; +- identifiers and requested format are non-empty and valid; +- `expires_at` is in the future; +- `expires_at` is no more than 30 minutes after receipt, and the content-transformer descriptor is supported (v1 accepts only `noop` version 1); +- `max_bytes` is positive and no larger than the protocol's 64-MiB ceiling; the effective output bound is the minimum of the request and configured ceilings; +- method is `PUT` or `POST`; +- PUT has no multipart fields, while POST has the fields supplied by the server; +- target/header/field values contain no control characters; +- target scheme satisfies the HTTPS/loopback rule. + +### Ownership and request state +Replace the single-purpose active-task map with a `TaskRegistry` that exposes exact compound-key lookup while retaining existing cancellation behavior: +- `active[(run_id, execution_id)]` stores task context/cancel state, backend kind, normalized backend resource identity, and its `TaskLogCapture`. +- `cleanup_grace[(run_id, execution_id)]` stores terminal outcome time, the already-resolved `idle-on-complete` deadline, backend/resource identity, and the direct-process `TaskLogCapture` when applicable. +- The existing run-ID cancellation lookup remains available because cancellation messages carry only `task_id`. + +At assignment, the worker creates the active registry entry and, for direct execution only, a bounded capture before starting the backend. If direct capture initialization fails, it records a non-fatal capture status and continues the task. + +Immediately after `Backend.ExecuteTask` returns, the worker computes the existing cleanup deadline from the task's resolved `idle-on-complete` duration and atomically moves ownership from active to `cleanup_grace` before `task_failed`, `task_completed`, or `task_cancelled` is enqueued. Docker/Kubernetes keep their exact registered resources until that deadline; direct execution finalizes a consistent capture watermark but retains the capture until the same deadline. A cleanup timer invokes normal backend resource cleanup and removes the entry at expiry. This ordering guarantees that an `ANY_FAILURE` server request triggered by the terminal message sees the grace entry rather than racing cleanup. + +Every connected process with the same worker ID receives the server broadcast. A process checks both IDs: +- Exact active/cleanup-grace match: it is the owner and handles the request. +- Run exists with a different execution, or neither exact key exists: silently ignore the request without creating request-cache state or sending an acknowledgement. +- Exact command-backend match: return `unavailable/backend_not_supported`. + +A request cache of at most 1,024 entries keyed by `request_id` makes duplicate delivery idempotent: +- one in-flight goroutine owns snapshot/upload; +- concurrent duplicates attach to that result; +- a completed duplicate replays the same acknowledgement without re-uploading; +- reuse of one request ID with different immutable fields returns `failed/invalid_request`. + +Completed request metadata remains until the later of request expiry and the cleanup-grace deadline. It contains IDs, transformer version, outcome, size/checksums, capture status/warnings, and reason only—not the target. Expired entries are removed first when the cache reaches its limit. If no expired entry is available, a new owning-worker request returns `failed/request_capacity_exhausted`; task execution and existing request handlers continue. + +### Asynchronous request coordinator +`handleMessage` parses and validates the envelope, then hands archive requests to a `DebugLogCoordinator` goroutine. It never performs provider reads or network uploads on the WebSocket read loop. + +The coordinator: +1. Performs ownership and idempotency checks. +2. Acquires the global upload semaphore and per-execution snapshot mutex. +3. Re-checks request expiry and ownership. +4. Creates one secure request-scoped bounded snapshot and instantiates the validated `ContentTransformer`. +5. For active or cleanup-grace Docker/Kubernetes executions, calls the exact live/retained backend snapshot method. For active or cleanup-grace direct executions, takes a consistent watermark copy of `TaskLogCapture`. +6. Encodes records through the transformer so message `data` is transformed while metadata remains structural. +7. Closes and verifies the local request snapshot. +8. Uploads it to the supplied target while computing CRC32C and SHA-256. +9. Enqueues one acknowledgement through the existing single WebSocket writer. +10. Records the completed request result and releases resources. + +Queueing never extends `expires_at`. A request that waits past expiry reports `upload_expired` without contacting the target. Worker shutdown cancels in-flight requests and does not delay ordinary shutdown beyond the existing bounded backend shutdown. + +### Backend contract +Extend `Backend` with: + +```go +SnapshotTaskLogs(ctx context.Context, taskID, executionID string, writer io.Writer) error +``` + +The method writes protocol-v1 NDJSON records and returns typed errors that the coordinator maps to stable reason codes. `PartialSnapshotError` carries bounded warning codes after valid sibling data has been written; the coordinator uploads those bytes with `capture_status=partial` rather than discarding them. Other errors produce a non-upload outcome when no valid snapshot exists. The method must: +- scope provider lookup to both exact IDs and this worker/backend; +- include all available stdout and stderr emitted by the sandbox entrypoint/launcher and client process without content-based filtering; +- preserve stream/container/phase identity the backend actually supplies, use `combined`/`unknown` when it does not, and never infer a process identity from log text; +- stream rather than return log bytes; +- honor cancellation and the supplied bounded writer; +- make no task lifecycle changes; +- be safe to call while `ExecuteTask` is running; +- never remove a provider resource; +- emit deterministic source ordering where the provider has multiple streams/containers. + +Add a direct-only task-scoped `TaskLogCapture` handle to `TaskParams`. Docker and Kubernetes retain their exact resource registries until cleanup grace expires and read provider logs on request; they do not copy a second terminal spool. The command backend implements `SnapshotTaskLogs` by returning the typed unsupported error. + +Add an idempotent backend cleanup method keyed by exact task/execution identity. Normal terminal completion schedules it at the existing cleanup-grace deadline; worker shutdown continues to use the existing bounded shutdown/preservation contract and does not wait for archive collection. A snapshot request never changes the cleanup deadline. A snapshot racing the deadline is serialized by the per-execution mutex: whichever acquires it first completes, then the other observes the resulting resource/entry state. + +### NDJSON format +Every output line is valid UTF-8 JSON with `schema_version: 1` and `kind`. + +Data records contain: +- `kind: "data"` +- monotonically increasing `sequence` within the snapshot +- `backend: "docker" | "kubernetes" | "direct"` +- `phase` (`setup`, `agent`, `teardown`, or `container`) +- `stream` (`stdout`, `stderr`, `combined`, or `unknown`) +- provider timestamp when supplied, plus worker `observed_at` +- `encoding: "utf8" | "base64"` +- `data` +- optional Docker `container_id` +- optional Kubernetes `namespace`, `pod`, `container`, `container_type`, `restart_attempt`, and `previous` + +`phase` is provider truth, not a required process classifier. Direct execution can label launcher/client/setup/teardown phases from the handles it owns. Docker/Kubernetes provider streams commonly combine entrypoint and child-client output inside one container; those records retain their real container and stdout/stderr identity and do not fabricate which process wrote a line. + +The encoder passes only each data record's decoded `data` content through `ContentTransformer` before selecting UTF-8 versus base64 encoding. It does not transform schema version, kind, sequence, backend, phase, stream, timestamps, identity fields, warning codes, or truncation metadata. V1's no-op transformer is byte preserving. Provider records are not first uploaded raw and transformed in cloud storage; the first object-store upload already contains the transformed encoding. + +Data is framed into bounded chunks no larger than 32 KiB before JSON encoding. Valid UTF-8 is stored directly; other bytes are base64. Empty streams do not produce placeholder records. + +A readable sibling plus an unreadable backend stream emits a `kind: "source_error"` record containing only safe source identity and one stable warning code. It never embeds the provider error string. The same warning appears in acknowledgement `warning_codes`, allowing `warp-server` to mark the source partial without parsing arbitrary provider text. + +When bytes are dropped, output includes one valid record: + +```json +{"schema_version":1,"kind":"truncation","policy":"first_last","omitted_bytes_at_least":123} +``` + +The first and last portions each end on an NDJSON record boundary. Final sequence numbers describe emitted order, not original provider byte offsets. The acknowledgement's `truncated` value agrees with the record. + +### Direct-process capture and request snapshots +Add a secure disk-backed `TaskLogCapture` only for the direct backend, independent of task workspaces: +- Uses a random/digested task-local path below `${TMPDIR}/oz-agent-worker/debug-logs`, with root `0700`, file `0600`, and symlink rejection. +- Uses bounded first/last segments with fixed record chunks so memory and disk are capped at 64 MiB per execution. +- Supports an atomic snapshot watermark while writers continue. +- Tracks bytes, truncation, capture errors, and the cleanup-grace deadline. +- Deletes orphan files on process startup because V1 does not reconstruct ownership across process replacement. +- Deletes the capture at normal cleanup-grace expiry and during best-effort shutdown cleanup. + +Request-scoped transformed snapshots use the same secure root, are capped by the request, and are deleted after upload/retry completion. They exist only to make a retry byte-identical; they do not extend execution retention and are not a second terminal archive spool. + +Add top-level bounds with no archive-specific retention field: + +```yaml +debug_log_capture: + directory: "" + max_total_bytes: 1073741824 + max_execution_bytes: 67108864 + max_concurrent_uploads: 2 +``` + +Retention comes exclusively from the existing resolved idle-on-complete grace. Under global pressure, expired-grace captures and completed request snapshots are removed first; active direct captures shrink their retained first/last window or become unavailable without blocking task output. Invalid/non-positive bounds, a per-execution limit over 64 MiB, a total below the per-execution limit, or an unwritable configured root fails debug-capture initialization but remains non-fatal to assigned task execution. + +The Helm chart mounts a dedicated 1-GiB `emptyDir` for the default capture root and renders these bounds, with overrides for operators. The volume is explicitly ephemeral and does not change the cleanup-grace contract. + +### Docker adapter +`DockerBackend` adds a mutex-protected exact `(task_id, execution_id) -> container_id` registry: +- Register after container creation and before start. +- Keep it through active execution and the existing cleanup grace. +- Remove it only after idempotent container cleanup at the grace deadline. + +`SnapshotTaskLogs`: +- resolves the exact registered container; +- calls Docker `ContainerLogs` with stdout, stderr, and timestamps enabled and no time-range filter; +- preserves the complete container streams containing both the entrypoint/launcher and spawned client process without line filtering or invented process labels; +- demultiplexes Docker's stream framing so each record has correct stdout/stderr identity; +- emits chunked NDJSON directly to the supplied writer. + +`ExecuteTask` stops/waits for the container but does not force-remove it at terminal return. Cleanup at the grace deadline removes the container and registry entry. This replaces the current immediate-removal defer and lets an `ANY_FAILURE` request read the provider logs without a duplicate terminal spool. Provider log failure is recorded but never changes the task result or cleanup. + +### Kubernetes adapter +`SnapshotTaskLogs`: +- lists pods using the execution, task, and configured worker hash labels and verifies returned labels; +- sorts pods deterministically by name; +- visits init containers and regular containers in declared order; +- requests all available logs with timestamps and no time-range filter; +- preserves both stdout and stderr containing each container's entrypoint/launcher and client output without line filtering or invented process labels; +- attempts `Previous: true` before current logs for a container whose restart count is non-zero; +- tags every output record with pod/container identity and whether it is previous/current; +- treats a vanished pod, unavailable previous stream, or one unreadable container as a typed per-snapshot partial error while retaining readable sibling streams. + +The adapter must not use the existing 1-MiB `collectPodLogs` helper for archive collection. It streams through the shared aggregate 64-MiB bound. + +`ExecuteTask` does not delete a successful Job at terminal return. It retains the exact Job/pod identity until cleanup-grace expiry, then applies the existing success/failure deletion and TTL policy. Operators must configure `ttl_seconds_after_finished` no shorter than cleanup grace if failed-Job logs need to remain available for the full window. Snapshot failure cannot change Job outcome, terminal reporting, preservation-on-shutdown, or cleanup policy. + +The chart's existing `get pods/log` RBAC is sufficient; implementation must not broaden it beyond the namespace or add secret-read permissions. + +If the Kubernetes worker process is disrupted while a preserved task Job continues, the replacement process does not inherit the old process's ownership registry in V1. A later request is partial unless a future reattach protocol reconstructs exact ownership. This limitation is explicit and does not weaken the existing preserve-on-shutdown behavior. + +### Direct adapter +Create one task capture before setup. Replace worker-global output assignment with phase-aware multi-writers: +- Setup and entrypoint/launcher stdout/stderr goes to the worker console and `phase=setup` capture. +- The launched client process's stdout/stderr goes to the worker console and `phase=agent` capture. +- Teardown stdout/stderr goes to the worker console and `phase=teardown` capture. +- Capture mechanics: the archive branch of each multiwriter is a bounded non-blocking queue, not a synchronous disk write. `Write` copies a bounded chunk into the queue and immediately reports the original byte count; when the queue is full it drops archive bytes, marks `output_dropped`, and leaves the existing console sink/child process behavior unchanged. A per-task background encoder drains the queue into `TaskLogCapture`. Terminal finalization drains it under a bounded deadline before workspace cleanup, but the capture remains until cleanup grace expires. `SnapshotTaskLogs` takes a consistent watermark copy without closing an active capture. No task workspace file content is stored. + +### Upload client +The upload client accepts the provider-neutral target from the authenticated server: +- PUT sends the snapshot as the request body with supplied headers. +- POST builds a streaming multipart form with supplied fields and one file part; it never loads the file into memory. +- Redirects are disabled. +- The request body is capped at the lower of request `max_bytes`, configured per-execution maximum, and actual snapshot size. +- Each attempt reopens the immutable local snapshot, allowing replay. +- Network errors, HTTP 408/429, and 5xx responses retry with bounded exponential backoff only before `expires_at`. +- Other 4xx responses are terminal `upload_rejected`. +- Only a 2xx response is `uploaded`. +- Response bodies are drained/closed under a small cap but never logged or returned in acknowledgement text. + +CRC32C and SHA-256 are computed over the exact uploaded file bytes. Byte count and digests are included in the acknowledgement so `warp-server` can compare them with object-store attributes before signaling its Temporal workflow. + +After a successful upload, the request-scoped transformed snapshot is removed; completed request metadata remains so an identical duplicate request can replay the acknowledgement without re-uploading. The execution's provider resource or direct capture remains available until cleanup grace expires so a later request ID can obtain a newer snapshot. After grace expiry, no acknowledgement is possible from that former owner; the server times out or uses another owning connection if one exists. + +### Immutable snapshots and future incremental collection +V1 does not add append semantics to hosted storage and does not keep a long-lived upload stream: +- Each request creates one immutable bounded snapshot and uploads one complete object. +- The server's later archive generation replaces the logical log source. +- Active `ALWAYS` collection is supported by requesting snapshots after assignment and later lifecycle events. + +For future continuous collection, `TaskLogCapture` and the NDJSON sequence field permit sealed immutable chunks. A future protocol can request/upload chunk keys plus a sequence watermark and checksums. It must not require mutating an existing GCS/S3 object. No continuous chunk scheduler, server chunk manifest, or upload bandwidth policy is in V1. + +### Observability +Add bounded-cardinality metrics: +- log collection requests by outcome, backend, and active/cleanup-grace ownership; +- snapshot duration and bytes buckets; +- direct-capture current bytes and cleanup-grace entry count; +- cleanup-grace duration/expiry and provider-cleanup results; +- truncation and direct-capture pressure counts; +- upload duration, retry, and result counts; +- request queue depth/in-flight count. + +Structured logs may include request ID, run ID, execution ID, backend, outcome, bytes, truncation, and stable reason. They never include captured data, local capture filenames, target coordinates, signed URL/query, signed headers/form fields, checksum input, or response bodies. + +### Rollout and compatibility +1. Land server schema/status/manual collection without claiming self-hosted completeness. +2. Land this worker implementation and publish a pinned immutable worker image. +3. Validate old-worker timeout and new-worker v1 upload against staging. +4. Validate active and cleanup-grace post-failure capture for Docker, Kubernetes, and direct, including a deliberately too-short grace producing a partial archive. +5. Document that self-hosted operators may need to lengthen idle-on-complete and Kubernetes Job TTL for reliable `ANY_FAILURE` capture. +6. Enable creator/team-admin access and automatic policies only after the coordinated server/worker metrics show bounded duration, grace retention, disk, and upload behavior. + +`warp-server` remains tolerant of old worker versions indefinitely. Documentation must state the minimum worker version required for self-hosted logs in debug archives. + +## Design alternatives + +### Server pulls from customer infrastructure +The server cannot safely reach an arbitrary customer Docker daemon or Kubernetes API and has no direct-process handle. Worker-pushed bytes to a narrow presigned target preserve customer network boundaries and avoid sending object-store credentials to the worker. + +### Send logs through the WebSocket +WebSocket transfer would burden the regular server, interfere with control messages, duplicate flow control, and risk loading large logs into memory. Direct upload keeps the control channel metadata-only and lets object-store limits enforce a hard boundary. + +### First worker connection wins +Multiple worker processes can use one worker ID. Selecting the first connection could leak the wrong execution or return misleading emptiness. Server-side GCP Pub/Sub fan-out plus exact active/cleanup-grace ownership makes only the process that executed the assignment authoritative; all non-owners are silent. + +### Query current worker version during collection +Reading a current connection, image tag, or metrics label when the archive is assembled is simpler but becomes wrong after reconnects and rolling deploys. Reporting the build on every authenticated connection and snapshotting the exact selected connection at successful claim was selected. It adds no collection-time round trip, keeps old workers execution-compatible, and ties provenance to what actually received the execution rather than what happens to be online later. + +### New archive retention versus existing cleanup grace +An independent terminal spool/TTL would duplicate lifecycle configuration and create conflicting cleanup clocks. V1 instead keeps Docker/Kubernetes resources and exact ownership through the already-resolved idle-on-complete grace, while direct execution retains only its bounded output capture for that same interval. This can consume provider resources longer than today's immediate cleanup, so the operator chooses the existing grace and may need to lengthen it for reliable `ANY_FAILURE` collection. + +### Capture continuously for every task +Continuous remote upload would maximize recovery but adds bandwidth, credentials, chunk publication, and policy complexity even when no archive is requested. V1 captures direct-process output locally within fixed bounds, retains provider resources for cleanup grace, and uploads on request. Immutable chunk streaming remains a compatible follow-up. + +### Append to one object +GCS/S3 and the server's existing provider-neutral presigned-target contract do not offer one portable append primitive. Immutable snapshots avoid partial object visibility, retry ambiguity, and provider-specific behavior. + +### Keep only the last 64 MiB +Tail-only capture often loses setup/image/initialization failures. First/last retention preserves both early context and the terminal failure while keeping a strict bound, at the cost of an explicit gap. + +### Memory buffering +64 MiB per task multiplied by concurrency can exhaust the long-lived worker. Secure disk-backed head/tail capture for direct execution and request-scoped snapshots keep memory bounded and permit retryable upload bodies. + +### One generic stdout hook for every backend +Direct execution can be teed, but Docker and Kubernetes already own provider-native logs with stream/container metadata. A shared `SnapshotTaskLogs` contract plus backend-specific adapters preserves this context and handles resources that outlive the local process. + +## Risks and mitigations +- Log bytes contain secrets: keep storage private, never log payload/targets, use exact presigned destinations, and apply the versioned content-transformer hook to message data before the first upload; actual redaction rules remain deferred. +- Disk pressure affects direct execution: strict per-execution bounds, secure task capture, cleanup-grace expiry, pressure metrics, and task-non-fatal degradation. +- Provider resources live longer: reuse the operator-selected existing cleanup grace, expose duration/cleanup metrics, and document Kubernetes TTL alignment. +- Terminal request races cleanup: the worker moves ownership to cleanup grace before terminal reporting, and snapshot/cleanup share a per-execution mutex. +- Multi-instance worker collision: exact compound ownership, server GCP Pub/Sub fan-out, silent non-owners, server identity/assignment verification, and idempotent request IDs. +- WebSocket read starvation: provider/upload work runs asynchronously with a semaphore. +- Signed URL expiry: immutable local snapshot, replayable attempts, bounded retry before expiry, and explicit expired outcome. +- Target abuse: authenticated control channel, method/scheme/header validation, no redirects, and no target logging. +- Kubernetes worker replacement loses ownership: explicit V1 limitation; the archive remains partial and a future Job reattach protocol can recover it. +- Protocol drift across repositories: explicit companion links, checked-in golden v1 fixtures on both sides, and coordinated staging tests before rollout. +- Worker-version drift during rolling deploys: version is connection-scoped and copied exactly once from the selected successful claimant; current connection/image/metrics state is never used to backfill an execution. +- Direct capture changes subprocess behavior: multiwriters keep the existing console sink, capture writes are non-blocking/best-effort, and exit status remains authoritative. +- Backend snapshot errors alter task result: capture errors are isolated and never replace execution errors, status, or cleanup. + +## Validation criteria + +### Wire contract and parsing +- WVC-001: A golden protocol-v1 request fixture, including the 30-minute expiry and no-op content-transformer descriptor, is accepted by both this worker and the server implementation from `warpdotdev/warp-server#13839` without field translation. +- WVC-002: Golden complete and partial `uploaded` acknowledgement fixtures produced by the worker are accepted by the server and preserve every identity, byte, checksum, truncation, transformer version, capture status/warning, backend, and outcome field. +- WVC-003: Unknown optional request fields are ignored; a missing required field, invalid ID/format/method/control character, non-positive/oversized bound, expiry over 30 minutes, or expired request yields `failed` with a stable sanitized reason and no target request. An unsupported transformer yields `failed/unsupported_content_transformer`. +- WVC-004: Unsupported protocol versions produce `failed/unsupported_protocol_version`; an old-worker fixture that ignores the new message still yields the server-spec timeout/partial behavior. +- WVC-005: No acknowledgement or worker log contains a target URL, signed query/header/form field, local path, captured byte, or response body. + +### Ownership and idempotency +- WVC-006: Among two connected processes with one worker ID, only the process with the exact active or cleanup-grace `(run_id, execution_id)` uploads; the other silently ignores the request and emits no acknowledgement, log, metric with IDs, or request-cache entry. +- WVC-007: A process that owns another execution of the same run cannot upload for the requested execution. +- WVC-008: An exact cleanup-grace entry remains authoritative after terminal execution until the existing resolved idle-on-complete deadline. +- WVC-009: A duplicate `request_id` with identical content causes at most one snapshot/upload and replays the same acknowledgement. +- WVC-010: Reusing a request ID with different archive, collection, run, execution, target, expiry, bound, or content-transformer descriptor is rejected. +- WVC-011: Requests run off the WebSocket read loop; while one provider read/upload is blocked, heartbeat, cancellation, and another assignment message are still processed. +- WVC-012: The upload semaphore limits total concurrent handlers and the per-execution mutex prevents overlapping provider snapshots without blocking unrelated executions. + +### Terminal ordering and partial behavior +- WVC-013: For success, process failure, timeout/cancellation, and backend setup failure, transition to cleanup-grace ownership occurs before terminal lifecycle message enqueue. +- WVC-014: A simulated `ANY_FAILURE` request arriving immediately after `task_failed` retrieves the cleanup-grace resource/capture without racing registry deletion or provider cleanup. +- WVC-015: Direct-capture initialization/write/finalization, provider snapshot, transform, and upload failures never change task claim, execution exit result, terminal message, cleanup deadline, or worker reconnect behavior. +- WVC-016: No owning entry is silently ignored; an owning entry with no captured bytes yields a classified `unavailable` result rather than an empty uploaded object. +- WVC-017: A command-backend owner reports `unavailable/backend_not_supported`, while a different command-worker process silently ignores the request. + +### NDJSON and truncation +- WVC-018: Docker, Kubernetes, and direct fixtures emit only valid schema-v1 NDJSON for all available entrypoint/launcher and client stdout/stderr; every data record has bounded data, encoding, backend, truthful phase/stream, sequence, and observed-time fields, and safe `source_error` records contain no raw provider error. +- WVC-019: Invalid UTF-8/binary output round-trips through base64 records without corrupting the NDJSON stream. +- WVC-020: Output below the bound is byte-complete after record decoding and reports `truncated=false`. +- WVC-021: Output above the bound preserves valid first/last record sets, contains one truncation record with an omitted-byte lower bound, stays within the request/configured limit, and reports `truncated=true`. +- WVC-022: Empty provider streams do not create misleading zero-byte data records. +- WVC-023: CRC32C, SHA-256, and byte count match the exact complete NDJSON uploaded, including its truncation record. + +### Secure direct capture and cleanup grace +- WVC-024: Direct captures and request snapshots create their root as `0700`, files as `0600`, use non-user-derived names, and reject symlink/path-traversal fixtures. +- WVC-025: Each direct execution and request snapshot cannot retain more than the requested/protocol 64-MiB ceiling, aggregate captures/snapshots cannot exceed the configured total budget, and Docker/Kubernetes create no duplicate terminal log file. +- WVC-026: Direct-capture pressure degrades only archive capture with a truthful truncated/unavailable state; active task output, console delivery, and process exit remain unaffected. +- WVC-027: A consistent active direct snapshot uses a fixed watermark while later writes continue into `TaskLogCapture`. +- WVC-028: Cleanup-grace entries, direct captures, and provider resources remain until the existing resolved idle-on-complete deadline and are removed idempotently at expiry; successful upload does not extend or shorten the execution cleanup deadline. +- WVC-029: Startup removes unowned orphan direct-capture/request-snapshot files without reconstructing or exposing their data. +- WVC-030: Task, worker, and default idle-on-complete values resolve with the existing precedence; invalid capture bounds/concurrency or an unwritable capture root records a non-sensitive archive-capture failure without failing assigned task execution. + +### Docker +- WVC-031: Exact task/execution ownership resolves exactly one container; a mismatched execution cannot read another container. +- WVC-032: Active Docker snapshot calls the container-log API with stdout, stderr, timestamps, and no time-range filter, retains entrypoint/launcher plus client output without line filtering, and demultiplexes stream identity into NDJSON. +- WVC-033: Success, non-zero exit, OOM, context cancellation, and wait failure each retain the exact stopped container and ownership through cleanup grace, then remove both idempotently at the deadline. +- WVC-034: A Docker request during cleanup grace reads the retained container; a request after container/grace cleanup is silently ignored by the former owner and becomes partial by server timeout. +- WVC-035: Archive capture replaces the existing unbounded `io.ReadAll` diagnostic path and passes a large-log test without memory scaling with output size. + +### Kubernetes +- WVC-036: Pod selection requires exact execution, task, and worker hash labels; pods with only a colliding/mismatched label set are excluded. +- WVC-037: Snapshot ordering is pod name, then declared init containers, then declared regular containers, with identity on every record. +- WVC-038: Every readable current stdout/stderr stream, including entrypoint/launcher and client output, is captured without line filtering with provider timestamps and no time filter; a restarted container also attempts and labels previous logs. +- WVC-039: One missing/unreadable container does not discard readable siblings, emits a safe source-error record, and uploads with `capture_status=partial` plus the matching bounded warning code. +- WVC-040: Success, failure, and cancellation retain exact Job/pod identity through cleanup grace; successful Job deletion and normal failed-Job cleanup happen at or after that deadline. +- WVC-041: Snapshot failure does not change Job cleanup deadline, configured failed-Job TTL, task result, or Kubernetes preserve-on-worker-shutdown semantics; a TTL shorter than cleanup grace is tested/documented as potentially causing a partial archive. +- WVC-042: Helm/RBAC rendering still grants only namespace-scoped Job/pod/event operations plus `get pods/log`; no secret-read or cluster-scoped permission is added. +- WVC-043: Worker-pod replacement with a preserved Job is documented/tested as loss of process-local cleanup-grace ownership in V1, not a false successful upload by the replacement. + +### Direct +- WVC-044: Entrypoint/launcher, launched client, setup, and teardown stdout/stderr continue to reach the worker console and are separately labeled from handles the direct backend actually owns. +- WVC-045: One task's direct output cannot appear in another concurrent task's snapshot. +- WVC-046: Active direct snapshot returns bytes only through its watermark while the process continues and later output remains available to a later snapshot. +- WVC-047: Direct terminal capture remains available after per-task workspace cleanup and stores no workspace file content. +- WVC-048: A full/failed archive queue drops bytes and records `output_dropped` while its `Write` returns the child byte count promptly; a capture write failure or slow archive request cannot block the subprocess pipe or alter its exit code. + +### Upload behavior +- WVC-049: A PUT fixture sends exactly the bounded snapshot with supplied safe headers, no redirect following, and matching digests. +- WVC-050: A POST fixture streams supplied multipart fields plus one file part without loading the snapshot into memory. +- WVC-051: Network error, 408, 429, and 5xx fixtures retry with bounded backoff before expiry; non-retryable 4xx fails immediately; no retry starts after expiry. +- WVC-052: Only 2xx produces `uploaded`; server-side object-attribute verification using the reported byte/checksum/transformer values passes for the candidate object, and the server marks partial capture/truncation truthfully from acknowledgement metadata. +- WVC-053: HTTP redirects are rejected without forwarding signed headers or body to the redirect destination. +- WVC-054: Upload memory remains bounded for a 64-MiB fixture and each retry reopens the same immutable local snapshot. +- WVC-055: After a successful upload and request-snapshot removal, an identical duplicate request replays the cached acknowledgement without another network request; execution resource/capture cleanup still follows only the existing grace deadline. + +### Active, cleanup-grace, and version-skew scenarios +- WVC-056: An active Docker, Kubernetes, and direct execution each satisfies an `ALWAYS`/manual request while continuing to run. +- WVC-057: A later request ID during cleanup grace uploads a newer terminal snapshot that includes output after the active watermark. +- WVC-058: A request before the backend resource exists reports `resource_not_ready` and does not interfere with later snapshots. +- WVC-059: A disconnected, shutting-down, expired-cleanup-grace, unsupported backend, and old-worker scenario each maps to the server's partial/stale behavior and never blocks archive publication. +- WVC-060: No V1 path appends to an existing object; every request uploads one complete immutable object. + +### Configuration, operations, and verification +- WVC-061: Tests prove cleanup grace resolves from task `idle_timeout_minutes`, then worker `idle_on_complete`, then the Oz default, and archive requests never extend that deadline; capture root/total/execution/concurrency bounds use validated defaults/overrides with strict unknown-field parsing and no separate retention duration. +- WVC-062: Helm lint/template proves the bounded ephemeral capture volume/config plus existing idle-on-complete and Kubernetes `ttl_seconds_after_finished` overrides render without a new archive-retention duration; a too-short TTL/grace fixture is documented as partial-prone. +- WVC-063: README/operator docs describe the sensitive-data implication, 30-minute request timeout, cleanup-grace sizing for `ANY_FAILURE`, Kubernetes TTL alignment, ephemeral worker-replacement limitation, supported backends, and minimum server/worker compatibility. +- WVC-064: Metrics cover request/result, active/cleanup-grace ownership, snapshot, bytes, truncation, direct-capture pressure/usage, provider cleanup, upload/retry, and queue state without run/execution/request IDs as metric labels. +- WVC-065: Structured-log tests or capture assertions prove no log bytes, target material, upload response body, or local capture path is emitted. +- WVC-066: Unit tests use fake/injected Docker, Kubernetes-log, filesystem, clock, WebSocket, and HTTP dependencies; no test requires customer infrastructure. +- WVC-067: Coordinated staging validation uses GCP Pub/Sub fan-out, server-generated GCS and S3-style targets, verifies server `GetAttrs`, and covers one owning plus one silent non-owning instance and old-worker timeout. +- WVC-068: `gofmt -s`, `go vet ./...`, `golangci-lint run`, `go test ./...`, `go build -v ./...`, `helm lint`, and `helm template` pass before the implementation PR is promoted. +- WVC-069: This change has no rendered UI; computer-use visual verification is not applicable. +- WVC-070: A no-op transformer preserves NDJSON message data exactly; a fake transformer changes only decoded `data` while timestamps, sequence, stream/container identity, warnings, and source metadata remain unchanged. The transformed NDJSON is the first cloud object uploaded, and an unsupported transformer uploads nothing. + +### Worker build provenance and local-log completeness +- WVC-071: Release, local `dev`, and arbitrary test builds pass their exact `main.Version` through `worker.Config` and send it as `X-Warp-Worker-Version` on every initial WebSocket dial and reconnect before the connection can receive an assignment. +- WVC-072: With two connected processes sharing one worker ID but reporting different versions, the server persists only the successfully selected claimant's version on the execution before sandbox launch; rejection, reconnect, later rollout, and archive collection do not overwrite it. +- WVC-073: Missing, overlong, or control-character-bearing version metadata does not fail connection or task execution, is never logged raw, and yields the coordinated server spec's explicit `not_reported` partial-provenance behavior rather than a value derived from current image, presence, or metrics state. +- WVC-074: Docker, Kubernetes, and direct integration fixtures emit distinguishable sentinels on both stdout and stderr from the entrypoint/launcher and launched client; every sentinel appears exactly once after NDJSON decoding unless an explicit bound/truncation or provider-partial record accounts for it. +- WVC-075: Docker/Kubernetes fixtures whose provider cannot distinguish entrypoint from child-client output retain real stdout/stderr/container identity and use `combined`/`unknown` where appropriate; no implementation or test infers process identity from log text. + +## Cross-repository completion +Implementation is complete only when this spec's WVC-001/WVC-002/WVC-052/WVC-067/WVC-070 through WVC-075 interoperate with the server-side VC-041, VC-076 through VC-082, and VC-087 through VC-094 in `warpdotdev/warp-server#13839`. A worker-only test double or server-only fake is not sufficient evidence for the final protocol gate. diff --git a/README.md b/README.md index 331a918..4d9e550 100644 --- a/README.md +++ b/README.md @@ -423,6 +423,25 @@ shows up as a distinct series. attempts; spikes indicate flapping workers. - `oz_worker_info{version,backend,worker_id}` (gauge, value `1`): build and runtime metadata, useful for joining other series by labels. +- `oz_worker_debug_archive_requests_total{backend,ownership,outcome,reason}` + (counter): debug-archive log requests this worker owned. Requests for + executions another instance ran are silent and are not counted here. +- `oz_worker_debug_archive_snapshot_duration_seconds{backend}` and + `oz_worker_debug_archive_snapshot_bytes{backend}` (histograms): cost and + size of producing a log snapshot. +- `oz_worker_debug_archive_truncations_total{backend}` (counter): snapshots + that dropped bytes to stay within their bound. +- `oz_worker_debug_archive_uploads_total{result}` and + `oz_worker_debug_archive_upload_duration_seconds{result}`: snapshot upload + outcomes and latency. +- `oz_worker_debug_archive_requests_in_flight` (gauge): requests currently + being snapshotted or uploaded, bounded by `debugLogCapture.maxConcurrentUploads`. +- `oz_worker_debug_archive_capture_bytes` (gauge): disk currently reserved by + direct-execution captures and request snapshots. +- `oz_worker_cleanup_grace_entries` (gauge) and + `oz_worker_cleanup_grace_results_total{backend,result}` (counter): + executions retained past terminal state and the backend cleanups performed + when their grace expired. ### Sample dashboards / alerts @@ -441,6 +460,109 @@ Direct mappings for the questions enterprise operators most commonly ask: - **Reconnect storms:** `sum(rate(oz_worker_websocket_reconnects_total[5m])) > 0.1` +## Debug archive log collection + +Warp can assemble a debug archive for a cloud-agent run. For self-hosted +executions the logs live inside your infrastructure, so Warp asks the worker +that actually ran the execution for a bounded snapshot instead of reaching into +your Docker daemon, Kubernetes cluster, or host. + +When Warp requests logs for an execution, the worker that ran it snapshots the +backed-up output, uploads it directly to a short-lived Warp-signed destination, +and reports the result. Every other worker process silently ignores the request. +Collection never blocks, delays, or fails a running agent: if logs cannot be +captured or uploaded, the archive is simply marked partial. + +### Sensitive data + +A snapshot contains whatever the execution wrote to stdout and stderr, which +can include prompts, source code, identifiers, and secrets a process printed. +Treat a debug archive as sensitive. The worker itself never logs captured +bytes, upload destinations, signed headers, local capture paths, or upload +response bodies. + +### Supported backends + +| Backend | Source | +| --- | --- | +| Docker | The execution container's stdout and stderr, covering the entrypoint script and the client process it starts. | +| Kubernetes | Every init and regular container in the execution's pods, including current and best-effort previous logs after a restart. | +| Direct | A bounded on-disk capture of setup, agent, and teardown stdout/stderr. | +| Command | Not supported. The dispatch command hands the task to an opaque runtime with no log API, so the worker reports the source unavailable rather than passing off dispatch output as the agent's log. | + +Docker and Kubernetes report one merged stream per container, so those records +are labeled `combined` rather than attributing a line to the entrypoint or the +client. Only direct execution owns distinct handles, so only it labels +`setup`, `agent`, and `teardown` phases. + +### Sizing the cleanup grace for failure capture + +The worker retains an execution's log source for the cleanup grace it already +resolves for `--idle-on-complete`, in this order: + +1. the run's `idle_timeout_minutes` +2. the worker's `idle_on_complete` (`worker.idleOnComplete` in the chart) +3. the Oz default of 45 minutes + +There is deliberately no separate archive retention setting: one clock governs +both how long the agent stays available for follow-ups and how long its logs +stay retrievable. + +Collection triggered by a failure has to reach the worker after the terminal +event propagates to Warp, and the request itself allows up to 30 minutes. If +you want reliable archives for failed runs, keep the grace comfortably longer +than that; a shorter grace still works but produces a partial archive when it +lapses first. + +**Kubernetes:** `kubernetesBackend.ttlSecondsAfterFinished` must not be shorter +than the effective grace. The Job TTL controller deletes a finished Job's pods, +and once they are gone their logs are gone with them regardless of the worker's +own retention. A successful Job is now deleted at the grace deadline rather +than immediately at task completion. + +**Worker replacement:** ownership is process-local. If a worker pod is replaced +while a preserved Job keeps running, the replacement does not inherit the old +process's ownership, and a later request for that execution yields a partial +archive rather than an incorrect upload. + +### Capture bounds + +Each execution's snapshot is capped at 64 MiB. Output above the cap keeps the +first and last portions with an explicit gap marker, so both the early setup +context and the terminal failure survive. + +The Helm chart mounts a dedicated 1 GiB ephemeral volume for the capture root +and renders the matching bounds: + +```yaml +debugLogCapture: + enabled: true + sizeLimit: 1Gi + directory: /var/lib/oz/debug-logs + maxTotalBytes: 1073741824 + maxExecutionBytes: 67108864 + maxConcurrentUploads: 2 +``` + +Outside Kubernetes, set the same bounds under `debug_log_capture` in the config +file; an unset value uses the default above, and an unwritable root or an +invalid bound disables archive capture without affecting task execution. + +The volume is scratch space, not storage: nothing in it survives worker +replacement, and it does not extend the cleanup grace. + +### Compatibility + +Self-hosted logs in debug archives require a worker built with this protocol. +An older worker ignores the request and Warp records the source as unavailable, +so upgrading is safe and never required for ordinary task execution. Each +authenticated connection reports its build version so Warp can show exactly +which worker ran an execution; a worker that reports no version still executes +tasks normally and is shown as not reported. + +The Kubernetes backend needs no additional permissions: the chart's existing +namespace-scoped `get pods/log` grant is sufficient. + ## License Copyright © 2026 Warp diff --git a/charts/oz-agent-worker/templates/configmap.yaml b/charts/oz-agent-worker/templates/configmap.yaml index 13fa3ea..4d86953 100644 --- a/charts/oz-agent-worker/templates/configmap.yaml +++ b/charts/oz-agent-worker/templates/configmap.yaml @@ -13,6 +13,14 @@ data: {{- if .Values.worker.idleOnComplete }} idle_on_complete: {{ .Values.worker.idleOnComplete | quote }} {{- end }} + debug_log_capture: + {{- if .Values.debugLogCapture.enabled }} + directory: {{ .Values.debugLogCapture.directory | quote }} + {{- end }} + {{- /* int64 keeps large byte counts out of YAML scientific notation. */}} + max_total_bytes: {{ .Values.debugLogCapture.maxTotalBytes | int64 }} + max_execution_bytes: {{ .Values.debugLogCapture.maxExecutionBytes | int64 }} + max_concurrent_uploads: {{ .Values.debugLogCapture.maxConcurrentUploads | int }} backend: kubernetes: namespace: {{ default .Release.Namespace .Values.kubernetesBackend.namespace | quote }} diff --git a/charts/oz-agent-worker/templates/deployment.yaml b/charts/oz-agent-worker/templates/deployment.yaml index aeb7e8b..e4d0d9e 100644 --- a/charts/oz-agent-worker/templates/deployment.yaml +++ b/charts/oz-agent-worker/templates/deployment.yaml @@ -98,6 +98,10 @@ spec: - name: config mountPath: /etc/oz-agent-worker readOnly: true + {{- if .Values.debugLogCapture.enabled }} + - name: debug-log-capture + mountPath: {{ .Values.debugLogCapture.directory | quote }} + {{- end }} {{- with .Values.worker.livenessProbe }} livenessProbe: {{- toYaml . | nindent 12 }} @@ -114,6 +118,14 @@ spec: - name: config configMap: name: {{ include "oz-agent-worker.fullname" . }}-config + {{- if .Values.debugLogCapture.enabled }} + # Scratch space for debug-archive log capture. It is explicitly + # ephemeral: nothing here survives worker pod replacement, and it does + # not change the execution cleanup-grace contract. + - name: debug-log-capture + emptyDir: + sizeLimit: {{ .Values.debugLogCapture.sizeLimit | quote }} + {{- end }} {{- with .Values.worker.nodeSelector }} nodeSelector: {{- toYaml . | nindent 8 }} diff --git a/charts/oz-agent-worker/values.yaml b/charts/oz-agent-worker/values.yaml index 25c47fe..c002210 100644 --- a/charts/oz-agent-worker/values.yaml +++ b/charts/oz-agent-worker/values.yaml @@ -91,9 +91,15 @@ kubernetesBackend: # How long finished task Jobs that the worker leaves in place are retained before # the Kubernetes Job TTL controller deletes them (and their pods). This covers # failed task Jobs, which are kept for post-mortem debugging, and Jobs orphaned by - # worker disruption. Successful task Jobs are deleted immediately by the worker, so - # this does not keep them around. Defaults to 24h (86400s). When worker.cleanup=false - # the worker sets no TTL, so these Jobs remain indefinitely. + # worker disruption. + # + # A successful task Job is deleted by the worker at its execution's cleanup-grace + # deadline (task idle_timeout_minutes, then worker.idleOnComplete, then the Oz + # default) rather than at task completion, so its logs remain readable for a debug + # archive. Keep this TTL at least as long as the effective grace: if the controller + # deletes the pods first, a debug archive collected for that run is partial. + # Defaults to 24h (86400s). When worker.cleanup=false the worker sets no TTL, so + # these Jobs remain indefinitely. ttlSecondsAfterFinished: 86400 workspaceSizeLimit: "" unschedulableTimeout: "30s" @@ -116,6 +122,30 @@ kubernetesBackend: # memory: 64Mi preflightResources: {} +# Bounds for debug-archive log capture. The worker uses this space for +# direct-execution output captures and for the request-scoped snapshot it +# uploads; Docker and Kubernetes executions read provider logs on demand and +# keep nothing here. +# +# There is deliberately no retention duration: how long an execution's logs stay +# retrievable is the existing idle-on-complete cleanup grace (task +# `idle_timeout_minutes`, then `worker.idleOnComplete`, then the Oz default). +debugLogCapture: + # Mounts a dedicated ephemeral volume for the capture root. Disable only if + # the container filesystem already provides suitable scratch space. + enabled: true + # Size of the ephemeral volume. Keep it at or above maxTotalBytes so the + # worker's own accounting, not the volume, is the binding limit. + sizeLimit: 1Gi + # Mount path for the capture volume, used as the worker's capture root. + directory: /var/lib/oz/debug-logs + # Process-local disk budget shared by captures and request snapshots. + maxTotalBytes: 1073741824 + # Per-execution retention bound. May not exceed the protocol's 64 MiB ceiling. + maxExecutionBytes: 67108864 + # Concurrent snapshot/upload handlers per worker process. + maxConcurrentUploads: 2 + # Metrics export uses the OpenTelemetry autoexport package, so the operator # selects the exporter via standard OpenTelemetry environment variables. When # `metrics.enabled` is false the worker emits no metrics, matching the diff --git a/internal/common/cleanup_grace_test.go b/internal/common/cleanup_grace_test.go new file mode 100644 index 0000000..bb72758 --- /dev/null +++ b/internal/common/cleanup_grace_test.go @@ -0,0 +1,100 @@ +package common + +import ( + "testing" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/types" +) + +func taskWithIdleTimeout(minutes int) *types.Task { + return &types.Task{ + ID: "task-1", + AgentConfigSnapshot: &types.AmbientAgentConfig{IdleTimeoutMinutes: &minutes}, + } +} + +func TestResolveCleanupGracePrecedence(t *testing.T) { + tests := []struct { + name string + task *types.Task + idleOnComplete string + want time.Duration + }{ + { + name: "task idle_timeout_minutes wins", + task: taskWithIdleTimeout(90), + idleOnComplete: "10m", + want: 90 * time.Minute, + }, + { + name: "worker idle_on_complete is next", + task: &types.Task{ID: "task-1"}, + idleOnComplete: "10m", + want: 10 * time.Minute, + }, + { + name: "oz default is the fallback", + task: &types.Task{ID: "task-1"}, + idleOnComplete: "", + want: DefaultIdleOnComplete, + }, + { + name: "a zero worker override disables retention", + task: &types.Task{ID: "task-1"}, + idleOnComplete: "0s", + want: 0, + }, + { + name: "a non-positive task timeout falls through to the worker override", + task: taskWithIdleTimeout(0), + idleOnComplete: "5m", + want: 5 * time.Minute, + }, + { + name: "an unparseable worker override falls back to the default", + task: &types.Task{ID: "task-1"}, + idleOnComplete: "not-a-duration", + want: DefaultIdleOnComplete, + }, + { + name: "a nil task uses the worker override", + task: nil, + idleOnComplete: "15m", + want: 15 * time.Minute, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := ResolveCleanupGrace(tc.task, tc.idleOnComplete); got != tc.want { + t.Fatalf("grace = %v, want %v", got, tc.want) + } + }) + } +} + +func TestResolveCleanupGraceMatchesTheEmittedIdleOnCompleteFlag(t *testing.T) { + // The retention window and the flag the agent receives must stay in step: + // an operator tuning one is tuning both. + task := taskWithIdleTimeout(30) + args := AugmentArgsForTask(task, nil, TaskAugmentOptions{IdleOnComplete: "10m"}) + + var emitted string + for i, arg := range args { + if arg == "--idle-on-complete" && i+1 < len(args) { + emitted = args[i+1] + } + } + if emitted != "30m" { + t.Fatalf("emitted --idle-on-complete %q, want 30m", emitted) + } + + emittedDuration, err := time.ParseDuration(emitted) + if err != nil { + t.Fatalf("the emitted flag %q is not a duration: %v", emitted, err) + } + if grace := ResolveCleanupGrace(task, "10m"); grace != emittedDuration { + t.Fatalf("cleanup grace = %s, want it to match the emitted flag %s", grace, emitted) + } +} diff --git a/internal/common/task_utils.go b/internal/common/task_utils.go index fe6b5df..bfd2f24 100644 --- a/internal/common/task_utils.go +++ b/internal/common/task_utils.go @@ -4,10 +4,15 @@ import ( "encoding/json" "fmt" "strings" + "time" "github.com/warpdotdev/oz-agent-worker/internal/types" ) +// DefaultIdleOnComplete mirrors the oz CLI's default --idle-on-complete value, +// used when neither the task nor the worker sets one. +const DefaultIdleOnComplete = 45 * time.Minute + // TaskAugmentOptions contains settings translated into oz CLI flags for every task. // Add new CLI overrides here rather than as extra parameters. type TaskAugmentOptions struct { @@ -148,6 +153,28 @@ func shareAccessLevelForEmission(access types.AccessLevel) string { } } +// ResolveCleanupGrace returns how long the agent stays alive after its +// conversation finishes, using the same precedence as the emitted +// --idle-on-complete flag: task idle_timeout_minutes, then the worker's +// idle_on_complete, then the oz CLI default. +// +// The worker reuses this window as its post-terminal ownership and +// log-retrieval deadline, so a debug-archive request has one cleanup clock to +// reason about instead of a second archive-specific retention setting. An +// unparseable worker override falls back to the default rather than dropping +// retention to zero. +func ResolveCleanupGrace(task *types.Task, workerIdleOnComplete string) time.Duration { + value, ok := resolveIdleOnComplete(task, TaskAugmentOptions{IdleOnComplete: workerIdleOnComplete}) + if !ok { + return DefaultIdleOnComplete + } + grace, err := time.ParseDuration(value) + if err != nil || grace < 0 { + return DefaultIdleOnComplete + } + return grace +} + func resolveIdleOnComplete(task *types.Task, opts TaskAugmentOptions) (string, bool) { if task != nil && task.AgentConfigSnapshot != nil && diff --git a/internal/config/config.go b/internal/config/config.go index 561f5b8..a7fdbf5 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -26,6 +26,26 @@ type FileConfig struct { // overrides are no longer needed. IdleOnComplete *string `yaml:"idle_on_complete"` Backend BackendConfig `yaml:"backend"` + // DebugLogCapture bounds the disk and concurrency used to collect debug + // archive logs. It carries no retention duration: how long an execution's + // logs stay retrievable is the existing idle_on_complete cleanup grace. + DebugLogCapture *DebugLogCaptureConfig `yaml:"debug_log_capture"` +} + +// DebugLogCaptureConfig bounds debug-archive log capture. Values left unset +// use the worker's built-in defaults. +type DebugLogCaptureConfig struct { + // Directory overrides the root beneath which captures and request + // snapshots are written. Empty uses the process temporary directory. + Directory string `yaml:"directory"` + // MaxTotalBytes is the process-local disk budget shared by direct-execution + // captures and request snapshots. + MaxTotalBytes *int64 `yaml:"max_total_bytes"` + // MaxExecutionBytes bounds one execution's retained output. It may not + // exceed the protocol's 64 MiB ceiling. + MaxExecutionBytes *int64 `yaml:"max_execution_bytes"` + // MaxConcurrentUploads bounds how many snapshot/upload handlers run at once. + MaxConcurrentUploads *int `yaml:"max_concurrent_uploads"` } // BackendConfig contains the backend selection. diff --git a/internal/config/debug_log_capture_test.go b/internal/config/debug_log_capture_test.go new file mode 100644 index 0000000..12d44e9 --- /dev/null +++ b/internal/config/debug_log_capture_test.go @@ -0,0 +1,91 @@ +package config + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func writeConfig(t *testing.T, contents string) string { + t.Helper() + path := filepath.Join(t.TempDir(), "config.yaml") + if err := os.WriteFile(path, []byte(contents), 0o600); err != nil { + t.Fatalf("failed to write config: %v", err) + } + return path +} + +// The Helm chart renders exactly this shape, so parsing it here keeps the chart +// and the worker's schema from drifting apart. +const chartRenderedConfig = ` +worker_id: "ci-worker" +cleanup: true +max_concurrent_tasks: 0 +idle_on_complete: "90m" +debug_log_capture: + directory: "/var/lib/oz/debug-logs" + max_total_bytes: 1073741824 + max_execution_bytes: 67108864 + max_concurrent_uploads: 2 +backend: + kubernetes: + namespace: "agents" + ttl_seconds_after_finished: 7200 +` + +func TestLoadParsesTheChartRenderedDebugLogCaptureBlock(t *testing.T) { + cfg, err := Load(writeConfig(t, chartRenderedConfig)) + if err != nil { + t.Fatalf("Load: %v", err) + } + + capture := cfg.DebugLogCapture + if capture == nil { + t.Fatal("debug_log_capture was not parsed") + } + if capture.Directory != "/var/lib/oz/debug-logs" { + t.Errorf("directory = %q, want the chart's capture root", capture.Directory) + } + if capture.MaxTotalBytes == nil || *capture.MaxTotalBytes != 1073741824 { + t.Errorf("max_total_bytes = %v, want 1073741824", capture.MaxTotalBytes) + } + if capture.MaxExecutionBytes == nil || *capture.MaxExecutionBytes != 67108864 { + t.Errorf("max_execution_bytes = %v, want 67108864", capture.MaxExecutionBytes) + } + if capture.MaxConcurrentUploads == nil || *capture.MaxConcurrentUploads != 2 { + t.Errorf("max_concurrent_uploads = %v, want 2", capture.MaxConcurrentUploads) + } + + // Retention comes from the existing cleanup grace, so the block must carry + // no duration of its own. + if cfg.IdleOnComplete == nil || *cfg.IdleOnComplete != "90m" { + t.Errorf("idle_on_complete = %v, want 90m", cfg.IdleOnComplete) + } +} + +func TestLoadOmittingDebugLogCaptureLeavesItUnset(t *testing.T) { + cfg, err := Load(writeConfig(t, "worker_id: \"w-1\"\n")) + if err != nil { + t.Fatalf("Load: %v", err) + } + if cfg.DebugLogCapture != nil { + t.Fatalf("debug_log_capture = %+v, want it unset so defaults apply", cfg.DebugLogCapture) + } +} + +func TestLoadRejectsAnArchiveRetentionSetting(t *testing.T) { + // Retention is deliberately not configurable here: strict field parsing is + // what stops an operator from silently setting a second cleanup clock. + _, err := Load(writeConfig(t, ` +worker_id: "w-1" +debug_log_capture: + retention: "2h" +`)) + if err == nil { + t.Fatal("expected an unknown debug_log_capture field to be rejected") + } + if !strings.Contains(err.Error(), "retention") { + t.Fatalf("error = %v, want it to name the unknown field", err) + } +} diff --git a/internal/debuglog/capture.go b/internal/debuglog/capture.go new file mode 100644 index 0000000..4d2a714 --- /dev/null +++ b/internal/debuglog/capture.go @@ -0,0 +1,240 @@ +package debuglog + +import ( + "encoding/base64" + "encoding/json" + "errors" + "io" + "os" + "sync" + "sync/atomic" + "time" +) + +// captureQueueDepth bounds how many chunks may await the background encoder. +// A full queue drops archive bytes rather than back-pressuring the subprocess +// pipe that feeds it. +const captureQueueDepth = 256 + +// captureRecord is the internal on-disk framing for captured output. It stores +// what the direct backend observed; the request's content transformer is +// applied later, when a snapshot re-encodes these records into NDJSON. +type captureRecord struct { + Phase string `json:"p"` + Stream string `json:"s"` + ObservedAt string `json:"t"` + Data string `json:"d"` +} + +// TaskLogCapture is a bounded, secure, disk-backed copy of one direct +// execution's stdout and stderr. Writers never block on it: a full queue drops +// archive bytes and marks the capture partial while the task's console output +// and process exit stay unaffected. +type TaskLogCapture struct { + store *Store + spool *boundedSpool + reserved int64 + now func() time.Time + + queue chan captureRecord + done chan struct{} + + dropped atomic.Bool + failed atomic.Bool + bytes atomic.Int64 + // pending counts records accepted but not yet written to disk. An empty + // queue is not enough for finalization: the encoder may have taken the + // last record and not written it yet. + pending atomic.Int64 + + closeOnce sync.Once +} + +// NewTaskLogCapture allocates a bounded capture against the store's budget. +func (s *Store) NewTaskLogCapture(now func() time.Time) (*TaskLogCapture, error) { + limit := s.config.MaxExecutionBytes + if err := s.reserve(limit); err != nil { + return nil, err + } + + spool, err := newBoundedSpool(limit, func(suffix string) (*os.File, error) { + return s.createFile("capture", suffix) + }) + if err != nil { + s.release(limit) + return nil, err + } + + if now == nil { + now = time.Now + } + capture := &TaskLogCapture{ + store: s, + spool: spool, + reserved: limit, + now: now, + queue: make(chan captureRecord, captureQueueDepth), + done: make(chan struct{}), + } + go capture.drain() + return capture, nil +} + +// Writer returns an io.Writer that copies output into the capture under the +// given phase and stream. Write always reports the full input length and never +// returns an error, so a capture problem can never alter the child process's +// view of its own pipe. +func (c *TaskLogCapture) Writer(phase Phase, stream Stream) io.Writer { + return &captureWriter{capture: c, phase: phase, stream: stream} +} + +type captureWriter struct { + capture *TaskLogCapture + phase Phase + stream Stream +} + +func (w *captureWriter) Write(p []byte) (int, error) { + w.capture.offer(w.phase, w.stream, p) + return len(p), nil +} + +// offer enqueues a bounded copy of p, dropping it when the queue is full. +func (c *TaskLogCapture) offer(phase Phase, stream Stream, p []byte) { + if len(p) == 0 { + return + } + for _, part := range splitChunk(p, MaxChunkBytes) { + // The caller owns p's backing array and may reuse it before the + // encoder drains the queue, so the chunk is copied here. + record := captureRecord{ + Phase: string(phase), + Stream: string(stream), + ObservedAt: c.now().UTC().Format(time.RFC3339Nano), + Data: base64.StdEncoding.EncodeToString(part), + } + c.pending.Add(1) + select { + case c.queue <- record: + default: + c.pending.Add(-1) + c.dropped.Store(true) + } + } +} + +func (c *TaskLogCapture) drain() { + defer close(c.done) + for record := range c.queue { + c.writeRecord(record) + c.pending.Add(-1) + } +} + +func (c *TaskLogCapture) writeRecord(record captureRecord) { + line, err := json.Marshal(record) + if err != nil { + c.failed.Store(true) + return + } + line = append(line, '\n') + if err := c.spool.WriteLine(line); err != nil { + if !errors.Is(err, ErrSpoolClosed) { + c.failed.Store(true) + } + return + } + c.bytes.Add(int64(len(line))) +} + +// Finalize drains queued output under a bounded deadline so a terminal +// snapshot sees the execution's last bytes. The capture stays readable +// afterwards; only Close releases it. +func (c *TaskLogCapture) Finalize(deadline time.Duration) { + timer := time.NewTimer(deadline) + defer timer.Stop() + + for { + if c.pending.Load() == 0 { + return + } + select { + case <-timer.C: + // Output still in flight past the deadline is reported as dropped + // so the snapshot is truthfully marked partial. + c.dropped.Store(true) + return + case <-time.After(time.Millisecond): + } + } +} + +// Bytes reports how many capture bytes have been written to disk. +func (c *TaskLogCapture) Bytes() int64 { return c.bytes.Load() } + +// SnapshotTo replays the capture's contents at a fixed watermark into sink +// while writers keep appending. Output after the watermark stays available to +// a later snapshot. +func (c *TaskLogCapture) SnapshotTo(sink Sink) error { + if c.dropped.Load() || c.failed.Load() { + // The exact dropped byte count is unknown because the queue discards + // chunks without accounting them; report a nonzero lower bound so the + // snapshot is truthfully marked truncated. + sink.NoteOmittedBytes(1) + } + + mark := c.spool.Watermark() + if mark.omitted > 0 { + sink.NoteOmittedBytes(mark.omitted) + } + + for _, segment := range c.spool.segmentsAt(mark) { + if err := replayCaptureSegment(segment, sink); err != nil { + return err + } + } + return nil +} + +func replayCaptureSegment(segment io.Reader, sink Sink) error { + decoder := json.NewDecoder(segment) + for { + var record captureRecord + if err := decoder.Decode(&record); err != nil { + if errors.Is(err, io.EOF) { + return nil + } + // A rotated segment can begin mid-line; the remainder of this + // segment is unparseable, and the truncation record already + // accounts for the gap. + return nil + } + + data, err := base64.StdEncoding.DecodeString(record.Data) + if err != nil { + continue + } + chunk := Chunk{ + Phase: Phase(record.Phase), + Stream: Stream(record.Stream), + Data: data, + } + if observed, parseErr := time.Parse(time.RFC3339Nano, record.ObservedAt); parseErr == nil { + chunk.ObservedAt = observed + } + if err := sink.WriteChunk(chunk); err != nil { + return err + } + } +} + +// Close stops the encoder, deletes the capture's bytes, and returns its share +// of the disk budget. It is idempotent. +func (c *TaskLogCapture) Close() { + c.closeOnce.Do(func() { + close(c.queue) + <-c.done + _ = c.spool.Close() + c.store.release(c.reserved) + }) +} diff --git a/internal/debuglog/coordinator.go b/internal/debuglog/coordinator.go new file mode 100644 index 0000000..13e36e9 --- /dev/null +++ b/internal/debuglog/coordinator.go @@ -0,0 +1,443 @@ +package debuglog + +import ( + "context" + "errors" + "sync" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/log" + "github.com/warpdotdev/oz-agent-worker/internal/metrics" + "github.com/warpdotdev/oz-agent-worker/internal/types" +) + +// maxCachedRequests bounds the idempotency cache. +const maxCachedRequests = 1024 + +// snapshotFinalizeDeadline bounds how long a terminal direct capture may keep +// draining before a snapshot reads it. +const snapshotFinalizeDeadline = 2 * time.Second + +// Ownership describes this process's claim over one exact (run, execution) +// pair. Only the process that executed the assignment produces one. +type Ownership struct { + // BackendKind is the backend that executed the assignment. + BackendKind string + // InCleanupGrace reports whether the execution has already reported + // terminal state and is being retained for its cleanup grace. + InCleanupGrace bool + // Capture is the direct backend's bounded output capture, if any. + Capture *TaskLogCapture +} + +// state names the ownership phase this request was served from, for metrics. +func (o Ownership) state() string { + if o.InCleanupGrace { + return metrics.DebugArchiveOwnershipCleanupGrace + } + return metrics.DebugArchiveOwnershipActive +} + +// OwnershipLookup resolves an exact (run, execution) pair to this process's +// claim. A false result means this process did not execute the assignment and +// must stay silent. +type OwnershipLookup interface { + LookupExecution(runID, executionID string) (Ownership, bool) +} + +// SnapshotSource writes an execution's provider logs into a sink. It is +// implemented by the worker's backends. +type SnapshotSource interface { + SnapshotLogs(ctx context.Context, runID, executionID string, sink Sink) error +} + +// Sender enqueues an acknowledgement through the worker's single WebSocket +// writer. +type Sender interface { + SendDebugArchiveAck(ack *types.DebugArchiveLogsUploadedMessage) error +} + +// Coordinator runs debug-archive log requests off the WebSocket read loop, so +// a slow provider read or upload can never delay heartbeats, cancellations, or +// later assignments. +type Coordinator struct { + ownership OwnershipLookup + source SnapshotSource + sender Sender + store *Store + uploader *Uploader + now func() time.Time + + uploadSlots chan struct{} + + mu sync.Mutex + cache map[string]*cacheEntry + perExec map[executionKey]*sync.Mutex + inflight sync.WaitGroup +} + +type executionKey struct { + runID string + executionID string +} + +// cacheEntry makes duplicate delivery of one request_id idempotent. It records +// only identifiers and result metadata, never the upload target. +type cacheEntry struct { + fingerprint string + retainUntil time.Time + + done chan struct{} + ack *types.DebugArchiveLogsUploadedMessage +} + +// CoordinatorOptions configures a Coordinator. +type CoordinatorOptions struct { + Ownership OwnershipLookup + Source SnapshotSource + Sender Sender + Store *Store + Uploader *Uploader + // Now supplies the coordinator's clock; nil uses time.Now. + Now func() time.Time +} + +// NewCoordinator builds a coordinator bounded by the store's configured upload +// concurrency. +func NewCoordinator(opts CoordinatorOptions) *Coordinator { + now := opts.Now + if now == nil { + now = time.Now + } + uploader := opts.Uploader + if uploader == nil { + uploader = NewUploader(nil, now) + } + return &Coordinator{ + ownership: opts.Ownership, + source: opts.Source, + sender: opts.Sender, + store: opts.Store, + uploader: uploader, + now: now, + uploadSlots: make(chan struct{}, opts.Store.Config().MaxConcurrentUploads), + cache: make(map[string]*cacheEntry), + perExec: make(map[executionKey]*sync.Mutex), + } +} + +// Handle dispatches a request to a background goroutine and returns +// immediately. The caller is the WebSocket read loop. +func (c *Coordinator) Handle(ctx context.Context, req *types.DebugArchiveLogsRequestedMessage) { + c.inflight.Add(1) + go func() { + defer c.inflight.Done() + c.process(ctx, req) + }() +} + +// Wait blocks until in-flight requests finish. Worker shutdown cancels their +// context first, so this does not extend shutdown beyond the bounded backend +// stop. +func (c *Coordinator) Wait() { c.inflight.Wait() } + +func (c *Coordinator) process(ctx context.Context, req *types.DebugArchiveLogsRequestedMessage) { + // Ownership is checked before validation: a process that did not execute + // this assignment must produce no acknowledgement, log, ID-bearing metric, + // or cache entry, even for a request it would otherwise reject. + owner, owns := c.ownership.LookupExecution(req.RunID, req.ExecutionID) + if !owns { + return + } + + effectiveMaxBytes, err := ValidateRequest(req, c.now(), c.store.Config().MaxExecutionBytes) + if err != nil { + var validation *ValidationError + if errors.As(err, &validation) { + c.respond(ctx, req, owner, failure(validation.ReasonCode, validation.Detail)) + return + } + c.respond(ctx, req, owner, failure(types.DebugArchiveReasonInvalidRequest, "request rejected")) + return + } + + entry, replay := c.admit(req, owner) + if replay != nil { + // A duplicate of a request already handled: replay the recorded + // acknowledgement without touching the provider or the network. + c.send(ctx, owner, replay) + return + } + if entry == nil { + c.respond(ctx, req, owner, failure( + types.DebugArchiveReasonRequestCapacityExhausted, + "worker request cache is full", + )) + return + } + + metrics.IncDebugArchiveRequestsInFlight() + ack := c.collect(ctx, req, owner, effectiveMaxBytes) + metrics.DecDebugArchiveRequestsInFlight() + + c.finish(entry, req, ack) + c.send(ctx, owner, ack) +} + +// collect performs the snapshot and upload for a request this process owns. +func (c *Coordinator) collect( + ctx context.Context, + req *types.DebugArchiveLogsRequestedMessage, + owner Ownership, + maxBytes int64, +) *types.DebugArchiveLogsUploadedMessage { + if owner.BackendKind == BackendCommand { + return c.ack(req, owner, failure( + types.DebugArchiveReasonBackendNotSupported, + "the command backend dispatches into an opaque runtime with no log API", + )) + } + + select { + case c.uploadSlots <- struct{}{}: + defer func() { <-c.uploadSlots }() + case <-ctx.Done(): + return c.ack(req, owner, failure(types.DebugArchiveReasonWorkerShuttingDown, "worker is shutting down")) + } + + execMutex := c.executionMutex(req.RunID, req.ExecutionID) + execMutex.Lock() + defer execMutex.Unlock() + + // Queueing never extends the request's deadline, and ownership can lapse + // while a slot is held, so both are re-checked before any provider work. + if !c.now().Before(req.ExpiresAt) { + return c.ack(req, owner, failure(types.DebugArchiveReasonUploadExpired, "request expired while queued")) + } + current, owns := c.ownership.LookupExecution(req.RunID, req.ExecutionID) + if !owns { + return c.ack(req, owner, failure(types.DebugArchiveReasonCleanupGraceExpired, "cleanup grace expired while queued")) + } + owner = current + + transformer, err := NewTransformer(req.ContentTransformer.Kind, req.ContentTransformer.Version) + if err != nil { + return c.ack(req, owner, failure(types.DebugArchiveReasonUnsupportedContentTransformer, "unsupported content transformer")) + } + + snapshot, err := c.store.NewSnapshot(owner.BackendKind, transformer, maxBytes) + if err != nil { + if errors.Is(err, ErrBudgetExhausted) { + return c.ack(req, owner, failure(types.DebugArchiveReasonRequestCapacityExhausted, "worker capture disk budget exhausted")) + } + return c.ack(req, owner, failure(types.DebugArchiveReasonSnapshotFailed, "failed to allocate a snapshot")) + } + defer snapshot.Close() + + if owner.Capture != nil && owner.InCleanupGrace { + owner.Capture.Finalize(snapshotFinalizeDeadline) + } + + start := c.now() + snapshotErr := c.source.SnapshotLogs(ctx, req.RunID, req.ExecutionID, snapshot.Sink()) + var partial *PartialSnapshotError + switch { + case snapshotErr == nil: + case errors.As(snapshotErr, &partial): + // A partial provider read still produced valid sibling data; it is + // uploaded with capture_status=partial rather than discarded. + default: + return c.ack(req, owner, failure(reasonForSnapshotError(snapshotErr), "log snapshot failed")) + } + + if err := snapshot.Finalize(); err != nil { + return c.ack(req, owner, failure(types.DebugArchiveReasonSnapshotFailed, "failed to finalize the snapshot")) + } + metrics.RecordDebugArchiveSnapshot(owner.BackendKind, c.now().Sub(start), snapshot.Bytes()) + if snapshot.Truncated() { + metrics.RecordDebugArchiveTruncation(owner.BackendKind) + } + + if snapshot.Bytes() == 0 { + return c.ack(req, owner, failure(types.DebugArchiveReasonCaptureUnavailable, "no log bytes were available for this execution")) + } + + uploadStart := c.now() + if err := c.uploader.Upload(ctx, req.UploadTarget, snapshot, req.ExpiresAt); err != nil { + metrics.RecordDebugArchiveUpload("failed", c.now().Sub(uploadStart)) + var uploadErr *UploadError + if errors.As(err, &uploadErr) { + return c.ack(req, owner, failure(uploadErr.ReasonCode, uploadErr.Detail)) + } + return c.ack(req, owner, failure(types.DebugArchiveReasonUploadFailed, "upload failed")) + } + metrics.RecordDebugArchiveUpload("uploaded", c.now().Sub(uploadStart)) + + warnings := combineWarnings(snapshot.Warnings(), partial) + captureStatus := types.DebugArchiveCaptureComplete + if len(warnings) > 0 { + captureStatus = types.DebugArchiveCapturePartial + } + + ack := c.ack(req, owner, result{outcome: types.DebugArchiveOutcomeUploaded}) + ack.Bytes = snapshot.Bytes() + ack.CRC32C = snapshot.CRC32C() + ack.SHA256 = snapshot.SHA256() + ack.Truncated = snapshot.Truncated() + ack.ContentTransformerVersion = transformer.Version() + ack.CaptureStatus = captureStatus + ack.WarningCodes = warnings + return ack +} + +// admit reserves the request's cache slot. It returns a replayable +// acknowledgement for a duplicate, or a nil entry when the cache is full. +func (c *Coordinator) admit(req *types.DebugArchiveLogsRequestedMessage, owner Ownership) (*cacheEntry, *types.DebugArchiveLogsUploadedMessage) { + fingerprint := requestFingerprint(req) + + c.mu.Lock() + existing, ok := c.cache[req.RequestID] + if ok { + c.mu.Unlock() + if existing.fingerprint != fingerprint { + // The same request ID with different immutable content is a + // server-side error, not a retry; honoring it would upload a + // second object under one identity. + return nil, c.ack(req, owner, failure( + types.DebugArchiveReasonInvalidRequest, + "request id was reused with different content", + )) + } + <-existing.done + return nil, existing.ack + } + + if len(c.cache) >= maxCachedRequests { + c.evictExpiredLocked() + if len(c.cache) >= maxCachedRequests { + c.mu.Unlock() + return nil, nil + } + } + + entry := &cacheEntry{fingerprint: fingerprint, done: make(chan struct{})} + c.cache[req.RequestID] = entry + c.mu.Unlock() + return entry, nil +} + +// finish records a completed request's result and publishes it to any +// duplicate waiting on the same request ID. +func (c *Coordinator) finish(entry *cacheEntry, req *types.DebugArchiveLogsRequestedMessage, ack *types.DebugArchiveLogsUploadedMessage) { + c.mu.Lock() + entry.ack = ack + entry.retainUntil = req.ExpiresAt + c.mu.Unlock() + close(entry.done) +} + +// evictExpiredLocked drops entries whose request has expired, so a long-lived +// worker reclaims cache space before refusing new requests. +func (c *Coordinator) evictExpiredLocked() { + now := c.now() + for id, entry := range c.cache { + select { + case <-entry.done: + default: + continue + } + if !entry.retainUntil.IsZero() && now.After(entry.retainUntil) { + delete(c.cache, id) + } + } +} + +func (c *Coordinator) executionMutex(runID, executionID string) *sync.Mutex { + key := executionKey{runID: runID, executionID: executionID} + + c.mu.Lock() + defer c.mu.Unlock() + if mutex, ok := c.perExec[key]; ok { + return mutex + } + mutex := &sync.Mutex{} + c.perExec[key] = mutex + return mutex +} + +// ForgetExecution drops the per-execution snapshot mutex once an execution's +// cleanup grace has expired. +func (c *Coordinator) ForgetExecution(runID, executionID string) { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.perExec, executionKey{runID: runID, executionID: executionID}) +} + +// result is the outcome-shaped part of an acknowledgement. +type result struct { + outcome string + reasonCode string + message string +} + +func failure(reasonCode, message string) result { + outcome := types.DebugArchiveOutcomeFailed + switch reasonCode { + case types.DebugArchiveReasonBackendNotSupported, + types.DebugArchiveReasonCaptureUnavailable, + types.DebugArchiveReasonCleanupGraceExpired, + types.DebugArchiveReasonResourceNotFound, + types.DebugArchiveReasonResourceNotReady: + outcome = types.DebugArchiveOutcomeUnavailable + } + return result{outcome: outcome, reasonCode: reasonCode, message: message} +} + +func (c *Coordinator) ack(req *types.DebugArchiveLogsRequestedMessage, owner Ownership, res result) *types.DebugArchiveLogsUploadedMessage { + return &types.DebugArchiveLogsUploadedMessage{ + ProtocolVersion: types.DebugArchiveProtocolVersion, + RequestID: req.RequestID, + ArchiveID: req.ArchiveID, + CollectionID: req.CollectionID, + RunID: req.RunID, + ExecutionID: req.ExecutionID, + Outcome: res.outcome, + BackendKind: owner.BackendKind, + WarningCodes: []string{}, + ReasonCode: res.reasonCode, + Message: SanitizeMessage(res.message), + } +} + +func (c *Coordinator) respond(ctx context.Context, req *types.DebugArchiveLogsRequestedMessage, owner Ownership, res result) { + c.send(ctx, owner, c.ack(req, owner, res)) +} + +func (c *Coordinator) send(ctx context.Context, owner Ownership, ack *types.DebugArchiveLogsUploadedMessage) { + metrics.RecordDebugArchiveRequest(ack.BackendKind, owner.state(), ack.Outcome, ack.ReasonCode) + if err := c.sender.SendDebugArchiveAck(ack); err != nil { + log.Warnf(ctx, "Failed to send debug archive acknowledgement: %v", err) + } +} + +// requestFingerprint captures the immutable fields a retry must repeat. It +// hashes nothing secret into a log line: the value stays in memory. +func requestFingerprint(req *types.DebugArchiveLogsRequestedMessage) string { + return req.ArchiveID + "|" + req.CollectionID + "|" + req.RunID + "|" + req.ExecutionID + "|" + + req.RequestedFormat + "|" + req.ExpiresAt.UTC().Format(time.RFC3339Nano) + "|" + + req.ContentTransformer.Kind + "|" + itoa(req.ContentTransformer.Version) + "|" + + itoa64(req.MaxBytes) + "|" + req.UploadTarget.Method + "|" + req.UploadTarget.URL +} + +func combineWarnings(encoderWarnings []string, partial *PartialSnapshotError) []string { + set := warningSet{} + for _, code := range encoderWarnings { + set.add(code) + } + if partial != nil { + for _, code := range partial.WarningCodes { + set.add(code) + } + } + return set.codes() +} diff --git a/internal/debuglog/coordinator_test.go b/internal/debuglog/coordinator_test.go new file mode 100644 index 0000000..82fc829 --- /dev/null +++ b/internal/debuglog/coordinator_test.go @@ -0,0 +1,531 @@ +package debuglog + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/types" +) + +// fakeOwnership answers exactly the executions a test says this process ran. +type fakeOwnership struct { + mu sync.Mutex + owned map[executionKey]Ownership +} + +func newFakeOwnership() *fakeOwnership { + return &fakeOwnership{owned: make(map[executionKey]Ownership)} +} + +func (o *fakeOwnership) own(runID, executionID string, ownership Ownership) { + o.mu.Lock() + defer o.mu.Unlock() + o.owned[executionKey{runID: runID, executionID: executionID}] = ownership +} + +func (o *fakeOwnership) release(runID, executionID string) { + o.mu.Lock() + defer o.mu.Unlock() + delete(o.owned, executionKey{runID: runID, executionID: executionID}) +} + +func (o *fakeOwnership) LookupExecution(runID, executionID string) (Ownership, bool) { + o.mu.Lock() + defer o.mu.Unlock() + ownership, ok := o.owned[executionKey{runID: runID, executionID: executionID}] + return ownership, ok +} + +// fakeSource writes canned output into whatever sink it is handed. +type fakeSource struct { + calls atomic.Int32 + // output, when non-empty, is emitted as one chunk. + output string + // err, when non-nil, is returned after any output is written. + err error + // block, when non-nil, gates the snapshot so a test can hold a slot. + block chan struct{} + // started reports the execution ID of each snapshot as it begins, so a + // test can wait for a specific request to hold the upload slot instead of + // racing the coordinator's goroutines. + started chan string +} + +func (s *fakeSource) SnapshotLogs(_ context.Context, _, executionID string, sink Sink) error { + s.calls.Add(1) + if s.started != nil { + s.started <- executionID + } + if s.block != nil { + <-s.block + } + if s.output != "" { + if err := sink.WriteChunk(Chunk{Phase: PhaseAgent, Stream: StreamStdout, Data: []byte(s.output)}); err != nil { + return err + } + } + return s.err +} + +// awaitStart blocks until a snapshot for executionID has begun. +func (s *fakeSource) awaitStart(t *testing.T, executionID string) { + t.Helper() + select { + case got := <-s.started: + if got != executionID { + t.Fatalf("snapshot started for %q, want %q", got, executionID) + } + case <-time.After(5 * time.Second): + t.Fatalf("timed out waiting for the snapshot of %q to start", executionID) + } +} + +// recordingSender captures the acknowledgements the coordinator enqueues. +type recordingSender struct { + mu sync.Mutex + acks []*types.DebugArchiveLogsUploadedMessage + sent chan struct{} +} + +func newRecordingSender() *recordingSender { + return &recordingSender{sent: make(chan struct{}, 16)} +} + +func (s *recordingSender) SendDebugArchiveAck(ack *types.DebugArchiveLogsUploadedMessage) error { + s.mu.Lock() + s.acks = append(s.acks, ack) + s.mu.Unlock() + s.sent <- struct{}{} + return nil +} + +func (s *recordingSender) await(t *testing.T) *types.DebugArchiveLogsUploadedMessage { + t.Helper() + select { + case <-s.sent: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for an acknowledgement") + } + s.mu.Lock() + defer s.mu.Unlock() + return s.acks[len(s.acks)-1] +} + +func (s *recordingSender) count() int { + s.mu.Lock() + defer s.mu.Unlock() + return len(s.acks) +} + +type coordinatorFixture struct { + coordinator *Coordinator + ownership *fakeOwnership + source *fakeSource + sender *recordingSender + uploads *atomic.Int32 + targetURL string +} + +func newCoordinatorFixture(t *testing.T, source *fakeSource, mutate func(*Config)) *coordinatorFixture { + t.Helper() + + var uploads atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + uploads.Add(1) + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(server.Close) + + ownership := newFakeOwnership() + sender := newRecordingSender() + store := newTestStore(t, mutate) + + uploader := NewUploader(&http.Client{Timeout: 5 * time.Second}, time.Now) + uploader.sleep = func(context.Context, time.Duration) error { return nil } + + return &coordinatorFixture{ + coordinator: NewCoordinator(CoordinatorOptions{ + Ownership: ownership, + Source: source, + Sender: sender, + Store: store, + Uploader: uploader, + }), + ownership: ownership, + source: source, + sender: sender, + uploads: &uploads, + targetURL: server.URL, + } +} + +func (f *coordinatorFixture) request() types.DebugArchiveLogsRequestedMessage { + request := goldenRequest() + request.ExpiresAt = time.Now().Add(MaxRequestLifetime) + request.MaxBytes = 1 << 15 + request.UploadTarget.URL = f.targetURL + return request +} + +func TestCoordinatorNonOwnerStaysSilent(t *testing.T) { + fixture := newCoordinatorFixture(t, &fakeSource{output: "should not be read"}, nil) + request := fixture.request() + + fixture.coordinator.Handle(context.Background(), &request) + fixture.coordinator.Wait() + + if fixture.sender.count() != 0 { + t.Fatalf("a non-owning process sent %d acknowledgements, want 0", fixture.sender.count()) + } + if fixture.source.calls.Load() != 0 { + t.Fatal("a non-owning process read the provider") + } + if fixture.uploads.Load() != 0 { + t.Fatal("a non-owning process contacted the upload target") + } +} + +func TestCoordinatorProcessOwningAnotherExecutionOfTheSameRunStaysSilent(t *testing.T) { + fixture := newCoordinatorFixture(t, &fakeSource{output: "other execution output"}, nil) + request := fixture.request() + fixture.ownership.own(request.RunID, "a-different-execution", Ownership{BackendKind: BackendDocker}) + + fixture.coordinator.Handle(context.Background(), &request) + fixture.coordinator.Wait() + + if fixture.sender.count() != 0 { + t.Fatalf("owning another execution of the same run produced %d acknowledgements, want 0", fixture.sender.count()) + } + if fixture.uploads.Load() != 0 { + t.Fatal("owning another execution of the same run uploaded a snapshot") + } +} + +func TestCoordinatorOwnerUploadsAndAcknowledges(t *testing.T) { + fixture := newCoordinatorFixture(t, &fakeSource{output: "captured output"}, nil) + request := fixture.request() + fixture.ownership.own(request.RunID, request.ExecutionID, Ownership{BackendKind: BackendDocker}) + + fixture.coordinator.Handle(context.Background(), &request) + fixture.coordinator.Wait() + + ack := fixture.sender.await(t) + if ack.Outcome != types.DebugArchiveOutcomeUploaded { + t.Fatalf("outcome = %q (%s), want %q", ack.Outcome, ack.ReasonCode, types.DebugArchiveOutcomeUploaded) + } + if ack.ProtocolVersion != types.DebugArchiveProtocolVersion { + t.Errorf("protocol version = %d, want %d", ack.ProtocolVersion, types.DebugArchiveProtocolVersion) + } + for field, got := range map[string]string{ + "request_id": ack.RequestID, + "archive_id": ack.ArchiveID, + "collection_id": ack.CollectionID, + "run_id": ack.RunID, + "execution_id": ack.ExecutionID, + } { + if got == "" { + t.Errorf("%s is empty; every identity field must round-trip", field) + } + } + if ack.BackendKind != BackendDocker { + t.Errorf("backend_kind = %q, want %q", ack.BackendKind, BackendDocker) + } + if ack.Bytes == 0 || ack.CRC32C == "" || ack.SHA256 == "" { + t.Errorf("expected byte count and both digests, got bytes=%d crc32c=%q sha256=%q", ack.Bytes, ack.CRC32C, ack.SHA256) + } + if ack.Truncated { + t.Error("a snapshot below its bound must report truncated=false") + } + if ack.CaptureStatus != types.DebugArchiveCaptureComplete { + t.Errorf("capture_status = %q, want %q", ack.CaptureStatus, types.DebugArchiveCaptureComplete) + } + if ack.ContentTransformerVersion != 1 { + t.Errorf("content_transformer_version = %d, want 1", ack.ContentTransformerVersion) + } + if len(ack.WarningCodes) != 0 { + t.Errorf("warning_codes = %v, want none", ack.WarningCodes) + } + if fixture.uploads.Load() != 1 { + t.Errorf("uploads = %d, want 1", fixture.uploads.Load()) + } +} + +func TestCoordinatorAcknowledgementNeverLeaksTargetOrContent(t *testing.T) { + fixture := newCoordinatorFixture(t, &fakeSource{output: "super secret log line"}, nil) + request := fixture.request() + fixture.ownership.own(request.RunID, request.ExecutionID, Ownership{BackendKind: BackendDocker}) + + fixture.coordinator.Handle(context.Background(), &request) + fixture.coordinator.Wait() + + ack := fixture.sender.await(t) + for _, forbidden := range []string{fixture.targetURL, "super secret log line", "/tmp/"} { + if strings.Contains(ack.Message, forbidden) { + t.Fatalf("acknowledgement message leaked %q: %q", forbidden, ack.Message) + } + } +} + +func TestCoordinatorPartialSnapshotUploadsWithWarnings(t *testing.T) { + source := &fakeSource{ + output: "readable sibling output", + err: &PartialSnapshotError{WarningCodes: []string{WarningContainerLogsUnavailable}}, + } + fixture := newCoordinatorFixture(t, source, nil) + request := fixture.request() + fixture.ownership.own(request.RunID, request.ExecutionID, Ownership{BackendKind: BackendKubernetes}) + + fixture.coordinator.Handle(context.Background(), &request) + fixture.coordinator.Wait() + + ack := fixture.sender.await(t) + if ack.Outcome != types.DebugArchiveOutcomeUploaded { + t.Fatalf("outcome = %q, want the readable siblings to still upload", ack.Outcome) + } + if ack.CaptureStatus != types.DebugArchiveCapturePartial { + t.Fatalf("capture_status = %q, want %q", ack.CaptureStatus, types.DebugArchiveCapturePartial) + } + if len(ack.WarningCodes) != 1 || ack.WarningCodes[0] != WarningContainerLogsUnavailable { + t.Fatalf("warning_codes = %v, want [%s]", ack.WarningCodes, WarningContainerLogsUnavailable) + } +} + +func TestCoordinatorOwningEntryWithNoBytesIsUnavailable(t *testing.T) { + fixture := newCoordinatorFixture(t, &fakeSource{}, nil) + request := fixture.request() + fixture.ownership.own(request.RunID, request.ExecutionID, Ownership{BackendKind: BackendDocker}) + + fixture.coordinator.Handle(context.Background(), &request) + fixture.coordinator.Wait() + + ack := fixture.sender.await(t) + if ack.Outcome != types.DebugArchiveOutcomeUnavailable { + t.Fatalf("outcome = %q, want %q", ack.Outcome, types.DebugArchiveOutcomeUnavailable) + } + if ack.ReasonCode != types.DebugArchiveReasonCaptureUnavailable { + t.Fatalf("reason = %q, want %q", ack.ReasonCode, types.DebugArchiveReasonCaptureUnavailable) + } + if fixture.uploads.Load() != 0 { + t.Fatal("an empty capture must not upload a zero-byte object") + } +} + +func TestCoordinatorCommandBackendReportsUnsupported(t *testing.T) { + fixture := newCoordinatorFixture(t, &fakeSource{}, nil) + request := fixture.request() + fixture.ownership.own(request.RunID, request.ExecutionID, Ownership{BackendKind: BackendCommand}) + + fixture.coordinator.Handle(context.Background(), &request) + fixture.coordinator.Wait() + + ack := fixture.sender.await(t) + if ack.Outcome != types.DebugArchiveOutcomeUnavailable { + t.Fatalf("outcome = %q, want %q", ack.Outcome, types.DebugArchiveOutcomeUnavailable) + } + if ack.ReasonCode != types.DebugArchiveReasonBackendNotSupported { + t.Fatalf("reason = %q, want %q", ack.ReasonCode, types.DebugArchiveReasonBackendNotSupported) + } + if fixture.source.calls.Load() != 0 { + t.Fatal("an unsupported backend must not be asked for logs") + } +} + +func TestCoordinatorRejectsUnsupportedProtocolVersion(t *testing.T) { + fixture := newCoordinatorFixture(t, &fakeSource{output: "x"}, nil) + request := fixture.request() + request.ProtocolVersion = 99 + fixture.ownership.own(request.RunID, request.ExecutionID, Ownership{BackendKind: BackendDocker}) + + fixture.coordinator.Handle(context.Background(), &request) + fixture.coordinator.Wait() + + ack := fixture.sender.await(t) + if ack.Outcome != types.DebugArchiveOutcomeFailed { + t.Fatalf("outcome = %q, want %q", ack.Outcome, types.DebugArchiveOutcomeFailed) + } + if ack.ReasonCode != types.DebugArchiveReasonUnsupportedProtocolVersion { + t.Fatalf("reason = %q, want %q", ack.ReasonCode, types.DebugArchiveReasonUnsupportedProtocolVersion) + } + if fixture.uploads.Load() != 0 { + t.Fatal("an unsupported protocol version must not contact the target") + } +} + +func TestCoordinatorUnsupportedTransformerUploadsNothing(t *testing.T) { + fixture := newCoordinatorFixture(t, &fakeSource{output: "x"}, nil) + request := fixture.request() + request.ContentTransformer = types.ContentTransformerDescriptor{Kind: "redact", Version: 3} + fixture.ownership.own(request.RunID, request.ExecutionID, Ownership{BackendKind: BackendDocker}) + + fixture.coordinator.Handle(context.Background(), &request) + fixture.coordinator.Wait() + + ack := fixture.sender.await(t) + if ack.ReasonCode != types.DebugArchiveReasonUnsupportedContentTransformer { + t.Fatalf("reason = %q, want %q", ack.ReasonCode, types.DebugArchiveReasonUnsupportedContentTransformer) + } + if fixture.uploads.Load() != 0 { + t.Fatal("an unsupported transformer must never upload untransformed data") + } +} + +func TestCoordinatorDuplicateRequestReplaysOneAcknowledgement(t *testing.T) { + fixture := newCoordinatorFixture(t, &fakeSource{output: "captured output"}, nil) + request := fixture.request() + fixture.ownership.own(request.RunID, request.ExecutionID, Ownership{BackendKind: BackendDocker}) + + fixture.coordinator.Handle(context.Background(), &request) + fixture.coordinator.Wait() + first := fixture.sender.await(t) + + duplicate := request + fixture.coordinator.Handle(context.Background(), &duplicate) + fixture.coordinator.Wait() + replayed := fixture.sender.await(t) + + if fixture.source.calls.Load() != 1 { + t.Fatalf("provider reads = %d, want exactly one", fixture.source.calls.Load()) + } + if fixture.uploads.Load() != 1 { + t.Fatalf("uploads = %d, want exactly one", fixture.uploads.Load()) + } + if replayed.SHA256 != first.SHA256 || replayed.Bytes != first.Bytes || replayed.Outcome != first.Outcome { + t.Fatalf("replayed acknowledgement %+v differs from the original %+v", replayed, first) + } +} + +func TestCoordinatorRejectsReusedRequestIDWithDifferentContent(t *testing.T) { + fixture := newCoordinatorFixture(t, &fakeSource{output: "captured output"}, nil) + request := fixture.request() + fixture.ownership.own(request.RunID, request.ExecutionID, Ownership{BackendKind: BackendDocker}) + + fixture.coordinator.Handle(context.Background(), &request) + fixture.coordinator.Wait() + fixture.sender.await(t) + + conflicting := request + conflicting.ArchiveID = "a-different-archive" + fixture.coordinator.Handle(context.Background(), &conflicting) + fixture.coordinator.Wait() + + ack := fixture.sender.await(t) + if ack.Outcome != types.DebugArchiveOutcomeFailed { + t.Fatalf("outcome = %q, want %q", ack.Outcome, types.DebugArchiveOutcomeFailed) + } + if ack.ReasonCode != types.DebugArchiveReasonInvalidRequest { + t.Fatalf("reason = %q, want %q", ack.ReasonCode, types.DebugArchiveReasonInvalidRequest) + } + if fixture.uploads.Load() != 1 { + t.Fatalf("uploads = %d, want the conflicting request to upload nothing", fixture.uploads.Load()) + } +} + +func TestCoordinatorReportsExpiryWhenQueuedPastTheDeadline(t *testing.T) { + blocked := make(chan struct{}) + source := &fakeSource{output: "captured output", block: blocked, started: make(chan string, 4)} + fixture := newCoordinatorFixture(t, source, func(c *Config) { c.MaxConcurrentUploads = 1 }) + + first := fixture.request() + fixture.ownership.own(first.RunID, first.ExecutionID, Ownership{BackendKind: BackendDocker}) + fixture.coordinator.Handle(context.Background(), &first) + source.awaitStart(t, first.ExecutionID) + + // A second request for a different execution has to wait for the only + // upload slot, and its deadline lapses while it does. + second := fixture.request() + second.RequestID = "second-request" + second.ExecutionID = "second-execution" + second.ExpiresAt = time.Now().Add(150 * time.Millisecond) + fixture.ownership.own(second.RunID, second.ExecutionID, Ownership{BackendKind: BackendDocker}) + fixture.coordinator.Handle(context.Background(), &second) + + time.Sleep(300 * time.Millisecond) + close(blocked) + fixture.coordinator.Wait() + + var expired *types.DebugArchiveLogsUploadedMessage + fixture.sender.mu.Lock() + for _, ack := range fixture.sender.acks { + if ack.RequestID == "second-request" { + expired = ack + } + } + fixture.sender.mu.Unlock() + + if expired == nil { + t.Fatal("the queued request produced no acknowledgement") + } + if expired.ReasonCode != types.DebugArchiveReasonUploadExpired { + t.Fatalf("reason = %q, want %q", expired.ReasonCode, types.DebugArchiveReasonUploadExpired) + } +} + +func TestCoordinatorReportsCleanupGraceExpiryWhenOwnershipLapses(t *testing.T) { + blocked := make(chan struct{}) + source := &fakeSource{output: "captured output", block: blocked, started: make(chan string, 4)} + fixture := newCoordinatorFixture(t, source, func(c *Config) { c.MaxConcurrentUploads = 1 }) + + holding := fixture.request() + fixture.ownership.own(holding.RunID, holding.ExecutionID, Ownership{BackendKind: BackendDocker}) + fixture.coordinator.Handle(context.Background(), &holding) + source.awaitStart(t, holding.ExecutionID) + + lapsing := fixture.request() + lapsing.RequestID = "lapsing-request" + lapsing.ExecutionID = "lapsing-execution" + fixture.ownership.own(lapsing.RunID, lapsing.ExecutionID, Ownership{BackendKind: BackendDocker, InCleanupGrace: true}) + fixture.coordinator.Handle(context.Background(), &lapsing) + + // The grace deadline passes while the request waits for a slot. + time.Sleep(100 * time.Millisecond) + fixture.ownership.release(lapsing.RunID, lapsing.ExecutionID) + close(blocked) + fixture.coordinator.Wait() + + var lapsed *types.DebugArchiveLogsUploadedMessage + fixture.sender.mu.Lock() + for _, ack := range fixture.sender.acks { + if ack.RequestID == "lapsing-request" { + lapsed = ack + } + } + fixture.sender.mu.Unlock() + + if lapsed == nil { + t.Fatal("the lapsing request produced no acknowledgement") + } + if lapsed.ReasonCode != types.DebugArchiveReasonCleanupGraceExpired { + t.Fatalf("reason = %q, want %q", lapsed.ReasonCode, types.DebugArchiveReasonCleanupGraceExpired) + } +} + +func TestCoordinatorSemaphoreBoundsConcurrentHandlers(t *testing.T) { + blocked := make(chan struct{}) + source := &fakeSource{output: "captured output", block: blocked, started: make(chan string, 8)} + fixture := newCoordinatorFixture(t, source, func(c *Config) { c.MaxConcurrentUploads = 1 }) + + for i := 0; i < 3; i++ { + request := fixture.request() + request.RequestID = "request-" + string(rune('a'+i)) + request.ExecutionID = "execution-" + string(rune('a'+i)) + fixture.ownership.own(request.RunID, request.ExecutionID, Ownership{BackendKind: BackendDocker}) + fixture.coordinator.Handle(context.Background(), &request) + } + + time.Sleep(200 * time.Millisecond) + if got := source.calls.Load(); got != 1 { + t.Fatalf("concurrent provider reads = %d, want the configured limit of 1", got) + } + + close(blocked) + fixture.coordinator.Wait() + if got := source.calls.Load(); got != 3 { + t.Fatalf("total provider reads = %d, want 3 once the slot frees", got) + } +} diff --git a/internal/debuglog/encoder.go b/internal/debuglog/encoder.go new file mode 100644 index 0000000..d09f9b9 --- /dev/null +++ b/internal/debuglog/encoder.go @@ -0,0 +1,311 @@ +package debuglog + +import ( + "encoding/base64" + "encoding/json" + "fmt" + "io" + "os" + "time" + "unicode/utf8" +) + +// Chunk is one bounded piece of provider output offered to a Sink. +type Chunk struct { + Phase Phase + // Stream must reflect what the provider actually reports. Backends that + // cannot separate streams use StreamCombined or StreamUnknown. + Stream Stream + // Timestamp is the provider's timestamp; the zero value omits it. + Timestamp time.Time + // ObservedAt is when the worker saw this output. The zero value uses the + // encoder's clock, which is correct for output read live from a provider; + // a replayed capture supplies its own recorded observation time. + ObservedAt time.Time + Source SourceIdentity + Data []byte +} + +// Sink receives snapshot records from a backend's SnapshotTaskLogs. Backends +// hand it provider output and identity; framing, chunk bounds, transformation, +// encoding selection, sequencing, and truncation belong to the sink. +type Sink interface { + // WriteChunk emits data, splitting it into records no larger than + // MaxChunkBytes. Empty data emits nothing, so an empty provider stream + // never produces a misleading zero-byte record. + WriteChunk(chunk Chunk) error + // WriteSourceError records that one source could not be read, using a + // stable warning code rather than the provider's error text. + WriteSourceError(phase Phase, stream Stream, source SourceIdentity, warningCode string) error + // NoteOmittedBytes reports bytes an upstream capture already dropped so + // the snapshot's truncation state stays truthful. + NoteOmittedBytes(n int64) +} + +// Encoder writes schema-v1 NDJSON into a bounded spool and finalizes it into +// one immutable snapshot file. +type Encoder struct { + backend string + transformer ContentTransformer + now func() time.Time + spool *boundedSpool + + sequence int64 + upstreamOmitted int64 + warnings warningSet +} + +// EncoderOptions configures a snapshot encoder. +type EncoderOptions struct { + // Backend is the backend kind stamped on every record. + Backend string + // Transformer rewrites decoded message data before encoding. + Transformer ContentTransformer + // MaxBytes is the effective output bound for this snapshot. + MaxBytes int64 + // CreateFile allocates one spool segment. The caller owns naming and mode. + CreateFile func(suffix string) (*os.File, error) + // Now supplies the worker observation clock; nil uses time.Now. + Now func() time.Time +} + +// NewEncoder creates a snapshot encoder backed by a bounded spool. +func NewEncoder(opts EncoderOptions) (*Encoder, error) { + if opts.Transformer == nil { + return nil, fmt.Errorf("debuglog: encoder requires a content transformer") + } + if opts.CreateFile == nil { + return nil, fmt.Errorf("debuglog: encoder requires a file allocator") + } + spool, err := newBoundedSpool(opts.MaxBytes, opts.CreateFile) + if err != nil { + return nil, err + } + now := opts.Now + if now == nil { + now = time.Now + } + return &Encoder{ + backend: opts.Backend, + transformer: opts.Transformer, + now: now, + spool: spool, + }, nil +} + +// WriteChunk implements Sink. +func (e *Encoder) WriteChunk(chunk Chunk) error { + for _, part := range splitChunk(chunk.Data, MaxChunkBytes) { + transformed, err := e.transformer.Transform(part) + if err != nil { + return fmt.Errorf("debuglog: content transform failed: %w", err) + } + if len(transformed) == 0 { + continue + } + + observedAt := chunk.ObservedAt + if observedAt.IsZero() { + observedAt = e.now() + } + + e.sequence++ + record := dataRecord{ + SchemaVersion: SchemaVersion, + Kind: KindData, + Sequence: e.sequence, + Backend: e.backend, + Phase: string(chunk.Phase), + Stream: string(chunk.Stream), + ObservedAt: observedAt.UTC().Format(time.RFC3339Nano), + } + if !chunk.Timestamp.IsZero() { + record.Timestamp = chunk.Timestamp.UTC().Format(time.RFC3339Nano) + } + applySourceIdentity(&record, chunk.Source) + + if utf8.Valid(transformed) { + record.Encoding = EncodingUTF8 + record.Data = string(transformed) + } else { + record.Encoding = EncodingBase64 + record.Data = base64.StdEncoding.EncodeToString(transformed) + } + + if err := e.writeRecord(record); err != nil { + return err + } + } + return nil +} + +// WriteSourceError implements Sink. +func (e *Encoder) WriteSourceError(phase Phase, stream Stream, source SourceIdentity, warningCode string) error { + e.sequence++ + record := sourceErrorRecord{ + SchemaVersion: SchemaVersion, + Kind: KindSourceError, + Sequence: e.sequence, + Backend: e.backend, + Phase: string(phase), + Stream: string(stream), + ObservedAt: e.now().UTC().Format(time.RFC3339Nano), + WarningCode: warningCode, + ContainerID: source.ContainerID, + Namespace: source.Namespace, + Pod: source.Pod, + Container: source.Container, + ContainerType: source.ContainerType, + RestartAttempt: source.RestartAttempt, + Previous: source.Previous, + } + e.warnings.add(warningCode) + return e.writeRecord(record) +} + +// NoteOmittedBytes implements Sink. +func (e *Encoder) NoteOmittedBytes(n int64) { + if n > 0 { + e.upstreamOmitted += n + e.warnings.add(WarningOutputDropped) + } +} + +// FinalizeResult describes the immutable snapshot an encoder produced. +type FinalizeResult struct { + // Bytes is the exact size of the written snapshot. + Bytes int64 + // Truncated reports whether the snapshot carries a truncation record. + Truncated bool + // Warnings are the deduplicated codes to report in the acknowledgement. + Warnings []string +} + +// Finalize streams the retained head, the truncation record when bytes were +// dropped, and the retained tail into out. The encoder's spool is released. +// +// The written object never exceeds the encoder's configured bound: the spool +// caps its segments against a budget that already reserves room for the +// truncation record. +func (e *Encoder) Finalize(out io.Writer) (FinalizeResult, error) { + defer func() { + _ = e.spool.Close() + }() + + mark := e.spool.Watermark() + segments := e.spool.segmentsAt(mark) + + omitted := mark.omitted + e.upstreamOmitted + if mark.omitted > 0 { + e.warnings.add(WarningOutputDropped) + } + + var written int64 + // The head reader holds the earliest retained records; the truncation + // record then names the gap before the retained tail segments. + n, err := io.Copy(out, segments[0]) + written += n + if err != nil { + return FinalizeResult{}, err + } + + truncated := false + if omitted > 0 { + line, marshalErr := marshalLine(truncationRecord{ + SchemaVersion: SchemaVersion, + Kind: KindTruncation, + Policy: TruncationPolicyFirstLast, + OmittedBytesAtLeast: omitted, + }) + if marshalErr != nil { + return FinalizeResult{}, marshalErr + } + // A budget too small to hold even this record retains nothing at all, + // so the object stays empty and the request reports the capture as + // unavailable rather than shipping an over-budget object. + if written+int64(len(line)) <= e.spool.maxBytes { + gap, writeErr := out.Write(line) + written += int64(gap) + if writeErr != nil { + return FinalizeResult{}, writeErr + } + truncated = true + } + } + + for _, segment := range segments[1:] { + n, err = io.Copy(out, segment) + written += n + if err != nil { + return FinalizeResult{}, err + } + } + + if written > e.spool.maxBytes { + return FinalizeResult{}, fmt.Errorf("debuglog: snapshot overran its %d byte bound", e.spool.maxBytes) + } + + return FinalizeResult{ + Bytes: written, + Truncated: truncated, + Warnings: e.warnings.codes(), + }, nil +} + +// Close releases the encoder's spool without producing a snapshot. +func (e *Encoder) Close() error { return e.spool.Close() } + +func (e *Encoder) writeRecord(record any) error { + line, err := marshalLine(record) + if err != nil { + return err + } + return e.spool.WriteLine(line) +} + +func marshalLine(record any) ([]byte, error) { + encoded, err := json.Marshal(record) + if err != nil { + return nil, fmt.Errorf("debuglog: failed to encode NDJSON record: %w", err) + } + return append(encoded, '\n'), nil +} + +func applySourceIdentity(record *dataRecord, source SourceIdentity) { + record.ContainerID = source.ContainerID + record.Namespace = source.Namespace + record.Pod = source.Pod + record.Container = source.Container + record.ContainerType = source.ContainerType + record.RestartAttempt = source.RestartAttempt + record.Previous = source.Previous +} + +// splitChunk divides data into pieces no larger than limit, preferring to cut +// on a UTF-8 boundary so text output stays UTF-8 encodable instead of falling +// back to base64 at every chunk seam. +func splitChunk(data []byte, limit int) [][]byte { + if len(data) == 0 { + return nil + } + if len(data) <= limit { + return [][]byte{data} + } + + var parts [][]byte + for len(data) > limit { + cut := limit + for back := 0; back < utf8.UTFMax && cut > 1; back++ { + if utf8.RuneStart(data[cut]) { + break + } + cut-- + } + parts = append(parts, data[:cut]) + data = data[cut:] + } + if len(data) > 0 { + parts = append(parts, data) + } + return parts +} diff --git a/internal/debuglog/encoder_test.go b/internal/debuglog/encoder_test.go new file mode 100644 index 0000000..bb5465c --- /dev/null +++ b/internal/debuglog/encoder_test.go @@ -0,0 +1,516 @@ +package debuglog + +import ( + "bytes" + "encoding/base64" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +// fileAllocator hands the spool real files under a test-owned directory. +func fileAllocator(t *testing.T) func(string) (*os.File, error) { + t.Helper() + dir := t.TempDir() + var index int + return func(suffix string) (*os.File, error) { + index++ + path := filepath.Join(dir, suffix+"-"+time.Now().Format("150405.000000000")+"-"+string(rune('a'+index))) + return os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0o600) + } +} + +func newTestEncoder(t *testing.T, maxBytes int64, transformer ContentTransformer) *Encoder { + t.Helper() + if transformer == nil { + transformer = noopTransformer{} + } + encoder, err := NewEncoder(EncoderOptions{ + Backend: BackendDocker, + Transformer: transformer, + MaxBytes: maxBytes, + CreateFile: fileAllocator(t), + Now: func() time.Time { return time.Unix(1700000000, 0).UTC() }, + }) + if err != nil { + t.Fatalf("NewEncoder: %v", err) + } + return encoder +} + +func decodeRecords(t *testing.T, out []byte) []map[string]any { + t.Helper() + var records []map[string]any + for _, line := range bytes.Split(bytes.TrimRight(out, "\n"), []byte("\n")) { + if len(line) == 0 { + continue + } + var record map[string]any + if err := json.Unmarshal(line, &record); err != nil { + t.Fatalf("snapshot line is not valid JSON (%q): %v", line, err) + } + records = append(records, record) + } + return records +} + +func TestEncoderEmitsSchemaV1DataRecords(t *testing.T) { + encoder := newTestEncoder(t, 1<<20, nil) + restarts := int32(2) + + if err := encoder.WriteChunk(Chunk{ + Phase: PhaseContainer, + Stream: StreamStderr, + Timestamp: time.Unix(1690000000, 0).UTC(), + Source: SourceIdentity{ + Namespace: "agents", + Pod: "oz-task-run-exec-abcd", + Container: "task", + ContainerType: "regular", + RestartAttempt: &restarts, + Previous: true, + }, + Data: []byte("hello world"), + }); err != nil { + t.Fatalf("WriteChunk: %v", err) + } + + var out bytes.Buffer + result, err := encoder.Finalize(&out) + if err != nil { + t.Fatalf("Finalize: %v", err) + } + if result.Truncated { + t.Error("a snapshot below its bound must not be marked truncated") + } + + records := decodeRecords(t, out.Bytes()) + if len(records) != 1 { + t.Fatalf("record count = %d, want 1", len(records)) + } + record := records[0] + + for field, want := range map[string]any{ + "schema_version": float64(SchemaVersion), + "kind": KindData, + "sequence": float64(1), + "backend": BackendDocker, + "phase": string(PhaseContainer), + "stream": string(StreamStderr), + "encoding": EncodingUTF8, + "data": "hello world", + "namespace": "agents", + "pod": "oz-task-run-exec-abcd", + "container": "task", + "container_type": "regular", + "restart_attempt": float64(2), + "previous": true, + } { + if record[field] != want { + t.Errorf("%s = %v, want %v", field, record[field], want) + } + } + if record["observed_at"] == "" || record["timestamp"] == "" { + t.Errorf("expected both provider and observed timestamps, got %v", record) + } +} + +func TestEncoderOmitsProviderTimestampWhenAbsent(t *testing.T) { + encoder := newTestEncoder(t, 1<<20, nil) + if err := encoder.WriteChunk(Chunk{Phase: PhaseAgent, Stream: StreamStdout, Data: []byte("x")}); err != nil { + t.Fatalf("WriteChunk: %v", err) + } + + var out bytes.Buffer + if _, err := encoder.Finalize(&out); err != nil { + t.Fatalf("Finalize: %v", err) + } + if _, present := decodeRecords(t, out.Bytes())[0]["timestamp"]; present { + t.Error("a chunk with no provider timestamp must not report one") + } +} + +func TestEncoderEmptyChunkProducesNoRecord(t *testing.T) { + encoder := newTestEncoder(t, 1<<20, nil) + if err := encoder.WriteChunk(Chunk{Phase: PhaseAgent, Stream: StreamStdout}); err != nil { + t.Fatalf("WriteChunk: %v", err) + } + + var out bytes.Buffer + result, err := encoder.Finalize(&out) + if err != nil { + t.Fatalf("Finalize: %v", err) + } + if result.Bytes != 0 { + t.Fatalf("empty stream produced %d bytes, want 0", result.Bytes) + } +} + +func TestEncoderBase64EncodesInvalidUTF8(t *testing.T) { + encoder := newTestEncoder(t, 1<<20, nil) + binary := []byte{0x00, 0xff, 0xfe, 0x41} + if err := encoder.WriteChunk(Chunk{Phase: PhaseAgent, Stream: StreamStdout, Data: binary}); err != nil { + t.Fatalf("WriteChunk: %v", err) + } + + var out bytes.Buffer + if _, err := encoder.Finalize(&out); err != nil { + t.Fatalf("Finalize: %v", err) + } + record := decodeRecords(t, out.Bytes())[0] + if record["encoding"] != EncodingBase64 { + t.Fatalf("encoding = %v, want %v", record["encoding"], EncodingBase64) + } + decoded, err := base64.StdEncoding.DecodeString(record["data"].(string)) + if err != nil { + t.Fatalf("base64 payload did not decode: %v", err) + } + if !bytes.Equal(decoded, binary) { + t.Fatalf("round-tripped %v, want %v", decoded, binary) + } +} + +func TestEncoderSplitsChunksAtMaxChunkBytes(t *testing.T) { + encoder := newTestEncoder(t, 8<<20, nil) + payload := bytes.Repeat([]byte("a"), MaxChunkBytes*2+7) + if err := encoder.WriteChunk(Chunk{Phase: PhaseAgent, Stream: StreamStdout, Data: payload}); err != nil { + t.Fatalf("WriteChunk: %v", err) + } + + var out bytes.Buffer + if _, err := encoder.Finalize(&out); err != nil { + t.Fatalf("Finalize: %v", err) + } + + records := decodeRecords(t, out.Bytes()) + if len(records) != 3 { + t.Fatalf("record count = %d, want 3", len(records)) + } + var reassembled strings.Builder + for i, record := range records { + if got := int(record["sequence"].(float64)); got != i+1 { + t.Errorf("sequence = %d, want %d", got, i+1) + } + data := record["data"].(string) + if len(data) > MaxChunkBytes { + t.Errorf("record %d carries %d bytes, above the %d chunk bound", i, len(data), MaxChunkBytes) + } + reassembled.WriteString(data) + } + if reassembled.String() != string(payload) { + t.Error("decoded records did not reassemble into the original payload") + } +} + +func TestEncoderTruncatesWithRecordBoundedFirstLastPolicy(t *testing.T) { + // A tight bound forces the spool to rotate, dropping the middle of the + // stream while keeping whole records at both ends. + encoder := newTestEncoder(t, 4096, nil) + const lines = 400 + for i := 0; i < lines; i++ { + if err := encoder.WriteChunk(Chunk{ + Phase: PhaseAgent, + Stream: StreamStdout, + Data: []byte(strings.Repeat("x", 40)), + }); err != nil { + t.Fatalf("WriteChunk: %v", err) + } + } + + var out bytes.Buffer + result, err := encoder.Finalize(&out) + if err != nil { + t.Fatalf("Finalize: %v", err) + } + if !result.Truncated { + t.Fatal("expected the snapshot to be marked truncated") + } + if result.Bytes > 4096 { + t.Fatalf("snapshot is %d bytes, above its %d bound", result.Bytes, 4096) + } + + records := decodeRecords(t, out.Bytes()) + truncations := 0 + var truncationIndex int + for i, record := range records { + if record["kind"] == KindTruncation { + truncations++ + truncationIndex = i + if record["policy"] != TruncationPolicyFirstLast { + t.Errorf("policy = %v, want %v", record["policy"], TruncationPolicyFirstLast) + } + if omitted, _ := record["omitted_bytes_at_least"].(float64); omitted <= 0 { + t.Errorf("omitted_bytes_at_least = %v, want a positive lower bound", record["omitted_bytes_at_least"]) + } + } + } + if truncations != 1 { + t.Fatalf("truncation record count = %d, want 1", truncations) + } + if truncationIndex == 0 || truncationIndex == len(records)-1 { + t.Fatalf("truncation record at index %d: expected records on both sides", truncationIndex) + } + + first := records[0]["sequence"].(float64) + last := records[len(records)-1]["sequence"].(float64) + if first != 1 { + t.Errorf("first retained sequence = %v, want 1", first) + } + if last != lines { + t.Errorf("last retained sequence = %v, want %d", last, lines) + } +} + +func TestEncoderSourceErrorCarriesOnlyWarningCode(t *testing.T) { + encoder := newTestEncoder(t, 1<<20, nil) + source := SourceIdentity{Pod: "pod-1", Container: "task"} + if err := encoder.WriteSourceError(PhaseContainer, StreamCombined, source, WarningContainerLogsUnavailable); err != nil { + t.Fatalf("WriteSourceError: %v", err) + } + + var out bytes.Buffer + result, err := encoder.Finalize(&out) + if err != nil { + t.Fatalf("Finalize: %v", err) + } + + record := decodeRecords(t, out.Bytes())[0] + if record["kind"] != KindSourceError { + t.Fatalf("kind = %v, want %v", record["kind"], KindSourceError) + } + if record["warning_code"] != WarningContainerLogsUnavailable { + t.Fatalf("warning_code = %v, want %v", record["warning_code"], WarningContainerLogsUnavailable) + } + if _, present := record["data"]; present { + t.Error("a source-error record must not carry provider data") + } + if len(result.Warnings) != 1 || result.Warnings[0] != WarningContainerLogsUnavailable { + t.Fatalf("warnings = %v, want [%s]", result.Warnings, WarningContainerLogsUnavailable) + } +} + +// upperTransformer proves the transformer touches only decoded message data. +type upperTransformer struct{} + +func (upperTransformer) Kind() string { return "test-upper" } + +func (upperTransformer) Version() int { return 7 } + +func (upperTransformer) Transform(data []byte) ([]byte, error) { + return bytes.ToUpper(data), nil +} + +func TestEncoderTransformsOnlyMessageData(t *testing.T) { + provider := time.Unix(1690000000, 0).UTC() + source := SourceIdentity{Pod: "pod-a", Container: "task"} + + encodeWith := func(transformer ContentTransformer) map[string]any { + encoder := newTestEncoder(t, 1<<20, transformer) + if err := encoder.WriteChunk(Chunk{ + Phase: PhaseContainer, + Stream: StreamStdout, + Timestamp: provider, + Source: source, + Data: []byte("secret value"), + }); err != nil { + t.Fatalf("WriteChunk: %v", err) + } + var out bytes.Buffer + if _, err := encoder.Finalize(&out); err != nil { + t.Fatalf("Finalize: %v", err) + } + return decodeRecords(t, out.Bytes())[0] + } + + preserved := encodeWith(noopTransformer{}) + transformed := encodeWith(upperTransformer{}) + + if preserved["data"] != "secret value" { + t.Fatalf("the no-op transformer changed data: %v", preserved["data"]) + } + if transformed["data"] != "SECRET VALUE" { + t.Fatalf("data = %v, want the transformed value", transformed["data"]) + } + for _, field := range []string{"schema_version", "kind", "sequence", "backend", "phase", "stream", "timestamp", "observed_at", "pod", "container"} { + if preserved[field] != transformed[field] { + t.Errorf("transformer changed structural field %s: %v vs %v", field, preserved[field], transformed[field]) + } + } +} + +func TestEncoderReportsUpstreamOmittedBytesAsTruncation(t *testing.T) { + encoder := newTestEncoder(t, 1<<20, nil) + encoder.NoteOmittedBytes(4096) + if err := encoder.WriteChunk(Chunk{Phase: PhaseAgent, Stream: StreamStdout, Data: []byte("tail")}); err != nil { + t.Fatalf("WriteChunk: %v", err) + } + + var out bytes.Buffer + result, err := encoder.Finalize(&out) + if err != nil { + t.Fatalf("Finalize: %v", err) + } + if !result.Truncated { + t.Fatal("upstream omission must mark the snapshot truncated") + } + if len(result.Warnings) != 1 || result.Warnings[0] != WarningOutputDropped { + t.Fatalf("warnings = %v, want [%s]", result.Warnings, WarningOutputDropped) + } +} + +// TestEncoderNeverExceedsItsBound sweeps bounds from pathologically small to +// comfortably large. The request's max_bytes is a hard ceiling on the object +// the worker uploads, so no combination of record size, tail overshoot, and +// truncation metadata may push the finalized snapshot past it. +func TestEncoderNeverExceedsItsBound(t *testing.T) { + bounds := []int64{ + 1, 2, 16, 64, + maxTruncationLineBytes - 1, + maxTruncationLineBytes, + maxTruncationLineBytes + 1, + 256, 512, 1024, 4096, 65536, + } + payloads := map[string][]byte{ + "tiny": []byte("x"), + "line": bytes.Repeat([]byte("y"), 200), + "chunk": bytes.Repeat([]byte("z"), MaxChunkBytes), + } + + for _, bound := range bounds { + for name, payload := range payloads { + t.Run(fmt.Sprintf("bound=%d/%s", bound, name), func(t *testing.T) { + encoder := newTestEncoder(t, bound, nil) + for i := 0; i < 50; i++ { + if err := encoder.WriteChunk(Chunk{ + Phase: PhaseAgent, + Stream: StreamStdout, + Data: payload, + }); err != nil { + t.Fatalf("WriteChunk: %v", err) + } + } + + var out bytes.Buffer + result, err := encoder.Finalize(&out) + if err != nil { + t.Fatalf("Finalize: %v", err) + } + + if result.Bytes > bound { + t.Fatalf("snapshot is %d bytes, above its %d bound", result.Bytes, bound) + } + if int64(out.Len()) != result.Bytes { + t.Fatalf("wrote %d bytes but reported %d", out.Len(), result.Bytes) + } + // Whatever survives must still be parseable NDJSON. + decodeRecords(t, out.Bytes()) + }) + } + } +} + +func TestEncoderKeepsWholeRecordsAndReportsTheGapAtATightBound(t *testing.T) { + // Room for the reserved truncation record plus a handful of data records, + // so the snapshot must retain some output at both ends and name the gap. + bound := maxTruncationLineBytes + 1200 + encoder := newTestEncoder(t, bound, nil) + for i := 0; i < 200; i++ { + if err := encoder.WriteChunk(Chunk{ + Phase: PhaseAgent, + Stream: StreamStdout, + Data: []byte(strings.Repeat("q", 40)), + }); err != nil { + t.Fatalf("WriteChunk: %v", err) + } + } + + var out bytes.Buffer + result, err := encoder.Finalize(&out) + if err != nil { + t.Fatalf("Finalize: %v", err) + } + if result.Bytes > bound { + t.Fatalf("snapshot is %d bytes, above its %d bound", result.Bytes, bound) + } + if !result.Truncated { + t.Fatal("expected the snapshot to be marked truncated") + } + + records := decodeRecords(t, out.Bytes()) + truncations := 0 + for _, record := range records { + if record["kind"] == KindTruncation { + truncations++ + } + } + if truncations != 1 { + t.Fatalf("truncation record count = %d, want exactly 1", truncations) + } + if records[0]["kind"] != KindData { + t.Fatalf("first record kind = %v, want the earliest data retained", records[0]["kind"]) + } + if records[len(records)-1]["kind"] != KindData { + t.Fatalf("last record kind = %v, want the newest data retained", records[len(records)-1]["kind"]) + } +} + +func TestEncoderYieldsAnEmptyObjectWhenTheBoundCannotHoldARecord(t *testing.T) { + // A bound this small cannot hold even the truncation record. Emitting + // nothing keeps the ceiling hard; the coordinator turns an empty snapshot + // into a classified unavailable outcome rather than uploading it. + encoder := newTestEncoder(t, 4, nil) + if err := encoder.WriteChunk(Chunk{Phase: PhaseAgent, Stream: StreamStdout, Data: []byte("dropped")}); err != nil { + t.Fatalf("WriteChunk: %v", err) + } + + var out bytes.Buffer + result, err := encoder.Finalize(&out) + if err != nil { + t.Fatalf("Finalize: %v", err) + } + if result.Bytes != 0 { + t.Fatalf("snapshot is %d bytes, want an empty object", result.Bytes) + } + if result.Truncated { + t.Fatal("truncated must not be reported without a truncation record") + } +} + +func TestEncoderReportsARecordTooLargeForItsBoundAsOmitted(t *testing.T) { + // A single record wider than a tail segment cannot be retained without + // overrunning the bound, so it is dropped and accounted for rather than + // written in full. + bound := maxTruncationLineBytes + 400 + encoder := newTestEncoder(t, bound, nil) + if err := encoder.WriteChunk(Chunk{ + Phase: PhaseAgent, + Stream: StreamStdout, + Data: bytes.Repeat([]byte("w"), 4096), + }); err != nil { + t.Fatalf("WriteChunk: %v", err) + } + + var out bytes.Buffer + result, err := encoder.Finalize(&out) + if err != nil { + t.Fatalf("Finalize: %v", err) + } + if result.Bytes > bound { + t.Fatalf("snapshot is %d bytes, above its %d bound", result.Bytes, bound) + } + if !result.Truncated { + t.Fatal("dropping an oversized record must mark the snapshot truncated") + } + + records := decodeRecords(t, out.Bytes()) + if len(records) != 1 || records[0]["kind"] != KindTruncation { + t.Fatalf("records = %v, want only the truncation record", records) + } + if omitted, _ := records[0]["omitted_bytes_at_least"].(float64); omitted <= 0 { + t.Fatalf("omitted_bytes_at_least = %v, want the dropped record accounted for", records[0]["omitted_bytes_at_least"]) + } +} diff --git a/internal/debuglog/errors.go b/internal/debuglog/errors.go new file mode 100644 index 0000000..0ac37d4 --- /dev/null +++ b/internal/debuglog/errors.go @@ -0,0 +1,57 @@ +package debuglog + +import ( + "errors" + "strconv" + + "github.com/warpdotdev/oz-agent-worker/internal/types" +) + +// PartialSnapshotError reports that a backend wrote valid data for some +// sources but could not read others. The coordinator uploads the bytes that +// were written and marks the capture partial rather than discarding them. +type PartialSnapshotError struct { + WarningCodes []string +} + +func (e *PartialSnapshotError) Error() string { + return "debuglog: snapshot completed with partial provider data" +} + +// SnapshotError is a backend snapshot failure that produced no usable data. Its +// reason code becomes the acknowledgement's, so it must stay in the bounded +// vocabulary and its detail must never carry provider text. +type SnapshotError struct { + ReasonCode string + Detail string +} + +func (e *SnapshotError) Error() string { + return "debuglog: " + e.Detail + " (" + e.ReasonCode + ")" +} + +// NewSnapshotError builds a typed snapshot failure. +func NewSnapshotError(reasonCode, detail string) error { + return &SnapshotError{ReasonCode: reasonCode, Detail: detail} +} + +// ErrBackendNotSupported is the typed error a backend without a log API +// returns from SnapshotLogs. +var ErrBackendNotSupported = &SnapshotError{ + ReasonCode: types.DebugArchiveReasonBackendNotSupported, + Detail: "backend does not expose execution logs", +} + +// reasonForSnapshotError maps a backend error onto the acknowledgement's +// bounded reason vocabulary. +func reasonForSnapshotError(err error) string { + var snapshotErr *SnapshotError + if errors.As(err, &snapshotErr) { + return snapshotErr.ReasonCode + } + return types.DebugArchiveReasonSnapshotFailed +} + +func itoa(value int) string { return strconv.Itoa(value) } + +func itoa64(value int64) string { return strconv.FormatInt(value, 10) } diff --git a/internal/debuglog/ndjson.go b/internal/debuglog/ndjson.go new file mode 100644 index 0000000..dc5eb67 --- /dev/null +++ b/internal/debuglog/ndjson.go @@ -0,0 +1,155 @@ +// Package debuglog implements the worker half of the REMOTE-2516 debug-archive +// log protocol: request validation, bounded NDJSON snapshot encoding, secure +// disk-backed capture for the direct backend, upload to a server-supplied +// target, and the asynchronous coordinator that ties them together. +// +// Nothing in this package may log captured bytes, upload targets, signed +// headers or form fields, local capture paths, or upload response bodies. +package debuglog + +import ( + "fmt" + "math" +) + +// SchemaVersion is the NDJSON schema every snapshot record carries. +const SchemaVersion = 1 + +// Record kinds emitted into a snapshot. +const ( + KindData = "data" + KindSourceError = "source_error" + KindTruncation = "truncation" +) + +// Data encodings. A chunk that is valid UTF-8 is stored directly; anything +// else is base64 so the NDJSON stream stays valid. +const ( + EncodingUTF8 = "utf8" + EncodingBase64 = "base64" +) + +// Stream identifies which of a source's streams a chunk came from. A provider +// that merges its streams reports StreamCombined, and one that cannot say at +// all reports StreamUnknown; neither is ever inferred from log text. +type Stream string + +const ( + StreamStdout Stream = "stdout" + StreamStderr Stream = "stderr" + StreamCombined Stream = "combined" + StreamUnknown Stream = "unknown" +) + +// Phase is provider truth about which part of an execution produced a chunk, +// not a process classifier. Only the direct backend owns distinct +// setup/agent/teardown handles; container providers report PhaseContainer. +type Phase string + +const ( + PhaseSetup Phase = "setup" + PhaseAgent Phase = "agent" + PhaseTeardown Phase = "teardown" + PhaseContainer Phase = "container" +) + +// Backend kinds reported on every record and in the acknowledgement. +const ( + BackendDocker = "docker" + BackendKubernetes = "kubernetes" + BackendDirect = "direct" + BackendCommand = "command" +) + +// MaxChunkBytes bounds how much decoded data one NDJSON record may carry, so +// a single provider read can never produce an unbounded line. +const MaxChunkBytes = 32 * 1024 + +// TruncationPolicyFirstLast names the only bytes-dropped policy V1 implements. +const TruncationPolicyFirstLast = "first_last" + +// SourceIdentity carries the provider identity a backend actually supplies for +// a chunk. Fields the provider does not supply stay empty rather than being +// guessed. +type SourceIdentity struct { + // ContainerID is set by the Docker backend. + ContainerID string + // Namespace, Pod, Container, ContainerType, RestartAttempt, and Previous + // are set by the Kubernetes backend. + Namespace string + Pod string + Container string + ContainerType string + RestartAttempt *int32 + Previous bool +} + +// dataRecord is the wire shape of a schema-v1 data record. Only Data passes +// through the request's content transformer; every other field is structural. +type dataRecord struct { + SchemaVersion int `json:"schema_version"` + Kind string `json:"kind"` + Sequence int64 `json:"sequence"` + Backend string `json:"backend"` + Phase string `json:"phase"` + Stream string `json:"stream"` + Timestamp string `json:"timestamp,omitempty"` + ObservedAt string `json:"observed_at"` + Encoding string `json:"encoding"` + Data string `json:"data"` + ContainerID string `json:"container_id,omitempty"` + Namespace string `json:"namespace,omitempty"` + Pod string `json:"pod,omitempty"` + Container string `json:"container,omitempty"` + ContainerType string `json:"container_type,omitempty"` + RestartAttempt *int32 `json:"restart_attempt,omitempty"` + Previous bool `json:"previous,omitempty"` +} + +// sourceErrorRecord reports that one source could not be read. It carries only +// safe source identity plus a stable warning code, never the provider's error. +type sourceErrorRecord struct { + SchemaVersion int `json:"schema_version"` + Kind string `json:"kind"` + Sequence int64 `json:"sequence"` + Backend string `json:"backend"` + Phase string `json:"phase"` + Stream string `json:"stream"` + ObservedAt string `json:"observed_at"` + WarningCode string `json:"warning_code"` + ContainerID string `json:"container_id,omitempty"` + Namespace string `json:"namespace,omitempty"` + Pod string `json:"pod,omitempty"` + Container string `json:"container,omitempty"` + ContainerType string `json:"container_type,omitempty"` + RestartAttempt *int32 `json:"restart_attempt,omitempty"` + Previous bool `json:"previous,omitempty"` +} + +// truncationRecord marks the gap between the retained first and last portions +// of a bounded snapshot. +type truncationRecord struct { + SchemaVersion int `json:"schema_version"` + Kind string `json:"kind"` + Policy string `json:"policy"` + OmittedBytesAtLeast int64 `json:"omitted_bytes_at_least"` +} + +// maxTruncationLineBytes is the largest a truncation record line can be, with +// the omitted-byte count at its widest. A bounded snapshot reserves this much +// of its budget up front so emitting the record can never push the finalized +// object past the size the request asked for. +var maxTruncationLineBytes = func() int64 { + line, err := marshalLine(truncationRecord{ + SchemaVersion: SchemaVersion, + Kind: KindTruncation, + Policy: TruncationPolicyFirstLast, + OmittedBytesAtLeast: math.MaxInt64, + }) + if err != nil { + // A fixed struct of scalars cannot fail to marshal; failing loudly + // here beats silently reserving nothing and overrunning the bound. + panic(fmt.Sprintf("debuglog: failed to size the truncation record: %v", err)) + } + return int64(len(line)) +}() diff --git a/internal/debuglog/request.go b/internal/debuglog/request.go new file mode 100644 index 0000000..aca46ae --- /dev/null +++ b/internal/debuglog/request.go @@ -0,0 +1,244 @@ +package debuglog + +import ( + "fmt" + "net/http" + "net/url" + "sort" + "strings" + "time" + "unicode/utf8" + + "github.com/warpdotdev/oz-agent-worker/internal/types" +) + +const ( + // ProtocolCeilingBytes is the largest snapshot the protocol permits for one + // execution. A server request may lower it but never raise it. + ProtocolCeilingBytes int64 = 64 << 20 + // MaxRequestLifetime bounds how far in the future a request's expiry may + // sit, matching the server's 30-minute presigned target. + MaxRequestLifetime = 30 * time.Minute + // MaxAckMessageBytes caps the sanitized human-readable acknowledgement text. + MaxAckMessageBytes = 256 + // MaxWarningCodes caps how many deduplicated warning codes one + // acknowledgement carries. + MaxWarningCodes = 16 +) + +// Warning codes re-exported so backends and the coordinator share one +// vocabulary with the wire contract. +const ( + WarningContainerLogsUnavailable = types.DebugArchiveWarningContainerLogsUnavailable + WarningPreviousLogsUnavailable = types.DebugArchiveWarningPreviousLogsUnavailable + WarningOutputDropped = types.DebugArchiveWarningOutputDropped + WarningProviderSnapshotIncomplete = types.DebugArchiveWarningProviderSnapshotIncomplete +) + +// ValidationError reports a request this worker refuses, carrying the stable +// reason code the acknowledgement reports. +type ValidationError struct { + ReasonCode string + Detail string +} + +func (e *ValidationError) Error() string { + return fmt.Sprintf("debuglog: %s (%s)", e.Detail, e.ReasonCode) +} + +func invalid(detail string) error { + return &ValidationError{ReasonCode: types.DebugArchiveReasonInvalidRequest, Detail: detail} +} + +// ValidateRequest checks every field the worker relies on before it touches a +// backend or the network. The effective output bound is the lower of the +// request's and the worker's configured ceilings. +// +// Callers must establish ownership first: a non-owning process stays silent +// even for a request that would fail validation. +func ValidateRequest(req *types.DebugArchiveLogsRequestedMessage, now time.Time, configuredCeiling int64) (effectiveMaxBytes int64, err error) { + if req.ProtocolVersion != types.DebugArchiveProtocolVersion { + return 0, &ValidationError{ + ReasonCode: types.DebugArchiveReasonUnsupportedProtocolVersion, + Detail: "unsupported protocol version", + } + } + + for name, value := range map[string]string{ + "request_id": req.RequestID, + "archive_id": req.ArchiveID, + "collection_id": req.CollectionID, + "run_id": req.RunID, + "execution_id": req.ExecutionID, + } { + if strings.TrimSpace(value) == "" { + return 0, invalid("missing " + name) + } + if containsControlCharacters(value) { + return 0, invalid(name + " contains control characters") + } + } + + if req.RequestedFormat != types.DebugArchiveFormatNDJSON { + return 0, invalid("unsupported requested format") + } + + if req.ExpiresAt.IsZero() { + return 0, invalid("missing expires_at") + } + if !req.ExpiresAt.After(now) { + return 0, &ValidationError{ + ReasonCode: types.DebugArchiveReasonRequestExpired, + Detail: "request already expired on receipt", + } + } + if req.ExpiresAt.Sub(now) > MaxRequestLifetime { + return 0, invalid("expires_at exceeds the protocol request lifetime") + } + + if _, transformerErr := NewTransformer(req.ContentTransformer.Kind, req.ContentTransformer.Version); transformerErr != nil { + return 0, &ValidationError{ + ReasonCode: types.DebugArchiveReasonUnsupportedContentTransformer, + Detail: "unsupported content transformer", + } + } + + if req.MaxBytes <= 0 { + return 0, invalid("max_bytes must be positive") + } + if req.MaxBytes > ProtocolCeilingBytes { + return 0, invalid("max_bytes exceeds the protocol ceiling") + } + effectiveMaxBytes = req.MaxBytes + if configuredCeiling > 0 && configuredCeiling < effectiveMaxBytes { + effectiveMaxBytes = configuredCeiling + } + + if err := validateUploadTarget(req.UploadTarget); err != nil { + return 0, err + } + + return effectiveMaxBytes, nil +} + +func validateUploadTarget(target types.UploadTarget) error { + switch target.Method { + case http.MethodPut: + if len(target.MultipartFields) > 0 { + return invalid("PUT target must not carry multipart fields") + } + case http.MethodPost: + if len(target.MultipartFields) == 0 { + return invalid("POST target requires multipart fields") + } + default: + return invalid("upload method must be PUT or POST") + } + + if containsControlCharacters(target.URL) { + return invalid("upload target contains control characters") + } + parsed, err := url.Parse(target.URL) + if err != nil { + return invalid("upload target is not a valid URL") + } + if !isPermittedTargetURL(parsed) { + return invalid("upload target must use HTTPS outside loopback") + } + + for key, value := range target.Headers { + if containsControlCharacters(key) || containsControlCharacters(value) { + return invalid("upload target header contains control characters") + } + } + for key, value := range target.MultipartFields { + if containsControlCharacters(key) || containsControlCharacters(value) { + return invalid("upload target field contains control characters") + } + } + return nil +} + +// isPermittedTargetURL allows HTTPS everywhere and plain HTTP only against +// loopback, so integration tests and local development can exercise the upload +// path without ever sending customer log bytes over an unencrypted network. +func isPermittedTargetURL(parsed *url.URL) bool { + switch parsed.Scheme { + case "https": + return parsed.Host != "" + case "http": + return isLoopbackHost(parsed.Hostname()) + default: + return false + } +} + +func isLoopbackHost(host string) bool { + switch host { + case "localhost", "127.0.0.1", "::1": + return true + default: + return strings.HasPrefix(host, "127.") + } +} + +func containsControlCharacters(value string) bool { + for _, r := range value { + if r < 0x20 || r == 0x7f { + return true + } + } + return false +} + +// SanitizeMessage bounds and strips human-readable acknowledgement text. It +// never receives provider output, URLs, headers, local paths, or response +// bodies; this is the last line of defense against accidentally forwarding one. +func SanitizeMessage(message string) string { + cleaned := strings.Map(func(r rune) rune { + if r < 0x20 || r == 0x7f { + return ' ' + } + return r + }, message) + cleaned = strings.Join(strings.Fields(cleaned), " ") + + if len(cleaned) <= MaxAckMessageBytes { + return cleaned + } + truncated := cleaned[:MaxAckMessageBytes] + for len(truncated) > 0 && !utf8.ValidString(truncated) { + truncated = truncated[:len(truncated)-1] + } + return truncated +} + +// warningSet deduplicates warning codes and caps them at MaxWarningCodes so an +// acknowledgement stays bounded regardless of how many sources degrade. +type warningSet struct { + seen map[string]struct{} + order []string +} + +func (w *warningSet) add(code string) { + if code == "" || len(w.order) >= MaxWarningCodes { + return + } + if w.seen == nil { + w.seen = make(map[string]struct{}, MaxWarningCodes) + } + if _, ok := w.seen[code]; ok { + return + } + w.seen[code] = struct{}{} + w.order = append(w.order, code) +} + +// codes returns the collected warnings sorted so acknowledgements for the same +// degradation are byte-identical across runs. +func (w *warningSet) codes() []string { + out := make([]string, len(w.order)) + copy(out, w.order) + sort.Strings(out) + return out +} diff --git a/internal/debuglog/request_test.go b/internal/debuglog/request_test.go new file mode 100644 index 0000000..ee9ca04 --- /dev/null +++ b/internal/debuglog/request_test.go @@ -0,0 +1,257 @@ +package debuglog + +import ( + "errors" + "net/http" + "strings" + "testing" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/types" +) + +var validationNow = time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC) + +// goldenRequest is the protocol-v1 request fixture both sides of the contract +// must accept without field translation. +func goldenRequest() types.DebugArchiveLogsRequestedMessage { + return types.DebugArchiveLogsRequestedMessage{ + ProtocolVersion: types.DebugArchiveProtocolVersion, + RequestID: "0f2f7c66-6f9d-4a05-9b4c-8d0f4f1e2a11", + ArchiveID: "3b3f6a1e-3f39-4a8d-9d7a-1f4c8f0a0b22", + CollectionID: "9c4f1a72-5d21-4a99-bb61-5b1d0f6a7c33", + RunID: "6c2c0f19-2c6f-4d1f-9f2a-0f6d1b3a4c44", + ExecutionID: "7d3d1f2a-3d7f-4e2f-8a3b-1f7e2c4b5d55", + RequestedFormat: types.DebugArchiveFormatNDJSON, + ExpiresAt: validationNow.Add(MaxRequestLifetime), + MaxBytes: ProtocolCeilingBytes, + ContentTransformer: types.ContentTransformerDescriptor{ + Kind: TransformerKindNoop, + Version: 1, + }, + UploadTarget: types.UploadTarget{ + URL: "https://storage.example.com/archives/candidate.ndjson?X-Signature=abc", + Method: http.MethodPut, + Headers: map[string]string{"Content-Type": types.DebugArchiveFormatNDJSON}, + }, + } +} + +func TestValidateRequestAcceptsGoldenFixture(t *testing.T) { + request := goldenRequest() + + effective, err := ValidateRequest(&request, validationNow, ProtocolCeilingBytes) + if err != nil { + t.Fatalf("golden protocol-v1 request was rejected: %v", err) + } + if effective != ProtocolCeilingBytes { + t.Fatalf("effective bound = %d, want %d", effective, ProtocolCeilingBytes) + } +} + +func TestValidateRequestUsesTheLowerOfRequestAndConfiguredBounds(t *testing.T) { + request := goldenRequest() + request.MaxBytes = 8 << 20 + + t.Run("request bound is lower", func(t *testing.T) { + effective, err := ValidateRequest(&request, validationNow, ProtocolCeilingBytes) + if err != nil { + t.Fatalf("ValidateRequest: %v", err) + } + if effective != 8<<20 { + t.Fatalf("effective bound = %d, want %d", effective, 8<<20) + } + }) + + t.Run("configured bound is lower", func(t *testing.T) { + effective, err := ValidateRequest(&request, validationNow, 1<<20) + if err != nil { + t.Fatalf("ValidateRequest: %v", err) + } + if effective != 1<<20 { + t.Fatalf("effective bound = %d, want %d", effective, 1<<20) + } + }) +} + +func TestValidateRequestIgnoresUnknownOptionalFields(t *testing.T) { + // json.Unmarshal into the request DTO discards fields this worker does not + // know, which is what lets a newer server add optional data safely. + request := goldenRequest() + if _, err := ValidateRequest(&request, validationNow, ProtocolCeilingBytes); err != nil { + t.Fatalf("ValidateRequest: %v", err) + } +} + +func TestValidateRequestRejections(t *testing.T) { + tests := []struct { + name string + mutate func(*types.DebugArchiveLogsRequestedMessage) + wantReason string + }{ + { + name: "unsupported protocol version", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { r.ProtocolVersion = 2 }, + wantReason: types.DebugArchiveReasonUnsupportedProtocolVersion, + }, + { + name: "missing request id", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { r.RequestID = " " }, + wantReason: types.DebugArchiveReasonInvalidRequest, + }, + { + name: "missing execution id", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { r.ExecutionID = "" }, + wantReason: types.DebugArchiveReasonInvalidRequest, + }, + { + name: "identifier with control characters", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { r.ArchiveID = "abc\ndef" }, + wantReason: types.DebugArchiveReasonInvalidRequest, + }, + { + name: "unsupported format", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { r.RequestedFormat = "application/json" }, + wantReason: types.DebugArchiveReasonInvalidRequest, + }, + { + name: "already expired", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { r.ExpiresAt = validationNow.Add(-time.Second) }, + wantReason: types.DebugArchiveReasonRequestExpired, + }, + { + name: "expiry beyond the protocol lifetime", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { + r.ExpiresAt = validationNow.Add(MaxRequestLifetime + time.Minute) + }, + wantReason: types.DebugArchiveReasonInvalidRequest, + }, + { + name: "non-positive bound", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { r.MaxBytes = 0 }, + wantReason: types.DebugArchiveReasonInvalidRequest, + }, + { + name: "bound above the protocol ceiling", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { r.MaxBytes = ProtocolCeilingBytes + 1 }, + wantReason: types.DebugArchiveReasonInvalidRequest, + }, + { + name: "unsupported content transformer", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { + r.ContentTransformer = types.ContentTransformerDescriptor{Kind: "redact", Version: 1} + }, + wantReason: types.DebugArchiveReasonUnsupportedContentTransformer, + }, + { + name: "unsupported transformer version", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { + r.ContentTransformer = types.ContentTransformerDescriptor{Kind: TransformerKindNoop, Version: 2} + }, + wantReason: types.DebugArchiveReasonUnsupportedContentTransformer, + }, + { + name: "unsupported upload method", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { r.UploadTarget.Method = http.MethodPatch }, + wantReason: types.DebugArchiveReasonInvalidRequest, + }, + { + name: "PUT target with multipart fields", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { + r.UploadTarget.MultipartFields = map[string]string{"key": "value"} + }, + wantReason: types.DebugArchiveReasonInvalidRequest, + }, + { + name: "POST target without multipart fields", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { + r.UploadTarget.Method = http.MethodPost + }, + wantReason: types.DebugArchiveReasonInvalidRequest, + }, + { + name: "plain HTTP target outside loopback", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { + r.UploadTarget.URL = "http://storage.example.com/candidate.ndjson" + }, + wantReason: types.DebugArchiveReasonInvalidRequest, + }, + { + name: "non-HTTP scheme", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { + r.UploadTarget.URL = "file:///etc/passwd" + }, + wantReason: types.DebugArchiveReasonInvalidRequest, + }, + { + name: "header with control characters", + mutate: func(r *types.DebugArchiveLogsRequestedMessage) { + r.UploadTarget.Headers = map[string]string{"X-Bad": "a\rb"} + }, + wantReason: types.DebugArchiveReasonInvalidRequest, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + request := goldenRequest() + tc.mutate(&request) + + _, err := ValidateRequest(&request, validationNow, ProtocolCeilingBytes) + var validation *ValidationError + if !errors.As(err, &validation) { + t.Fatalf("error = %v, want a *ValidationError", err) + } + if validation.ReasonCode != tc.wantReason { + t.Fatalf("reason = %q, want %q", validation.ReasonCode, tc.wantReason) + } + }) + } +} + +func TestValidateRequestAllowsLoopbackHTTPForLocalTesting(t *testing.T) { + request := goldenRequest() + request.UploadTarget.URL = "http://127.0.0.1:8080/candidate.ndjson" + + if _, err := ValidateRequest(&request, validationNow, ProtocolCeilingBytes); err != nil { + t.Fatalf("a loopback HTTP target must be accepted for local testing: %v", err) + } +} + +func TestSanitizeMessageBoundsAndStripsControlCharacters(t *testing.T) { + sanitized := SanitizeMessage("upload\nrejected\ttarget") + if strings.ContainsAny(sanitized, "\n\t") { + t.Fatalf("sanitized message retained control characters: %q", sanitized) + } + if sanitized != "upload rejected target" { + t.Fatalf("sanitized = %q, want %q", sanitized, "upload rejected target") + } + + long := SanitizeMessage(strings.Repeat("a", MaxAckMessageBytes*2)) + if len(long) != MaxAckMessageBytes { + t.Fatalf("sanitized length = %d, want %d", len(long), MaxAckMessageBytes) + } +} + +func TestWarningSetDeduplicatesSortsAndCaps(t *testing.T) { + var set warningSet + set.add(WarningOutputDropped) + set.add(WarningContainerLogsUnavailable) + set.add(WarningOutputDropped) + set.add("") + + codes := set.codes() + if len(codes) != 2 { + t.Fatalf("codes = %v, want two deduplicated entries", codes) + } + if codes[0] != WarningContainerLogsUnavailable || codes[1] != WarningOutputDropped { + t.Fatalf("codes = %v, want them sorted", codes) + } + + for i := 0; i < MaxWarningCodes*2; i++ { + set.add(strings.Repeat("w", i+1)) + } + if got := len(set.codes()); got != MaxWarningCodes { + t.Fatalf("codes = %d, want the %d cap", got, MaxWarningCodes) + } +} diff --git a/internal/debuglog/snapshot.go b/internal/debuglog/snapshot.go new file mode 100644 index 0000000..4444834 --- /dev/null +++ b/internal/debuglog/snapshot.go @@ -0,0 +1,134 @@ +package debuglog + +import ( + "crypto/sha256" + "encoding/base64" + "encoding/binary" + "encoding/hex" + "fmt" + "hash/crc32" + "io" + "os" + "sync" +) + +// Snapshot is one request's immutable, transformed NDJSON object on local +// disk. Keeping it on disk is what makes an upload retry byte-identical; it +// never extends an execution's retention and is deleted once the request +// finishes. +type Snapshot struct { + store *Store + encoder *Encoder + file *os.File + reserved int64 + + result FinalizeResult + crc32c string + sha256 string + closeOnce sync.Once +} + +// NewSnapshot allocates a request-scoped snapshot bounded by maxBytes. +func (s *Store) NewSnapshot(backend string, transformer ContentTransformer, maxBytes int64) (*Snapshot, error) { + if maxBytes <= 0 || maxBytes > s.config.MaxExecutionBytes { + maxBytes = s.config.MaxExecutionBytes + } + // The reservation covers the bounded spool plus the finalized object, + // which briefly coexist while Finalize streams one into the other. + reserved := maxBytes * 2 + if err := s.reserve(reserved); err != nil { + return nil, err + } + + encoder, err := NewEncoder(EncoderOptions{ + Backend: backend, + Transformer: transformer, + MaxBytes: maxBytes, + CreateFile: func(suffix string) (*os.File, error) { + return s.createFile("snapshot", suffix) + }, + }) + if err != nil { + s.release(reserved) + return nil, err + } + + file, err := s.createFile("snapshot", "object") + if err != nil { + _ = encoder.Close() + s.release(reserved) + return nil, err + } + + return &Snapshot{store: s, encoder: encoder, file: file, reserved: reserved}, nil +} + +// Sink is the destination a backend writes provider output into. +func (s *Snapshot) Sink() Sink { return s.encoder } + +// Finalize writes the bounded NDJSON object and computes its digests over the +// exact bytes that will be uploaded. +func (s *Snapshot) Finalize() error { + result, err := s.encoder.Finalize(s.file) + if err != nil { + return err + } + s.result = result + + if err := s.file.Sync(); err != nil { + return fmt.Errorf("debuglog: failed to flush snapshot: %w", err) + } + info, err := s.file.Stat() + if err != nil { + return fmt.Errorf("debuglog: failed to stat snapshot: %w", err) + } + if info.Size() != result.Bytes { + return fmt.Errorf("debuglog: snapshot size mismatch") + } + + if _, err := s.file.Seek(0, io.SeekStart); err != nil { + return fmt.Errorf("debuglog: failed to rewind snapshot: %w", err) + } + crcHash := crc32.New(crc32.MakeTable(crc32.Castagnoli)) + shaHash := sha256.New() + if _, err := io.Copy(io.MultiWriter(crcHash, shaHash), s.file); err != nil { + return fmt.Errorf("debuglog: failed to checksum snapshot: %w", err) + } + + crcBytes := make([]byte, 4) + binary.BigEndian.PutUint32(crcBytes, crcHash.Sum32()) + s.crc32c = base64.StdEncoding.EncodeToString(crcBytes) + s.sha256 = hex.EncodeToString(shaHash.Sum(nil)) + return nil +} + +// Bytes is the finalized object's exact size. +func (s *Snapshot) Bytes() int64 { return s.result.Bytes } + +// Truncated reports whether the object carries a truncation record. +func (s *Snapshot) Truncated() bool { return s.result.Truncated } + +// Warnings are the deduplicated warning codes collected while encoding. +func (s *Snapshot) Warnings() []string { return s.result.Warnings } + +// CRC32C is the Castagnoli checksum, base64-encoded to match object-store +// attribute formats. +func (s *Snapshot) CRC32C() string { return s.crc32c } + +// SHA256 is the hex-encoded digest of the finalized object. +func (s *Snapshot) SHA256() string { return s.sha256 } + +// Open returns an independent reader over the finalized object so every upload +// attempt replays exactly the same bytes. +func (s *Snapshot) Open() io.Reader { + return io.NewSectionReader(s.file, 0, s.result.Bytes) +} + +// Close deletes the snapshot's bytes and returns its share of the disk budget. +func (s *Snapshot) Close() { + s.closeOnce.Do(func() { + _ = s.encoder.Close() + _ = closeAndRemove(s.file) + s.store.release(s.reserved) + }) +} diff --git a/internal/debuglog/spool.go b/internal/debuglog/spool.go new file mode 100644 index 0000000..b65d24b --- /dev/null +++ b/internal/debuglog/spool.go @@ -0,0 +1,223 @@ +package debuglog + +import ( + "errors" + "fmt" + "io" + "os" + "sync" +) + +// ErrSpoolClosed is returned when a caller writes to a spool that has already +// been closed. +var ErrSpoolClosed = errors.New("debuglog: spool closed") + +// boundedSpool stores newline-terminated lines on disk under a fixed byte +// budget using a first/last retention policy: the head segment keeps the +// earliest lines and two rotating tail segments keep the most recent ones. +// Retention is always whole lines, so a reader never sees a partial record. +// +// Head-and-tail retention keeps both the early context of a failing execution +// (image pull, setup) and its terminal output, at the cost of an explicit gap +// that the caller reports as a truncation record. +// +// The budget is a hard ceiling on what the caller can finalize, not a target: +// headLimit + 2*tailLimit plus the reserved truncation record never exceeds +// maxBytes, so the object uploaded to a server-signed target can never overrun +// the size the request asked for. +type boundedSpool struct { + mu sync.Mutex + + head *os.File + tailOld *os.File + tailNew *os.File + + maxBytes int64 + headLimit int64 + tailLimit int64 + + headBytes int64 + tailOldBytes int64 + tailNewBytes int64 + omitted int64 + + closed bool +} + +// spoolWatermark is a consistent view of a spool's contents at one instant. +// A reader replays exactly these byte prefixes while writers keep appending. +type spoolWatermark struct { + headBytes int64 + tailOldBytes int64 + tailNewBytes int64 + omitted int64 +} + +// newBoundedSpool creates a spool whose finalized output never exceeds +// maxBytes. Files are created through create so callers control mode and +// naming. +func newBoundedSpool(maxBytes int64, create func(suffix string) (*os.File, error)) (*boundedSpool, error) { + if maxBytes <= 0 { + return nil, fmt.Errorf("debuglog: spool budget must be positive, got %d", maxBytes) + } + + head, err := create("head") + if err != nil { + return nil, err + } + tailOld, err := create("tail-a") + if err != nil { + _ = closeAndRemove(head) + return nil, err + } + tailNew, err := create("tail-b") + if err != nil { + _ = closeAndRemove(head) + _ = closeAndRemove(tailOld) + return nil, err + } + + // The truncation record is reserved up front rather than added on top at + // finalization, so emitting it can never push the object past maxBytes. + contentBudget := maxBytes - maxTruncationLineBytes + if contentBudget < 0 { + contentBudget = 0 + } + + // Half the content budget preserves the earliest output; the remaining + // half is split across two rotating tail segments so the newest output + // survives while each rotation drops at most a quarter. The limits are + // deliberately allowed to reach zero: a budget too small to hold a whole + // record retains nothing rather than overrunning. + headLimit := contentBudget / 2 + tailLimit := (contentBudget - headLimit) / 2 + + return &boundedSpool{ + head: head, + tailOld: tailOld, + tailNew: tailNew, + maxBytes: maxBytes, + headLimit: headLimit, + tailLimit: tailLimit, + }, nil +} + +// WriteLine appends one complete line, rotating tail segments when the budget +// is exhausted. +// +// A line that cannot fit a segment on its own is dropped and counted as +// omitted. Writing it anyway would produce an object larger than the request +// asked for, and truncating it would produce invalid NDJSON; dropping it keeps +// the stream parseable and the bound hard, and the truncation record reports +// the loss. +func (s *boundedSpool) WriteLine(line []byte) error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.closed { + return ErrSpoolClosed + } + + if s.headBytes+int64(len(line)) <= s.headLimit { + if _, err := s.head.Write(line); err != nil { + return err + } + s.headBytes += int64(len(line)) + return nil + } + + if int64(len(line)) > s.tailLimit { + s.omitted += int64(len(line)) + return nil + } + + if s.tailNewBytes+int64(len(line)) > s.tailLimit { + if err := s.rotateTailLocked(); err != nil { + return err + } + } + if _, err := s.tailNew.Write(line); err != nil { + return err + } + s.tailNewBytes += int64(len(line)) + return nil +} + +// rotateTailLocked retires the older tail segment and reuses its file for new +// output. The retired bytes become the reported omission lower bound. +func (s *boundedSpool) rotateTailLocked() error { + s.omitted += s.tailOldBytes + + retired := s.tailOld + if err := retired.Truncate(0); err != nil { + return err + } + if _, err := retired.Seek(0, io.SeekStart); err != nil { + return err + } + + s.tailOld = s.tailNew + s.tailOldBytes = s.tailNewBytes + s.tailNew = retired + s.tailNewBytes = 0 + return nil +} + +// Watermark captures the current retained extents so a snapshot reads a fixed +// prefix of each segment even while writers continue appending. +func (s *boundedSpool) Watermark() spoolWatermark { + s.mu.Lock() + defer s.mu.Unlock() + return spoolWatermark{ + headBytes: s.headBytes, + tailOldBytes: s.tailOldBytes, + tailNewBytes: s.tailNewBytes, + omitted: s.omitted, + } +} + +// segmentsAt returns readers over the exact byte extents recorded in mark, in +// chronological order. The returned readers are independent of the write +// offsets, so reading never disturbs a concurrent writer. +func (s *boundedSpool) segmentsAt(mark spoolWatermark) []io.Reader { + s.mu.Lock() + defer s.mu.Unlock() + + return []io.Reader{ + io.NewSectionReader(s.head, 0, mark.headBytes), + io.NewSectionReader(s.tailOld, 0, mark.tailOldBytes), + io.NewSectionReader(s.tailNew, 0, mark.tailNewBytes), + } +} + +// Close releases the spool's files and deletes them. +func (s *boundedSpool) Close() error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.closed { + return nil + } + s.closed = true + + var errs []error + for _, f := range []*os.File{s.head, s.tailOld, s.tailNew} { + if err := closeAndRemove(f); err != nil { + errs = append(errs, err) + } + } + return errors.Join(errs...) +} + +func closeAndRemove(f *os.File) error { + if f == nil { + return nil + } + name := f.Name() + closeErr := f.Close() + removeErr := os.Remove(name) + if removeErr != nil && os.IsNotExist(removeErr) { + removeErr = nil + } + return errors.Join(closeErr, removeErr) +} diff --git a/internal/debuglog/store.go b/internal/debuglog/store.go new file mode 100644 index 0000000..21396a1 --- /dev/null +++ b/internal/debuglog/store.go @@ -0,0 +1,225 @@ +package debuglog + +import ( + "crypto/rand" + "encoding/hex" + "fmt" + "os" + "path/filepath" + "sync" + "sync/atomic" +) + +const ( + // DefaultMaxTotalBytes is the process-local disk budget shared by direct + // captures and request snapshots. + DefaultMaxTotalBytes int64 = 1 << 30 + // DefaultMaxExecutionBytes is the per-execution retention bound. + DefaultMaxExecutionBytes = ProtocolCeilingBytes + // DefaultMaxConcurrentUploads bounds how many snapshot/upload handlers run + // at once per worker process. + DefaultMaxConcurrentUploads = 2 + + // captureDirName is the fixed sub-tree the worker owns beneath the + // configured (or temporary) capture root. + captureDirName = "oz-agent-worker/debug-logs" + + captureDirMode os.FileMode = 0o700 + captureFileMode os.FileMode = 0o600 +) + +// Config bounds debug-log capture. It deliberately carries no retention +// duration: retention comes exclusively from the execution's already-resolved +// idle-on-complete cleanup grace. +type Config struct { + // Directory overrides the capture root. Empty uses ${TMPDIR}. + Directory string + // MaxTotalBytes is the aggregate budget for captures plus request + // snapshots in this process. + MaxTotalBytes int64 + // MaxExecutionBytes bounds one execution's capture or snapshot. + MaxExecutionBytes int64 + // MaxConcurrentUploads bounds concurrent snapshot/upload handlers. + MaxConcurrentUploads int +} + +// DefaultConfig returns the built-in bounds. +func DefaultConfig() Config { + return Config{ + MaxTotalBytes: DefaultMaxTotalBytes, + MaxExecutionBytes: DefaultMaxExecutionBytes, + MaxConcurrentUploads: DefaultMaxConcurrentUploads, + } +} + +// WithDefaults fills unset bounds from DefaultConfig. Only a bound the operator +// actually set can fail validation, so an embedder that leaves the whole +// configuration zero still gets working capture. +func (c Config) WithDefaults() Config { + defaults := DefaultConfig() + if c.MaxTotalBytes == 0 { + c.MaxTotalBytes = defaults.MaxTotalBytes + } + if c.MaxExecutionBytes == 0 { + c.MaxExecutionBytes = defaults.MaxExecutionBytes + } + if c.MaxConcurrentUploads == 0 { + c.MaxConcurrentUploads = defaults.MaxConcurrentUploads + } + return c +} + +// Validate rejects bounds that cannot produce a usable capture. Callers treat +// a validation failure as a non-fatal loss of archive capture, never as a +// reason to refuse assigned task execution. +func (c Config) Validate() error { + if c.MaxExecutionBytes <= 0 { + return fmt.Errorf("debuglog: max_execution_bytes must be positive") + } + if c.MaxExecutionBytes > ProtocolCeilingBytes { + return fmt.Errorf("debuglog: max_execution_bytes must not exceed %d", ProtocolCeilingBytes) + } + if c.MaxTotalBytes <= 0 { + return fmt.Errorf("debuglog: max_total_bytes must be positive") + } + if c.MaxTotalBytes < c.MaxExecutionBytes { + return fmt.Errorf("debuglog: max_total_bytes must be at least max_execution_bytes") + } + if c.MaxConcurrentUploads <= 0 { + return fmt.Errorf("debuglog: max_concurrent_uploads must be positive") + } + return nil +} + +// Store owns the secure capture root and the process-local disk budget. It is +// the only component that creates files under that root. +type Store struct { + config Config + root string + + reserved atomic.Int64 + + mu sync.Mutex + nextIndex uint64 +} + +// NewStore validates the configuration, prepares the capture root, and removes +// files orphaned by a previous process. V1 does not reconstruct ownership +// across process replacement, so anything already present is unowned garbage. +func NewStore(config Config) (*Store, error) { + config = config.WithDefaults() + if err := config.Validate(); err != nil { + return nil, err + } + + base := config.Directory + if base == "" { + base = os.TempDir() + } + root := filepath.Join(base, captureDirName) + + if err := os.MkdirAll(root, captureDirMode); err != nil { + return nil, fmt.Errorf("debuglog: failed to create capture root: %w", err) + } + // MkdirAll honors the process umask and leaves an existing directory's + // mode alone, so tighten the leaf explicitly. + if err := os.Chmod(root, captureDirMode); err != nil { + return nil, fmt.Errorf("debuglog: failed to secure capture root: %w", err) + } + info, err := os.Lstat(root) + if err != nil { + return nil, fmt.Errorf("debuglog: failed to inspect capture root: %w", err) + } + if !info.IsDir() || info.Mode()&os.ModeSymlink != 0 { + return nil, fmt.Errorf("debuglog: capture root is not a regular directory") + } + + store := &Store{config: config, root: root} + if err := store.removeOrphans(); err != nil { + return nil, err + } + return store, nil +} + +// Config returns the validated bounds this store enforces. +func (s *Store) Config() Config { return s.config } + +// removeOrphans deletes every entry left in the capture root. Contents are +// never read, so nothing from a previous process is exposed. +func (s *Store) removeOrphans() error { + entries, err := os.ReadDir(s.root) + if err != nil { + return fmt.Errorf("debuglog: failed to scan capture root: %w", err) + } + for _, entry := range entries { + if err := os.RemoveAll(filepath.Join(s.root, entry.Name())); err != nil { + return fmt.Errorf("debuglog: failed to remove orphaned capture data: %w", err) + } + } + return nil +} + +// ErrBudgetExhausted reports that the process-local disk budget cannot admit +// another capture or snapshot. +var ErrBudgetExhausted = fmt.Errorf("debuglog: capture disk budget exhausted") + +// reserve claims bytes against the shared budget, or reports +// ErrBudgetExhausted when the budget is already committed. +func (s *Store) reserve(bytes int64) error { + for { + current := s.reserved.Load() + if current+bytes > s.config.MaxTotalBytes { + return ErrBudgetExhausted + } + if s.reserved.CompareAndSwap(current, current+bytes) { + return nil + } + } +} + +func (s *Store) release(bytes int64) { + s.reserved.Add(-bytes) +} + +// ReservedBytes reports the currently committed share of the disk budget. +func (s *Store) ReservedBytes() int64 { return s.reserved.Load() } + +// createFile allocates one file with a non-user-derived name under the capture +// root. Opening with O_EXCL rejects a pre-existing entry, including a symlink +// planted to redirect writes outside the root. +func (s *Store) createFile(prefix, suffix string) (*os.File, error) { + token, err := randomToken() + if err != nil { + return nil, err + } + + s.mu.Lock() + s.nextIndex++ + index := s.nextIndex + s.mu.Unlock() + + name := fmt.Sprintf("%s-%s-%d-%s", prefix, token, index, suffix) + if filepath.Base(name) != name { + return nil, fmt.Errorf("debuglog: refusing unsafe capture file name") + } + + path := filepath.Join(s.root, name) + file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_EXCL, captureFileMode) // #nosec G304 -- path is a random name the store generates beneath its own 0700 root. + if err != nil { + return nil, fmt.Errorf("debuglog: failed to create capture file: %w", err) + } + if err := file.Chmod(captureFileMode); err != nil { + _ = file.Close() + _ = os.Remove(path) + return nil, fmt.Errorf("debuglog: failed to secure capture file: %w", err) + } + return file, nil +} + +func randomToken() (string, error) { + buf := make([]byte, 16) + if _, err := rand.Read(buf); err != nil { + return "", fmt.Errorf("debuglog: failed to generate capture file name: %w", err) + } + return hex.EncodeToString(buf), nil +} diff --git a/internal/debuglog/store_test.go b/internal/debuglog/store_test.go new file mode 100644 index 0000000..c152d1f --- /dev/null +++ b/internal/debuglog/store_test.go @@ -0,0 +1,390 @@ +package debuglog + +import ( + "errors" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +func newTestStore(t *testing.T, mutate func(*Config)) *Store { + t.Helper() + config := DefaultConfig() + config.Directory = t.TempDir() + config.MaxExecutionBytes = 1 << 16 + config.MaxTotalBytes = 1 << 20 + if mutate != nil { + mutate(&config) + } + store, err := NewStore(config) + if err != nil { + t.Fatalf("NewStore: %v", err) + } + return store +} + +func TestConfigValidateRejectsUnusableBounds(t *testing.T) { + tests := map[string]Config{ + "non-positive execution bound": {MaxExecutionBytes: 0, MaxTotalBytes: 1 << 20, MaxConcurrentUploads: 2}, + "execution bound above the protocol ceiling": { + MaxExecutionBytes: ProtocolCeilingBytes + 1, + MaxTotalBytes: ProtocolCeilingBytes * 2, + MaxConcurrentUploads: 2, + }, + "total below the execution bound": {MaxExecutionBytes: 1 << 20, MaxTotalBytes: 1 << 10, MaxConcurrentUploads: 2}, + "non-positive concurrency": {MaxExecutionBytes: 1 << 20, MaxTotalBytes: 1 << 21, MaxConcurrentUploads: -1}, + } + + for name, config := range tests { + t.Run(name, func(t *testing.T) { + if err := config.Validate(); err == nil { + t.Fatal("expected the configuration to be rejected") + } + }) + } +} + +func TestConfigWithDefaultsFillsUnsetBounds(t *testing.T) { + filled := Config{}.WithDefaults() + if filled.MaxTotalBytes != DefaultMaxTotalBytes || + filled.MaxExecutionBytes != DefaultMaxExecutionBytes || + filled.MaxConcurrentUploads != DefaultMaxConcurrentUploads { + t.Fatalf("defaults were not applied: %+v", filled) + } + if err := filled.Validate(); err != nil { + t.Fatalf("the default configuration must validate: %v", err) + } +} + +func TestNewStoreCreatesASecureRoot(t *testing.T) { + base := t.TempDir() + config := DefaultConfig() + config.Directory = base + store, err := NewStore(config) + if err != nil { + t.Fatalf("NewStore: %v", err) + } + + info, err := os.Stat(store.root) + if err != nil { + t.Fatalf("capture root was not created: %v", err) + } + if perm := info.Mode().Perm(); perm != captureDirMode { + t.Fatalf("capture root mode = %o, want %o", perm, captureDirMode) + } + if !strings.HasPrefix(store.root, base) { + t.Fatalf("capture root %q escaped its configured base %q", store.root, base) + } +} + +func TestNewStoreRemovesOrphansFromAPreviousProcess(t *testing.T) { + base := t.TempDir() + root := filepath.Join(base, captureDirName) + if err := os.MkdirAll(root, captureDirMode); err != nil { + t.Fatalf("failed to seed capture root: %v", err) + } + orphan := filepath.Join(root, "capture-deadbeef-1-head") + if err := os.WriteFile(orphan, []byte("stale bytes"), captureFileMode); err != nil { + t.Fatalf("failed to seed orphan: %v", err) + } + + config := DefaultConfig() + config.Directory = base + if _, err := NewStore(config); err != nil { + t.Fatalf("NewStore: %v", err) + } + + if _, err := os.Stat(orphan); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("orphaned capture file survived startup: %v", err) + } +} + +func TestStoreCreatesFilesWithNonUserDerivedNamesAndTightMode(t *testing.T) { + store := newTestStore(t, nil) + + file, err := store.createFile("capture", "head") + if err != nil { + t.Fatalf("createFile: %v", err) + } + defer func() { _ = closeAndRemove(file) }() + + info, err := file.Stat() + if err != nil { + t.Fatalf("Stat: %v", err) + } + if perm := info.Mode().Perm(); perm != captureFileMode { + t.Fatalf("capture file mode = %o, want %o", perm, captureFileMode) + } + if filepath.Dir(file.Name()) != store.root { + t.Fatalf("capture file %q was created outside the capture root", file.Name()) + } +} + +func TestStoreRefusesToOverwriteAPlantedSymlink(t *testing.T) { + store := newTestStore(t, nil) + + // A predictable name is what a symlink attack needs; the store's random + // names plus O_EXCL are what defeat it. Recreate the exact name the store + // would use and prove it refuses rather than following the link. + target := filepath.Join(t.TempDir(), "outside") + victim := filepath.Join(store.root, "planted") + if err := os.Symlink(target, victim); err != nil { + t.Skipf("symlinks are unavailable on this platform: %v", err) + } + + if _, err := os.OpenFile(victim, os.O_RDWR|os.O_CREATE|os.O_EXCL, captureFileMode); err == nil { + t.Fatal("expected O_EXCL to refuse an existing symlink") + } + if _, err := os.Stat(target); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("the symlink target was written through: %v", err) + } +} + +func TestStoreBudgetIsSharedAndReleased(t *testing.T) { + store := newTestStore(t, func(c *Config) { + c.MaxExecutionBytes = 1 << 16 + c.MaxTotalBytes = 1 << 17 + }) + + first, err := store.NewTaskLogCapture(nil) + if err != nil { + t.Fatalf("first capture: %v", err) + } + second, err := store.NewTaskLogCapture(nil) + if err != nil { + t.Fatalf("second capture: %v", err) + } + if _, err := store.NewTaskLogCapture(nil); !errors.Is(err, ErrBudgetExhausted) { + t.Fatalf("third capture error = %v, want %v", err, ErrBudgetExhausted) + } + + first.Close() + third, err := store.NewTaskLogCapture(nil) + if err != nil { + t.Fatalf("a capture must be admitted after budget is released: %v", err) + } + + second.Close() + third.Close() + if reserved := store.ReservedBytes(); reserved != 0 { + t.Fatalf("reserved bytes = %d after closing every capture, want 0", reserved) + } +} + +// collectingSink records what a capture replays without touching disk. +type collectingSink struct { + chunks []Chunk + omitted int64 +} + +func (s *collectingSink) WriteChunk(chunk Chunk) error { + copied := chunk + copied.Data = append([]byte(nil), chunk.Data...) + s.chunks = append(s.chunks, copied) + return nil +} + +func (s *collectingSink) WriteSourceError(Phase, Stream, SourceIdentity, string) error { return nil } + +func (s *collectingSink) NoteOmittedBytes(n int64) { s.omitted += n } + +func (s *collectingSink) text() string { + var builder strings.Builder + for _, chunk := range s.chunks { + builder.Write(chunk.Data) + } + return builder.String() +} + +func TestTaskLogCaptureLabelsPhaseAndStream(t *testing.T) { + store := newTestStore(t, nil) + capture, err := store.NewTaskLogCapture(nil) + if err != nil { + t.Fatalf("NewTaskLogCapture: %v", err) + } + defer capture.Close() + + if _, err := capture.Writer(PhaseSetup, StreamStdout).Write([]byte("setup-out\n")); err != nil { + t.Fatalf("setup write: %v", err) + } + if _, err := capture.Writer(PhaseAgent, StreamStderr).Write([]byte("agent-err\n")); err != nil { + t.Fatalf("agent write: %v", err) + } + if _, err := capture.Writer(PhaseTeardown, StreamStdout).Write([]byte("teardown-out\n")); err != nil { + t.Fatalf("teardown write: %v", err) + } + capture.Finalize(2 * time.Second) + + sink := &collectingSink{} + if err := capture.SnapshotTo(sink); err != nil { + t.Fatalf("SnapshotTo: %v", err) + } + + if len(sink.chunks) != 3 { + t.Fatalf("chunk count = %d, want 3", len(sink.chunks)) + } + want := []struct { + phase Phase + stream Stream + data string + }{ + {PhaseSetup, StreamStdout, "setup-out\n"}, + {PhaseAgent, StreamStderr, "agent-err\n"}, + {PhaseTeardown, StreamStdout, "teardown-out\n"}, + } + for i, expected := range want { + got := sink.chunks[i] + if got.Phase != expected.phase || got.Stream != expected.stream || string(got.Data) != expected.data { + t.Errorf("chunk %d = {%s %s %q}, want {%s %s %q}", + i, got.Phase, got.Stream, got.Data, expected.phase, expected.stream, expected.data) + } + if got.ObservedAt.IsZero() { + t.Errorf("chunk %d has no worker observation time", i) + } + } +} + +func TestTaskLogCaptureWriteAlwaysReportsTheChildByteCount(t *testing.T) { + store := newTestStore(t, nil) + capture, err := store.NewTaskLogCapture(nil) + if err != nil { + t.Fatalf("NewTaskLogCapture: %v", err) + } + defer capture.Close() + + writer := capture.Writer(PhaseAgent, StreamStdout) + payload := []byte(strings.Repeat("z", 4096)) + // Far more writes than the archive queue can hold, so drops are certain. + for i := 0; i < captureQueueDepth*4; i++ { + n, err := writer.Write(payload) + if err != nil { + t.Fatalf("write %d returned an error, which would surface on the child's pipe: %v", i, err) + } + if n != len(payload) { + t.Fatalf("write %d reported %d bytes, want %d", i, n, len(payload)) + } + } +} + +func TestTaskLogCaptureSnapshotUsesAFixedWatermark(t *testing.T) { + store := newTestStore(t, nil) + capture, err := store.NewTaskLogCapture(nil) + if err != nil { + t.Fatalf("NewTaskLogCapture: %v", err) + } + defer capture.Close() + + writer := capture.Writer(PhaseAgent, StreamStdout) + if _, err := writer.Write([]byte("before-watermark\n")); err != nil { + t.Fatalf("write: %v", err) + } + capture.Finalize(2 * time.Second) + + first := &collectingSink{} + if err := capture.SnapshotTo(first); err != nil { + t.Fatalf("first SnapshotTo: %v", err) + } + if !strings.Contains(first.text(), "before-watermark") { + t.Fatalf("first snapshot = %q, want the pre-watermark output", first.text()) + } + if strings.Contains(first.text(), "after-watermark") { + t.Fatal("first snapshot leaked output written after its watermark") + } + + if _, err := writer.Write([]byte("after-watermark\n")); err != nil { + t.Fatalf("write: %v", err) + } + capture.Finalize(2 * time.Second) + + second := &collectingSink{} + if err := capture.SnapshotTo(second); err != nil { + t.Fatalf("second SnapshotTo: %v", err) + } + if !strings.Contains(second.text(), "after-watermark") { + t.Fatalf("second snapshot = %q, want the later output", second.text()) + } +} + +func TestTaskLogCaptureIsPerExecution(t *testing.T) { + store := newTestStore(t, nil) + first, err := store.NewTaskLogCapture(nil) + if err != nil { + t.Fatalf("first capture: %v", err) + } + defer first.Close() + second, err := store.NewTaskLogCapture(nil) + if err != nil { + t.Fatalf("second capture: %v", err) + } + defer second.Close() + + if _, err := first.Writer(PhaseAgent, StreamStdout).Write([]byte("task-one-output\n")); err != nil { + t.Fatalf("write: %v", err) + } + first.Finalize(2 * time.Second) + second.Finalize(2 * time.Second) + + sink := &collectingSink{} + if err := second.SnapshotTo(sink); err != nil { + t.Fatalf("SnapshotTo: %v", err) + } + if strings.Contains(sink.text(), "task-one-output") { + t.Fatalf("one task's output appeared in another's snapshot: %q", sink.text()) + } +} + +func TestTaskLogCaptureCloseIsIdempotent(t *testing.T) { + store := newTestStore(t, nil) + capture, err := store.NewTaskLogCapture(nil) + if err != nil { + t.Fatalf("NewTaskLogCapture: %v", err) + } + + capture.Close() + capture.Close() + + if reserved := store.ReservedBytes(); reserved != 0 { + t.Fatalf("reserved bytes = %d after a double close, want 0", reserved) + } +} + +func TestSnapshotComputesChecksumsOverTheUploadedBytes(t *testing.T) { + store := newTestStore(t, nil) + snapshot, err := store.NewSnapshot(BackendDirect, noopTransformer{}, 1<<15) + if err != nil { + t.Fatalf("NewSnapshot: %v", err) + } + defer snapshot.Close() + + if err := snapshot.Sink().WriteChunk(Chunk{Phase: PhaseAgent, Stream: StreamStdout, Data: []byte("payload")}); err != nil { + t.Fatalf("WriteChunk: %v", err) + } + if err := snapshot.Finalize(); err != nil { + t.Fatalf("Finalize: %v", err) + } + + if snapshot.Bytes() == 0 { + t.Fatal("expected the snapshot to carry bytes") + } + if snapshot.CRC32C() == "" || snapshot.SHA256() == "" { + t.Fatalf("expected both digests, got crc32c=%q sha256=%q", snapshot.CRC32C(), snapshot.SHA256()) + } + + // Every attempt must replay identical bytes for a retry to be safe. + firstBytes, err := readAll(snapshot) + if err != nil { + t.Fatalf("first read: %v", err) + } + secondBytes, err := readAll(snapshot) + if err != nil { + t.Fatalf("second read: %v", err) + } + if string(firstBytes) != string(secondBytes) { + t.Fatal("re-opening the snapshot produced different bytes") + } + if int64(len(firstBytes)) != snapshot.Bytes() { + t.Fatalf("read %d bytes, want the reported %d", len(firstBytes), snapshot.Bytes()) + } +} diff --git a/internal/debuglog/transformer.go b/internal/debuglog/transformer.go new file mode 100644 index 0000000..08c2aea --- /dev/null +++ b/internal/debuglog/transformer.go @@ -0,0 +1,49 @@ +package debuglog + +import "fmt" + +// TransformerKindNoop is the only content-transformer descriptor V1 accepts. +const TransformerKindNoop = "noop" + +// ContentTransformer rewrites log message data while a snapshot is encoded. +// Applying it at encode time means the first object uploaded to cloud storage +// already carries the transformed bytes, so no later raw-to-transformed copy +// is required. V1 ships only a byte-preserving implementation; real redaction +// rules arrive as a new descriptor kind or version. +type ContentTransformer interface { + // Kind is the descriptor kind this transformer implements. + Kind() string + // Version is the descriptor version reported in the acknowledgement. + Version() int + // Transform rewrites one decoded data chunk. It must not be given + // timestamps, sequence numbers, identity fields, or warning codes. + Transform(data []byte) ([]byte, error) +} + +// ErrUnsupportedTransformer reports a descriptor this worker cannot honor. +// The coordinator maps it to failed/unsupported_content_transformer and +// uploads nothing, so untransformed data can never reach the target. +type ErrUnsupportedTransformer struct { + Kind string + Version int +} + +func (e *ErrUnsupportedTransformer) Error() string { + return fmt.Sprintf("debuglog: unsupported content transformer %q version %d", e.Kind, e.Version) +} + +type noopTransformer struct{} + +func (noopTransformer) Kind() string { return TransformerKindNoop } + +func (noopTransformer) Version() int { return 1 } + +func (noopTransformer) Transform(data []byte) ([]byte, error) { return data, nil } + +// NewTransformer resolves a descriptor to its implementation. +func NewTransformer(kind string, version int) (ContentTransformer, error) { + if kind == TransformerKindNoop && version == 1 { + return noopTransformer{}, nil + } + return nil, &ErrUnsupportedTransformer{Kind: kind, Version: version} +} diff --git a/internal/debuglog/upload.go b/internal/debuglog/upload.go new file mode 100644 index 0000000..8853534 --- /dev/null +++ b/internal/debuglog/upload.go @@ -0,0 +1,243 @@ +package debuglog + +import ( + "context" + "errors" + "fmt" + "io" + "mime/multipart" + "net/http" + "sort" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/types" +) + +const ( + // uploadFilePartName is the multipart field name for the snapshot object. + // Presigned POST policies conventionally sign the file part last, under + // this name. + uploadFilePartName = "file" + // uploadFileName is the filename reported in the multipart file part. + uploadFileName = "worker-logs.ndjson" + + uploadInitialBackoff = 500 * time.Millisecond + uploadMaxBackoff = 15 * time.Second + uploadBackoffRate = 2.0 + // uploadResponseDrainLimit bounds how much of a response body is read + // before the connection is reused. The bytes are discarded, never logged. + uploadResponseDrainLimit = 4 << 10 +) + +// UploadError reports a terminal upload failure with the stable reason code +// the acknowledgement carries. It deliberately excludes the target, the signed +// query, request headers, and the response body. +type UploadError struct { + ReasonCode string + Detail string +} + +func (e *UploadError) Error() string { + return fmt.Sprintf("debuglog: %s (%s)", e.Detail, e.ReasonCode) +} + +// Uploader sends a finalized snapshot to the server-supplied target. +type Uploader struct { + client *http.Client + now func() time.Time + // sleep waits between retries; tests substitute it to avoid real delays. + sleep func(ctx context.Context, d time.Duration) error +} + +// NewUploader builds an uploader whose HTTP client refuses redirects, so a +// signed target can never bounce the snapshot, its headers, or its body to +// another host. +func NewUploader(client *http.Client, now func() time.Time) *Uploader { + if client == nil { + client = &http.Client{Timeout: MaxRequestLifetime} + } + client.CheckRedirect = func(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse + } + if now == nil { + now = time.Now + } + return &Uploader{client: client, now: now, sleep: sleepContext} +} + +func sleepContext(ctx context.Context, d time.Duration) error { + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +// Upload sends the snapshot, retrying transient failures with bounded +// exponential backoff until expiresAt. Every attempt reopens the immutable +// local object so a retry is byte-identical. +func (u *Uploader) Upload(ctx context.Context, target types.UploadTarget, snapshot *Snapshot, expiresAt time.Time) error { + backoff := uploadInitialBackoff + var lastDetail string + + for { + if !u.now().Before(expiresAt) { + return &UploadError{ + ReasonCode: types.DebugArchiveReasonUploadExpired, + Detail: "upload target expired before the snapshot was accepted", + } + } + + status, err := u.attempt(ctx, target, snapshot) + var terminal *UploadError + switch { + case err == nil && status >= 200 && status < 300: + return nil + case err == nil && isRetryableStatus(status): + lastDetail = fmt.Sprintf("upload target returned a retryable status (%d)", status) + case err == nil: + return &UploadError{ + ReasonCode: types.DebugArchiveReasonUploadRejected, + Detail: fmt.Sprintf("upload target rejected the snapshot (%d)", status), + } + case errors.As(err, &terminal): + // The attempt already classified this as unrecoverable, such as a + // target that tried to redirect the snapshot elsewhere. + return terminal + case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded): + return &UploadError{ + ReasonCode: types.DebugArchiveReasonWorkerShuttingDown, + Detail: "upload cancelled before completion", + } + default: + lastDetail = "upload transport error" + } + + // Waiting past the deadline is pointless: the target is already + // unusable, so report expiry instead of burning another attempt. + if !u.now().Add(backoff).Before(expiresAt) { + return &UploadError{ + ReasonCode: types.DebugArchiveReasonUploadFailed, + Detail: lastDetail, + } + } + if err := u.sleep(ctx, backoff); err != nil { + return &UploadError{ + ReasonCode: types.DebugArchiveReasonWorkerShuttingDown, + Detail: "upload cancelled before completion", + } + } + backoff = min(time.Duration(float64(backoff)*uploadBackoffRate), uploadMaxBackoff) + } +} + +func isRetryableStatus(status int) bool { + return status == http.StatusRequestTimeout || status == http.StatusTooManyRequests || status >= 500 +} + +// attempt performs one upload and returns the response status. The response +// body is drained under a small cap and discarded without being read into an +// error or acknowledgement. +func (u *Uploader) attempt(ctx context.Context, target types.UploadTarget, snapshot *Snapshot) (int, error) { + req, err := u.buildRequest(ctx, target, snapshot) + if err != nil { + return 0, err + } + + resp, err := u.client.Do(req) + if err != nil { + return 0, err + } + defer func() { + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, uploadResponseDrainLimit)) + _ = resp.Body.Close() + }() + + // CheckRedirect surfaces redirects as ordinary responses, so a 3xx here + // means the target tried to send the snapshot elsewhere. It never gets a + // second request carrying the signed headers or body. + if resp.StatusCode >= 300 && resp.StatusCode < 400 { + return 0, &UploadError{ + ReasonCode: types.DebugArchiveReasonUploadRejected, + Detail: "upload target attempted a redirect", + } + } + return resp.StatusCode, nil +} + +func (u *Uploader) buildRequest(ctx context.Context, target types.UploadTarget, snapshot *Snapshot) (*http.Request, error) { + if target.Method == http.MethodPost { + return u.buildMultipartRequest(ctx, target, snapshot) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPut, target.URL, io.NopCloser(snapshot.Open())) + if err != nil { + return nil, err + } + req.ContentLength = snapshot.Bytes() + applyTargetHeaders(req, target.Headers) + if req.Header.Get("Content-Type") == "" { + req.Header.Set("Content-Type", types.DebugArchiveFormatNDJSON) + } + return req, nil +} + +// buildMultipartRequest streams the form so the snapshot is never held in +// memory. The file part is written last, matching presigned POST policies. +func (u *Uploader) buildMultipartRequest(ctx context.Context, target types.UploadTarget, snapshot *Snapshot) (*http.Request, error) { + reader, writer := io.Pipe() + form := multipart.NewWriter(writer) + + go func() { + err := func() error { + for _, field := range sortedPairs(target.MultipartFields) { + if err := form.WriteField(field.key, field.value); err != nil { + return err + } + } + part, err := form.CreateFormFile(uploadFilePartName, uploadFileName) + if err != nil { + return err + } + if _, err := io.Copy(part, snapshot.Open()); err != nil { + return err + } + return form.Close() + }() + _ = writer.CloseWithError(err) + }() + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, target.URL, reader) + if err != nil { + _ = reader.Close() + return nil, err + } + applyTargetHeaders(req, target.Headers) + req.Header.Set("Content-Type", form.FormDataContentType()) + return req, nil +} + +func applyTargetHeaders(req *http.Request, headers map[string]string) { + for _, header := range sortedPairs(headers) { + req.Header.Set(header.key, header.value) + } +} + +type keyValue struct { + key string + value string +} + +// sortedPairs orders a target's headers and fields by key so repeated attempts +// build byte-identical requests instead of depending on map iteration order. +func sortedPairs(values map[string]string) []keyValue { + pairs := make([]keyValue, 0, len(values)) + for key, value := range values { + pairs = append(pairs, keyValue{key: key, value: value}) + } + sort.Slice(pairs, func(i, j int) bool { return pairs[i].key < pairs[j].key }) + return pairs +} diff --git a/internal/debuglog/upload_test.go b/internal/debuglog/upload_test.go new file mode 100644 index 0000000..3c65fa2 --- /dev/null +++ b/internal/debuglog/upload_test.go @@ -0,0 +1,321 @@ +package debuglog + +import ( + "context" + "errors" + "io" + "mime" + "mime/multipart" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/types" +) + +func readAll(snapshot *Snapshot) ([]byte, error) { + return io.ReadAll(snapshot.Open()) +} + +func newTestSnapshot(t *testing.T, payload string) *Snapshot { + t.Helper() + store := newTestStore(t, nil) + snapshot, err := store.NewSnapshot(BackendDirect, noopTransformer{}, 1<<15) + if err != nil { + t.Fatalf("NewSnapshot: %v", err) + } + t.Cleanup(snapshot.Close) + + if err := snapshot.Sink().WriteChunk(Chunk{Phase: PhaseAgent, Stream: StreamStdout, Data: []byte(payload)}); err != nil { + t.Fatalf("WriteChunk: %v", err) + } + if err := snapshot.Finalize(); err != nil { + t.Fatalf("Finalize: %v", err) + } + return snapshot +} + +// newTestUploader returns an uploader whose retry backoff is instant, so the +// retry policy is exercised without slowing the suite. +func newTestUploader(now func() time.Time) *Uploader { + uploader := NewUploader(&http.Client{Timeout: 5 * time.Second}, now) + uploader.sleep = func(ctx context.Context, _ time.Duration) error { + return ctx.Err() + } + return uploader +} + +func TestUploadPutSendsExactlyTheSnapshotWithSuppliedHeaders(t *testing.T) { + snapshot := newTestSnapshot(t, "put payload") + want, err := readAll(snapshot) + if err != nil { + t.Fatalf("readAll: %v", err) + } + + var gotBody []byte + var gotHeader, gotMethod string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotMethod = r.Method + gotHeader = r.Header.Get("X-Provider-Token") + gotBody, _ = io.ReadAll(r.Body) + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + uploader := newTestUploader(time.Now) + target := types.UploadTarget{ + URL: server.URL, + Method: http.MethodPut, + Headers: map[string]string{"X-Provider-Token": "signed"}, + } + if err := uploader.Upload(context.Background(), target, snapshot, time.Now().Add(time.Minute)); err != nil { + t.Fatalf("Upload: %v", err) + } + + if gotMethod != http.MethodPut { + t.Errorf("method = %q, want PUT", gotMethod) + } + if gotHeader != "signed" { + t.Errorf("X-Provider-Token = %q, want %q", gotHeader, "signed") + } + if string(gotBody) != string(want) { + t.Errorf("uploaded body did not match the snapshot bytes") + } +} + +func TestUploadPostStreamsMultipartFieldsAndOneFilePart(t *testing.T) { + snapshot := newTestSnapshot(t, "post payload") + want, err := readAll(snapshot) + if err != nil { + t.Fatalf("readAll: %v", err) + } + + gotFields := map[string]string{} + var gotFile []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, params, parseErr := mime.ParseMediaType(r.Header.Get("Content-Type")) + if parseErr != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + reader := multipart.NewReader(r.Body, params["boundary"]) + for { + part, partErr := reader.NextPart() + if partErr != nil { + break + } + body, _ := io.ReadAll(part) + if part.FileName() != "" { + gotFile = body + } else { + gotFields[part.FormName()] = string(body) + } + } + w.WriteHeader(http.StatusNoContent) + })) + defer server.Close() + + uploader := newTestUploader(time.Now) + target := types.UploadTarget{ + URL: server.URL, + Method: http.MethodPost, + MultipartFields: map[string]string{"key": "archives/candidate.ndjson", "policy": "signed-policy"}, + } + if err := uploader.Upload(context.Background(), target, snapshot, time.Now().Add(time.Minute)); err != nil { + t.Fatalf("Upload: %v", err) + } + + if gotFields["key"] != "archives/candidate.ndjson" || gotFields["policy"] != "signed-policy" { + t.Errorf("multipart fields = %v, want the supplied provider fields", gotFields) + } + if string(gotFile) != string(want) { + t.Error("multipart file part did not match the snapshot bytes") + } +} + +func TestUploadRetriesTransientStatusesThenSucceeds(t *testing.T) { + snapshot := newTestSnapshot(t, "retry payload") + + var attempts atomic.Int32 + var bodies []string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + bodies = append(bodies, string(body)) + switch attempts.Add(1) { + case 1: + w.WriteHeader(http.StatusRequestTimeout) + case 2: + w.WriteHeader(http.StatusTooManyRequests) + case 3: + w.WriteHeader(http.StatusBadGateway) + default: + w.WriteHeader(http.StatusOK) + } + })) + defer server.Close() + + uploader := NewUploader(&http.Client{Timeout: 5 * time.Second}, time.Now) + uploader.sleep = func(context.Context, time.Duration) error { return nil } + + target := types.UploadTarget{URL: server.URL, Method: http.MethodPut} + if err := uploader.Upload(context.Background(), target, snapshot, time.Now().Add(time.Minute)); err != nil { + t.Fatalf("Upload: %v", err) + } + if attempts.Load() != 4 { + t.Fatalf("attempts = %d, want 4", attempts.Load()) + } + for i, body := range bodies { + if body != bodies[0] { + t.Fatalf("attempt %d replayed different bytes than the first attempt", i) + } + } +} + +func TestUploadTreatsOther4xxAsTerminalRejection(t *testing.T) { + snapshot := newTestSnapshot(t, "rejected payload") + + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + attempts.Add(1) + w.WriteHeader(http.StatusForbidden) + })) + defer server.Close() + + uploader := newTestUploader(time.Now) + target := types.UploadTarget{URL: server.URL, Method: http.MethodPut} + err := uploader.Upload(context.Background(), target, snapshot, time.Now().Add(time.Minute)) + + var uploadErr *UploadError + if !errors.As(err, &uploadErr) { + t.Fatalf("error = %v, want an *UploadError", err) + } + if uploadErr.ReasonCode != types.DebugArchiveReasonUploadRejected { + t.Fatalf("reason = %q, want %q", uploadErr.ReasonCode, types.DebugArchiveReasonUploadRejected) + } + if attempts.Load() != 1 { + t.Fatalf("attempts = %d, want no retry after a terminal rejection", attempts.Load()) + } +} + +func TestUploadRejectsRedirectsWithoutForwardingSignedMaterial(t *testing.T) { + snapshot := newTestSnapshot(t, "redirect payload") + + var redirectTargetHits atomic.Int32 + redirectTarget := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + redirectTargetHits.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer redirectTarget.Close() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, redirectTarget.URL, http.StatusTemporaryRedirect) + })) + defer server.Close() + + uploader := newTestUploader(time.Now) + target := types.UploadTarget{ + URL: server.URL, + Method: http.MethodPut, + Headers: map[string]string{"X-Provider-Token": "signed"}, + } + err := uploader.Upload(context.Background(), target, snapshot, time.Now().Add(time.Minute)) + + var uploadErr *UploadError + if !errors.As(err, &uploadErr) { + t.Fatalf("error = %v, want an *UploadError", err) + } + if uploadErr.ReasonCode != types.DebugArchiveReasonUploadRejected { + t.Fatalf("reason = %q, want %q", uploadErr.ReasonCode, types.DebugArchiveReasonUploadRejected) + } + if redirectTargetHits.Load() != 0 { + t.Fatal("the redirect destination received the snapshot") + } + if strings.Contains(uploadErr.Detail, server.URL) { + t.Fatalf("upload error leaked the target URL: %q", uploadErr.Detail) + } +} + +func TestUploadDoesNotStartAfterExpiry(t *testing.T) { + snapshot := newTestSnapshot(t, "expired payload") + + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + attempts.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + now := time.Now() + uploader := newTestUploader(func() time.Time { return now }) + target := types.UploadTarget{URL: server.URL, Method: http.MethodPut} + err := uploader.Upload(context.Background(), target, snapshot, now.Add(-time.Second)) + + var uploadErr *UploadError + if !errors.As(err, &uploadErr) { + t.Fatalf("error = %v, want an *UploadError", err) + } + if uploadErr.ReasonCode != types.DebugArchiveReasonUploadExpired { + t.Fatalf("reason = %q, want %q", uploadErr.ReasonCode, types.DebugArchiveReasonUploadExpired) + } + if attempts.Load() != 0 { + t.Fatal("an expired request must not contact the target") + } +} + +func TestUploadDoesNotRetryPastExpiry(t *testing.T) { + snapshot := newTestSnapshot(t, "late payload") + + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + attempts.Add(1) + w.WriteHeader(http.StatusServiceUnavailable) + })) + defer server.Close() + + now := time.Now() + uploader := newTestUploader(func() time.Time { return now }) + target := types.UploadTarget{URL: server.URL, Method: http.MethodPut} + // The deadline is closer than the first backoff, so no retry may start. + err := uploader.Upload(context.Background(), target, snapshot, now.Add(uploadInitialBackoff/2)) + + var uploadErr *UploadError + if !errors.As(err, &uploadErr) { + t.Fatalf("error = %v, want an *UploadError", err) + } + if uploadErr.ReasonCode != types.DebugArchiveReasonUploadFailed { + t.Fatalf("reason = %q, want %q", uploadErr.ReasonCode, types.DebugArchiveReasonUploadFailed) + } + if attempts.Load() != 1 { + t.Fatalf("attempts = %d, want exactly one before the deadline", attempts.Load()) + } +} + +func TestUploadRetriesTransportErrors(t *testing.T) { + snapshot := newTestSnapshot(t, "transport payload") + + // A closed server produces a connection error rather than a status. + server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + url := server.URL + server.Close() + + var sleeps atomic.Int32 + uploader := NewUploader(&http.Client{Timeout: time.Second}, time.Now) + uploader.sleep = func(context.Context, time.Duration) error { + if sleeps.Add(1) >= 2 { + return context.Canceled + } + return nil + } + + target := types.UploadTarget{URL: url, Method: http.MethodPut} + err := uploader.Upload(context.Background(), target, snapshot, time.Now().Add(time.Minute)) + if err == nil { + t.Fatal("expected an upload failure against a closed target") + } + if sleeps.Load() < 2 { + t.Fatalf("backoff attempts = %d, want the transport error to be retried", sleeps.Load()) + } +} diff --git a/internal/metrics/debug_archive.go b/internal/metrics/debug_archive.go new file mode 100644 index 0000000..97d1c3b --- /dev/null +++ b/internal/metrics/debug_archive.go @@ -0,0 +1,190 @@ +package metrics + +import ( + "context" + "time" + + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" +) + +// Debug-archive metrics deliberately carry no run, execution, or request +// identifier as a label. Those identifiers belong on task span events, where +// they cannot blow up metric cardinality. + +// Ownership states reported alongside a debug-archive request outcome. +const ( + DebugArchiveOwnershipActive = "active" + DebugArchiveOwnershipCleanupGrace = "cleanup_grace" +) + +// debugArchiveInstruments holds the debug-archive metric set. It is built and +// swapped alongside the main instrument set. +type debugArchiveInstruments struct { + requests metric.Int64Counter + snapshotDuration metric.Float64Histogram + snapshotBytes metric.Int64Histogram + uploads metric.Int64Counter + uploadDuration metric.Float64Histogram + captureBytes metric.Int64Gauge + cleanupGraceCount metric.Int64Gauge + cleanupResults metric.Int64Counter + requestsInFlight metric.Int64UpDownCounter + truncations metric.Int64Counter +} + +func buildDebugArchiveInstruments(m metric.Meter) (*debugArchiveInstruments, error) { + requests, err := m.Int64Counter( + "oz_worker_debug_archive_requests_total", + metric.WithDescription("Debug-archive log requests this worker owned, labeled by backend, outcome, and stable reason."), + ) + if err != nil { + return nil, err + } + snapshotDuration, err := m.Float64Histogram( + "oz_worker_debug_archive_snapshot_duration_seconds", + metric.WithDescription("Wall-clock duration of a debug-archive log snapshot, labeled by backend."), + metric.WithUnit("s"), + ) + if err != nil { + return nil, err + } + snapshotBytes, err := m.Int64Histogram( + "oz_worker_debug_archive_snapshot_bytes", + metric.WithDescription("Size of the NDJSON object a debug-archive snapshot produced, labeled by backend."), + metric.WithUnit("By"), + ) + if err != nil { + return nil, err + } + uploads, err := m.Int64Counter( + "oz_worker_debug_archive_uploads_total", + metric.WithDescription("Debug-archive snapshot uploads, labeled by result."), + ) + if err != nil { + return nil, err + } + uploadDuration, err := m.Float64Histogram( + "oz_worker_debug_archive_upload_duration_seconds", + metric.WithDescription("Wall-clock duration of a debug-archive snapshot upload, labeled by result."), + metric.WithUnit("s"), + ) + if err != nil { + return nil, err + } + captureBytes, err := m.Int64Gauge( + "oz_worker_debug_archive_capture_bytes", + metric.WithDescription("Bytes currently reserved by direct-execution captures and request snapshots."), + metric.WithUnit("By"), + ) + if err != nil { + return nil, err + } + cleanupGraceCount, err := m.Int64Gauge( + "oz_worker_cleanup_grace_entries", + metric.WithDescription("Executions this worker is retaining through their idle-on-complete cleanup grace."), + ) + if err != nil { + return nil, err + } + cleanupResults, err := m.Int64Counter( + "oz_worker_cleanup_grace_results_total", + metric.WithDescription("Backend resource cleanups performed at cleanup-grace expiry, labeled by backend and result."), + ) + if err != nil { + return nil, err + } + requestsInFlight, err := m.Int64UpDownCounter( + "oz_worker_debug_archive_requests_in_flight", + metric.WithDescription("Debug-archive log requests currently being snapshotted or uploaded."), + ) + if err != nil { + return nil, err + } + truncations, err := m.Int64Counter( + "oz_worker_debug_archive_truncations_total", + metric.WithDescription("Debug-archive snapshots that dropped bytes to stay within their bound, labeled by backend."), + ) + if err != nil { + return nil, err + } + + return &debugArchiveInstruments{ + requests: requests, + snapshotDuration: snapshotDuration, + snapshotBytes: snapshotBytes, + uploads: uploads, + uploadDuration: uploadDuration, + captureBytes: captureBytes, + cleanupGraceCount: cleanupGraceCount, + cleanupResults: cleanupResults, + requestsInFlight: requestsInFlight, + truncations: truncations, + }, nil +} + +// RecordDebugArchiveRequest records one owned request's terminal outcome. +func RecordDebugArchiveRequest(backend, ownership, outcome, reasonCode string) { + current().debugArchive.requests.Add(context.Background(), 1, + metric.WithAttributes( + attribute.String("backend", backend), + attribute.String("ownership", ownership), + attribute.String("outcome", outcome), + attribute.String("reason", reasonCode), + ), + ) +} + +// RecordDebugArchiveSnapshot records the cost of producing one snapshot. +func RecordDebugArchiveSnapshot(backend string, duration time.Duration, bytes int64) { + attrs := metric.WithAttributes(attribute.String("backend", backend)) + current().debugArchive.snapshotDuration.Record(context.Background(), duration.Seconds(), attrs) + current().debugArchive.snapshotBytes.Record(context.Background(), bytes, attrs) +} + +// RecordDebugArchiveTruncation records a snapshot that had to drop bytes. +func RecordDebugArchiveTruncation(backend string) { + current().debugArchive.truncations.Add(context.Background(), 1, + metric.WithAttributes(attribute.String("backend", backend)), + ) +} + +// RecordDebugArchiveUpload records one upload attempt's terminal result. +func RecordDebugArchiveUpload(result string, duration time.Duration) { + attrs := metric.WithAttributes(attribute.String("result", result)) + current().debugArchive.uploads.Add(context.Background(), 1, attrs) + current().debugArchive.uploadDuration.Record(context.Background(), duration.Seconds(), attrs) +} + +// SetDebugArchiveCaptureBytes records the disk budget currently committed to +// captures and request snapshots. +func SetDebugArchiveCaptureBytes(bytes int64) { + current().debugArchive.captureBytes.Record(context.Background(), bytes) +} + +// SetCleanupGraceEntries records how many executions are being retained for +// post-terminal log retrieval. +func SetCleanupGraceEntries(count int) { + current().debugArchive.cleanupGraceCount.Record(context.Background(), int64(count)) +} + +// RecordCleanupGraceResult records a backend resource cleanup performed when an +// execution's cleanup grace expired. +func RecordCleanupGraceResult(backend, result string) { + current().debugArchive.cleanupResults.Add(context.Background(), 1, + metric.WithAttributes( + attribute.String("backend", backend), + attribute.String("result", result), + ), + ) +} + +// IncDebugArchiveRequestsInFlight marks a request as entering snapshot/upload. +func IncDebugArchiveRequestsInFlight() { + current().debugArchive.requestsInFlight.Add(context.Background(), 1) +} + +// DecDebugArchiveRequestsInFlight marks a request as leaving snapshot/upload. +func DecDebugArchiveRequestsInFlight() { + current().debugArchive.requestsInFlight.Add(context.Background(), -1) +} diff --git a/internal/metrics/metrics.go b/internal/metrics/metrics.go index 36147ba..af2743a 100644 --- a/internal/metrics/metrics.go +++ b/internal/metrics/metrics.go @@ -77,6 +77,7 @@ type instruments struct { taskFailures metric.Int64Counter wsReconnects metric.Int64Counter workerInfo metric.Int64Gauge + debugArchive *debugArchiveInstruments } // activeInstruments is the current instrument set. It always points to a @@ -325,6 +326,9 @@ func primeInstruments(ctx context.Context, set *instruments) { metric.WithAttributes(attribute.String("reason", r)), ) } + set.debugArchive.captureBytes.Record(ctx, 0) + set.debugArchive.cleanupGraceCount.Record(ctx, 0) + set.debugArchive.requestsInFlight.Add(ctx, 0) } func newResource(ctx context.Context, cfg Config) (*resource.Resource, error) { @@ -419,6 +423,10 @@ func buildInstruments(m metric.Meter) (*instruments, error) { if err != nil { return nil, err } + debugArchive, err := buildDebugArchiveInstruments(m) + if err != nil { + return nil, err + } return &instruments{ connected: connected, tasksActive: tasksActive, @@ -430,6 +438,7 @@ func buildInstruments(m metric.Meter) (*instruments, error) { taskFailures: taskFailures, wsReconnects: wsReconnects, workerInfo: workerInfo, + debugArchive: debugArchive, }, nil } diff --git a/internal/types/messages.go b/internal/types/messages.go index 5f59927..0814d19 100644 --- a/internal/types/messages.go +++ b/internal/types/messages.go @@ -16,8 +16,19 @@ const ( MessageTypeTaskRejected MessageType = "task_rejected" MessageTypeTaskCancellation MessageType = "task_cancellation" MessageTypeHeartbeat MessageType = "heartbeat" + // MessageTypeDebugArchiveLogsRequested is sent from server to worker to ask + // the process that executed an assignment for a bounded log snapshot. + MessageTypeDebugArchiveLogsRequested MessageType = "debug_archive_logs_requested" + // MessageTypeDebugArchiveLogsUploaded is the owning worker's acknowledgement + // of a debug-archive log request. + MessageTypeDebugArchiveLogsUploaded MessageType = "debug_archive_logs_uploaded" ) +// WorkerVersionHeader carries the worker's build-time version on every +// authenticated WebSocket dial so warp-server can snapshot the exact build +// that claims an execution. +const WorkerVersionHeader = "X-Warp-Worker-Version" + // WebSocketMessage is the base structure for all WebSocket messages type WebSocketMessage struct { Type MessageType `json:"type"` @@ -101,6 +112,114 @@ type TaskCancellationMessage struct { TaskID string `json:"task_id"` } +// DebugArchiveProtocolVersion is the only worker-log protocol version this +// worker implements. A request carrying any other version is refused with +// DebugArchiveReasonUnsupportedProtocolVersion. +const DebugArchiveProtocolVersion = 1 + +// DebugArchiveFormatNDJSON is the only snapshot encoding this worker produces. +const DebugArchiveFormatNDJSON = "application/x-ndjson" + +// Debug-archive acknowledgement outcomes. +const ( + DebugArchiveOutcomeUploaded = "uploaded" + DebugArchiveOutcomeUnavailable = "unavailable" + DebugArchiveOutcomeFailed = "failed" +) + +// Debug-archive capture statuses, reported only alongside an uploaded outcome. +const ( + DebugArchiveCaptureComplete = "complete" + DebugArchiveCapturePartial = "partial" +) + +// Debug-archive reason codes. These are the complete, stable set warp-server +// keys off for non-upload outcomes. +const ( + DebugArchiveReasonUnsupportedProtocolVersion = "unsupported_protocol_version" + DebugArchiveReasonUnsupportedContentTransformer = "unsupported_content_transformer" + DebugArchiveReasonInvalidRequest = "invalid_request" + DebugArchiveReasonRequestExpired = "request_expired" + DebugArchiveReasonBackendNotSupported = "backend_not_supported" + DebugArchiveReasonResourceNotReady = "resource_not_ready" + DebugArchiveReasonResourceNotFound = "resource_not_found" + DebugArchiveReasonCleanupGraceExpired = "cleanup_grace_expired" + DebugArchiveReasonCaptureUnavailable = "capture_unavailable" + DebugArchiveReasonSnapshotFailed = "snapshot_failed" + DebugArchiveReasonUploadRejected = "upload_rejected" + DebugArchiveReasonUploadExpired = "upload_expired" + DebugArchiveReasonUploadFailed = "upload_failed" + DebugArchiveReasonWorkerShuttingDown = "worker_shutting_down" + DebugArchiveReasonRequestCapacityExhausted = "request_capacity_exhausted" +) + +// Debug-archive warning codes. A partial capture reports these so warp-server +// can mark a source partial without parsing provider text. +const ( + DebugArchiveWarningContainerLogsUnavailable = "container_logs_unavailable" + DebugArchiveWarningPreviousLogsUnavailable = "previous_logs_unavailable" + DebugArchiveWarningOutputDropped = "output_dropped" + DebugArchiveWarningProviderSnapshotIncomplete = "provider_snapshot_incomplete" +) + +// ContentTransformerDescriptor names the versioned transform applied to log +// message data while the snapshot is encoded. V1 defines only the byte +// preserving "noop" transformer. +type ContentTransformerDescriptor struct { + Kind string `json:"kind"` + Version int `json:"version"` +} + +// UploadTarget is the provider-neutral destination warp-server signs for one +// immutable snapshot object. Its URL, headers, and multipart fields are +// credential-bearing and must never be logged or echoed in an acknowledgement. +type UploadTarget struct { + URL string `json:"url"` + Method string `json:"method"` + Headers map[string]string `json:"headers,omitempty"` + MultipartFields map[string]string `json:"multipart_fields,omitempty"` +} + +// DebugArchiveLogsRequestedMessage is the server's request for a bounded log +// snapshot of one exact execution. Unknown fields are ignored so a newer +// server can add optional data without breaking this worker. +type DebugArchiveLogsRequestedMessage struct { + ProtocolVersion int `json:"protocol_version"` + RequestID string `json:"request_id"` + ArchiveID string `json:"archive_id"` + CollectionID string `json:"collection_id"` + RunID string `json:"run_id"` + ExecutionID string `json:"execution_id"` + RequestedFormat string `json:"requested_format"` + ExpiresAt time.Time `json:"expires_at"` + MaxBytes int64 `json:"max_bytes"` + ContentTransformer ContentTransformerDescriptor `json:"content_transformer"` + UploadTarget UploadTarget `json:"upload_target"` +} + +// DebugArchiveLogsUploadedMessage is the owning worker's acknowledgement. Byte, +// checksum, and capture fields are present only for an uploaded outcome; the +// reason code and sanitized message describe every other outcome. +type DebugArchiveLogsUploadedMessage struct { + ProtocolVersion int `json:"protocol_version"` + RequestID string `json:"request_id"` + ArchiveID string `json:"archive_id"` + CollectionID string `json:"collection_id"` + RunID string `json:"run_id"` + ExecutionID string `json:"execution_id"` + Outcome string `json:"outcome"` + BackendKind string `json:"backend_kind,omitempty"` + Bytes int64 `json:"bytes,omitempty"` + CRC32C string `json:"crc32c,omitempty"` + SHA256 string `json:"sha256,omitempty"` + Truncated bool `json:"truncated"` + ContentTransformerVersion int `json:"content_transformer_version,omitempty"` + CaptureStatus string `json:"capture_status,omitempty"` + WarningCodes []string `json:"warning_codes"` + ReasonCode string `json:"reason_code"` + Message string `json:"message"` +} + // TaskState is the serialized terminal task state accepted by warp-server. type TaskState string diff --git a/internal/worker/backend.go b/internal/worker/backend.go index 2120ec8..9f656e7 100644 --- a/internal/worker/backend.go +++ b/internal/worker/backend.go @@ -3,6 +3,7 @@ package worker import ( "context" + "github.com/warpdotdev/oz-agent-worker/internal/debuglog" "github.com/warpdotdev/oz-agent-worker/internal/metrics" "github.com/warpdotdev/oz-agent-worker/internal/types" ) @@ -85,6 +86,12 @@ type TaskParams struct { // Containerized backends apply it as CPU/memory limits (Docker) or resource // requests/limits (Kubernetes). Backends that cannot enforce a shape (direct) ignore it. InstanceShape *types.InstanceShape + + // LogCapture, when non-nil, receives the task's stdout and stderr for a + // later debug-archive snapshot. Only the direct backend uses it: Docker and + // Kubernetes read provider-native logs on demand instead of keeping a + // second copy. + LogCapture *debuglog.TaskLogCapture } // Backend defines the interface for task execution backends. @@ -100,10 +107,35 @@ type Backend interface { // PreservesTasksOnShutdown reports whether active task execution units can // safely outlive the worker process during shutdown. PreservesTasksOnShutdown() bool + // SnapshotTaskLogs writes protocol-v1 NDJSON records for one exact + // (task, execution) pair's stdout and stderr into sink, covering the + // sandbox entrypoint/launcher and the client process it starts. + // + // It streams rather than buffering, makes no lifecycle change, removes no + // provider resource, and is safe to call while ExecuteTask is running. + // A backend that wrote valid data for some sources but could not read + // others returns *debuglog.PartialSnapshotError so the caller can upload + // what it has; other failures return *debuglog.SnapshotError. + SnapshotTaskLogs(ctx context.Context, params *SnapshotParams) error + // CleanupTaskResources idempotently releases the backend resources an + // execution retained past terminal state for log retrieval. The worker + // calls it when the execution's cleanup grace expires. + CleanupTaskResources(ctx context.Context, params *CancelParams) error // Shutdown cleans up backend resources. Shutdown(ctx context.Context) } +// SnapshotParams identifies the execution whose logs to snapshot and where to +// write them. +type SnapshotParams struct { + TaskID string + ExecutionID string + // Sink owns record framing, chunk bounds, content transformation, encoding + // selection, sequencing, and truncation. Backends supply provider output + // and whatever stream and container identity the provider actually reports. + Sink debuglog.Sink +} + // CancelParams carries the minimal, non-secret identifiers a backend needs to // cancel a task. It deliberately excludes env/secrets so the worker need not // retain secrets for the lifetime of a spawned task. diff --git a/internal/worker/backend_testing_test.go b/internal/worker/backend_testing_test.go new file mode 100644 index 0000000..4239be4 --- /dev/null +++ b/internal/worker/backend_testing_test.go @@ -0,0 +1,27 @@ +package worker + +import ( + "context" + + "github.com/warpdotdev/oz-agent-worker/internal/debuglog" +) + +// noLogSnapshotBackend satisfies the debug-archive half of the Backend +// contract for fakes whose test does not exercise log collection. +type noLogSnapshotBackend struct{} + +func (noLogSnapshotBackend) SnapshotTaskLogs(context.Context, *SnapshotParams) error { + return debuglog.ErrBackendNotSupported +} + +func (noLogSnapshotBackend) CleanupTaskResources(context.Context, *CancelParams) error { return nil } + +// registryWith builds a task registry pre-populated with active tasks, so a +// test can start a worker mid-execution without going through assignment. +func registryWith(tasks map[string]activeTask) *TaskRegistry { + registry := newTaskRegistry() + for taskID, task := range tasks { + registry.StartTask(taskID, task, debuglog.BackendDocker, nil) + } + return registry +} diff --git a/internal/worker/cleanup_durability_test.go b/internal/worker/cleanup_durability_test.go new file mode 100644 index 0000000..a969edf --- /dev/null +++ b/internal/worker/cleanup_durability_test.go @@ -0,0 +1,287 @@ +package worker + +import ( + "context" + "errors" + "path/filepath" + "sync/atomic" + "testing" + "time" + + "github.com/moby/moby/client" + "github.com/warpdotdev/oz-agent-worker/internal/debuglog" + "github.com/warpdotdev/oz-agent-worker/internal/types" + batchv1 "k8s.io/api/batch/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/runtime/schema" + "k8s.io/client-go/kubernetes/fake" + k8stesting "k8s.io/client-go/testing" +) + +// A resource whose deletion was not confirmed must stay registered. Dropping +// the identifier first leaves nothing to retry, so a transient API failure +// silently becomes a resource that outlives its cleanup grace. +func TestKubernetesCleanupRetainsTheJobWhenDeletionFails(t *testing.T) { + client := fake.NewSimpleClientset() + var deletes atomic.Int32 + client.PrependReactor("delete", "jobs", func(k8stesting.Action) (bool, runtime.Object, error) { + if deletes.Add(1) == 1 { + return true, nil, errors.New("etcdserver: request timed out") + } + return false, nil, nil + }) + + backend := &KubernetesBackend{ + config: KubernetesBackendConfig{WorkerID: "worker-123", Namespace: "agents"}, + clientset: client, + jobs: make(map[executionKey]*retainedJob), + } + if _, err := client.BatchV1().Jobs("agents").Create(context.Background(), &batchv1.Job{ + ObjectMeta: metav1.ObjectMeta{Name: "oz-task-task-1-exec-ution-1", Namespace: "agents"}, + }, metav1.CreateOptions{}); err != nil { + t.Fatalf("failed to seed Job: %v", err) + } + backend.registerJob("task-1", "execution-1", "oz-task-task-1-exec-ution-1") + backend.setJobDeleteAtCleanup("task-1", "execution-1", true) + + params := &CancelParams{TaskID: "task-1", ExecutionID: "execution-1"} + if err := backend.CleanupTaskResources(context.Background(), params); err == nil { + t.Fatal("expected the failed deletion to be surfaced") + } + if !backend.ownsJob("task-1", "execution-1") { + t.Fatal("the Job was forgotten despite an unconfirmed deletion, so nothing can retry it") + } + + // The retry succeeds, and only then is the identifier released. + if err := backend.CleanupTaskResources(context.Background(), params); err != nil { + t.Fatalf("retry: %v", err) + } + if backend.ownsJob("task-1", "execution-1") { + t.Fatal("the Job should be forgotten once deletion is confirmed") + } + + jobs, err := client.BatchV1().Jobs("agents").List(context.Background(), metav1.ListOptions{}) + if err != nil { + t.Fatalf("failed to list jobs: %v", err) + } + if len(jobs.Items) != 0 { + t.Fatalf("expected the Job to be deleted on retry, got %d", len(jobs.Items)) + } +} + +func TestKubernetesCleanupTreatsAnAbsentJobAsDeleted(t *testing.T) { + client := fake.NewSimpleClientset() + client.PrependReactor("delete", "jobs", func(k8stesting.Action) (bool, runtime.Object, error) { + return true, nil, apierrors.NewNotFound(schema.GroupResource{Group: "batch", Resource: "jobs"}, "oz-task-task-1-exec-ution-1") + }) + + backend := &KubernetesBackend{ + config: KubernetesBackendConfig{WorkerID: "worker-123", Namespace: "agents"}, + clientset: client, + jobs: make(map[executionKey]*retainedJob), + } + backend.registerJob("task-1", "execution-1", "oz-task-task-1-exec-ution-1") + backend.setJobDeleteAtCleanup("task-1", "execution-1", true) + + if err := backend.CleanupTaskResources(context.Background(), &CancelParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + }); err != nil { + t.Fatalf("an already-absent Job must count as deleted: %v", err) + } + if backend.ownsJob("task-1", "execution-1") { + t.Fatal("an already-absent Job should be forgotten") + } +} + +func TestKubernetesCleanupReleasesAFailedJobWithoutDeletingIt(t *testing.T) { + // A failed Job is deliberately left for the TTL controller, so cleanup has + // nothing to confirm and must still release its registry entry. + client := fake.NewSimpleClientset() + client.PrependReactor("delete", "jobs", func(k8stesting.Action) (bool, runtime.Object, error) { + t.Error("a failed Job must not be deleted at the cleanup-grace deadline") + return true, nil, nil + }) + + backend := &KubernetesBackend{ + config: KubernetesBackendConfig{WorkerID: "worker-123", Namespace: "agents"}, + clientset: client, + jobs: make(map[executionKey]*retainedJob), + } + backend.registerJob("task-1", "execution-1", "oz-task-task-1-exec-ution-1") + backend.setJobDeleteAtCleanup("task-1", "execution-1", false) + + if err := backend.CleanupTaskResources(context.Background(), &CancelParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + }); err != nil { + t.Fatalf("CleanupTaskResources: %v", err) + } + if backend.ownsJob("task-1", "execution-1") { + t.Fatal("a TTL-owned Job should still be released from the registry") + } +} + +// failingCleanupBackend refuses cleanup a fixed number of times so a test can +// exercise the shutdown sweep's retry and its give-up path. +type failingCleanupBackend struct { + noLogSnapshotBackend + failures atomic.Int32 + attempts atomic.Int32 + // delay stalls each attempt so a test can prove one slow entry does not + // starve another. + delay time.Duration +} + +func (b *failingCleanupBackend) ExecuteTask(context.Context, *TaskParams) ExecuteResult { + return executeCompleted() +} + +func (b *failingCleanupBackend) CancelTask(context.Context, *CancelParams) error { return nil } + +func (b *failingCleanupBackend) PreservesTasksOnShutdown() bool { return false } + +func (b *failingCleanupBackend) Shutdown(context.Context) {} + +func (b *failingCleanupBackend) CleanupTaskResources(ctx context.Context, _ *CancelParams) error { + b.attempts.Add(1) + if b.delay > 0 { + select { + case <-time.After(b.delay): + case <-ctx.Done(): + return ctx.Err() + } + } + if b.failures.Load() > 0 { + b.failures.Add(-1) + return errors.New("transient backend failure") + } + return nil +} + +func TestShutdownRetriesATransientCleanupFailure(t *testing.T) { + backend := &failingCleanupBackend{} + backend.failures.Store(2) + w := newDebugArchiveWorker(t, backend, "60m") + + w.tasks.StartTask("task-1", activeTask{cancel: func() {}, executionID: "exec-1"}, debuglog.BackendKubernetes, nil) + w.beginCleanupGrace(&types.TaskAssignmentMessage{ + TaskID: "task-1", + ExecutionID: "exec-1", + Task: &types.Task{ID: "task-1"}, + }) + w.tasks.Delete("task-1") + + w.Shutdown() + + if got := backend.attempts.Load(); got != 3 { + t.Fatalf("cleanup attempts = %d, want the transient failures retried up to %d", got, shutdownCleanupAttempts) + } + if backend.failures.Load() != 0 { + t.Fatal("expected the retry to eventually succeed") + } +} + +func TestShutdownGivesUpAfterTheRetryBudget(t *testing.T) { + backend := &failingCleanupBackend{} + backend.failures.Store(100) + w := newDebugArchiveWorker(t, backend, "60m") + + w.tasks.StartTask("task-1", activeTask{cancel: func() {}, executionID: "exec-1"}, debuglog.BackendDocker, nil) + w.beginCleanupGrace(&types.TaskAssignmentMessage{ + TaskID: "task-1", + ExecutionID: "exec-1", + Task: &types.Task{ID: "task-1"}, + }) + w.tasks.Delete("task-1") + + w.Shutdown() + + if got := backend.attempts.Load(); got != shutdownCleanupAttempts { + t.Fatalf("cleanup attempts = %d, want the bounded %d", got, shutdownCleanupAttempts) + } +} + +func TestShutdownGivesEveryEntryItsOwnBudget(t *testing.T) { + // A shared budget let one slow backend call consume the whole allowance + // and starve the entries behind it. Each entry now gets its own. + backend := &failingCleanupBackend{delay: 300 * time.Millisecond} + w := newDebugArchiveWorker(t, backend, "60m") + + const entries = 8 + for i := 0; i < entries; i++ { + taskID := "task-" + string(rune('a'+i)) + w.tasks.StartTask(taskID, activeTask{cancel: func() {}, executionID: "exec-1"}, debuglog.BackendDocker, nil) + w.beginCleanupGrace(&types.TaskAssignmentMessage{ + TaskID: taskID, + ExecutionID: "exec-1", + Task: &types.Task{ID: taskID}, + }) + w.tasks.Delete(taskID) + } + + w.Shutdown() + + if got := backend.attempts.Load(); got != entries { + t.Fatalf("cleanup attempts = %d, want every one of the %d entries attempted", got, entries) + } +} + +// unreachableDockerBackend points at a socket that does not exist, so every +// daemon call fails the way a transient outage would. +func unreachableDockerBackend(t *testing.T) *DockerBackend { + t.Helper() + dockerClient, err := client.New(client.WithHost("unix://" + filepath.Join(t.TempDir(), "absent.sock"))) + if err != nil { + t.Fatalf("failed to build a Docker client: %v", err) + } + t.Cleanup(func() { _ = dockerClient.Close() }) + return &DockerBackend{dockerClient: dockerClient, containers: make(map[executionKey]string)} +} + +func TestDockerCleanupRetainsTheContainerWhenRemovalFails(t *testing.T) { + backend := unreachableDockerBackend(t) + backend.registerContainer("task-1", "execution-1", "container-abc") + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + err := backend.CleanupTaskResources(ctx, &CancelParams{TaskID: "task-1", ExecutionID: "execution-1"}) + if err == nil { + t.Fatal("expected the failed removal to be surfaced rather than swallowed") + } + if _, ok := backend.lookupContainer("task-1", "execution-1"); !ok { + t.Fatal("the container was forgotten despite an unconfirmed removal, so nothing can retry it") + } +} + +func TestDockerCleanupReleasesTheContainerWhenCleanupIsDisabled(t *testing.T) { + // With cleanup disabled the container is intentionally left running, so + // there is nothing to confirm and the registry entry must still be freed. + backend := unreachableDockerBackend(t) + backend.config.NoCleanup = true + backend.registerContainer("task-1", "execution-1", "container-abc") + + if err := backend.CleanupTaskResources(context.Background(), &CancelParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + }); err != nil { + t.Fatalf("CleanupTaskResources: %v", err) + } + if _, ok := backend.lookupContainer("task-1", "execution-1"); ok { + t.Fatal("the registry entry should be released when cleanup is disabled") + } +} + +func TestDockerCleanupIsANoOpForAnUnknownExecution(t *testing.T) { + backend := &DockerBackend{containers: make(map[executionKey]string)} + + if err := backend.CleanupTaskResources(context.Background(), &CancelParams{ + TaskID: "unknown-task", + ExecutionID: "unknown-execution", + }); err != nil { + t.Fatalf("CleanupTaskResources: %v", err) + } +} diff --git a/internal/worker/command.go b/internal/worker/command.go index 39758e6..de185b5 100644 --- a/internal/worker/command.go +++ b/internal/worker/command.go @@ -9,6 +9,7 @@ import ( "os/exec" "time" + "github.com/warpdotdev/oz-agent-worker/internal/debuglog" "github.com/warpdotdev/oz-agent-worker/internal/log" "github.com/warpdotdev/oz-agent-worker/internal/metrics" ) @@ -136,6 +137,17 @@ func (b *CommandBackend) CancelTask(ctx context.Context, params *CancelParams) e // remote runtime, independent of this worker, so worker shutdown must not cancel them. func (b *CommandBackend) PreservesTasksOnShutdown() bool { return true } +// SnapshotTaskLogs is unsupported. The dispatch command hands the task to an +// opaque operator-owned runtime with no log API, and the command's own stdout +// is dispatch output, not the remote agent's execution log. +func (b *CommandBackend) SnapshotTaskLogs(context.Context, *SnapshotParams) error { + return debuglog.ErrBackendNotSupported +} + +// CleanupTaskResources has nothing to release: the backend retains no local +// resource for a dispatched task. +func (b *CommandBackend) CleanupTaskResources(context.Context, *CancelParams) error { return nil } + // Shutdown has nothing to clean up; the backend owns no local resources. func (b *CommandBackend) Shutdown(ctx context.Context) { log.Debugf(ctx, "Command backend shutdown") diff --git a/internal/worker/command_integration_test.go b/internal/worker/command_integration_test.go index d7e5b8c..c321e6e 100644 --- a/internal/worker/command_integration_test.go +++ b/internal/worker/command_integration_test.go @@ -25,9 +25,7 @@ func waitFor(t *testing.T, timeout time.Duration, cond func() bool) { } func (w *Worker) activeTaskCount() int { - w.tasksMutex.Lock() - defer w.tasksMutex.Unlock() - return len(w.activeTasks) + return w.tasks.Len() } func drainMessages(t *testing.T, ch <-chan []byte) []types.WebSocketMessage { @@ -92,9 +90,7 @@ func TestIntegrationCommandBackendDispatchSuppressesTerminalMessage(t *testing.T return err == nil }) waitFor(t, 5*time.Second, func() bool { - w.tasksMutex.Lock() - defer w.tasksMutex.Unlock() - task, ok := w.activeTasks["task-1"] + task, ok := w.tasks.Get("task-1") return ok && task.spawned }) diff --git a/internal/worker/debug_archive_test.go b/internal/worker/debug_archive_test.go new file mode 100644 index 0000000..c352b96 --- /dev/null +++ b/internal/worker/debug_archive_test.go @@ -0,0 +1,488 @@ +package worker + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/warpdotdev/oz-agent-worker/internal/debuglog" + "github.com/warpdotdev/oz-agent-worker/internal/types" + "go.opentelemetry.io/otel/trace" +) + +func TestWorkerVersionHeaderValue(t *testing.T) { + tests := []struct { + name string + version string + wantOK bool + }{ + {name: "release build", version: "v2026-08-04-15-14-28", wantOK: true}, + {name: "local dev build", version: "dev", wantOK: true}, + {name: "arbitrary test build", version: "test-build-1234", wantOK: true}, + {name: "empty", version: "", wantOK: false}, + {name: "overlong", version: strings.Repeat("v", maxWorkerVersionBytes+1), wantOK: false}, + {name: "carries a newline", version: "v1\ninjected: header", wantOK: false}, + {name: "carries a NUL", version: "v1\x00", wantOK: false}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got, ok := workerVersionHeaderValue(tc.version) + if ok != tc.wantOK { + t.Fatalf("ok = %v, want %v", ok, tc.wantOK) + } + if ok && got != tc.version { + t.Fatalf("value = %q, want the build identifier unchanged", got) + } + if !ok && got != "" { + t.Fatalf("value = %q, want it omitted", got) + } + }) + } +} + +// dialWorker connects a worker to a test WebSocket server and returns the +// headers the server observed on the upgrade request. +func dialWorker(t *testing.T, version string) http.Header { + t.Helper() + + observed := make(chan http.Header, 1) + upgrader := websocket.Upgrader{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + observed <- r.Header.Clone() + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + _ = conn.Close() + })) + defer server.Close() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + w := &Worker{ + ctx: ctx, + config: Config{ + WebSocketURL: "ws" + strings.TrimPrefix(server.URL, "http"), + WorkerID: "test-worker", + APIKey: "wk-test", + Version: version, + }, + } + if err := w.connect(); err != nil { + t.Fatalf("connect: %v", err) + } + w.connMutex.Lock() + conn := w.conn + w.connMutex.Unlock() + if conn != nil { + _ = conn.Close() + } + + select { + case headers := <-observed: + return headers + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for the upgrade request") + return nil + } +} + +func TestConnectSendsTheWorkerVersionHeader(t *testing.T) { + headers := dialWorker(t, "v2026-08-04-15-14-28") + + if got := headers.Get(types.WorkerVersionHeader); got != "v2026-08-04-15-14-28" { + t.Fatalf("%s = %q, want the build identifier", types.WorkerVersionHeader, got) + } +} + +func TestConnectOmitsInvalidWorkerVersionWithoutFailing(t *testing.T) { + // An unusable build identifier must not cost the worker its connection: + // the server records provenance as not reported and the worker still runs + // tasks. + headers := dialWorker(t, "v1\ninjected: header") + + if got := headers.Get(types.WorkerVersionHeader); got != "" { + t.Fatalf("%s = %q, want the invalid value omitted", types.WorkerVersionHeader, got) + } + if got := headers.Get("Injected"); got != "" { + t.Fatalf("a control character in the version smuggled in header %q", got) + } +} + +func TestConnectOmitsTheHeaderWhenNoVersionIsStamped(t *testing.T) { + headers := dialWorker(t, "") + + if _, present := headers[types.WorkerVersionHeader]; present { + t.Fatal("an unstamped build must not report a version") + } +} + +// snapshotRecordingBackend records the debug-archive calls the worker makes. +type snapshotRecordingBackend struct { + outcome ExecuteResult + + snapshots atomic.Int32 + cleanups chan CancelParams + output string +} + +func newSnapshotRecordingBackend(outcome ExecuteResult) *snapshotRecordingBackend { + return &snapshotRecordingBackend{outcome: outcome, cleanups: make(chan CancelParams, 4)} +} + +func (b *snapshotRecordingBackend) ExecuteTask(context.Context, *TaskParams) ExecuteResult { + return b.outcome +} + +func (b *snapshotRecordingBackend) CancelTask(context.Context, *CancelParams) error { return nil } + +func (b *snapshotRecordingBackend) PreservesTasksOnShutdown() bool { return false } + +func (b *snapshotRecordingBackend) Shutdown(context.Context) {} + +func (b *snapshotRecordingBackend) SnapshotTaskLogs(_ context.Context, params *SnapshotParams) error { + b.snapshots.Add(1) + if b.output == "" { + return nil + } + return params.Sink.WriteChunk(debuglog.Chunk{ + Phase: debuglog.PhaseContainer, + Stream: debuglog.StreamCombined, + Data: []byte(b.output), + }) +} + +func (b *snapshotRecordingBackend) CleanupTaskResources(_ context.Context, params *CancelParams) error { + b.cleanups <- *params + return nil +} + +func newDebugArchiveWorker(t *testing.T, backend Backend, idleOnComplete string) *Worker { + t.Helper() + + captureConfig := debuglog.DefaultConfig() + captureConfig.Directory = t.TempDir() + store, err := debuglog.NewStore(captureConfig) + if err != nil { + t.Fatalf("NewStore: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + w := &Worker{ + ctx: ctx, + cancel: cancel, + config: Config{BackendType: "docker", IdleOnComplete: idleOnComplete}, + sendChan: make(chan []byte, 8), + tasks: newTaskRegistry(), + backend: backend, + debugLogStore: store, + reconnectDelay: InitialReconnectDelay, + } + w.debugLogs = debuglog.NewCoordinator(debuglog.CoordinatorOptions{ + Ownership: w.tasks, + Source: backendSnapshotSource{backend: backend}, + Sender: w, + Store: store, + }) + return w +} + +func TestExecuteTaskMovesOwnershipToCleanupGraceBeforeTerminalMessage(t *testing.T) { + tests := []struct { + name string + outcome ExecuteResult + }{ + {name: "success", outcome: executeCompleted()}, + {name: "backend failure", outcome: executeError(newBackendFailure("backend", "container_exit", context.DeadlineExceeded))}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + backend := newSnapshotRecordingBackend(tc.outcome) + // A long grace keeps the entry in place for the assertion. + w := newDebugArchiveWorker(t, backend, "60m") + + assignment := &types.TaskAssignmentMessage{ + TaskID: "task-1", + ExecutionID: "exec-1", + Task: &types.Task{ID: "task-1", Title: "test task"}, + } + w.tasks.StartTask("task-1", activeTask{cancel: func() {}, executionID: "exec-1"}, debuglog.BackendDocker, nil) + + w.executeTask(context.Background(), func() {}, trace.SpanFromContext(context.Background()), assignment, time.Now()) + + // The terminal message is only enqueued after the grace entry + // exists, so a request triggered by it always finds an owner. + if len(w.sendChan) == 0 { + t.Fatal("expected a terminal lifecycle message") + } + owner, owned := w.tasks.LookupExecution("task-1", "exec-1") + if !owned { + t.Fatal("execution ownership was released before the cleanup grace expired") + } + if !owner.InCleanupGrace { + t.Fatal("ownership should have moved to cleanup grace") + } + }) + } +} + +func TestCleanupGraceExpiryReleasesBackendResourcesOnce(t *testing.T) { + backend := newSnapshotRecordingBackend(executeCompleted()) + // A zero grace expires immediately, so the timer fires without a wait. + w := newDebugArchiveWorker(t, backend, "0s") + + w.tasks.StartTask("task-1", activeTask{cancel: func() {}, executionID: "exec-1"}, debuglog.BackendDocker, nil) + w.beginCleanupGrace(&types.TaskAssignmentMessage{ + TaskID: "task-1", + ExecutionID: "exec-1", + Task: &types.Task{ID: "task-1"}, + }) + + select { + case params := <-backend.cleanups: + if params.TaskID != "task-1" || params.ExecutionID != "exec-1" { + t.Fatalf("cleanup params = %+v, want {task-1 exec-1}", params) + } + case <-time.After(5 * time.Second): + t.Fatal("backend cleanup was not invoked at the grace deadline") + } + + if _, owned := w.tasks.LookupExecution("task-1", "exec-1"); owned { + t.Fatal("ownership survived the cleanup-grace deadline") + } + + // A second expiry must be a no-op rather than a duplicate cleanup. + w.expireCleanupGrace("task-1", "exec-1") + select { + case params := <-backend.cleanups: + t.Fatalf("cleanup ran twice for %+v", params) + case <-time.After(200 * time.Millisecond): + } +} + +// A worker restart must not strand the resources an execution was holding only +// for log retrieval. Ownership and the expiry timer are process-local, so a +// replacement worker cannot finish the job; without this the resource survives +// until an unrelated backstop collects it, far past the operator's grace. +func TestShutdownReleasesTerminalCleanupGraceResources(t *testing.T) { + backend := newSnapshotRecordingBackend(executeCompleted()) + // A long grace guarantees the timer has not fired, so shutdown is the only + // thing that can release the entry. + w := newDebugArchiveWorker(t, backend, "60m") + + w.tasks.StartTask("task-1", activeTask{cancel: func() {}, executionID: "exec-1"}, debuglog.BackendKubernetes, nil) + w.beginCleanupGrace(&types.TaskAssignmentMessage{ + TaskID: "task-1", + ExecutionID: "exec-1", + Task: &types.Task{ID: "task-1"}, + }) + w.tasks.Delete("task-1") + + if _, owned := w.tasks.LookupExecution("task-1", "exec-1"); !owned { + t.Fatal("the execution should be in cleanup grace before shutdown") + } + + w.Shutdown() + + select { + case params := <-backend.cleanups: + if params.TaskID != "task-1" || params.ExecutionID != "exec-1" { + t.Fatalf("cleanup params = %+v, want {task-1 exec-1}", params) + } + case <-time.After(5 * time.Second): + t.Fatal("shutdown did not release the cleanup-grace execution's backend resources") + } + + if _, owned := w.tasks.LookupExecution("task-1", "exec-1"); owned { + t.Fatal("ownership survived shutdown") + } +} + +func TestShutdownLeavesActiveExecutionsToTheBackendContract(t *testing.T) { + // An execution still running has not reported terminal state, so its task + // unit may legitimately outlive this process. Only each backend's own + // shutdown contract may decide its fate. + backend := newSnapshotRecordingBackend(executeCompleted()) + w := newDebugArchiveWorker(t, backend, "60m") + w.tasks.StartTask("task-1", activeTask{cancel: func() {}, executionID: "exec-1"}, debuglog.BackendKubernetes, nil) + + w.Shutdown() + + select { + case params := <-backend.cleanups: + t.Fatalf("shutdown cleaned up an active execution: %+v", params) + case <-time.After(200 * time.Millisecond): + } +} + +func TestShutdownReleasesTheDirectCaptureItWasHolding(t *testing.T) { + backend := newSnapshotRecordingBackend(executeCompleted()) + w := newDebugArchiveWorker(t, backend, "60m") + + capture, err := w.debugLogStore.NewTaskLogCapture(nil) + if err != nil { + t.Fatalf("NewTaskLogCapture: %v", err) + } + w.tasks.StartTask("task-1", activeTask{cancel: func() {}, executionID: "exec-1"}, debuglog.BackendDirect, capture) + w.beginCleanupGrace(&types.TaskAssignmentMessage{ + TaskID: "task-1", + ExecutionID: "exec-1", + Task: &types.Task{ID: "task-1"}, + }) + w.tasks.Delete("task-1") + + if reserved := w.debugLogStore.ReservedBytes(); reserved == 0 { + t.Fatal("expected the capture to hold part of the disk budget") + } + + w.Shutdown() + + if reserved := w.debugLogStore.ReservedBytes(); reserved != 0 { + t.Fatalf("reserved bytes = %d after shutdown, want the capture released", reserved) + } +} + +func TestRegistryDistinguishesExecutionsOfTheSameRun(t *testing.T) { + registry := newTaskRegistry() + registry.StartTask("run-1", activeTask{cancel: func() {}, executionID: "exec-1"}, debuglog.BackendDocker, nil) + + if _, owned := registry.LookupExecution("run-1", "exec-1"); !owned { + t.Fatal("the exact execution must be owned") + } + if _, owned := registry.LookupExecution("run-1", "exec-2"); owned { + t.Fatal("a different execution of the same run must not be owned") + } + if _, owned := registry.LookupExecution("run-2", "exec-1"); owned { + t.Fatal("a different run must not be owned") + } +} + +func TestHandleMessageDispatchesArchiveRequestsOffTheReadLoop(t *testing.T) { + backend := newSnapshotRecordingBackend(executeCompleted()) + backend.output = "captured container output" + w := newDebugArchiveWorker(t, backend, "60m") + w.tasks.StartTask("run-1", activeTask{cancel: func() {}, executionID: "exec-1"}, debuglog.BackendDocker, nil) + + uploaded := make(chan struct{}, 1) + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, _ *http.Request) { + uploaded <- struct{}{} + rw.WriteHeader(http.StatusOK) + })) + defer server.Close() + + request := types.DebugArchiveLogsRequestedMessage{ + ProtocolVersion: types.DebugArchiveProtocolVersion, + RequestID: "request-1", + ArchiveID: "archive-1", + CollectionID: "collection-1", + RunID: "run-1", + ExecutionID: "exec-1", + RequestedFormat: types.DebugArchiveFormatNDJSON, + ExpiresAt: time.Now().Add(5 * time.Minute), + MaxBytes: 1 << 15, + ContentTransformer: types.ContentTransformerDescriptor{ + Kind: debuglog.TransformerKindNoop, + Version: 1, + }, + UploadTarget: types.UploadTarget{URL: server.URL, Method: http.MethodPut}, + } + data, err := json.Marshal(request) + if err != nil { + t.Fatalf("marshal request: %v", err) + } + message, err := json.Marshal(types.WebSocketMessage{ + Type: types.MessageTypeDebugArchiveLogsRequested, + Data: data, + }) + if err != nil { + t.Fatalf("marshal envelope: %v", err) + } + + // handleMessage returning promptly is what keeps heartbeats and later + // assignments flowing while a snapshot runs. + w.handleMessage(message) + + select { + case <-uploaded: + case <-time.After(10 * time.Second): + t.Fatal("the archive request never reached the upload target") + } + + w.debugLogs.Wait() + ack := readAckMessage(t, w.sendChan) + if ack.Outcome != types.DebugArchiveOutcomeUploaded { + t.Fatalf("outcome = %q (%s), want %q", ack.Outcome, ack.ReasonCode, types.DebugArchiveOutcomeUploaded) + } + if ack.RunID != "run-1" || ack.ExecutionID != "exec-1" { + t.Fatalf("acknowledgement identity = %s/%s, want run-1/exec-1", ack.RunID, ack.ExecutionID) + } +} + +func TestHandleMessageIgnoresArchiveRequestsForUnownedExecutions(t *testing.T) { + backend := newSnapshotRecordingBackend(executeCompleted()) + w := newDebugArchiveWorker(t, backend, "60m") + + request := types.DebugArchiveLogsRequestedMessage{ + ProtocolVersion: types.DebugArchiveProtocolVersion, + RequestID: "request-1", + ArchiveID: "archive-1", + CollectionID: "collection-1", + RunID: "run-1", + ExecutionID: "exec-1", + RequestedFormat: types.DebugArchiveFormatNDJSON, + ExpiresAt: time.Now().Add(5 * time.Minute), + MaxBytes: 1 << 15, + ContentTransformer: types.ContentTransformerDescriptor{ + Kind: debuglog.TransformerKindNoop, + Version: 1, + }, + UploadTarget: types.UploadTarget{URL: "https://storage.example.com/candidate", Method: http.MethodPut}, + } + data, _ := json.Marshal(request) + message, _ := json.Marshal(types.WebSocketMessage{ + Type: types.MessageTypeDebugArchiveLogsRequested, + Data: data, + }) + + w.handleMessage(message) + w.debugLogs.Wait() + + if len(w.sendChan) != 0 { + t.Fatalf("a non-owning process enqueued %d messages, want 0", len(w.sendChan)) + } + if backend.snapshots.Load() != 0 { + t.Fatal("a non-owning process read the backend") + } +} + +func readAckMessage(t *testing.T, ch <-chan []byte) types.DebugArchiveLogsUploadedMessage { + t.Helper() + select { + case raw := <-ch: + var envelope types.WebSocketMessage + if err := json.Unmarshal(raw, &envelope); err != nil { + t.Fatalf("failed to unmarshal websocket message: %v", err) + } + if envelope.Type != types.MessageTypeDebugArchiveLogsUploaded { + t.Fatalf("message type = %q, want %q", envelope.Type, types.MessageTypeDebugArchiveLogsUploaded) + } + var ack types.DebugArchiveLogsUploadedMessage + if err := json.Unmarshal(envelope.Data, &ack); err != nil { + t.Fatalf("failed to unmarshal acknowledgement: %v", err) + } + return ack + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for an acknowledgement message") + return types.DebugArchiveLogsUploadedMessage{} + } +} diff --git a/internal/worker/direct.go b/internal/worker/direct.go index 7f93bbb..5b7d367 100644 --- a/internal/worker/direct.go +++ b/internal/worker/direct.go @@ -4,15 +4,19 @@ import ( "context" "errors" "fmt" + "io" "os" "os/exec" "path/filepath" "strings" + "sync" "syscall" "github.com/joho/godotenv" + "github.com/warpdotdev/oz-agent-worker/internal/debuglog" "github.com/warpdotdev/oz-agent-worker/internal/log" "github.com/warpdotdev/oz-agent-worker/internal/metrics" + "github.com/warpdotdev/oz-agent-worker/internal/types" "go.opentelemetry.io/otel/attribute" ) @@ -87,6 +91,12 @@ type DirectBackendConfig struct { type DirectBackend struct { config DirectBackendConfig ozPath string // resolved path to the oz CLI + + // capturesMutex guards captures, the exact (task, execution) to capture + // registry a debug-archive request resolves against. Direct execution has + // no provider log API, so its bounded capture is the only log source. + capturesMutex sync.Mutex + captures map[executionKey]*debuglog.TaskLogCapture } // NewDirectBackend creates a new direct backend, verifying the oz CLI is available. @@ -123,8 +133,9 @@ func NewDirectBackend(ctx context.Context, config DirectBackendConfig) (*DirectB } return &DirectBackend{ - config: config, - ozPath: ozPath, + config: config, + ozPath: ozPath, + captures: make(map[executionKey]*debuglog.TaskLogCapture), }, nil } @@ -135,6 +146,15 @@ func (b *DirectBackend) ExecuteTask(ctx context.Context, params *TaskParams) Exe return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonWorkspaceSetup, fmt.Errorf("invalid task ID for workspace path: %w", err))) } + // Registering the capture before any setup work runs means output from the + // very first phase is archivable, and a request arriving mid-execution + // resolves this exact execution's capture. + executionID := taskExecutionID(params) + phases := newDirectPhaseWriters(params.LogCapture) + if params.LogCapture != nil { + b.registerCapture(taskID, executionID, params.LogCapture) + } + // Determine working directory: shared target dir or per-task workspace. var workspaceDir string usingTargetDir := b.config.TargetDir != "" @@ -159,14 +179,14 @@ func (b *DirectBackend) ExecuteTask(ctx context.Context, params *TaskParams) Exe defer func() { if usingTargetDir { // Don't clean up the shared target directory. - b.runTeardownIfConfigured(ctx, taskID, workspaceDir, gitConfigPath) + b.runTeardownIfConfigured(ctx, taskID, workspaceDir, gitConfigPath, phases) return } if b.config.NoCleanup { log.Infof(ctx, "Skipping cleanup for workspace: %s", workspaceDir) return } - b.cleanup(ctx, taskID, workspaceDir, gitConfigPath) + b.cleanup(ctx, taskID, workspaceDir, gitConfigPath, phases) }() // 2. Create temp environment file for setup script to write to. @@ -203,7 +223,7 @@ func (b *DirectBackend) ExecuteTask(ctx context.Context, params *TaskParams) Exe ) log.Infof(ctx, "Running setup command: %s", b.config.SetupCommand) - if err := b.runCommand(ctx, b.config.SetupCommand, workspaceDir, setupEnv); err != nil { + if err := b.runCommand(ctx, b.config.SetupCommand, workspaceDir, setupEnv, phases.setup); err != nil { if ctx.Err() != nil { return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonTaskCancelled, ctx.Err())) } @@ -230,8 +250,8 @@ func (b *DirectBackend) ExecuteTask(ctx context.Context, params *TaskParams) Exe cmd := exec.CommandContext(ctx, b.ozPath, params.BaseArgs...) // #nosec G204 -- ozPath is resolved at backend startup and args are generated by the worker. cmd.Dir = workspaceDir cmd.Env = mergeEnvVars(hostBaseEnv(), envVars) - cmd.Stdout = os.Stdout - cmd.Stderr = os.Stderr + cmd.Stdout = phases.agent.stdout + cmd.Stderr = phases.agent.stderr log.Infof(ctx, "Running oz agent in workspace %s", workspaceDir) log.Debugf(ctx, "Command: %s %s", b.ozPath, strings.Join(params.BaseArgs, " ")) @@ -303,7 +323,7 @@ func (b *DirectBackend) PreservesTasksOnShutdown() bool { } // runTeardownIfConfigured runs the teardown command if one is configured. -func (b *DirectBackend) runTeardownIfConfigured(ctx context.Context, taskID, workspaceDir, gitConfigPath string) { +func (b *DirectBackend) runTeardownIfConfigured(ctx context.Context, taskID, workspaceDir, gitConfigPath string, phases directPhaseWriters) { if b.config.TeardownCommand == "" { return } @@ -314,7 +334,7 @@ func (b *DirectBackend) runTeardownIfConfigured(ctx context.Context, taskID, wor fmt.Sprintf("OZ_RUN_ID=%s", taskID), } log.Infof(ctx, "Running teardown command: %s", b.config.TeardownCommand) - if err := b.runCommand(ctx, b.config.TeardownCommand, workspaceDir, teardownEnv); err != nil { + if err := b.runCommand(ctx, b.config.TeardownCommand, workspaceDir, teardownEnv, phases.teardown); err != nil { metrics.AddTaskEvent(ctx, "cleanup.failed", attribute.String("operation", "teardown"), attribute.String("error.message", err.Error()), @@ -324,8 +344,8 @@ func (b *DirectBackend) runTeardownIfConfigured(ctx context.Context, taskID, wor } // cleanup runs the teardown command (if configured) and removes the workspace directory. -func (b *DirectBackend) cleanup(ctx context.Context, taskID, workspaceDir, gitConfigPath string) { - b.runTeardownIfConfigured(ctx, taskID, workspaceDir, gitConfigPath) +func (b *DirectBackend) cleanup(ctx context.Context, taskID, workspaceDir, gitConfigPath string, phases directPhaseWriters) { + b.runTeardownIfConfigured(ctx, taskID, workspaceDir, gitConfigPath, phases) log.Infof(ctx, "Removing workspace: %s", workspaceDir) if err := os.RemoveAll(workspaceDir); err != nil { @@ -340,15 +360,84 @@ func (b *DirectBackend) cleanup(ctx context.Context, taskID, workspaceDir, gitCo // runCommand executes a shell command with the given working directory and environment. // Setup/teardown commands inherit the full worker environment so they can access // tools and credentials (e.g. aws, docker) needed for workspace provisioning. -func (b *DirectBackend) runCommand(ctx context.Context, command, dir string, env []string) error { +func (b *DirectBackend) runCommand(ctx context.Context, command, dir string, env []string, output phaseWriters) error { cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command) // #nosec G204 -- setup/teardown commands are explicit operator configuration. cmd.Dir = dir cmd.Env = mergeEnvVars(os.Environ(), env) - cmd.Stdout = os.Stdout - cmd.Stderr = os.Stderr + cmd.Stdout = output.stdout + cmd.Stderr = output.stderr return cmd.Run() } +// phaseWriters are the stdout and stderr sinks for one execution phase. +type phaseWriters struct { + stdout io.Writer + stderr io.Writer +} + +// directPhaseWriters routes each phase's output to the worker console and, +// when archive capture is enabled, to the task's bounded capture as well. +type directPhaseWriters struct { + setup phaseWriters + agent phaseWriters + teardown phaseWriters +} + +func newDirectPhaseWriters(capture *debuglog.TaskLogCapture) directPhaseWriters { + console := phaseWriters{stdout: os.Stdout, stderr: os.Stderr} + if capture == nil { + return directPhaseWriters{setup: console, agent: console, teardown: console} + } + + phase := func(p debuglog.Phase) phaseWriters { + return phaseWriters{ + stdout: io.MultiWriter(os.Stdout, capture.Writer(p, debuglog.StreamStdout)), + stderr: io.MultiWriter(os.Stderr, capture.Writer(p, debuglog.StreamStderr)), + } + } + return directPhaseWriters{ + setup: phase(debuglog.PhaseSetup), + agent: phase(debuglog.PhaseAgent), + teardown: phase(debuglog.PhaseTeardown), + } +} + +func (b *DirectBackend) registerCapture(taskID, executionID string, capture *debuglog.TaskLogCapture) { + b.capturesMutex.Lock() + defer b.capturesMutex.Unlock() + b.captures[executionKey{runID: taskID, executionID: executionID}] = capture +} + +func (b *DirectBackend) lookupCapture(taskID, executionID string) (*debuglog.TaskLogCapture, bool) { + b.capturesMutex.Lock() + defer b.capturesMutex.Unlock() + capture, ok := b.captures[executionKey{runID: taskID, executionID: executionID}] + return capture, ok +} + +// SnapshotTaskLogs replays this execution's capture into sink at a fixed +// watermark, so an active task keeps writing while the snapshot is taken and +// its later output remains available to a later request. +func (b *DirectBackend) SnapshotTaskLogs(_ context.Context, params *SnapshotParams) error { + capture, ok := b.lookupCapture(params.TaskID, params.ExecutionID) + if !ok || capture == nil { + return debuglog.NewSnapshotError(types.DebugArchiveReasonCaptureUnavailable, "no output capture exists for this execution") + } + if err := capture.SnapshotTo(params.Sink); err != nil { + return debuglog.NewSnapshotError(types.DebugArchiveReasonSnapshotFailed, "failed to replay the output capture") + } + return nil +} + +// CleanupTaskResources forgets the execution's capture. The worker owns the +// capture handle and closes it, which is what deletes its bytes from disk. +func (b *DirectBackend) CleanupTaskResources(_ context.Context, params *CancelParams) error { + b.capturesMutex.Lock() + defer b.capturesMutex.Unlock() + delete(b.captures, executionKey{runID: params.TaskID, executionID: params.ExecutionID}) + return nil +} + // mergeEnvVars merges base and override env var slices (KEY=VALUE format). // Override entries take precedence over base entries with the same key. func mergeEnvVars(base, override []string) []string { diff --git a/internal/worker/direct_capture_test.go b/internal/worker/direct_capture_test.go new file mode 100644 index 0000000..0e30706 --- /dev/null +++ b/internal/worker/direct_capture_test.go @@ -0,0 +1,216 @@ +package worker + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/debuglog" + "github.com/warpdotdev/oz-agent-worker/internal/types" +) + +// writeShellScript writes an executable script for the direct backend to run. +func writeShellScript(t *testing.T, path, body string) string { + t.Helper() + if err := os.WriteFile(path, []byte(body), 0o700); err != nil { // #nosec G306 -- an executable test fixture. + t.Fatalf("failed to write %s: %v", path, err) + } + return path +} + +func newCaptureStore(t *testing.T) *debuglog.Store { + t.Helper() + config := debuglog.DefaultConfig() + config.Directory = t.TempDir() + store, err := debuglog.NewStore(config) + if err != nil { + t.Fatalf("NewStore: %v", err) + } + return store +} + +// TestDirectBackendCapturesEveryPhaseOnBothStreams runs a real direct execution +// whose setup, agent, and teardown each write a distinguishable sentinel to +// stdout and stderr, then proves every sentinel survives a snapshot taken after +// the per-task workspace has been removed. +func TestDirectBackendCapturesEveryPhaseOnBothStreams(t *testing.T) { + testDir := t.TempDir() + ozPath := writeShellScript(t, filepath.Join(testDir, "oz"), `#!/bin/sh +echo "AGENT-STDOUT-SENTINEL" +echo "AGENT-STDERR-SENTINEL" >&2 +exit 0 +`) + + workspaceRoot := filepath.Join(testDir, "workspaces") + backend, err := NewDirectBackend(context.Background(), DirectBackendConfig{ + WorkspaceRoot: workspaceRoot, + OzPath: ozPath, + SetupCommand: strings.Join([]string{ + `echo "SETUP-STDOUT-SENTINEL"`, + `echo "SETUP-STDERR-SENTINEL" >&2`, + }, "\n"), + TeardownCommand: strings.Join([]string{ + `echo "TEARDOWN-STDOUT-SENTINEL"`, + `echo "TEARDOWN-STDERR-SENTINEL" >&2`, + }, "\n"), + }) + if err != nil { + t.Fatalf("NewDirectBackend: %v", err) + } + + store := newCaptureStore(t) + capture, err := store.NewTaskLogCapture(nil) + if err != nil { + t.Fatalf("NewTaskLogCapture: %v", err) + } + defer capture.Close() + + result := backend.ExecuteTask(context.Background(), &TaskParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + LogCapture: capture, + }) + if result.Error != nil { + t.Fatalf("ExecuteTask: %v", result.Error) + } + + // The per-task workspace is gone by now; only the capture remains. + if _, err := os.Stat(filepath.Join(workspaceRoot, "task-1")); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("expected the workspace to be cleaned up, stat err = %v", err) + } + + capture.Finalize(5 * time.Second) + sink := &captureSink{} + if err := backend.SnapshotTaskLogs(context.Background(), &SnapshotParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + Sink: sink, + }); err != nil { + t.Fatalf("SnapshotTaskLogs: %v", err) + } + + // Each sentinel must appear exactly once, under the phase and stream the + // backend actually owns the handle for. + want := map[string]struct { + phase debuglog.Phase + stream debuglog.Stream + }{ + "SETUP-STDOUT-SENTINEL": {debuglog.PhaseSetup, debuglog.StreamStdout}, + "SETUP-STDERR-SENTINEL": {debuglog.PhaseSetup, debuglog.StreamStderr}, + "AGENT-STDOUT-SENTINEL": {debuglog.PhaseAgent, debuglog.StreamStdout}, + "AGENT-STDERR-SENTINEL": {debuglog.PhaseAgent, debuglog.StreamStderr}, + "TEARDOWN-STDOUT-SENTINEL": {debuglog.PhaseTeardown, debuglog.StreamStdout}, + "TEARDOWN-STDERR-SENTINEL": {debuglog.PhaseTeardown, debuglog.StreamStderr}, + } + for sentinel, expected := range want { + occurrences := 0 + for _, chunk := range sink.chunks { + if !strings.Contains(string(chunk.Data), sentinel) { + continue + } + occurrences++ + if chunk.Phase != expected.phase { + t.Errorf("%s phase = %q, want %q", sentinel, chunk.Phase, expected.phase) + } + if chunk.Stream != expected.stream { + t.Errorf("%s stream = %q, want %q", sentinel, chunk.Stream, expected.stream) + } + } + if occurrences != 1 { + t.Errorf("%s appeared %d times, want exactly once", sentinel, occurrences) + } + } +} + +func TestDirectBackendStoresNoWorkspaceFileContent(t *testing.T) { + testDir := t.TempDir() + ozPath := writeShellScript(t, filepath.Join(testDir, "oz"), `#!/bin/sh +printf 'WORKSPACE-FILE-SECRET\n' > secret.txt +echo "AGENT-OUTPUT" +exit 0 +`) + + backend, err := NewDirectBackend(context.Background(), DirectBackendConfig{ + WorkspaceRoot: filepath.Join(testDir, "workspaces"), + OzPath: ozPath, + }) + if err != nil { + t.Fatalf("NewDirectBackend: %v", err) + } + + store := newCaptureStore(t) + capture, err := store.NewTaskLogCapture(nil) + if err != nil { + t.Fatalf("NewTaskLogCapture: %v", err) + } + defer capture.Close() + + if result := backend.ExecuteTask(context.Background(), &TaskParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + LogCapture: capture, + }); result.Error != nil { + t.Fatalf("ExecuteTask: %v", result.Error) + } + capture.Finalize(5 * time.Second) + + sink := &captureSink{} + if err := backend.SnapshotTaskLogs(context.Background(), &SnapshotParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + Sink: sink, + }); err != nil { + t.Fatalf("SnapshotTaskLogs: %v", err) + } + + var captured strings.Builder + for _, chunk := range sink.chunks { + captured.Write(chunk.Data) + } + if !strings.Contains(captured.String(), "AGENT-OUTPUT") { + t.Fatalf("captured = %q, want the agent's stdout", captured.String()) + } + if strings.Contains(captured.String(), "WORKSPACE-FILE-SECRET") { + t.Fatal("the capture stored workspace file content, not just stdout/stderr") + } +} + +func TestDirectBackendRunsWithoutACapture(t *testing.T) { + // Archive capture is optional: an execution must run identically when the + // store could not allocate one. + testDir := t.TempDir() + ozPath := writeShellScript(t, filepath.Join(testDir, "oz"), "#!/bin/sh\necho ok\nexit 0\n") + + backend, err := NewDirectBackend(context.Background(), DirectBackendConfig{ + WorkspaceRoot: filepath.Join(testDir, "workspaces"), + OzPath: ozPath, + }) + if err != nil { + t.Fatalf("NewDirectBackend: %v", err) + } + + if result := backend.ExecuteTask(context.Background(), &TaskParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + Task: &types.Task{ID: "task-1"}, + }); result.Error != nil { + t.Fatalf("ExecuteTask: %v", result.Error) + } + + err = backend.SnapshotTaskLogs(context.Background(), &SnapshotParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + Sink: &captureSink{}, + }) + var snapshotErr *debuglog.SnapshotError + if !errors.As(err, &snapshotErr) { + t.Fatalf("error = %v, want a *debuglog.SnapshotError", err) + } + if snapshotErr.ReasonCode != types.DebugArchiveReasonCaptureUnavailable { + t.Fatalf("reason = %q, want %q", snapshotErr.ReasonCode, types.DebugArchiveReasonCaptureUnavailable) + } +} diff --git a/internal/worker/dispatch_test.go b/internal/worker/dispatch_test.go index ae986ec..b444a84 100644 --- a/internal/worker/dispatch_test.go +++ b/internal/worker/dispatch_test.go @@ -12,7 +12,9 @@ import ( // dispatchBackend is a fake Backend that reports a successful fire-and-forget // dispatch. Its CancelTask is a no-op. -type dispatchBackend struct{} +type dispatchBackend struct { + noLogSnapshotBackend +} func (b *dispatchBackend) ExecuteTask(context.Context, *TaskParams) ExecuteResult { return executeSpawned() @@ -34,11 +36,11 @@ func (b *cancelableDispatchBackend) CancelTask(_ context.Context, params *Cancel func newDispatchWorker(backend Backend) *Worker { return &Worker{ - ctx: context.Background(), - config: Config{}, - sendChan: make(chan []byte, 4), - activeTasks: map[string]activeTask{"task-1": {cancel: func() {}, executionID: "exec-1"}}, - backend: backend, + ctx: context.Background(), + config: Config{}, + sendChan: make(chan []byte, 4), + tasks: registryWith(map[string]activeTask{"task-1": {cancel: func() {}, executionID: "exec-1"}}), + backend: backend, } } @@ -60,11 +62,9 @@ func TestExecuteTaskDispatchedSuppressesTerminalMessage(t *testing.T) { msg := readWebSocketMessage(t, w.sendChan) t.Fatalf("expected no terminal message after dispatch, got %q", msg.Type) } - w.tasksMutex.Lock() - task, ok := w.activeTasks["task-1"] - w.tasksMutex.Unlock() + task, ok := w.tasks.Get("task-1") if !ok { - t.Fatal("spawned task should remain in activeTasks") + t.Fatal("spawned task should remain tracked") } if !task.spawned { t.Error("spawned task entry should be marked spawned") @@ -102,10 +102,10 @@ func spawnedActiveTask(executionID string) activeTask { func TestHandleTaskCancellationRoutesToBackendCancelTask(t *testing.T) { backend := &cancelableDispatchBackend{cancelCalled: make(chan *CancelParams, 1)} w := &Worker{ - ctx: context.Background(), - sendChan: make(chan []byte, 1), - activeTasks: map[string]activeTask{"task-1": spawnedActiveTask("exec-1")}, - backend: backend, + ctx: context.Background(), + sendChan: make(chan []byte, 1), + tasks: registryWith(map[string]activeTask{"task-1": spawnedActiveTask("exec-1")}), + backend: backend, } w.handleTaskCancellation(&types.TaskCancellationMessage{TaskID: "task-1"}) @@ -119,10 +119,7 @@ func TestHandleTaskCancellationRoutesToBackendCancelTask(t *testing.T) { t.Fatal("CancelTask was not invoked for a spawned task") } - w.tasksMutex.Lock() - _, ok := w.activeTasks["task-1"] - w.tasksMutex.Unlock() - if ok { + if _, ok := w.tasks.Get("task-1"); ok { t.Error("spawned task should be removed after cancellation is routed") } } @@ -134,11 +131,11 @@ func TestHandleTaskCancellationRunningTaskCancelsContextAndBackend(t *testing.T) w := &Worker{ ctx: context.Background(), sendChan: make(chan []byte, 1), - activeTasks: map[string]activeTask{"task-1": { + tasks: registryWith(map[string]activeTask{"task-1": { ctx: taskCtx, cancel: taskCancel, executionID: "exec-1", - }}, + }}), backend: backend, } @@ -157,11 +154,9 @@ func TestHandleTaskCancellationRunningTaskCancelsContextAndBackend(t *testing.T) } // The entry stays until executeTask's deferred cleanup removes it. - w.tasksMutex.Lock() - task, ok := w.activeTasks["task-1"] - w.tasksMutex.Unlock() + task, ok := w.tasks.Get("task-1") if !ok { - t.Fatal("running task should remain in activeTasks until executeTask returns") + t.Fatal("running task should remain tracked until executeTask returns") } if task.cancellationSource != taskCancellationSourceUser { t.Fatalf("cancellation source = %q, want %q", task.cancellationSource, taskCancellationSourceUser) @@ -170,10 +165,10 @@ func TestHandleTaskCancellationRunningTaskCancelsContextAndBackend(t *testing.T) func TestHandleTaskCancellationSpawnedNoopCancelEmitsNoMessage(t *testing.T) { w := &Worker{ - ctx: context.Background(), - sendChan: make(chan []byte, 1), - activeTasks: map[string]activeTask{"task-1": spawnedActiveTask("exec-1")}, - backend: &dispatchBackend{}, + ctx: context.Background(), + sendChan: make(chan []byte, 1), + tasks: registryWith(map[string]activeTask{"task-1": spawnedActiveTask("exec-1")}), + backend: &dispatchBackend{}, } // Must not panic and must not emit any task status message when the diff --git a/internal/worker/docker.go b/internal/worker/docker.go index 9da79b5..0912f06 100644 --- a/internal/worker/docker.go +++ b/internal/worker/docker.go @@ -5,8 +5,10 @@ import ( "fmt" "io" "strings" + "sync" "time" + "github.com/containerd/errdefs" "github.com/distribution/reference" cliconfig "github.com/docker/cli/cli/config" "github.com/moby/moby/api/pkg/authconfig" @@ -15,6 +17,7 @@ import ( "github.com/moby/moby/client" ocispec "github.com/opencontainers/image-spec/specs-go/v1" "github.com/rs/zerolog" + "github.com/warpdotdev/oz-agent-worker/internal/debuglog" "github.com/warpdotdev/oz-agent-worker/internal/log" "github.com/warpdotdev/oz-agent-worker/internal/metrics" "github.com/warpdotdev/oz-agent-worker/internal/types" @@ -22,6 +25,10 @@ import ( const dockerHubAuthConfigKey = "https://index.docker.io/v1/" +// maxDiagnosticLogBytes caps how much container output the worker reads into +// memory for its own failure logging. +const maxDiagnosticLogBytes = 1 << 20 + // DockerBackendConfig holds configuration specific to the Docker backend. type DockerBackendConfig struct { NoCleanup bool @@ -43,6 +50,12 @@ type DockerBackend struct { dockerClient *client.Client platform string // Docker daemon platform (e.g., "linux/amd64" or "linux/arm64") platformSpec ocispec.Platform + + // containersMutex guards containers, the exact (task, execution) to + // container-ID registry that makes a debug-archive request resolve one + // container and no other. + containersMutex sync.Mutex + containers map[executionKey]string } // NewDockerBackend creates a new Docker backend, connecting to the Docker daemon. @@ -92,6 +105,7 @@ func NewDockerBackend(ctx context.Context, config DockerBackendConfig) (*DockerB OS: versionInfo.Os, Architecture: versionInfo.Arch, }, + containers: make(map[executionKey]string), }, nil } @@ -152,15 +166,12 @@ func (b *DockerBackend) ExecuteTask(ctx context.Context, params *TaskParams) Exe containerID := resp.ID log.Debugf(ctx, "Created Docker container: %s", containerID) - defer func() { - if containerID != "" && !b.config.NoCleanup { - cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), BackendShutdownTimeout) - defer cleanupCancel() - if _, removeErr := dockerClient.ContainerRemove(cleanupCtx, containerID, client.ContainerRemoveOptions{Force: true}); removeErr != nil { - log.Debugf(ctx, "Container %s already removed or removal failed: %v", containerID, removeErr) - } - } - }() + // Registering before start means a debug-archive request that arrives while + // the container is still coming up resolves the right container. The + // container is deliberately not removed when ExecuteTask returns: it is the + // retained log source until the execution's cleanup grace expires, which is + // what makes post-failure collection possible. + b.registerContainer(params.TaskID, taskExecutionID(params), containerID) if _, err := dockerClient.ContainerStart(ctx, containerID, client.ContainerStartOptions{}); err != nil { return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonContainerStart, fmt.Errorf("failed to start container: %w", err))) @@ -231,8 +242,110 @@ func dockerResourcesForShape(shape *types.InstanceShape) container.Resources { // Docker-backend task. func (b *DockerBackend) CancelTask(context.Context, *CancelParams) error { return nil } -// Shutdown closes the Docker client. +func (b *DockerBackend) registerContainer(taskID, executionID, containerID string) { + b.containersMutex.Lock() + defer b.containersMutex.Unlock() + b.containers[executionKey{runID: taskID, executionID: executionID}] = containerID +} + +func (b *DockerBackend) lookupContainer(taskID, executionID string) (string, bool) { + b.containersMutex.Lock() + defer b.containersMutex.Unlock() + containerID, ok := b.containers[executionKey{runID: taskID, executionID: executionID}] + return containerID, ok +} + +func (b *DockerBackend) forgetContainer(taskID, executionID string) { + b.containersMutex.Lock() + defer b.containersMutex.Unlock() + delete(b.containers, executionKey{runID: taskID, executionID: executionID}) +} + +// SnapshotTaskLogs streams the exact registered container's stdout and stderr +// into sink. Docker's log stream carries everything the entrypoint/launcher and +// the client process it starts wrote, so no line filtering is applied and no +// process identity is inferred. +func (b *DockerBackend) SnapshotTaskLogs(ctx context.Context, params *SnapshotParams) error { + containerID, ok := b.lookupContainer(params.TaskID, params.ExecutionID) + if !ok { + return debuglog.NewSnapshotError(types.DebugArchiveReasonResourceNotFound, "no container is registered for this execution") + } + + stream, err := b.dockerClient.ContainerLogs(ctx, containerID, client.ContainerLogsOptions{ + ShowStdout: true, + ShowStderr: true, + Timestamps: true, + }) + if err != nil { + return debuglog.NewSnapshotError(types.DebugArchiveReasonCaptureUnavailable, "container logs are unavailable") + } + defer func() { + if closeErr := stream.Close(); closeErr != nil { + log.Warnf(ctx, "Failed to close container log stream: %v", closeErr) + } + }() + + source := debuglog.SourceIdentity{ContainerID: containerID} + if err := copyDockerLogStream(stream, params.Sink, source); err != nil { + return debuglog.NewSnapshotError(types.DebugArchiveReasonSnapshotFailed, "failed to read container logs") + } + return nil +} + +// CleanupTaskResources removes the container retained for log retrieval. It is +// idempotent: an already-removed container or an unregistered execution is not +// an error. +// +// The registry entry is dropped only once removal is confirmed. Forgetting it +// first would strand a container the daemon refused to delete, because nothing +// would remain for the caller or Shutdown to retry. +func (b *DockerBackend) CleanupTaskResources(ctx context.Context, params *CancelParams) error { + containerID, ok := b.lookupContainer(params.TaskID, params.ExecutionID) + if !ok || containerID == "" { + return nil + } + if b.config.NoCleanup { + b.forgetContainer(params.TaskID, params.ExecutionID) + return nil + } + + if err := b.removeContainer(ctx, containerID); err != nil { + return err + } + b.forgetContainer(params.TaskID, params.ExecutionID) + return nil +} + +// removeContainer deletes a container, treating an already-absent container as +// success. Any other failure is surfaced so the caller can retry rather than +// silently losing track of the container. +func (b *DockerBackend) removeContainer(ctx context.Context, containerID string) error { + _, err := b.dockerClient.ContainerRemove(ctx, containerID, client.ContainerRemoveOptions{Force: true}) + if err == nil || errdefs.IsNotFound(err) { + return nil + } + return fmt.Errorf("failed to remove container %s: %w", containerID, err) +} + +// Shutdown removes any containers still retained for log retrieval and closes +// the Docker client. Docker task containers do not outlive the worker, so +// leaving them behind at shutdown would leak them until an operator intervened. +// This is the last chance to remove them, so a failure is reported loudly +// rather than swallowed. func (b *DockerBackend) Shutdown(ctx context.Context) { + b.containersMutex.Lock() + retained := b.containers + b.containers = make(map[executionKey]string) + b.containersMutex.Unlock() + + if !b.config.NoCleanup { + for key, containerID := range retained { + if err := b.removeContainer(ctx, containerID); err != nil { + log.Warnf(ctx, "Leaving container for task %s behind after worker shutdown: %v", key.runID, err) + } + } + } + if b.dockerClient != nil { if err := b.dockerClient.Close(); err != nil { log.Warnf(ctx, "Failed to close Docker client: %v", err) @@ -330,6 +443,11 @@ func (b *DockerBackend) getRegistryAuth(ctx context.Context, imageName string) s return authStr } +// getContainerLogs reads a bounded prefix of a container's output for the +// worker's own diagnostic logging. Debug-archive collection uses +// SnapshotTaskLogs instead, which streams under the request's byte bound; this +// path is capped so a chatty container cannot scale worker memory with its +// output. func (b *DockerBackend) getContainerLogs(ctx context.Context, dockerClient *client.Client, containerID string) (string, error) { out, err := dockerClient.ContainerLogs(ctx, containerID, client.ContainerLogsOptions{ ShowStdout: true, @@ -345,7 +463,7 @@ func (b *DockerBackend) getContainerLogs(ctx context.Context, dockerClient *clie } }() - logBytes, err := io.ReadAll(out) + logBytes, err := io.ReadAll(io.LimitReader(out, maxDiagnosticLogBytes)) if err != nil { return "", err } diff --git a/internal/worker/docker_logs.go b/internal/worker/docker_logs.go new file mode 100644 index 0000000..f4e82fb --- /dev/null +++ b/internal/worker/docker_logs.go @@ -0,0 +1,143 @@ +package worker + +import ( + "bufio" + "bytes" + "encoding/binary" + "errors" + "io" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/debuglog" +) + +const ( + // dockerStreamHeaderLen is the length of the frame header the Docker + // daemon prepends to each chunk of a non-TTY log stream. + dockerStreamHeaderLen = 8 + // dockerStreamTypeStdout and dockerStreamTypeStderr are the stream + // identifiers in that header's first byte. + dockerStreamTypeStdout = 1 + dockerStreamTypeStderr = 2 + // dockerMaxFrameBytes bounds a single frame so a malformed header cannot + // make the worker allocate an arbitrary buffer. + dockerMaxFrameBytes = 4 << 20 +) + +// copyDockerLogStream demultiplexes Docker's framed log stream into snapshot +// chunks, preserving each frame's real stdout/stderr identity and the +// per-line provider timestamp. It never filters lines and never attributes +// output to the entrypoint or the client process: the container's stream is +// the only identity Docker reports. +func copyDockerLogStream(stream io.Reader, sink debuglog.Sink, source debuglog.SourceIdentity) error { + reader := bufio.NewReader(stream) + header := make([]byte, dockerStreamHeaderLen) + + for { + if _, err := io.ReadFull(reader, header); err != nil { + if errors.Is(err, io.EOF) { + return nil + } + if errors.Is(err, io.ErrUnexpectedEOF) { + return nil + } + return err + } + + streamID := debuglog.StreamCombined + switch header[0] { + case dockerStreamTypeStdout: + streamID = debuglog.StreamStdout + case dockerStreamTypeStderr: + streamID = debuglog.StreamStderr + } + + size := binary.BigEndian.Uint32(header[4:]) + if size == 0 { + continue + } + if size > dockerMaxFrameBytes { + // A frame this large means the header was not real framing (for + // example a TTY-attached container). Report the remainder as + // combined output rather than guessing at boundaries. + return copyUnframedLogStream(io.MultiReader(bytes.NewReader(header), reader), sink, source) + } + + payload := make([]byte, size) + if _, err := io.ReadFull(reader, payload); err != nil { + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { + return nil + } + return err + } + + if err := emitTimestampedLines(payload, streamID, source, sink); err != nil { + return err + } + } +} + +// copyUnframedLogStream handles a stream Docker did not multiplex. The output +// is truthfully labeled combined rather than split by guesswork. +func copyUnframedLogStream(stream io.Reader, sink debuglog.Sink, source debuglog.SourceIdentity) error { + buf := make([]byte, debuglog.MaxChunkBytes) + for { + n, err := stream.Read(buf) + if n > 0 { + if emitErr := emitTimestampedLines(buf[:n], debuglog.StreamCombined, source, sink); emitErr != nil { + return emitErr + } + } + if err != nil { + if errors.Is(err, io.EOF) { + return nil + } + return err + } + } +} + +// emitTimestampedLines splits payload into lines and lifts each line's leading +// RFC3339 timestamp, which the Docker log API prepends when timestamps are +// requested, into the record's provider timestamp field. +func emitTimestampedLines(payload []byte, stream debuglog.Stream, source debuglog.SourceIdentity, sink debuglog.Sink) error { + for len(payload) > 0 { + line := payload + if idx := bytes.IndexByte(payload, '\n'); idx >= 0 { + line = payload[:idx+1] + payload = payload[idx+1:] + } else { + payload = nil + } + + timestamp, content := splitProviderTimestamp(line) + if len(content) == 0 { + continue + } + if err := sink.WriteChunk(debuglog.Chunk{ + Phase: debuglog.PhaseContainer, + Stream: stream, + Timestamp: timestamp, + Source: source, + Data: content, + }); err != nil { + return err + } + } + return nil +} + +// splitProviderTimestamp separates a provider-prefixed RFC3339 timestamp from +// the log content. A line without one keeps all of its bytes and reports a zero +// timestamp, so the record simply omits the provider timestamp field. +func splitProviderTimestamp(line []byte) (time.Time, []byte) { + space := bytes.IndexByte(line, ' ') + if space <= 0 { + return time.Time{}, line + } + timestamp, err := time.Parse(time.RFC3339Nano, string(line[:space])) + if err != nil { + return time.Time{}, line + } + return timestamp, line[space+1:] +} diff --git a/internal/worker/docker_logs_test.go b/internal/worker/docker_logs_test.go new file mode 100644 index 0000000..01d1e21 --- /dev/null +++ b/internal/worker/docker_logs_test.go @@ -0,0 +1,217 @@ +package worker + +import ( + "bytes" + "encoding/binary" + "io" + "strings" + "testing" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/debuglog" +) + +// dockerFrame builds one frame of the Docker daemon's multiplexed log stream. +func dockerFrame(streamType byte, payload string) []byte { + header := make([]byte, dockerStreamHeaderLen) + header[0] = streamType + binary.BigEndian.PutUint32(header[4:], uint32(len(payload))) + return append(header, payload...) +} + +// captureSink records the chunks a backend writes without touching disk. +type captureSink struct { + chunks []debuglog.Chunk + warnings []string + omitted int64 +} + +func (s *captureSink) WriteChunk(chunk debuglog.Chunk) error { + copied := chunk + copied.Data = append([]byte(nil), chunk.Data...) + s.chunks = append(s.chunks, copied) + return nil +} + +func (s *captureSink) WriteSourceError(_ debuglog.Phase, _ debuglog.Stream, _ debuglog.SourceIdentity, warningCode string) error { + s.warnings = append(s.warnings, warningCode) + return nil +} + +func (s *captureSink) NoteOmittedBytes(n int64) { s.omitted += n } + +func TestCopyDockerLogStreamPreservesStreamIdentityAndTimestamps(t *testing.T) { + stream := bytes.NewReader(bytes.Join([][]byte{ + dockerFrame(dockerStreamTypeStdout, "2026-01-02T03:04:05.000000000Z entrypoint starting\n"), + dockerFrame(dockerStreamTypeStderr, "2026-01-02T03:04:06.000000000Z client warning\n"), + dockerFrame(dockerStreamTypeStdout, "2026-01-02T03:04:07.000000000Z client done\n"), + }, nil)) + + sink := &captureSink{} + source := debuglog.SourceIdentity{ContainerID: "container-abc"} + if err := copyDockerLogStream(stream, sink, source); err != nil { + t.Fatalf("copyDockerLogStream: %v", err) + } + + if len(sink.chunks) != 3 { + t.Fatalf("chunk count = %d, want 3", len(sink.chunks)) + } + want := []struct { + stream debuglog.Stream + data string + }{ + {debuglog.StreamStdout, "entrypoint starting\n"}, + {debuglog.StreamStderr, "client warning\n"}, + {debuglog.StreamStdout, "client done\n"}, + } + for i, expected := range want { + got := sink.chunks[i] + if got.Stream != expected.stream { + t.Errorf("chunk %d stream = %q, want %q", i, got.Stream, expected.stream) + } + if string(got.Data) != expected.data { + t.Errorf("chunk %d data = %q, want %q", i, got.Data, expected.data) + } + if got.Phase != debuglog.PhaseContainer { + t.Errorf("chunk %d phase = %q, want %q", i, got.Phase, debuglog.PhaseContainer) + } + if got.Source.ContainerID != "container-abc" { + t.Errorf("chunk %d container = %q, want container-abc", i, got.Source.ContainerID) + } + if got.Timestamp.IsZero() { + t.Errorf("chunk %d lost the provider timestamp", i) + } + } + if !sink.chunks[0].Timestamp.Equal(time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)) { + t.Errorf("first timestamp = %v, want the provider value", sink.chunks[0].Timestamp) + } +} + +func TestCopyDockerLogStreamKeepsUntimestampedLinesIntact(t *testing.T) { + stream := bytes.NewReader(dockerFrame(dockerStreamTypeStdout, "plain line with no timestamp\n")) + + sink := &captureSink{} + if err := copyDockerLogStream(stream, sink, debuglog.SourceIdentity{}); err != nil { + t.Fatalf("copyDockerLogStream: %v", err) + } + + if len(sink.chunks) != 1 { + t.Fatalf("chunk count = %d, want 1", len(sink.chunks)) + } + if string(sink.chunks[0].Data) != "plain line with no timestamp\n" { + t.Fatalf("data = %q, want the whole line retained", sink.chunks[0].Data) + } + if !sink.chunks[0].Timestamp.IsZero() { + t.Fatal("a line without a provider timestamp must not report one") + } +} + +func TestCopyDockerLogStreamDoesNotFilterOrAttributeLines(t *testing.T) { + // Both the entrypoint and the client it starts write to the same container + // streams. Every line must survive, with no invented process identity. + stream := bytes.NewReader(bytes.Join([][]byte{ + dockerFrame(dockerStreamTypeStdout, "[entrypoint] preparing workspace\n"), + dockerFrame(dockerStreamTypeStdout, "[oz] agent turn 1\n"), + dockerFrame(dockerStreamTypeStderr, "[oz] agent warning\n"), + dockerFrame(dockerStreamTypeStderr, "[entrypoint] teardown\n"), + }, nil)) + + sink := &captureSink{} + if err := copyDockerLogStream(stream, sink, debuglog.SourceIdentity{}); err != nil { + t.Fatalf("copyDockerLogStream: %v", err) + } + + if len(sink.chunks) != 4 { + t.Fatalf("chunk count = %d, want every line retained", len(sink.chunks)) + } + for i, chunk := range sink.chunks { + if chunk.Phase != debuglog.PhaseContainer { + t.Errorf("chunk %d phase = %q, want the provider's own %q", i, chunk.Phase, debuglog.PhaseContainer) + } + } +} + +func TestCopyDockerLogStreamHandlesAnUnframedStream(t *testing.T) { + // A TTY-attached container returns raw bytes with no frame header. The + // output is reported as combined rather than split by guesswork. + stream := strings.NewReader(strings.Repeat("tty output line\n", 4)) + + sink := &captureSink{} + if err := copyUnframedLogStream(stream, sink, debuglog.SourceIdentity{}); err != nil { + t.Fatalf("copyUnframedLogStream: %v", err) + } + + if len(sink.chunks) == 0 { + t.Fatal("expected the unframed output to be captured") + } + for i, chunk := range sink.chunks { + if chunk.Stream != debuglog.StreamCombined { + t.Errorf("chunk %d stream = %q, want %q", i, chunk.Stream, debuglog.StreamCombined) + } + } +} + +func TestCopyDockerLogStreamStopsCleanlyOnATruncatedFrame(t *testing.T) { + // A stream cut mid-frame must end the snapshot rather than error out and + // discard everything already read. + complete := dockerFrame(dockerStreamTypeStdout, "first line\n") + partial := dockerFrame(dockerStreamTypeStdout, "second line\n")[:6] + + sink := &captureSink{} + if err := copyDockerLogStream(bytes.NewReader(append(complete, partial...)), sink, debuglog.SourceIdentity{}); err != nil { + t.Fatalf("copyDockerLogStream: %v", err) + } + if len(sink.chunks) != 1 { + t.Fatalf("chunk count = %d, want the complete frame retained", len(sink.chunks)) + } +} + +// countingReader reports how much of a stream was pulled into memory at once. +type countingReader struct { + remaining int + frame []byte + offset int + maxRead int +} + +func (r *countingReader) Read(p []byte) (int, error) { + if r.remaining == 0 && r.offset >= len(r.frame) { + return 0, io.EOF + } + if r.offset >= len(r.frame) { + r.offset = 0 + r.remaining-- + if r.remaining < 0 { + return 0, io.EOF + } + } + n := copy(p, r.frame[r.offset:]) + r.offset += n + if n > r.maxRead { + r.maxRead = n + } + return n, nil +} + +func TestCopyDockerLogStreamMemoryDoesNotScaleWithOutputSize(t *testing.T) { + // The archive path streams rather than reading the whole container log + // into memory, so a very large log is bounded by the chunk size. + frame := dockerFrame(dockerStreamTypeStdout, strings.Repeat("x", 4096)+"\n") + reader := &countingReader{remaining: 4096, frame: frame} + + sink := &captureSink{} + if err := copyDockerLogStream(reader, sink, debuglog.SourceIdentity{}); err != nil { + t.Fatalf("copyDockerLogStream: %v", err) + } + + total := 0 + for _, chunk := range sink.chunks { + if len(chunk.Data) > debuglog.MaxChunkBytes { + t.Fatalf("a chunk carried %d bytes, above the %d bound", len(chunk.Data), debuglog.MaxChunkBytes) + } + total += len(chunk.Data) + } + if total < 4096 { + t.Fatalf("captured %d bytes, want the large log to be streamed through", total) + } +} diff --git a/internal/worker/kubernetes.go b/internal/worker/kubernetes.go index 4621863..55dd4e1 100644 --- a/internal/worker/kubernetes.go +++ b/internal/worker/kubernetes.go @@ -7,9 +7,11 @@ import ( "io" "sort" "strings" + "sync" "time" "github.com/rs/zerolog" + "github.com/warpdotdev/oz-agent-worker/internal/debuglog" "github.com/warpdotdev/oz-agent-worker/internal/log" "github.com/warpdotdev/oz-agent-worker/internal/metrics" "github.com/warpdotdev/oz-agent-worker/internal/types" @@ -100,6 +102,21 @@ func terminatedExitCode(terminated *corev1.ContainerStateTerminated) int { type KubernetesBackend struct { config KubernetesBackendConfig clientset kubernetes.Interface + + // jobsMutex guards jobs, the exact (task, execution) to retained-Job + // registry a debug-archive request resolves against. + jobsMutex sync.Mutex + jobs map[executionKey]*retainedJob +} + +// retainedJob is a task Job kept past terminal state so its pods' logs stay +// readable through the execution's cleanup grace. +type retainedJob struct { + name string + // deleteAtCleanup mirrors the existing policy: a successful Job is deleted + // to keep the namespace clean, while a failed one is left for the Job TTL + // controller so operators can inspect it after the fact. + deleteAtCleanup bool } func (b *KubernetesBackend) PreservesTasksOnShutdown() bool { @@ -138,6 +155,7 @@ func NewKubernetesBackend(ctx context.Context, config KubernetesBackendConfig) ( backend := &KubernetesBackend{ config: config, clientset: clientset, + jobs: make(map[executionKey]*retainedJob), } if err := backend.runStartupPreflight(ctx); err != nil { return nil, err @@ -309,26 +327,16 @@ func (b *KubernetesBackend) ExecuteTask(ctx context.Context, params *TaskParams) if _, err := b.clientset.BatchV1().Jobs(b.config.Namespace).Create(ctx, job, metav1.CreateOptions{}); err != nil { return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonJobCreate, fmt.Errorf("failed to create Kubernetes Job: %w", err))) } + b.registerJob(params.TaskID, executionID, jobName) defer func() { - if ctx.Err() != nil { - log.Infof(ctx, "Leaving Kubernetes Job %s in place after task context cancellation", jobName) - return - } - if b.config.NoCleanup { - return - } - // Preserve failed task Jobs (and their pods) so operators can inspect logs - // and pod state after the fact. They are garbage-collected by the Job's - // TTLSecondsAfterFinished (see taskJobTTLSecondsAfterFinished). Successful - // Jobs are deleted immediately to keep the namespace clean. - if res.Error != nil { - log.Infof(ctx, "Leaving failed Kubernetes Job %s in place for TTL-based cleanup", jobName) - return - } - if err := b.deleteJob(context.Background(), jobName); err != nil { - log.Warnf(ctx, "Failed to delete Job %s: %v", jobName, err) - } + // The Job and its pods are the retained log source until the + // execution's cleanup grace expires, so nothing is deleted here. + // CleanupTaskResources applies the deletion policy at that deadline: + // a successful Job is deleted to keep the namespace clean, while a + // failed one is left for the Job TTL controller (see + // taskJobTTLSecondsAfterFinished) so operators can inspect it. + b.setJobDeleteAtCleanup(params.TaskID, executionID, res.Error == nil && ctx.Err() == nil) }() jobWatcher, err := b.watchJob(ctx, jobName) @@ -438,6 +446,232 @@ func (b *KubernetesBackend) ExecuteTask(ctx context.Context, params *TaskParams) // Kubernetes-backend task. func (b *KubernetesBackend) CancelTask(context.Context, *CancelParams) error { return nil } +func (b *KubernetesBackend) registerJob(taskID, executionID, jobName string) { + b.jobsMutex.Lock() + defer b.jobsMutex.Unlock() + if b.jobs == nil { + b.jobs = make(map[executionKey]*retainedJob) + } + b.jobs[executionKey{runID: taskID, executionID: executionID}] = &retainedJob{name: jobName} +} + +func (b *KubernetesBackend) setJobDeleteAtCleanup(taskID, executionID string, deleteAtCleanup bool) { + b.jobsMutex.Lock() + defer b.jobsMutex.Unlock() + if entry, ok := b.jobs[executionKey{runID: taskID, executionID: executionID}]; ok { + entry.deleteAtCleanup = deleteAtCleanup + } +} + +func (b *KubernetesBackend) lookupJob(taskID, executionID string) (retainedJob, bool) { + b.jobsMutex.Lock() + defer b.jobsMutex.Unlock() + entry, ok := b.jobs[executionKey{runID: taskID, executionID: executionID}] + if !ok { + return retainedJob{}, false + } + return *entry, true +} + +func (b *KubernetesBackend) forgetJob(taskID, executionID string) { + b.jobsMutex.Lock() + defer b.jobsMutex.Unlock() + delete(b.jobs, executionKey{runID: taskID, executionID: executionID}) +} + +func (b *KubernetesBackend) ownsJob(taskID, executionID string) bool { + b.jobsMutex.Lock() + defer b.jobsMutex.Unlock() + _, ok := b.jobs[executionKey{runID: taskID, executionID: executionID}] + return ok +} + +// SnapshotTaskLogs streams every container's stdout and stderr for the exact +// execution's pods into sink, in deterministic pod-name then declared-container +// order. Kubernetes merges a container's streams and does not say whether a +// line came from the entrypoint or the client it starts, so records are +// labeled combined rather than given a fabricated process identity. +func (b *KubernetesBackend) SnapshotTaskLogs(ctx context.Context, params *SnapshotParams) error { + if !b.ownsJob(params.TaskID, params.ExecutionID) { + return debuglog.NewSnapshotError(types.DebugArchiveReasonResourceNotFound, "no Job is registered for this execution") + } + + pods, err := b.listExecutionPods(ctx, params.TaskID, params.ExecutionID) + if err != nil { + return debuglog.NewSnapshotError(types.DebugArchiveReasonCaptureUnavailable, "failed to list the execution's pods") + } + if len(pods) == 0 { + return debuglog.NewSnapshotError(types.DebugArchiveReasonResourceNotFound, "the execution's pods are no longer present") + } + sort.Slice(pods, func(i, j int) bool { return pods[i].Name < pods[j].Name }) + + var warnings []string + noteWarning := func(code string) { + for _, existing := range warnings { + if existing == code { + return + } + } + warnings = append(warnings, code) + } + + for i := range pods { + pod := &pods[i] + for _, container := range pod.Spec.InitContainers { + if err := b.snapshotContainer(ctx, pod, container.Name, "init", params.Sink, noteWarning); err != nil { + return err + } + } + for _, container := range pod.Spec.Containers { + if err := b.snapshotContainer(ctx, pod, container.Name, "regular", params.Sink, noteWarning); err != nil { + return err + } + } + } + + if len(warnings) > 0 { + return &debuglog.PartialSnapshotError{WarningCodes: warnings} + } + return nil +} + +// listExecutionPods selects pods by the exact execution, task, and worker +// label hashes and re-verifies the returned labels, so a colliding label set +// from another worker or execution can never contribute log bytes. +func (b *KubernetesBackend) listExecutionPods(ctx context.Context, taskID, executionID string) ([]corev1.Pod, error) { + selector := strings.Join([]string{ + fmt.Sprintf("%s=%s", kubernetesExecutionHashLabel, kubernetesLabelHash(executionID)), + fmt.Sprintf("%s=%s", kubernetesTaskHashLabel, kubernetesLabelHash(taskID)), + fmt.Sprintf("%s=%s", kubernetesWorkerHashLabel, kubernetesLabelHash(b.config.WorkerID)), + }, ",") + + podList, err := b.clientset.CoreV1().Pods(b.config.Namespace).List(ctx, metav1.ListOptions{LabelSelector: selector}) + if err != nil { + return nil, err + } + + verified := make([]corev1.Pod, 0, len(podList.Items)) + for _, pod := range podList.Items { + if pod.Labels[kubernetesExecutionHashLabel] != kubernetesLabelHash(executionID) || + pod.Labels[kubernetesTaskHashLabel] != kubernetesLabelHash(taskID) || + pod.Labels[kubernetesWorkerHashLabel] != kubernetesLabelHash(b.config.WorkerID) { + continue + } + verified = append(verified, pod) + } + return verified, nil +} + +// snapshotContainer streams one container's current logs and, when it has +// restarted, its previous logs first. An unreadable stream becomes a safe +// source-error record plus a warning code so readable siblings still upload. +func (b *KubernetesBackend) snapshotContainer( + ctx context.Context, + pod *corev1.Pod, + containerName string, + containerType string, + sink debuglog.Sink, + noteWarning func(string), +) error { + source := debuglog.SourceIdentity{ + Namespace: b.objectNamespace(pod.Namespace), + Pod: pod.Name, + Container: containerName, + ContainerType: containerType, + } + if restarts, ok := containerRestartCount(pod, containerName); ok { + source.RestartAttempt = &restarts + if restarts > 0 { + previous := source + previous.Previous = true + if err := b.streamContainerLogs(ctx, pod.Name, containerName, true, previous, sink); err != nil { + noteWarning(debuglog.WarningPreviousLogsUnavailable) + if sinkErr := sink.WriteSourceError(debuglog.PhaseContainer, debuglog.StreamCombined, previous, debuglog.WarningPreviousLogsUnavailable); sinkErr != nil { + return sinkErr + } + } + } + } + + if err := b.streamContainerLogs(ctx, pod.Name, containerName, false, source, sink); err != nil { + noteWarning(debuglog.WarningContainerLogsUnavailable) + if sinkErr := sink.WriteSourceError(debuglog.PhaseContainer, debuglog.StreamCombined, source, debuglog.WarningContainerLogsUnavailable); sinkErr != nil { + return sinkErr + } + } + return nil +} + +// streamContainerLogs requests all of a container's logs with provider +// timestamps and no time-range filter, and streams them into sink under the +// snapshot's aggregate byte bound. +func (b *KubernetesBackend) streamContainerLogs( + ctx context.Context, + podName string, + containerName string, + previous bool, + source debuglog.SourceIdentity, + sink debuglog.Sink, +) error { + request := b.clientset.CoreV1().Pods(b.config.Namespace).GetLogs(podName, &corev1.PodLogOptions{ + Container: containerName, + Timestamps: true, + Previous: previous, + }) + stream, err := request.Stream(ctx) + if err != nil { + return err + } + defer func() { + if closeErr := stream.Close(); closeErr != nil { + log.Warnf(ctx, "Failed to close log stream for %s/%s: %v", podName, containerName, closeErr) + } + }() + + // The Kubernetes log API returns one merged stream per container, so the + // records are labeled combined rather than split into stdout and stderr. + return copyUnframedLogStream(stream, sink, source) +} + +func containerRestartCount(pod *corev1.Pod, containerName string) (int32, bool) { + for _, status := range pod.Status.InitContainerStatuses { + if status.Name == containerName { + return status.RestartCount, true + } + } + for _, status := range pod.Status.ContainerStatuses { + if status.Name == containerName { + return status.RestartCount, true + } + } + return 0, false +} + +// CleanupTaskResources applies the Job deletion policy once the execution's +// cleanup grace has expired. It is idempotent and never deletes a Job the +// worker chose to leave for TTL-based cleanup. +// +// The registry entry is dropped only once the deletion is confirmed (an +// already-absent Job counts as deleted). Forgetting it first would strand a +// Job the API server refused to delete, leaving nothing for the caller to +// retry and pushing its removal out to the Job TTL controller. +func (b *KubernetesBackend) CleanupTaskResources(ctx context.Context, params *CancelParams) error { + entry, ok := b.lookupJob(params.TaskID, params.ExecutionID) + if !ok { + return nil + } + if b.config.NoCleanup || !entry.deleteAtCleanup { + b.forgetJob(params.TaskID, params.ExecutionID) + return nil + } + + if err := b.deleteJob(ctx, entry.name); err != nil { + return fmt.Errorf("failed to delete Job %s: %w", entry.name, err) + } + b.forgetJob(params.TaskID, params.ExecutionID) + return nil +} + // Shutdown intentionally does not delete task Jobs. // // Kubernetes Jobs are the durable execution unit for this backend. During @@ -451,9 +685,13 @@ func (b *KubernetesBackend) Shutdown(ctx context.Context) { // TTLSecondsAfterFinished. It bounds how long finished Jobs that the worker // leaves in place - failed Jobs (kept for post-mortem debugging) and Jobs // orphaned by worker disruption - and their pods survive before the Kubernetes -// Job TTL controller deletes them. Successful Jobs are deleted immediately by the -// worker, so this does not delay their cleanup. Returns nil when cleanup is -// disabled, so those Jobs are retained indefinitely. +// Job TTL controller deletes them. +// +// A successful Job is deleted by the worker at its execution's cleanup-grace +// deadline, not at task completion, so its logs stay readable for a debug +// archive. This TTL must therefore be at least as long as the effective grace, +// or the controller deletes the pods first and the archive is partial. Returns +// nil when cleanup is disabled, so those Jobs are retained indefinitely. func (b *KubernetesBackend) taskJobTTLSecondsAfterFinished() *int32 { if b.config.NoCleanup { return nil diff --git a/internal/worker/kubernetes_logs_test.go b/internal/worker/kubernetes_logs_test.go new file mode 100644 index 0000000..c5c8f27 --- /dev/null +++ b/internal/worker/kubernetes_logs_test.go @@ -0,0 +1,221 @@ +package worker + +import ( + "context" + "errors" + "testing" + + "github.com/warpdotdev/oz-agent-worker/internal/debuglog" + "github.com/warpdotdev/oz-agent-worker/internal/types" + corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/client-go/kubernetes/fake" +) + +func executionPod(backend *KubernetesBackend, name, taskID, executionID string, restarts int32) *corev1.Pod { + return &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: name, + Namespace: backend.config.Namespace, + Labels: backend.baseLabels(taskID, executionID), + }, + Spec: corev1.PodSpec{ + InitContainers: []corev1.Container{{Name: "sidecar-init"}}, + Containers: []corev1.Container{{Name: "task"}}, + }, + Status: corev1.PodStatus{ + ContainerStatuses: []corev1.ContainerStatus{{Name: "task", RestartCount: restarts}}, + }, + } +} + +func newKubernetesLogBackend(t *testing.T, pods ...*corev1.Pod) *KubernetesBackend { + t.Helper() + backend := &KubernetesBackend{ + config: KubernetesBackendConfig{WorkerID: "worker-123", Namespace: "agents"}, + jobs: make(map[executionKey]*retainedJob), + } + client := fake.NewSimpleClientset() + for _, pod := range pods { + if _, err := client.CoreV1().Pods(backend.config.Namespace).Create(context.Background(), pod, metav1.CreateOptions{}); err != nil { + t.Fatalf("failed to seed pod %s: %v", pod.Name, err) + } + } + backend.clientset = client + return backend +} + +func TestKubernetesSnapshotRequiresARegisteredJob(t *testing.T) { + backend := newKubernetesLogBackend(t) + + err := backend.SnapshotTaskLogs(context.Background(), &SnapshotParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + Sink: &captureSink{}, + }) + + var snapshotErr *debuglog.SnapshotError + if !errors.As(err, &snapshotErr) { + t.Fatalf("error = %v, want a *debuglog.SnapshotError", err) + } + if snapshotErr.ReasonCode != types.DebugArchiveReasonResourceNotFound { + t.Fatalf("reason = %q, want %q", snapshotErr.ReasonCode, types.DebugArchiveReasonResourceNotFound) + } +} + +func TestKubernetesSnapshotReportsMissingPodsAsNotFound(t *testing.T) { + backend := newKubernetesLogBackend(t) + backend.registerJob("task-1", "execution-1", "oz-task-task-1-exec-ution-1") + + err := backend.SnapshotTaskLogs(context.Background(), &SnapshotParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + Sink: &captureSink{}, + }) + + var snapshotErr *debuglog.SnapshotError + if !errors.As(err, &snapshotErr) { + t.Fatalf("error = %v, want a *debuglog.SnapshotError", err) + } + if snapshotErr.ReasonCode != types.DebugArchiveReasonResourceNotFound { + t.Fatalf("reason = %q, want %q", snapshotErr.ReasonCode, types.DebugArchiveReasonResourceNotFound) + } +} + +func TestListExecutionPodsRequiresEveryIdentityLabel(t *testing.T) { + backend := &KubernetesBackend{ + config: KubernetesBackendConfig{WorkerID: "worker-123", Namespace: "agents"}, + } + wanted := executionPod(backend, "wanted", "task-1", "execution-1", 0) + + // A pod from another execution, another task, and another worker must all + // be excluded even though each shares one label with the target. + otherExecution := executionPod(backend, "other-execution", "task-1", "execution-2", 0) + otherTask := executionPod(backend, "other-task", "task-2", "execution-1", 0) + + otherWorker := &KubernetesBackend{ + config: KubernetesBackendConfig{WorkerID: "worker-999", Namespace: "agents"}, + } + foreignWorkerPod := executionPod(otherWorker, "other-worker", "task-1", "execution-1", 0) + + backend = newKubernetesLogBackend(t, wanted, otherExecution, otherTask, foreignWorkerPod) + + pods, err := backend.listExecutionPods(context.Background(), "task-1", "execution-1") + if err != nil { + t.Fatalf("listExecutionPods: %v", err) + } + if len(pods) != 1 { + names := make([]string, 0, len(pods)) + for _, pod := range pods { + names = append(names, pod.Name) + } + t.Fatalf("selected pods = %v, want only the exactly-matching pod", names) + } + if pods[0].Name != "wanted" { + t.Fatalf("selected pod = %q, want %q", pods[0].Name, "wanted") + } +} + +func TestKubernetesSnapshotVisitsPodsAndContainersDeterministically(t *testing.T) { + backend := &KubernetesBackend{ + config: KubernetesBackendConfig{WorkerID: "worker-123", Namespace: "agents"}, + } + // Seeded out of order so the snapshot's own sort is what produces the + // deterministic result. + second := executionPod(backend, "pod-b", "task-1", "execution-1", 0) + first := executionPod(backend, "pod-a", "task-1", "execution-1", 0) + backend = newKubernetesLogBackend(t, second, first) + backend.registerJob("task-1", "execution-1", "oz-task-task-1-exec-ution-1") + + sink := &captureSink{} + // The fake clientset returns a canned log body for every container, so + // each visited container contributes at least one chunk. + if err := backend.SnapshotTaskLogs(context.Background(), &SnapshotParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + Sink: sink, + }); err != nil { + t.Fatalf("SnapshotTaskLogs: %v", err) + } + + var visited []string + for _, chunk := range sink.chunks { + entry := chunk.Source.Pod + "/" + chunk.Source.Container + "/" + chunk.Source.ContainerType + if len(visited) == 0 || visited[len(visited)-1] != entry { + visited = append(visited, entry) + } + } + want := []string{ + "pod-a/sidecar-init/init", + "pod-a/task/regular", + "pod-b/sidecar-init/init", + "pod-b/task/regular", + } + if len(visited) != len(want) { + t.Fatalf("visited = %v, want %v", visited, want) + } + for i := range want { + if visited[i] != want[i] { + t.Fatalf("visited = %v, want %v", visited, want) + } + } + + for _, chunk := range sink.chunks { + if chunk.Stream != debuglog.StreamCombined { + t.Errorf("stream = %q, want %q: Kubernetes merges a container's streams", chunk.Stream, debuglog.StreamCombined) + } + if chunk.Source.Namespace != "agents" { + t.Errorf("namespace = %q, want agents", chunk.Source.Namespace) + } + } +} + +func TestKubernetesSnapshotAttemptsPreviousLogsForARestartedContainer(t *testing.T) { + backend := &KubernetesBackend{ + config: KubernetesBackendConfig{WorkerID: "worker-123", Namespace: "agents"}, + } + pod := executionPod(backend, "pod-a", "task-1", "execution-1", 2) + backend = newKubernetesLogBackend(t, pod) + backend.registerJob("task-1", "execution-1", "oz-task-task-1-exec-ution-1") + + sink := &captureSink{} + if err := backend.SnapshotTaskLogs(context.Background(), &SnapshotParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + Sink: sink, + }); err != nil { + t.Fatalf("SnapshotTaskLogs: %v", err) + } + + var sawPrevious, sawCurrent bool + for _, chunk := range sink.chunks { + if chunk.Source.Container != "task" { + continue + } + if chunk.Source.RestartAttempt == nil || *chunk.Source.RestartAttempt != 2 { + t.Errorf("restart_attempt = %v, want 2", chunk.Source.RestartAttempt) + } + if chunk.Source.Previous { + sawPrevious = true + } else { + sawCurrent = true + } + } + if !sawPrevious { + t.Error("a restarted container must also contribute its previous logs") + } + if !sawCurrent { + t.Error("a restarted container must still contribute its current logs") + } +} + +func TestKubernetesCleanupIsANoOpForAnUnknownExecution(t *testing.T) { + backend := newKubernetesLogBackend(t) + + if err := backend.CleanupTaskResources(context.Background(), &CancelParams{ + TaskID: "unknown-task", + ExecutionID: "unknown-execution", + }); err != nil { + t.Fatalf("CleanupTaskResources: %v", err) + } +} diff --git a/internal/worker/kubernetes_test.go b/internal/worker/kubernetes_test.go index 95017a8..ad258ea 100644 --- a/internal/worker/kubernetes_test.go +++ b/internal/worker/kubernetes_test.go @@ -1461,9 +1461,10 @@ func TestTaskJobTTLDefaultsToTwentyFourHours(t *testing.T) { } } -// On success the worker actively deletes the task Job (and its pod) so the -// namespace stays clean. -func TestExecuteTaskDeletesJobOnSuccess(t *testing.T) { +// A successful task Job is retained past terminal state so a debug-archive +// request can still read its pods' logs, then deleted at the cleanup-grace +// deadline to keep the namespace clean. +func TestExecuteTaskDeletesSuccessfulJobAtCleanupGrace(t *testing.T) { fakeClient := fake.NewSimpleClientset() jobWatch := watch.NewFake() podWatch := watch.NewFake() @@ -1511,6 +1512,7 @@ func TestExecuteTaskDeletesJobOnSuccess(t *testing.T) { if result := backend.ExecuteTask(context.Background(), &TaskParams{ TaskID: "task-1", + ExecutionID: "execution-1", DockerImage: "ubuntu:22.04", BaseArgs: []string{"run"}, }); result.Error != nil { @@ -1521,8 +1523,31 @@ func TestExecuteTaskDeletesJobOnSuccess(t *testing.T) { if err != nil { t.Fatalf("failed to list jobs: %v", err) } + if len(jobs.Items) != 1 { + t.Fatalf("expected the Job to be retained for log retrieval, got %d", len(jobs.Items)) + } + + if err := backend.CleanupTaskResources(context.Background(), &CancelParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + }); err != nil { + t.Fatalf("CleanupTaskResources: %v", err) + } + + jobs, err = fakeClient.BatchV1().Jobs("agents").List(context.Background(), metav1.ListOptions{}) + if err != nil { + t.Fatalf("failed to list jobs: %v", err) + } if len(jobs.Items) != 0 { - t.Fatalf("expected successful task Job to be deleted, got %d", len(jobs.Items)) + t.Fatalf("expected successful task Job to be deleted at the grace deadline, got %d", len(jobs.Items)) + } + + // A repeated expiry must not fail or delete anything else. + if err := backend.CleanupTaskResources(context.Background(), &CancelParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + }); err != nil { + t.Fatalf("repeated CleanupTaskResources: %v", err) } } @@ -1576,6 +1601,7 @@ func TestExecuteTaskPreservesJobOnFailure(t *testing.T) { result := backend.ExecuteTask(context.Background(), &TaskParams{ TaskID: "task-1", + ExecutionID: "execution-1", DockerImage: "ubuntu:22.04", BaseArgs: []string{"run"}, }) @@ -1590,6 +1616,22 @@ func TestExecuteTaskPreservesJobOnFailure(t *testing.T) { if len(jobs.Items) != 1 { t.Fatalf("expected failed task Job to be preserved, got %d", len(jobs.Items)) } + + // The grace deadline must not delete a failed Job either: it is left for + // the Job TTL controller so operators can inspect it. + if err := backend.CleanupTaskResources(context.Background(), &CancelParams{ + TaskID: "task-1", + ExecutionID: "execution-1", + }); err != nil { + t.Fatalf("CleanupTaskResources: %v", err) + } + jobs, listErr = fakeClient.BatchV1().Jobs("agents").List(context.Background(), metav1.ListOptions{}) + if listErr != nil { + t.Fatalf("failed to list jobs: %v", listErr) + } + if len(jobs.Items) != 1 { + t.Fatalf("expected failed task Job to survive the grace deadline, got %d", len(jobs.Items)) + } } func TestRunStartupPreflightCreatesLegacyRootInitJobAndWaitsForPodCreationByDefault(t *testing.T) { diff --git a/internal/worker/registry.go b/internal/worker/registry.go new file mode 100644 index 0000000..58a5fe0 --- /dev/null +++ b/internal/worker/registry.go @@ -0,0 +1,216 @@ +package worker + +import ( + "sync" + "time" + + "github.com/warpdotdev/oz-agent-worker/internal/debuglog" + "github.com/warpdotdev/oz-agent-worker/internal/metrics" +) + +// executionKey identifies one execution exactly. A run can be executed more +// than once (follow-ups, handoffs), so debug-archive ownership is only ever +// resolved through both identifiers. +type executionKey struct { + runID string + executionID string +} + +// ownedExecution is the worker's claim over one execution: the backend that +// ran it and the direct-backend capture, if any. +type ownedExecution struct { + backendKind string + capture *debuglog.TaskLogCapture + // cleanupTimer fires backend resource cleanup at the grace deadline. It is + // nil while the execution is still active, and is only ever touched under + // the registry's mutex: a zero grace fires the timer before the arming call + // returns, so its callback races an unguarded field. + cleanupTimer *time.Timer +} + +// TaskRegistry tracks the worker's in-progress tasks. It answers two +// questions: which task ID a cancellation should route to, and whether this +// process executed one exact (run, execution) pair — the ownership test that +// makes only the executing instance answer a debug-archive log request. +type TaskRegistry struct { + mu sync.Mutex + // active is keyed by task ID because cancellation messages carry only + // task_id. + active map[string]activeTask + // owned and grace are keyed by the exact execution identity a + // debug-archive request names. + owned map[executionKey]*ownedExecution + grace map[executionKey]*ownedExecution +} + +func newTaskRegistry() *TaskRegistry { + return &TaskRegistry{ + active: make(map[string]activeTask), + owned: make(map[executionKey]*ownedExecution), + grace: make(map[executionKey]*ownedExecution), + } +} + +// StartTask records an assignment as active and claims ownership of its exact +// execution identity. +func (r *TaskRegistry) StartTask(taskID string, task activeTask, backendKind string, capture *debuglog.TaskLogCapture) { + r.mu.Lock() + defer r.mu.Unlock() + + r.active[taskID] = task + r.owned[executionKey{runID: taskID, executionID: task.executionID}] = &ownedExecution{ + backendKind: backendKind, + capture: capture, + } +} + +// Get returns the active task for a task ID. +func (r *TaskRegistry) Get(taskID string) (activeTask, bool) { + r.mu.Lock() + defer r.mu.Unlock() + task, ok := r.active[taskID] + return task, ok +} + +// Update replaces an active task's record when it is still tracked. +func (r *TaskRegistry) Update(taskID string, mutate func(*activeTask)) (activeTask, bool) { + r.mu.Lock() + defer r.mu.Unlock() + + task, ok := r.active[taskID] + if !ok { + return activeTask{}, false + } + mutate(&task) + r.active[taskID] = task + return task, true +} + +// Delete stops tracking an active task. Any cleanup-grace entry survives, so a +// terminal execution stays reachable for log collection; only an execution +// that never reached cleanup grace loses its ownership here. +func (r *TaskRegistry) Delete(taskID string) { + r.mu.Lock() + defer r.mu.Unlock() + + task, ok := r.active[taskID] + if !ok { + return + } + delete(r.active, taskID) + delete(r.owned, executionKey{runID: taskID, executionID: task.executionID}) +} + +// Len reports how many tasks are active. +func (r *TaskRegistry) Len() int { + r.mu.Lock() + defer r.mu.Unlock() + return len(r.active) +} + +// Snapshot copies the active task map for iteration outside the lock. +func (r *TaskRegistry) Snapshot() map[string]activeTask { + r.mu.Lock() + defer r.mu.Unlock() + + out := make(map[string]activeTask, len(r.active)) + for taskID, task := range r.active { + out[taskID] = task + } + return out +} + +// LookupExecution implements debuglog.OwnershipLookup. A miss means this +// process did not execute the assignment, so the caller must stay silent. +func (r *TaskRegistry) LookupExecution(runID, executionID string) (debuglog.Ownership, bool) { + key := executionKey{runID: runID, executionID: executionID} + + r.mu.Lock() + defer r.mu.Unlock() + + if entry, ok := r.owned[key]; ok { + return debuglog.Ownership{BackendKind: entry.backendKind, Capture: entry.capture}, true + } + if entry, ok := r.grace[key]; ok { + return debuglog.Ownership{ + BackendKind: entry.backendKind, + InCleanupGrace: true, + Capture: entry.capture, + }, true + } + return debuglog.Ownership{}, false +} + +// MoveToCleanupGrace transfers an execution from active to cleanup-grace +// ownership and schedules onExpiry at the deadline. It runs before the terminal +// lifecycle message is enqueued, so a request the server triggers off that +// message always finds the grace entry rather than racing its deletion. +func (r *TaskRegistry) MoveToCleanupGrace(runID, executionID string, grace time.Duration, onExpiry func()) { + key := executionKey{runID: runID, executionID: executionID} + + r.mu.Lock() + entry, ok := r.owned[key] + if !ok { + r.mu.Unlock() + return + } + delete(r.owned, key) + r.grace[key] = entry + // The timer is armed while the lock is held because a zero grace fires + // onExpiry immediately; its callback blocks on this same lock until the + // entry is fully published. + entry.cleanupTimer = time.AfterFunc(grace, onExpiry) + graceCount := len(r.grace) + r.mu.Unlock() + + metrics.SetCleanupGraceEntries(graceCount) +} + +// ReleaseCleanupGrace drops an execution's cleanup-grace entry. It is +// idempotent so an expiry timer and a shutdown sweep can both call it. +func (r *TaskRegistry) ReleaseCleanupGrace(runID, executionID string) (*ownedExecution, bool) { + key := executionKey{runID: runID, executionID: executionID} + + r.mu.Lock() + entry, ok := r.grace[key] + if ok { + delete(r.grace, key) + if entry.cleanupTimer != nil { + entry.cleanupTimer.Stop() + } + } + graceCount := len(r.grace) + r.mu.Unlock() + + if !ok { + return nil, false + } + metrics.SetCleanupGraceEntries(graceCount) + return entry, true +} + +// ReleaseOwnership drops an execution's active ownership without moving it to +// cleanup grace. It covers the paths where no terminal state was ever reached, +// such as a rejected assignment. +func (r *TaskRegistry) ReleaseOwnership(runID, executionID string) { + r.mu.Lock() + defer r.mu.Unlock() + delete(r.owned, executionKey{runID: runID, executionID: executionID}) +} + +// PendingCleanups returns every cleanup-grace entry and clears the registry's +// record of them, so shutdown can release their resources exactly once. +func (r *TaskRegistry) PendingCleanups() map[executionKey]*ownedExecution { + r.mu.Lock() + entries := r.grace + r.grace = make(map[executionKey]*ownedExecution) + for _, entry := range entries { + if entry.cleanupTimer != nil { + entry.cleanupTimer.Stop() + } + } + r.mu.Unlock() + + metrics.SetCleanupGraceEntries(0) + return entries +} diff --git a/internal/worker/worker.go b/internal/worker/worker.go index 5bbc054..dddd30e 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -8,9 +8,11 @@ import ( "strings" "sync" "time" + "unicode/utf8" "github.com/gorilla/websocket" "github.com/warpdotdev/oz-agent-worker/internal/common" + "github.com/warpdotdev/oz-agent-worker/internal/debuglog" "github.com/warpdotdev/oz-agent-worker/internal/log" "github.com/warpdotdev/oz-agent-worker/internal/metrics" "github.com/warpdotdev/oz-agent-worker/internal/types" @@ -40,6 +42,14 @@ type Config struct { ServerRootURL string LogLevel string BackendType string // "docker", "direct", or "kubernetes" + // Version is the worker's build identifier, reported to warp-server on + // every authenticated WebSocket dial so the server can snapshot the exact + // build that claims an execution. + Version string + // DebugLogCapture bounds the disk and concurrency debug-archive log + // collection may use. Retention comes from the execution's existing + // idle-on-complete cleanup grace, not from this configuration. + DebugLogCapture debuglog.Config // MaxConcurrentTasks caps how many tasks may execute locally at once // (0 means unlimited). A task's slot is released when the backend's // ExecuteTask returns, so for backends that spawn tasks fire-and-forget @@ -68,13 +78,17 @@ type Worker struct { reconnectDelay time.Duration lastHeartbeat time.Time sendChan chan []byte - activeTasks map[string]activeTask - tasksMutex sync.Mutex + tasks *TaskRegistry backend Backend taskSemaphore *semaphore.Weighted // nil when unlimited // heartbeatInterval is how often the worker pings the server. It defaults // to HeartbeatInterval and is overridable in tests. heartbeatInterval time.Duration + // debugLogStore and debugLogs are nil when debug-log capture failed to + // initialize. That is non-fatal: ordinary task execution never depends on + // archive availability. + debugLogStore *debuglog.Store + debugLogs *debuglog.Coordinator } type taskCancellationSource string @@ -140,17 +154,48 @@ func New(ctx context.Context, config Config) (*Worker, error) { taskSemaphore = semaphore.NewWeighted(int64(config.MaxConcurrentTasks)) } - return &Worker{ + w := &Worker{ config: config, ctx: workerCtx, cancel: cancel, reconnectDelay: InitialReconnectDelay, sendChan: make(chan []byte, 256), - activeTasks: make(map[string]activeTask), + tasks: newTaskRegistry(), backend: backend, taskSemaphore: taskSemaphore, heartbeatInterval: HeartbeatInterval, - }, nil + } + + // Losing debug-log capture must never cost the operator task execution, so + // an invalid bound or an unwritable capture root degrades to "no archive + // capture" and the worker carries on. + store, err := debuglog.NewStore(config.DebugLogCapture) + if err != nil { + log.Errorf(ctx, "Debug archive log capture is disabled: %v", err) + return w, nil + } + w.debugLogStore = store + w.debugLogs = debuglog.NewCoordinator(debuglog.CoordinatorOptions{ + Ownership: w.tasks, + Source: backendSnapshotSource{backend: backend}, + Sender: w, + Store: store, + }) + return w, nil +} + +// backendSnapshotSource adapts the worker's backend to the coordinator's +// snapshot contract. +type backendSnapshotSource struct { + backend Backend +} + +func (s backendSnapshotSource) SnapshotLogs(ctx context.Context, runID, executionID string, sink debuglog.Sink) error { + return s.backend.SnapshotTaskLogs(ctx, &SnapshotParams{ + TaskID: runID, + ExecutionID: executionID, + Sink: sink, + }) } func (w *Worker) Start() error { @@ -195,6 +240,16 @@ func (w *Worker) connect() error { headers := make(map[string][]string) headers["Authorization"] = []string{fmt.Sprintf("Bearer %s", w.config.APIKey)} + // The version travels as connection metadata rather than a later message so + // it reaches the server before this connection is eligible for a task, and + // a reconnect re-reports it for future assignments. + if version, ok := workerVersionHeaderValue(w.config.Version); ok { + headers[types.WorkerVersionHeader] = []string{version} + } else if w.config.Version != "" { + // The value itself is never logged: an invalid build identifier is + // still attacker-influenced input. + log.Warnf(w.ctx, "Omitting invalid worker version metadata from the connection") + } log.Infof(w.ctx, "Connecting to %s", u.String()) @@ -346,7 +401,6 @@ func (w *Worker) handleMessage(message []byte) { return } - // Currently there is only one message type, but we anticipate needing more in the future. switch msg.Type { case types.MessageTypeTaskAssignment: var assignment types.TaskAssignmentMessage @@ -364,27 +418,80 @@ func (w *Worker) handleMessage(message []byte) { } w.handleTaskCancellation(&cancellation) + case types.MessageTypeDebugArchiveLogsRequested: + var request types.DebugArchiveLogsRequestedMessage + if err := json.Unmarshal(msg.Data, &request); err != nil { + // A request whose envelope will not parse cannot be attributed to + // an execution, so this process cannot know whether it owns it and + // must not answer. + log.Errorf(w.ctx, "Failed to unmarshal debug archive log request: %v", err) + return + } + if w.debugLogs == nil { + return + } + // Handing off to the coordinator keeps provider reads and uploads off + // this loop, so heartbeats, cancellations, and later assignments are + // still processed while a snapshot is in flight. + w.debugLogs.Handle(w.ctx, &request) + default: log.Warnf(w.ctx, "Unknown message type: %s", msg.Type) } } +// maxWorkerVersionBytes bounds the build identifier the worker reports. +const maxWorkerVersionBytes = 128 + +// workerVersionHeaderValue reports whether a build identifier is safe to send +// as connection metadata. It is an opaque, non-secret display value, so the +// only requirements are that it fits the bound, is valid UTF-8, and carries no +// control characters that could break header framing. An empty or invalid +// value is simply omitted: the server records provenance as not reported and +// the connection still executes tasks. +func workerVersionHeaderValue(version string) (string, bool) { + if version == "" || len(version) > maxWorkerVersionBytes || !utf8.ValidString(version) { + return "", false + } + for _, r := range version { + if r < 0x20 || r == 0x7f { + return "", false + } + } + return version, true +} + +// SendDebugArchiveAck enqueues a debug-archive acknowledgement through the +// worker's single WebSocket writer. +func (w *Worker) SendDebugArchiveAck(ack *types.DebugArchiveLogsUploadedMessage) error { + data, err := json.Marshal(ack) + if err != nil { + return fmt.Errorf("failed to marshal debug archive acknowledgement: %w", err) + } + + msgBytes, err := json.Marshal(types.WebSocketMessage{ + Type: types.MessageTypeDebugArchiveLogsUploaded, + Data: data, + }) + if err != nil { + return fmt.Errorf("failed to marshal websocket message: %w", err) + } + + return w.sendMessage(msgBytes) +} + func (w *Worker) handleTaskCancellation(cancellation *types.TaskCancellationMessage) { - w.tasksMutex.Lock() - task, ok := w.activeTasks[cancellation.TaskID] - if ok { + task, ok := w.tasks.Update(cancellation.TaskID, func(task *activeTask) { if task.cancellationSource == "" { task.cancellationSource = taskCancellationSourceUser - w.activeTasks[cancellation.TaskID] = task - } - if task.spawned { - // executeTask has already returned for a spawned task, so no - // deferred cleanup will remove the entry; drop it now that the - // cancellation is being routed to the backend. - delete(w.activeTasks, cancellation.TaskID) } + }) + if ok && task.spawned { + // executeTask has already returned for a spawned task, so no deferred + // cleanup will remove the entry; drop it now that the cancellation is + // being routed to the backend. + w.tasks.Delete(cancellation.TaskID) } - w.tasksMutex.Unlock() if !ok { log.Warnf(w.ctx, "Received cancellation for inactive task: taskID=%s", cancellation.TaskID) @@ -470,16 +577,41 @@ func (w *Worker) handleTaskAssignment(assignment *types.TaskAssignmentMessage) { } taskCtx, taskCancel := context.WithCancel(executionCtx) - w.tasksMutex.Lock() - w.activeTasks[assignment.TaskID] = activeTask{ + w.tasks.StartTask(assignment.TaskID, activeTask{ ctx: taskCtx, cancel: taskCancel, executionID: assignment.ExecutionID, - } - w.tasksMutex.Unlock() + }, w.backendKind(), w.newTaskLogCapture()) go w.executeTask(taskCtx, taskCancel, span, assignment, receivedAt) } +// backendKind is the resolved backend name reported on debug-archive records +// and acknowledgements. It normalizes the empty default to docker so ownership +// lookups and NDJSON records always name a real backend. +func (w *Worker) backendKind() string { + if w.config.BackendType == "" { + return debuglog.BackendDocker + } + return w.config.BackendType +} + +// newTaskLogCapture allocates a bounded output capture for a direct execution. +// Every other backend reads provider-native logs on demand instead of keeping +// a second copy, and a capture that cannot be allocated is recorded and +// skipped rather than failing the assignment. +func (w *Worker) newTaskLogCapture() *debuglog.TaskLogCapture { + if w.debugLogStore == nil || w.backendKind() != debuglog.BackendDirect { + return nil + } + capture, err := w.debugLogStore.NewTaskLogCapture(nil) + if err != nil { + log.Warnf(w.ctx, "Debug archive output capture is unavailable for this task: %v", err) + return nil + } + metrics.SetDebugArchiveCaptureBytes(w.debugLogStore.ReservedBytes()) + return capture +} + // prepareTaskParams converts a TaskAssignmentMessage into backend-agnostic TaskParams, // resolving common environment variables, default images, and base CLI arguments. func (w *Worker) prepareTaskParams(assignment *types.TaskAssignmentMessage) *TaskParams { @@ -614,13 +746,11 @@ func (w *Worker) executeTask(ctx context.Context, taskCancel context.CancelFunc, defer func() { taskCancel() span.End() - w.tasksMutex.Lock() - // Spawned tasks stay in activeTasks so a later cancellation can be - // routed to the backend's CancelTask; everything else is done. - if task, tracked := w.activeTasks[assignment.TaskID]; !tracked || !task.spawned { - delete(w.activeTasks, assignment.TaskID) + // Spawned tasks stay tracked so a later cancellation can be routed to + // the backend's CancelTask; everything else is done. + if task, tracked := w.tasks.Get(assignment.TaskID); !tracked || !task.spawned { + w.tasks.Delete(assignment.TaskID) } - w.tasksMutex.Unlock() if w.taskSemaphore != nil { w.taskSemaphore.Release(1) @@ -635,12 +765,24 @@ func (w *Worker) executeTask(ctx context.Context, taskCancel context.CancelFunc, metrics.AddTaskEvent(ctx, "task.started") params := w.prepareTaskParams(assignment) + if owner, owned := w.tasks.LookupExecution(assignment.TaskID, assignment.ExecutionID); owned { + params.LogCapture = owner.Capture + } metrics.AddTaskEvent(ctx, "backend.started", attribute.String("backend", w.config.BackendType), attribute.String("docker.image", params.DockerImage), ) executeResult := w.backend.ExecuteTask(ctx, params) + + // Ownership moves to cleanup grace before any terminal lifecycle message is + // enqueued. A server that reacts to task_failed by requesting logs then + // always finds the grace entry instead of racing registry deletion and + // backend cleanup, which is the whole point of the ANY_FAILURE path. + if executeResult.Outcome != ExecuteOutcomeSpawned { + w.beginCleanupGrace(assignment) + } + if executeResult.Error != nil { err := executeResult.Error if ctx.Err() == context.Canceled && w.cancellationSource(taskID) == taskCancellationSourceUser { @@ -687,12 +829,11 @@ func (w *Worker) executeTask(ctx context.Context, taskCancel context.CancelFunc, // so that cancellation can be routed to the backend's CancelTask // implementation later. result = metrics.TaskResultDispatched - w.tasksMutex.Lock() - if task, tracked := w.activeTasks[taskID]; tracked && task.cancellationSource == "" { - task.spawned = true - w.activeTasks[taskID] = task - } - w.tasksMutex.Unlock() + w.tasks.Update(taskID, func(task *activeTask) { + if task.cancellationSource == "" { + task.spawned = true + } + }) metrics.AddTaskEvent(ctx, "task.dispatched") span.SetStatus(codes.Ok, "task dispatched to remote runtime") log.Infof(ctx, "Task %s dispatched", taskID) @@ -708,15 +849,57 @@ func (w *Worker) executeTask(ctx context.Context, taskCancel context.CancelFunc, } func (w *Worker) cancellationSource(taskID string) taskCancellationSource { - w.tasksMutex.Lock() - defer w.tasksMutex.Unlock() - - task, ok := w.activeTasks[taskID] + task, ok := w.tasks.Get(taskID) if !ok { return "" } return task.cancellationSource } + +// beginCleanupGrace retains the execution's backend resources and output +// capture for the same window the agent itself stays idle, then releases them. +// Reusing the resolved idle-on-complete duration keeps one cleanup clock: +// operators size retention with the setting they already tune, and an archive +// request never extends it. +func (w *Worker) beginCleanupGrace(assignment *types.TaskAssignmentMessage) { + grace := common.ResolveCleanupGrace(assignment.Task, w.config.IdleOnComplete) + runID := assignment.TaskID + executionID := assignment.ExecutionID + + w.tasks.MoveToCleanupGrace(runID, executionID, grace, func() { + w.expireCleanupGrace(runID, executionID) + }) +} + +// expireCleanupGrace performs the backend's normal resource cleanup and frees +// the execution's capture. It is idempotent: the registry hands the entry over +// exactly once, so a racing shutdown sweep finds nothing left to do. +func (w *Worker) expireCleanupGrace(runID, executionID string) { + entry, ok := w.tasks.ReleaseCleanupGrace(runID, executionID) + if !ok { + return + } + + ctx, cancel := context.WithTimeout(context.WithoutCancel(w.ctx), BackendShutdownTimeout) + defer cancel() + + result := "succeeded" + if err := w.backend.CleanupTaskResources(ctx, &CancelParams{TaskID: runID, ExecutionID: executionID}); err != nil { + result = "failed" + log.Warnf(w.ctx, "Backend cleanup failed after cleanup grace for task %s: %v", runID, err) + } + metrics.RecordCleanupGraceResult(entry.backendKind, result) + + if entry.capture != nil { + entry.capture.Close() + } + if w.debugLogStore != nil { + metrics.SetDebugArchiveCaptureBytes(w.debugLogStore.ReservedBytes()) + } + if w.debugLogs != nil { + w.debugLogs.ForgetExecution(runID, executionID) + } +} func (w *Worker) sendTaskClaimed(taskID string) error { claimed := types.TaskClaimedMessage{ TaskID: taskID, @@ -852,21 +1035,100 @@ func (w *Worker) sendMessage(message []byte) error { } } +// releaseCleanupGraceEntries performs the backend cleanup each cleanup-grace +// entry was waiting for and deletes the bytes held for log retrieval. +// +// These executions have already reported terminal state; only their log source +// is being retained. Ownership is process-local, so a replacement worker cannot +// inherit the pending timer — leaving the resources behind would strand them +// until an unrelated backstop (the Kubernetes Job TTL) eventually collected +// them, far past the operator's chosen cleanup grace. Running the cleanup early +// here is the same work the expiry timer would have done. Active executions are +// untouched: they are not in cleanup grace, so each backend's own shutdown +// contract still decides whether their task units may outlive this process. +func (w *Worker) releaseCleanupGraceEntries() { + pending := w.tasks.PendingCleanups() + if len(pending) == 0 { + return + } + + log.Infof(w.ctx, "Releasing %d cleanup-grace executions during worker shutdown", len(pending)) + + // Each resource gets its own budget and runs concurrently. Sharing one + // budget across the sweep let a single slow backend call consume it and + // starve every entry behind it, which on a busy worker is exactly when + // there is the most to release. Per-resource budgets keep the whole sweep + // bounded by one timeout regardless of how many entries there are. + var wg sync.WaitGroup + for key, entry := range pending { + wg.Add(1) + go func(key executionKey, entry *ownedExecution) { + defer wg.Done() + w.releaseCleanupGraceEntry(key, entry) + }(key, entry) + } + wg.Wait() + + if w.debugLogStore != nil { + metrics.SetDebugArchiveCaptureBytes(w.debugLogStore.ReservedBytes()) + } +} + +// shutdownCleanupAttempts bounds how many times shutdown retries one backend +// resource. A transient API error should not cost the resource its deletion, +// but shutdown cannot retry indefinitely either. +const shutdownCleanupAttempts = 3 + +// releaseCleanupGraceEntry performs one execution's backend cleanup under its +// own budget, retrying a transient failure, and then frees its capture. +func (w *Worker) releaseCleanupGraceEntry(key executionKey, entry *ownedExecution) { + // The capture is local disk this process owns. Releasing it unconditionally + // is safe: a replacement worker's startup sweep removes any file left + // behind, so it can never accumulate the way a provider resource can. + defer func() { + if entry.capture != nil { + entry.capture.Close() + } + }() + + ctx, cancel := context.WithTimeout(context.Background(), BackendShutdownTimeout) + defer cancel() + + params := &CancelParams{TaskID: key.runID, ExecutionID: key.executionID} + var err error + for attempt := 1; attempt <= shutdownCleanupAttempts; attempt++ { + if err = w.backend.CleanupTaskResources(ctx, params); err == nil { + metrics.RecordCleanupGraceResult(entry.backendKind, "succeeded") + return + } + if ctx.Err() != nil { + break + } + } + + // The backend keeps the resource registered when deletion is unconfirmed, + // so its own shutdown still gets a final attempt. Naming what was left + // behind gives an operator something to act on if that attempt also fails. + metrics.RecordCleanupGraceResult(entry.backendKind, "failed") + log.Warnf(w.ctx, "Backend cleanup failed during shutdown for task %s; the resource remains registered for the backend's own shutdown: %v", key.runID, err) +} + func (w *Worker) Shutdown() { log.Infof(w.ctx, "Shutting down worker...") preserveActiveTasks := w.backend.PreservesTasksOnShutdown() - w.tasksMutex.Lock() - activeTaskCount := len(w.activeTasks) + active := w.tasks.Snapshot() + activeTaskCount := len(active) if activeTaskCount > 0 && preserveActiveTasks { log.Infof(w.ctx, "Preserving %d active tasks during worker shutdown", activeTaskCount) } else if activeTaskCount > 0 { log.Infof(w.ctx, "Cancelling %d active tasks", activeTaskCount) - for taskID, task := range w.activeTasks { - if task.cancellationSource == "" { - task.cancellationSource = taskCancellationSourceShutdown - w.activeTasks[taskID] = task - } + for taskID, task := range active { + w.tasks.Update(taskID, func(tracked *activeTask) { + if tracked.cancellationSource == "" { + tracked.cancellationSource = taskCancellationSourceShutdown + } + }) log.Debugf(w.ctx, "Cancelling task: %s", taskID) metrics.AddTaskEvent(task.ctx, "task.cancellation_requested", attribute.String("source", "signal"), @@ -875,13 +1137,20 @@ func (w *Worker) Shutdown() { task.cancel() } } - w.tasksMutex.Unlock() if activeTaskCount > 0 && !preserveActiveTasks { time.Sleep(500 * time.Millisecond) } w.cancel() + + // Cancelling the worker context aborts in-flight archive requests, so + // waiting for them adds no delay beyond the bounded backend shutdown below. + if w.debugLogs != nil { + w.debugLogs.Wait() + } + w.releaseCleanupGraceEntries() + backendShutdownCtx, backendShutdownCancel := context.WithTimeout(context.Background(), BackendShutdownTimeout) defer backendShutdownCancel() w.backend.Shutdown(backendShutdownCtx) diff --git a/internal/worker/worker_test.go b/internal/worker/worker_test.go index f744836..40b4590 100644 --- a/internal/worker/worker_test.go +++ b/internal/worker/worker_test.go @@ -21,6 +21,7 @@ import ( ) type shutdownRecordingBackend struct { + noLogSnapshotBackend shutdownCalled bool shutdownCtxErr error } @@ -48,6 +49,7 @@ func (b *preservingShutdownRecordingBackend) PreservesTasksOnShutdown() bool { } type recordingBackend struct { + noLogSnapshotBackend err error } @@ -186,10 +188,10 @@ func TestExecuteTaskReportsGracefulShutdownOnWorkerShutdown(t *testing.T) { ctx: context.Background(), config: Config{}, sendChan: make(chan []byte, 1), - activeTasks: map[string]activeTask{"task-1": { + tasks: registryWith(map[string]activeTask{"task-1": { cancel: func() {}, cancellationSource: taskCancellationSourceShutdown, - }}, + }}), backend: &recordingBackend{err: newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonTaskCancelled, context.Canceled)}, } w.executeTask(context.Background(), func() {}, trace.SpanFromContext(context.Background()), &types.TaskAssignmentMessage{ @@ -217,10 +219,10 @@ func TestExecuteTaskReportsTaskCancelledOnUserCancellation(t *testing.T) { ctx: context.Background(), config: Config{}, sendChan: make(chan []byte, 1), - activeTasks: map[string]activeTask{"task-1": { + tasks: registryWith(map[string]activeTask{"task-1": { cancel: func() {}, cancellationSource: taskCancellationSourceUser, - }}, + }}), backend: &recordingBackend{err: context.Canceled}, } w.executeTask(taskCtx, func() {}, trace.SpanFromContext(taskCtx), &types.TaskAssignmentMessage{ @@ -246,18 +248,18 @@ func TestExecuteTaskReportsTaskCancelledOnUserCancellation(t *testing.T) { if completed.Message != "Task cancelled by user request." { t.Errorf("message = %q, want %q", completed.Message, "Task cancelled by user request.") } - if _, ok := w.activeTasks["task-1"]; ok { + if _, ok := w.tasks.Get("task-1"); ok { t.Fatal("task should be removed from active tasks") } } func TestExecuteTaskDoesNotReportTaskCancelledOnBackendCancellationError(t *testing.T) { w := &Worker{ - ctx: context.Background(), - config: Config{}, - sendChan: make(chan []byte, 1), - activeTasks: map[string]activeTask{"task-1": {cancel: func() {}}}, - backend: &recordingBackend{err: fmt.Errorf("backend request failed: %w", context.Canceled)}, + ctx: context.Background(), + config: Config{}, + sendChan: make(chan []byte, 1), + tasks: registryWith(map[string]activeTask{"task-1": {cancel: func() {}}}), + backend: &recordingBackend{err: fmt.Errorf("backend request failed: %w", context.Canceled)}, } w.executeTask(context.Background(), func() {}, trace.SpanFromContext(context.Background()), &types.TaskAssignmentMessage{ @@ -278,12 +280,12 @@ func TestHandleMessageCancelsActiveTask(t *testing.T) { w := &Worker{ ctx: context.Background(), sendChan: make(chan []byte, 1), - activeTasks: map[string]activeTask{ + tasks: registryWith(map[string]activeTask{ "task-1": { ctx: taskCtx, cancel: taskCancel, }, - }, + }), backend: &recordingBackend{}, } @@ -304,18 +306,18 @@ func TestHandleMessageCancelsActiveTask(t *testing.T) { if taskCtx.Err() != context.Canceled { t.Fatalf("task context error = %v, want %v", taskCtx.Err(), context.Canceled) } - if task := w.activeTasks["task-1"]; task.cancellationSource != taskCancellationSourceUser { + if task, _ := w.tasks.Get("task-1"); task.cancellationSource != taskCancellationSourceUser { t.Fatalf("task cancellation source = %q, want %q", task.cancellationSource, taskCancellationSourceUser) } } func TestExecuteTaskReportsTaskCompletedOnSuccess(t *testing.T) { w := &Worker{ - ctx: context.Background(), - config: Config{}, - sendChan: make(chan []byte, 1), - activeTasks: map[string]activeTask{"task-1": {cancel: func() {}}}, - backend: &recordingBackend{}, + ctx: context.Background(), + config: Config{}, + sendChan: make(chan []byte, 1), + tasks: registryWith(map[string]activeTask{"task-1": {cancel: func() {}}}), + backend: &recordingBackend{}, } w.executeTask(context.Background(), func() {}, trace.SpanFromContext(context.Background()), &types.TaskAssignmentMessage{ @@ -338,18 +340,18 @@ func TestExecuteTaskReportsTaskCompletedOnSuccess(t *testing.T) { if completed.Message != "Task completed successfully" { t.Errorf("message = %q, want %q", completed.Message, "Task completed successfully") } - if _, ok := w.activeTasks["task-1"]; ok { + if _, ok := w.tasks.Get("task-1"); ok { t.Fatal("task should be removed from active tasks") } } func TestExecuteTaskReportsTaskFailedOnBackendError(t *testing.T) { w := &Worker{ - ctx: context.Background(), - config: Config{}, - sendChan: make(chan []byte, 1), - activeTasks: map[string]activeTask{"task-1": {cancel: func() {}}}, - backend: &recordingBackend{err: errors.New("boom")}, + ctx: context.Background(), + config: Config{}, + sendChan: make(chan []byte, 1), + tasks: registryWith(map[string]activeTask{"task-1": {cancel: func() {}}}), + backend: &recordingBackend{err: errors.New("boom")}, } w.executeTask(context.Background(), func() {}, trace.SpanFromContext(context.Background()), &types.TaskAssignmentMessage{ @@ -372,18 +374,18 @@ func TestExecuteTaskReportsTaskFailedOnBackendError(t *testing.T) { if failed.Message != "Failed to execute task: boom" { t.Errorf("message = %q, want %q", failed.Message, "Failed to execute task: boom") } - if _, ok := w.activeTasks["task-1"]; ok { + if _, ok := w.tasks.Get("task-1"); ok { t.Fatal("task should be removed from active tasks") } } func TestExecuteTaskReportsUserFriendlyMessageOnDeadlineExceeded(t *testing.T) { w := &Worker{ - ctx: context.Background(), - config: Config{}, - sendChan: make(chan []byte, 1), - activeTasks: map[string]activeTask{"task-1": {cancel: func() {}}}, - backend: &recordingBackend{err: context.DeadlineExceeded}, + ctx: context.Background(), + config: Config{}, + sendChan: make(chan []byte, 1), + tasks: registryWith(map[string]activeTask{"task-1": {cancel: func() {}}}), + backend: &recordingBackend{err: context.DeadlineExceeded}, } w.executeTask(context.Background(), func() {}, trace.SpanFromContext(context.Background()), &types.TaskAssignmentMessage{ @@ -478,7 +480,7 @@ func TestRunHeartbeatAndWritesAreConcurrencySafe(t *testing.T) { return } defer close(serverConnClosed) - defer conn.Close() + defer func() { _ = conn.Close() }() conn.SetPingHandler(func(string) error { pingsReceived.Add(1) return nil @@ -498,7 +500,7 @@ func TestRunHeartbeatAndWritesAreConcurrencySafe(t *testing.T) { t.Fatalf("failed to dial test server: %v", err) } if resp != nil && resp.Body != nil { - resp.Body.Close() + _ = resp.Body.Close() } ctx, cancel := context.WithCancel(context.Background()) @@ -509,7 +511,7 @@ func TestRunHeartbeatAndWritesAreConcurrencySafe(t *testing.T) { ctx: ctx, cancel: cancel, sendChan: make(chan []byte, 256), - activeTasks: make(map[string]activeTask), + tasks: newTaskRegistry(), heartbeatInterval: time.Millisecond, } @@ -535,7 +537,7 @@ flood: // Tear down: cancelling the context stops writeLoop/heartbeatLoop, and // closing the connection unblocks readLoop so run() returns. cancel() - conn.Close() + _ = conn.Close() select { case <-runDone: case <-time.After(5 * time.Second): @@ -976,10 +978,10 @@ func TestWorkerShutdownUsesFreshContextForBackendCleanup(t *testing.T) { workerCtx, cancel := context.WithCancel(context.Background()) backend := &shutdownRecordingBackend{} w := &Worker{ - ctx: workerCtx, - cancel: cancel, - activeTasks: make(map[string]activeTask), - backend: backend, + ctx: workerCtx, + cancel: cancel, + tasks: newTaskRegistry(), + backend: backend, } w.Shutdown() @@ -999,11 +1001,11 @@ func TestWorkerShutdownPreservesActiveTasksForPreservingBackend(t *testing.T) { w := &Worker{ ctx: workerCtx, cancel: cancel, - activeTasks: map[string]activeTask{ + tasks: registryWith(map[string]activeTask{ "task-1": {cancel: func() { cancelledTask = true }}, - }, + }), backend: backend, } @@ -1021,12 +1023,12 @@ func TestHandleTaskAssignmentDoesNotStartTaskAfterShutdownDuringClaim(t *testing workerCtx, cancel := context.WithCancel(context.Background()) cancel() w := &Worker{ - ctx: workerCtx, - cancel: cancel, - config: Config{}, - sendChan: make(chan []byte, 1), - activeTasks: make(map[string]activeTask), - backend: &preservingShutdownRecordingBackend{}, + ctx: workerCtx, + cancel: cancel, + config: Config{}, + sendChan: make(chan []byte, 1), + tasks: newTaskRegistry(), + backend: &preservingShutdownRecordingBackend{}, } w.handleTaskAssignment(&types.TaskAssignmentMessage{ @@ -1034,7 +1036,7 @@ func TestHandleTaskAssignmentDoesNotStartTaskAfterShutdownDuringClaim(t *testing Task: &types.Task{ID: "task-1", Title: "test task"}, }) - if len(w.activeTasks) != 0 { - t.Fatalf("expected no active tasks to start after shutdown during claim, got %d", len(w.activeTasks)) + if w.tasks.Len() != 0 { + t.Fatalf("expected no active tasks to start after shutdown during claim, got %d", w.tasks.Len()) } } diff --git a/main.go b/main.go index 424396f..1504c57 100644 --- a/main.go +++ b/main.go @@ -11,6 +11,7 @@ import ( "github.com/alecthomas/kong" "github.com/warpdotdev/oz-agent-worker/internal/config" + "github.com/warpdotdev/oz-agent-worker/internal/debuglog" "github.com/warpdotdev/oz-agent-worker/internal/log" "github.com/warpdotdev/oz-agent-worker/internal/metrics" "github.com/warpdotdev/oz-agent-worker/internal/worker" @@ -178,9 +179,11 @@ func mergeConfig(fileConfig *config.FileConfig) (worker.Config, error) { ServerRootURL: CLI.ServerRootURL, LogLevel: CLI.LogLevel, BackendType: backendType, + Version: Version, MaxConcurrentTasks: maxConcurrentTasks, IdleOnComplete: idleOnComplete, SessionSharingServerURL: CLI.SessionSharingServerURL, + DebugLogCapture: debugLogCaptureConfig(fileConfig), } switch backendType { @@ -404,6 +407,29 @@ func parseEnvFlags(raw []string) (map[string]string, error) { return result, nil } +// debugLogCaptureConfig overlays any operator-supplied debug-log capture bounds +// onto the built-in defaults. Invalid values are rejected later by the worker, +// which disables archive capture without failing task execution. +func debugLogCaptureConfig(fileConfig *config.FileConfig) debuglog.Config { + captureConfig := debuglog.DefaultConfig() + if fileConfig == nil || fileConfig.DebugLogCapture == nil { + return captureConfig + } + + overrides := fileConfig.DebugLogCapture + captureConfig.Directory = overrides.Directory + if overrides.MaxTotalBytes != nil { + captureConfig.MaxTotalBytes = *overrides.MaxTotalBytes + } + if overrides.MaxExecutionBytes != nil { + captureConfig.MaxExecutionBytes = *overrides.MaxExecutionBytes + } + if overrides.MaxConcurrentUploads != nil { + captureConfig.MaxConcurrentUploads = *overrides.MaxConcurrentUploads + } + return captureConfig +} + func copyStringMap(values map[string]string) map[string]string { if len(values) == 0 { return nil