stageflow-framework 0.1.2__tar.gz → 0.1.4__tar.gz

{stageflow_framework-0.1.2 → stageflow_framework-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stageflow-framework
-Version: 0.1.2
+Version: 0.1.4
 Summary: StageFlow: Pipeline framework for stages
 Author: lethargy
 Requires-Python: >=3.10

{stageflow_framework-0.1.2 → stageflow_framework-0.1.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "stageflow-framework"
-version = "0.1.2"
+version = "0.1.4"
 description = "StageFlow: Pipeline framework for stages"
 readme = "README.md"
 authors = [

{stageflow_framework-0.1.2 → stageflow_framework-0.1.4}/stageflow/core/session.py

@@ -87,15 +87,20 @@ class Session:
             self._pending_inputs.setdefault(type_, []).append(entry)
         return entry
 
-    async def wait_input(self, type_: str, timeout: float | None = None):
-        # Deliver buffered input if it arrived before waiter.
+    def start_wait_input(self, type_: str) -> asyncio.Future:
         pending = self._pending_inputs.get(type_)
         if pending:
-            return pending.pop(0)
+            loop = asyncio.get_running_loop()
+            fut = loop.create_future()
+            fut.set_result(pending.pop(0))
+            return fut
         loop = asyncio.get_running_loop()
         fut = loop.create_future()
         self._waiting.setdefault(type_, []).append(fut)
         self.emit(Event(type="waiting_for_input", session_id=self.id, payload={"type": type_}))
+        return fut
+
+    async def finish_wait_input(self, type_: str, fut: asyncio.Future, timeout: float | None = None):
         try:
             return await asyncio.wait_for(fut, timeout=timeout)
         except asyncio.TimeoutError:
@@ -108,6 +113,10 @@ class Session:
                 if not waiters:
                     del self._waiting[type_]
 
+    async def wait_input(self, type_: str, timeout: float | None = None):
+        fut = self.start_wait_input(type_)
+        return await self.finish_wait_input(type_, fut, timeout=timeout)
+
     def last_input(self, type_: str | None = None):
         if not self.input_history:
             return None

{stageflow_framework-0.1.2 → stageflow_framework-0.1.4}/stageflow/core/stage.py

@@ -1,3 +1,4 @@
+import asyncio
 import copy
 from typing import Any, TYPE_CHECKING
 import yaml
@@ -121,14 +122,29 @@ class BaseStage:
             payload=payload or {},
         ))
 
-    async def wait_input(self, type_: str, timeout: float | None = None):
+    def _get_allowed_input(self, type_: str) -> InputSpec | None:
         if self.allowed_inputs:
             allowed = {spec.type for spec in self.allowed_inputs if spec.type}
             if allowed and type_ not in allowed:
                 raise ValueError(f"Input type '{type_}' is not allowed for stage '{self.stage_name}'")
-            matching = next((spec for spec in self.allowed_inputs if spec.type == type_), None)
-        else:
-            matching = None
+            return next((spec for spec in self.allowed_inputs if spec.type == type_), None)
+        return None
+
+    def start_wait_input(self, type_: str) -> asyncio.Future:
+        self._get_allowed_input(type_)
+        return self.session.start_wait_input(type_)
+
+    async def finish_wait_input(self, type_: str, fut: asyncio.Future, timeout: float | None = None):
+        matching = self._get_allowed_input(type_)
+        result = await self.session.finish_wait_input(type_, fut, timeout=timeout)
+        if result is None:
+            return None
+        if matching and matching.payload_schema is not None:
+            validate_schema(result.get("payload", {}), matching.payload_schema, "Input payload")
+        return result
+
+    async def wait_input(self, type_: str, timeout: float | None = None):
+        matching = self._get_allowed_input(type_)
         result = await self.session.wait_input(type_, timeout=timeout)
         if result is None:
             return None

{stageflow_framework-0.1.2 → stageflow_framework-0.1.4}/stageflow_framework.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stageflow-framework
-Version: 0.1.2
+Version: 0.1.4
 Summary: StageFlow: Pipeline framework for stages
 Author: lethargy
 Requires-Python: >=3.10

{stageflow_framework-0.1.2 → stageflow_framework-0.1.4}/stageflow_framework.egg-info/SOURCES.txt

@@ -34,6 +34,8 @@ tests/test_full_pipeline.py
 tests/test_jsonlogic.py
 tests/test_payload_validation.py
 tests/test_pipeline_tester.py
+tests/test_session_wait_input.py
 tests/test_snapshot.py
+tests/test_stage_wait_input.py
 tests/test_std_stages.py
 tests/test_subpipeline.py

stageflow_framework-0.1.4/tests/test_session_wait_input.py

@@ -0,0 +1,40 @@
+import unittest
+
+from stageflow.core.context import Context
+from stageflow.core.node import TerminalNode
+from stageflow.core.pipeline import Pipeline
+from stageflow.core.session import Session
+
+
+def _make_pipeline():
+    end = TerminalNode(id="end", type="terminal")
+    return Pipeline(entry="end", nodes=[end])
+
+
+class SessionWaitInputTests(unittest.IsolatedAsyncioTestCase):
+    async def test_start_and_finish_wait_input_success(self):
+        session = Session(id="s1", pipeline=_make_pipeline(), context=Context())
+        waiter = session.start_wait_input("user")
+
+        await session.input("user", {"id": 1})
+        result = await session.finish_wait_input("user", waiter, timeout=1)
+
+        self.assertEqual(result["payload"]["id"], 1)
+        self.assertIn("waiting_for_input", [e.type for e in session.event_history])
+
+    async def test_start_finish_timeout_and_pending_consumption(self):
+        session = Session(id="s2", pipeline=_make_pipeline(), context=Context())
+        waiter = session.start_wait_input("later")
+        result = await session.finish_wait_input("later", waiter, timeout=0.01)
+
+        self.assertIsNone(result)
+        self.assertIn("input_timeout", [e.type for e in session.event_history])
+
+        await session.input("later", {"value": 2})
+        waiter2 = session.start_wait_input("later")
+        result2 = await session.finish_wait_input("later", waiter2, timeout=1)
+        self.assertEqual(result2["payload"]["value"], 2)
+
+
+if __name__ == "__main__":
+    unittest.main()

stageflow_framework-0.1.4/tests/test_stage_wait_input.py

@@ -0,0 +1,64 @@
+import unittest
+
+from stageflow.core.context import Context
+from stageflow.core.event import InputSpec
+from stageflow.core.node import TerminalNode
+from stageflow.core.pipeline import Pipeline
+from stageflow.core.session import Session
+from stageflow.core.stage import BaseStage
+
+
+def _make_pipeline():
+    end = TerminalNode(id="end", type="terminal")
+    return Pipeline(entry="end", nodes=[end])
+
+
+class StageWaitInputTests(unittest.IsolatedAsyncioTestCase):
+    async def test_start_and_finish_wait_input_success(self):
+        session = Session(id="stage1", pipeline=_make_pipeline(), context=Context())
+
+        class WaitStage(BaseStage):
+            allowed_inputs = [InputSpec(type="user", payload_schema={"id": int})]
+
+            async def run(self):
+                return None
+
+        stage = WaitStage(stage_id="s1", config={}, arguments={}, outputs={}, session=session)
+        fut = stage.start_wait_input("user")
+        await session.input("user", {"id": 10})
+
+        result = await stage.finish_wait_input("user", fut, timeout=1)
+        self.assertEqual(result["payload"]["id"], 10)
+
+    async def test_finish_wait_input_validates_payload(self):
+        session = Session(id="stage2", pipeline=_make_pipeline(), context=Context())
+
+        class WaitStage(BaseStage):
+            allowed_inputs = [InputSpec(type="user", payload_schema={"id": int})]
+
+            async def run(self):
+                return None
+
+        stage = WaitStage(stage_id="s1", config={}, arguments={}, outputs={}, session=session)
+        fut = stage.start_wait_input("user")
+        await session.input("user", {"id": "bad"})
+
+        with self.assertRaises(ValueError):
+            await stage.finish_wait_input("user", fut, timeout=1)
+
+    async def test_start_wait_input_rejects_disallowed_type(self):
+        session = Session(id="stage3", pipeline=_make_pipeline(), context=Context())
+
+        class LimitedStage(BaseStage):
+            allowed_inputs = [InputSpec(type="only")]
+
+            async def run(self):
+                return None
+
+        stage = LimitedStage(stage_id="s1", config={}, arguments={}, outputs={}, session=session)
+        with self.assertRaises(ValueError):
+            stage.start_wait_input("other")
+
+
+if __name__ == "__main__":
+    unittest.main()