Abort Handling and Revamp Test

DEVELOPMENT.md

@@ -37,11 +37,37 @@ uv build --sdist
 ## Run tests
 
 ```sh
-# Run the tests
-cd tests/release/integration
-uv run python -m unittest test_wrapper.py
+# Run the whole suite
+uv run python -m pytest
+
+# Skip the slower subprocess tests, or the tests that need a Release server
+uv run python -m pytest -m "not e2e"
+uv run python -m pytest -m "not live"
 ```
 
+Tests are grouped by module under `tests/release/integration/` (`test_<module>.py`), with the
+wrapper split into `test_wrapper_io.py`, `test_wrapper_run.py`, `test_wrapper_execute_task.py`,
+`test_wrapper_abort.py` and the subprocess runs in `test_wrapper_e2e.py`. Shared ids and input
+contexts are in `tests/support.py`; fixtures are in `tests/release/integration/conftest.py`, and
+the task classes the wrapper loads in tests are in `tests/release/integration/tasks/`.
+
+Markers:
+
+- `e2e`: runs the wrapper with `python -m` in a subprocess. The two real-`SIGABRT` tests are
+  skipped on Windows, which cannot deliver `SIGABRT` to a child process.
+- `live`: needs a real server, and is skipped when none is reachable.
+  - `test_api_base_task_live.py` creates and deletes a folder, template, release and variables
+    on a Release server, by default `http://localhost:5516` as `admin`/`admin`; set
+    `RELEASE_LIVE_URL`, `RELEASE_LIVE_USER` and `RELEASE_LIVE_PASSWORD` to use another one.
+  - `test_wrapper_k8s_live.py` runs the wrapper's Kubernetes path against a cluster: input and
+    result secrets, the callback, the input-context watcher, and a full wrapper run. It uses only
+    the kubeconfig context named by `RELEASE_LIVE_K8S_CONTEXT` (default `docker-desktop`, e.g.
+    Docker Desktop with Kubernetes enabled), whatever the current context is, and works in a
+    temporary `sdk-live-<random>` namespace that it deletes afterwards.
+
+Known bugs are pinned by tests marked `xfail`. `xfail_strict` is on, so when a fix makes such a
+test pass, the run fails until the marker is removed.
+
 ## Publish to PyPI
 
 ```sh

README.md

@@ -76,6 +76,21 @@ class ShowVersion(ApiBaseTask):
 
 ## Changelog
 
+### Unreleased
+
+#### 🚀 Features
+
+- Tasks can be aborted. On abort, the SDK calls the task's `abort()` method; long-running tasks can check `is_abort_requested()` and stop early.
+- After an abort, the SDK waits briefly (`ABORT_GRACE_PERIOD_SECONDS`, default 3) so the task's outputs and reporting records are still sent.
+- Added the `AiUsageRecord` reporting record for AI agent runs (model, tokens, cost and more). Requires a Release version that supports `udm.AiUsageRecord`.
+
+#### 🐛 Fixes
+
+- The task result is no longer cut off if an abort signal arrives while it is being reported.
+- Aborted tasks no longer report exit code 0; `sys.exit(n)` now reports exit code `n`.
+- Tasks that don't call `super().__init__()` work again.
+- `setReleaseVariable()`, `setFolderVariable()` and `setGlobalVariable()` now accept a list for list variables.
+
 ### Version 26.3.1
 
 #### 🚀 Features

digitalai/release/integration/__init__.py

@@ -3,5 +3,5 @@
 from .input_context import InputContext
 from .output_context import OutputContext
 from .exceptions import AbortException
-from .reporting_records import BuildRecord, PlanRecord, ItsmRecord,CodeComplianceRecord, DeploymentRecord
+from .reporting_records import BuildRecord, PlanRecord, ItsmRecord, CodeComplianceRecord, DeploymentRecord, AiUsageRecord
 from .logger import dai_logger

digitalai/release/integration/api_base_task.py

@@ -556,8 +556,9 @@ def _value_matches_type(vtype: str, value: Any) -> bool:
         The accepted Python types mirror :meth:`_variable_type_for_value`.
         ``vtype`` is matched by suffix so subtypes are treated as their base
         kind (e.g. ``PasswordStringVariable`` is a string). More specific
-        suffixes are tested first because ``SetStringVariable`` and
-        ``MapStringStringVariable`` also end in ``StringVariable``. A type the
+        suffixes are tested first because ``ListStringVariable``,
+        ``SetStringVariable`` and ``MapStringStringVariable`` also end in
+        ``StringVariable``. A type the
         SDK does not recognise returns ``True`` -- it is left for the server to
         validate rather than blocked here.
         """
@@ -571,7 +572,7 @@ def _value_matches_type(vtype: str, value: Any) -> bool:
             return isinstance(value, (date, str))
         if vtype.endswith("MapStringStringVariable"):
             return isinstance(value, dict)
-        if vtype.endswith("SetStringVariable"):
+        if vtype.endswith(("ListStringVariable", "SetStringVariable")):
             return isinstance(value, (list, set, tuple))
         if vtype.endswith("StringVariable"):
             return isinstance(value, str)

digitalai/release/integration/base_task.py

@@ -58,6 +58,12 @@ def abort(self) -> None:
         immediately. When called from the SDK's signal handler (outside of
         execute_task()), the caller is responsible for catching this exception
         and performing its own cleanup/exit.
+
+        The SDK's signal handler runs abort() on the main thread while execute()
+        may still be running on the worker thread, so an override must be
+        thread-safe. Output properties and reporting records set in an override
+        are sent with the aborted result; set them before calling super().abort(),
+        since it raises.
         """
         self.set_exit_code(1)
         self.set_error_message("Abort requested")
@@ -69,26 +75,38 @@ def is_abort_requested(self) -> bool:
         (e.g. the process received SIGABRT/SIGTERM). Tasks may poll this from
         within execute() to cooperatively stop long-running work promptly.
         """
-        return self._abort_event.is_set()
+        abort_event = getattr(self, "_abort_event", None)
+        return abort_event is not None and abort_event.is_set()
 
     def _request_abort(self) -> None:
         """
         Internal SDK hook: signals is_abort_requested() to start returning True.
         Not intended to be called by task authors.
         """
+        if getattr(self, "_abort_event", None) is None:
+            self._abort_event = threading.Event()
         self._abort_event.set()
 
+    def _ensure_output_context(self) -> None:
+        """
+        Creates the output context if it is missing, for subclasses whose
+        __init__ does not call super().__init__().
+        """
+        if getattr(self, "output_context", None) is None:
+            self.output_context = OutputContext(0, "", {}, [])
+
     def get_output_context(self) -> OutputContext:
         """
         Returns the OutputContext object associated with the task.
         """
+        self._ensure_output_context()
         return self.output_context
 
     def get_output_properties(self) -> Dict[str, Any]:
         """
         Returns the output properties dictionary of the task's OutputContext object.
         """
-        return self.output_context.output_properties
+        return self.get_output_context().output_properties
 
     def get_input_properties(self) -> Dict[str, Any]:
         """
@@ -112,19 +130,19 @@ def set_output_property(self, name: str, value: Any) -> None:
                 f"Invalid data type for value '{value}' in name '{name}' in set_output_property. Accepted data types "
                 f"are: str, int, list, dict, bool")
 
-        self.output_context.output_properties[name] = value
+        self.get_output_context().output_properties[name] = value
 
     def set_exit_code(self, value) -> None:
         """
         Sets the exit code of the task's OutputContext object.
         """
-        self.output_context.exit_code = value
+        self.get_output_context().exit_code = value
 
     def set_error_message(self, value) -> None:
         """
         Sets the error message of the task's OutputContext object.
         """
-        self.output_context.job_error_message = value
+        self.get_output_context().job_error_message = value
 
     def add_comment(self, comment: str) -> None:
         """
@@ -142,7 +160,7 @@ def add_reporting_record(self, reporting_record: Any) -> None:
         """
         Adds a reporting record to the OutputContext
         """
-        self.output_context.reporting_records.append(reporting_record)
+        self.get_output_context().reporting_records.append(reporting_record)
 
     def get_release_server_url(self) -> str:
         """

digitalai/release/integration/reporting_records.py

@@ -106,3 +106,40 @@ class DeploymentRecord:
     environment_name: Optional[str] = None
     server_url: Optional[str] = None
     server_user: Optional[str] = None
+
+
+def _decimal_string(value) -> Optional[str]:
+    """Sends a cost as the decimal string the server stores, e.g. Decimal("0.0123") -> "0.0123"."""
+    return None if value is None else str(value)
+
+
+@dataclass_json(letter_case=LetterCase.CAMEL)
+@dataclass
+class AiUsageRecord:
+    """
+     AiUsageRecord is a data model class which represents one AI agent execution attempt of a task:
+     the language model provider and model, the agent, why it stopped, and its token usage and cost.
+
+     The record's start and end times are taken from the task. The provider's endpoint and user go in
+     server_url and server_user. estimated_cost is a decimal string; pass a str or a Decimal, since a
+     float can carry rounding noise (str(0.1 + 0.2) is "0.30000000000000004").
+
+     Only add it on a Release version that defines udm.AiUsageRecord: an older server cannot read
+     the task's result and fails the task, without its output properties.
+    """
+    target_id: str
+    llm_provider: Optional[str] = None
+    agent: Optional[str] = None
+    model: Optional[str] = None
+    termination_reason: Optional[str] = None
+    iterations: Optional[int] = None
+    tool_calls: Optional[int] = None
+    input_tokens: Optional[int] = None
+    output_tokens: Optional[int] = None
+    total_tokens: Optional[int] = None
+    estimated_cost: Optional[str] = field(default=None, metadata=config(encoder=_decimal_string))
+    currency: str = "USD"
+    type: str = "udm.AiUsageRecord"
+    id: str = ""
+    server_url: Optional[str] = None
+    server_user: Optional[str] = None