diff --git a/src/langbot/pkg/agent/runner/orchestrator.py b/src/langbot/pkg/agent/runner/orchestrator.py index f650f686f..48ca75cf9 100644 --- a/src/langbot/pkg/agent/runner/orchestrator.py +++ b/src/langbot/pkg/agent/runner/orchestrator.py @@ -148,6 +148,7 @@ class AgentRunOrchestrator: event=event, binding=binding, descriptor=descriptor, + tool_catalog=getattr(execution_query, '_host_tool_catalog', None), ) # Telemetry: report what this run was granted (never what it used). The # materialized resource list is authoritative for tools/skills/models; diff --git a/src/langbot/pkg/agent/runner/resource_builder.py b/src/langbot/pkg/agent/runner/resource_builder.py index 4e113f8f0..6cec4ba6a 100644 --- a/src/langbot/pkg/agent/runner/resource_builder.py +++ b/src/langbot/pkg/agent/runner/resource_builder.py @@ -55,6 +55,7 @@ class AgentResourceBuilder: event: AgentEventEnvelope, binding: AgentBinding, descriptor: RunnerDescriptor, + tool_catalog: typing.Iterable[dict[str, typing.Any]] | None = None, ) -> AgentResources: """Build AgentResources from event and binding. @@ -86,6 +87,7 @@ class AgentResourceBuilder: resource_policy, descriptor, runner_config, + tool_catalog, ) runner_uses_host_tools = config_schema.uses_host_tools(descriptor) platform_tools, platform_capabilities = build_platform_tool_resources( @@ -178,6 +180,7 @@ class AgentResourceBuilder: resource_policy: typing.Any, descriptor: RunnerDescriptor, runner_config: dict[str, typing.Any], + tool_catalog: typing.Iterable[dict[str, typing.Any]] | None = None, ) -> list[ToolResource]: """Build tools list from binding.""" tools: list[ToolResource] = [] @@ -190,6 +193,7 @@ class AgentResourceBuilder: allowed_names = resource_policy.allowed_tool_names allowed_sources = resource_policy.allowed_tool_sources + catalog_by_name = {item['name']: item for item in (tool_catalog or ())} if resource_policy.allow_all_tools or allowed_sources is None: get_catalog = getattr(getattr(self.ap, 'tool_mgr', None), 'get_resolved_tool_catalog', None) if get_catalog is None: @@ -208,6 +212,7 @@ class AgentResourceBuilder: selected_names = set(allowed_names or []) catalog = [item for item in catalog if item.get('name') in selected_names] allowed_names = ResourcePolicyProjector.extract_tool_names(catalog) + catalog_by_name = {item['name']: item for item in catalog} allowed_sources = { item['name']: ref for item in catalog if (ref := self._catalog_source_ref(item)) is not None } @@ -238,7 +243,15 @@ class AgentResourceBuilder: if source_ref is None: self.ap.logger.warning(f'Tool {tool_name} is not authorized because its Host source is unresolved') continue - if get_tool_schema is not None: + catalog_item = catalog_by_name.get(tool_name) + if ( + catalog_item is not None + and self._catalog_source_ref(catalog_item) == source_ref + and 'parameters' in catalog_item + ): + description = catalog_item.get('description') + parameters = catalog_item['parameters'] or None + elif get_tool_schema is not None: description, parameters = await get_tool_schema( execution_context, tool_name, diff --git a/src/langbot/pkg/pipeline/pipelinemgr.py b/src/langbot/pkg/pipeline/pipelinemgr.py index d1cfd5c62..f30507d15 100644 --- a/src/langbot/pkg/pipeline/pipelinemgr.py +++ b/src/langbot/pkg/pipeline/pipelinemgr.py @@ -346,7 +346,6 @@ class RuntimePipeline: if isinstance(result, typing.Coroutine): result = await result - await self._assert_execution_active(query) if isinstance(result, pipeline_entities.StageProcessResult): # 直接返回结果 self.ap.logger.debug( @@ -369,7 +368,6 @@ class RuntimePipeline: sub_result = await anext(iterator) except StopAsyncIteration: break - await self._assert_execution_active(query) self.ap.logger.debug( f'Stage {stage_container.inst_name} processed query {query.query_id} res {sub_result.result_type}' ) diff --git a/src/langbot/pkg/plugin/connector.py b/src/langbot/pkg/plugin/connector.py index d0a531355..706d9e1b8 100644 --- a/src/langbot/pkg/plugin/connector.py +++ b/src/langbot/pkg/plugin/connector.py @@ -558,8 +558,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): async def _load_workspace_desired_states( self, execution_context: ExecutionContext, + *, + settings: list[persistence_plugin.PluginSetting] | None = None, ) -> tuple[PluginInstallationDesiredState, ...]: - settings = await self._load_workspace_settings(execution_context) + if settings is None: + settings = await self._load_workspace_settings(execution_context) desired_states: list[PluginInstallationDesiredState] = [] for setting in settings: binding = self._binding_from_setting(execution_context, setting) @@ -928,16 +931,19 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): entitlement_revision=getattr(context, 'entitlement_revision', 0), ) - async def _synchronize_workspace(self, execution_context: ExecutionContext) -> None: + async def _synchronize_workspace( + self, execution_context: ExecutionContext + ) -> list[persistence_plugin.PluginSetting] | None: if not self.is_enable_plugin or not hasattr(self, 'handler'): - return + return None runtime_handler = self._runtime_handler() async with self._state_lock: # Read the durable desired state while holding the same gate used by # install/remove bookkeeping. Otherwise a request can load a stale # pre-install snapshot, wait for the installer to publish its # in-memory state, and then incorrectly remove that new binding. - desired_states = await self._load_workspace_desired_states(execution_context) + settings = await self._load_workspace_settings(execution_context) + desired_states = await self._load_workspace_desired_states(execution_context, settings=settings) desired_by_uuid = {state.binding.installation_uuid: state for state in desired_states} previous_ids = set(self._workspace_installations.get(execution_context.workspace_uuid, set())) for installation_uuid in previous_ids - set(desired_by_uuid): @@ -965,6 +971,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): execution_context.workspace_uuid, None, ) + return settings async def _current_execution_context(self) -> ExecutionContext: current = self._execution_context.get() @@ -1252,16 +1259,21 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): async def ping_plugin_runtime(self): return await self._runtime_handler().ping() - async def require_workspace_context(self, context: TenantContext) -> ExecutionContext: - """Validate and select one Workspace for this asyncio request task.""" - + async def _prepare_workspace_context( + self, context: TenantContext + ) -> tuple[ExecutionContext, list[persistence_plugin.PluginSetting] | None]: execution_context = await self._validate_execution_context(context) self._execution_context.set(execution_context) if not self.is_enable_plugin: - return execution_context + return execution_context, None if not hasattr(self, 'handler'): raise PluginRuntimeNotConnectedError('Plugin runtime is not connected') - await self._synchronize_workspace(execution_context) + settings = await self._synchronize_workspace(execution_context) + return execution_context, settings + + async def require_workspace_context(self, context: TenantContext) -> ExecutionContext: + """Validate and select one Workspace for this asyncio request task.""" + execution_context, _settings = await self._prepare_workspace_context(context) return execution_context def _inspect_plugin_package( @@ -1298,9 +1310,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): *, include_plugins: list[str] | None = None, include_disabled: bool = False, + settings: list[persistence_plugin.PluginSetting] | None = None, ) -> list[InstallationBinding]: execution_context = await self._current_execution_context() - settings = await self._load_workspace_settings(execution_context) + if settings is None: + settings = await self._load_workspace_settings(execution_context) bindings: list[InstallationBinding] = ( [self._legacy_oss_bridge_binding(execution_context)] if self.runtime_profile == 'oss_dev' else [] ) @@ -2347,10 +2361,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): bound_plugins: list[str] | None = None, ) -> context.EventContext: query = getattr(event, 'query', None) + settings = None if query is not None: from ..pipeline.pool import get_query_execution_context - await self.require_workspace_context(get_query_execution_context(query)) + _execution_context, settings = await self._prepare_workspace_context(get_query_execution_context(query)) event_ctx = context.EventContext.from_event(event) if not self.is_enable_plugin or not self._runtime_available(): @@ -2361,7 +2376,8 @@ class PluginRuntimeConnector(ManagedRuntimeConnector): runtime_handler = self._runtime_handler() emitted_plugins: list[Any] = [] response_sources: list[dict[str, Any]] = [] - for binding in await self._operation_bindings(include_plugins=bound_plugins): + # Reuse the settings read under the reconciliation lock for this dispatch. + for binding in await self._operation_bindings(include_plugins=bound_plugins, settings=settings): with runtime_handler.installation_scope(binding): result = await runtime_handler.emit_event( event_ctx.model_dump(serialize_as_any=False), diff --git a/src/langbot/pkg/provider/tools/toolmgr.py b/src/langbot/pkg/provider/tools/toolmgr.py index b4116a9e0..f2ff12b62 100644 --- a/src/langbot/pkg/provider/tools/toolmgr.py +++ b/src/langbot/pkg/provider/tools/toolmgr.py @@ -271,8 +271,11 @@ class ToolManager: query: pipeline_query.Query, catalog: typing.Iterable[dict[str, typing.Any]], ) -> None: + catalog = list(catalog) query.variables = query.variables or {} query.variables[TOOL_SOURCE_REFS_QUERY_KEY] = cls.source_refs_from_catalog(catalog) + # Keep Host metadata off the plugin wire; invocation still resolves its live source. + object.__setattr__(query, '_host_tool_catalog', tuple(catalog)) @staticmethod def get_query_tool_source( diff --git a/src/langbot/pkg/workspace/repository.py b/src/langbot/pkg/workspace/repository.py index 374fa9b0b..6f1df32c3 100644 --- a/src/langbot/pkg/workspace/repository.py +++ b/src/langbot/pkg/workspace/repository.py @@ -1,5 +1,7 @@ from __future__ import annotations +import typing + import sqlalchemy from sqlalchemy.ext.asyncio import AsyncSession @@ -13,6 +15,26 @@ from ..entity.persistence.workspace import ( ) +_WORKSPACE_COLUMNS = Workspace.__table__.c +_EXECUTION_COLUMNS = WorkspaceExecutionState.__table__.c +_EXECUTION_PROJECTION = ( + sqlalchemy.select( + _WORKSPACE_COLUMNS.uuid, + _WORKSPACE_COLUMNS.instance_uuid, + _WORKSPACE_COLUMNS.source, + _WORKSPACE_COLUMNS.status, + _WORKSPACE_COLUMNS.created_at, + _EXECUTION_COLUMNS.instance_uuid.label('execution_instance_uuid'), + _EXECUTION_COLUMNS.active_generation, + _EXECUTION_COLUMNS.state, + _EXECUTION_COLUMNS.write_fenced, + _EXECUTION_COLUMNS.source.label('execution_source'), + ) + .outerjoin(WorkspaceExecutionState.__table__, _EXECUTION_COLUMNS.workspace_uuid == _WORKSPACE_COLUMNS.uuid) + .where(_WORKSPACE_COLUMNS.uuid == sqlalchemy.bindparam('workspace_uuid')) +) + + class WorkspaceRepository: """Transaction-bound persistence operations for the workspace directory.""" @@ -46,6 +68,49 @@ class WorkspaceRepository: async def get_workspace(self, workspace_uuid: str) -> Workspace | None: return await self.session.get(Workspace, workspace_uuid) + async def get_execution_projection(self, workspace_uuid: str) -> dict[str, typing.Any] | None: + """Read current fencing records without flushing caller-owned writes.""" + with self.session.no_autoflush: + result = await self.session.execute(_EXECUTION_PROJECTION, {'workspace_uuid': workspace_uuid}) + row = result.mappings().first() + if row is None: + return None + projection = dict(row) + for entity in self.session.new | self.session.deleted: + if isinstance(entity, Workspace) and ( + entity.uuid == workspace_uuid or sqlalchemy.inspect(entity).identity == (workspace_uuid,) + ): + # New and deleted directory records cannot authorize execution. + return None + if isinstance(entity, WorkspaceExecutionState) and ( + entity.workspace_uuid == workspace_uuid or sqlalchemy.inspect(entity).identity == (workspace_uuid,) + ): + projection['execution_instance_uuid'] = None + for entity in self.session.dirty: + if not isinstance(entity, (Workspace, WorkspaceExecutionState)): + continue + state = sqlalchemy.inspect(entity) + if state.identity != (workspace_uuid,): + continue + identity_field = 'uuid' if isinstance(entity, Workspace) else 'workspace_uuid' + if state.attrs[identity_field].history.has_changes(): + return None + if isinstance(entity, Workspace): + fields = {name: name for name in ('uuid', 'instance_uuid', 'source', 'status', 'created_at')} + else: + fields = { + 'instance_uuid': 'execution_instance_uuid', + 'active_generation': 'active_generation', + 'state': 'state', + 'write_fenced': 'write_fenced', + 'source': 'execution_source', + } + for name, key in fields.items(): + history = state.attrs[name].history + if history.has_changes(): + projection[key] = history.added[0] if history.added else None + return projection + def add_workspace(self, workspace: Workspace) -> None: self.session.add(workspace) diff --git a/src/langbot/pkg/workspace/service.py b/src/langbot/pkg/workspace/service.py index c20fdc228..dc7cf84dc 100644 --- a/src/langbot/pkg/workspace/service.py +++ b/src/langbot/pkg/workspace/service.py @@ -231,7 +231,8 @@ class WorkspaceService: ) async def operation(repository: WorkspaceRepository) -> WorkspaceExecutionBinding: - if workspace_uuid is None: + selected_uuid = workspace_uuid + if selected_uuid is None: workspaces = await repository.list_local_workspaces(self.instance_uuid) if not workspaces: raise WorkspaceNotFoundError('The local workspace has not been initialized') @@ -239,52 +240,51 @@ class WorkspaceService: raise WorkspaceInvariantError( f'Expected one local workspace for {self.instance_uuid!r}, found {len(workspaces)}' ) - workspace = workspaces[0] - else: - workspace = await repository.get_workspace(workspace_uuid) - if workspace is None: - raise WorkspaceNotFoundError(f'Workspace {workspace_uuid!r} does not exist') + selected_uuid = workspaces[0].uuid - if workspace.instance_uuid != self.instance_uuid: - raise WorkspaceInvariantError(f'Workspace {workspace.uuid!r} belongs to another instance') - if _require_local and workspace.source != WorkspaceSource.LOCAL.value: - raise WorkspaceInvariantError(f'Workspace {workspace.uuid!r} is not an OSS local workspace') - if workspace.status != WorkspaceStatus.ACTIVE.value: - raise WorkspaceExecutionUnavailableError(f'Workspace {workspace.uuid!r} is not active') + workspace = await repository.get_execution_projection(selected_uuid) + if workspace is None: + raise WorkspaceNotFoundError(f'Workspace {selected_uuid!r} does not exist') - execution_state = await repository.get_execution_state(workspace.uuid) - if execution_state is None: - raise WorkspaceExecutionUnavailableError(f'Workspace {workspace.uuid!r} has no execution state') - if execution_state.instance_uuid != self.instance_uuid: + if workspace['instance_uuid'] != self.instance_uuid: + raise WorkspaceInvariantError(f'Workspace {workspace["uuid"]!r} belongs to another instance') + if _require_local and workspace['source'] != WorkspaceSource.LOCAL.value: + raise WorkspaceInvariantError(f'Workspace {workspace["uuid"]!r} is not an OSS local workspace') + if workspace['status'] != WorkspaceStatus.ACTIVE.value: + raise WorkspaceExecutionUnavailableError(f'Workspace {workspace["uuid"]!r} is not active') + + if workspace['execution_instance_uuid'] is None: + raise WorkspaceExecutionUnavailableError(f'Workspace {workspace["uuid"]!r} has no execution state') + if workspace['execution_instance_uuid'] != self.instance_uuid: raise WorkspaceInvariantError( - f'Workspace {workspace.uuid!r} execution state belongs to another instance' + f'Workspace {workspace["uuid"]!r} execution state belongs to another instance' ) expected_source = ( WorkspaceExecutionSource.LOCAL.value - if workspace.source == WorkspaceSource.LOCAL.value + if workspace['source'] == WorkspaceSource.LOCAL.value else WorkspaceExecutionSource.CLOUD.value ) - if execution_state.source != expected_source: + if workspace['execution_source'] != expected_source: raise WorkspaceInvariantError( - f'Workspace {workspace.uuid!r} execution source does not match its directory source' + f'Workspace {workspace["uuid"]!r} execution source does not match its directory source' ) - if execution_state.state != WorkspaceExecutionStatus.ACTIVE.value or execution_state.write_fenced: - raise WorkspaceExecutionUnavailableError(f'Workspace {workspace.uuid!r} execution is unavailable') - if execution_state.active_generation <= 0: - raise WorkspaceInvariantError(f'Workspace {workspace.uuid!r} has an invalid execution generation') - if expected_generation is not None and execution_state.active_generation != expected_generation: + if workspace['state'] != WorkspaceExecutionStatus.ACTIVE.value or workspace['write_fenced']: + raise WorkspaceExecutionUnavailableError(f'Workspace {workspace["uuid"]!r} execution is unavailable') + if workspace['active_generation'] <= 0: + raise WorkspaceInvariantError(f'Workspace {workspace["uuid"]!r} has an invalid execution generation') + if expected_generation is not None and workspace['active_generation'] != expected_generation: raise WorkspaceGenerationMismatchError( - f'Workspace {workspace.uuid!r} generation {execution_state.active_generation} ' + f'Workspace {workspace["uuid"]!r} generation {workspace["active_generation"]} ' f'does not match expected generation {expected_generation}' ) return WorkspaceExecutionBinding( instance_uuid=self.instance_uuid, - workspace_uuid=workspace.uuid, - placement_generation=execution_state.active_generation, - write_fenced=execution_state.write_fenced, - state=execution_state.state, - workspace_created_at=workspace.created_at, + workspace_uuid=workspace['uuid'], + placement_generation=workspace['active_generation'], + write_fenced=workspace['write_fenced'], + state=workspace['state'], + workspace_created_at=workspace['created_at'], ) binding = await self._run(operation, session=session) diff --git a/tests/unit_tests/agent/test_resource_builder.py b/tests/unit_tests/agent/test_resource_builder.py index 03c2fc69a..3273a6293 100644 --- a/tests/unit_tests/agent/test_resource_builder.py +++ b/tests/unit_tests/agent/test_resource_builder.py @@ -14,6 +14,7 @@ from langbot.pkg.agent.runner.query_entry_adapter import QueryEntryAdapter from langbot.pkg.agent.runner.resource_builder import AgentResourceBuilder from langbot.pkg.agent.runner.host_models import AgentBinding, AgentEventEnvelope, BindingScope, ResourcePolicy from langbot.pkg.api.http.context import ExecutionContext +from langbot.pkg.provider.tools.toolmgr import ToolManager RUNNER_ID = 'plugin:test/runner/default' @@ -100,6 +101,7 @@ async def build_resources(app, query, descriptor): event=event, binding=binding, descriptor=descriptor, + tool_catalog=getattr(query, '_host_tool_catalog', None), ) @@ -419,6 +421,38 @@ async def test_build_tools_authorizes_query_declared_tools(app): ] +@pytest.mark.asyncio +@pytest.mark.parametrize('catalog_source', ['builtin', 'plugin']) +async def test_build_tools_reuses_only_matching_host_catalog_schema(app, catalog_source): + query = make_query({}, use_funcs=[{'name': 'exec'}]) + ToolManager.bind_query_tool_sources( + query, + [ + { + 'name': 'exec', + 'source': catalog_source, + 'description': 'Catalog command', + 'parameters': {'type': 'object', 'properties': {'command': {'type': 'string'}}}, + } + ], + ) + query.variables['_host_tool_source_refs'] = {'exec': {'source': 'builtin', 'source_id': None}} + app.tool_mgr.get_tool_schema.return_value = ('Live command', {'type': 'object'}) + + resources = await build_resources(app, query, make_descriptor(capabilities={'tool_calling': True})) + tool = next(tool for tool in resources['tools'] if tool['tool_name'] == 'exec') + if catalog_source == 'builtin': + assert tool['description'] == 'Catalog command' + assert tool['parameters'] == {'type': 'object', 'properties': {'command': {'type': 'string'}}} + app.tool_mgr.get_tool_schema.assert_not_awaited() + else: + assert tool['description'] == 'Live command' + assert tool['parameters'] == {'type': 'object'} + app.tool_mgr.get_tool_schema.assert_awaited_once_with( + TEST_CONTEXT, 'exec', source_ref={'source': 'builtin', 'source_id': None} + ) + + @pytest.mark.asyncio async def test_build_tools_manifest_permission_denies_binding_tools(app): """Binding tool grants should be removed when manifest does not request tools.""" @@ -446,8 +480,19 @@ async def test_build_tools_materializes_independent_agent_all_tools_policy(app): """Independent Agents resolve an all-tools grant against the live Host catalog.""" app.tool_mgr.get_resolved_tool_catalog = AsyncMock( return_value=[ - {'name': 'exec', 'source': 'builtin'}, - {'name': 'plugin_tool', 'source': 'plugin', 'source_id': 'test/plugin'}, + { + 'name': 'exec', + 'source': 'builtin', + 'description': 'Execute a command', + 'parameters': {'type': 'object', 'properties': {'command': {'type': 'string'}}}, + }, + { + 'name': 'plugin_tool', + 'source': 'plugin', + 'source_id': 'test/plugin', + 'description': 'Plugin tool', + 'parameters': {}, + }, ] ) descriptor = make_descriptor(capabilities={'tool_calling': True}) @@ -467,6 +512,10 @@ async def test_build_tools_materializes_independent_agent_all_tools_policy(app): ) assert [tool['tool_name'] for tool in resources['tools']] == ['exec', 'plugin_tool'] + assert resources['tools'][0]['description'] == 'Execute a command' + assert resources['tools'][0]['parameters'] == {'type': 'object', 'properties': {'command': {'type': 'string'}}} + assert resources['tools'][1]['parameters'] is None + app.tool_mgr.get_tool_schema.assert_not_awaited() app.tool_mgr.get_resolved_tool_catalog.assert_awaited_once_with( TEST_CONTEXT, include_skill_authoring=True, diff --git a/tests/unit_tests/pipeline/test_pipelinemgr.py b/tests/unit_tests/pipeline/test_pipelinemgr.py index 10bf0ce70..0e56d3697 100644 --- a/tests/unit_tests/pipeline/test_pipelinemgr.py +++ b/tests/unit_tests/pipeline/test_pipelinemgr.py @@ -349,9 +349,11 @@ async def test_runtime_pipeline_rejects_stale_generation_before_side_effects( @pytest.mark.asyncio +@pytest.mark.parametrize('streaming', [False, True]) async def test_runtime_pipeline_revalidates_after_awaited_stage( mock_app, sample_query, + streaming, ): pipelinemgr = get_pipelinemgr_module() stage = get_stage_module() @@ -378,8 +380,11 @@ async def test_runtime_pipeline_revalidates_after_awaited_stage( ) return result + async def streaming_process(*args): + yield await stage_process(*args) + mock_stage = Mock(spec=stage.PipelineStage) - mock_stage.process = Mock(side_effect=stage_process) + mock_stage.process = Mock(side_effect=streaming_process if streaming else stage_process) runtime_pipeline = pipelinemgr.RuntimePipeline( mock_app, pipeline_entity, diff --git a/tests/unit_tests/plugin/test_connector_methods.py b/tests/unit_tests/plugin/test_connector_methods.py index c971324ba..82450a033 100644 --- a/tests/unit_tests/plugin/test_connector_methods.py +++ b/tests/unit_tests/plugin/test_connector_methods.py @@ -68,6 +68,7 @@ def create_mock_connector(): instance._target_binding = AsyncMock(return_value=TEST_INSTALLATION_BINDING) instance._load_workspace_settings = AsyncMock(return_value=[]) instance.require_workspace_context = AsyncMock(side_effect=lambda context: context) + instance._prepare_workspace_context = AsyncMock(return_value=(TEST_EXECUTION_CONTEXT, [])) return instance diff --git a/tests/unit_tests/plugin/test_connector_reconcile.py b/tests/unit_tests/plugin/test_connector_reconcile.py index bfb7c2431..84564963b 100644 --- a/tests/unit_tests/plugin/test_connector_reconcile.py +++ b/tests/unit_tests/plugin/test_connector_reconcile.py @@ -490,6 +490,7 @@ async def test_workspace_reads_do_not_wait_for_an_installation_apply(): connector._store_artifact_package = AsyncMock() connector._persist_installation_package = AsyncMock(return_value=(binding, None, False)) connector._wait_for_installed_plugin_ready = AsyncMock() + connector._load_workspace_settings = AsyncMock(return_value=[]) connector._load_workspace_desired_states = AsyncMock( return_value=[ PluginInstallationDesiredState( @@ -880,3 +881,49 @@ async def test_marketplace_download_honors_migration_version_without_latest_look SimpleNamespace(), 'langbot-team', 'LocalAgent', None, version=version ) assert download.await_count == 1 + + +@pytest.mark.asyncio +async def test_event_dispatch_reuses_reconciled_settings_and_refreshes_next_event(): + from langbot_plugin.api.entities.events import PersonMessageReceived + from tests.factories import text_query + + binding = execution_binding('workspace-a') + setting = plugin_setting('01', 'a' * 64) + settings = {'workspace-a': [setting]} + connector = shared_connector([[binding]], settings) + connector.handler = runtime_handler() + connector.handler.remove_plugin_installation = AsyncMock() + connector._apply_desired_state = AsyncMock() + + async def emit(event_context, include_plugins=None): + return {'event_context': event_context, 'emitted_plugins': []} + + connector.handler.emit_event = AsyncMock(side_effect=emit) + query = text_query('hello') + object.__setattr__( + query, + '_execution_context', + ExecutionContext(instance_uuid='instance-a', workspace_uuid='workspace-a', placement_generation=1), + ) + event = PersonMessageReceived( + query=query, + launcher_type=query.launcher_type.value, + launcher_id=query.launcher_id, + sender_id=query.sender_id, + message_event=query.message_event, + message_chain=query.message_chain, + ) + + await connector.emit_event(event) + connector._load_workspace_settings.assert_awaited_once() + connector.handler.emit_event.assert_awaited_once() + first_scope = connector.handler.installation_scope.call_args.args[0] + assert first_scope.workspace_uuid == 'workspace-a' + assert first_scope.installation_uuid == setting.installation_uuid + + setting.enabled = False + await connector.emit_event(event) + assert connector._load_workspace_settings.await_count == 2 + connector.handler.emit_event.assert_awaited_once() + connector._apply_desired_state.assert_awaited() diff --git a/tests/unit_tests/workspace/test_workspace_service.py b/tests/unit_tests/workspace/test_workspace_service.py index 5d6633d32..f4de6d051 100644 --- a/tests/unit_tests/workspace/test_workspace_service.py +++ b/tests/unit_tests/workspace/test_workspace_service.py @@ -23,6 +23,7 @@ from langbot.pkg.workspace import ( WorkspaceGenerationMismatchError, WorkspaceInvariantError, WorkspaceLimitExceededError, + WorkspaceNotFoundError, WorkspaceOwnerAlreadyExistsError, WorkspaceService, ) @@ -217,6 +218,154 @@ async def test_execution_binding_fails_closed_when_fenced(workspace_test_context await service.get_local_execution_context(workspace.uuid) +@pytest.mark.parametrize( + ('model', 'changes', 'error'), + [ + (Workspace, {'status': 'suspended'}, WorkspaceExecutionUnavailableError), + (Workspace, {'instance_uuid': 'another-instance'}, WorkspaceInvariantError), + (WorkspaceExecutionState, {'write_fenced': True}, WorkspaceExecutionUnavailableError), + (WorkspaceExecutionState, {'active_generation': 2}, WorkspaceGenerationMismatchError), + (WorkspaceExecutionState, {'instance_uuid': 'another-instance'}, WorkspaceInvariantError), + (WorkspaceExecutionState, {'source': 'cloud'}, WorkspaceInvariantError), + (WorkspaceExecutionState, None, WorkspaceExecutionUnavailableError), + ], +) +async def test_execution_binding_observes_projection_changes(workspace_test_context, model, changes, error): + service, session_factory = workspace_test_context + workspace = await service.ensure_singleton_workspace() + assert (await service.get_execution_binding(workspace.uuid, expected_generation=1)).placement_generation == 1 + + key = Workspace.uuid if model is Workspace else WorkspaceExecutionState.workspace_uuid + statement = sqlalchemy.delete(model) if changes is None else sqlalchemy.update(model).values(**changes) + async with session_factory.begin() as session: + await session.execute(statement.where(key == workspace.uuid)) + + with pytest.raises(error): + await service.get_execution_binding(workspace.uuid, expected_generation=1) + + +@pytest.mark.parametrize( + ('changes', 'error'), + [ + ({'write_fenced': True}, WorkspaceExecutionUnavailableError), + ({'active_generation': 2}, WorkspaceGenerationMismatchError), + ], +) +async def test_execution_binding_observes_unflushed_projection_changes(workspace_test_context, changes, error): + service, session_factory = workspace_test_context + workspace = await service.ensure_singleton_workspace() + async with session_factory.begin() as session: + state = await session.get(WorkspaceExecutionState, workspace.uuid) + for name, value in changes.items(): + setattr(state, name, value) + with pytest.raises(error): + await service.get_execution_binding(workspace.uuid, expected_generation=1, session=session) + + +async def test_execution_binding_does_not_lock_unrelated_pending_writes(workspace_test_context): + service, session_factory = workspace_test_context + workspace = await service.ensure_singleton_workspace() + engine = service.ap.persistence_mgr.get_db_engine() + async with engine.begin() as connection: + await connection.exec_driver_sql('PRAGMA journal_mode=WAL') + writer_engine = create_async_engine(engine.url, connect_args={'timeout': 0.05}) + writer_factory = async_sessionmaker(writer_engine, expire_on_commit=False) + try: + async with session_factory.begin() as session: + directory = [ + await session.get(Workspace, workspace.uuid), + await session.get(WorkspaceExecutionState, workspace.uuid), + ] + pending_account = User( + uuid=str(uuid.uuid4()), + user='pending@example.com', + normalized_email='pending@example.com', + password='hashed-password', + account_type='local', + ) + session.add(pending_account) + binding = await service.get_execution_binding(workspace.uuid, expected_generation=1, session=session) + assert binding.placement_generation == 1 + assert pending_account in session.new + async with writer_factory.begin() as writer: + await _insert_account(writer, 'independent@example.com') + assert all(directory) + finally: + await writer_engine.dispose() + + +@pytest.mark.parametrize( + ('changes', 'error'), + [ + ({'write_fenced': True}, WorkspaceExecutionUnavailableError), + ({'active_generation': 2}, WorkspaceGenerationMismatchError), + ], +) +async def test_execution_binding_keeps_external_updates_over_stale_orm_values(workspace_test_context, changes, error): + service, session_factory = workspace_test_context + workspace = await service.ensure_singleton_workspace() + async with session_factory() as session: + state = await session.get(WorkspaceExecutionState, workspace.uuid) + state.desired_state_revision = 1 + async with session_factory.begin() as writer: + await writer.execute( + sqlalchemy.update(WorkspaceExecutionState) + .where(WorkspaceExecutionState.workspace_uuid == workspace.uuid) + .values(**changes) + ) + with pytest.raises(error): + await service.get_execution_binding(workspace.uuid, expected_generation=1, session=session) + assert state in session.dirty + assert state.write_fenced is False + assert state.active_generation == 1 + + +@pytest.mark.parametrize('model', [Workspace, WorkspaceExecutionState]) +async def test_execution_binding_rejects_unflushed_directory_deletion(workspace_test_context, model): + service, session_factory = workspace_test_context + workspace = await service.ensure_singleton_workspace() + async with session_factory() as session: + entity = await session.get(model, workspace.uuid) + await session.delete(entity) + error = WorkspaceNotFoundError if model is Workspace else WorkspaceExecutionUnavailableError + with pytest.raises(error): + await service.get_execution_binding(workspace.uuid, expected_generation=1, session=session) + assert entity in session.deleted + + +async def test_execution_binding_does_not_publish_new_execution_state(workspace_test_context): + service, session_factory = workspace_test_context + workspace = await service.ensure_singleton_workspace() + async with session_factory.begin() as session: + await session.execute( + sqlalchemy.delete(WorkspaceExecutionState).where(WorkspaceExecutionState.workspace_uuid == workspace.uuid) + ) + async with session_factory() as session: + state = WorkspaceExecutionState( + workspace_uuid=workspace.uuid, + instance_uuid=service.instance_uuid, + active_generation=1, + state='active', + write_fenced=False, + source='local', + desired_state_revision=0, + ) + session.add(state) + with pytest.raises(WorkspaceExecutionUnavailableError): + await service.get_execution_binding(workspace.uuid, expected_generation=1, session=session) + assert state in session.new + + +async def test_execution_binding_does_not_publish_new_workspace(workspace_test_context): + service, session_factory = workspace_test_context + workspace = service._new_local_workspace(name='Pending', slug='pending') + async with session_factory() as session: + session.add(workspace) + with pytest.raises(WorkspaceNotFoundError): + await service.get_execution_binding(workspace.uuid, expected_generation=1, session=session) + assert workspace in session.new + + async def test_cloud_projection_requires_explicit_general_binding(workspace_test_context): service, session_factory = workspace_test_context await service.ensure_singleton_workspace()