5 Commits

Author SHA1 Message Date
Nan Gao
c542185a7f
feat(extensions): add gateway contribution points and packaged extension management (#4780)
* feat(extensions): add gateway services and routers

* feat(extensions): add standalone reference extension

* fix(extensions): harden contributed gateway routes

* docs(extensions): document gateway contribution points

* feat(extensions): add operator CLI for packaged extension management

Add `deerflow extensions install/list/enable/disable/remove` plus the root
`make extension-*` wrappers, backed by an `ExtensionManager` that owns one
transaction over backend/pyproject.toml, backend/uv.lock, the managed source
snapshot, the uv environment, and the `plugins:` block in config.yaml.

Install accepts a package requirement, a public HTTPS Git URL, or a local
directory. Local directories are copied to backend/extensions/sources/ as
deployable snapshots rather than editable installs, and the root .dockerignore
re-includes that tree so snapshots reach the backend builder. Remote sources are
HTTPS-only; SSH Git, file:// and local wheels are rejected because the stock
Docker builder cannot reproduce them.

Because environment configuration can still resolve a plain package name to a
local wheel (a UV_FIND_LINKS wheelhouse, say), every uv add/remove is followed
by an audit of the new lock: any local reference the stock image build cannot
reproduce rolls back the whole transaction. A config carrying duplicate
top-level `plugins:` keys is rejected outright rather than managed against one
block while the Gateway reads another.

Dependency synchronization now has one lock authority. The `extensions`
dependency group joins [tool.uv].default-groups, every startup path syncs the
same lock with --locked and launches with --no-sync, and the Docker images move
to uv 0.11.1 for the --no-workspace boundary the manager needs.

Loader gains `enabled`, `name` and `package` fields so a disabled extension is
skipped before resolution and import.

Co-authored-by: Codex <codex@openai.com>

* fix(extensions): stop the managed plugins rewrite from destroying config

Two data-safety defects in the managed `plugins:` block writer.

The "next top-level key" boundary was a regex matching only
`[A-Za-z_][A-Za-z0-9_-]*` or a quoted key. `AppConfig` is `extra="allow"`, so a
config may legally carry any top-level key, and a key the pattern cannot
recognize did not fail loudly — it read as "no next section", and the rewrite
replaced that neighbour and its entire subtree with the managed block. `my.key`,
`2fa`, `$schema`, `my key` and non-ASCII keys were all silently deleted by a
plain `extension-enable`/`disable`. Both boundaries now come from the YAML
parser's node marks, so key shape is irrelevant.

The file-final branch never consulted the trailing-comment scan the has-next-key
branch used, so any comment below the block was dropped. Since the manager
appends `plugins:` at end of file, that is the steady-state shape for most
installs: an operator note below the block was destroyed on the next toggle.

Separately, every managed install wrote `required: true` while the loader
defaults to false. That turned any later load failure — broken wheel, missing
native library, deleted snapshot — into a Gateway startup abort recoverable only
with shell access. New records are now written `required: false`, with an
explicit `install --required` opt-in; adopting an existing hand-written record
still preserves the operator's own choice.

* fix(extensions): harden the manager transaction and correct its docs

Follow-up hardening on the extension package manager.

Security posture, which the docs already claimed:
- Scrub `UV_PYTHON`, `UV_INSECURE_HOST`, `UV_CONSTRAINT` and
  `UV_NO_BUILD_ISOLATION` from the controlled uv environment. `UV_PYTHON` swaps
  the interpreter that the entry-point probe then imports and calls, and every
  later `uv run --no-sync` startup uses; `UV_INSECURE_HOST` removes the TLS
  validation the HTTPS-only source rule depends on. Neither is an index, proxy,
  cache or credential-provider setting, so neither was covered by the carve-out.
- Recognize run-together and all-caps secret query parameters (`accesstoken`,
  `ACCESSTOKEN`, `key`, `pw`, `sas`, `code`). The camel-case splitter only fires
  on case transitions, so only the separated spellings were caught. Short
  generic words stay boundary-anchored, so `?keyword=` remains installable.
- Validate the config before running any uv command. `uv add`/`uv sync` execute
  the package's build backend, so a config the manager could never write to must
  fail before that code runs rather than afterwards through rollback.

Transaction integrity:
- Run the second dependency-file restore from a `finally`. The recovery sync
  runs without `--locked` when the checkout had no lock, so uv writes one while
  resolving; if that sync then failed, the restore was skipped and the operator
  kept a lock file they never had. A failing recovery sync now also reports the
  original failure instead of replacing it.
- Skip the recovery sync on cancellation. Answering Ctrl-C with a full
  dependency resolve invites a second interrupt that escapes the handler and
  strands the checkout mid-transaction; the declarations are already restored
  and the next locked startup sync reconciles the environment.
- Retry a non-blocking lock on Windows instead of using `msvcrt.LK_LOCK`, which
  gives up after ~10s — far shorter than a real `uv add` plus `uv sync`, so
  contention surfaced as `Permission denied` rather than serializing.
- Locate the entry-point probe's JSON payload instead of parsing stdout's first
  line, so a `sitecustomize`/`.pth` banner cannot roll back a good install.
- Warn when the lock records a loopback source. `127.0.0.1` inside the image
  builder is a different machine, but unlike an environment-driven wheelhouse
  resolution this is a source the operator typed deliberately, so it is reported
  rather than rolled back. Private-network indexes are untouched: a builder on
  that network can reach them.

Docs: the blanket claim that failed operations restore the config file was
wrong — the conflict branches deliberately preserve a concurrent external edit
and leave `remove` deactivated. Document that, the `required: false` default,
the config preflight, the interrupt behaviour, and where the plugins-block
boundaries come from.

* test(gateway): pin the request-path projection agreement

`get_request_route_path()` imports the private
`starlette._utils.get_route_path` so the auth and CSRF predicates classify
the exact string Starlette's router matches on. Its requirement is not
"strip root_path correctly" but "return what the dispatcher is matching",
so delegating to the router's own implementation keeps the two in lockstep
by construction. Keep the private import rather than vendoring a copy: an
import that disappears fails loudly at startup, while a stale copy diverges
silently at a security boundary.

Cover the property directly instead of the mechanism, so the tests survive
a future reimplementation:

- projection edge cases, including the segment-boundary guard that keeps
  root_path="/api" from slicing "/apifoo/models" into a string the router
  would never match
- agreement with the router under nested mounts
- the two bypasses these predicates exist to prevent: a protected route
  mounted under the "/health" public prefix must still 401, and a POST
  mounted under "/api/webhooks" must still require a CSRF token

Both are verified to fail when the projection is reverted to
`request.url.path` (9/13 red) and when a plausible vendored copy omits the
boundary guard (the 2 boundary cases red).

Declare starlette as a bounded direct dependency so a bump — which is
security-relevant here — shows up in review rather than arriving silently
through FastAPI.

* ci: pin uv to the version production ships

ExtensionManager is not a consumer of uv the build tool -- it is a program
whose whole job is driving `uv` as a subprocess, depending on its CLI
behavior (`--no-workspace`, `--no-sync`, what `uv add` writes into
`[dependency-groups] extensions`) and on the `uv.lock` serialization format.
uv is closer to a runtime dependency with a contract than to incidental
tooling.

backend/Dockerfile pins that binary to 0.11.1, but all eight
astral-sh/setup-uv steps installed whatever was latest at run time, so CI
exercised the manager against a uv that is not the uv production runs. The
sharpest failure that allows: a newer uv bumps uv.lock's `revision`, CI
stays green because the same uv reads back what it wrote, and the pinned uv
in the production image cannot read the committed lock. `uv lock --check`
is version-sensitive for the same reason -- it verifies the lock is what
*this* uv would produce, and two versions can emit equivalent but
non-identical output.

Pin every step to 0.11.1 and lift the one lingering setup-uv@v3 to v7 so
the steps share input and caching behavior.

Pinning alone drifts apart again on the next bump, so add a constraint test
in the style of test_compose_default_bind_host.py: the Dockerfile's
UV_IMAGE tag is the single source of truth, and both compose defaults plus
every setup-uv step must match it. Verified to fail when a pin drifts, when
a step omits `version`, and -- the real scenario -- when the Dockerfile is
bumped alone, which lights up the workflows and both compose files at once.

* fix(gateway): state the extension route auth limit and abort a failed dev sync

Two scoped review follow-ups.

README: contributed routers cannot enter the host's reserved public prefixes,
which makes every extension endpoint session-authenticated -- there is no way
to expose an unauthenticated route. The rejection rule was documented but its
consequence was not, so inbound provider webhooks and public status endpoints
read as merely undocumented rather than out of scope for this release.

docker/dev-entrypoint.sh: the self-heal retry reuses `--locked`, so it repairs
a corrupt .venv but never a lock that disagrees with pyproject.toml. `set -e`
already stopped the script there -- uvicorn was not being started against a
stale environment -- but it exited on a bare uv exit code with no indication of
what to do. Abort explicitly with the cause and the fix.

Tests slice the sync block out of the real script and run it against a stub uv,
so they exercise the shipped code rather than a copy of it (/app/backend only
exists inside the container). They cover the success path, the retry that
recovers, the abort, and the guidance. Verified against the pre-fix script:
only the guidance case goes red, confirming the abort itself was already
correct.

* fix(extensions): point Git SSH shorthand at the HTTPS correction

Git's SCP-like shorthand carries no URL scheme, so `git+git@host:org/repo.git`
reached the scheme rules looking like a bare path and was rejected with
"local path references are not deployable; pass a local directory so DeerFlow
can snapshot it". The operator asked for a remote source, so that guidance
points at the wrong fix. Detect the shorthand ahead of the scheme rules and
report the public-HTTPS correction instead.

The bare `git@host:org/repo.git` spelling took a different wrong turn: packaging
parses it as a direct reference named `git`, leaving `host:org/repo.git`, whose
`host` reads as a URL scheme and produced the generic HTTPS message. Both
spellings now share one message, as does the PEP 508 named form.

* docs: keep the root extension summary within its new budget

#4799 split the depth out of the module guides and added a size gate; the root
file's job is now orientation, and this branch had pushed it 192 bytes past the
soft limit. The manager transaction, source rules, and lock discipline are
already stated in full in the extensions guide, so the root keeps the one-line
orientation and points there instead of restating them.

---------

Co-authored-by: Codex <codex@openai.com>
2026-08-13 23:55:30 +08:00
Nan Gao
7389331e65
feat(extensions): observe task lifecycle and system model calls (#4684)
* feat(extensions): observe task lifecycle and system model calls

PR 1 (#4636) gave extensions a middleware chain, and a middleware only sees
what passes through the agent graph. Two runtime surfaces stay invisible to
it: when a lead run or a subagent begins and ends, and the DeerFlow-owned
model calls made outside the graph. This slice adds both, with no new
Gateway surface -- routers, services, and the reference extension stay in
PR 3.

Contract (deerflow-extension-api 0.1.1)
---------------------------------------
Two contribution kinds join `middlewares` on the registry:
`task_lifecycle` (`on_task_start` / `on_task_stop`, receiving a `TaskInfo`
and a conservative `TaskOutcome` of completed / aborted / failed) and
`system_model_observer` (`on_system_model_call`, receiving a
`SystemOperationKind`, a `SystemModelRequest` snapshot, and a
`SystemModelResult` carrying either the response or the provider exception
plus a duration).

`SystemModelRequest.messages` normalizes to a tuple at construction. Goal
evaluation and memory extraction pass a message list while title generation
and summarization pass one prompt string, and a bare `str` already satisfies
`Sequence` -- without normalization an observer iterating `request.messages`
would silently walk characters. Copying also makes the frozen snapshot
immutable in fact rather than only by declaration, since observations may run
after the call site returns and keeps mutating its own list.

Registry marks and rollbacks become per-bucket and positional, so an
`install()` that fails after registering two different kinds cannot leave one
of them behind. `needs_task_store` now covers all three kinds: a deployment
that registers only lifecycle hooks still gets a task store.

Task lifecycle
--------------
The lead worker notifies start after the run has started and stop after
completion persistence and the completion hook, but before clearing the
finalizing barrier and publishing the stream end -- holding the barrier
across stop is what keeps a same-thread replacement run from overlapping this
task's lifecycle. Cancellation raised out of the stop notification is
deferred, not propagated in place, so a cancelled run still clears the
barrier and emits its end frame. A subagent with a parent `run_id` wraps its
execution in the same pair inside `finally`, reporting `parent_task_id` so a
delegation tree is reconstructable; a subagent without a `run_id` (embedded
client, standalone LangGraph Server) logs and skips rather than inventing a
parent. Contributors run in registration order inside one shared 3s budget
and every failure is logged and failed open.

System model calls
------------------
Four kinds cover the model calls the middleware chain cannot see: goal
evaluation, memory extraction, title generation, and summarization. Each site
reports both terminal paths without changing the provider exception the host
observes, short-circuits on `has_system_model_observers`, and passes the live
task store when the runtime has one (detached work gets an isolated store).
The sync summarization half stays unobserved on purpose -- it and its only
host caller are the sync side of an async-only runtime, so notifying there
would block a thread on a call site the host never reaches; the reason is
recorded at the call site.

The DeerMem backend must stay vendorable and cannot import the extension API,
so it reports through a new `MemoryCallbacks.on_memory_llm_result` host hook
that the DeerFlow-side callbacks translate into an observation.

Notification loop
-----------------
Extension resources must be touched on the loop that created them, but
subagents can execute on isolated loops and DeerMem runs on a worker thread.
The Gateway registers its serving loop before any runtime dependency starts
and resets it last through the exit stack, so every startup-failure and
cancellation path is covered. Awaited hooks raised on another loop are
dispatched across with `run_coroutine_threadsafe` and awaited under the same
budget; synchronous sites submit fire-and-forget work. Shutdown stops
accepting detached observations before the memory flush -- that flush runs on
a worker thread and can emit memory observations -- while keeping the loop
alive for awaited task hooks until run and subagent drain completes.

Tests
-----
`test_extension_task_lifecycle.py`, `test_extension_subagent_lifecycle.py`,
and `test_extension_system_model_calls.py` cover ordering, fail-open, budget
exhaustion, snapshot binding under a concurrent singleton replacement, the
loop-dispatch and shutdown-suspension paths, and both terminal paths at every
call site. `test_gateway_run_drain_shutdown.py` pins the stop-before-barrier
and drain ordering.

* fix(extensions): decide notification fail-open by origin, observe cancellation

`_notify_each` only guarded `Exception`, so a contributor letting a
`CancelledError` escape — an extension implementing an internal timeout with
cancellation, say — skipped its successors and reached the worker's
deferred-interrupt path, ending an otherwise successful run as cancelled.
Fail-open is about where a failure came from, not its base class: only a
genuine cancellation of the host task increments `Task.cancelling()`, so
propagate on that and contain everything else. `KeyboardInterrupt` /
`SystemExit` still propagate.

`observe_system_model_call` skipped observers on cancellation for the same
base-class reason, leaving goal / title / summarization silent on a terminal
path that is routine — interrupt/rollback admission and shutdown both cancel
the run task, with the provider tokens already spent. Awaiting observers there
is unreliable (a repeated cancel interrupts that await before any of them
runs), so report through the same non-blocking submission the synchronous
memory bridge uses, then propagate the cancellation untouched.

DeerMem keeps `BaseException` around its provider call, now with the reason
recorded: that path runs on a worker thread, where cancelling the awaiting
side never interrupts the running thread, so `CancelledError` cannot arrive
at all. Its host-hook wrapper narrows to `Exception` — only the hook's own
failures are non-fatal, and an observability path must not swallow a process
teardown signal.

* fix(extensions): warn on budget exhaustion, scope observer logs by task, propagate teardown

Review response on #4684:

- The memory observation bridge caught BaseException, which would swallow
  a teardown signal raised while dispatching; it now catches Exception,
  matching the boundary the DeerMem-side call site documents and tests.
- A notification-budget timeout raised mid-hook fell into the generic
  hook-failure path and logged an asyncio-internal traceback; it now logs
  a warning like the pre-hook budget skip, while a TimeoutError a
  contributor raises on its own stays classified as a hook failure.
- System model observer logs passed the operation kind as the task id,
  so log lines said "task goal/title/..."; they now carry the task
  scope id alongside the kind.
2026-08-11 16:33:22 +08:00
Vanzeren
c48de5e70b
feat(checkpoint): make delta snapshot_frequency configurable (#4516)
* feat(checkpoint): make delta snapshot_frequency configurable

* fix(config): carry legacy checkpoint_delta_snapshot_frequency with warning

Addresses review on #4516: the rename from the flat
database.checkpoint_delta_snapshot_frequency key to nested
database.checkpoint_delta.snapshot_frequency silently dropped the old
value (pydantic extra="ignore"). Add a before-validator that maps the
legacy key onto the nested one with a deprecation warning (nested key
wins when both are set), plus a CHANGELOG breaking-change note covering
the rename and the 1000 -> 10 default change.

* fix(checkpoint): validate frozen snapshot frequency
2026-07-28 23:21:23 +08:00
Vanzeren
42baed8c8c
feat(checkpoint): dual-mode checkpoint storage with LangGraph DeltaChannel (#4292)
* feat(checkpoint): dual-mode checkpoint storage with LangGraph DeltaChannel

Add a restart-required database.checkpoint_channel_mode ("full" default,
"delta") that stores the messages channel via LangGraph 1.2 DeltaChannel,
cutting checkpoint storage from O(n^2) to O(n) for append-only history.
Existing full checkpoints seed delta state transparently; no data migration.

- config: mode schema + freeze-on-first-use with
  CheckpointModeReconfigurationError; mode marker persisted in checkpoint
  metadata; unsafe delta->full downgrade rejected fail-closed with
  CheckpointModeMismatchError (run-level error, failed state read)
- state: delta message state schema; CheckpointStateAccessor centralizes
  materialized reads for all consumers (threads API, branches,
  regeneration, compaction, state updates, memory, goal workers)
- runtime: raw writers (run durations, interrupted title, thread goal)
  parent their checkpoints to the checkpoint they derive from, preserving
  delta ancestry; rollback forks the pre-run lineage through a state
  mutation graph with Overwrite restores; InMemorySaver delta-history
  override delegates to the base walk (fixes dropped first write after
  migration, also present upstream)
- tests: conformance suite over {memory, sqlite, postgres} covering
  migration replay, stable message IDs, storage shape and writer
  preservation; conftest fixture isolates the frozen mode between tests;
  stale config fakes refreshed
- ci: backend unit tests gain a postgres service

* fix(checkpoint): close materialization gaps in goal flow, guard public factory

- Route goal-continuation message reads through CheckpointStateAccessor:
  raw channel_values reads see the delta sentinel in delta mode, which
  disabled goal continuation (stand_down=no_durable_end_of_turn) after
  durable assistant turns. Raw tuples remain for tuple-only metadata
  (checkpoint id, pending_writes).
- Reject checkpoint_channel_mode='delta' + checkpointer in
  create_deerflow_agent at construction: factory-built persisted graphs
  bypass mode-marker injection and the fail-closed gate, reproducing
  silent mixed-mode state loss. Delta without persistence stays allowed.
- Import the postgres saver lazily (pytest.importorskip in the fixture)
  so the documented default install collects the suite; add a CI job
  running pytest --collect-only on uv sync --group dev without extras.
- Fix test_checkpointer fallback test to patch get_app_config at its
  use site (provider module), making it deterministic when a local
  config.yaml selects a persistent backend.

* fix(gateway): preserve extension-owned channels in state mutations, bump config version

- build_state_mutation_graph / build_checkpoint_state_mutation_accessor
  accept an explicit state_schema; branch and POST /state now compile the
  mutation graph from the thread's effective schema
  (graph_state_schema on the assistant graph). The base-ThreadState
  fallback silently discarded channels contributed by custom
  AgentMiddleware.state_schema on branch (data loss) and returned a
  false-success 200 on POST /state.
- POST /state validates values keys against the mutation graph's
  channels and rejects unknown fields with 422 instead of ignoring
  them; reducer detection covers extension channels
  (BinaryOperatorAggregate or DeltaChannel) so Overwrite replace
  semantics work for middleware reducers in both modes.
- Endpoint regression: custom AgentMiddleware.state_schema value
  survives branch, updates through POST /state, and an unknown field
  receives 422.
- config_version 26 -> 27 for the new database.checkpoint_channel_mode
  (example, Helm chart values + README, support-bundle fixture), so
  existing installs get the outdated-config warning and
  make config-upgrade merges the field; covered by a test driving the
  real example file and the real config-upgrade script.

* fix(gateway): resolve assistant schema via one boundary, copy branch reducer values with Overwrite

GET /threads/{id}/state now resolves the thread's assistant_id through a
single reusable boundary (thread metadata -> assistant_id -> effective
graph), so channels contributed by a custom AgentMiddleware.state_schema
are materialized instead of dropped by the default lead schema. POST
/state uses the same boundary instead of resolving the schema ad hoc.

Branch writes wrap every copied reducer channel in Overwrite (derived
from the effective mutation graph: BinaryOperatorAggregate + DeltaChannel),
not just messages, so already-aggregated values are never re-merged.

Regression tests use a real AgentMiddleware.state_schema with a
non-identity reducer in both full and delta modes: GET /state returns the
extension value, POST /state replaces it, branch preserves it
byte-for-byte; the unknown-field 422 is a separate assertion.

* refactor(checkpoint): collapse read-path round-trips and ship dual-mode parity tests

Address review round 4 on PR #4292:

- Push ahistory/history limit through Pregel into checkpointer.alist
  (SQL LIMIT) instead of materializing all rows and breaking in Python
- Fold the read-side mode-compat gate onto the returned snapshot's
  metadata; only writes keep the pre-write tuple fetch (fail-closed)
- Cache factory-built accessor graphs per (assistant_id, mode) with
  factory-identity revalidation; state reads no longer build a lead
  agent per request
- get_thread: one snapshot fetch + one raw pending_writes fetch on the
  resolved checkpoint (post-checkpoint __error__ writes never surface
  in snapshot.tasks; verified empirically)
- DeerFlowClient.get_thread: single checkpointer.list walk collects
  pending_writes per checkpoint instead of N get_tuple calls
- InMemorySaver delta-history patch: stand-down when the upstream
  override disappears, try/except guard, validated-version warning,
  guard tests
- make_lead_agent mode precedence: first freeze is owned by app_config
  (client-supplied configurable key ignored); once frozen, injected
  key/app_config must match or fail closed
- Rollback: lock in non-message channel restoration via fork
  inheritance with a dedicated reducer-channel test
- Add tests/test_threads_checkpoint_mode.py and
  tests/test_gateway_checkpoint_mode.py referenced by AGENTS.md and
  the PR validation section: lifecycle parity (memory + sqlite),
  per-step blob-count storage guard, gateway endpoint parity

Counted-saver tests pin checkpoint round-trips for aget/ahistory so
these regressions cannot silently return.

* fix(checkpoint): precise mode-mismatch HTTP mapping, gate E2E, and accessor resilience

- threads router: map CheckpointModeMismatchError to 409 (with cause and
  thread id) and CheckpointModeReconfigurationError to 503 across all state
  endpoints instead of swallowing both into a generic 500
- gate coverage: seed a real delta checkpoint into AsyncSqliteSaver and
  assert aget/aupdate/ahistory fail closed; assert 409 at the HTTP boundary
  through the real route stack
- rollback: compile the restore mutation graph with the thread's effective
  state schema per the build_state_mutation_graph contract
- inheritance contract locks: rollback and manual compaction preserve
  middleware-contributed channels via checkpoint fork cloning
- services: revalidate the accessor-graph cache against app_config identity
  so config.yaml hot-reloads never serve a stale compiled graph
- services: degrade full-mode state reads to raw checkpointer reads when the
  agent factory is unavailable (delta gate still applies; delta mode has no
  fallback)
- deps: override websockets==16.0 (langgraph-sdk 0.4.2's <16 pin silently
  downgraded 16.0 -> 15.0.1; pin is not grounded in any API incompatibility)
  and bump the langchain lower bound to what the lockfile actually resolves

* fix(checkpoint): include anchor checkpoint in degraded history walk + cover get_thread

- _RawCheckpointReadAccessor.ahistory: alist(before=...) is exclusive while
  pregel's get_state_history treats config.checkpoint_id as the inclusive
  start; fetch the anchor explicitly so both read paths paginate identically
- extend the degraded-path gateway test: GET /thread returns raw values, and
  POST /history with before starts at the anchor checkpoint

* fix(gateway): preserve degraded checkpoint timestamps

* fix(gateway): harden degraded checkpoint access

* fix(gateway): resolve assistants for checkpoint reads

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
2026-07-22 08:33:29 +08:00
Xinmin Zeng
268fdd6968
fix(gateway): drain in-flight runs before closing checkpointer on shutdown (#3381)
* fix(gateway): drain in-flight runs before closing checkpointer on shutdown

Chat runs execute in fire-and-forget background asyncio tasks that write
checkpoints through a shared checkpointer. On shutdown, langgraph_runtime's
AsyncExitStack tore down the checkpointer's postgres connection pool while
those run tasks were still mid-graph. langgraph's
AsyncPregelLoop._checkpointer_put_after_previous then ran its
`finally: await checkpointer.aput(...)` against the closed pool, raising
psycopg_pool.PoolClosed. Because that put runs in a langgraph-internal task
(not on run_agent's call stack), run_agent's try/except cannot catch it and it
surfaces as "unhandled exception during asyncio.run() shutdown".

Add RunManager.shutdown() to cancel and bounded-await all in-flight runs, and
call it from langgraph_runtime BEFORE the AsyncExitStack closes the
checkpointer, so the final checkpoint write lands while the pool is still open.
The drain is bounded by a timeout so a stuck run cannot hang worker shutdown,
and is shielded so a second shutdown signal cannot abandon it mid-drain and
reopen the race.

Closes #3373

* fix(gateway): address review — preserve completed-run status, bound drain persistence

Addresses Copilot review on #3381:

- RunManager.shutdown(): decide run status AFTER the drain. Under the lock it
  now only requests cancellation; after asyncio.wait it marks/persists
  `interrupted` only for runs still pending or ended cancelled. A run that
  completes (e.g. `success`) during the drain window keeps its real terminal
  status instead of being unconditionally overwritten.
- Bound the trailing status persistence within the timeout budget
  (deadline = loop.time()+timeout; gather wrapped in asyncio.wait_for) so a slow
  store backing off under DB pressure cannot push shutdown past the deadline.
- deps: use asyncio.create_task instead of asyncio.ensure_future.
- tests: wait deterministically for the run to be in-flight (poll the first
  checkpoint) instead of a fixed sleep; init shutdown_calls explicitly in the
  recovery test double; add regression test asserting a run completing during
  the drain keeps its status (in memory and in the store).

* fix(gateway): address maintainer review — surface failed drain persists, clarify timeout constant

Addresses @WillemJiang review on #3381:

- shutdown(): inspect the gather result of the trailing interrupted-status
  persistence. _persist_status is best-effort (it catches + logs its own
  failure with exc_info and returns False, so it never raises out of the
  gather), but the aggregate result was never checked — a partial failure had
  no shutdown-level visibility. Now any escaped Exception is logged, and any
  False (a persist that did not confirm) is logged with the run_id. Added
  regression test test_shutdown_surfaces_failed_interrupted_persist.
- deps: clarify the _RUN_DRAIN_TIMEOUT_SECONDS comment — state the actual value
  of _SHUTDOWN_HOOK_TIMEOUT_SECONDS (5.0s) and that both count toward the
  lifespan shutdown window. Kept as two separate constants (independent teardown
  steps that may diverge) rather than one shared "must match" value.
- Verified no other test fake needs the shutdown stub: _FakeRunManager in
  test_worker_langfuse_metadata.py is a run_agent() argument (worker path),
  never injected into langgraph_runtime, so it never receives shutdown().
2026-06-07 11:24:30 +08:00