diff --git a/src/langbot/pkg/api/http/controller/groups/system.py b/src/langbot/pkg/api/http/controller/groups/system.py index 55af6c4bd..2c243ea36 100644 --- a/src/langbot/pkg/api/http/controller/groups/system.py +++ b/src/langbot/pkg/api/http/controller/groups/system.py @@ -7,7 +7,7 @@ from .. import group from .....utils import constants from .....entity.persistence.metadata import WorkspaceMetadata from ...authz import Permission -from ...context import RequestContext +from ...context import PrincipalType, RequestContext from .....provider.tools.loaders.mcp_policy import stdio_mcp_enabled from .....workspace.invitation_delivery import InvitationDeliveryService @@ -24,6 +24,11 @@ SYSTEM_CAPABILITY_OPERATIONS = ( 'pipeline.update', 'pipeline.delete', 'pipeline.copy', + 'task.list', + 'task.get', + 'knowledge_base.get', + 'knowledge_base.file.store', + 'file.document.upload', ) @@ -258,7 +263,7 @@ class SystemRouterGroup(group.RouterGroup): @self.route( '/tasks', methods=['GET'], - auth_type=group.AuthType.USER_TOKEN, + auth_type=group.AuthType.USER_TOKEN_OR_API_KEY, permission=Permission.RESOURCE_VIEW, ) async def _(request_context: RequestContext) -> str: @@ -277,18 +282,23 @@ class SystemRouterGroup(group.RouterGroup): instance_uuid=request_context.instance_uuid, workspace_uuid=request_context.workspace_uuid, placement_generation=request_context.placement_generation, + public=request_context.principal.principal_type == PrincipalType.API_KEY, ) ) @self.route( '/tasks/', methods=['GET'], - auth_type=group.AuthType.USER_TOKEN, + auth_type=group.AuthType.USER_TOKEN_OR_API_KEY, permission=Permission.RESOURCE_VIEW, ) async def _(task_id: str, request_context: RequestContext) -> str: + try: + task_index = int(task_id) + except (TypeError, ValueError): + return self.http_status(404, 404, 'Task not found') task = self.ap.task_mgr.get_task_by_id( - int(task_id), + task_index, instance_uuid=request_context.instance_uuid, workspace_uuid=request_context.workspace_uuid, placement_generation=request_context.placement_generation, @@ -297,6 +307,8 @@ class SystemRouterGroup(group.RouterGroup): if task is None: return self.http_status(404, 404, 'Task not found') + if request_context.principal.principal_type == PrincipalType.API_KEY: + return self.success(data=task.to_public_dict()) return self.success(data=task.to_dict()) @self.route( diff --git a/src/langbot/pkg/core/taskmgr.py b/src/langbot/pkg/core/taskmgr.py index 25cc38e0d..6c5df7b81 100644 --- a/src/langbot/pkg/core/taskmgr.py +++ b/src/langbot/pkg/core/taskmgr.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import json import typing import datetime import time @@ -197,6 +198,41 @@ class TaskWrapper: }, } + def to_public_dict(self) -> dict: + """Return the stable task projection exposed to API-key callers.""" + if self.task.cancelled(): + status = 'cancelled' + error = {'type': 'task_cancelled', 'message': 'Task was cancelled'} + result = None + elif not self.task.done(): + status = 'running' + error = None + result = None + else: + exception = self.assume_exception() + if exception is not None: + status = 'failed' + error = {'type': 'task_failed', 'message': 'Task execution failed'} + result = None + else: + status = 'succeeded' + error = None + result = self.assume_result() + try: + json.dumps(result) + except (TypeError, ValueError): + result = None + + return { + 'id': self.id, + 'task_type': self.task_type, + 'kind': self.kind, + 'status': status, + 'error': error, + 'result': result, + 'created_at': self.created_at, + } + def cancel(self): self.task.cancel() @@ -325,19 +361,20 @@ class AsyncTaskManager: instance_uuid: str | None = None, workspace_uuid: str | None = None, placement_generation: int | None = None, + public: bool = False, ) -> dict: - return { - 'tasks': [ - t.to_dict() - for t in self.tasks - if (type is None or t.task_type == type) - and (kind is None or t.kind == kind) - and (instance_uuid is None or t.instance_uuid == instance_uuid) - and (workspace_uuid is None or t.workspace_uuid == workspace_uuid) - and (placement_generation is None or t.placement_generation == placement_generation) - ], - 'id_index': TaskWrapper._id_index, - } + tasks = [ + t.to_public_dict() if public else t.to_dict() + for t in self.tasks + if (type is None or t.task_type == type) + and (kind is None or t.kind == kind) + and (instance_uuid is None or t.instance_uuid == instance_uuid) + and (workspace_uuid is None or t.workspace_uuid == workspace_uuid) + and (placement_generation is None or t.placement_generation == placement_generation) + ] + if public: + return {'tasks': tasks} + return {'tasks': tasks, 'id_index': TaskWrapper._id_index} def get_stats(self) -> dict: completed = sum(1 for t in self.tasks if t.task.done()) diff --git a/tests/integration/api/test_workspaces.py b/tests/integration/api/test_workspaces.py index b861ce2da..656ae5941 100644 --- a/tests/integration/api/test_workspaces.py +++ b/tests/integration/api/test_workspaces.py @@ -509,6 +509,11 @@ async def test_api_key_context_returns_bound_identity_without_workspace_permissi 'pipeline.update', 'pipeline.delete', 'pipeline.copy', + 'task.list', + 'task.get', + 'knowledge_base.get', + 'knowledge_base.file.store', + 'file.document.upload', ] ) assert all(item == {'supported': True} for item in capabilities['operations'].values()) @@ -556,6 +561,83 @@ async def test_api_key_context_returns_bound_identity_without_workspace_permissi assert revoked_capabilities.status_code == 401 +async def test_api_key_can_query_tasks_with_public_contract_and_resource_permission(workspace_api): + application, client, _, owner_token = workspace_api + task_query = {} + task_lookup = {} + fake_task = SimpleNamespace( + to_public_dict=lambda: {'id': 7, 'status': 'running', 'error': None, 'result': None}, + to_dict=lambda: {'id': 7, 'runtime': {'state': 'PENDING'}}, + ) + + def get_tasks_dict(*args, **kwargs): + task_query.update(kwargs) + if kwargs.get('public'): + return {'tasks': []} + return {'tasks': [], 'id_index': 1} + + def get_task_by_id(*args, **kwargs): + task_lookup.update(kwargs) + return fake_task if args and args[0] == 7 else None + + application.task_mgr = SimpleNamespace( + get_tasks_dict=get_tasks_dict, + get_task_by_id=get_task_by_id, + ) + current_response = await client.get('/api/v1/workspaces/current', headers=_auth(owner_token)) + workspace_uuid = (await current_response.get_json())['data']['workspace']['uuid'] + create_response = await client.post( + '/api/v1/apikeys', + headers=_auth(owner_token, workspace_uuid), + json={'name': 'Task reader', 'scopes': ['resource.view']}, + ) + assert create_response.status_code == 200 + key = (await create_response.get_json())['data']['key']['key'] + + listing = await client.get('/api/v1/system/tasks', headers={'X-API-Key': key}) + assert listing.status_code == 200 + assert (await listing.get_json())['data'] == {'tasks': []} + assert task_query['instance_uuid'] == application.workspace_service.instance_uuid + assert task_query['workspace_uuid'] == workspace_uuid + assert task_query['placement_generation'] == 1 + assert task_query['public'] is True + + bearer_listing = await client.get('/api/v1/system/tasks', headers=_auth(owner_token, workspace_uuid)) + assert bearer_listing.status_code == 200 + assert (await bearer_listing.get_json())['data'] == {'tasks': [], 'id_index': 1} + + public_task = await client.get('/api/v1/system/tasks/7', headers={'X-API-Key': key}) + assert public_task.status_code == 200 + assert (await public_task.get_json())['data'] == { + 'id': 7, + 'status': 'running', + 'error': None, + 'result': None, + } + assert task_lookup == { + 'instance_uuid': application.workspace_service.instance_uuid, + 'workspace_uuid': workspace_uuid, + 'placement_generation': 1, + } + + legacy_task = await client.get('/api/v1/system/tasks/7', headers=_auth(owner_token, workspace_uuid)) + assert legacy_task.status_code == 200 + assert (await legacy_task.get_json())['data'] == {'id': 7, 'runtime': {'state': 'PENDING'}} + + missing = await client.get('/api/v1/system/tasks/not-an-id', headers={'X-API-Key': key}) + assert missing.status_code == 404 + + no_permission_response = await client.post( + '/api/v1/apikeys', + headers=_auth(owner_token, workspace_uuid), + json={'name': 'Task denied', 'scopes': []}, + ) + assert no_permission_response.status_code == 200 + no_permission_key = (await no_permission_response.get_json())['data']['key']['key'] + denied = await client.get('/api/v1/system/tasks', headers={'X-API-Key': no_permission_key}) + assert denied.status_code == 403 + + async def test_cloud_projection_is_selected_explicitly_and_collaboration_runs_in_core( workspace_api, ): diff --git a/tests/unit_tests/core/test_taskmgr.py b/tests/unit_tests/core/test_taskmgr.py index 44503de9a..4c23fb2ac 100644 --- a/tests/unit_tests/core/test_taskmgr.py +++ b/tests/unit_tests/core/test_taskmgr.py @@ -338,6 +338,69 @@ class TestTaskWrapper: assert result['runtime']['exception'] == 'Test error' assert 'exception_traceback' in result['runtime'] + @pytest.mark.asyncio + async def test_public_dict_has_stable_success_projection(self): + _, TaskWrapper, _ = get_taskmgr_classes() + mock_app = create_mock_app() + + async def successful_coro(): + return {'file_id': 'file-a'} + + wrapper = TaskWrapper(mock_app, successful_coro(), kind='knowledge_base.store') + await wrapper.task + + result = wrapper.to_public_dict() + + assert result == { + 'id': wrapper.id, + 'task_type': 'system', + 'kind': 'knowledge_base.store', + 'status': 'succeeded', + 'error': None, + 'result': {'file_id': 'file-a'}, + 'created_at': result['created_at'], + } + assert 'runtime' not in result + assert 'traceback' not in str(result).lower() + + @pytest.mark.asyncio + async def test_public_dict_hides_exception_traceback(self): + _, TaskWrapper, _ = get_taskmgr_classes() + mock_app = create_mock_app() + + async def failing_coro(): + raise ValueError('private failure') + + wrapper = TaskWrapper(mock_app, failing_coro()) + try: + await wrapper.task + except ValueError: + pass + + result = wrapper.to_public_dict() + + assert result['status'] == 'failed' + assert result['error'] == {'type': 'task_failed', 'message': 'Task execution failed'} + assert 'runtime' not in result + assert 'traceback' not in str(result).lower() + + @pytest.mark.asyncio + async def test_public_dict_does_not_change_success_when_result_is_not_json_serializable(self): + _, TaskWrapper, _ = get_taskmgr_classes() + mock_app = create_mock_app() + + async def successful_coro(): + return object() + + wrapper = TaskWrapper(mock_app, successful_coro()) + await wrapper.task + + result = wrapper.to_public_dict() + + assert result['status'] == 'succeeded' + assert result['error'] is None + assert result['result'] is None + @pytest.mark.asyncio async def test_cancel_task(self): """Test cancel method cancels the asyncio task.""" @@ -487,6 +550,66 @@ class TestAsyncTaskManager: w2.cancel() w3.cancel() + @pytest.mark.asyncio + async def test_public_task_queries_keep_workspace_and_generation_isolation(self): + _, _, AsyncTaskManager = get_taskmgr_classes() + mock_app = create_mock_app() + manager = AsyncTaskManager(mock_app) + + async def dummy_coro(): + await asyncio.sleep(10) + + current = manager.create_user_task( + dummy_coro(), + instance_uuid='instance-a', + workspace_uuid='workspace-a', + placement_generation=2, + ) + other_workspace = manager.create_user_task( + dummy_coro(), + instance_uuid='instance-a', + workspace_uuid='workspace-b', + placement_generation=2, + ) + stale_generation = manager.create_user_task( + dummy_coro(), + instance_uuid='instance-a', + workspace_uuid='workspace-a', + placement_generation=1, + ) + + result = manager.get_tasks_dict( + instance_uuid='instance-a', + workspace_uuid='workspace-a', + placement_generation=2, + public=True, + ) + + assert [task['id'] for task in result['tasks']] == [current.id] + assert 'id_index' not in result + assert ( + manager.get_task_by_id( + other_workspace.id, + instance_uuid='instance-a', + workspace_uuid='workspace-a', + placement_generation=2, + ) + is None + ) + assert ( + manager.get_task_by_id( + stale_generation.id, + instance_uuid='instance-a', + workspace_uuid='workspace-a', + placement_generation=2, + ) + is None + ) + + current.cancel() + other_workspace.cancel() + stale_generation.cancel() + @pytest.mark.asyncio async def test_cancel_by_scope(self): """Test cancel_by_scope cancels matching tasks."""