langgraph-api 0.0.22__py3-none-any.whl → 0.0.23__py3-none-any.whl

langgraph_api/queue.py

@@ -8,7 +8,7 @@ import structlog
 from langgraph.pregel.debug import CheckpointPayload, TaskResultPayload
 
 from langgraph_api.auth.custom import SimpleUser, normalize_user
-from langgraph_api.config import BG_JOB_DELAY, STATS_INTERVAL_SECS
+from langgraph_api.config import BG_JOB_HEARTBEAT, STATS_INTERVAL_SECS
 from langgraph_api.errors import (
     UserInterrupt,
     UserRollback,
@@ -21,7 +21,7 @@ from langgraph_api.stream import (
     astream_state,
     consume,
 )
-from langgraph_api.utils import AsyncConnectionProto, set_auth_ctx, with_user
+from langgraph_api.utils import set_auth_ctx, with_user
 from langgraph_storage.database import connect
 from langgraph_storage.ops import Runs, Threads
 from langgraph_storage.retry import RETRIABLE_EXCEPTIONS
@@ -46,6 +46,7 @@ def ms(after: datetime, before: datetime) -> int:
 async def queue(concurrency: int, timeout: float):
     loop = asyncio.get_running_loop()
     last_stats_secs: int | None = None
+    last_sweep_secs: int | None = None
     semaphore = asyncio.Semaphore(concurrency)
 
     def cleanup(task: asyncio.Task):
@@ -95,8 +96,15 @@ async def queue(concurrency: int, timeout: float):
 
     await logger.ainfo(f"Starting {concurrency} background workers")
     try:
+        tup: tuple[Run, int] | None = None
         while True:
             try:
+                # check if we need to sweep runs
+                do_sweep = (
+                    last_sweep_secs is None
+                    or loop.time() - last_sweep_secs > BG_JOB_HEARTBEAT * 2
+                )
+                # check if we need to update stats
                 if calc_stats := (
                     last_stats_secs is None
                     or loop.time() - last_stats_secs > STATS_INTERVAL_SECS
@@ -109,16 +117,16 @@ async def queue(concurrency: int, timeout: float):
                         available=concurrency - active,
                         active=active,
                     )
+                # wait for semaphore to respect concurrency
                 await semaphore.acquire()
                 exit = AsyncExitStack()
-                conn = await exit.enter_async_context(connect())
-                if calc_stats:
-                    stats = await Runs.stats(conn)
-                    await logger.ainfo("Queue stats", **stats)
-                if tup := await exit.enter_async_context(Runs.next(conn)):
+                # skip the wait, if 1st time, or got a run last time
+                wait = tup is None and last_stats_secs is not None
+                # try to get a run, handle it
+                if tup := await exit.enter_async_context(Runs.next(wait=wait)):
                     run_, attempt_ = tup
                     task = asyncio.create_task(
-                        worker(timeout, exit, conn, run_, attempt_),
+                        worker(timeout, exit, run_, attempt_),
                         name=f"run-{run_['run_id']}-attempt-{attempt_}",
                     )
                     task.add_done_callback(cleanup)
@@ -126,13 +134,23 @@ async def queue(concurrency: int, timeout: float):
                 else:
                     semaphore.release()
                     await exit.aclose()
-                    await asyncio.sleep(BG_JOB_DELAY)
+                # run stats and sweep if needed
+                if calc_stats or do_sweep:
+                    async with connect() as conn:
+                        # update stats if needed
+                        if calc_stats:
+                            stats = await Runs.stats(conn)
+                            await logger.ainfo("Queue stats", **stats)
+                        # sweep runs if needed
+                        if do_sweep:
+                            last_sweep_secs = loop.time()
+                            run_ids = await Runs.sweep(conn)
+                            logger.info("Sweeped runs", run_ids=run_ids)
             except Exception as exc:
                 # keep trying to run the scheduler indefinitely
                 logger.exception("Background worker scheduler failed", exc_info=exc)
                 semaphore.release()
                 await exit.aclose()
-                await asyncio.sleep(BG_JOB_DELAY)
     finally:
         logger.info("Shutting down background workers")
         for task in WORKERS:
@@ -187,7 +205,6 @@ async def set_auth_ctx_for_run(
 async def worker(
     timeout: float,
     exit: AsyncExitStack,
-    conn: AsyncConnectionProto,
     run: Run,
     attempt: int,
 ) -> WorkerResult:
@@ -201,7 +218,12 @@ async def worker(
     run_started_at = datetime.now(UTC)
     run_ended_at: str | None = None
 
-    async with set_auth_ctx_for_run(run["kwargs"]), Runs.enter(run_id) as done, exit:
+    async with (
+        connect() as conn,
+        set_auth_ctx_for_run(run["kwargs"]),
+        Runs.enter(run_id) as done,
+        exit,
+    ):
         temporary = run["kwargs"].get("temporary", False)
         run_created_at = run["created_at"].isoformat()
         await logger.ainfo(
@@ -323,9 +345,8 @@ async def worker(
                 run_ended_at=run_ended_at,
                 run_exec_ms=ms(datetime.now(UTC), run_started_at),
             )
+            await Runs.set_status(conn, run_id, "pending")
             raise
-            # Note we re-raise here, thus marking the run
-            # as available to be picked up by another worker
         except Exception as exc:
             exception = exc
             status = "error"

langgraph_api/schema.py

@@ -188,5 +188,6 @@ class ThreadUpdateResponse(TypedDict):
 
 class QueueStats(TypedDict):
     n_pending: int
+    n_running: int
     max_age_secs: datetime | None
     med_age_secs: datetime | None

{langgraph_api-0.0.22.dist-info → langgraph_api-0.0.23.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langgraph-api
-Version: 0.0.22
+Version: 0.0.23
 Summary: 
 License: Elastic-2.0
 Author: Nuno Campos

{langgraph_api-0.0.22.dist-info → langgraph_api-0.0.23.dist-info}/RECORD

@@ -17,7 +17,7 @@ langgraph_api/auth/middleware.py,sha256=qc7SbaFoeWaqxS1wbjZ2PPQ4iI2p9T0shWL7c6g0
 langgraph_api/auth/noop.py,sha256=Bk6Nf3p8D_iMVy_OyfPlyiJp_aEwzL-sHrbxoXpCbac,586
 langgraph_api/auth/studio_user.py,sha256=FzFQRROKDlA9JjtBuwyZvk6Mbwno5M9RVYjDO6FU3F8,186
 langgraph_api/cli.py,sha256=r7NJVIdTQ9mQ6_X01tk_I0ktlgn9odH0B8J53oaySz4,12022
-langgraph_api/config.py,sha256=vVANO2F7O94XvnXc0krwQrE8o7E-nEJtZQVGLbMoxZc,4272
+langgraph_api/config.py,sha256=mkuhRHjUHA5cjzcSy0nYYS6K28FZ_3_FRl9s0sg_030,4225
 langgraph_api/cron_scheduler.py,sha256=MW41-TSGUe5OuXycFTy7Ax7ypxHVAv-0ImLonRT8h8o,2629
 langgraph_api/errors.py,sha256=Bu_i5drgNTyJcLiyrwVE_6-XrSU50BHf9TDpttki9wQ,1690
 langgraph_api/graph.py,sha256=FombjYQkqj8jrXJFEVkl3m2UyFcq5nSVNswR2HoRsQY,16385
@@ -25,14 +25,14 @@ langgraph_api/http.py,sha256=XrbyxpjtfSvnaWWh5ZLGpgZmY83WoDCrP_1GPguNiXI,4712
 langgraph_api/js/.gitignore,sha256=qAah3Fq0HWAlfRj5ktZyC6QRQIsAolGLRGcRukA1XJI,33
 langgraph_api/js/base.py,sha256=BpE8-xkUp8HFPRjSKx1tfUQubvoV4jYl6OwZdre3veI,209
 langgraph_api/js/build.mts,sha256=43_LQDjmtEyV6Sj6IZo7Mp6Y8zrsBbMP_5sWXQ_xBsY,1372
-langgraph_api/js/client.mts,sha256=AWDSJHAFFtCKU8JkSg0NEPt85jtovk64SYDesIALtW0,22684
+langgraph_api/js/client.mts,sha256=5cJgJLmER3fpf-QQU8v2cjR593ROXcQk6vXd71cntsw,22772
 langgraph_api/js/client.new.mts,sha256=FskIvfdS4MvzoD07fJuHMoz26Puh_Mkwh2Lh6jKFj60,23658
 langgraph_api/js/errors.py,sha256=Cm1TKWlUCwZReDC5AQ6SgNIVGD27Qov2xcgHyf8-GXo,361
 langgraph_api/js/global.d.ts,sha256=zR_zLYfpzyPfxpEFth5RgZoyfGulIXyZYPRf7cU0K0Y,106
-langgraph_api/js/package.json,sha256=Z7Yluyxl4WbrcagLzBNB2ZL0cihX9dWkUn08vYrLoUw,840
+langgraph_api/js/package.json,sha256=AmpkMzr96yF9xZ7bCrSApF-j7PJH6WeALn9HpPGBnmQ,840
 langgraph_api/js/remote.py,sha256=D9cqcEgXau-fm_trpNwCHMra5BXntgUa469lgs_a9JQ,622
 langgraph_api/js/remote_new.py,sha256=T_Vr8459bax1C9xxqz_ZYmGivq5Vhspg2Iu9TL0Qc-Q,22707
-langgraph_api/js/remote_old.py,sha256=2a-3ooAYUZs8aPsfnXafbBd4pP7lRmokoU7TiO7P9Js,22546
+langgraph_api/js/remote_old.py,sha256=wC0wpYR4xv6Xqqq1PboIluViVeJ9ETFVTUEl53FSZyo,22578
 langgraph_api/js/schema.py,sha256=7idnv7URlYUdSNMBXQcw7E4SxaPxCq_Oxwnlml8q5ik,408
 langgraph_api/js/server_sent_events.py,sha256=DLgXOHauemt7706vnfDUCG1GI3TidKycSizccdz9KgA,3702
 langgraph_api/js/src/graph.mts,sha256=J-M-vYHj1G5tyonPUym3ePNGqGYtspPCrZOgr92xKb4,3171
@@ -44,7 +44,7 @@ langgraph_api/js/src/schema/types.template.mts,sha256=oAwnwWOgkEAQ3EouB8dG5Mdg4H
 langgraph_api/js/src/utils/importMap.mts,sha256=pX4TGOyUpuuWF82kXcxcv3-8mgusRezOGe6Uklm2O5A,1644
 langgraph_api/js/src/utils/pythonSchemas.mts,sha256=98IW7Z_VP7L_CHNRMb3_MsiV3BgLE2JsWQY_PQcRR3o,685
 langgraph_api/js/src/utils/serde.mts,sha256=OuyyO9btvwWd55rU_H4x91dFEJiaPxL-lL9O6Zgo908,742
-langgraph_api/js/tests/api.test.mts,sha256=qJdISmuKo7AWZmR4thY9eq8h4EL5bWKlhI5x7LCJ8dQ,56919
+langgraph_api/js/tests/api.test.mts,sha256=fiQDSdfJj6DhOqQGghFq0RknIYnTPQ7nH_QfBjs7kAk,56931
 langgraph_api/js/tests/compose-postgres.yml,sha256=pbNfeqVUqhWILBuUdwAgQOYsVU_fgkCVm0YlTgU8og8,1721
 langgraph_api/js/tests/graphs/.gitignore,sha256=26J8MarZNXh7snXD5eTpV3CPFTht5Znv8dtHYCLNfkw,12
 langgraph_api/js/tests/graphs/agent.mts,sha256=E9WMv0alMv0njUEECqEsqoRk9NXJUgXW7SyQJ3GOZ8k,5396
@@ -55,9 +55,9 @@ langgraph_api/js/tests/graphs/nested.mts,sha256=4G7jSOSaFVQAza-_ARbK-Iai1biLlF2D
 langgraph_api/js/tests/graphs/package.json,sha256=U1e03aGujf7LHM9j_01fwbD7hcC59kxJM4ZvUECVF3o,118
 langgraph_api/js/tests/graphs/weather.mts,sha256=A7mLK3xW8h5B-ZyJNAyX2M2fJJwzPJzXs4DYesJwreQ,1655
 langgraph_api/js/tests/graphs/yarn.lock,sha256=N-LBwqNAve0B_VvVZvSAGVkCCo9AQO6VDM-AznGVVLQ,10407
-langgraph_api/js/tests/parser.test.mts,sha256=3zAbboUNhI-cY3hj4Ssr7J-sQXCBTeeI1ItrkG0Ftuk,26257
+langgraph_api/js/tests/parser.test.mts,sha256=dEC8KTqKygeb1u39ZvpPqCT4HtfPD947nLmITt2buxA,27883
 langgraph_api/js/tests/utils.mts,sha256=2kTybJ3O7Yfe1q3ehDouqV54ibXkNzsPZ_wBZLJvY-4,421
-langgraph_api/js/yarn.lock,sha256=slQ8Gr_G-_8tyLgjhE6ABtvVAnv0PTZzei5Eo6Xh-1g,105950
+langgraph_api/js/yarn.lock,sha256=SaJYDS1TmC5J94zGeoi2gHC3_lnF-h8xQcVND6wWd1c,102827
 langgraph_api/lifespan.py,sha256=Uj7NV-NqxxD1fgx_umM9pVqclcy-VlqrIxDljyj2he0,1820
 langgraph_api/logging.py,sha256=tiDNrEFwqaIdL5ywZv908OXlzzfXsPCws9GXeoFtBV8,3367
 langgraph_api/metadata.py,sha256=wvNCHvejBiO_VaAF12kbdlNDn1QXH4Fh_buvHD9J60U,3276
@@ -67,9 +67,9 @@ langgraph_api/middleware/private_network.py,sha256=eYgdyU8AzU2XJu362i1L8aSFoQRiV
 langgraph_api/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langgraph_api/models/run.py,sha256=1TzTmagDXFQD_LhIMRdguZHmrPSzztq1wiMjF63d2fc,9843
 langgraph_api/patch.py,sha256=94ddcTSZJe22JcpjxiSNjFZdYVnmeoWjk4IX4iBSoyk,1249
-langgraph_api/queue.py,sha256=qRuM09mz8o7CkrBIwg-9lV3SW0TehdVGgeXvHc_adYk,13647
+langgraph_api/queue.py,sha256=2sw9HB2cYVBhYUNA3F7lcJAgRjhQJXhA_HNGhFt2BW8,14508
 langgraph_api/route.py,sha256=fM4qYCGbmH0a3_cV8uKocb1sLklehxO6HhdRXqLK6OM,4421
-langgraph_api/schema.py,sha256=4aZCFY-dxd_nTot71bdcd9S8QCIgKajuRyj0p2QfgJ4,5291
+langgraph_api/schema.py,sha256=mgam5lpuqZnrNWMm_0nQ95683gCnCvQNRKbiuFj7z8Q,5310
 langgraph_api/serde.py,sha256=VoJ7Z1IuqrQGXFzEP1qijAITtWCrmjtVqlCRuScjXJI,3533
 langgraph_api/server.py,sha256=mKJWBuHN5HFHCQIL_FtH04wyFabR4mR6WmQcioI-_Ns,2071
 langgraph_api/sse.py,sha256=2wNodCOP2eg7a9mpSu0S3FQ0CHk2BBV_vv0UtIgJIcc,4034
@@ -83,15 +83,15 @@ langgraph_license/validation.py,sha256=Uu_G8UGO_WTlLsBEY0gTVWjRR4czYGfw5YAD3HLZo
 langgraph_storage/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langgraph_storage/checkpoint.py,sha256=V4t2GwYEJdPCHbhq_4Udhlv0TWKDzlMu_rlNPdTDc50,3589
 langgraph_storage/database.py,sha256=Nr5zE9Fur3-tESkqe7xNXMf2QlBuw3H0CUie7jVa6Q4,6003
-langgraph_storage/ops.py,sha256=Pv829coa6wCYKKdYBR0IZaKMtkgw1g9EmOIayI_GJXg,68242
+langgraph_storage/ops.py,sha256=Lp6ICiPq-Td8mE7dxBRRvkw6vLh1-hkIJBw-b6XrN84,68533
 langgraph_storage/queue.py,sha256=6cTZ0ubHu3S1T43yxHMVOwsQsDaJupByiU0sTUFFls8,3261
 langgraph_storage/retry.py,sha256=uvYFuXJ-T6S1QY1ZwkZHyZQbsvS-Ab68LSbzbUUSI2E,696
 langgraph_storage/store.py,sha256=D-p3cWc_umamkKp-6Cz3cAriSACpvM5nxUIvND6PuxE,2710
 langgraph_storage/ttl_dict.py,sha256=FlpEY8EANeXWKo_G5nmIotPquABZGyIJyk6HD9u6vqY,1533
 logging.json,sha256=3RNjSADZmDq38eHePMm1CbP6qZ71AmpBtLwCmKU9Zgo,379
 openapi.json,sha256=L1Ap1o6oAg2vKqjWohq4gl1-GpxezjlddS_oZHjS-xE,125227
-langgraph_api-0.0.22.dist-info/LICENSE,sha256=ZPwVR73Biwm3sK6vR54djCrhaRiM4cAD2zvOQZV8Xis,3859
-langgraph_api-0.0.22.dist-info/METADATA,sha256=7fHD4Lr32sZxoJaG4mBi-sUDq4WF6l-6cfK6BGHi-jI,4038
-langgraph_api-0.0.22.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-langgraph_api-0.0.22.dist-info/entry_points.txt,sha256=3EYLgj89DfzqJHHYGxPH4A_fEtClvlRbWRUHaXO7hj4,77
-langgraph_api-0.0.22.dist-info/RECORD,,
+langgraph_api-0.0.23.dist-info/LICENSE,sha256=ZPwVR73Biwm3sK6vR54djCrhaRiM4cAD2zvOQZV8Xis,3859
+langgraph_api-0.0.23.dist-info/METADATA,sha256=mwWP3njPolhdWlLEDl7uQZMF-ptzkDXFRfU_BRaDSF0,4038
+langgraph_api-0.0.23.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+langgraph_api-0.0.23.dist-info/entry_points.txt,sha256=3EYLgj89DfzqJHHYGxPH4A_fEtClvlRbWRUHaXO7hj4,77
+langgraph_api-0.0.23.dist-info/RECORD,,

langgraph_storage/ops.py

@@ -1160,63 +1160,65 @@ class Runs(Authenticated):
 
     @asynccontextmanager
     @staticmethod
-    async def next(conn: InMemConnectionProto) -> AsyncIterator[tuple[Run, int] | None]:
+    async def next(wait: bool) -> AsyncIterator[tuple[Run, int] | None]:
         """Get the next run from the queue, and the attempt number.
         1 is the first attempt, 2 is the first retry, etc."""
         now = datetime.now(UTC)
 
-        pending_runs = sorted(
-            [
-                run
-                for run in conn.store["runs"]
-                if run["status"] == "pending" and run.get("created_at", now) < now
-            ],
-            key=lambda x: x.get("created_at", datetime.min),
-        )
+        if wait:
+            await asyncio.sleep(0.5)
+        else:
+            await asyncio.sleep(0)
 
-        if not pending_runs:
-            yield None
-            return
+        async with connect() as conn:
+            pending_runs = sorted(
+                [
+                    run
+                    for run in conn.store["runs"]
+                    if run["status"] == "pending" and run.get("created_at", now) < now
+                ],
+                key=lambda x: x.get("created_at", datetime.min),
+            )
 
-        # Try to lock and get the first available run
-        for run in pending_runs:
-            run_id = run["run_id"]
-            thread_id = run["thread_id"]
-            lock = conn.locks[thread_id]
-            acquired = lock.acquire(blocking=False)
-            if not acquired:
-                continue
-            try:
-                if run["status"] != "pending":
+            if not pending_runs:
+                yield None
+                return
+
+            # Try to lock and get the first available run
+            for run in pending_runs:
+                run_id = run["run_id"]
+                thread_id = run["thread_id"]
+                lock = conn.locks[thread_id]
+                acquired = lock.acquire(blocking=False)
+                if not acquired:
                     continue
+                try:
+                    if run["status"] != "pending":
+                        continue
 
-                thread = next(
-                    (
-                        t
-                        for t in conn.store["threads"]
-                        if t["thread_id"] == run["thread_id"]
-                    ),
-                    None,
-                )
-
-                if thread is None:
-                    await logger.awarning(
-                        "Unexpected missing thread in Runs.next",
-                        thread_id=run["thread_id"],
+                    thread = next(
+                        (
+                            t
+                            for t in conn.store["threads"]
+                            if t["thread_id"] == run["thread_id"]
+                        ),
+                        None,
                     )
-                    continue
 
-                # Increment attempt counter
-                attempt = await conn.retry_counter.increment(run_id)
-                enriched_run = {
-                    **run,
-                    "thread_created_at": thread.get("created_at", now),
-                }
-                yield enriched_run, attempt
-            finally:
-                lock.release()
-            return
-        yield None
+                    if thread is None:
+                        await logger.awarning(
+                            "Unexpected missing thread in Runs.next",
+                            thread_id=run["thread_id"],
+                        )
+                        continue
+
+                    # Increment attempt counter
+                    attempt = await conn.retry_counter.increment(run_id)
+                    yield run, attempt
+                finally:
+                    lock.release()
+                return
+            yield None
 
     @asynccontextmanager
     @staticmethod
@@ -1247,6 +1249,11 @@ class Runs(Authenticated):
                 # Clean up this queue
                 await stream_manager.remove_queue(run_id, queue)
 
+    @staticmethod
+    async def sweep(conn: InMemConnectionProto) -> list[UUID]:
+        """Sweep runs that are no longer running"""
+        return []
+
     @staticmethod
     def _merge_jsonb(*objects: dict) -> dict:
         """Mimics PostgreSQL's JSONB merge behavior"""