diff --git a/docs/core/action.mdx b/docs/core/action.mdx index e6c8cbe..36bd25d 100644 --- a/docs/core/action.mdx +++ b/docs/core/action.mdx @@ -113,6 +113,10 @@ Each manager is a key-value store backed by a `Map`. See [BaseManager](./base-ma description: 'Stores registered runtime function definitions keyed by runtimeName.', type: 'RuntimeFunctionManager', }, + subFlowFunctions: { + description: 'Stores sub flow functions returned by a runtime function at execution time, keyed by the identifier reported to Aquila. Entries expire 30 minutes after they were registered.', + type: 'SubFlowFunctionManager', + }, dataTypes: { description: 'Stores registered data type definitions keyed by identifier.', type: 'DataTypeManager', @@ -174,6 +178,19 @@ entry is also created automatically. action.registerRuntimeFunction(FibonacciRuntimeFunction); ``` +### `registerSubFlowFunction` + +```ts +action.registerSubFlowFunction(definition: SubFlowFunctionDefinition): string +``` + +Registers a sub flow function returned from a `run()` block and returns the identifier +Aquila invokes it by. Called automatically while a result is sent, so a runtime +function only has to return the definition — see +[Returning a sub flow function](./runtime-function.mdx#returning-a-sub-flow-function). +The registration expires after 30 minutes +([`SubFlowFunctionManager`](./sub-flow-function-manager.mdx#lifetime)). + ### `registerDataTypeClass` ```ts diff --git a/docs/core/meta.json b/docs/core/meta.json index 6090c2e..f50e6c0 100644 --- a/docs/core/meta.json +++ b/docs/core/meta.json @@ -10,6 +10,7 @@ "base-manager", "function-manager", "runtime-function-manager", + "sub-flow-function-manager", "event-manager", "runtime-event-manager", "data-type-manager", diff --git a/docs/core/runtime-function-manager.mdx b/docs/core/runtime-function-manager.mdx index b0e42e8..f4373a9 100644 --- a/docs/core/runtime-function-manager.mdx +++ b/docs/core/runtime-function-manager.mdx @@ -22,13 +22,23 @@ The interface implemented by all RuntimeFunction classes. ```ts interface RuntimeFunctionRunnable { - run(...args: (PlainValue | undefined)[]): Promise | PlainValue; + run(...args: (PlainValue | undefined)[]): Promise | RuntimeFunctionReturn; } ``` The first argument to `run()` is always a [`FunctionContext`](#functioncontext) object. Subsequent arguments correspond to the declared parameters in order. +### `RuntimeFunctionReturn` + +What a `run()` may produce: raw data, or a sub flow function that Hercules registers +and reports to Aquila. See +[Returning a sub flow function](./runtime-function.mdx#returning-a-sub-flow-function). + +```ts +type RuntimeFunctionReturn = PlainValue | SubFlowFunctionDefinition; +``` + ### `RuntimeFunctionProps` The descriptor stored for each registered runtime function. @@ -47,7 +57,7 @@ interface RuntimeFunctionProps { alias?: Translation[]; displayIcon?: string; design?: string; - handler: (...args: (PlainValue | undefined)[]) => Promise | PlainValue; + handler: (...args: (PlainValue | undefined)[]) => Promise | RuntimeFunctionReturn; } ``` @@ -107,7 +117,7 @@ interface RuntimeFunctionProps { }, handler: { description: 'Populated automatically from the class run() method during registration.', - type: '(...args: (PlainValue | undefined)[]) => Promise | PlainValue', + type: '(...args: (PlainValue | undefined)[]) => Promise | RuntimeFunctionReturn', required: true, }, }} diff --git a/docs/core/runtime-function.mdx b/docs/core/runtime-function.mdx index 49ccee1..2946a9c 100644 --- a/docs/core/runtime-function.mdx +++ b/docs/core/runtime-function.mdx @@ -42,6 +42,58 @@ export class FibonacciRuntimeFunction { } ``` +## Returning a sub flow function + +A `run()` block may return a **sub flow function** instead of raw data. Hercules +registers its `run` internally and reports the function to Aquila as a sub flow +function result, so a later node of the flow can invoke it like any other sub flow. + +Because nothing in the module definition describes a function that only comes into +existence at runtime, its name, description and I/O schema have to be declared with it. +The schemas are written as [Zod](https://zod.dev) schemas — the same way +[`@Schema`](./decorator-schema.mdx) describes a data type — and are converted to JSON +schemas while the result is sent: + +```ts +import { z } from 'zod'; +import { SubFlowFunctionDefinition } from '@code0-tech/hercules'; + +@Identifier('greeter_runtime') +@Signature('(greeting: STRING): TRANSFORM') +export class GreeterRuntimeFunction { + run(context: FunctionContext, greeting: string): SubFlowFunctionDefinition { + return { + name: 'greet', + description: `Greets a name with "${greeting}"`, + inputSchema: z.object({ name: z.string() }), + outputSchema: z.string(), + run: (name) => `${greeting} ${String(name)}`, + }; + } +} +``` + +The input schema is an object keyed by parameter name, the output schema describes the +value the function returns. Each is generated from its own side of the zod schema — +the input schema from the schema's input side, the output schema from its output side +— so a schema with a transform or a default describes each direction correctly. + +The [`@Signature`](./decorator-signature.mdx) return type is the function data type +the returned function matches — not the raw data type it eventually produces: + +| Returned function | Signature return type | +|---|---| +| `(value) => result` | `TRANSFORM` | +| `(value) => void` | `CONSUMER` | +| `(value) => boolean` | `PREDICATE` | +| `(a, b) => number` | `COMPARATOR` | +| `() => void` | `RUNNABLE` | + +Each returned function is registered under its own identifier, so returning one per +execution never overwrites a function a running flow still holds. A registration has a +maximum lifetime of **30 minutes** and is then dropped from memory — see +[`action.subFlowFunctions`](./sub-flow-function-manager.mdx). + ## Error Handling Throw a `RuntimeError` to signal a recoverable error back to the platform: diff --git a/docs/core/sub-flow-function-manager.mdx b/docs/core/sub-flow-function-manager.mdx new file mode 100644 index 0000000..faa1381 --- /dev/null +++ b/docs/core/sub-flow-function-manager.mdx @@ -0,0 +1,99 @@ +--- +title: SubFlowFunctionManager +description: Stores sub flow functions returned at execution time +--- + +import { TypeTable } from 'fumadocs-ui/components/type-table'; + +`SubFlowFunctionManager` stores the sub flow functions a runtime function returned from +its `run()` block (see +[Returning a sub flow function](./runtime-function.mdx#returning-a-sub-flow-function)). +It extends [`BaseManager`](./base-manager.mdx) with key type `string` (the identifier +reported to Aquila) and value type `SubFlowFunctionDefinition`. + +Unlike the other managers it is not filled by registering classes up front: entries +appear while the action runs, each time a `run()` returns a function. + +```ts +action.subFlowFunctions.size; +action.subFlowFunctions.delete(identifier); +``` + +## Lifetime + +Aquila gives no signal that a flow is done with a sub flow function, so each entry has +a maximum lifetime of `SUB_FLOW_FUNCTION_TTL_MS` — **30 minutes** — counted from when +it was registered and never extended by an invocation. Once it passes, the function is +dropped from memory and can no longer be invoked, so the registry cannot grow without +bound while the action runs. + +Expired entries are cleared by a sweep scheduled for the next expiry, so memory is +freed without waiting for another lookup, and additionally on every read and write, so +an expired function is never invoked even if the sweep could not run. Delete an entry +to drop it earlier: + +```ts +action.subFlowFunctions.delete(identifier); +``` + +An execution request naming a function that has expired fails the node with a +`SUB_FLOW_FUNCTION_EXPIRED` error rather than going unanswered, so the flow reports a +clear failure instead of waiting for a result that cannot come. Identifiers this +action never minted are still only logged — they address a function the action does +not know at all. + +## Interfaces + +### `SubFlowFunctionDefinition` + +Returned from a `run()` block in place of raw data. + +```ts +interface SubFlowFunctionDefinition { + name: string; + description: string; + inputSchema: SubFlowFunctionSchema; + outputSchema: SubFlowFunctionSchema; + run(...args: (PlainValue | undefined)[]): Promise | R; +} +``` + + Promise | PlainValue', + required: true, + }, + }} +/> + +### `SubFlowFunctionSchema` + +Normally a Zod schema, converted to a JSON schema when the result is sent. A plain +JSON schema object or a ready-made protobuf `Struct` — e.g. the `inputSchema` of a +[`SubFlow`](./runtime-function.mdx) parameter — is passed through unchanged. + +```ts +type SubFlowFunctionSchema = ZodTypeAny | Struct | Record; +``` diff --git a/py/examples/simple-example-py/functions/greeter_runtime_function.py b/py/examples/simple-example-py/functions/greeter_runtime_function.py new file mode 100644 index 0000000..91c6534 --- /dev/null +++ b/py/examples/simple-example-py/functions/greeter_runtime_function.py @@ -0,0 +1,58 @@ +"""A runtime function whose run block returns another function instead of raw data. + +Hercules registers the returned ``run`` internally and reports it to Aquila as a sub +flow function, so the flow can invoke it later. Name, description and the I/O schema +are declared here because nothing in the module definition describes a function that +is only created at runtime. The Pydantic models are converted to JSON schemas while +the result is sent. +""" +from hercules import ( + Description, + DisplayMessage, + FunctionContext, + Identifier, + Name, + Parameter, + RuntimeFunctionRunnable, + Signature, + SubFlowFunctionDefinition, +) +from hercules.schema import BaseModel, RootModel + + +class GreetInput(BaseModel): + """The parameters the returned function accepts, keyed by parameter name.""" + + name: str + + +class GreetOutput(RootModel[str]): + """The value the returned function produces.""" + + +@Identifier("greeter_runtime") +@Signature("(greeting: STRING): TRANSFORM") +@Name({"code": "en-US", "content": "Greeter"}) +@DisplayMessage({"code": "en-US", "content": "Build a greeter saying ${greeting}"}) +@Parameter( + { + "runtime_name": "greeting", + "name": [{"code": "en-US", "content": "Greeting"}], + "description": [ + { + "code": "en-US", + "content": "The greeting the returned function prefixes a name with", + } + ], + } +) +@Description({"code": "en-US", "content": "Returns a function that greets a name"}) +class GreeterRuntimeFunction(RuntimeFunctionRunnable): + def run(self, context: FunctionContext, greeting: str) -> SubFlowFunctionDefinition: + return SubFlowFunctionDefinition( + name="greet", + description=f'Greets a name with "{greeting}"', + input_schema=GreetInput, + output_schema=GreetOutput, + run=lambda name: f"{greeting} {name}", + ) diff --git a/py/examples/simple-example-py/index.py b/py/examples/simple-example-py/index.py index f00b10a..6997af2 100644 --- a/py/examples/simple-example-py/index.py +++ b/py/examples/simple-example-py/index.py @@ -9,6 +9,7 @@ from functions.fibonacci_runtime_function import FibonacciRuntimeFunction from functions.for_each_runtime_function import ForEachRuntimeFunction from functions.for_each_consumers_runtime_function import ForEachConsumersRuntimeFunction +from functions.greeter_runtime_function import GreeterRuntimeFunction action = Action( os.environ.get("ACTION_ID", "testing-action"), @@ -39,6 +40,9 @@ # Runtime function that takes a list of consumers (each an inline ${signature} sub flow reference). action.register_runtime_function(ForEachConsumersRuntimeFunction) +# Runtime function that returns a sub flow function instead of raw data. +action.register_runtime_function(GreeterRuntimeFunction) + # Data type: derived from a schema. action.register_data_type_class(EmailDataType) diff --git a/py/hercules/__init__.py b/py/hercules/__init__.py index b99d61b..96904d7 100644 --- a/py/hercules/__init__.py +++ b/py/hercules/__init__.py @@ -18,6 +18,8 @@ RuntimeError, RuntimeFlowTypeSetting_UniquenessScope, SubFlow, + SubFlowFunctionDefinition, + SubFlowFunctionSchema, Translation, ) from hercules.events import CodeZeroEvent @@ -55,6 +57,10 @@ from hercules.manager.function import FunctionManager from hercules.manager.runtime_event import RuntimeEventManager from hercules.manager.runtime_function import RuntimeFunctionManager +from hercules.manager.sub_flow_function import ( + SUB_FLOW_FUNCTION_TTL_SECONDS, + SubFlowFunctionManager, +) from hercules.models.datatype import DataTypeProps, DataTypeRunnable from hercules.models.event import EventModel, EventSettingProps from hercules.models.function import FunctionParameterProps, FunctionProps @@ -63,7 +69,8 @@ from hercules.action import Action, is_export_mode, registered_actions __all__ = [ - "Translation", "PlainValue", "SubFlow", "FunctionContext", + "Translation", "PlainValue", "SubFlow", "SubFlowFunctionDefinition", + "SubFlowFunctionSchema", "FunctionContext", "ProjectConfiguration", "ConfigurationDefinition", "RuntimeError", "FlowTypeSetting_UniquenessScope", "RuntimeFlowTypeSetting_UniquenessScope", "CodeZeroEvent", "SchemaType", @@ -73,7 +80,9 @@ "OmitFunction", "Design", "ThrowsError", "Parameter", "Schema", "TypeString", "EventSetting", "OmitEvent", "OmitRuntimeFunction", "BaseManager", "ConfigManager", "FlowManager", "FunctionManager", - "RuntimeFunctionManager", "DataTypeManager", "EventManager", "RuntimeEventManager", + "RuntimeFunctionManager", "SubFlowFunctionManager", + "SUB_FLOW_FUNCTION_TTL_SECONDS", "DataTypeManager", + "EventManager", "RuntimeEventManager", "FunctionProps", "FunctionParameterProps", "RuntimeFunctionProps", "RuntimeFunctionRunnable", "DataTypeProps", "DataTypeRunnable", "EventModel", "EventSettingProps", "RuntimeEventProps", "RuntimeEventRunnable", diff --git a/py/hercules/action.py b/py/hercules/action.py index 50bfd71..1c1dc59 100644 --- a/py/hercules/action.py +++ b/py/hercules/action.py @@ -22,6 +22,7 @@ from hercules.manager.function import FunctionManager from hercules.manager.runtime_event import RuntimeEventManager from hercules.manager.runtime_function import RuntimeFunctionManager +from hercules.manager.sub_flow_function import SubFlowFunctionManager from hercules.map.datatype import data_type_map from hercules.map.event import event_map from hercules.map.function import function_map @@ -29,7 +30,12 @@ from hercules.map.runtime_function import runtime_function_map from hercules.models.event import EventModel from hercules.models.function import FunctionProps -from hercules.types import ConfigurationDefinition, RuntimeError, Translation +from hercules.types import ( + ConfigurationDefinition, + RuntimeError, + SubFlowFunctionDefinition, + Translation, +) # Registry of actions constructed while exporting, so the CLI can find them. _REGISTERED_ACTIONS: List["Action"] = [] @@ -114,6 +120,7 @@ def __init__( self.flows = FlowManager() self.functions = FunctionManager() self.runtime_functions = RuntimeFunctionManager() + self.sub_flow_functions = SubFlowFunctionManager() self.data_types = DataTypeManager() self.events = EventManager() self.runtime_events = RuntimeEventManager() @@ -165,6 +172,35 @@ def register_runtime_function(self, klass: type) -> None: ), ) + def register_sub_flow_function(self, definition: SubFlowFunctionDefinition) -> str: + """Register a sub flow function returned from a ``run`` block. + + Returns the identifier Aquila invokes it by. Each returned function gets its + own identifier, so repeatedly returning one (e.g. once per execution) never + overwrites a function a running flow still holds. + + Aquila gives no signal that a flow is done with a function, so a registration + expires on its own 30 minutes after it was registered and can no longer be + invoked (see :class:`~hercules.manager.sub_flow_function.SubFlowFunctionManager`). + Call ``sub_flow_functions.delete(identifier)`` to drop one earlier. + """ + identifier = f"{self._sub_flow_function_prefix}{definition.name}::{uuid.uuid4()}" + self.sub_flow_functions.set(identifier, definition) + return identifier + + def is_sub_flow_function_identifier(self, identifier: str) -> bool: + """True for an identifier this action minted for a sub flow function. + + Holds whether or not it is still registered. An execution request naming one + that is gone addresses a function that expired (or was deleted) rather than an + unknown one, which lets the execution handler report it as such. + """ + return identifier.startswith(self._sub_flow_function_prefix) + + @property + def _sub_flow_function_prefix(self) -> str: + return f"{self._identifier}::" + def register_data_type_class(self, klass: type) -> None: definition = data_type_map(klass) self.data_types.set(definition.identifier, definition) diff --git a/py/hercules/actions/execution.py b/py/hercules/actions/execution.py index 2066a26..e5024fa 100644 --- a/py/hercules/actions/execution.py +++ b/py/hercules/actions/execution.py @@ -4,25 +4,60 @@ import asyncio import inspect import time +from dataclasses import dataclass +from typing import Callable, Optional from hercules._tucana.helpers import construct_value from tucana.generated.aquila import action_pb2 from tucana.generated.shared import errors_pb2, execution_result_pb2 from hercules.events import CodeZeroEvent from hercules.internal.literal import resolve_node_value +from hercules.internal.sub_flow_function import ( + as_sub_flow_function_definition, + build_sub_flow_function_result, +) from hercules.types import ProjectConfiguration, RuntimeError packet_type = "execution" +@dataclass +class _ExecutionTarget: + """A function an execution request may address. + + Either a registered runtime function or a sub flow function returned by one. + ``parameter_count`` is ``None`` for a sub flow function, which has no definition + declaring its parameters, so every parameter of the request is passed through. + """ + + handler: Callable[..., object] + parameter_count: Optional[int] = None + + def _now_micros() -> int: return time.time_ns() // 1000 -def _build_params(action, execution, func): +def _resolve_target(action, identifier: str) -> Optional[_ExecutionTarget]: + func = action.runtime_functions.get(identifier) + if func is not None: + return _ExecutionTarget(func.handler, len(func.parameters or [])) + + sub_flow_function = action.sub_flow_functions.get(identifier) + if sub_flow_function is not None: + return _ExecutionTarget(sub_flow_function.run) + + return None + + +def _build_params(action, execution, target: _ExecutionTarget): + count = ( + target.parameter_count + if target.parameter_count is not None + else len(execution.parameters) + ) params = [] - parameters = func.parameters or [] - for index, _param in enumerate(parameters): + for index in range(count): field = execution.parameters[index] if index < len(execution.parameters) else None params.append(resolve_node_value(action, field)) return params @@ -42,15 +77,48 @@ def _wants_context(handler, param_count: int) -> bool: return len(positional) == param_count + 1 +def _to_node_value(action, value) -> execution_result_pb2.NodeValueExecutionResult: + """Wrap a handler's return value in a ``NodeValueExecutionResult``. + + Raw data goes out as a value; a returned sub flow function is registered on the + action and reported as a ``SubFlowFunctionResult`` so the flow can invoke it later. + """ + definition = as_sub_flow_function_definition(value) + if definition is not None: + return execution_result_pb2.NodeValueExecutionResult( + sub_flow_function_result=build_sub_flow_function_result(action, definition) + ) + return execution_result_pb2.NodeValueExecutionResult( + raw_value=construct_value(value if value is not None else None) + ) + + def handle(action, execution) -> None: action.emit(CodeZeroEvent.execution_request_received, execution) - func = action.runtime_functions.get(execution.function_identifier) - if func is None: + target = _resolve_target(action, execution.function_identifier) + if target is None: + if action.is_sub_flow_function_identifier(execution.function_identifier): + # A sub flow function this action returned earlier, whose maximum + # lifetime has passed (or which was deleted). Fail the node rather than + # staying silent, so the flow does not wait for a result that cannot come. + asyncio.ensure_future( + _send_error( + action, + execution, + RuntimeError( + "SUB_FLOW_FUNCTION_EXPIRED", + f"Sub flow function {execution.function_identifier} is no " + f"longer registered: its maximum lifetime has passed", + ), + _now_micros(), + ) + ) + return print("Received execution request for unknown function:", execution.function_identifier) return - params = _build_params(action, execution, func) + params = _build_params(action, execution, target) conf = action.configs.get(execution.project_id) or ProjectConfiguration( project_id=0, @@ -68,7 +136,7 @@ def handle(action, execution) -> None: ) started_at = _now_micros() - handler = func.handler + handler = target.handler all_params = ( [context, *params] if _wants_context(handler, len(params)) else params ) @@ -89,33 +157,38 @@ async def _run(action, execution, handler, all_params, started_at) -> None: started_at=started_at, finished_at=finished_at, parameter_results=[], - success=construct_value(value if value is not None else None), + success=_to_node_value(action, value), ), ) ) await action.send(request) except Exception as error: # noqa: BLE001 - finished_at = _now_micros() - is_runtime_error = isinstance(error, RuntimeError) - proto_error = errors_pb2.Error( - code=error.code if is_runtime_error else "UNKNOWN_ERROR", - category="RUNTIME", - message=(error.description or str(error)) if is_runtime_error else str(error), - timestamp=finished_at, - version=action.version, - ) - request = action_pb2.ActionTransferRequest( - result=action_pb2.ActionExecutionResponse( - execution_identifier=execution.execution_identifier, - node_result=execution_result_pb2.NodeExecutionResult( - started_at=started_at, - finished_at=finished_at, - parameter_results=[], - error=proto_error, - ), - ) + await _send_error(action, execution, error, started_at) + + +async def _send_error(action, execution, error, started_at) -> None: + """Report a failed execution back to Aquila.""" + finished_at = _now_micros() + is_runtime_error = isinstance(error, RuntimeError) + proto_error = errors_pb2.Error( + code=error.code if is_runtime_error else "UNKNOWN_ERROR", + category="RUNTIME", + message=(error.description or str(error)) if is_runtime_error else str(error), + timestamp=finished_at, + version=action.version, + ) + request = action_pb2.ActionTransferRequest( + result=action_pb2.ActionExecutionResponse( + execution_identifier=execution.execution_identifier, + node_result=execution_result_pb2.NodeExecutionResult( + started_at=started_at, + finished_at=finished_at, + parameter_results=[], + error=proto_error, + ), ) - try: - await action.send(request) - except Exception as err: # noqa: BLE001 - print("Failed to send execution error:", err) + ) + try: + await action.send(request) + except Exception as err: # noqa: BLE001 + print("Failed to send execution error:", err) diff --git a/py/hercules/internal/schema.py b/py/hercules/internal/schema.py index ad4ad1a..7d75989 100644 --- a/py/hercules/internal/schema.py +++ b/py/hercules/internal/schema.py @@ -133,6 +133,16 @@ def pydantic_to_type_string(model: SchemaModel) -> str: return json_schema_to_ts(resolved, name) +def pydantic_to_json_schema(model: SchemaModel, mode: str) -> Dict[str, Any]: + """Convert a Pydantic model into a plain JSON schema dict. + + ``mode`` picks which side of the model is described: ``"validation"`` for the + input it accepts, ``"serialization"`` for the output it produces (the Python + counterpart of the TypeScript SDK's zod ``io`` option). + """ + return model.model_json_schema(mode=mode) + + def pydantic_to_rules(model: SchemaModel) -> List[data_type_pb2.DefinitionDataTypeRule]: schema = model.model_json_schema() rules: List[data_type_pb2.DefinitionDataTypeRule] = [] diff --git a/py/hercules/internal/sub_flow_function.py b/py/hercules/internal/sub_flow_function.py new file mode 100644 index 0000000..91c9bce --- /dev/null +++ b/py/hercules/internal/sub_flow_function.py @@ -0,0 +1,86 @@ +"""Sub flow function support (port of ``src/internal/sub-flow-function.ts``). + +A runtime function's ``run`` may return a sub flow function instead of raw data. +Hercules registers it on the action and reports it to Aquila as a +``shared.SubFlowFunctionResult``, which the flow can then invoke like any sub flow. +""" +from __future__ import annotations + +from typing import Mapping, Optional + +from pydantic import BaseModel + +from hercules._tucana.helpers import construct_value +from hercules.internal.schema import pydantic_to_json_schema +from hercules.types import ( + RuntimeError, + SubFlowFunctionDefinition, + SubFlowFunctionSchema, +) +from tucana.generated.shared import execution_result_pb2, flow_pb2, struct_pb2 + +__all__ = ["as_sub_flow_function_definition", "build_sub_flow_function_result", "to_struct"] + + +def as_sub_flow_function_definition(value) -> Optional[SubFlowFunctionDefinition]: + """Return ``value`` as a definition if ``run`` returned a sub flow function. + + A sub flow function is told apart from raw data by its callable ``run``, which a + plain value never carries. Both a :class:`SubFlowFunctionDefinition` and a plain + mapping with the same keys are accepted. + """ + if isinstance(value, SubFlowFunctionDefinition): + return value + if isinstance(value, Mapping) and callable(value.get("run")): + return SubFlowFunctionDefinition( + name=value.get("name"), + description=value.get("description") or "", + input_schema=value.get("input_schema") or {}, + output_schema=value.get("output_schema") or {}, + run=value["run"], + ) + return None + + +def _is_model(schema) -> bool: + return isinstance(schema, type) and issubclass(schema, BaseModel) + + +def to_struct(schema: Optional[SubFlowFunctionSchema], mode: str) -> struct_pb2.Struct: + """Convert a declared schema into the ``shared.Struct`` the protobuf field expects. + + A Pydantic model becomes a JSON schema first: ``mode`` is ``"validation"`` for a + function's input schema (what it accepts) and ``"serialization"`` for its output + schema (what it produces). + """ + if isinstance(schema, struct_pb2.Struct): + return schema + if _is_model(schema): + return construct_value(pydantic_to_json_schema(schema, mode)).struct_value + return construct_value(schema or {}).struct_value + + +def build_sub_flow_function_result( + action, definition: SubFlowFunctionDefinition +) -> execution_result_pb2.SubFlowFunctionResult: + """Register the sub flow function and build the result reported back to Aquila. + + The identifier in the result is the one Aquila sends back on an execution request + when the flow invokes the function. + """ + if not definition.name: + raise RuntimeError( + "SUB_FLOW_FUNCTION_MISSING_NAME", + "A sub flow function returned from a run block must declare a name", + ) + function_identifier = action.register_sub_flow_function(definition) + return execution_result_pb2.SubFlowFunctionResult( + function=flow_pb2.SubFlowFunction( + function_identifier=function_identifier, + definition_source="action", + ), + name=definition.name, + description=definition.description or "", + input_schema=to_struct(definition.input_schema, "validation"), + output_schema=to_struct(definition.output_schema, "serialization"), + ) diff --git a/py/hercules/manager/sub_flow_function.py b/py/hercules/manager/sub_flow_function.py new file mode 100644 index 0000000..ff1e251 --- /dev/null +++ b/py/hercules/manager/sub_flow_function.py @@ -0,0 +1,93 @@ +"""Sub flow function manager (port of ``src/manager/SubFlowFunctionManager.ts``).""" +from __future__ import annotations + +import asyncio +import time +from typing import Dict, Optional + +from hercules.manager.base import BaseManager + +#: How long a returned sub flow function stays in memory and invocable: 30 minutes. +SUB_FLOW_FUNCTION_TTL_SECONDS = 30 * 60 + + +class SubFlowFunctionManager(BaseManager): + """Sub flow functions a runtime function returned at execution time. + + Keyed by the identifier reported to Aquila. Unlike the other managers this one is + not filled by registering classes up front: entries appear while the action runs + (see :meth:`hercules.action.Action.register_sub_flow_function`). + + Each entry has a maximum lifetime of :data:`SUB_FLOW_FUNCTION_TTL_SECONDS`, + counted from when it was registered and never extended by an invocation, so a flow + that holds a function forever cannot grow the registry without bound. Expired + entries are dropped by a sweep scheduled for the next expiry, and on any read or + write, so an expired function is never invoked even if the sweep could not run + (there is no running event loop to schedule it on, for instance). + """ + + def __init__(self, ttl_seconds: float = SUB_FLOW_FUNCTION_TTL_SECONDS) -> None: + super().__init__() + self._ttl_seconds = ttl_seconds + # Expiry per key. Insertion order is expiry order, since the TTL is fixed. + self._expiries: Dict[str, float] = {} + self._sweep_handle: Optional[asyncio.TimerHandle] = None + + @staticmethod + def _now() -> float: + # Monotonic, so a system clock change cannot postpone or trigger an expiry. + return time.monotonic() + + def set(self, key, value): + # Re-insert the expiry so a replaced key moves to the back and the dict stays + # ordered by expiry. + self._expiries.pop(key, None) + super().set(key, value) + self._expiries[key] = self._now() + self._ttl_seconds + self.prune() + return self + + def get(self, key): + self.prune() + return super().get(key) + + def has(self, key) -> bool: + self.prune() + return super().has(key) + + def delete(self, key) -> bool: + self._expiries.pop(key, None) + deleted = super().delete(key) + self._schedule_sweep(self._now()) + return deleted + + def clear(self) -> None: + self._expiries.clear() + super().clear() + self._schedule_sweep(self._now()) + + def prune(self, now: Optional[float] = None) -> None: + """Drop every entry whose lifetime has passed, then schedule the next sweep.""" + if now is None: + now = self._now() + for key, expires_at in list(self._expiries.items()): + # Ordered by expiry, so the first entry still alive ends the sweep. + if expires_at > now: + break + del self._expiries[key] + super().delete(key) + self._schedule_sweep(now) + + def _schedule_sweep(self, now: float) -> None: + if self._sweep_handle is not None: + self._sweep_handle.cancel() + self._sweep_handle = None + next_expiry = next(iter(self._expiries.values()), None) + if next_expiry is None: + return + try: + loop = asyncio.get_running_loop() + except RuntimeError: + # No loop to schedule on; pruning on the next read or write still applies. + return + self._sweep_handle = loop.call_later(max(0.0, next_expiry - now), self.prune) diff --git a/py/hercules/models/runtime_function.py b/py/hercules/models/runtime_function.py index 7612a94..7881f53 100644 --- a/py/hercules/models/runtime_function.py +++ b/py/hercules/models/runtime_function.py @@ -2,11 +2,16 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Callable, List, Optional +from typing import Callable, List, Optional, Union from hercules._tucana.helpers import PlainValue from hercules.models.function import FunctionParameterProps -from hercules.types import Translation +from hercules.types import SubFlowFunctionDefinition, Translation + +#: What a runtime function's ``run`` may produce: raw data, or a +#: :class:`~hercules.types.SubFlowFunctionDefinition` that Hercules registers and +#: reports to Aquila as a ``shared.SubFlowFunctionResult``. +RuntimeFunctionReturn = Union[PlainValue, SubFlowFunctionDefinition] class RuntimeFunctionRunnable: @@ -20,7 +25,7 @@ def run(self, *args): # pragma: no cover - overridden by subclasses class RuntimeFunctionProps: runtime_name: str signature: str - handler: Callable[..., PlainValue] + handler: Callable[..., RuntimeFunctionReturn] parameters: Optional[List[FunctionParameterProps]] = None throws_error: Optional[bool] = None name: Optional[List[Translation]] = None diff --git a/py/hercules/types.py b/py/hercules/types.py index 1b5017f..6ec2304 100644 --- a/py/hercules/types.py +++ b/py/hercules/types.py @@ -2,7 +2,9 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any, Awaitable, Callable, List, Optional, Protocol, Union +from typing import Any, Awaitable, Callable, List, Optional, Protocol, Type, Union + +from pydantic import BaseModel from hercules._tucana.helpers import PlainValue from tucana.generated.shared.flow_type_pb2 import FlowTypeSetting @@ -20,6 +22,8 @@ "RuntimeFlowTypeSetting_UniquenessScope", "Translation", "SubFlow", + "SubFlowFunctionSchema", + "SubFlowFunctionDefinition", "FunctionContext", "ProjectConfiguration", "ConfigurationDefinition", @@ -48,6 +52,40 @@ class SubFlow(Protocol): def __call__(self, *args: PlainValue) -> Awaitable[PlainValue]: ... +#: Describes a sub flow function's input or output. Normally a Pydantic model, which +#: is converted to a JSON schema while the result is sent (the counterpart of the +#: TypeScript SDK's zod schemas). A plain JSON schema dict or a ready-made ``Struct`` +#: (e.g. one taken from a :class:`SubFlow`) is passed through unchanged. +SubFlowFunctionSchema = Union[Type[BaseModel], Struct, dict] + + +@dataclass +class SubFlowFunctionDefinition: + """A sub flow function returned from a runtime function's ``run`` in place of raw data. + + Hercules registers :attr:`run` internally and reports the function to Aquila as a + ``shared.SubFlowFunctionResult``, so the flow can invoke it later like any other + sub flow. Name, description and I/O schema have to be declared here, because + nothing in the action's module definition describes a function created at runtime. + + May also be returned as a plain dict with the same keys:: + + return SubFlowFunctionDefinition( + name="greet", + description="Greets a name", + input_schema=GreetInput, + output_schema=GreetOutput, + run=lambda name: f"Hello {name}", + ) + """ + + name: str + description: str + input_schema: SubFlowFunctionSchema + output_schema: SubFlowFunctionSchema + run: Callable[..., Any] + + @dataclass class ProjectConfiguration: project_id: int diff --git a/py/pyproject.toml b/py/pyproject.toml index f4159e9..44f49d6 100644 --- a/py/pyproject.toml +++ b/py/pyproject.toml @@ -12,7 +12,7 @@ license = "MIT" dependencies = [ "grpcio>=1.64", "protobuf>=5.0", - "tucana>=0.0.82", + "tucana>=0.0.83", "pydantic>=2", ] diff --git a/py/tests/test_execution.py b/py/tests/test_execution.py new file mode 100644 index 0000000..b8d6783 --- /dev/null +++ b/py/tests/test_execution.py @@ -0,0 +1,309 @@ +"""Port of ``ts/test/execution.test.ts``.""" +import asyncio + +import pytest + +from hercules import Action, SubFlowFunctionDefinition +from hercules.schema import BaseModel, RootModel +from hercules.actions.execution import handle +from hercules.manager.sub_flow_function import ( + SUB_FLOW_FUNCTION_TTL_SECONDS, + SubFlowFunctionManager, +) +from hercules.models.runtime_function import RuntimeFunctionProps +from tucana.generated.aquila import action_pb2 +from tucana.generated.shared import struct_pb2 + + +class GreetInput(BaseModel): + name: str + + +class GreetOutput(RootModel[str]): + pass + + +def _str(string_value): + return struct_pb2.Value(string_value=string_value) + + +def _literal_param(value): + return action_pb2.ActionNodeValue( + literal_value=action_pb2.ActionLiteralValue(value=value, references=[]) + ) + + +class _StubStream: + """Records everything the action would send to Aquila.""" + + def __init__(self): + self.sent = [] + + async def write(self, request): + self.sent.append(request) + + +@pytest.fixture +def action(): + action = Action( + "my_action", "1.0.0", None, "me", "icon", "docs", + [{"code": "en-US", "content": "My Action"}], + ) + action._stream = _StubStream() + return action + + +@pytest.fixture +def sent(action): + return action._stream.sent + + +def _runtime_function(handler, parameters=()): + return RuntimeFunctionProps( + runtime_name="fn", + signature="(): STRING", + handler=handler, + parameters=list(parameters), + ) + + +def _request(function_identifier, parameters=()): + return action_pb2.ActionExecutionRequest( + execution_identifier="exec-1", + function_identifier=function_identifier, + parameters=list(parameters), + project_id=1, + ) + + +def _sub_flow_function(run): + return SubFlowFunctionDefinition( + name="greet", + description="Greets someone", + input_schema=GreetInput, + output_schema=GreetOutput, + run=run, + ) + + +def _run_handle(action, request): + """Dispatch a request and drain the task the handler schedules.""" + + async def main(): + handle(action, request) + # Let the task the handler created run to completion. + await asyncio.sleep(0) + await asyncio.sleep(0) + + asyncio.new_event_loop().run_until_complete(main()) + + +def _node_result(sent): + assert len(sent) == 1 + return sent[0].result.node_result + + +def _success(sent): + node_result = _node_result(sent) + assert node_result.WhichOneof("result") == "success" + return node_result.success + + +def test_reports_raw_data_as_a_raw_value(action, sent): + action.runtime_functions.set("fn", _runtime_function(lambda: "hello")) + + _run_handle(action, _request("fn")) + + success = _success(sent) + assert success.WhichOneof("value") == "raw_value" + assert success.raw_value == _str("hello") + + +def test_registers_a_returned_sub_flow_function_and_reports_it(action, sent): + action.runtime_functions.set( + "fn", _runtime_function(lambda: _sub_flow_function(lambda name: f"Hi {name}")) + ) + + _run_handle(action, _request("fn")) + + success = _success(sent) + assert success.WhichOneof("value") == "sub_flow_function_result" + result = success.sub_flow_function_result + assert result.name == "greet" + assert result.description == "Greets someone" + # The Pydantic models are converted to JSON schemas before they are sent. + assert result.input_schema == struct_pb2.Struct( + fields={ + "title": _str("GreetInput"), + "type": _str("object"), + "properties": struct_pb2.Value( + struct_value=struct_pb2.Struct( + fields={ + "name": struct_pb2.Value( + struct_value=struct_pb2.Struct( + fields={"title": _str("Name"), "type": _str("string")} + ) + ) + } + ) + ), + "required": struct_pb2.Value( + list_value=struct_pb2.ListValue(values=[_str("name")]) + ), + } + ) + assert result.output_schema == struct_pb2.Struct( + fields={"title": _str("GreetOutput"), "type": _str("string")} + ) + assert result.function.definition_source == "action" + # The identifier Aquila invokes the function by is registered on the action. + assert action.sub_flow_functions.has(result.function.function_identifier) + + +def test_accepts_a_sub_flow_function_returned_as_a_dict(action, sent): + action.runtime_functions.set("fn", _runtime_function(lambda: { + "name": "greet", + "description": "Greets someone", + "input_schema": {}, + "output_schema": {}, + "run": lambda: None, + })) + + _run_handle(action, _request("fn")) + + success = _success(sent) + assert success.WhichOneof("value") == "sub_flow_function_result" + assert success.sub_flow_function_result.name == "greet" + + +def test_passes_a_plain_json_schema_dict_through_unchanged(action, sent): + definition = _sub_flow_function(lambda: None) + definition.output_schema = {"type": "string"} + action.runtime_functions.set("fn", _runtime_function(lambda: definition)) + + _run_handle(action, _request("fn")) + + result = _success(sent).sub_flow_function_result + assert result.output_schema == struct_pb2.Struct(fields={"type": _str("string")}) + + +def test_gives_each_returned_sub_flow_function_its_own_identifier(action, sent): + action.runtime_functions.set( + "fn", _runtime_function(lambda: _sub_flow_function(lambda: None)) + ) + + _run_handle(action, _request("fn")) + _run_handle(action, _request("fn")) + + assert len(sent) == 2 + assert action.sub_flow_functions.size == 2 + + +def test_executes_a_registered_sub_flow_function_with_every_request_parameter(action, sent): + identifier = action.register_sub_flow_function( + _sub_flow_function(lambda name: f"Hi {name}") + ) + + _run_handle(action, _request(identifier, [_literal_param(_str("Ada"))])) + + success = _success(sent) + assert success.WhichOneof("value") == "raw_value" + assert success.raw_value == _str("Hi Ada") + + +def test_reports_an_unnamed_sub_flow_function_as_an_error(action, sent): + definition = _sub_flow_function(lambda: None) + definition.name = "" + action.runtime_functions.set("fn", _runtime_function(lambda: definition)) + + _run_handle(action, _request("fn")) + + node_result = _node_result(sent) + assert node_result.WhichOneof("result") == "error" + assert node_result.error.code == "SUB_FLOW_FUNCTION_MISSING_NAME" + assert action.sub_flow_functions.size == 0 + + +class TestSubFlowFunctionLifetime: + """The registry drops a sub flow function once its maximum lifetime has passed.""" + + @staticmethod + def _manager(monkeypatch, clock): + manager = SubFlowFunctionManager() + monkeypatch.setattr(SubFlowFunctionManager, "_now", staticmethod(lambda: clock[0])) + return manager + + def test_default_lifetime_is_30_minutes(self): + assert SUB_FLOW_FUNCTION_TTL_SECONDS == 30 * 60 + + def test_keeps_a_function_invocable_right_up_to_its_lifetime(self, monkeypatch): + clock = [0.0] + manager = self._manager(monkeypatch, clock) + manager.set("a", "fn") + + clock[0] += SUB_FLOW_FUNCTION_TTL_SECONDS - 0.001 + + assert manager.get("a") == "fn" + + def test_drops_a_function_once_its_lifetime_has_passed(self, monkeypatch): + clock = [0.0] + manager = self._manager(monkeypatch, clock) + manager.set("a", "fn") + + clock[0] += SUB_FLOW_FUNCTION_TTL_SECONDS + + assert manager.get("a") is None + assert manager.size == 0 + + def test_does_not_extend_a_lifetime_when_the_function_is_invoked(self, monkeypatch): + clock = [0.0] + manager = self._manager(monkeypatch, clock) + manager.set("a", "fn") + + clock[0] += SUB_FLOW_FUNCTION_TTL_SECONDS / 2 + assert manager.get("a") == "fn" + clock[0] += SUB_FLOW_FUNCTION_TTL_SECONDS / 2 + + assert manager.get("a") is None + + def test_expires_entries_in_registration_order(self, monkeypatch): + clock = [0.0] + manager = self._manager(monkeypatch, clock) + manager.set("first", "fn") + clock[0] += SUB_FLOW_FUNCTION_TTL_SECONDS / 2 + manager.set("second", "fn") + + clock[0] += SUB_FLOW_FUNCTION_TTL_SECONDS / 2 + + assert manager.get("first") is None + assert manager.get("second") == "fn" + + def test_frees_an_expired_function_without_waiting_for_a_read(self): + async def main(): + manager = SubFlowFunctionManager(ttl_seconds=0.05) + manager.set("a", "fn") + assert manager.size == 1 + # The scheduled sweep clears the entry on its own, so nothing is held + # just because no one looked the function up again. + await asyncio.sleep(0.1) + assert len(manager.cache) == 0 + + asyncio.new_event_loop().run_until_complete(main()) + + def test_an_expired_function_fails_the_node_instead_of_running(self, action, sent): + action.sub_flow_functions = SubFlowFunctionManager(ttl_seconds=0.0) + identifier = action.register_sub_flow_function( + _sub_flow_function(lambda name: f"Hi {name}") + ) + + _run_handle(action, _request(identifier, [_literal_param(_str("Ada"))])) + + node_result = _node_result(sent) + assert node_result.WhichOneof("result") == "error" + assert node_result.error.code == "SUB_FLOW_FUNCTION_EXPIRED" + assert identifier in node_result.error.message + + def test_an_unknown_function_sends_nothing(self, action, sent): + _run_handle(action, _request("not_a_function")) + + assert sent == [] diff --git a/py/uv.lock b/py/uv.lock index 24dd373..d40eb72 100644 --- a/py/uv.lock +++ b/py/uv.lock @@ -118,7 +118,7 @@ requires-dist = [ { name = "grpcio", specifier = ">=1.64" }, { name = "protobuf", specifier = ">=5.0" }, { name = "pydantic", specifier = ">=2" }, - { name = "tucana", specifier = ">=0.0.82" }, + { name = "tucana", specifier = ">=0.0.83" }, ] provides-extras = ["codegen"] @@ -768,15 +768,15 @@ wheels = [ [[package]] name = "tucana" -version = "0.0.82" +version = "0.0.83" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "grpcio" }, { name = "protobuf" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b2/7e/b483593c7b0480a73515848819d6a55c9937c7772dde2f22217b682d18eb/tucana-0.0.82.tar.gz", hash = "sha256:1714911505bc9755e4e80f85dcc8fc755f6e03ad742a57c7de384bb6966e7ad2", size = 22076, upload-time = "2026-08-14T23:18:41.235Z" } +sdist = { url = "https://files.pythonhosted.org/packages/2b/73/5f0e92f9acb879ab73a0945ee04a3e285e4c7616e69b098f5a810605a25c/tucana-0.0.83.tar.gz", hash = "sha256:b6806c1cbe0414fec9c9025f5068cf8574b55d060030162d65506b75b71346a9", size = 23142, upload-time = "2026-10-10T12:57:26.269Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/98/84/a3436373627d5d78488b9a3797c40053a69cfdb51cc0490617ba8efe4f04/tucana-0.0.82-py3-none-any.whl", hash = "sha256:4d87dbe8987889fe8071907850a2d37e5c65d2160e9a8d1bc75f9aac1af06480", size = 66322, upload-time = "2026-08-14T23:18:40.214Z" }, + { url = "https://files.pythonhosted.org/packages/16/2c/1bebb8ef0901c409b0fdff8ff8aecdc568dd7d679da8d1f80b5ed0dfbfb6/tucana-0.0.83-py3-none-any.whl", hash = "sha256:76c9f9c568150c9ddf81cace039d7d20991126fb89089e5ee4564d914530492b", size = 72866, upload-time = "2026-10-10T12:57:25.206Z" }, ] [[package]] diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 5115838..ad19914 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -1901,9 +1901,9 @@ checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" [[package]] name = "tucana" -version = "0.0.82" +version = "0.0.83" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f2eb0829308e8695969b580f1ff4c80f0e2c5e6b6b7f49aeb751d458410e93a" +checksum = "b4810f214ea212f895371454f99eb87b1af072b3275f89fa6b4d3aba7780dbef" dependencies = [ "pbjson", "pbjson-build", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 8c7cf8f..867e905 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -13,7 +13,7 @@ license = "ISC" repository = "https://github.com/code0-tech/hercules" [workspace.dependencies] -tucana = { version = "0.0.82", default-features = false, features = ["aquila"] } +tucana = { version = "0.0.83", default-features = false, features = ["aquila"] } tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "time"] } tonic = "0.14" async-trait = "0.1" diff --git a/rust/examples/simple-example-rs/src/functions/fibonacci_runtime_function.rs b/rust/examples/simple-example-rs/src/functions/fibonacci_runtime_function.rs index ffd0f97..8e7143f 100644 --- a/rust/examples/simple-example-rs/src/functions/fibonacci_runtime_function.rs +++ b/rust/examples/simple-example-rs/src/functions/fibonacci_runtime_function.rs @@ -1,5 +1,5 @@ use hercules_sdk::{ - async_trait, Arguments, FunctionContext, PlainValue, Result, RuntimeFunctionHandler, + async_trait, Arguments, FunctionContext, Result, RuntimeFunctionHandler, RuntimeFunctionReturn, }; #[hercules_sdk::runtime_function( @@ -14,7 +14,11 @@ pub struct FibonacciRuntimeFunction; #[async_trait] impl RuntimeFunctionHandler for FibonacciRuntimeFunction { - async fn run(&self, context: &FunctionContext, args: &Arguments) -> Result { + async fn run( + &self, + context: &FunctionContext, + args: &Arguments, + ) -> Result { // Missing or malformed, `args.get` reports it back to Aquila as a // runtime error on its own — nothing to check by hand here. let n: u64 = args.get("test")?; @@ -24,7 +28,7 @@ impl RuntimeFunctionHandler for FibonacciRuntimeFunction { context.project_id, context.execution_id ); - Ok(fib(n).into()) + Ok(RuntimeFunctionReturn::Value(fib(n).into())) } } diff --git a/rust/examples/simple-example-rs/src/functions/greeter_runtime_function.rs b/rust/examples/simple-example-rs/src/functions/greeter_runtime_function.rs new file mode 100644 index 0000000..8fb6b1c --- /dev/null +++ b/rust/examples/simple-example-rs/src/functions/greeter_runtime_function.rs @@ -0,0 +1,64 @@ +use hercules_sdk::{ + async_trait, Arguments, FunctionContext, JsonSchema, PlainValue, Result, + RuntimeFunctionHandler, RuntimeFunctionReturn, SubFlowFunctionDefinition, +}; + +/// The input the returned `greet` function accepts. Declared as its own type +/// because nothing in the module definition describes a function created at +/// runtime; its `schemars` schema is sent to Aquila as the function's input +/// schema. +#[derive(JsonSchema)] +struct GreetInput { + #[allow(dead_code)] + name: String, +} + +/// A runtime function whose handler returns another function instead of raw +/// data. +/// +/// Hercules registers the returned function internally and reports it to Aquila +/// as a sub flow function, so the flow can invoke it later. Name, description +/// and the I/O schema are declared on the returned +/// [`SubFlowFunctionDefinition`] because nothing in the module definition +/// describes a function that is only created at runtime. +#[hercules_sdk::runtime_function( + identifier = "greeter_runtime", + signature = "(greeting: STRING): TRANSFORM", + name(en_US = "Greeter"), + display_message(en_US = "Build a greeter saying ${greeting}"), + description(en_US = "Returns a function that greets a name") +)] +#[parameter( + runtime_name = "greeting", + name(en_US = "Greeting"), + description(en_US = "The greeting the returned function prefixes a name with") +)] +pub struct GreeterRuntimeFunction; + +#[async_trait] +impl RuntimeFunctionHandler for GreeterRuntimeFunction { + async fn run( + &self, + _context: &FunctionContext, + args: &Arguments, + ) -> Result { + let greeting: String = args.get("greeting")?; + let description = format!("Greets a name with \"{greeting}\""); + + let definition = SubFlowFunctionDefinition::new( + "greet", + move |_context: &FunctionContext, parameters: Vec| -> Result { + let name = parameters + .first() + .and_then(PlainValue::as_str) + .unwrap_or_default(); + Ok(format!("{greeting} {name}").into()) + }, + ) + .description(description) + .input_schema::() + .output_schema::(); + + Ok(definition.into()) + } +} diff --git a/rust/examples/simple-example-rs/src/functions/mod.rs b/rust/examples/simple-example-rs/src/functions/mod.rs index e1f67dd..2e2ef20 100644 --- a/rust/examples/simple-example-rs/src/functions/mod.rs +++ b/rust/examples/simple-example-rs/src/functions/mod.rs @@ -1,2 +1,3 @@ pub mod fibonacci_function; pub mod fibonacci_runtime_function; +pub mod greeter_runtime_function; diff --git a/rust/hercules/src/action.rs b/rust/hercules/src/action.rs index db9f2af..9f49c4e 100644 --- a/rust/hercules/src/action.rs +++ b/rust/hercules/src/action.rs @@ -281,6 +281,8 @@ impl Action { request_tx: connection.request_tx, queue_saturation_count: Default::default(), events_tx: self.events_tx, + sub_flow_functions: Default::default(), + sub_flow_function_ttl: crate::SUB_FLOW_FUNCTION_TTL, }); let _ = inner.events_tx.send(HerculesEvent::Connected); diff --git a/rust/hercules/src/connected.rs b/rust/hercules/src/connected.rs index aca2cd3..35f9508 100644 --- a/rust/hercules/src/connected.rs +++ b/rust/hercules/src/connected.rs @@ -4,6 +4,7 @@ use std::sync::{Arc, Mutex, RwLock}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tokio::sync::{broadcast, mpsc, oneshot}; +use tokio::time::Instant; use tokio_stream::{Stream, StreamExt}; use tonic::Streaming; use tucana::aquila::{ @@ -13,14 +14,20 @@ use tucana::aquila::{ action_flow_execution_response, action_flow_update, action_node_value, action_sub_flow_execution_response, action_transfer_request, action_transfer_response, }; -use tucana::shared::{Error as WireError, NodeExecutionResult, node_execution_result}; +use tucana::shared::{ + Error as WireError, NodeExecutionResult, NodeValueExecutionResult, node_execution_result, + node_value_execution_result, +}; use crate::arguments::Arguments; use crate::error::{HerculesError, Result}; use crate::events::{HerculesEvent, event_stream}; -use crate::function::RuntimeFunctionHandler; +use crate::function::{RuntimeFunctionHandler, RuntimeFunctionReturn}; use crate::literal::{ResolvedValue, resolve_literal}; use crate::meta::{ParameterMeta, RuntimeFunctionMeta}; +use crate::sub_flow_function::{ + SubFlowFunctionDefinition, SubFlowFunctionState, build_sub_flow_function_result, +}; use crate::sync; use crate::types::{FunctionContext, PlainValue, ProjectConfiguration}; use crate::value::{construct_value, to_allowed_value}; @@ -78,6 +85,14 @@ pub(crate) struct ConnectedInner { pub request_tx: mpsc::Sender, pub queue_saturation_count: AtomicU64, pub events_tx: broadcast::Sender, + /// Sub flow functions this action's runtime functions returned at + /// execution time, keyed by the identifier reported to Aquila. Filled + /// while the action runs, not up front — see + /// [`ConnectedInner::register_sub_flow_function`]. + pub(crate) sub_flow_functions: Mutex, + /// How long a registered sub flow function stays invocable. Configurable + /// only so tests can shorten it; production uses [`SUB_FLOW_FUNCTION_TTL`]. + pub(crate) sub_flow_function_ttl: Duration, } impl ConnectedInner { @@ -96,6 +111,101 @@ impl ConnectedInner { fn next_pending_token(&self) -> u64 { self.next_pending_token.fetch_add(1, Ordering::Relaxed) } + + /// Registers a sub flow function returned from a runtime function's handler + /// and returns the identifier Aquila invokes it by. Each returned function + /// gets its own identifier, so repeatedly returning one (e.g. once per + /// execution) never overwrites a function a running flow still holds. + /// + /// Aquila gives no signal that a flow is done with a function, so a + /// registration expires on its own after [`SUB_FLOW_FUNCTION_TTL`] and can + /// no longer be invoked. + pub(crate) fn register_sub_flow_function( + self: &Arc, + definition: SubFlowFunctionDefinition, + ) -> String { + let identifier = format!( + "{}::{}::{}", + self.identifier, + definition.name, + uuid::Uuid::new_v4() + ); + let now = Instant::now(); + { + let mut state = sync::lock(&self.sub_flow_functions); + state.insert( + identifier.clone(), + definition, + now + self.sub_flow_function_ttl, + ); + state.prune(now); + } + self.schedule_sub_flow_sweep(); + identifier + } + + /// Looks up a registered sub flow function, pruning any that have expired + /// first so an expired one is never returned even if its sweep could not + /// run. + fn get_sub_flow_function(&self, identifier: &str) -> Option { + let mut state = sync::lock(&self.sub_flow_functions); + state.prune(Instant::now()); + state.get(identifier).cloned() + } + + /// True for an identifier this action minted for a sub flow function, + /// whether or not it is still registered. An execution request naming one + /// that is gone addresses a function that expired rather than an unknown + /// one, which lets the execution handler report it as such. + fn is_sub_flow_function_identifier(&self, identifier: &str) -> bool { + identifier.starts_with(&format!("{}::", self.identifier)) + } + + /// Ensures a single background task is dropping expired sub flow functions + /// at each expiry, so memory is freed without waiting for a read. Holds + /// only a [`Weak`] reference, so it stops on its own once the connection is + /// gone. A no-op when no Tokio runtime is available (pruning on the next + /// read or write still applies). + fn schedule_sub_flow_sweep(self: &Arc) { + { + let mut state = sync::lock(&self.sub_flow_functions); + if state.sweeper_running() || state.next_expiry().is_none() { + return; + } + state.set_sweeper_running(true); + } + let weak = Arc::downgrade(self); + if tokio::runtime::Handle::try_current().is_err() { + // No runtime to spawn on; undo the flag so a later registration + // under a runtime can start the sweep. + if let Some(inner) = weak.upgrade() { + sync::lock(&inner.sub_flow_functions).set_sweeper_running(false); + } + return; + } + tokio::spawn(async move { + loop { + let Some(inner) = weak.upgrade() else { return }; + let next = { + let state = sync::lock(&inner.sub_flow_functions); + state.next_expiry() + }; + let Some(next) = next else { + sync::lock(&inner.sub_flow_functions).set_sweeper_running(false); + return; + }; + drop(inner); + tokio::time::sleep_until(next).await; + let Some(inner) = weak.upgrade() else { return }; + let mut state = sync::lock(&inner.sub_flow_functions); + state.prune(Instant::now()); + if state.next_expiry().is_none() { + state.set_sweeper_running(false); + return; + } + } + }); + } } /// Removes an unresolved pending entry when its caller's future is dropped. @@ -557,8 +667,49 @@ fn now_micros() -> i64 { .unwrap_or_default() } +/// A function an execution request may address: a registered runtime +/// function, or a sub flow function one returned earlier. +enum ExecutionTarget { + /// A registered runtime function, together with its declared parameters — + /// needed to recover each positional wire parameter's name. + RuntimeFunction(Arc, Vec), + /// A sub flow function returned by a runtime function. It has no parameter + /// declaration, so every parameter of the request is passed through. + SubFlowFunction(SubFlowFunctionDefinition), +} + async fn handle_execution(inner: Arc, execution: ActionExecutionRequest) { - let Some(entry) = inner.runtime_functions.get(&execution.function_identifier) else { + let started_at = now_micros(); + + let target = if let Some(entry) = inner.runtime_functions.get(&execution.function_identifier) { + ExecutionTarget::RuntimeFunction(entry.handler.clone(), entry.meta.parameters.clone()) + } else if let Some(definition) = inner.get_sub_flow_function(&execution.function_identifier) { + ExecutionTarget::SubFlowFunction(definition) + } else if inner.is_sub_flow_function_identifier(&execution.function_identifier) { + // A sub flow function this action returned earlier, whose maximum + // lifetime has passed (or which was deleted). Fail the node rather than + // staying silent, so the flow does not wait for a result that cannot + // come. + let finished_at = now_micros(); + let result = node_execution_result::Result::Error(wire_error( + "SUB_FLOW_FUNCTION_EXPIRED".into(), + format!( + "Sub flow function {} is no longer registered: its maximum lifetime has passed", + execution.function_identifier + ), + finished_at, + &inner.version, + )); + send_execution_result( + &inner, + &execution.execution_identifier, + started_at, + finished_at, + result, + ) + .await; + return; + } else { log::error!( "execution request for unknown function {:?}", execution.function_identifier @@ -573,48 +724,120 @@ async fn handle_execution(inner: Arc, execution: ActionExecution project_id: execution.project_id, config_values: HashMap::new(), }); - let execution_identifier = execution.execution_identifier.clone(); - let started_at = now_micros(); - let outcome = match resolve_parameters(&entry.meta.parameters, execution.parameters) { - Ok(ResolvedParameters { - literals, - sub_flow_parameters, - }) => { + + let outcome = run_target(&inner, target, &execution, matched_config).await; + let finished_at = now_micros(); + + let result = match outcome.and_then(|value| to_node_value(&inner, value)) { + Ok(node_value) => node_execution_result::Result::Success(node_value), + Err(err) => error_result(err, finished_at, &inner.version), + }; + + send_execution_result( + &inner, + &execution.execution_identifier, + started_at, + finished_at, + result, + ) + .await; +} + +/// Runs the resolved target against the request's parameters. A runtime +/// function receives its parameters as name-keyed [`Arguments`]; a sub flow +/// function receives them positionally. +async fn run_target( + inner: &Arc, + target: ExecutionTarget, + execution: &ActionExecutionRequest, + matched_config: ProjectConfiguration, +) -> Result { + match target { + ExecutionTarget::RuntimeFunction(handler, parameters) => { + let ResolvedParameters { + literals, + sub_flow_parameters, + } = resolve_parameters(¶meters, execution.parameters.clone())?; let context = FunctionContext { project_id: execution.project_id, - execution_id: execution_identifier.clone(), + execution_id: execution.execution_identifier.clone(), matched_config, connected: Connected::new(inner.clone()), sub_flow_parameters, }; - entry.handler.run(&context, &Arguments::new(literals)).await + handler.run(&context, &Arguments::new(literals)).await } - Err(err) => Err(err), - }; - let finished_at = now_micros(); + ExecutionTarget::SubFlowFunction(definition) => { + let parameters = resolve_positional(execution.parameters.clone())?; + let context = FunctionContext { + project_id: execution.project_id, + execution_id: execution.execution_identifier.clone(), + matched_config, + connected: Connected::new(inner.clone()), + sub_flow_parameters: HashMap::new(), + }; + definition + .handler + .run(&context, parameters) + .await + .map(RuntimeFunctionReturn::Value) + } + } +} - let result = match outcome { - Ok(value) => node_execution_result::Result::Success(construct_value(&value)), - Err(HerculesError::Runtime { code, description }) => { - node_execution_result::Result::Error(wire_error( - code, - description.unwrap_or_default(), - finished_at, - &inner.version, - )) +/// Wraps a handler's return value in a `NodeValueExecutionResult`: raw data +/// goes out as a raw value; a returned sub flow function is registered on the +/// connection and reported as a `SubFlowFunctionResult` so the flow can invoke +/// it later. +fn to_node_value( + inner: &Arc, + value: RuntimeFunctionReturn, +) -> Result { + let value = match value { + RuntimeFunctionReturn::Value(value) => { + node_value_execution_result::Value::RawValue(construct_value(&value)) } - Err(other) => node_execution_result::Result::Error(wire_error( + RuntimeFunctionReturn::SubFlowFunction(definition) => { + node_value_execution_result::Value::SubFlowFunctionResult( + build_sub_flow_function_result(inner, definition)?, + ) + } + }; + Ok(NodeValueExecutionResult { value: Some(value) }) +} + +/// Renders a handler failure as the error variant of a node result. +fn error_result( + err: HerculesError, + finished_at: i64, + version: &str, +) -> node_execution_result::Result { + match err { + HerculesError::Runtime { code, description } => node_execution_result::Result::Error( + wire_error(code, description.unwrap_or_default(), finished_at, version), + ), + other => node_execution_result::Result::Error(wire_error( "UNKNOWN_ERROR".into(), other.to_string(), finished_at, - &inner.version, + version, )), - }; + } +} +/// Sends a single node execution result back to Aquila, logging a send +/// failure rather than propagating it (nothing upstream can act on it). +async fn send_execution_result( + inner: &Arc, + execution_identifier: &str, + started_at: i64, + finished_at: i64, + result: node_execution_result::Result, +) { let request = ActionTransferRequest { data: Some(action_transfer_request::Data::Result( ActionExecutionResponse { - execution_identifier: execution.execution_identifier, + execution_identifier: execution_identifier.to_string(), node_result: Some(NodeExecutionResult { started_at, finished_at, @@ -625,7 +848,7 @@ async fn handle_execution(inner: Arc, execution: ActionExecution }, )), }; - if let Err(err) = send_mandatory(&inner, request).await { + if let Err(err) = send_mandatory(inner, request).await { log::error!( "failed to send execution result for {:?}: {}", execution_identifier, @@ -634,6 +857,23 @@ async fn handle_execution(inner: Arc, execution: ActionExecution } } +/// Resolves each positional wire parameter of a sub flow function invocation +/// to a plain value, in order. A sub flow function has no parameter +/// declaration, so every parameter of the request is passed through; one that +/// resolved to a sub flow reference, or was unset, becomes `null`. +fn resolve_positional(parameters: Vec) -> Result> { + parameters + .into_iter() + .map(|param| match param.value { + Some(action_node_value::Value::LiteralValue(value)) => match resolve_literal(value)? { + ResolvedValue::Literal(value) => Ok(value), + ResolvedValue::SubFlow(_) => Ok(PlainValue::Null), + }, + Some(action_node_value::Value::SubFlow(_)) | None => Ok(PlainValue::Null), + }) + .collect() +} + /// The result of splitting a function execution's positional wire /// parameters by kind. See [`resolve_parameters`]. struct ResolvedParameters { @@ -810,8 +1050,12 @@ mod tests { #[async_trait::async_trait] impl RuntimeFunctionHandler for FixedResultHandler { - async fn run(&self, _context: &FunctionContext, _args: &Arguments) -> Result { - Ok(serde_json::json!("completed")) + async fn run( + &self, + _context: &FunctionContext, + _args: &Arguments, + ) -> Result { + Ok(serde_json::json!("completed").into()) } } @@ -846,6 +1090,8 @@ mod tests { request_tx, queue_saturation_count: Default::default(), events_tx, + sub_flow_functions: Default::default(), + sub_flow_function_ttl: crate::sub_flow_function::SUB_FLOW_FUNCTION_TTL, }); (inner, request_rx) } @@ -1327,3 +1573,391 @@ mod tests { assert_eq!(connected.metrics().pending_flow_executions, 0); } } + +/// Sub flow functions: a runtime function returning a function in place of raw +/// data, and the lifetime of the functions it leaves registered (port of +/// `ts/test/execution.test.ts`). +#[cfg(test)] +mod sub_flow_function_tests { + use std::time::Duration; + + use schemars::JsonSchema; + use tucana::aquila::{ + ActionExecutionRequest, ActionLiteralValue, ActionNodeValue, action_node_value, + }; + use tucana::shared::{ + Error as WireError, Value, node_execution_result, node_value_execution_result, + }; + + use super::*; + use crate::sub_flow_function::SUB_FLOW_FUNCTION_TTL; + + /// The input of the sample `greet` sub flow function. `name` is only read by + /// the generated schema, not by Rust code. + #[derive(JsonSchema)] + #[allow(dead_code)] + struct GreetInput { + name: String, + } + + /// The sample sub flow function the tests return and invoke, parameterized + /// by its behavior. + fn greet_definition(run: F) -> SubFlowFunctionDefinition + where + F: Fn(&FunctionContext, Vec) -> Result + Send + Sync + 'static, + { + SubFlowFunctionDefinition::new("greet", run) + .description("Greets someone") + .input_schema::() + .output_schema::() + } + + /// A runtime function whose handler returns whatever `factory` builds. + struct Returns(Box SubFlowFunctionDefinition + Send + Sync>); + + #[async_trait::async_trait] + impl RuntimeFunctionHandler for Returns { + async fn run( + &self, + _context: &FunctionContext, + _args: &Arguments, + ) -> Result { + Ok((self.0)().into()) + } + } + + fn returns(factory: impl Fn() -> SubFlowFunctionDefinition + Send + Sync + 'static) -> Returns { + Returns(Box::new(factory)) + } + + /// A runtime function returning raw data. + struct ReturnsValue(PlainValue); + + #[async_trait::async_trait] + impl RuntimeFunctionHandler for ReturnsValue { + async fn run( + &self, + _context: &FunctionContext, + _args: &Arguments, + ) -> Result { + Ok(self.0.clone().into()) + } + } + + fn entry(handler: impl RuntimeFunctionHandler) -> RuntimeFunctionEntry { + RuntimeFunctionEntry { + meta: RuntimeFunctionMeta { + identifier: "fn".into(), + ..Default::default() + }, + handler: Arc::new(handler), + } + } + + fn test_inner_with( + runtime_functions: HashMap, + ) -> (Arc, mpsc::Receiver) { + let (request_tx, request_rx) = mpsc::channel(16); + let (events_tx, _events_rx) = broadcast::channel(16); + let inner = Arc::new(ConnectedInner { + identifier: "test-action".into(), + version: "0.0.0".into(), + runtime_functions, + configs: Default::default(), + flows: Default::default(), + pending_flow_executions: Default::default(), + pending_sub_flow_executions: Default::default(), + next_pending_token: Default::default(), + request_tx, + queue_saturation_count: Default::default(), + events_tx, + sub_flow_functions: Default::default(), + sub_flow_function_ttl: SUB_FLOW_FUNCTION_TTL, + }); + (inner, request_rx) + } + + fn test_inner() -> (Arc, mpsc::Receiver) { + test_inner_with(HashMap::new()) + } + + fn with_function( + handler: impl RuntimeFunctionHandler, + ) -> (Arc, mpsc::Receiver) { + test_inner_with(HashMap::from([("fn".to_string(), entry(handler))])) + } + + fn str_value(s: &str) -> Value { + construct_value(&serde_json::json!(s)) + } + + fn literal_param(value: PlainValue) -> ActionNodeValue { + ActionNodeValue { + value: Some(action_node_value::Value::LiteralValue(ActionLiteralValue { + value: Some(construct_value(&value)), + references: vec![], + })), + } + } + + fn request( + function_identifier: &str, + parameters: Vec, + ) -> ActionExecutionRequest { + ActionExecutionRequest { + execution_identifier: "exec-1".into(), + function_identifier: function_identifier.into(), + parameters, + project_id: 1, + } + } + + /// The single node result the action sent, as its success value. + async fn sent_success( + rx: &mut mpsc::Receiver, + ) -> node_value_execution_result::Value { + let request = rx.recv().await.expect("a result was sent"); + let Some(action_transfer_request::Data::Result(response)) = request.data else { + panic!("expected a result"); + }; + match response.node_result.unwrap().result.unwrap() { + node_execution_result::Result::Success(value) => value.value.unwrap(), + other => panic!("expected a success, got {other:?}"), + } + } + + /// The single node result the action sent, as its error. + async fn sent_error(rx: &mut mpsc::Receiver) -> WireError { + let request = rx.recv().await.expect("a result was sent"); + let Some(action_transfer_request::Data::Result(response)) = request.data else { + panic!("expected a result"); + }; + match response.node_result.unwrap().result.unwrap() { + node_execution_result::Result::Error(error) => error, + other => panic!("expected an error, got {other:?}"), + } + } + + #[tokio::test] + async fn reports_raw_data_as_a_raw_value() { + let (inner, mut rx) = with_function(ReturnsValue(serde_json::json!("hello"))); + + handle_execution(inner, request("fn", vec![])).await; + + assert_eq!( + sent_success(&mut rx).await, + node_value_execution_result::Value::RawValue(str_value("hello")) + ); + } + + #[tokio::test] + async fn registers_a_returned_sub_flow_function_and_reports_it() { + let (inner, mut rx) = with_function(returns(|| { + greet_definition(|_ctx, params| { + let name = params + .first() + .and_then(PlainValue::as_str) + .unwrap_or_default(); + Ok(format!("Hi {name}").into()) + }) + })); + + handle_execution(inner.clone(), request("fn", vec![])).await; + + let node_value_execution_result::Value::SubFlowFunctionResult(result) = + sent_success(&mut rx).await + else { + panic!("expected a sub flow function result"); + }; + + assert_eq!(result.name, "greet"); + assert_eq!(result.description, "Greets someone"); + let function = result.function.expect("a function reference"); + assert_eq!(function.definition_source.as_deref(), Some("action")); + // The input schema describes an object with a `name` property. + let input = result.input_schema.expect("an input schema"); + assert_eq!(input.fields.get("type"), Some(&str_value("object"))); + assert!(input.fields.contains_key("properties")); + // The output schema describes a string. + let output = result.output_schema.expect("an output schema"); + assert_eq!(output.fields.get("type"), Some(&str_value("string"))); + // The identifier Aquila invokes the function by is registered. + assert!( + inner + .get_sub_flow_function(&function.function_identifier) + .is_some() + ); + } + + #[tokio::test] + async fn gives_each_returned_sub_flow_function_its_own_identifier() { + let (inner, mut rx) = with_function(returns(|| { + greet_definition(|_ctx, _params| Ok(PlainValue::Null)) + })); + + handle_execution(inner.clone(), request("fn", vec![])).await; + handle_execution(inner.clone(), request("fn", vec![])).await; + + // Two distinct results, two distinct registrations. + drop(sent_success(&mut rx).await); + drop(sent_success(&mut rx).await); + assert_eq!(sync::lock(&inner.sub_flow_functions).len(), 2); + } + + #[tokio::test] + async fn executes_a_registered_sub_flow_function_with_every_parameter() { + let (inner, mut rx) = with_function(ReturnsValue(PlainValue::Null)); + let identifier = inner.register_sub_flow_function(greet_definition(|_ctx, params| { + let name = params + .first() + .and_then(PlainValue::as_str) + .unwrap_or_default(); + Ok(format!("Hi {name}").into()) + })); + + handle_execution( + inner, + request(&identifier, vec![literal_param(serde_json::json!("Ada"))]), + ) + .await; + + assert_eq!( + sent_success(&mut rx).await, + node_value_execution_result::Value::RawValue(str_value("Hi Ada")) + ); + } + + #[tokio::test] + async fn passes_a_plain_json_schema_object_through_unchanged() { + let (inner, mut rx) = with_function(returns(|| { + SubFlowFunctionDefinition::new( + "greet", + |_ctx: &FunctionContext, _params: Vec| -> Result { + Ok(PlainValue::Null) + }, + ) + .output(serde_json::json!({"type": "string"})) + })); + + handle_execution(inner, request("fn", vec![])).await; + + let node_value_execution_result::Value::SubFlowFunctionResult(result) = + sent_success(&mut rx).await + else { + panic!("expected a sub flow function result"); + }; + let output = result.output_schema.expect("an output schema"); + assert_eq!(output.fields.get("type"), Some(&str_value("string"))); + assert_eq!(output.fields.len(), 1); + } + + #[tokio::test] + async fn reports_an_unnamed_sub_flow_function_as_an_error() { + let (inner, mut rx) = with_function(returns(|| { + SubFlowFunctionDefinition::new( + "", + |_ctx: &FunctionContext, _params: Vec| -> Result { + Ok(PlainValue::Null) + }, + ) + })); + + handle_execution(inner.clone(), request("fn", vec![])).await; + + assert_eq!( + sent_error(&mut rx).await.code, + "SUB_FLOW_FUNCTION_MISSING_NAME" + ); + assert_eq!(sync::lock(&inner.sub_flow_functions).len(), 0); + } + + #[tokio::test] + async fn sends_nothing_for_a_function_it_never_minted() { + let (inner, mut rx) = test_inner(); + + handle_execution(inner, request("not_a_function", vec![])).await; + + assert!(rx.try_recv().is_err()); + } + + #[tokio::test(start_paused = true)] + async fn fails_the_node_instead_of_running_an_expired_function() { + let (inner, mut rx) = test_inner(); + let identifier = inner + .register_sub_flow_function(greet_definition(|_ctx, _params| Ok(PlainValue::Null))); + + tokio::time::advance(SUB_FLOW_FUNCTION_TTL).await; + handle_execution( + inner, + request(&identifier, vec![literal_param(serde_json::json!("Ada"))]), + ) + .await; + + let error = sent_error(&mut rx).await; + assert_eq!(error.code, "SUB_FLOW_FUNCTION_EXPIRED"); + assert!(error.message.contains(&identifier)); + } + + #[tokio::test(start_paused = true)] + async fn keeps_a_function_invocable_right_up_to_its_lifetime() { + let (inner, _rx) = test_inner(); + let identifier = inner + .register_sub_flow_function(greet_definition(|_ctx, _params| Ok(PlainValue::Null))); + + tokio::time::advance(SUB_FLOW_FUNCTION_TTL - Duration::from_millis(1)).await; + + assert!(inner.get_sub_flow_function(&identifier).is_some()); + } + + #[tokio::test(start_paused = true)] + async fn drops_a_function_once_its_lifetime_has_passed() { + let (inner, _rx) = test_inner(); + let identifier = inner + .register_sub_flow_function(greet_definition(|_ctx, _params| Ok(PlainValue::Null))); + + tokio::time::advance(SUB_FLOW_FUNCTION_TTL).await; + + assert!(inner.get_sub_flow_function(&identifier).is_none()); + } + + #[tokio::test(start_paused = true)] + async fn frees_an_expired_function_without_waiting_for_a_read() { + let (inner, _rx) = test_inner(); + inner.register_sub_flow_function(greet_definition(|_ctx, _params| Ok(PlainValue::Null))); + assert_eq!(sync::lock(&inner.sub_flow_functions).len(), 1); + + // The scheduled sweep clears the entry on its own — len() does not prune. + tokio::time::advance(SUB_FLOW_FUNCTION_TTL).await; + tokio::task::yield_now().await; + + assert_eq!(sync::lock(&inner.sub_flow_functions).len(), 0); + } + + #[tokio::test(start_paused = true)] + async fn does_not_extend_a_lifetime_when_the_function_is_invoked() { + let (inner, _rx) = test_inner(); + let identifier = inner + .register_sub_flow_function(greet_definition(|_ctx, _params| Ok(PlainValue::Null))); + + tokio::time::advance(SUB_FLOW_FUNCTION_TTL / 2).await; + assert!(inner.get_sub_flow_function(&identifier).is_some()); + tokio::time::advance(SUB_FLOW_FUNCTION_TTL / 2).await; + + assert!(inner.get_sub_flow_function(&identifier).is_none()); + } + + #[tokio::test(start_paused = true)] + async fn expires_entries_in_registration_order() { + let (inner, _rx) = test_inner(); + let first = inner + .register_sub_flow_function(greet_definition(|_ctx, _params| Ok(PlainValue::Null))); + tokio::time::advance(SUB_FLOW_FUNCTION_TTL / 2).await; + let second = inner + .register_sub_flow_function(greet_definition(|_ctx, _params| Ok(PlainValue::Null))); + + tokio::time::advance(SUB_FLOW_FUNCTION_TTL / 2).await; + + assert!(inner.get_sub_flow_function(&first).is_none()); + assert!(inner.get_sub_flow_function(&second).is_some()); + } +} diff --git a/rust/hercules/src/function.rs b/rust/hercules/src/function.rs index 42cd0b7..de1ad97 100644 --- a/rust/hercules/src/function.rs +++ b/rust/hercules/src/function.rs @@ -3,8 +3,34 @@ use async_trait::async_trait; use crate::arguments::Arguments; use crate::error::Result; use crate::meta::{FunctionMeta, RuntimeFunctionMeta}; +use crate::sub_flow_function::SubFlowFunctionDefinition; use crate::types::{FunctionContext, PlainValue}; +/// What a runtime function's handler may produce: raw data +/// ([`Value`](RuntimeFunctionReturn::Value)), or a +/// [`SubFlowFunctionDefinition`] that Hercules registers and reports to Aquila +/// as a `shared.SubFlowFunctionResult` for the flow to invoke later. +/// +/// A handler returning raw data converts into `Value` with `.into()` +/// (`Ok(json!(...).into())`); one returning a function returns the definition +/// (`Ok(definition.into())`). +pub enum RuntimeFunctionReturn { + Value(PlainValue), + SubFlowFunction(SubFlowFunctionDefinition), +} + +impl From for RuntimeFunctionReturn { + fn from(value: PlainValue) -> Self { + Self::Value(value) + } +} + +impl From for RuntimeFunctionReturn { + fn from(definition: SubFlowFunctionDefinition) -> Self { + Self::SubFlowFunction(definition) + } +} + /// The actual behavior of a runtime function. Split out from /// [`RuntimeFunction`] so it can be boxed as `Arc` /// for dynamic dispatch when an execution request comes in off the wire — @@ -12,7 +38,16 @@ use crate::types::{FunctionContext, PlainValue}; /// associated function with no `self`. #[async_trait] pub trait RuntimeFunctionHandler: Send + Sync + 'static { - async fn run(&self, context: &FunctionContext, args: &Arguments) -> Result; + /// Runs the function. Return raw data as a + /// [`Value`](RuntimeFunctionReturn::Value), or a + /// [`SubFlowFunctionDefinition`] to hand the flow a function it can invoke + /// later — see [`RuntimeFunctionReturn`]. + #[allow(clippy::double_must_use)] // `async_trait` adds `#[must_use]` to its boxed future. + async fn run( + &self, + context: &FunctionContext, + args: &Arguments, + ) -> Result; } /// A function Aquila can invoke. Implemented by attaching diff --git a/rust/hercules/src/lib.rs b/rust/hercules/src/lib.rs index 12dc504..88187a5 100644 --- a/rust/hercules/src/lib.rs +++ b/rust/hercules/src/lib.rs @@ -12,9 +12,9 @@ //! //! #[async_trait] //! impl RuntimeFunctionHandler for Greet { -//! async fn run(&self, _ctx: &FunctionContext, args: &Arguments) -> hercules_sdk::Result { +//! async fn run(&self, _ctx: &FunctionContext, args: &Arguments) -> hercules_sdk::Result { //! let name: String = args.get("name")?; -//! Ok(format!("Hello, {name}!").into()) +//! Ok(hercules_sdk::PlainValue::from(format!("Hello, {name}!")).into()) //! } //! } //! @@ -41,6 +41,7 @@ mod literal; mod meta; mod module_builder; mod registration; +mod sub_flow_function; mod sync; mod types; mod value; @@ -52,13 +53,16 @@ pub use data_type::{DataType, DataTypeDef, DataTypeRule}; pub use error::{HerculesError, Result}; pub use event::{Event, RuntimeEvent}; pub use events::HerculesEvent; -pub use function::{Function, RuntimeFunction, RuntimeFunctionHandler}; +pub use function::{Function, RuntimeFunction, RuntimeFunctionHandler, RuntimeFunctionReturn}; pub use meta::DataTypeMeta; pub use meta::{ EventDef, EventMeta, EventSettingMeta, FunctionDef, FunctionMeta, ParameterMeta, RuntimeEventMeta, RuntimeFunctionMeta, }; pub use registration::Registration; +pub use sub_flow_function::{ + SUB_FLOW_FUNCTION_TTL, SubFlowFunctionDefinition, SubFlowFunctionHandler, SubFlowFunctionSchema, +}; pub use types::{ ConfigurationDefinition, FunctionContext, PlainValue, ProjectConfiguration, ScalingOption, Translation, UniquenessScope, diff --git a/rust/hercules/src/sub_flow_function.rs b/rust/hercules/src/sub_flow_function.rs new file mode 100644 index 0000000..6607016 --- /dev/null +++ b/rust/hercules/src/sub_flow_function.rs @@ -0,0 +1,273 @@ +//! Sub flow function support (port of `ts/src/internal/sub-flow-function.ts` +//! and `ts/src/manager/SubFlowFunctionManager.ts`). +//! +//! A runtime function's handler may return a [`SubFlowFunctionDefinition`] +//! instead of raw data. Hercules registers it on the connection and reports it +//! to Aquila as a `shared.SubFlowFunctionResult`, so the flow can invoke it +//! later like any other sub flow. Because nothing in the module definition +//! describes a function created at runtime, the returned definition has to +//! carry its own name, description and I/O schema. + +use std::collections::{HashMap, VecDeque}; +use std::sync::Arc; + +use async_trait::async_trait; +use schemars::{JsonSchema, SchemaGenerator, generate::SchemaSettings}; +use tokio::time::{Duration, Instant}; +use tucana::shared::value::Kind; +use tucana::shared::{Struct, SubFlowFunction as WireSubFlowFunction, SubFlowFunctionResult}; + +use crate::connected::ConnectedInner; +use crate::error::{HerculesError, Result}; +use crate::types::{FunctionContext, PlainValue}; +use crate::value::construct_value; + +/// How long a returned sub flow function stays in memory and invocable: 30 minutes. +pub const SUB_FLOW_FUNCTION_TTL: Duration = Duration::from_secs(30 * 60); + +/// The behavior of a sub flow function returned at execution time. Invoked +/// with every parameter of the request passed through positionally (a sub flow +/// function has no declaration naming its parameters). +/// +/// A plain `Fn(&FunctionContext, Vec) -> Result` +/// closure implements this automatically, so the common case — returning a +/// function that closes over values computed in the outer `run` — needs no +/// extra type. Implement the trait by hand for a handler that has to `.await`. +#[async_trait] +pub trait SubFlowFunctionHandler: Send + Sync + 'static { + #[allow(clippy::double_must_use)] // `async_trait` adds `#[must_use]` to its boxed future. + async fn run( + &self, + context: &FunctionContext, + parameters: Vec, + ) -> Result; +} + +#[async_trait] +impl SubFlowFunctionHandler for F +where + F: Fn(&FunctionContext, Vec) -> Result + Send + Sync + 'static, +{ + async fn run( + &self, + context: &FunctionContext, + parameters: Vec, + ) -> Result { + (self)(context, parameters) + } +} + +/// Describes a sub flow function's input or output. A plain JSON schema object +/// (e.g. produced from a Rust type's [`JsonSchema`] impl via +/// [`SubFlowFunctionDefinition::input_schema`]) is converted to the wire +/// `Struct` while the result is sent; a ready-made [`Struct`] (e.g. one taken +/// from a `SubFlow`) is passed through unchanged. +#[derive(Clone, Debug)] +pub enum SubFlowFunctionSchema { + Json(PlainValue), + Struct(Struct), +} + +impl From for SubFlowFunctionSchema { + fn from(value: Struct) -> Self { + Self::Struct(value) + } +} + +impl From for SubFlowFunctionSchema { + fn from(value: PlainValue) -> Self { + Self::Json(value) + } +} + +/// A sub flow function returned from a runtime function's handler in place of +/// raw data. See the module docs. +#[derive(Clone)] +pub struct SubFlowFunctionDefinition { + pub name: String, + pub description: String, + pub input_schema: Option, + pub output_schema: Option, + pub(crate) handler: Arc, +} + +impl SubFlowFunctionDefinition { + /// A new definition with the given `name` and behavior. Name a non-empty + /// string: a function returned without one is reported to Aquila as a + /// `SUB_FLOW_FUNCTION_MISSING_NAME` error rather than registered. + pub fn new(name: impl Into, handler: impl SubFlowFunctionHandler) -> Self { + Self { + name: name.into(), + description: String::new(), + input_schema: None, + output_schema: None, + handler: Arc::new(handler), + } + } + + pub fn description(mut self, description: impl Into) -> Self { + self.description = description.into(); + self + } + + /// Declares the input the function accepts, from a Rust type's + /// [`JsonSchema`] impl under the deserialize (validation) contract — what + /// the function reads. + pub fn input_schema(mut self) -> Self { + self.input_schema = Some(SubFlowFunctionSchema::Json(json_schema_for::(false))); + self + } + + /// Declares the output the function produces, from a Rust type's + /// [`JsonSchema`] impl under the serialize contract — what the function + /// writes. + pub fn output_schema(mut self) -> Self { + self.output_schema = Some(SubFlowFunctionSchema::Json(json_schema_for::(true))); + self + } + + /// Sets the input schema from a ready-made [`Struct`] or plain JSON value, + /// for callers that already have one (e.g. taken from a `SubFlow`). + pub fn input(mut self, schema: impl Into) -> Self { + self.input_schema = Some(schema.into()); + self + } + + /// Sets the output schema from a ready-made [`Struct`] or plain JSON value. + pub fn output(mut self, schema: impl Into) -> Self { + self.output_schema = Some(schema.into()); + self + } +} + +/// Converts a Rust type's JSON schema under the contract matching the side it +/// describes: the deserialize contract for an input schema (what the function +/// reads) and the serialize contract for an output schema (what it writes). +fn json_schema_for(serialize: bool) -> PlainValue { + let settings = if serialize { + SchemaSettings::default().for_serialize() + } else { + SchemaSettings::default().for_deserialize() + }; + SchemaGenerator::new(settings) + .into_root_schema_for::() + .to_value() +} + +/// Converts a declared schema into the `shared.Struct` the protobuf field +/// expects, or `None` when there is nothing structured to send. +fn to_struct(schema: Option<&SubFlowFunctionSchema>) -> Option { + match schema { + Some(SubFlowFunctionSchema::Struct(s)) => Some(s.clone()), + Some(SubFlowFunctionSchema::Json(json)) => struct_from_json(json), + None => struct_from_json(&serde_json::json!({})), + } +} + +fn struct_from_json(json: &PlainValue) -> Option { + match construct_value(json).kind { + Some(Kind::StructValue(s)) => Some(s), + _ => None, + } +} + +/// Registers the sub flow function on the connection and builds the +/// `shared.SubFlowFunctionResult` that reports it back to Aquila. The +/// identifier in the result is the one Aquila sends back on an execution +/// request when the flow invokes the function. +pub(crate) fn build_sub_flow_function_result( + inner: &Arc, + definition: SubFlowFunctionDefinition, +) -> Result { + if definition.name.is_empty() { + return Err(HerculesError::runtime( + "SUB_FLOW_FUNCTION_MISSING_NAME", + Some("A sub flow function returned from a run block must declare a name".to_string()), + )); + } + let name = definition.name.clone(); + let description = definition.description.clone(); + let input_schema = to_struct(definition.input_schema.as_ref()); + let output_schema = to_struct(definition.output_schema.as_ref()); + let function_identifier = inner.register_sub_flow_function(definition); + Ok(SubFlowFunctionResult { + function: Some(WireSubFlowFunction { + function_identifier, + definition_source: Some("action".to_string()), + }), + input_schema, + output_schema, + name, + description, + }) +} + +/// Sub flow functions a runtime function returned at execution time, keyed by +/// the identifier reported to Aquila. +/// +/// Each entry has a maximum lifetime of [`SUB_FLOW_FUNCTION_TTL`], counted from +/// when it was registered and never extended by an invocation, so a flow that +/// holds a function forever cannot grow the registry without bound. Expired +/// entries are dropped by a sweep scheduled for the next expiry (see +/// [`ConnectedInner::schedule_sub_flow_sweep`]), and on any read or write, so +/// an expired function is never invoked even if the sweep could not run. +#[derive(Default)] +pub(crate) struct SubFlowFunctionState { + functions: HashMap, + /// Expiry per key, in expiry order — which is insertion order, since the + /// TTL is fixed. + expiries: VecDeque<(String, Instant)>, + /// Whether a sweep task is already running, so a burst of registrations + /// schedules at most one. + sweeper_running: bool, +} + +impl SubFlowFunctionState { + pub(crate) fn insert( + &mut self, + key: String, + value: SubFlowFunctionDefinition, + expires_at: Instant, + ) { + // Drop any prior expiry for the key so a replaced key moves to the back + // and the queue stays ordered by expiry. + self.expiries.retain(|(existing, _)| existing != &key); + self.functions.insert(key.clone(), value); + self.expiries.push_back((key, expires_at)); + } + + pub(crate) fn get(&self, key: &str) -> Option<&SubFlowFunctionDefinition> { + self.functions.get(key) + } + + #[cfg(test)] + pub(crate) fn len(&self) -> usize { + self.functions.len() + } + + /// Drops every entry whose lifetime has passed. + pub(crate) fn prune(&mut self, now: Instant) { + while let Some((key, expires_at)) = self.expiries.front() { + // Ordered by expiry, so the first entry still alive ends the sweep. + if *expires_at > now { + break; + } + let key = key.clone(); + self.expiries.pop_front(); + self.functions.remove(&key); + } + } + + /// The earliest expiry still pending, if any. + pub(crate) fn next_expiry(&self) -> Option { + self.expiries.front().map(|(_, expires_at)| *expires_at) + } + + pub(crate) fn set_sweeper_running(&mut self, running: bool) { + self.sweeper_running = running; + } + + pub(crate) fn sweeper_running(&self) -> bool { + self.sweeper_running + } +} diff --git a/ts/examples/simple-example-ts/package-lock.json b/ts/examples/simple-example-ts/package-lock.json index bdf3d28..8fbe670 100644 --- a/ts/examples/simple-example-ts/package-lock.json +++ b/ts/examples/simple-example-ts/package-lock.json @@ -26,7 +26,7 @@ "hercules": "bin/hercules.js" }, "devDependencies": { - "@code0-tech/tucana": "0.0.74", + "@code0-tech/tucana": "0.0.83", "@emnapi/core": "^1.11.1", "@emnapi/runtime": "^1.11.1", "@grpc/grpc-js": "^1.14.3", @@ -46,7 +46,7 @@ "zod-to-ts": "^2.1.0" }, "peerDependencies": { - "@code0-tech/tucana": "0.0.74", + "@code0-tech/tucana": "0.0.83", "@grpc/grpc-js": "^1.14.3", "@protobuf-ts/grpc-backend": "^2.11.1", "@protobuf-ts/grpc-transport": "^2.11.1", diff --git a/ts/examples/simple-example-ts/src/functions/greeterRuntimeFunction.ts b/ts/examples/simple-example-ts/src/functions/greeterRuntimeFunction.ts new file mode 100644 index 0000000..9c48966 --- /dev/null +++ b/ts/examples/simple-example-ts/src/functions/greeterRuntimeFunction.ts @@ -0,0 +1,42 @@ +import {z} from "zod"; +import { + Description, + DisplayMessage, + FunctionContext, + Identifier, + Name, + Parameter, + Signature, + SubFlowFunctionDefinition, +} from "@code0-tech/hercules"; + +/** + * A runtime function whose run block returns another function instead of raw data. + * + * Hercules registers the returned `run` internally and reports it to Aquila as a sub + * flow function, so the flow can invoke it later. Name, description and the I/O + * schema are declared here because nothing in the module definition describes a + * function that is only created at runtime. The zod schemas are converted to JSON + * schemas while the result is sent. + */ +@Identifier("greeter_runtime") +@Signature("(greeting: STRING): TRANSFORM") +@Name({code: "en-US", content: "Greeter"}) +@DisplayMessage({code: "en-US", content: "Build a greeter saying ${greeting}"}) +@Parameter({ + runtimeName: "greeting", + name: [{code: "en-US", content: "Greeting"}], + description: [{code: "en-US", content: "The greeting the returned function prefixes a name with"}], +}) +@Description({code: "en-US", content: "Returns a function that greets a name"}) +export class GreeterRuntimeFunction { + run(_context: FunctionContext, greeting: string): SubFlowFunctionDefinition { + return { + name: "greet", + description: `Greets a name with "${greeting}"`, + inputSchema: z.object({name: z.string()}), + outputSchema: z.string(), + run: (name) => `${greeting} ${String(name)}`, + }; + } +} diff --git a/ts/examples/simple-example-ts/src/index.ts b/ts/examples/simple-example-ts/src/index.ts index 898fcdc..df1f99e 100644 --- a/ts/examples/simple-example-ts/src/index.ts +++ b/ts/examples/simple-example-ts/src/index.ts @@ -3,6 +3,7 @@ import {FibonacciRuntimeFunction} from "./functions/fibonacciRuntimeFunction.js" import {FibonacciFunction} from "./functions/fibonacciFunction.js"; import {ForEachRuntimeFunction} from "./functions/forEachRuntimeFunction.js"; import {ForEachConsumersRuntimeFunction} from "./functions/forEachConsumersRuntimeFunction.js"; +import {GreeterRuntimeFunction} from "./functions/greeterRuntimeFunction.js"; import {UserCreatedRuntimeEvent} from "./events/userCreatedRuntimeEvent.js"; import {EmailDataType} from "./data_types/emailDataType.js"; @@ -33,6 +34,9 @@ action.registerRuntimeFunction(ForEachRuntimeFunction); // Runtime function that takes a list of consumers (each an inline ${signature} sub flow reference) action.registerRuntimeFunction(ForEachConsumersRuntimeFunction); +// Runtime function that returns a sub flow function instead of raw data +action.registerRuntimeFunction(GreeterRuntimeFunction); + // Data type: derived from Zod schema action.registerDataTypeClass(EmailDataType); diff --git a/ts/package-lock.json b/ts/package-lock.json index 3bed6c6..ae95ba7 100644 --- a/ts/package-lock.json +++ b/ts/package-lock.json @@ -12,7 +12,7 @@ "hercules": "bin/hercules.js" }, "devDependencies": { - "@code0-tech/tucana": "0.0.82", + "@code0-tech/tucana": "0.0.83", "@emnapi/core": "^1.11.1", "@emnapi/runtime": "^1.11.1", "@grpc/grpc-js": "^1.14.3", @@ -32,7 +32,7 @@ "zod-to-ts": "^2.1.0" }, "peerDependencies": { - "@code0-tech/tucana": "0.0.82", + "@code0-tech/tucana": "0.0.83", "@grpc/grpc-js": "^1.14.3", "@protobuf-ts/grpc-backend": "^2.11.1", "@protobuf-ts/grpc-transport": "^2.11.1", @@ -119,9 +119,9 @@ } }, "node_modules/@code0-tech/tucana": { - "version": "0.0.82", - "resolved": "https://registry.npmjs.org/@code0-tech/tucana/-/tucana-0.0.82.tgz", - "integrity": "sha512-9ImPNeGV09BeLT0M8/F08DHDe10MO6UI/VfE7i2aEYtRsB9DS3jIAGrcMEV67RXdmAj5ZWTtvxWVGOkem0CljQ==", + "version": "0.0.83", + "resolved": "https://registry.npmjs.org/@code0-tech/tucana/-/tucana-0.0.83.tgz", + "integrity": "sha512-scBOjgf3oIQF8pUlskqPHqfzPvgCJgG1+/tt9RAKWZ1IdCdSqTRG768uEVeTjpuryiGRiO0p8r4bCXNEQNRtRg==", "dev": true, "license": "MIT" }, diff --git a/ts/package.json b/ts/package.json index c6f2f23..84f6c9d 100644 --- a/ts/package.json +++ b/ts/package.json @@ -39,7 +39,7 @@ "devDependencies": { "@emnapi/core": "^1.11.1", "@emnapi/runtime": "^1.11.1", - "@code0-tech/tucana": "0.0.82", + "@code0-tech/tucana": "0.0.83", "@grpc/grpc-js": "^1.14.3", "@protobuf-ts/grpc-backend": "^2.11.1", "@protobuf-ts/grpc-transport": "^2.11.1", @@ -57,7 +57,7 @@ "vitest": "^4.1.2" }, "peerDependencies": { - "@code0-tech/tucana": "0.0.82", + "@code0-tech/tucana": "0.0.83", "@grpc/grpc-js": "^1.14.3", "@protobuf-ts/grpc-backend": "^2.11.1", "@protobuf-ts/grpc-transport": "^2.11.1", diff --git a/ts/src/action.ts b/ts/src/action.ts index ba63e15..cbbd4c6 100644 --- a/ts/src/action.ts +++ b/ts/src/action.ts @@ -16,7 +16,7 @@ import type {EventClass} from "./models/event.model"; import {eventMap} from "./map/event.map"; import type {RuntimeEventClass} from "./models/runtime_event.model"; import {runtimeEventMap} from "./map/runtime_event.map"; -import type {ConfigurationDefinition, Translation} from "./types"; +import type {ConfigurationDefinition, SubFlowFunctionDefinition, Translation} from "./types"; import {CodeZeroEvent, type CodeZeroEventMap} from "./events"; import {createConnection} from "./internal/connection"; import {buildModule} from "./internal/module-builder"; @@ -24,6 +24,7 @@ import {ConfigManager} from "./manager/config-manager"; import {FlowManager} from "./manager/FlowManager"; import {FunctionManager} from "./manager/FunctionManager"; import {RuntimeFunctionManager} from "./manager/RuntimeFunctionManager"; +import {SubFlowFunctionManager} from "./manager/SubFlowFunctionManager"; import {DataTypeManager} from "./manager/DataTypeManager"; import {EventManager} from "./manager/EventManager"; import {RuntimeEventManager} from "./manager/RuntimeEventManager"; @@ -59,6 +60,7 @@ export class Action extends EventEmitter { readonly flows = new FlowManager(); readonly functions = new FunctionManager(); readonly runtimeFunctions = new RuntimeFunctionManager(); + readonly subFlowFunctions = new SubFlowFunctionManager(); readonly dataTypes = new DataTypeManager(); readonly events = new EventManager(); readonly runtimeEvents = new RuntimeEventManager(); @@ -103,6 +105,37 @@ export class Action extends EventEmitter { } } + /** + * Register a sub flow function a runtime function returned from its `run` block + * and return the identifier Aquila invokes it by. Each returned function gets its + * own identifier, so repeatedly returning one (e.g. once per execution) never + * overwrites a function a running flow still holds. + * + * Aquila gives no signal that a flow is done with a function, so a registration + * expires on its own 30 minutes after it was registered and can no longer be + * invoked (see {@link SubFlowFunctionManager}). Call + * `subFlowFunctions.delete(identifier)` to drop one earlier. + */ + registerSubFlowFunction(definition: SubFlowFunctionDefinition): string { + const identifier = `${this._subFlowFunctionPrefix}${definition.name}::${randomUUID()}`; + this.subFlowFunctions.set(identifier, definition); + return identifier; + } + + /** + * True for an identifier this action minted for a sub flow function, whether or + * not it is still registered. An execution request naming one that is gone + * addresses a function that expired (or was deleted) rather than an unknown one, + * which lets the execution handler report it as such. + */ + isSubFlowFunctionIdentifier(identifier: string): boolean { + return identifier.startsWith(this._subFlowFunctionPrefix); + } + + private get _subFlowFunctionPrefix(): string { + return `${this._identifier}::`; + } + registerDataTypeClass(klass: DataTypeClass) { const def = dataTypeMap(klass); this.dataTypes.set(def.identifier, def); diff --git a/ts/src/actions/Execution.ts b/ts/src/actions/Execution.ts index 2accd9a..ce90eaa 100644 --- a/ts/src/actions/Execution.ts +++ b/ts/src/actions/Execution.ts @@ -1,34 +1,120 @@ import {constructValue, PlainValue} from "@code0-tech/tucana/helpers"; import {ActionExecutionRequest, ActionExecutionResponse, ActionTransferRequest} from "@code0-tech/tucana/aquila"; -import {NodeExecutionResult, Error as ProtoError} from "@code0-tech/tucana/shared"; +import {NodeExecutionResult, NodeValueExecutionResult, Error as ProtoError} from "@code0-tech/tucana/shared"; import {FunctionContext, RuntimeError} from "../types"; -import {RuntimeFunctionProps} from "../models/runtime_function.model"; +import type {RuntimeFunctionReturn} from "../models/runtime_function.model"; import {ResolvedValue, resolveNodeValue} from "../internal/literal"; +import {buildSubFlowFunctionResult, isSubFlowFunctionDefinition} from "../internal/sub-flow-function"; import {CodeZeroEvent} from "../events"; import type {Action} from "../action"; export const packetType = "execution"; +/** A function an execution request may address: a registered runtime function, or a sub flow function returned by one. */ +interface ExecutionTarget { + handler: (...args: (PlainValue | undefined)[]) => Promise | RuntimeFunctionReturn; + /** + * Number of declared parameters. Undefined for a sub flow function, which has no + * definition declaring them, so every parameter of the request is passed through. + */ + parameterCount?: number; +} + function nowMicros(): bigint { return BigInt(Math.floor((performance.timeOrigin + performance.now()) * 1000)); } -function buildParams(action: Action, execution: ActionExecutionRequest, func: RuntimeFunctionProps): (ResolvedValue | undefined)[] { - return (func.parameters || []).map((_param, index) => +function resolveTarget(action: Action, identifier: string): ExecutionTarget | undefined { + const func = action.runtimeFunctions.get(identifier); + if (func) return {handler: func.handler, parameterCount: func.parameters?.length ?? 0}; + + const subFlowFunction = action.subFlowFunctions.get(identifier); + // Bind rather than wrap so the handler keeps its arity, which decides whether a + // function context is passed in (see below). + if (subFlowFunction) return {handler: subFlowFunction.run.bind(subFlowFunction)}; + + return undefined; +} + +function buildParams(action: Action, execution: ActionExecutionRequest, target: ExecutionTarget): (ResolvedValue | undefined)[] { + const count = target.parameterCount ?? execution.parameters.length; + return Array.from({length: count}, (_value, index) => resolveNodeValue(action, execution.parameters?.[index]), ); } +/** + * Wrap a handler's return value in a `NodeValueExecutionResult`: raw data goes out as + * a value, a returned sub flow function is registered on the action and reported as a + * `SubFlowFunctionResult` so the flow can invoke it later. + */ +function toNodeValue(action: Action, value: RuntimeFunctionReturn): NodeValueExecutionResult { + if (isSubFlowFunctionDefinition(value)) { + return NodeValueExecutionResult.create({ + value: { + oneofKind: "subFlowFunctionResult", + subFlowFunctionResult: buildSubFlowFunctionResult(action, value), + }, + }); + } + return NodeValueExecutionResult.create({ + value: {oneofKind: "rawValue", rawValue: constructValue(value ?? null)}, + }); +} + +/** Report a failed execution back to Aquila. */ +function sendError(action: Action, execution: ActionExecutionRequest, error: unknown, startedAt: bigint): void { + const finishedAt = nowMicros(); + const isRuntimeError = error instanceof RuntimeError; + action.stream!.requests.send(ActionTransferRequest.create({ + data: { + oneofKind: "result", + result: ActionExecutionResponse.create({ + executionIdentifier: execution.executionIdentifier, + nodeResult: NodeExecutionResult.create({ + id: {oneofKind: undefined}, + startedAt, + finishedAt, + parameterResults: [], + result: { + oneofKind: "error", + error: ProtoError.create({ + code: isRuntimeError ? (error as RuntimeError).code : "UNKNOWN_ERROR", + category: "RUNTIME", + message: isRuntimeError + ? ((error as RuntimeError).description || (error as Error).message) + : String(error), + timestamp: finishedAt, + version: action.version, + dependencies: {}, + }), + }, + }), + }), + }, + })).catch(err => console.error("Failed to send execution error:", err)); +} + export function handle(action: Action, execution: ActionExecutionRequest): void { action.emit(CodeZeroEvent.executionRequestReceived, execution); - const func = action.runtimeFunctions.get(execution.functionIdentifier); - if (!func) { + const target = resolveTarget(action, execution.functionIdentifier); + if (!target) { + if (action.isSubFlowFunctionIdentifier(execution.functionIdentifier)) { + // A sub flow function this action returned earlier, whose maximum + // lifetime has passed (or which was deleted). Fail the node rather than + // staying silent, so the flow does not wait for a result that cannot come. + sendError(action, execution, new RuntimeError( + "SUB_FLOW_FUNCTION_EXPIRED", + `Sub flow function ${execution.functionIdentifier} is no longer registered: its maximum lifetime has passed`, + ), nowMicros()); + return; + } console.error("Received execution request for unknown function:", execution.functionIdentifier); return; } - const params = buildParams(action, execution, func); + const params = buildParams(action, execution, target); const conf = action.configs.get(execution.projectId) ?? { projectId: 0n, @@ -44,11 +130,11 @@ export function handle(action: Action, execution: ActionExecutionRequest): void }; const startedAt = nowMicros(); - const funcHandler = func.handler; + const funcHandler = target.handler; const allParams: (PlainValue | undefined)[] = funcHandler.length === params.length + 1 ? [context as PlainValue, ...params] : params; - Promise.resolve(funcHandler(...allParams)).then((value: PlainValue) => { + Promise.resolve(funcHandler(...allParams)).then((value: RuntimeFunctionReturn) => { const finishedAt = nowMicros(); action.stream!.requests.send(ActionTransferRequest.create({ data: { @@ -60,41 +146,11 @@ export function handle(action: Action, execution: ActionExecutionRequest): void startedAt, finishedAt, parameterResults: [], - result: {oneofKind: "success", success: constructValue(value ?? null)}, + result: {oneofKind: "success", success: toNodeValue(action, value)}, }), }), }, })).catch(err => console.error("Failed to send execution result:", err)); - }).catch((error: unknown) => { - const finishedAt = nowMicros(); - const isRuntimeError = error instanceof RuntimeError; - action.stream!.requests.send(ActionTransferRequest.create({ - data: { - oneofKind: "result", - result: ActionExecutionResponse.create({ - executionIdentifier: execution.executionIdentifier, - nodeResult: NodeExecutionResult.create({ - id: {oneofKind: undefined}, - startedAt, - finishedAt, - parameterResults: [], - result: { - oneofKind: "error", - error: ProtoError.create({ - code: isRuntimeError ? (error as RuntimeError).code : "UNKNOWN_ERROR", - category: "RUNTIME", - message: isRuntimeError - ? ((error as RuntimeError).description || (error as Error).message) - : String(error), - timestamp: finishedAt, - version: action.version, - dependencies: {}, - }), - }, - }), - }), - }, - })).catch(err => console.error("Failed to send execution error:", err)); - }); + }).catch((error: unknown) => sendError(action, execution, error, startedAt)); } diff --git a/ts/src/index.ts b/ts/src/index.ts index 4545b74..67e4cdd 100644 --- a/ts/src/index.ts +++ b/ts/src/index.ts @@ -12,6 +12,7 @@ export * from "./manager/config-manager" export * from "./manager/FlowManager" export * from "./manager/FunctionManager" export * from "./manager/RuntimeFunctionManager" +export * from "./manager/SubFlowFunctionManager" export * from "./manager/DataTypeManager" export * from "./manager/EventManager" export * from "./manager/RuntimeEventManager" diff --git a/ts/src/internal/sub-flow-function.ts b/ts/src/internal/sub-flow-function.ts new file mode 100644 index 0000000..c287b01 --- /dev/null +++ b/ts/src/internal/sub-flow-function.ts @@ -0,0 +1,72 @@ +import {constructValue, type PlainValue} from "@code0-tech/tucana/helpers"; +import {SubFlowFunctionResult, type Struct} from "@code0-tech/tucana/shared"; +import type {ZodTypeAny} from "zod"; +import {RuntimeError, type SubFlowFunctionDefinition, type SubFlowFunctionSchema} from "../types"; +import {zodToJsonSchema} from "./zod-schema"; +import type {Action} from "../action"; + +/** + * True when a runtime function's `run` block returned a sub flow function instead of + * raw data. A sub flow function is told apart from a plain object by its callable + * `run` property — plain values never carry functions. + */ +export function isSubFlowFunctionDefinition(value: unknown): value is SubFlowFunctionDefinition { + return typeof value === "object" + && value !== null + && typeof (value as SubFlowFunctionDefinition).run === "function"; +} + +/** + * True for a zod schema. Duck-typed rather than checked with `instanceof`, so a + * schema built by a second copy of zod in the dependency tree still matches. + */ +function isZodSchema(schema: SubFlowFunctionSchema): schema is ZodTypeAny { + const candidate = schema as {parse?: unknown; safeParse?: unknown}; + return typeof candidate.parse === "function" && typeof candidate.safeParse === "function"; +} + +/** True for a ready-made protobuf {@link Struct} rather than a plain schema object. */ +function isStruct(schema: SubFlowFunctionSchema): schema is Struct { + const fields = (schema as Struct).fields; + if (fields == null || typeof fields !== "object") return false; + return Object.values(fields).every(field => + field != null && typeof field === "object" && "kind" in field); +} + +/** + * Convert a declared schema into the `shared.Struct` the protobuf field expects. A + * zod schema becomes a JSON schema first — describing the input it accepts for a + * function's input schema, and the output it produces for its output schema. + */ +export function toStruct(schema: SubFlowFunctionSchema, io: "input" | "output"): Struct | undefined { + if (isStruct(schema)) return schema; + const json = isZodSchema(schema) ? zodToJsonSchema(schema, io) : schema; + const value = constructValue(json as PlainValue); + return value.kind.oneofKind === "structValue" ? value.kind.structValue : undefined; +} + +/** + * Register the sub flow function on the action and build the + * `shared.SubFlowFunctionResult` that reports it back to Aquila. The identifier in + * the result is the one Aquila sends back on an execution request when the flow + * invokes the function. + */ +export function buildSubFlowFunctionResult( + action: Action, + definition: SubFlowFunctionDefinition, +): SubFlowFunctionResult { + if (!definition.name) { + throw new RuntimeError( + "SUB_FLOW_FUNCTION_MISSING_NAME", + "A sub flow function returned from a run block must declare a name", + ); + } + const functionIdentifier = action.registerSubFlowFunction(definition); + return SubFlowFunctionResult.create({ + function: {functionIdentifier, definitionSource: "action"}, + name: definition.name, + description: definition.description ?? "", + inputSchema: toStruct(definition.inputSchema ?? {}, "input"), + outputSchema: toStruct(definition.outputSchema ?? {}, "output"), + }); +} diff --git a/ts/src/internal/zod-schema.ts b/ts/src/internal/zod-schema.ts index 2d4c98b..4723334 100644 --- a/ts/src/internal/zod-schema.ts +++ b/ts/src/internal/zod-schema.ts @@ -1,4 +1,5 @@ import {toJSONSchema, type ZodTypeAny} from "zod"; +import type {PlainValue} from "@code0-tech/tucana/helpers"; import {zodToTs, printNode, createAuxiliaryTypeStore, type TypeOverrideMap} from "zod-to-ts"; import type {DefinitionDataTypeRule} from "@code0-tech/tucana/shared"; @@ -56,6 +57,17 @@ export function zodToTypeString(schema: ZodTypeAny): string { return type; } +/** + * Convert a zod schema into a plain JSON schema object. + * + * `io` picks which side of a transforming schema is described: the input it accepts + * or the output it produces. Cyclic schemas are emitted as `$ref`s rather than + * rejected, so a self-referencing schema still converts. + */ +export function zodToJsonSchema(schema: ZodTypeAny, io: "input" | "output"): Record { + return toJSONSchema(schema, {cycles: "ref", io}) as Record; +} + export function zodToRules(schema: ZodTypeAny): DefinitionDataTypeRule[] { const json = toJSONSchema(schema, {cycles: "ref"}) as JsonSchema; const rules: DefinitionDataTypeRule[] = []; diff --git a/ts/src/manager/SubFlowFunctionManager.ts b/ts/src/manager/SubFlowFunctionManager.ts new file mode 100644 index 0000000..66e7623 --- /dev/null +++ b/ts/src/manager/SubFlowFunctionManager.ts @@ -0,0 +1,86 @@ +import {BaseManager} from "./BaseManager"; +import type {SubFlowFunctionDefinition} from "../types"; + +/** How long a returned sub flow function stays in memory and invocable: 30 minutes. */ +export const SUB_FLOW_FUNCTION_TTL_MS = 30 * 60 * 1000; + +/** + * Sub flow functions a runtime function returned at execution time, keyed by the + * identifier reported to Aquila. Unlike the other managers this one is not filled by + * registering classes up front: entries appear while the action runs (see + * {@link Action.registerSubFlowFunction}). + * + * Each entry has a maximum lifetime of {@link SUB_FLOW_FUNCTION_TTL_MS}, counted from + * when it was registered and never extended by an invocation, so a flow that holds a + * function forever cannot grow the registry without bound. Expired entries are + * dropped by a sweep scheduled for the next expiry, and on any read or write, so an + * expired function is never invoked even if the sweep could not run. + */ +export class SubFlowFunctionManager extends BaseManager { + /** Expiry timestamp per key. Insertion order is expiry order, since the TTL is fixed. */ + private readonly _expiries = new Map(); + private _sweepTimer?: ReturnType; + + constructor(private readonly _ttlMs: number = SUB_FLOW_FUNCTION_TTL_MS) { + super(); + } + + set(key: string, value: SubFlowFunctionDefinition): this { + // Re-insert the expiry so a replaced key moves to the back and the map stays + // ordered by expiry. + this._expiries.delete(key); + super.set(key, value); + this._expiries.set(key, Date.now() + this._ttlMs); + this.prune(); + return this; + } + + get(key: string): SubFlowFunctionDefinition | undefined { + this.prune(); + return super.get(key); + } + + has(key: string): boolean { + this.prune(); + return super.has(key); + } + + delete(key: string): boolean { + this._expiries.delete(key); + const deleted = super.delete(key); + this._scheduleSweep(Date.now()); + return deleted; + } + + clear(): void { + this._expiries.clear(); + super.clear(); + this._scheduleSweep(Date.now()); + } + + /** Drop every entry whose lifetime has passed, then schedule the next sweep. */ + prune(now: number = Date.now()): void { + for (const [key, expiresAt] of this._expiries) { + // Ordered by expiry, so the first entry still alive ends the sweep. + if (expiresAt > now) break; + this._expiries.delete(key); + super.delete(key); + } + this._scheduleSweep(now); + } + + private _scheduleSweep(now: number): void { + if (this._sweepTimer !== undefined) { + clearTimeout(this._sweepTimer); + this._sweepTimer = undefined; + } + const next = this._expiries.values().next(); + if (next.done) return; + this._sweepTimer = setTimeout(() => { + this._sweepTimer = undefined; + this.prune(); + }, Math.max(0, next.value - now)); + // Expiring entries must not keep the process alive on its own. + (this._sweepTimer as {unref?: () => void}).unref?.(); + } +} diff --git a/ts/src/map/runtime_function.map.ts b/ts/src/map/runtime_function.map.ts index 925caaa..6de2a86 100644 --- a/ts/src/map/runtime_function.map.ts +++ b/ts/src/map/runtime_function.map.ts @@ -1,4 +1,3 @@ -import type {PlainValue} from "@code0-tech/tucana/helpers"; import {RuntimeFunctionProps, RuntimeFunctionClass} from "../models/runtime_function.model"; import {FunctionParameterProps} from "../models/function.model"; @@ -20,7 +19,7 @@ export const runtimeFunctionMap = (klass: RuntimeFunctionClass): RuntimeFunction if (!signature) throw new Error(`Runtime function class ${klass.name} is missing a signature. Add @Signature("(param1: TYPE_1): RETURN_TYPE") to the class.`); const instance = new klass(); - const handler = instance.run.bind(instance) as (...args: (PlainValue | undefined)[]) => Promise | PlainValue; + const handler = instance.run.bind(instance) as RuntimeFunctionProps["handler"]; return { runtimeName: identifier, diff --git a/ts/src/models/runtime_function.model.ts b/ts/src/models/runtime_function.model.ts index 2c62fbf..46092f8 100644 --- a/ts/src/models/runtime_function.model.ts +++ b/ts/src/models/runtime_function.model.ts @@ -1,7 +1,14 @@ import {PlainValue} from "@code0-tech/tucana/helpers"; -import {Translation} from "../types"; +import {SubFlowFunctionDefinition, Translation} from "../types"; import {FunctionParameterProps} from "./function.model"; +/** + * What a runtime function's `run` block may produce: raw data, or a + * {@link SubFlowFunctionDefinition} that Hercules registers and reports to Aquila as + * a `shared.SubFlowFunctionResult`. + */ +export type RuntimeFunctionReturn = PlainValue | SubFlowFunctionDefinition; + export interface RuntimeFunctionProps { runtimeName: string, parameters?: FunctionParameterProps[], @@ -15,11 +22,11 @@ export interface RuntimeFunctionProps { alias?: Translation[], displayIcon?: string, design?: string, - handler: (...args: (PlainValue | undefined)[]) => Promise | PlainValue, + handler: (...args: (PlainValue | undefined)[]) => Promise | RuntimeFunctionReturn, } export interface RuntimeFunctionRunnable { - run(...args: (PlainValue | undefined)[]): Promise | PlainValue; + run(...args: (PlainValue | undefined)[]): Promise | RuntimeFunctionReturn; } export type RuntimeFunctionClass = diff --git a/ts/src/types.ts b/ts/src/types.ts index b2f82c6..dad1cf6 100644 --- a/ts/src/types.ts +++ b/ts/src/types.ts @@ -1,6 +1,7 @@ import {FlowTypeSetting_UniquenessScope, RuntimeFlowTypeSetting_UniquenessScope} from "@code0-tech/tucana/shared"; import type {Struct} from "@code0-tech/tucana/shared"; import type {ActionNodeSubFlowValue} from "@code0-tech/tucana/aquila"; +import type {ZodTypeAny} from "zod"; import {PlainValue} from "@code0-tech/tucana/helpers"; import 'reflect-metadata'; @@ -24,6 +25,39 @@ export interface SubFlow { readonly outputSchema?: Struct; } +/** + * Describes a sub flow function's input or output. Normally a zod schema, which is + * converted to a JSON schema while the result is sent. A plain JSON schema object or + * a ready-made {@link Struct} (e.g. one taken from a {@link SubFlow}) is passed + * through unchanged. + */ +export type SubFlowFunctionSchema = ZodTypeAny | Struct | Record; + +/** + * A sub flow function returned from a runtime function's `run` block in place of raw + * data. Hercules registers `run` internally and reports the function to Aquila as a + * `shared.SubFlowFunctionResult`, so the flow can invoke it later like any other sub + * flow. Its name, description and I/O schema have to be declared here, because + * nothing in the action's module definition describes a function created at runtime. + * + * ```ts + * return { + * name: "greet", + * description: "Greets a name", + * inputSchema: z.object({name: z.string()}), + * outputSchema: z.string(), + * run: (name) => `Hello ${String(name)}`, + * }; + * ``` + */ +export interface SubFlowFunctionDefinition { + name: string; + description: string; + inputSchema: SubFlowFunctionSchema; + outputSchema: SubFlowFunctionSchema; + run(...args: (PlainValue | undefined)[]): Promise | R; +} + export interface FunctionContext { projectId: number | bigint, executionId: string, diff --git a/ts/test/execution.test.ts b/ts/test/execution.test.ts new file mode 100644 index 0000000..b86b1f0 --- /dev/null +++ b/ts/test/execution.test.ts @@ -0,0 +1,237 @@ +import "reflect-metadata"; +import {afterEach, beforeEach, describe, expect, it, vi} from "vitest"; +import {z} from "zod"; +import type {ActionExecutionRequest, ActionTransferRequest} from "@code0-tech/tucana/aquila"; +import type {Value} from "@code0-tech/tucana/shared"; +import {Action} from "../src/action"; +import {SUB_FLOW_FUNCTION_TTL_MS} from "../src/manager/SubFlowFunctionManager"; +import {handle} from "../src/actions/Execution"; +import type {RuntimeFunctionProps} from "../src/models/runtime_function.model"; +import type {SubFlowFunctionDefinition} from "../src/types"; + +const str = (stringValue: string): Value => ({kind: {oneofKind: "stringValue", stringValue}}); +const list = (...values: Value[]): Value => ({kind: {oneofKind: "listValue", listValue: {values}}}); +const struct = (fields: Record): Value => ({kind: {oneofKind: "structValue", structValue: {fields}}}); + +const literalParam = (value: Value) => ({value: {oneofKind: "literalValue" as const, literalValue: {value, references: []}}}); + +function runtimeFunction(props: Partial & Pick): RuntimeFunctionProps { + return {runtimeName: "fn", signature: "(): STRING", parameters: [], ...props}; +} + +/** An action with a stubbed stream that records everything it would send to Aquila. */ +function stubAction(): {action: Action; sent: ActionTransferRequest[]} { + const action = new Action("my_action", "1.0.0", undefined, "me", "icon", "docs", []); + const sent: ActionTransferRequest[] = []; + (action as unknown as {_stream: unknown})._stream = { + requests: {send: (request: ActionTransferRequest) => (sent.push(request), Promise.resolve())}, + }; + return {action, sent}; +} + +const request = (functionIdentifier: string, parameters: ActionExecutionRequest["parameters"] = []): ActionExecutionRequest => ({ + executionIdentifier: "exec-1", + functionIdentifier, + parameters, + projectId: 1n, +}); + +/** The NodeValueExecutionResult of the single result the action sent. */ +function sentResult(sent: ActionTransferRequest[]) { + expect(sent).toHaveLength(1); + const {data} = sent[0]; + if (data.oneofKind !== "result") throw new Error(`expected a result, got ${data.oneofKind}`); + const {result} = data.result.nodeResult!; + if (result.oneofKind !== "success") throw new Error(`expected a success, got ${result.oneofKind}`); + return result.success.value; +} + +/** The ProtoError of the single error result the action sent. */ +function sentError(sent: ActionTransferRequest[]) { + expect(sent).toHaveLength(1); + const {data} = sent[0]; + if (data.oneofKind !== "result") throw new Error(`expected a result, got ${data.oneofKind}`); + const {result} = data.result.nodeResult!; + if (result.oneofKind !== "error") throw new Error(`expected an error, got ${result.oneofKind}`); + return result.error; +} + +const subFlowFunction = (run: SubFlowFunctionDefinition["run"]): SubFlowFunctionDefinition => ({ + name: "greet", + description: "Greets someone", + inputSchema: z.object({name: z.string()}), + outputSchema: z.string(), + run, +}); + +describe("execution results", () => { + let action: Action; + let sent: ActionTransferRequest[]; + + beforeEach(() => ({action, sent} = stubAction())); + + it("reports raw data as a raw value", async () => { + action.runtimeFunctions.set("fn", runtimeFunction({handler: () => "hello"})); + + handle(action, request("fn")); + await new Promise(resolve => setImmediate(resolve)); + + expect(sentResult(sent)).toEqual({oneofKind: "rawValue", rawValue: str("hello")}); + }); + + it("registers a returned sub flow function and reports it as a sub flow function result", async () => { + action.runtimeFunctions.set("fn", runtimeFunction({handler: () => subFlowFunction(name => `Hi ${name}`)})); + + handle(action, request("fn")); + await new Promise(resolve => setImmediate(resolve)); + + const value = sentResult(sent); + if (value.oneofKind !== "subFlowFunctionResult") throw new Error("expected a sub flow function result"); + const result = value.subFlowFunctionResult; + + expect(result.name).toBe("greet"); + expect(result.description).toBe("Greets someone"); + // The zod schemas are converted to JSON schemas before they are sent. The + // input schema describes what the function accepts, so it carries no + // "additionalProperties: false" the way the output side would. + expect(result.inputSchema).toEqual({ + fields: { + $schema: str("https://json-schema.org/draft/2020-12/schema"), + type: str("object"), + properties: struct({name: struct({type: str("string")})}), + required: list(str("name")), + }, + }); + expect(result.outputSchema).toEqual({ + fields: { + $schema: str("https://json-schema.org/draft/2020-12/schema"), + type: str("string"), + }, + }); + expect(result.function?.definitionSource).toBe("action"); + // The identifier Aquila invokes the function by is registered on the action. + expect(action.subFlowFunctions.has(result.function!.functionIdentifier)).toBe(true); + }); + + it("gives each returned sub flow function its own identifier", async () => { + action.runtimeFunctions.set("fn", runtimeFunction({handler: () => subFlowFunction(() => null)})); + + handle(action, request("fn")); + handle(action, request("fn")); + await new Promise(resolve => setImmediate(resolve)); + + expect(sent).toHaveLength(2); + expect(action.subFlowFunctions.size).toBe(2); + }); + + it("executes a registered sub flow function with every parameter of the request", async () => { + const identifier = action.registerSubFlowFunction(subFlowFunction(name => `Hi ${name}`)); + + handle(action, request(identifier, [literalParam(str("Ada"))])); + await new Promise(resolve => setImmediate(resolve)); + + expect(sentResult(sent)).toEqual({oneofKind: "rawValue", rawValue: str("Hi Ada")}); + }); + + it("passes a plain JSON schema object through unchanged", async () => { + action.runtimeFunctions.set("fn", runtimeFunction({ + handler: () => ({...subFlowFunction(() => null), outputSchema: {type: "string"}}), + })); + + handle(action, request("fn")); + await new Promise(resolve => setImmediate(resolve)); + + const value = sentResult(sent); + if (value.oneofKind !== "subFlowFunctionResult") throw new Error("expected a sub flow function result"); + expect(value.subFlowFunctionResult.outputSchema).toEqual({fields: {type: str("string")}}); + }); + + it("reports an unnamed sub flow function as an error", async () => { + action.runtimeFunctions.set("fn", runtimeFunction({ + handler: () => ({...subFlowFunction(() => null), name: ""}), + })); + + handle(action, request("fn")); + await new Promise(resolve => setImmediate(resolve)); + + expect(sentError(sent).code).toBe("SUB_FLOW_FUNCTION_MISSING_NAME"); + expect(action.subFlowFunctions.size).toBe(0); + }); +}); + +describe("sub flow function lifetime", () => { + let action: Action; + let sent: ActionTransferRequest[]; + + beforeEach(() => { + vi.useFakeTimers(); + ({action, sent} = stubAction()); + }); + + afterEach(() => vi.useRealTimers()); + + it("keeps a function invocable right up to its lifetime", () => { + const identifier = action.registerSubFlowFunction(subFlowFunction(() => null)); + + vi.advanceTimersByTime(SUB_FLOW_FUNCTION_TTL_MS - 1); + + expect(action.subFlowFunctions.get(identifier)).toBeDefined(); + }); + + it("drops a function once its lifetime has passed", () => { + const identifier = action.registerSubFlowFunction(subFlowFunction(() => null)); + + vi.advanceTimersByTime(SUB_FLOW_FUNCTION_TTL_MS); + + expect(action.subFlowFunctions.get(identifier)).toBeUndefined(); + }); + + it("frees an expired function without waiting for a read", () => { + action.registerSubFlowFunction(subFlowFunction(() => null)); + expect(action.subFlowFunctions.size).toBe(1); + + // The scheduled sweep clears the entry on its own, so nothing is held just + // because no one looked the function up again. + vi.advanceTimersByTime(SUB_FLOW_FUNCTION_TTL_MS); + + expect(action.subFlowFunctions.size).toBe(0); + }); + + it("does not extend a lifetime when the function is invoked", () => { + const identifier = action.registerSubFlowFunction(subFlowFunction(() => null)); + + vi.advanceTimersByTime(SUB_FLOW_FUNCTION_TTL_MS / 2); + expect(action.subFlowFunctions.get(identifier)).toBeDefined(); + vi.advanceTimersByTime(SUB_FLOW_FUNCTION_TTL_MS / 2); + + expect(action.subFlowFunctions.get(identifier)).toBeUndefined(); + }); + + it("fails the node instead of running an expired function", () => { + const identifier = action.registerSubFlowFunction(subFlowFunction(() => null)); + vi.advanceTimersByTime(SUB_FLOW_FUNCTION_TTL_MS); + + handle(action, request(identifier, [literalParam(str("Ada"))])); + + const error = sentError(sent); + expect(error.code).toBe("SUB_FLOW_FUNCTION_EXPIRED"); + expect(error.message).toContain(identifier); + }); + + it("sends nothing for a function it never minted", () => { + handle(action, request("not_a_function")); + + expect(sent).toEqual([]); + }); + + it("expires entries in registration order", () => { + const first = action.registerSubFlowFunction(subFlowFunction(() => null)); + vi.advanceTimersByTime(SUB_FLOW_FUNCTION_TTL_MS / 2); + const second = action.registerSubFlowFunction(subFlowFunction(() => null)); + + vi.advanceTimersByTime(SUB_FLOW_FUNCTION_TTL_MS / 2); + + expect(action.subFlowFunctions.get(first)).toBeUndefined(); + expect(action.subFlowFunctions.get(second)).toBeDefined(); + }); +});