mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
Patrick Buckley
c328bebecd
A scheduled task could pin the model and skill of the workstream each
firing creates; it can now also pin its persona and project, so a
schedule can run under, e.g., the researcher persona attached to a
specific project's memory bucket.
The two values live on scheduled_tasks (migration 066, Text NOT NULL
default '') and are passed verbatim to create_workstream at dispatch,
where the node resolves the persona for the workstream kind and gates
the project attach. Empty means "kind-default persona / no project",
resolved late at each firing (mirrors how empty model/skill already
behave) -- existing schedules keep byte-identical dispatch behaviour,
so there is no backfill.
Also fixes a latent bug this feature depends on: admin_create_schedule
read created_by from request.state.user_id, which AuthMiddleware never
sets, so every scheduled task stored created_by=''. It now reads
auth_result.user_id like every other console endpoint. This is now
load-bearing -- the scheduler dispatches under created_by and the node
gates the project attach against it. admin_update_schedule adopts the
editing admin as owner when a project is assigned to a pre-fix orphaned
('') schedule, and re-validates persona/project only when they change
so a since-disabled persona or lost membership does not block unrelated
edits (the node re-checks at dispatch either way).
Wired through: schema + migration (up/down + parity tested), both
storage backends, API schemas, SDK create_workstream and console
create_schedule/update_schedule, scheduler dispatch, and the admin
schedule shelf (persona + project pickers, current value preserved so
an edit cannot silently clear a filtered-out selection).
435 lines
15 KiB
Python
435 lines
15 KiB
Python
"""Tests for turnstone.console.scheduler — TaskScheduler tick and dispatch."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from turnstone.console.scheduler import TaskScheduler
|
|
from turnstone.sdk._types import TurnstoneAPIError
|
|
|
|
|
|
def _wire_lock_storage(storage: MagicMock, initial: dict[str, str] | None = None) -> None:
|
|
"""Configure *storage* mock so upsert/get track scheduler_lock state.
|
|
|
|
The scheduler's ``_try_acquire_lock`` now writes then reads back to
|
|
verify ownership. The mock must reflect what was most recently
|
|
upserted so the read-back succeeds.
|
|
"""
|
|
state: dict[str, dict[str, str] | None] = {"scheduler_lock": initial}
|
|
|
|
def _get(key: str, **_kw: object) -> dict[str, str] | None:
|
|
return state.get(key)
|
|
|
|
def _upsert(key: str, value: str, **_kw: object) -> None:
|
|
state[key] = {"value": value}
|
|
|
|
def _delete(key: str, **_kw: object) -> None:
|
|
state.pop(key, None)
|
|
|
|
storage.get_system_setting.side_effect = _get
|
|
storage.upsert_system_setting.side_effect = _upsert
|
|
storage.delete_system_setting.side_effect = _delete
|
|
|
|
|
|
@pytest.fixture
|
|
def mocks():
|
|
"""Collector and storage mocks for scheduler tests."""
|
|
collector = MagicMock()
|
|
storage = MagicMock()
|
|
# Default: no existing lock
|
|
_wire_lock_storage(storage, initial=None)
|
|
return collector, storage
|
|
|
|
|
|
def _make_task(**overrides):
|
|
"""Build a minimal task dict matching storage row format."""
|
|
defaults = {
|
|
"task_id": "task_001",
|
|
"name": "Test task",
|
|
"description": "",
|
|
"schedule_type": "cron",
|
|
"cron_expr": "0 9 * * *",
|
|
"at_time": "",
|
|
"target_mode": "auto",
|
|
"model": "gpt-5",
|
|
"initial_message": "Run the tests",
|
|
"auto_approve": 0,
|
|
"auto_approve_tools": "",
|
|
"enabled": 1,
|
|
"created_by": "u_admin",
|
|
"next_run": "2020-01-01T09:00:00",
|
|
"last_run": "",
|
|
"created": "2020-01-01T00:00:00",
|
|
"updated": "2020-01-01T00:00:00",
|
|
}
|
|
defaults.update(overrides)
|
|
return defaults
|
|
|
|
|
|
def _make_node(node_id="node-001", reachable=True, ws_total=2, max_ws=10):
|
|
"""Build a minimal node dict matching collector output."""
|
|
return {
|
|
"node_id": node_id,
|
|
"reachable": reachable,
|
|
"ws_total": ws_total,
|
|
"max_ws": max_ws,
|
|
}
|
|
|
|
|
|
def _mock_create_response(ws_id: str = "ws_abc123") -> MagicMock:
|
|
"""Build a mock CreateWorkstreamResponse with the given ws_id."""
|
|
resp = MagicMock()
|
|
resp.ws_id = ws_id
|
|
return resp
|
|
|
|
|
|
class TestSchedulerTick:
|
|
"""Tests for _tick() lock acquisition and dispatch logic."""
|
|
|
|
def test_tick_acquires_lock(self, mocks):
|
|
collector, storage = mocks
|
|
storage.list_due_tasks.return_value = []
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
scheduler._tick()
|
|
|
|
storage.get_system_setting.assert_called()
|
|
storage.upsert_system_setting.assert_called()
|
|
storage.list_due_tasks.assert_called_once()
|
|
|
|
def test_tick_skips_when_locked(self, mocks):
|
|
collector, storage = mocks
|
|
# Another instance holds the lock (recent timestamp)
|
|
from datetime import UTC, datetime
|
|
|
|
now_str = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S")
|
|
_wire_lock_storage(
|
|
storage,
|
|
initial={"value": json.dumps({"owner": "other-instance", "acquired": now_str})},
|
|
)
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
scheduler._tick()
|
|
|
|
storage.list_due_tasks.assert_not_called()
|
|
|
|
def test_tick_takes_expired_lock(self, mocks):
|
|
"""An expired lock from another instance should be taken over."""
|
|
collector, storage = mocks
|
|
_wire_lock_storage(
|
|
storage,
|
|
initial={
|
|
"value": json.dumps({"owner": "other-instance", "acquired": "2020-01-01T00:00:00"})
|
|
},
|
|
)
|
|
storage.list_due_tasks.return_value = []
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
scheduler._tick()
|
|
|
|
storage.list_due_tasks.assert_called_once()
|
|
|
|
def test_dispatch_auto_mode(self, mocks):
|
|
collector, storage = mocks
|
|
|
|
task = _make_task(target_mode="auto")
|
|
storage.list_due_tasks.return_value = [task]
|
|
collector.get_nodes.return_value = ([_make_node()], 1)
|
|
collector.get_node_detail.return_value = {
|
|
"server_url": "http://node-001:8080",
|
|
}
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
with patch(
|
|
"turnstone.console.scheduler.TurnstoneServer.create_workstream",
|
|
return_value=_mock_create_response(),
|
|
) as mock_create:
|
|
scheduler._tick()
|
|
|
|
mock_create.assert_called_once()
|
|
storage.record_task_run.assert_called_once()
|
|
run_kwargs = storage.record_task_run.call_args[1]
|
|
assert run_kwargs["node_id"] == "node-001"
|
|
assert run_kwargs["status"] == "dispatched"
|
|
assert run_kwargs["ws_id"] == "ws_abc123"
|
|
|
|
def test_dispatch_passes_persona_and_project(self, mocks):
|
|
"""persona + project_id ride to create_workstream; created_by becomes
|
|
the user_id the node gates the project attach against."""
|
|
collector, storage = mocks
|
|
|
|
task = _make_task(persona="researcher", project_id="proj_42")
|
|
storage.list_due_tasks.return_value = [task]
|
|
collector.get_nodes.return_value = ([_make_node()], 1)
|
|
collector.get_node_detail.return_value = {"server_url": "http://node-001:8080"}
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
with patch(
|
|
"turnstone.console.scheduler.TurnstoneServer.create_workstream",
|
|
return_value=_mock_create_response(),
|
|
) as mock_create:
|
|
scheduler._tick()
|
|
|
|
mock_create.assert_called_once()
|
|
call_kwargs = mock_create.call_args[1]
|
|
assert call_kwargs["persona"] == "researcher"
|
|
assert call_kwargs["project_id"] == "proj_42"
|
|
assert call_kwargs["user_id"] == "u_admin"
|
|
|
|
def test_dispatch_defaults_persona_project_empty(self, mocks):
|
|
"""A task row without persona/project keys dispatches with empty
|
|
strings — the node then resolves the current kind default / no attach."""
|
|
collector, storage = mocks
|
|
|
|
task = _make_task()
|
|
task.pop("persona", None)
|
|
task.pop("project_id", None)
|
|
storage.list_due_tasks.return_value = [task]
|
|
collector.get_nodes.return_value = ([_make_node()], 1)
|
|
collector.get_node_detail.return_value = {"server_url": "http://node-001:8080"}
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
with patch(
|
|
"turnstone.console.scheduler.TurnstoneServer.create_workstream",
|
|
return_value=_mock_create_response(),
|
|
) as mock_create:
|
|
scheduler._tick()
|
|
|
|
call_kwargs = mock_create.call_args[1]
|
|
assert call_kwargs["persona"] == ""
|
|
assert call_kwargs["project_id"] == ""
|
|
|
|
def test_dispatch_pool_mode(self, mocks):
|
|
collector, storage = mocks
|
|
|
|
task = _make_task(target_mode="pool")
|
|
storage.list_due_tasks.return_value = [task]
|
|
collector.get_nodes.return_value = ([_make_node("node-001")], 1)
|
|
collector.get_node_detail.return_value = {
|
|
"server_url": "http://node-001:8080",
|
|
}
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
with patch(
|
|
"turnstone.console.scheduler.TurnstoneServer.create_workstream",
|
|
return_value=_mock_create_response(),
|
|
) as mock_create:
|
|
scheduler._tick()
|
|
|
|
mock_create.assert_called_once()
|
|
storage.record_task_run.assert_called_once()
|
|
|
|
def test_dispatch_all_mode(self, mocks):
|
|
collector, storage = mocks
|
|
|
|
task = _make_task(target_mode="all")
|
|
storage.list_due_tasks.return_value = [task]
|
|
collector.get_nodes.return_value = (
|
|
[_make_node("node-001"), _make_node("node-002")],
|
|
2,
|
|
)
|
|
collector.get_node_detail.side_effect = lambda nid: {
|
|
"server_url": f"http://{nid}:8080",
|
|
}
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
with patch(
|
|
"turnstone.console.scheduler.TurnstoneServer.create_workstream",
|
|
return_value=_mock_create_response(),
|
|
) as mock_create:
|
|
scheduler._tick()
|
|
|
|
assert mock_create.call_count == 2
|
|
assert storage.record_task_run.call_count == 2
|
|
|
|
def test_dispatch_specific_node(self, mocks):
|
|
collector, storage = mocks
|
|
|
|
task = _make_task(target_mode="node-001")
|
|
storage.list_due_tasks.return_value = [task]
|
|
collector.get_node_detail.return_value = {
|
|
"server_url": "http://node-001:8080",
|
|
}
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
with patch(
|
|
"turnstone.console.scheduler.TurnstoneServer.create_workstream",
|
|
return_value=_mock_create_response(),
|
|
) as mock_create:
|
|
scheduler._tick()
|
|
|
|
mock_create.assert_called_once()
|
|
run_kwargs = storage.record_task_run.call_args[1]
|
|
assert run_kwargs["node_id"] == "node-001"
|
|
|
|
def test_at_task_disables_after_dispatch(self, mocks):
|
|
collector, storage = mocks
|
|
|
|
task = _make_task(schedule_type="at", cron_expr="", at_time="2099-01-01T00:00:00")
|
|
storage.list_due_tasks.return_value = [task]
|
|
collector.get_nodes.return_value = ([_make_node()], 1)
|
|
collector.get_node_detail.return_value = {
|
|
"server_url": "http://node-001:8080",
|
|
}
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
with patch(
|
|
"turnstone.console.scheduler.TurnstoneServer.create_workstream",
|
|
return_value=_mock_create_response(),
|
|
):
|
|
scheduler._tick()
|
|
|
|
# At-task should be disabled after dispatch
|
|
update_calls = storage.update_scheduled_task.call_args_list
|
|
assert len(update_calls) == 1
|
|
args, kwargs = update_calls[0]
|
|
assert args[0] == "task_001"
|
|
assert kwargs["enabled"] is False
|
|
assert kwargs["next_run"] == ""
|
|
|
|
def test_cron_task_updates_next_run(self, mocks):
|
|
collector, storage = mocks
|
|
|
|
task = _make_task(schedule_type="cron", cron_expr="0 9 * * *")
|
|
storage.list_due_tasks.return_value = [task]
|
|
collector.get_nodes.return_value = ([_make_node()], 1)
|
|
collector.get_node_detail.return_value = {
|
|
"server_url": "http://node-001:8080",
|
|
}
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
with patch(
|
|
"turnstone.console.scheduler.TurnstoneServer.create_workstream",
|
|
return_value=_mock_create_response(),
|
|
):
|
|
scheduler._tick()
|
|
|
|
update_calls = storage.update_scheduled_task.call_args_list
|
|
assert len(update_calls) == 1
|
|
_, kwargs = update_calls[0]
|
|
assert kwargs["next_run"] != ""
|
|
assert "enabled" not in kwargs # cron tasks stay enabled
|
|
|
|
def test_no_reachable_nodes_records_failure(self, mocks):
|
|
collector, storage = mocks
|
|
|
|
task = _make_task(target_mode="auto")
|
|
storage.list_due_tasks.return_value = [task]
|
|
# No reachable nodes
|
|
collector.get_nodes.return_value = (
|
|
[_make_node("node-001", reachable=False)],
|
|
1,
|
|
)
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
scheduler._tick()
|
|
|
|
storage.record_task_run.assert_called_once()
|
|
run_kwargs = storage.record_task_run.call_args[1]
|
|
assert run_kwargs["status"] == "failed"
|
|
assert run_kwargs["error"] != ""
|
|
|
|
def test_failure_does_not_advance_schedule(self, mocks):
|
|
"""When dispatch fails, last_run/next_run should not be updated."""
|
|
collector, storage = mocks
|
|
|
|
task = _make_task(target_mode="auto")
|
|
storage.list_due_tasks.return_value = [task]
|
|
collector.get_nodes.return_value = ([], 0) # no nodes at all
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
scheduler._tick()
|
|
|
|
# update_scheduled_task should NOT be called (no last_run/next_run advance)
|
|
storage.update_scheduled_task.assert_not_called()
|
|
|
|
def test_fan_out_capped(self, mocks):
|
|
"""Fan-out 'all' mode should respect max_fan_out limit."""
|
|
collector, storage = mocks
|
|
|
|
task = _make_task(target_mode="all")
|
|
storage.list_due_tasks.return_value = [task]
|
|
# 10 reachable nodes but max_fan_out=3
|
|
nodes = [_make_node(f"node-{i:03d}") for i in range(10)]
|
|
collector.get_nodes.return_value = (nodes, 10)
|
|
collector.get_node_detail.side_effect = lambda nid: {
|
|
"server_url": f"http://{nid}:8080",
|
|
}
|
|
|
|
scheduler = TaskScheduler(collector, storage, max_fan_out=3)
|
|
with patch(
|
|
"turnstone.console.scheduler.TurnstoneServer.create_workstream",
|
|
return_value=_mock_create_response(),
|
|
) as mock_create:
|
|
scheduler._tick()
|
|
|
|
assert mock_create.call_count == 3
|
|
assert storage.record_task_run.call_count == 3
|
|
|
|
def test_specific_node_target(self, mocks):
|
|
"""Non-enum target_mode is treated as a specific node_id."""
|
|
collector, storage = mocks
|
|
|
|
task = _make_task(target_mode="node-custom-123")
|
|
storage.list_due_tasks.return_value = [task]
|
|
collector.get_node_detail.return_value = {
|
|
"server_url": "http://node-custom-123:8080",
|
|
}
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
with patch(
|
|
"turnstone.console.scheduler.TurnstoneServer.create_workstream",
|
|
return_value=_mock_create_response(),
|
|
) as mock_create:
|
|
scheduler._tick()
|
|
|
|
mock_create.assert_called_once()
|
|
run_kwargs = storage.record_task_run.call_args[1]
|
|
assert run_kwargs["node_id"] == "node-custom-123"
|
|
|
|
def test_user_id_in_dispatched_call(self, mocks):
|
|
"""Dispatched SDK call should include created_by as user_id."""
|
|
collector, storage = mocks
|
|
|
|
task = _make_task(target_mode="auto", created_by="u_scheduler_admin")
|
|
storage.list_due_tasks.return_value = [task]
|
|
collector.get_nodes.return_value = ([_make_node()], 1)
|
|
collector.get_node_detail.return_value = {
|
|
"server_url": "http://node-001:8080",
|
|
}
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
with patch(
|
|
"turnstone.console.scheduler.TurnstoneServer.create_workstream",
|
|
return_value=_mock_create_response(),
|
|
) as mock_create:
|
|
scheduler._tick()
|
|
|
|
_, kwargs = mock_create.call_args
|
|
assert kwargs["user_id"] == "u_scheduler_admin"
|
|
|
|
def test_sdk_failure_records_failure(self, mocks):
|
|
"""SDK errors during dispatch should record a failure."""
|
|
collector, storage = mocks
|
|
|
|
task = _make_task(target_mode="auto")
|
|
storage.list_due_tasks.return_value = [task]
|
|
collector.get_nodes.return_value = ([_make_node()], 1)
|
|
collector.get_node_detail.return_value = {
|
|
"server_url": "http://node-001:8080",
|
|
}
|
|
|
|
scheduler = TaskScheduler(collector, storage)
|
|
with patch(
|
|
"turnstone.console.scheduler.TurnstoneServer.create_workstream",
|
|
side_effect=TurnstoneAPIError(502, "Bad Gateway"),
|
|
):
|
|
scheduler._tick()
|
|
|
|
storage.record_task_run.assert_called_once()
|
|
run_kwargs = storage.record_task_run.call_args[1]
|
|
assert run_kwargs["status"] == "failed"
|