mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-10 15:56:03 +00:00
fix(agent-runner): harden state and event APIs
This commit is contained in:
huanghuoguoguo
@@ -12,6 +12,9 @@ from datetime import datetime
|
||||
import sqlalchemy
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
from sqlalchemy import select, delete, update
|
||||
from sqlalchemy.dialects.postgresql import insert as postgresql_insert
|
||||
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
|
||||
from .descriptor import AgentRunnerDescriptor
|
||||
from .host_models import AgentEventEnvelope, AgentBinding
|
||||
@@ -87,6 +90,49 @@ class PersistentStateStore:
|
||||
|
||||
return json_str, None
|
||||
|
||||
async def _upsert_state_row(
|
||||
self,
|
||||
conn: typing.Any,
|
||||
values: dict[str, typing.Any],
|
||||
) -> None:
|
||||
"""Insert or update a state row by the logical scope/key identity."""
|
||||
update_values = {
|
||||
'value_json': values['value_json'],
|
||||
'updated_at': values['updated_at'],
|
||||
}
|
||||
constraint_columns = ['scope_key', 'state_key']
|
||||
dialect_name = self._db_engine.dialect.name
|
||||
|
||||
if dialect_name == 'sqlite':
|
||||
stmt = sqlite_insert(AgentRunnerState).values(**values)
|
||||
await conn.execute(
|
||||
stmt.on_conflict_do_update(
|
||||
index_elements=constraint_columns,
|
||||
set_=update_values,
|
||||
)
|
||||
)
|
||||
return
|
||||
|
||||
if dialect_name == 'postgresql':
|
||||
stmt = postgresql_insert(AgentRunnerState).values(**values)
|
||||
await conn.execute(
|
||||
stmt.on_conflict_do_update(
|
||||
index_elements=constraint_columns,
|
||||
set_=update_values,
|
||||
)
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
await conn.execute(sqlalchemy.insert(AgentRunnerState).values(**values))
|
||||
except IntegrityError:
|
||||
await conn.execute(
|
||||
update(AgentRunnerState)
|
||||
.where(AgentRunnerState.scope_key == values['scope_key'])
|
||||
.where(AgentRunnerState.state_key == values['state_key'])
|
||||
.values(**update_values)
|
||||
)
|
||||
|
||||
# ========== Async DB Operations ==========
|
||||
|
||||
async def build_snapshot_from_event(
|
||||
@@ -195,49 +241,29 @@ class PersistentStateStore:
|
||||
# Build context fields
|
||||
binding_identity = get_binding_identity(binding)
|
||||
|
||||
now = datetime.utcnow()
|
||||
async with self._db_engine.begin() as conn:
|
||||
# Check if entry exists
|
||||
result = await conn.execute(
|
||||
select(AgentRunnerState.id)
|
||||
.where(AgentRunnerState.scope_key == scope_key)
|
||||
.where(AgentRunnerState.state_key == key)
|
||||
await self._upsert_state_row(
|
||||
conn,
|
||||
{
|
||||
'runner_id': descriptor.id,
|
||||
'binding_identity': binding_identity,
|
||||
'scope': scope,
|
||||
'scope_key': scope_key,
|
||||
'state_key': key,
|
||||
'value_json': value_json,
|
||||
'bot_id': event.bot_id,
|
||||
'workspace_id': event.workspace_id,
|
||||
'conversation_id': event.conversation_id,
|
||||
'thread_id': event.thread_id,
|
||||
'actor_type': event.actor.actor_type if event.actor else None,
|
||||
'actor_id': event.actor.actor_id if event.actor else None,
|
||||
'subject_type': event.subject.subject_type if event.subject else None,
|
||||
'subject_id': event.subject.subject_id if event.subject else None,
|
||||
'created_at': now,
|
||||
'updated_at': now,
|
||||
},
|
||||
)
|
||||
existing = result.first()
|
||||
|
||||
now = datetime.utcnow()
|
||||
|
||||
if existing:
|
||||
# Update existing entry
|
||||
await conn.execute(
|
||||
update(AgentRunnerState)
|
||||
.where(AgentRunnerState.id == existing.id)
|
||||
.values(
|
||||
value_json=value_json,
|
||||
updated_at=now,
|
||||
)
|
||||
)
|
||||
else:
|
||||
# Insert new entry
|
||||
await conn.execute(
|
||||
sqlalchemy.insert(AgentRunnerState).values(
|
||||
runner_id=descriptor.id,
|
||||
binding_identity=binding_identity,
|
||||
scope=scope,
|
||||
scope_key=scope_key,
|
||||
state_key=key,
|
||||
value_json=value_json,
|
||||
bot_id=event.bot_id,
|
||||
workspace_id=event.workspace_id,
|
||||
conversation_id=event.conversation_id,
|
||||
thread_id=event.thread_id,
|
||||
actor_type=event.actor.actor_type if event.actor else None,
|
||||
actor_id=event.actor.actor_id if event.actor else None,
|
||||
subject_type=event.subject.subject_type if event.subject else None,
|
||||
subject_id=event.subject.subject_id if event.subject else None,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
)
|
||||
|
||||
return True, None
|
||||
|
||||
@@ -293,49 +319,29 @@ class PersistentStateStore:
|
||||
|
||||
context = context or {}
|
||||
|
||||
now = datetime.utcnow()
|
||||
async with self._db_engine.begin() as conn:
|
||||
# Check if entry exists
|
||||
result = await conn.execute(
|
||||
select(AgentRunnerState.id)
|
||||
.where(AgentRunnerState.scope_key == scope_key)
|
||||
.where(AgentRunnerState.state_key == state_key)
|
||||
await self._upsert_state_row(
|
||||
conn,
|
||||
{
|
||||
'runner_id': runner_id,
|
||||
'binding_identity': binding_identity,
|
||||
'scope': scope,
|
||||
'scope_key': scope_key,
|
||||
'state_key': state_key,
|
||||
'value_json': value_json,
|
||||
'bot_id': context.get('bot_id'),
|
||||
'workspace_id': context.get('workspace_id'),
|
||||
'conversation_id': context.get('conversation_id'),
|
||||
'thread_id': context.get('thread_id'),
|
||||
'actor_type': context.get('actor_type'),
|
||||
'actor_id': context.get('actor_id'),
|
||||
'subject_type': context.get('subject_type'),
|
||||
'subject_id': context.get('subject_id'),
|
||||
'created_at': now,
|
||||
'updated_at': now,
|
||||
},
|
||||
)
|
||||
existing = result.first()
|
||||
|
||||
now = datetime.utcnow()
|
||||
|
||||
if existing:
|
||||
# Update existing entry
|
||||
await conn.execute(
|
||||
update(AgentRunnerState)
|
||||
.where(AgentRunnerState.id == existing.id)
|
||||
.values(
|
||||
value_json=value_json,
|
||||
updated_at=now,
|
||||
)
|
||||
)
|
||||
else:
|
||||
# Insert new entry
|
||||
await conn.execute(
|
||||
sqlalchemy.insert(AgentRunnerState).values(
|
||||
runner_id=runner_id,
|
||||
binding_identity=binding_identity,
|
||||
scope=scope,
|
||||
scope_key=scope_key,
|
||||
state_key=state_key,
|
||||
value_json=value_json,
|
||||
bot_id=context.get('bot_id'),
|
||||
workspace_id=context.get('workspace_id'),
|
||||
conversation_id=context.get('conversation_id'),
|
||||
thread_id=context.get('thread_id'),
|
||||
actor_type=context.get('actor_type'),
|
||||
actor_id=context.get('actor_id'),
|
||||
subject_type=context.get('subject_type'),
|
||||
subject_id=context.get('subject_id'),
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
)
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
@@ -243,6 +243,33 @@ def _resolve_run_conversation(
|
||||
return session_conversation_id, None
|
||||
|
||||
|
||||
def _project_event_record_for_api(event: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Project EventLogStore rows onto the SDK AgentEventRecord DTO."""
|
||||
seq = event.get('seq') or event.get('id')
|
||||
return {
|
||||
'event_id': event.get('event_id'),
|
||||
'event_type': event.get('event_type'),
|
||||
'event_time': event.get('event_time'),
|
||||
'source': event.get('source'),
|
||||
'bot_id': event.get('bot_id'),
|
||||
'workspace_id': event.get('workspace_id'),
|
||||
'conversation_id': event.get('conversation_id'),
|
||||
'thread_id': event.get('thread_id'),
|
||||
'actor_type': event.get('actor_type'),
|
||||
'actor_id': event.get('actor_id'),
|
||||
'actor_name': event.get('actor_name'),
|
||||
'subject_type': event.get('subject_type'),
|
||||
'subject_id': event.get('subject_id'),
|
||||
'input_summary': event.get('input_summary'),
|
||||
'input_ref': event.get('input_ref'),
|
||||
'raw_ref': event.get('raw_ref'),
|
||||
'seq': seq,
|
||||
'cursor': event.get('cursor') or (str(seq) if seq is not None else None),
|
||||
'created_at': event.get('created_at'),
|
||||
'metadata': event.get('metadata') or {},
|
||||
}
|
||||
|
||||
|
||||
def _normalize_uuid_list(values: Any) -> list[str]:
|
||||
"""Normalize a user/config supplied UUID list while preserving order."""
|
||||
if not isinstance(values, list):
|
||||
@@ -1619,13 +1646,13 @@ class RuntimeConnectionHandler(handler.Handler):
|
||||
session_conversation_id = _get_run_authorization(session).get('conversation_id')
|
||||
event_run_id = event.get('run_id')
|
||||
if event_run_id and event_run_id == run_id:
|
||||
return handler.ActionResponse.success(data=event)
|
||||
return handler.ActionResponse.success(data=_project_event_record_for_api(event))
|
||||
if not session_conversation_id or event.get('conversation_id') != session_conversation_id:
|
||||
return handler.ActionResponse.error(
|
||||
message=f'Event {event_id} is not accessible by this run'
|
||||
)
|
||||
|
||||
return handler.ActionResponse.success(data=event)
|
||||
return handler.ActionResponse.success(data=_project_event_record_for_api(event))
|
||||
except Exception as e:
|
||||
self.ap.logger.error(f'EVENT_GET error: {e}', exc_info=True)
|
||||
return handler.ActionResponse.error(message=f'Event get error: {e}')
|
||||
@@ -1689,7 +1716,7 @@ class RuntimeConnectionHandler(handler.Handler):
|
||||
)
|
||||
|
||||
return handler.ActionResponse.success(data={
|
||||
'items': items,
|
||||
'items': [_project_event_record_for_api(item) for item in items],
|
||||
'next_cursor': str(next_seq) if next_seq else None,
|
||||
'prev_cursor': None,
|
||||
'has_more': has_more,
|
||||
|
||||
Reference in New Issue
Block a user