Persist API chat lifetime and add cleanup job
Validate and persist API chat lifetime: lifetime_hours is validated as a positive number and stored in the AgentContext data, and context.last_message is set using UTC. Removed the in-class threading-based cleanup state and old _cleanup_expired_chats method. Introduced a new job-loop extension (extensions/python/job_loop/_20_cleanup_expired_api_chats.py) that periodically scans AgentContext instances and removes expired API chats (using persist_chat.remove_chat) in a UTC-aware manner. Added tests (tests/test_api_chat_lifetime.py) to verify lifetime persistence and that the job loop removes expired chats.
frdel committed
May 11, 2026 at 08:49 UTC
904a0f4a258c6e6f83fa96739eec1c7afcc3dbac
3 files changed
+161
-33
api/api_message.py
+14
-33
@@ -1,7 +1,7 @@
1
import base64
2
import os
3
import uuid
4
-from datetime import datetime, timedelta
4
+from datetime import datetime, timezone
5
from agent import AgentContext, UserMessage, AgentContextType
6
from helpers.api import ApiHandler, Request, Response
7
from helpers import files, projects
@@ -9,14 +9,9 @@ from helpers.print_style import PrintStyle
9
from helpers.projects import activate_project
10
from helpers.security import safe_filename
11
from initialize import initialize_agent
12
-import threading
12
13
14
class ApiMessage(ApiHandler):
16
- # Track chat lifetimes for cleanup
17
- _chat_lifetimes = {}
18
- _cleanup_lock = threading.Lock()
19
-
15
@classmethod
16
def requires_auth(cls) -> bool:
17
return False # No web auth required
@@ -37,6 +32,16 @@ class ApiMessage(ApiHandler):
32
lifetime_hours = input.get("lifetime_hours", 24) # Default 24 hours
33
project_name = input.get("project_name", None)
34
agent_profile = input.get("agent_profile", None)
35
+ try:
36
+ lifetime_hours = float(lifetime_hours)
37
+ if lifetime_hours <= 0:
38
+ raise ValueError("lifetime_hours must be greater than 0")
39
+ except (TypeError, ValueError):
40
+ return Response(
41
+ '{"error": "lifetime_hours must be a positive number"}',
42
+ status=400,
43
+ mimetype="application/json",
44
+ )
45
46
# Set an agent if profile provided
47
override_settings = {}
@@ -116,9 +121,9 @@ class ApiMessage(ApiHandler):
121
except Exception as e:
122
return Response(f'{{"error": "Failed to activate project: {str(e)}"}}', status=400, mimetype="application/json")
123
119
- # Update chat lifetime
120
- with self._cleanup_lock:
121
- self._chat_lifetimes[context_id] = datetime.now() + timedelta(hours=lifetime_hours)
124
+ # Persist API chat lifetime in context data so cleanup survives restarts.
125
+ context.set_data("lifetime_hours", lifetime_hours)
126
+ context.last_message = datetime.now(timezone.utc)
127
128
# Process message
129
try:
@@ -148,9 +153,6 @@ class ApiMessage(ApiHandler):
153
task = context.communicate(UserMessage(message=message, attachments=attachment_paths, id=msg_id))
154
result = await task.result()
155
151
- # Clean up expired chats
152
- self._cleanup_expired_chats()
153
-
156
return {
157
"context_id": context_id,
158
"response": result
@@ -159,24 +161,3 @@ class ApiMessage(ApiHandler):
161
except Exception as e:
162
PrintStyle.error(f"External API error: {e}")
163
return Response(f'{{"error": "{str(e)}"}}', status=500, mimetype="application/json")
162
-
163
- @classmethod
164
- def _cleanup_expired_chats(cls):
165
- """Clean up expired chats"""
166
- with cls._cleanup_lock:
167
- now = datetime.now()
168
- expired_contexts = [
169
- context_id for context_id, expiry in cls._chat_lifetimes.items()
170
- if now > expiry
171
- ]
172
-
173
- for context_id in expired_contexts:
174
- try:
175
- context = AgentContext.get(context_id)
176
- if context:
177
- context.reset()
178
- AgentContext.remove(context_id)
179
- del cls._chat_lifetimes[context_id]
180
- PrintStyle().print(f"Cleaned up expired chat: {context_id}")
181
- except Exception as e:
182
- PrintStyle.error(f"Failed to cleanup chat {context_id}: {e}")
extensions/python/job_loop/_20_cleanup_expired_api_chats.py
new
+63
@@ -0,0 +1,63 @@
1
+from datetime import datetime, timedelta, timezone
2
+from typing import Any
3
+
4
+from agent import AgentContext
5
+from helpers import persist_chat
6
+from helpers.extension import Extension
7
+from helpers.print_style import PrintStyle
8
+from helpers.state_monitor_integration import mark_dirty_all
9
+
10
+
11
+CHECK_INTERVAL = timedelta(hours=1)
12
+LIFETIME_KEY = "lifetime_hours"
13
+
14
+
15
+class CleanupExpiredApiChats(Extension):
16
+ _last_check: datetime | None = None
17
+
18
+ async def execute(self, data: dict[str, Any] | None = None, **kwargs):
19
+ now = datetime.now(timezone.utc)
20
+ if type(self)._last_check and now - type(self)._last_check < CHECK_INTERVAL:
21
+ return
22
+ type(self)._last_check = now
23
+
24
+ removed = 0
25
+ for context in list(AgentContext.all()):
26
+ lifetime_hours = context.get_data(LIFETIME_KEY)
27
+ if lifetime_hours is None:
28
+ continue
29
+
30
+ try:
31
+ lifetime = timedelta(hours=float(lifetime_hours))
32
+ except (TypeError, ValueError):
33
+ PrintStyle.error(
34
+ f"Invalid chat lifetime for {context.id}: {lifetime_hours}"
35
+ )
36
+ continue
37
+
38
+ if lifetime <= timedelta(0) or context.is_running():
39
+ continue
40
+
41
+ last_message = _as_utc(context.last_message)
42
+ if now - last_message <= lifetime:
43
+ continue
44
+
45
+ try:
46
+ context.reset()
47
+ AgentContext.remove(context.id)
48
+ persist_chat.remove_chat(context.id)
49
+ removed += 1
50
+ PrintStyle().print(f"Cleaned up expired API chat: {context.id}")
51
+ except Exception as e:
52
+ PrintStyle.error(f"Failed to cleanup expired API chat {context.id}: {e}")
53
+
54
+ if removed:
55
+ mark_dirty_all(reason="job_loop.CleanupExpiredApiChats")
56
+
57
+
58
+def _as_utc(value: datetime | None) -> datetime:
59
+ if value is None:
60
+ return datetime.fromtimestamp(0, timezone.utc)
61
+ if value.tzinfo is None:
62
+ return value.replace(tzinfo=timezone.utc)
63
+ return value.astimezone(timezone.utc)
tests/test_api_chat_lifetime.py
new
+84
@@ -0,0 +1,84 @@
1
+from datetime import datetime, timedelta, timezone
2
+import json
3
+from pathlib import Path
4
+import sys
5
+import threading
6
+
7
+import pytest
8
+
9
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
10
+if str(PROJECT_ROOT) not in sys.path:
11
+ sys.path.insert(0, str(PROJECT_ROOT))
12
+
13
+from agent import AgentContext
14
+from initialize import initialize_agent
15
+
16
+
17
+class _CompletedTask:
18
+ async def result(self):
19
+ return "ok"
20
+
21
+
22
+@pytest.mark.asyncio
23
+async def test_api_message_persists_lifetime_hours_in_context_data(monkeypatch):
24
+ from api.api_message import ApiMessage
25
+ from helpers import persist_chat
26
+
27
+ monkeypatch.setattr(AgentContext, "communicate", lambda self, msg: _CompletedTask())
28
+
29
+ handler = ApiMessage(app=None, thread_lock=threading.RLock()) # type: ignore[arg-type]
30
+ output = await handler.process(
31
+ {
32
+ "message": "hello",
33
+ "lifetime_hours": 1,
34
+ },
35
+ request=None, # type: ignore[arg-type]
36
+ )
37
+
38
+ context_id = output["context_id"] # type: ignore[index]
39
+ context = AgentContext.get(context_id)
40
+ restored = None
41
+ try:
42
+ assert context is not None
43
+ assert context.get_data("lifetime_hours") == 1.0
44
+
45
+ serialized = json.loads(persist_chat.export_json_chat(context))
46
+ assert serialized["data"]["lifetime_hours"] == 1.0
47
+
48
+ AgentContext.remove(context_id)
49
+ restored = persist_chat._deserialize_context(serialized)
50
+ assert restored.get_data("lifetime_hours") == 1.0
51
+ finally:
52
+ AgentContext.remove(context_id)
53
+ if restored:
54
+ AgentContext.remove(restored.id)
55
+
56
+
57
+@pytest.mark.asyncio
58
+async def test_job_loop_removes_expired_lifetime_chat(monkeypatch):
59
+ from extensions.python.job_loop._20_cleanup_expired_api_chats import (
60
+ CleanupExpiredApiChats,
61
+ )
62
+ import extensions.python.job_loop._20_cleanup_expired_api_chats as cleanup_module
63
+
64
+ removed_chats = []
65
+ dirty_reasons = []
66
+ monkeypatch.setattr(cleanup_module.persist_chat, "remove_chat", removed_chats.append)
67
+ monkeypatch.setattr(
68
+ cleanup_module,
69
+ "mark_dirty_all",
70
+ lambda reason: dirty_reasons.append(reason),
71
+ )
72
+
73
+ context = AgentContext(
74
+ config=initialize_agent(),
75
+ last_message=datetime.now(timezone.utc) - timedelta(hours=2),
76
+ )
77
+ context.set_data("lifetime_hours", 1)
78
+ CleanupExpiredApiChats._last_check = None
79
+
80
+ await CleanupExpiredApiChats(agent=None).execute()
81
+
82
+ assert AgentContext.get(context.id) is None
83
+ assert removed_chats == [context.id]
84
+ assert dirty_reasons == ["job_loop.CleanupExpiredApiChats"]