diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260429000000_add_scheduled_tasks/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260429000000_add_scheduled_tasks/migration.sql index 9a564b4a05b..489814e0d7a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260429000000_add_scheduled_tasks/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260429000000_add_scheduled_tasks/migration.sql @@ -10,6 +10,7 @@ CREATE TABLE "LiteLLM_ScheduledTaskTable" ( "action_args" JSONB, "check_prompt" TEXT, "format_prompt" TEXT, + "metadata" JSONB, "schedule_kind" TEXT NOT NULL, "schedule_spec" TEXT NOT NULL, "schedule_tz" TEXT, diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 24334fc04c1..8137c616b4b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1309,6 +1309,7 @@ model LiteLLM_ScheduledTaskTable { action_args Json? check_prompt String? format_prompt String? + metadata Json? schedule_kind String schedule_spec String diff --git a/litellm/proxy/scheduled_tasks/endpoints.py b/litellm/proxy/scheduled_tasks/endpoints.py index fe0506cd1cd..5c2e17ae14f 100644 --- a/litellm/proxy/scheduled_tasks/endpoints.py +++ b/litellm/proxy/scheduled_tasks/endpoints.py @@ -57,6 +57,7 @@ def _row_to_response(row) -> ScheduledTaskResponse: action_args=row.action_args, check_prompt=row.check_prompt, format_prompt=row.format_prompt, + metadata=getattr(row, "metadata", None), schedule_kind=row.schedule_kind, schedule_spec=row.schedule_spec, schedule_tz=row.schedule_tz, @@ -136,6 +137,7 @@ async def create_scheduled_task( action_args=data.action_args, check_prompt=data.check_prompt, format_prompt=data.format_prompt, + metadata=data.metadata, schedule_kind=data.schedule_kind, schedule_spec=data.schedule_spec, schedule_tz=data.schedule_tz, @@ -207,6 +209,7 @@ async def get_due_tasks( action_args=r["action_args"], check_prompt=r["check_prompt"], format_prompt=r["format_prompt"], + metadata=r.get("metadata"), scheduled_for=r["scheduled_for"], ) for r in rows diff --git a/litellm/proxy/scheduled_tasks/store.py b/litellm/proxy/scheduled_tasks/store.py index 217984da6e0..45314195ce8 100644 --- a/litellm/proxy/scheduled_tasks/store.py +++ b/litellm/proxy/scheduled_tasks/store.py @@ -40,9 +40,14 @@ UPDATABLE_FIELDS = frozenset( "action", "action_args", "format_prompt", + "metadata", } ) +# Json? columns that go through Prisma — must be JSON-encoded if present, +# omitted entirely if None. +JSON_FIELDS = frozenset({"action_args", "metadata"}) + MAX_ACTIVE_TASKS_PER_KEY = 10 TERMINAL_STATUSES = ("fired", "expired", "cancelled") @@ -65,6 +70,7 @@ async def create_task( action_args: Optional[Dict[str, Any]], check_prompt: Optional[str], format_prompt: Optional[str], + metadata: Optional[Any], schedule_kind: str, schedule_spec: str, schedule_tz: Optional[str], @@ -90,6 +96,8 @@ async def create_task( } if action_args is not None: data["action_args"] = _encode_json(action_args) + if metadata is not None: + data["metadata"] = _encode_json(metadata) return await prisma_client.db.litellm_scheduledtasktable.create(data=data) @@ -149,13 +157,17 @@ async def update_task_for_owner( # Mirror the create path: Json columns must be json-encoded for # prisma-client-python. Drop the key entirely when caller sent null; # passing None to a Json? column is rejected. - if "action_args" in fields: - if fields["action_args"] is None: - fields = {k: v for k, v in fields.items() if k != "action_args"} - if not fields: - raise ValueError("no fields to update") + cleaned: Dict[str, Any] = {} + for k, v in fields.items(): + if k in JSON_FIELDS: + if v is None: + continue # skip — Prisma rejects None on Json? + cleaned[k] = _encode_json(v) else: - fields = {**fields, "action_args": _encode_json(fields["action_args"])} + cleaned[k] = v + if not cleaned: + raise ValueError("no fields to update") + fields = cleaned return await prisma_client.db.litellm_scheduledtasktable.update( where={"task_id": task_id}, @@ -205,7 +217,8 @@ async def claim_due( """ SELECT task_id, schedule_kind, schedule_spec, schedule_tz, fire_once, expires_at, next_run_at, - action, action_args, check_prompt, format_prompt, title + action, action_args, check_prompt, format_prompt, + metadata, title FROM "LiteLLM_ScheduledTaskTable" WHERE status = 'pending' AND next_run_at <= now() @@ -264,6 +277,7 @@ async def claim_due( "action_args": r["action_args"], "check_prompt": r["check_prompt"], "format_prompt": r["format_prompt"], + "metadata": r["metadata"], "scheduled_for": r["next_run_at"], } for r in rows diff --git a/litellm/proxy/scheduled_tasks/types.py b/litellm/proxy/scheduled_tasks/types.py index dcb60e2641e..e415ccc17d6 100644 --- a/litellm/proxy/scheduled_tasks/types.py +++ b/litellm/proxy/scheduled_tasks/types.py @@ -13,6 +13,7 @@ class CreateScheduledTaskRequest(BaseModel): action_args: Optional[Dict[str, Any]] = None check_prompt: Optional[str] = None format_prompt: Optional[str] = None + metadata: Optional[Any] = None schedule_kind: ScheduleKind schedule_spec: str @@ -34,6 +35,7 @@ class UpdateScheduledTaskRequest(BaseModel): action: Optional[str] = None action_args: Optional[Dict[str, Any]] = None format_prompt: Optional[str] = None + metadata: Optional[Any] = None class ScheduledTaskResponse(BaseModel): @@ -48,6 +50,7 @@ class ScheduledTaskResponse(BaseModel): action_args: Optional[Dict[str, Any]] check_prompt: Optional[str] format_prompt: Optional[str] + metadata: Optional[Any] schedule_kind: str schedule_spec: str @@ -73,6 +76,7 @@ class DueTaskResponse(BaseModel): action_args: Optional[Dict[str, Any]] check_prompt: Optional[str] format_prompt: Optional[str] + metadata: Optional[Any] scheduled_for: datetime diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 24334fc04c1..8137c616b4b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1309,6 +1309,7 @@ model LiteLLM_ScheduledTaskTable { action_args Json? check_prompt String? format_prompt String? + metadata Json? schedule_kind String schedule_spec String diff --git a/schema.prisma b/schema.prisma index 24334fc04c1..8137c616b4b 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1309,6 +1309,7 @@ model LiteLLM_ScheduledTaskTable { action_args Json? check_prompt String? format_prompt String? + metadata Json? schedule_kind String schedule_spec String diff --git a/tests/test_litellm/proxy/scheduled_tasks/test_endpoints.py b/tests/test_litellm/proxy/scheduled_tasks/test_endpoints.py index 75b03f92cf1..31161b3ecb0 100644 --- a/tests/test_litellm/proxy/scheduled_tasks/test_endpoints.py +++ b/tests/test_litellm/proxy/scheduled_tasks/test_endpoints.py @@ -44,6 +44,7 @@ def _make_row(**kwargs) -> MagicMock: "action_args": None, "check_prompt": "is it done?", "format_prompt": None, + "metadata": None, "schedule_kind": "interval", "schedule_spec": "5m", "schedule_tz": None, @@ -84,18 +85,19 @@ class _FakeScheduledTaskTable: # Mirror real Prisma write semantics: # - relation field `owner_key: {connect: {token: X}}` becomes the # scalar `owner_token` column on the row; - # - Json columns (`action_args`) come in as JSON strings and are - # deserialized back to Python on read. + # - Json? columns (`action_args`, `metadata`) come in as JSON + # strings and are deserialised back to Python on read. normalized = dict(data) if "owner_key" in normalized: connect = normalized.pop("owner_key", {}).get("connect", {}) if "token" in connect: normalized["owner_token"] = connect["token"] - if isinstance(normalized.get("action_args"), str): - try: - normalized["action_args"] = json.loads(normalized["action_args"]) - except ValueError: - pass + for json_field in ("action_args", "metadata"): + if isinstance(normalized.get(json_field), str): + try: + normalized[json_field] = json.loads(normalized[json_field]) + except ValueError: + pass self._counter += 1 row = _make_row(task_id=f"task-{self._counter}", **normalized) self.rows.append(row) @@ -122,7 +124,7 @@ class _FakeScheduledTaskTable: for r in self.rows: if r.task_id == where["task_id"]: for k, v in data.items(): - if k == "action_args" and isinstance(v, str): + if k in ("action_args", "metadata") and isinstance(v, str): try: v = json.loads(v) except ValueError: @@ -193,6 +195,7 @@ class _FakeTx: "action_args": r.action_args, "check_prompt": r.check_prompt, "format_prompt": r.format_prompt, + "metadata": r.metadata, "title": r.title, } ) @@ -340,6 +343,20 @@ class TestCreate: r = self.client.post("/v1/tasks", json=_create_payload(title="overflow")) assert r.status_code == 429 + def test_metadata_roundtrips(self): + payload = _create_payload(metadata={"agent_session": "abc", "tags": [1, 2, 3]}) + with _patch_prisma(self.prisma): + r = self.client.post("/v1/tasks", json=payload) + assert r.status_code == 200, r.text + body = r.json() + assert body["metadata"] == {"agent_session": "abc", "tags": [1, 2, 3]} + + def test_metadata_omitted_when_not_supplied(self): + with _patch_prisma(self.prisma): + r = self.client.post("/v1/tasks", json=_create_payload()) + assert r.status_code == 200, r.text + assert r.json()["metadata"] is None + def test_expires_in_past_rejected(self): payload = _create_payload( expires_at=(_now() - timedelta(seconds=1)).isoformat(),