diff --git a/litellm/proxy/scheduled_tasks/store.py b/litellm/proxy/scheduled_tasks/store.py index f88b239ac3c..382f469c8fe 100644 --- a/litellm/proxy/scheduled_tasks/store.py +++ b/litellm/proxy/scheduled_tasks/store.py @@ -8,11 +8,26 @@ is required for multi-pod tick safety. from __future__ import annotations +import json from datetime import datetime, timedelta, timezone from typing import Any, Dict, List, Optional from litellm.proxy.scheduled_tasks.schedule import compute_next_run + +def _serialize_json_for_prisma(value: Any) -> Optional[str]: + """ + Encode a value bound for a Prisma `Json?` column. prisma-client-python + rejects raw Python dicts/lists on Json columns + (`MissingRequiredValueError` / `DataError`); always json.dumps so the + driver hands a string to Postgres jsonb. Read path round-trips back + to native Python. + """ + if value is None: + return None + return json.dumps(value) + + # Whitelist of fields PATCH can touch. Anything outside this set would let # a buggy caller rewrite scheduling state (status, fired flags, ...). UPDATABLE_FIELDS = frozenset( @@ -62,13 +77,15 @@ async def create_task( ) -> Any: return await prisma_client.db.litellm_scheduledtasktable.create( data={ - "owner_token": owner_token, + # Prisma exposes the FK via the relation field (`owner_key`), + # not the scalar column. Use connect to satisfy the input type. + "owner_key": {"connect": {"token": owner_token}}, "user_id": user_id, "team_id": team_id, "agent_id": agent_id, "title": title, "action": action, - "action_args": action_args, + "action_args": _serialize_json_for_prisma(action_args), "check_prompt": check_prompt, "format_prompt": format_prompt, "schedule_kind": schedule_kind, @@ -133,6 +150,15 @@ async def update_task_for_owner( ) if existing is None: return None + + # Mirror the create path: Json columns must be json-encoded for + # prisma-client-python. + if "action_args" in fields: + fields = { + **fields, + "action_args": _serialize_json_for_prisma(fields["action_args"]), + } + return await prisma_client.db.litellm_scheduledtasktable.update( where={"task_id": task_id}, data=fields, diff --git a/tests/test_litellm/proxy/scheduled_tasks/test_endpoints.py b/tests/test_litellm/proxy/scheduled_tasks/test_endpoints.py index d3a735fda17..75b03f92cf1 100644 --- a/tests/test_litellm/proxy/scheduled_tasks/test_endpoints.py +++ b/tests/test_litellm/proxy/scheduled_tasks/test_endpoints.py @@ -4,6 +4,7 @@ prisma scheduled-tasks table. The /due test path uses a fake of query_raw + tx() that exercises the same Python-side schedule advance logic. """ +import json import os import sys from datetime import datetime, timedelta, timezone @@ -80,8 +81,23 @@ class _FakeScheduledTaskTable: return [r for r in self.rows if self._matches(r, where)] async def create(self, data: Dict[str, Any]) -> MagicMock: + # 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. + 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 self._counter += 1 - row = _make_row(task_id=f"task-{self._counter}", **data) + row = _make_row(task_id=f"task-{self._counter}", **normalized) self.rows.append(row) return row @@ -106,6 +122,11 @@ 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): + try: + v = json.loads(v) + except ValueError: + pass setattr(r, k, v) return r raise Exception("Not found")