amfs-http-server 0.3.39__tar.gz → 0.3.40__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: amfs-http-server
3
- Version: 0.3.39
3
+ Version: 0.3.40
4
4
  Summary: AMFS HTTP/REST API server with SSE support
5
5
  License-Expression: Apache-2.0
6
6
  Requires-Python: >=3.11
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "amfs-http-server"
3
- version = "0.3.39"
3
+ version = "0.3.40"
4
4
  description = "AMFS HTTP/REST API server with SSE support"
5
5
  requires-python = ">=3.11"
6
6
  license = "Apache-2.0"
@@ -134,8 +134,19 @@ _async_adapter = None # AsyncPostgresAdapter | None, set in lifespan
134
134
 
135
135
  @asynccontextmanager
136
136
  async def _lifespan(application: FastAPI): # noqa: ARG001
137
- """Open/close the async connection pool for hot-path DB access."""
138
- global _async_adapter
137
+ """Open/close the async connection pool for hot-path DB access, and run
138
+ the embedded Cortex worker for the life of this process.
139
+
140
+ The Cortex worker starts here rather than in ``main()`` because this is the
141
+ hook that runs once in every process that serves requests. With
142
+ ``--workers`` above one, uvicorn imports ``amfs_http.server:app`` afresh
143
+ in each child; ``main()`` runs only in the supervisor, so a worker started
144
+ there is invisible to the ``/api/v1/cortex/*`` routes in the children —
145
+ they read the module global and saw ``None``. Starting it per process
146
+ keeps those routes truthful and costs one LISTEN connection and a
147
+ two-connection pool per worker instead of per instance.
148
+ """
149
+ global _async_adapter, _cortex_worker
139
150
  dsn = os.environ.get("AMFS_POSTGRES_DSN")
140
151
  if dsn and not os.environ.get("AMFS_HTTP_URL"):
141
152
  try:
@@ -147,12 +158,24 @@ async def _lifespan(application: FastAPI): # noqa: ARG001
147
158
  except Exception:
148
159
  logger.warning("Failed to start async adapter — falling back to sync", exc_info=True)
149
160
  _async_adapter = None
161
+ if _env_flag("AMFS_WITH_CORTEX"):
162
+ _start_embedded_cortex()
150
163
  yield
164
+ if _cortex_worker is not None:
165
+ try:
166
+ _cortex_worker.stop()
167
+ except Exception: # noqa: BLE001 - shutting down; nothing to do about it
168
+ logger.debug("Cortex worker stop raised", exc_info=True)
151
169
  if _async_adapter is not None:
152
170
  await _async_adapter.close()
153
171
  logger.info("Async Postgres adapter closed")
154
172
 
155
173
 
174
+ def _env_flag(name: str) -> bool:
175
+ """Whether environment variable *name* is set to a truthy value."""
176
+ return (os.environ.get(name) or "").strip().lower() in {"1", "true", "yes", "on"}
177
+
178
+
156
179
  app = FastAPI(
157
180
  title="AMFS HTTP API",
158
181
  description="Agent Memory File System — REST API with SSE support",
@@ -233,8 +256,72 @@ _bg_executor = ThreadPoolExecutor(max_workers=8, thread_name_prefix="amfs-bg")
233
256
  # checkouts from the sync pool, not a throughput target. ``_offload`` copies
234
257
  # the calling context into the thread so the tenant ContextVars the RLS pool
235
258
  # reads are the request's, not the thread's leftovers.
259
+ #
260
+ # Routes whose whole body is synchronous — the admin, team, API-key and
261
+ # pattern routes that open ``pool.connection()`` and return — are plain
262
+ # ``def``, on purpose. FastAPI runs a ``def`` endpoint on the anyio worker
263
+ # pool, with the caller's contextvars copied in, so the body needs no
264
+ # ``_offload`` wrapping; declared ``async`` the same body holds the loop for
265
+ # the length of the query, and for however long the pool checkout has to wait
266
+ # when the pool is busy. Do not "fix" one back to ``async def`` unless it
267
+ # gains an ``await``.
268
+
269
+
270
+ def _effective_cpu_count() -> int:
271
+ """The CPUs this process may actually use, not the ones the host has.
272
+
273
+ ``os.cpu_count()`` reports the machine. In a container it is the node's
274
+ core count, while the scheduler holds the process to whatever the cgroup
275
+ quota or affinity mask says — on Cloud Run a 2-vCPU service can read 8 or
276
+ more here. Sizing a thread pool for the host's cores then puts several
277
+ ONNX threads on every core the container is allowed, which is the thrash
278
+ the pool was sized to avoid.
279
+
280
+ Takes the smallest of: the affinity mask (Linux), the cgroup v2 quota
281
+ (``cpu.max``, ``quota/period`` rounded up), the cgroup v1 quota, and
282
+ ``os.cpu_count()``. ``AMFS_MODEL_THREADS`` overrides all of it for the
283
+ deployment that knows better. Never below one.
284
+ """
285
+ override = os.environ.get("AMFS_MODEL_THREADS")
286
+ if override:
287
+ try:
288
+ return max(1, int(override))
289
+ except ValueError:
290
+ logger.warning("AMFS_MODEL_THREADS=%r is not an integer; ignoring", override)
291
+
292
+ candidates: list[int] = []
293
+ host = os.cpu_count()
294
+ if host:
295
+ candidates.append(host)
296
+ try:
297
+ candidates.append(len(os.sched_getaffinity(0))) # type: ignore[attr-defined]
298
+ except (AttributeError, OSError):
299
+ pass
300
+ for quota_path, period_path in (
301
+ ("/sys/fs/cgroup/cpu.max", None),
302
+ ("/sys/fs/cgroup/cpu/cpu.cfs_quota_us", "/sys/fs/cgroup/cpu/cpu.cfs_period_us"),
303
+ ):
304
+ try:
305
+ with open(quota_path) as fh:
306
+ first = fh.read().split()
307
+ if period_path is None:
308
+ quota_s, period_s = first[0], first[1]
309
+ else:
310
+ quota_s = first[0]
311
+ with open(period_path) as fh:
312
+ period_s = fh.read().split()[0]
313
+ if quota_s in ("max", "-1"):
314
+ continue
315
+ quota, period = int(quota_s), int(period_s)
316
+ if quota > 0 and period > 0:
317
+ candidates.append(max(1, -(-quota // period)))
318
+ except (OSError, ValueError, IndexError):
319
+ continue
320
+ return max(1, min(candidates) if candidates else 2)
321
+
322
+
236
323
  _model_executor = ThreadPoolExecutor(
237
- max_workers=max(1, os.cpu_count() or 2), thread_name_prefix="amfs-model"
324
+ max_workers=_effective_cpu_count(), thread_name_prefix="amfs-model"
238
325
  )
239
326
  _db_executor = ThreadPoolExecutor(max_workers=8, thread_name_prefix="amfs-db")
240
327
 
@@ -248,6 +335,19 @@ async def _offload(executor: ThreadPoolExecutor, fn: Any, /, *args: Any, **kwarg
248
335
  )
249
336
 
250
337
 
338
+ def _sync_search(adapter: Any, sq: Any, branch: Any) -> list[Any]:
339
+ """``adapter.search`` for adapters with and without a ``branch`` parameter.
340
+
341
+ Meant to run on ``_db_executor`` via ``_offload``; the sync search is the
342
+ fallback the async path takes when it returns nothing, and it must not
343
+ hold the loop any more than the primary path does.
344
+ """
345
+ try:
346
+ return adapter.search(sq, branch=branch)
347
+ except TypeError:
348
+ return adapter.search(sq)
349
+
350
+
251
351
  _known_agents: set[str] = set()
252
352
  # Tracks (agent, namespace, user) triples whose owner linkage was already
253
353
  # upserted, so the hot write path doesn't repeat the DB call. Kept separate
@@ -604,8 +704,28 @@ except ImportError:
604
704
  from amfs_http.pro_proxy import mount_pro_proxy
605
705
  mount_pro_proxy(app)
606
706
 
707
+ _memory_lock = threading.Lock()
708
+
709
+
607
710
  def _get_memory() -> AgentMemory:
608
- """Lazily initialise the shared AgentMemory singleton."""
711
+ """Return the shared AgentMemory singleton, building it on first call.
712
+
713
+ Double-checked under a lock: the sync-bodied routes run on the threadpool,
714
+ so the first requests after a worker starts — a dashboard load fires
715
+ several at once — can arrive here together, and without the lock each
716
+ would build its own ``AgentMemory`` and adapter pool, with every loser
717
+ dropped unclosed. The fast path stays lock-free.
718
+ """
719
+ if _memory is not None:
720
+ return _memory
721
+ with _memory_lock:
722
+ if _memory is not None:
723
+ return _memory
724
+ return _build_memory()
725
+
726
+
727
+ def _build_memory() -> AgentMemory:
728
+ """Construct the shared AgentMemory. Call through ``_get_memory``."""
609
729
  global _memory
610
730
  if _memory is not None:
611
731
  return _memory
@@ -1635,7 +1755,7 @@ async def _read_entry(
1635
1755
 
1636
1756
 
1637
1757
  @app.get("/api/v1/quality/{entity_path:path}/{key}")
1638
- async def entry_quality(
1758
+ def entry_quality(
1639
1759
  request: Request,
1640
1760
  entity_path: str,
1641
1761
  key: str,
@@ -1736,11 +1856,10 @@ async def _scope_block(
1736
1856
  vis = _get_visibility_filter(request) if request is not None else None
1737
1857
  if vis is not None and vis.should_filter():
1738
1858
  try:
1739
- visible = {
1740
- e.key for e in vis.filter_entries(
1741
- _get_memory()._adapter.list(entity_path, branch=branch)
1742
- )
1743
- }
1859
+ listed = await _offload(
1860
+ _db_executor, _get_memory()._adapter.list, entity_path, branch=branch
1861
+ )
1862
+ visible = {e.key for e in vis.filter_entries(listed)}
1744
1863
  counts = {**counts, "keys": [k for k in counts.get("keys", []) if k in visible]}
1745
1864
  except Exception: # noqa: BLE001 - drop the sample rather than leak it
1746
1865
  counts = {**counts, "keys": []}
@@ -2404,10 +2523,9 @@ async def search_entries(
2404
2523
  logger.warning("Async search failed — falling back to sync", exc_info=True)
2405
2524
  results = []
2406
2525
  if not results:
2407
- try:
2408
- sync_results = mem._adapter.search(sq, branch=branch)
2409
- except TypeError:
2410
- sync_results = mem._adapter.search(sq)
2526
+ # Runs on every search that finds nothing, so it goes off the loop
2527
+ # like the primary path did.
2528
+ sync_results = await _offload(_db_executor, _sync_search, mem._adapter, sq, branch)
2411
2529
  if sync_results:
2412
2530
  logger.warning(
2413
2531
  "Async search returned 0 results but sync found %d — RLS context mismatch",
@@ -2415,10 +2533,7 @@ async def search_entries(
2415
2533
  )
2416
2534
  results = sync_results
2417
2535
  else:
2418
- try:
2419
- results = mem._adapter.search(sq, branch=branch)
2420
- except TypeError:
2421
- results = mem._adapter.search(sq)
2536
+ results = await _offload(_db_executor, _sync_search, mem._adapter, sq, branch)
2422
2537
 
2423
2538
  vis = _get_visibility_filter(request)
2424
2539
  if vis is not None and vis.should_filter():
@@ -2593,9 +2708,7 @@ async def retrieve_entries(
2593
2708
  if not candidates:
2594
2709
  mem = _get_memory()
2595
2710
  try:
2596
- results = mem._adapter.search(sq_lex, branch=branch)
2597
- except TypeError:
2598
- results = mem._adapter.search(sq_lex)
2711
+ results = await _offload(_db_executor, _sync_search, mem._adapter, sq_lex, branch)
2599
2712
  except Exception:
2600
2713
  results = []
2601
2714
  for entry in results:
@@ -5076,7 +5189,7 @@ Return ONLY valid JSON, no markdown formatting."""
5076
5189
 
5077
5190
 
5078
5191
  @app.get("/api/v1/admin/usage")
5079
- async def get_usage(
5192
+ def get_usage(
5080
5193
  request: Request,
5081
5194
  _auth: str | None = Depends(verify_api_key),
5082
5195
  ) -> dict[str, Any]:
@@ -6618,7 +6731,7 @@ def _audit_log(
6618
6731
 
6619
6732
 
6620
6733
  @app.get("/api/v1/admin/api-keys")
6621
- async def list_api_keys(
6734
+ def list_api_keys(
6622
6735
  request: Request,
6623
6736
  _auth: str | None = Depends(verify_api_key),
6624
6737
  ) -> dict[str, Any]:
@@ -6675,7 +6788,7 @@ async def list_api_keys(
6675
6788
 
6676
6789
 
6677
6790
  @app.post("/api/v1/admin/api-keys")
6678
- async def create_api_key(
6791
+ def create_api_key(
6679
6792
  req: CreateAPIKeyRequest,
6680
6793
  request: Request,
6681
6794
  _auth: str | None = Depends(verify_api_key),
@@ -6729,7 +6842,7 @@ async def create_api_key(
6729
6842
 
6730
6843
 
6731
6844
  @app.delete("/api/v1/admin/api-keys/{key_id}")
6732
- async def revoke_api_key(
6845
+ def revoke_api_key(
6733
6846
  key_id: str,
6734
6847
  request: Request,
6735
6848
  _auth: str | None = Depends(verify_api_key),
@@ -6776,7 +6889,7 @@ async def revoke_api_key(
6776
6889
 
6777
6890
 
6778
6891
  @app.get("/api/v1/admin/audit")
6779
- async def list_audit_log(
6892
+ def list_audit_log(
6780
6893
  request: Request,
6781
6894
  action: str | None = Query(None),
6782
6895
  limit: int = Query(200),
@@ -6876,7 +6989,7 @@ async def list_patterns(
6876
6989
 
6877
6990
 
6878
6991
  @app.get("/api/v1/admin/teams")
6879
- async def list_teams(
6992
+ def list_teams(
6880
6993
  request: Request,
6881
6994
  _auth: str | None = Depends(verify_api_key),
6882
6995
  ) -> dict[str, Any]:
@@ -6919,7 +7032,7 @@ async def list_teams(
6919
7032
 
6920
7033
 
6921
7034
  @app.post("/api/v1/admin/teams")
6922
- async def create_team(
7035
+ def create_team(
6923
7036
  req: CreateTeamRequest,
6924
7037
  request: Request,
6925
7038
  _auth: str | None = Depends(verify_api_key),
@@ -6954,7 +7067,7 @@ async def create_team(
6954
7067
 
6955
7068
 
6956
7069
  @app.patch("/api/v1/admin/teams/{team_id}")
6957
- async def update_team(
7070
+ def update_team(
6958
7071
  team_id: str,
6959
7072
  req: UpdateTeamRequest,
6960
7073
  request: Request,
@@ -7009,7 +7122,7 @@ async def update_team(
7009
7122
 
7010
7123
 
7011
7124
  @app.delete("/api/v1/admin/teams/{team_id}")
7012
- async def delete_team(
7125
+ def delete_team(
7013
7126
  team_id: str,
7014
7127
  request: Request,
7015
7128
  _auth: str | None = Depends(verify_api_key),
@@ -7044,7 +7157,7 @@ async def delete_team(
7044
7157
 
7045
7158
 
7046
7159
  @app.get("/api/v1/admin/teams/{team_id}/members")
7047
- async def list_team_members(
7160
+ def list_team_members(
7048
7161
  request: Request,
7049
7162
  team_id: str,
7050
7163
  include_removed: bool = Query(False),
@@ -7100,7 +7213,7 @@ async def list_team_members(
7100
7213
 
7101
7214
 
7102
7215
  @app.post("/api/v1/admin/teams/{team_id}/members")
7103
- async def add_team_member(
7216
+ def add_team_member(
7104
7217
  team_id: str,
7105
7218
  req: AddTeamMemberRequest,
7106
7219
  request: Request,
@@ -7157,7 +7270,7 @@ async def add_team_member(
7157
7270
 
7158
7271
 
7159
7272
  @app.patch("/api/v1/admin/teams/{team_id}/members/{member_id}")
7160
- async def update_team_member(
7273
+ def update_team_member(
7161
7274
  team_id: str,
7162
7275
  member_id: str,
7163
7276
  req: UpdateTeamMemberRequest,
@@ -7213,7 +7326,7 @@ async def update_team_member(
7213
7326
 
7214
7327
 
7215
7328
  @app.delete("/api/v1/admin/teams/{team_id}/members/{member_id}")
7216
- async def remove_team_member(
7329
+ def remove_team_member(
7217
7330
  team_id: str,
7218
7331
  member_id: str,
7219
7332
  request: Request,
@@ -7247,7 +7360,7 @@ async def remove_team_member(
7247
7360
 
7248
7361
 
7249
7362
  @app.post("/api/v1/admin/teams/{team_id}/members/{member_id}/reinstate")
7250
- async def reinstate_team_member(
7363
+ def reinstate_team_member(
7251
7364
  team_id: str,
7252
7365
  member_id: str,
7253
7366
  request: Request,
@@ -7291,7 +7404,7 @@ async def reinstate_team_member(
7291
7404
 
7292
7405
 
7293
7406
  @app.get("/api/v1/admin/members/check-email")
7294
- async def check_member_email(
7407
+ def check_member_email(
7295
7408
  request: Request,
7296
7409
  email: str = Query(...),
7297
7410
  _auth: str | None = Depends(verify_api_key),
@@ -7356,7 +7469,7 @@ async def check_member_email(
7356
7469
 
7357
7470
 
7358
7471
  @app.get("/api/v1/admin/patterns")
7359
- async def list_detected_patterns(
7472
+ def list_detected_patterns(
7360
7473
  pattern_type: str | None = Query(None),
7361
7474
  category: str | None = Query(None),
7362
7475
  severity: str | None = Query(None),
@@ -7437,7 +7550,7 @@ async def list_detected_patterns(
7437
7550
 
7438
7551
 
7439
7552
  @app.post("/api/v1/admin/patterns/scan")
7440
- async def run_pattern_scan(
7553
+ def run_pattern_scan(
7441
7554
  req: RunPatternDetectionRequest,
7442
7555
  request: Request,
7443
7556
  _auth: str | None = Depends(verify_api_key),
@@ -7543,7 +7656,7 @@ async def run_pattern_scan(
7543
7656
 
7544
7657
 
7545
7658
  @app.patch("/api/v1/admin/patterns/{pattern_id}/resolve")
7546
- async def resolve_pattern(
7659
+ def resolve_pattern(
7547
7660
  pattern_id: str,
7548
7661
  request: Request,
7549
7662
  _auth: str | None = Depends(verify_api_key),
@@ -8162,7 +8275,9 @@ async def _credit_briefing_reuse(
8162
8275
  if _async_adapter is not None:
8163
8276
  entry = await _async_adapter.read(entity_path, key, branch=branch)
8164
8277
  else:
8165
- entry = _get_memory()._adapter.read(entity_path, key, branch=branch)
8278
+ entry = await _offload(
8279
+ _db_executor, _get_memory()._adapter.read, entity_path, key, branch=branch
8280
+ )
8166
8281
  except Exception: # noqa: BLE001 - the value block is reporting, not the answer
8167
8282
  logger.debug("briefing reuse lookup failed", exc_info=True)
8168
8283
 
@@ -8778,86 +8893,104 @@ def _make_tenant_provider(dsn: str):
8778
8893
  return _provider
8779
8894
 
8780
8895
 
8781
- def main() -> None:
8782
- """Run the AMFS HTTP server via uvicorn."""
8896
+ def _start_embedded_cortex() -> None:
8897
+ """Start the embedded Cortex worker in this process, if it can be.
8898
+
8899
+ Called from the app lifespan, so it runs once per serving process — see
8900
+ ``_lifespan`` for why not from ``main()``. Idempotent: a second call in the
8901
+ same process (a test re-entering the lifespan) is a no-op.
8902
+ """
8783
8903
  global _cortex_worker
8784
- args = _parse_args()
8904
+ if _cortex_worker is not None:
8905
+ return
8906
+ dsn = os.environ.get("AMFS_POSTGRES_DSN")
8907
+ if not dsn:
8908
+ logger.warning("AMFS_WITH_CORTEX requires AMFS_POSTGRES_DSN")
8909
+ return
8785
8910
 
8786
- if args.with_cortex:
8787
- dsn = os.environ.get("AMFS_POSTGRES_DSN")
8788
- if dsn:
8789
- import threading
8911
+ try:
8912
+ from amfs_postgres.adapter import PostgresAdapter
8913
+ from amfs_cortex.compiler import DigestCompiler
8914
+ from amfs_cortex.worker import CortexWorker
8915
+
8916
+ namespace = os.environ.get("AMFS_NAMESPACE", "default")
8917
+ # One background thread compiling digests on a timer, sharing a
8918
+ # process with the request path. Sized for that rather than left
8919
+ # on the pool default, which would have it hold as many
8920
+ # connections as the endpoints do while using one at a time —
8921
+ # multiplied by every process the service scales out to.
8922
+ adapter = PostgresAdapter(
8923
+ dsn=dsn, namespace=namespace, min_pool_size=1, max_pool_size=2
8924
+ )
8790
8925
 
8791
- try:
8792
- from amfs_postgres.adapter import PostgresAdapter
8793
- from amfs_cortex.compiler import DigestCompiler
8794
- from amfs_cortex.worker import CortexWorker
8795
-
8796
- namespace = os.environ.get("AMFS_NAMESPACE", "default")
8797
- # One background thread compiling digests on a timer, sharing a
8798
- # process with the request path. Sized for that rather than left
8799
- # on the pool default, which would have it hold as many
8800
- # connections as the endpoints do while using one at a time —
8801
- # multiplied by every instance the service scales out to.
8802
- adapter = PostgresAdapter(
8803
- dsn=dsn, namespace=namespace, min_pool_size=1, max_pool_size=2
8804
- )
8926
+ strategies = []
8927
+ try:
8928
+ from amfs_cortex_pro import get_pro_strategies
8929
+ strategies = get_pro_strategies()
8930
+ logger.info("Pro compilation strategies loaded")
8931
+ except ImportError:
8932
+ pass
8805
8933
 
8806
- strategies = []
8807
- try:
8808
- from amfs_cortex_pro import get_pro_strategies
8809
- strategies = get_pro_strategies()
8810
- logger.info("Pro compilation strategies loaded")
8811
- except ImportError:
8812
- pass
8934
+ compiler = DigestCompiler(
8935
+ adapter=adapter,
8936
+ strategies=strategies or None,
8937
+ namespace=namespace,
8938
+ )
8939
+ tenant_provider = _make_tenant_provider(dsn)
8940
+ # The catch-up scan runs in every serving process (no advisory lock),
8941
+ # so its interval sets the fleet-wide scan rate: 36 processes at the
8942
+ # 300 s default is one tenant-wide scan every ~8 s. A deployment that
8943
+ # runs several workers per instance should raise this in proportion
8944
+ # to keep the rate where it was. The scan is a GROUP BY since the
8945
+ # adapter grew list_scopes(), but the knob stays: 0 disables it on
8946
+ # deployments where the event path is trusted to compile every scope.
8947
+ catchup_s = float(os.environ.get("AMFS_CORTEX_CATCHUP_INTERVAL_S", "300"))
8948
+ worker = CortexWorker(
8949
+ dsn=dsn,
8950
+ compiler=compiler,
8951
+ use_advisory_lock=False,
8952
+ catchup_interval_s=catchup_s,
8953
+ tenant_provider=tenant_provider,
8954
+ )
8813
8955
 
8814
- compiler = DigestCompiler(
8815
- adapter=adapter,
8816
- strategies=strategies or None,
8817
- namespace=namespace,
8818
- )
8819
- tenant_provider = _make_tenant_provider(dsn)
8820
- # The catch-up scan runs on every instance (no advisory lock), so
8821
- # its interval sets the fleet-wide scan rate: 36 instances at the
8822
- # 300 s default is one tenant-wide scan every ~8 s. The scan is
8823
- # a GROUP BY since the adapter grew list_scopes(), but the knob
8824
- # stays: 0 disables it on deployments where the event path is
8825
- # trusted to compile every scope.
8826
- catchup_s = float(os.environ.get("AMFS_CORTEX_CATCHUP_INTERVAL_S", "300"))
8827
- _cortex_worker = CortexWorker(
8828
- dsn=dsn,
8829
- compiler=compiler,
8830
- use_advisory_lock=False,
8831
- catchup_interval_s=catchup_s,
8832
- tenant_provider=tenant_provider,
8833
- )
8956
+ try:
8957
+ from amfs_cortex_pro import get_outcome_wiring, HotContextTracker
8958
+ wiring = get_outcome_wiring(adapter, namespace)
8959
+ if wiring:
8960
+ worker._outcome_wiring = wiring
8961
+ logger.info("Outcome wiring attached to embedded Cortex worker")
8962
+ tracker = HotContextTracker()
8963
+ worker._hot_tracker = tracker
8964
+ logger.info("Hot context tracker attached to embedded Cortex worker")
8965
+ except ImportError:
8966
+ pass
8834
8967
 
8835
- try:
8836
- from amfs_cortex_pro import get_outcome_wiring, HotContextTracker
8837
- wiring = get_outcome_wiring(adapter, namespace)
8838
- if wiring:
8839
- _cortex_worker._outcome_wiring = wiring
8840
- logger.info("Outcome wiring attached to embedded Cortex worker")
8841
- tracker = HotContextTracker()
8842
- _cortex_worker._hot_tracker = tracker
8843
- logger.info("Hot context tracker attached to embedded Cortex worker")
8844
- except ImportError:
8845
- pass
8968
+ from amfs_http.pro_proxy import create_forwarder
8969
+ forwarder = create_forwarder()
8970
+ if forwarder:
8971
+ worker._pro_forwarder = forwarder
8846
8972
 
8847
- from amfs_http.pro_proxy import create_forwarder
8848
- forwarder = create_forwarder()
8849
- if forwarder:
8850
- _cortex_worker._pro_forwarder = forwarder
8973
+ t = threading.Thread(target=worker.run, daemon=True, name="cortex-embedded")
8974
+ t.start()
8975
+ _cortex_worker = worker
8976
+ logger.info("Embedded Cortex worker started (pid=%d)", os.getpid())
8977
+ except ImportError:
8978
+ logger.warning("AMFS_WITH_CORTEX requires amfs-cortex package")
8979
+ except Exception:
8980
+ logger.exception("Failed to start embedded Cortex worker — server will run without Cortex")
8851
8981
 
8852
- t = threading.Thread(target=_cortex_worker.run, daemon=True, name="cortex-embedded")
8853
- t.start()
8854
- logger.info("Embedded Cortex worker started")
8855
- except ImportError:
8856
- logger.warning("--with-cortex requires amfs-cortex package")
8857
- except Exception:
8858
- logger.exception("Failed to start embedded Cortex worker — server will run without Cortex")
8859
- else:
8860
- logger.warning("--with-cortex requires AMFS_POSTGRES_DSN")
8982
+
8983
+ def main() -> None:
8984
+ """Run the AMFS HTTP server via uvicorn."""
8985
+ args = _parse_args()
8986
+
8987
+ # The flag becomes an environment variable rather than a start here, so
8988
+ # that every uvicorn worker process — which inherits the environment but
8989
+ # not this function's locals — starts its own Cortex from the lifespan.
8990
+ # Setting the variable directly (docker-compose, a systemd unit) works the
8991
+ # same without the flag.
8992
+ if args.with_cortex:
8993
+ os.environ["AMFS_WITH_CORTEX"] = "1"
8861
8994
 
8862
8995
  workers = args.workers
8863
8996
  if args.reload and workers > 1:
@@ -0,0 +1,189 @@
1
+ """Server-Sent Events for real-time AMFS watch notifications.
2
+
3
+ Subscribers are ``asyncio.Queue`` objects owned by the event loop that opened
4
+ the stream. Broadcasts, though, arrive from wherever the write happened: an
5
+ ``async`` route on that loop, or — since the sync-bodied routes run on FastAPI's
6
+ threadpool — a worker thread. ``asyncio.Queue.put_nowait`` is not thread-safe:
7
+ called from another thread it appends the item but wakes the waiting consumer
8
+ through ``call_soon``, which does not rouse a sleeping loop, so the event sits
9
+ until something else happens to wake it (and raises under asyncio debug mode).
10
+ The manager records each queue's loop at subscribe time and, when a broadcast
11
+ comes from outside it, hands the put to ``call_soon_threadsafe``.
12
+ """
13
+ from __future__ import annotations
14
+
15
+ import asyncio
16
+ import json
17
+ import logging
18
+ import threading
19
+ from typing import Any
20
+
21
+ from amfs_core.models import MemoryEntry
22
+
23
+ logger = logging.getLogger(__name__)
24
+
25
+
26
+ class SSEManager:
27
+ def __init__(self) -> None:
28
+ self._subscribers: dict[str, list[asyncio.Queue[dict[str, Any]]]] = {}
29
+ self._room_subscribers: dict[str, list[asyncio.Queue[dict[str, Any]]]] = {}
30
+ # The loop each queue belongs to, keyed by the queue's id. Subscribe and
31
+ # unsubscribe run on that loop; broadcasts may run on any thread, so the
32
+ # registries are read under the lock and iterated as copies.
33
+ self._loops: dict[int, asyncio.AbstractEventLoop] = {}
34
+ self._lock = threading.Lock()
35
+
36
+ # ── Subscriptions ────────────────────────────────────────────────
37
+
38
+ def _new_queue(self) -> asyncio.Queue[dict[str, Any]]:
39
+ queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
40
+ try:
41
+ loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop()
42
+ except RuntimeError: # subscribed outside a loop (tests); deliver directly
43
+ loop = None
44
+ if loop is not None:
45
+ with self._lock:
46
+ self._loops[id(queue)] = loop
47
+ return queue
48
+
49
+ def _forget_queue(self, queue: asyncio.Queue[dict[str, Any]]) -> None:
50
+ with self._lock:
51
+ self._loops.pop(id(queue), None)
52
+
53
+ def subscribe(self, entity_path: str = "*") -> asyncio.Queue[dict[str, Any]]:
54
+ queue = self._new_queue()
55
+ with self._lock:
56
+ self._subscribers.setdefault(entity_path, []).append(queue)
57
+ return queue
58
+
59
+ def unsubscribe(self, entity_path: str, queue: asyncio.Queue[dict[str, Any]]) -> None:
60
+ with self._lock:
61
+ subs = self._subscribers.get(entity_path, [])
62
+ if queue in subs:
63
+ subs.remove(queue)
64
+ self._forget_queue(queue)
65
+
66
+ def subscribe_room(self, room_id: str) -> asyncio.Queue[dict[str, Any]]:
67
+ """Subscribe to all events for a specific room."""
68
+ queue = self._new_queue()
69
+ with self._lock:
70
+ self._room_subscribers.setdefault(room_id, []).append(queue)
71
+ return queue
72
+
73
+ def unsubscribe_room(self, room_id: str, queue: asyncio.Queue[dict[str, Any]]) -> None:
74
+ with self._lock:
75
+ subs = self._room_subscribers.get(room_id, [])
76
+ if queue in subs:
77
+ subs.remove(queue)
78
+ self._forget_queue(queue)
79
+
80
+ # ── Delivery ─────────────────────────────────────────────────────
81
+
82
+ @staticmethod
83
+ def _put(queue: asyncio.Queue[dict[str, Any]], event: dict[str, Any], what: str) -> None:
84
+ try:
85
+ queue.put_nowait(event)
86
+ except asyncio.QueueFull:
87
+ logger.debug("SSE queue full, dropping %s", what)
88
+
89
+ def _deliver(
90
+ self, queue: asyncio.Queue[dict[str, Any]], event: dict[str, Any], what: str
91
+ ) -> None:
92
+ """Put ``event`` on ``queue`` from whatever thread this is.
93
+
94
+ On the queue's own loop the put is direct. From any other thread — a
95
+ threadpool worker running a sync route — it is scheduled onto that loop
96
+ with ``call_soon_threadsafe``, which is the one asyncio entry point that
97
+ may be called from outside the loop and that wakes it.
98
+ """
99
+ with self._lock:
100
+ owner = self._loops.get(id(queue))
101
+ if owner is None:
102
+ self._put(queue, event, what)
103
+ return
104
+ try:
105
+ running: asyncio.AbstractEventLoop | None = asyncio.get_running_loop()
106
+ except RuntimeError:
107
+ running = None
108
+ if running is owner:
109
+ self._put(queue, event, what)
110
+ return
111
+ if owner.is_closed():
112
+ return
113
+ try:
114
+ owner.call_soon_threadsafe(self._put, queue, event, what)
115
+ except RuntimeError: # the loop closed between the check and the call
116
+ logger.debug("SSE subscriber loop closed, dropping %s", what)
117
+
118
+ def broadcast(self, entry: MemoryEntry) -> None:
119
+ data = entry.model_dump(mode="json")
120
+ data.pop("embedding", None)
121
+ event = {"type": "write", "entry": data}
122
+ with self._lock:
123
+ targets = [
124
+ queue
125
+ for pattern, queues in self._subscribers.items()
126
+ if (
127
+ pattern == "*"
128
+ or entry.entity_path == pattern
129
+ or entry.entity_path.startswith(pattern + "/")
130
+ )
131
+ for queue in queues
132
+ ]
133
+ for queue in targets:
134
+ self._deliver(queue, event, "event")
135
+
136
+ def broadcast_room_event(
137
+ self,
138
+ room_id: str,
139
+ event_type: str,
140
+ data: dict[str, Any],
141
+ ) -> None:
142
+ """Broadcast a room-scoped event (join, leave, write, read, etc.)."""
143
+ event = {"type": event_type, "room_id": room_id, **data}
144
+ with self._lock:
145
+ targets = list(self._room_subscribers.get(room_id, []))
146
+ for queue in targets:
147
+ # Each subscriber pops "type" from its own copy in the generator.
148
+ self._deliver(queue, dict(event), f"event for room {room_id}")
149
+
150
+ # ── Generators ───────────────────────────────────────────────────
151
+
152
+ async def event_generator(self, entity_path: str = "*", predicate=None):
153
+ """Yield SSE events, optionally filtered per subscriber.
154
+
155
+ ``predicate`` receives the entry dict and returns True when the
156
+ subscriber may see it (used for per-user visibility). Events that
157
+ fail the predicate — or whose predicate raises — are dropped.
158
+ """
159
+ queue = self.subscribe(entity_path)
160
+ try:
161
+ while True:
162
+ event = await queue.get()
163
+ if predicate is not None:
164
+ try:
165
+ if not predicate(event.get("entry", {})):
166
+ continue
167
+ except Exception:
168
+ logger.debug("SSE visibility predicate failed — dropping event", exc_info=True)
169
+ continue
170
+ yield {
171
+ "event": event["type"],
172
+ "data": json.dumps(event["entry"], default=str),
173
+ }
174
+ finally:
175
+ self.unsubscribe(entity_path, queue)
176
+
177
+ async def room_event_generator(self, room_id: str):
178
+ """SSE generator for room-scoped events."""
179
+ queue = self.subscribe_room(room_id)
180
+ try:
181
+ while True:
182
+ event = await queue.get()
183
+ event_type = event.pop("type", "room_event")
184
+ yield {
185
+ "event": event_type,
186
+ "data": json.dumps(event, default=str),
187
+ }
188
+ finally:
189
+ self.unsubscribe_room(room_id, queue)
@@ -1,112 +0,0 @@
1
- """Server-Sent Events for real-time AMFS watch notifications."""
2
- from __future__ import annotations
3
-
4
- import asyncio
5
- import json
6
- import logging
7
- from typing import Any
8
-
9
- from amfs_core.models import MemoryEntry
10
-
11
- logger = logging.getLogger(__name__)
12
-
13
-
14
- class SSEManager:
15
- def __init__(self) -> None:
16
- self._subscribers: dict[str, list[asyncio.Queue[dict[str, Any]]]] = {}
17
- self._room_subscribers: dict[str, list[asyncio.Queue[dict[str, Any]]]] = {}
18
-
19
- def subscribe(self, entity_path: str = "*") -> asyncio.Queue[dict[str, Any]]:
20
- queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
21
- if entity_path not in self._subscribers:
22
- self._subscribers[entity_path] = []
23
- self._subscribers[entity_path].append(queue)
24
- return queue
25
-
26
- def unsubscribe(self, entity_path: str, queue: asyncio.Queue[dict[str, Any]]) -> None:
27
- subs = self._subscribers.get(entity_path, [])
28
- if queue in subs:
29
- subs.remove(queue)
30
-
31
- def subscribe_room(self, room_id: str) -> asyncio.Queue[dict[str, Any]]:
32
- """Subscribe to all events for a specific room."""
33
- queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
34
- if room_id not in self._room_subscribers:
35
- self._room_subscribers[room_id] = []
36
- self._room_subscribers[room_id].append(queue)
37
- return queue
38
-
39
- def unsubscribe_room(self, room_id: str, queue: asyncio.Queue[dict[str, Any]]) -> None:
40
- subs = self._room_subscribers.get(room_id, [])
41
- if queue in subs:
42
- subs.remove(queue)
43
-
44
- def broadcast(self, entry: MemoryEntry) -> None:
45
- data = entry.model_dump(mode="json")
46
- data.pop("embedding", None)
47
- event = {"type": "write", "entry": data}
48
- for pattern, queues in self._subscribers.items():
49
- if (
50
- pattern == "*"
51
- or entry.entity_path == pattern
52
- or entry.entity_path.startswith(pattern + "/")
53
- ):
54
- for queue in queues:
55
- try:
56
- queue.put_nowait(event)
57
- except asyncio.QueueFull:
58
- logger.debug("SSE queue full, dropping event")
59
-
60
- def broadcast_room_event(
61
- self,
62
- room_id: str,
63
- event_type: str,
64
- data: dict[str, Any],
65
- ) -> None:
66
- """Broadcast a room-scoped event (join, leave, write, read, etc.)."""
67
- event = {"type": event_type, "room_id": room_id, **data}
68
- queues = self._room_subscribers.get(room_id, [])
69
- for queue in queues:
70
- try:
71
- queue.put_nowait(event)
72
- except asyncio.QueueFull:
73
- logger.debug("Room SSE queue full, dropping event for room %s", room_id)
74
-
75
- async def event_generator(self, entity_path: str = "*", predicate=None):
76
- """Yield SSE events, optionally filtered per subscriber.
77
-
78
- ``predicate`` receives the entry dict and returns True when the
79
- subscriber may see it (used for per-user visibility). Events that
80
- fail the predicate — or whose predicate raises — are dropped.
81
- """
82
- queue = self.subscribe(entity_path)
83
- try:
84
- while True:
85
- event = await queue.get()
86
- if predicate is not None:
87
- try:
88
- if not predicate(event.get("entry", {})):
89
- continue
90
- except Exception:
91
- logger.debug("SSE visibility predicate failed — dropping event", exc_info=True)
92
- continue
93
- yield {
94
- "event": event["type"],
95
- "data": json.dumps(event["entry"], default=str),
96
- }
97
- finally:
98
- self.unsubscribe(entity_path, queue)
99
-
100
- async def room_event_generator(self, room_id: str):
101
- """SSE generator for room-scoped events."""
102
- queue = self.subscribe_room(room_id)
103
- try:
104
- while True:
105
- event = await queue.get()
106
- event_type = event.pop("type", "room_event")
107
- yield {
108
- "event": event_type,
109
- "data": json.dumps(event, default=str),
110
- }
111
- finally:
112
- self.unsubscribe_room(room_id, queue)