jiang
be6317eb48
Add project-bound OAuth MCP queries and controlled analysis with scoped evidence, limits, timeouts and a nearby-model spatial query using EPSG:4326 inputs. Default nearby searches to 200 meters and five results. Integrate investigation messaging/presentation, health/readiness and typed history; remove duplicate web search and fix online simulation and compressed baseline writes. Include project migrations 023-025; keep local credentials and demo seeding excluded. Validation: 319 focused pytest checks; real PostGIS temporary-table regression; read-only Jinqian bridge radius comparison; cross-language MCP discovery/calls.
89 lines
4.5 KiB
Python
89 lines
4.5 KiB
Python
"""Project scoped private message storage, sharing task transaction locks."""
|
|
from psycopg import AsyncConnection
|
|
|
|
from app.infra.db.postgresql.operational_tasks import OperationalTaskRepository
|
|
|
|
|
|
class AgentMessageRepository(OperationalTaskRepository):
|
|
def __init__(self, connection: AsyncConnection, project_id):
|
|
super().__init__(connection, project_id)
|
|
|
|
async def get_message(self, message_id: str) -> dict | None:
|
|
cursor = await self.connection.execute(
|
|
"SELECT * FROM workbench.agent_messages WHERE project_id=%s AND message_id=%s",
|
|
(self.project_id, message_id),
|
|
)
|
|
row = await cursor.fetchone()
|
|
return self._message(row) if row else None
|
|
|
|
async def get_message_request(self, owner_id: str, session_id: str, request_id: str) -> dict | None:
|
|
cursor = await self.connection.execute(
|
|
"""SELECT * FROM workbench.agent_messages WHERE project_id=%s
|
|
AND source_owner_id=%s AND source_session_id=%s AND request_id=%s""",
|
|
(self.project_id, owner_id, session_id, request_id),
|
|
)
|
|
row = await cursor.fetchone()
|
|
return self._message(row) if row else None
|
|
|
|
@staticmethod
|
|
def _message(row) -> dict:
|
|
return dict(row)
|
|
|
|
async def insert_message(self, message: dict) -> None:
|
|
keys = [key for key in message if key != "project_id"]
|
|
await self.connection.execute(
|
|
f"INSERT INTO workbench.agent_messages (project_id,{','.join(keys)}) "
|
|
f"VALUES ({','.join(['%s'] * (len(keys) + 1))})",
|
|
(self.project_id, *(message[key] for key in keys)),
|
|
)
|
|
|
|
async def list_session_messages(
|
|
self, owner_id: str, session_id: str, limit: int, *, cursor_at=None, cursor_id=None, pending_only=False,
|
|
) -> list[dict]:
|
|
cursor = await self.connection.execute(
|
|
"""SELECT * FROM workbench.agent_messages WHERE project_id=%s
|
|
AND conversation_owner_id=%s AND (source_session_id=%s OR
|
|
(direction='to_chat' AND target_session_id=%s))
|
|
AND (%s::timestamptz IS NULL OR (created_at,message_id)<(%s::timestamptz,%s::uuid))
|
|
AND (NOT %s OR status IN ('pending','enqueued'))
|
|
ORDER BY created_at DESC,message_id DESC LIMIT %s""",
|
|
(self.project_id, owner_id, session_id, session_id, cursor_at, cursor_at, cursor_id, pending_only, limit),
|
|
)
|
|
return [self._message(row) for row in await cursor.fetchall()]
|
|
|
|
async def list_run_session_messages(
|
|
self, session_id: str, limit: int, *, cursor_at=None, cursor_id=None, pending_only=False,
|
|
) -> list[dict]:
|
|
cursor = await self.connection.execute(
|
|
"""SELECT m.* FROM workbench.agent_messages m
|
|
JOIN workbench.operational_task_runs r ON r.project_id=m.project_id AND r.run_id=m.run_id
|
|
WHERE m.project_id=%s AND r.agent_session_id=%s
|
|
AND r.status IN ('claimed','running')
|
|
AND (%s::timestamptz IS NULL OR (m.created_at,m.message_id)<(%s::timestamptz,%s::uuid))
|
|
AND (NOT %s OR (m.direction='to_task' AND m.status IN ('pending','enqueued')))
|
|
ORDER BY m.created_at DESC,m.message_id DESC LIMIT %s""",
|
|
(self.project_id, session_id, cursor_at, cursor_at, cursor_id, pending_only, limit),
|
|
)
|
|
return [self._message(row) for row in await cursor.fetchall()]
|
|
|
|
async def list_pending_messages(self, limit: int, *, cursor_at=None, cursor_id=None) -> list[dict]:
|
|
cursor = await self.connection.execute(
|
|
"""SELECT m.* FROM workbench.agent_messages m
|
|
JOIN workbench.operational_task_runs r
|
|
ON r.project_id=m.project_id AND r.run_id=m.run_id
|
|
WHERE m.project_id=%s AND m.status IN ('pending','enqueued')
|
|
AND (m.direction='to_chat' OR r.status IN ('claimed','running'))
|
|
AND (%s::timestamptz IS NULL OR (m.created_at,m.message_id)>(%s::timestamptz,%s::uuid))
|
|
ORDER BY m.created_at,m.message_id LIMIT %s""",
|
|
(self.project_id, cursor_at, cursor_at, cursor_id, limit),
|
|
)
|
|
return [self._message(row) for row in await cursor.fetchall()]
|
|
|
|
async def update_message_delivery(self, message_id: str, update: dict) -> None:
|
|
keys = list(update)
|
|
await self.connection.execute(
|
|
f"UPDATE workbench.agent_messages SET {','.join(key+'=%s' for key in keys)} "
|
|
"WHERE project_id=%s AND message_id=%s",
|
|
(*(update[key] for key in keys), self.project_id, message_id),
|
|
)
|