diff --git a/.github/workflows/deploy-prod.yml b/.github/workflows/deploy-prod.yml index acf33d0e5..c8500c4e2 100644 --- a/.github/workflows/deploy-prod.yml +++ b/.github/workflows/deploy-prod.yml @@ -15,7 +15,7 @@ concurrency: env: CORE_IMAGE: ${{ secrets.DOCKER_USERNAME }}/langbot CLOUD_IMAGE: ${{ secrets.DOCKER_USERNAME }}/langbot-cloud-core - SPACE_REF: 178b1634c104e635c706e1fb4ca37cb3de073cf9 + SPACE_REF: 58253c53933f95d81b035fbe2efedb55b6c1a82b jobs: build-and-deploy: diff --git a/src/langbot/pkg/api/http/controller/groups/user.py b/src/langbot/pkg/api/http/controller/groups/user.py index e8e450911..6cfd5f8a7 100644 --- a/src/langbot/pkg/api/http/controller/groups/user.py +++ b/src/langbot/pkg/api/http/controller/groups/user.py @@ -1,6 +1,7 @@ import quart import argon2 import asyncio +import datetime import uuid from urllib.parse import parse_qs, urlsplit @@ -218,7 +219,22 @@ class UserRouterGroup(group.RouterGroup): try: consumed_state = await self.ap.user_service.consume_space_oauth_state_details(state, 'login') # Exchange code for tokens - token_data = await self.ap.space_service.exchange_oauth_code(code) + launch_workspace_uuid = consumed_state.launch_workspace_uuid + workspace_uuids = [launch_workspace_uuid] if launch_workspace_uuid else [] + workspace_created_ats: dict[str, int] = {} + if not workspace_uuids and getattr(getattr(self.ap, 'deployment', None), 'mode', 'oss') != 'cloud': + binding = await self.ap.workspace_service.get_execution_binding() + workspace_uuids = [binding.workspace_uuid] + workspace_created_at = binding.workspace_created_at + if workspace_created_at is not None: + if workspace_created_at.tzinfo is None: + workspace_created_at = workspace_created_at.replace(tzinfo=datetime.UTC) + workspace_created_ats[binding.workspace_uuid] = int(workspace_created_at.timestamp()) + token_data = await self.ap.space_service.exchange_oauth_code( + code, + workspace_uuids, + workspace_created_ats, + ) access_token = token_data.get('access_token') refresh_token = token_data.get('refresh_token') expires_in = token_data.get('expires_in', 0) @@ -231,7 +247,6 @@ class UserRouterGroup(group.RouterGroup): access_token, refresh_token, expires_in ) - launch_workspace_uuid = consumed_state.launch_workspace_uuid if launch_workspace_uuid: try: access = await self.ap.workspace_collaboration_service.resolve_account_workspace( diff --git a/src/langbot/pkg/api/http/service/space.py b/src/langbot/pkg/api/http/service/space.py index 855b009c1..a0fc7dfb8 100644 --- a/src/langbot/pkg/api/http/service/space.py +++ b/src/langbot/pkg/api/http/service/space.py @@ -59,6 +59,10 @@ class SpaceService: result_list = result.all() return result_list[0] if result_list else None + async def get_valid_access_token(self, user_email: str) -> str | None: + """Return a current Space bearer, refreshing and persisting it when needed.""" + return await self._ensure_valid_token(user_email) + async def _ensure_valid_token(self, user_email: str) -> str | None: """Ensure access token is valid, refresh if expired. Returns valid access_token or None.""" user_obj = await self._get_user_by_email(user_email) @@ -117,7 +121,12 @@ class SpaceService: params['state'] = state return f'{authorize_url}?{urlencode(params)}' - async def exchange_oauth_code(self, code: str) -> typing.Dict: + async def exchange_oauth_code( + self, + code: str, + workspace_uuids: list[str] | None = None, + workspace_created_ats: dict[str, int] | None = None, + ) -> typing.Dict: """Exchange OAuth authorization code for tokens""" from langbot.pkg.utils import constants @@ -127,7 +136,14 @@ class SpaceService: session = httpclient.get_session() async with session.post( f'{space_url}/api/v1/accounts/oauth/token', - json={'code': code, 'instance_id': constants.instance_id}, + json={ + 'code': code, + 'instance_id': constants.instance_id, + # Sending an explicit empty list tells new Space servers not to + # synthesize a legacy instance-derived Workspace binding. + 'workspace_uuids': workspace_uuids if workspace_uuids is not None else [], + 'workspace_created_ats': workspace_created_ats or {}, + }, ) as response: if response.status != 200: error = await httpclient.read_text_limited(response) diff --git a/src/langbot/pkg/api/http/service/user.py b/src/langbot/pkg/api/http/service/user.py index d1ae53608..28c23b32e 100644 --- a/src/langbot/pkg/api/http/service/user.py +++ b/src/langbot/pkg/api/http/service/user.py @@ -779,8 +779,27 @@ class UserService: local_account = await self.get_user_by_email(user_email) if local_account is None: raise ValueError('User not found') - # Exchange code for tokens - token_data = await self.ap.space_service.exchange_oauth_code(code) + # Exchange code for tokens and bind both installation and the active + # OSS Workspace as independent identities. + workspace_service = getattr(self.ap, 'workspace_service', None) + if workspace_service is not None: + binding = await workspace_service.get_execution_binding() + created_at = binding.workspace_created_at + created_ts = ( + int(created_at.replace(tzinfo=datetime.timezone.utc).timestamp()) + if created_at.tzinfo is None + else int(created_at.timestamp()) + ) + token_data = await self.ap.space_service.exchange_oauth_code( + code, + [binding.workspace_uuid], + {binding.workspace_uuid: created_ts}, + ) + else: + # Compatibility for early/bootstrap call sites that have not wired + # WorkspaceService yet; old Space servers still derive the legacy + # Workspace identity from instance_id when the field is omitted. + token_data = await self.ap.space_service.exchange_oauth_code(code) access_token = token_data.get('access_token') refresh_token = token_data.get('refresh_token') expires_in = token_data.get('expires_in', 0) diff --git a/src/langbot/pkg/telemetry/heartbeat.py b/src/langbot/pkg/telemetry/heartbeat.py index 2dd2b2afd..45dcbf268 100644 --- a/src/langbot/pkg/telemetry/heartbeat.py +++ b/src/langbot/pkg/telemetry/heartbeat.py @@ -120,6 +120,7 @@ async def build_heartbeat_payload( ap: core_app.Application, *, workspace_uuid: str, + workspace_create_ts: int = 0, workspace_resource: WorkspaceResourceSnapshot | None = None, ) -> dict: """Collect one anonymous Workspace profile snapshot.""" @@ -212,7 +213,9 @@ async def build_heartbeat_payload( 'event_type': 'instance_heartbeat', 'query_id': '', 'version': constants.semantic_version, + 'instance_id': constants.instance_id, 'workspace_uuid': workspace_uuid, + 'workspace_create_ts': workspace_create_ts, 'instance_create_ts': constants.instance_create_ts, 'edition': constants.edition, 'features': features, @@ -220,10 +223,24 @@ async def build_heartbeat_payload( } +def _workspace_created_timestamp(created_at: datetime | None) -> int: + if created_at is None: + return 0 + if created_at.tzinfo is None: + # SQLAlchemy may return persisted UTC values without tzinfo. Never + # reinterpret them in the host's local timezone. + created_at = created_at.replace(tzinfo=timezone.utc) + return int(created_at.timestamp()) + + async def build_heartbeat_payloads(ap: core_app.Application) -> list[dict]: """Build one heartbeat per active Workspace.""" bindings = await ap.workspace_service.list_active_execution_bindings() workspace_uuids = sorted({binding.workspace_uuid for binding in bindings}) + workspace_create_ts = { + binding.workspace_uuid: _workspace_created_timestamp(getattr(binding, 'workspace_created_at', None)) + for binding in bindings + } resources = { resource['workspace_uuid']: resource for resource in await _cloud_workspace_resource_counts(ap, bindings) } @@ -231,6 +248,7 @@ async def build_heartbeat_payloads(ap: core_app.Application) -> list[dict]: await build_heartbeat_payload( ap, workspace_uuid=workspace_uuid, + workspace_create_ts=workspace_create_ts.get(workspace_uuid, 0), workspace_resource=resources.get(workspace_uuid), ) for workspace_uuid in workspace_uuids diff --git a/src/langbot/pkg/telemetry/identity.py b/src/langbot/pkg/telemetry/identity.py index a1320f1d4..68e1ba8ad 100644 --- a/src/langbot/pkg/telemetry/identity.py +++ b/src/langbot/pkg/telemetry/identity.py @@ -4,13 +4,19 @@ import typing class WorkspaceExecutionContext(typing.Protocol): + @property + def instance_uuid(self) -> str: ... + @property def workspace_uuid(self) -> str: ... def workspace_identity(execution_context: WorkspaceExecutionContext) -> dict[str, str]: - """Build the canonical telemetry identity for one Workspace execution.""" + """Build both first-class telemetry identities for one execution.""" + instance_id = execution_context.instance_uuid.strip() workspace_uuid = execution_context.workspace_uuid.strip() + if not instance_id: + raise ValueError('Telemetry execution instance ID is empty') if not workspace_uuid: raise ValueError('Telemetry execution Workspace UUID is empty') - return {'workspace_uuid': workspace_uuid} + return {'instance_id': instance_id, 'workspace_uuid': workspace_uuid} diff --git a/src/langbot/pkg/telemetry/telemetry.py b/src/langbot/pkg/telemetry/telemetry.py index 28bcf1a0a..3d3e29898 100644 --- a/src/langbot/pkg/telemetry/telemetry.py +++ b/src/langbot/pkg/telemetry/telemetry.py @@ -136,12 +136,31 @@ class TelemetryManager: try: # Use asyncio.wait_for to ensure we always bound the total time telemetry_token = os.getenv('LANGBOT_TELEMETRY_INGEST_TOKEN', '').strip() + headers: dict[str, str] = {} if telemetry_token: - request = client.post( - url, - json=sanitized, - headers={'X-LangBot-Telemetry-Token': telemetry_token}, - ) + headers['X-LangBot-Telemetry-Token'] = telemetry_token + else: + workspace_uuid = str(sanitized.get('workspace_uuid', '')).strip() + user_service = getattr(self.ap, 'user_service', None) + if workspace_uuid and user_service is not None: + try: + owner = await user_service.get_workspace_owner(workspace_uuid) + owner_email = str(getattr(owner, 'user', '') or '').strip() + space_service = getattr(self.ap, 'space_service', None) + access_token = ( + await space_service.get_valid_access_token(owner_email) + if owner_email and space_service is not None + else None + ) + access_token = str(access_token or '').strip() + if access_token: + headers['Authorization'] = f'Bearer {access_token}' + except Exception: + self.ap.logger.debug( + 'Could not resolve authenticated telemetry reporter', exc_info=True + ) + if headers: + request = client.post(url, json=sanitized, headers=headers) else: request = client.post(url, json=sanitized) resp = await asyncio.wait_for(request, timeout=10 + 1) diff --git a/src/langbot/pkg/workspace/entities.py b/src/langbot/pkg/workspace/entities.py index 77ef4914f..f50d7542d 100644 --- a/src/langbot/pkg/workspace/entities.py +++ b/src/langbot/pkg/workspace/entities.py @@ -1,5 +1,6 @@ from __future__ import annotations +import datetime from dataclasses import dataclass @@ -12,3 +13,4 @@ class WorkspaceExecutionBinding: placement_generation: int write_fenced: bool state: str + workspace_created_at: datetime.datetime | None = None diff --git a/src/langbot/pkg/workspace/service.py b/src/langbot/pkg/workspace/service.py index b4eb0a845..96439b2af 100644 --- a/src/langbot/pkg/workspace/service.py +++ b/src/langbot/pkg/workspace/service.py @@ -283,6 +283,7 @@ class WorkspaceService: placement_generation=execution_state.active_generation, write_fenced=execution_state.write_fenced, state=execution_state.state, + workspace_created_at=workspace.created_at, ) binding = await self._run(operation, session=session) diff --git a/tests/integration/api/test_user_space_oauth.py b/tests/integration/api/test_user_space_oauth.py index 16c86156d..6fc1dfe07 100644 --- a/tests/integration/api/test_user_space_oauth.py +++ b/tests/integration/api/test_user_space_oauth.py @@ -2,6 +2,7 @@ from __future__ import annotations +import datetime from types import SimpleNamespace from unittest.mock import AsyncMock, Mock from urllib.parse import parse_qs, urlsplit @@ -14,6 +15,7 @@ from langbot.pkg.api.http.controller.groups.user import UserRouterGroup pytestmark = pytest.mark.integration WORKSPACE_UUID = '11111111-1111-4111-8111-111111111111' +WORKSPACE_CREATED_AT = datetime.datetime(2026, 1, 2, 3, 4, 5, tzinfo=datetime.UTC) @pytest.fixture @@ -58,6 +60,12 @@ async def space_oauth_api(): return_value={'account_uuid': 'account-a', 'workspace_uuid': WORKSPACE_UUID} ) application.workspace_collaboration_service.resolve_account_workspace = AsyncMock(return_value=access) + application.workspace_service.get_execution_binding = AsyncMock( + return_value=SimpleNamespace( + workspace_uuid=WORKSPACE_UUID, + workspace_created_at=WORKSPACE_CREATED_AT, + ) + ) application.space_service.get_oauth_authorize_url = Mock( side_effect=lambda redirect_uri, state: f'https://space.example/authorize?state={state}' ) @@ -234,7 +242,11 @@ async def test_login_callback_requires_and_consumes_server_state(space_oauth_api assert response.status_code == 200 assert (await response.get_json())['data']['token'] == 'space-login-token' application.user_service.consume_space_oauth_state_details.assert_awaited_once_with('opaque-login-state', 'login') - application.space_service.exchange_oauth_code.assert_awaited_once_with('oauth-code') + application.space_service.exchange_oauth_code.assert_awaited_once_with( + 'oauth-code', + [WORKSPACE_UUID], + {WORKSPACE_UUID: int(WORKSPACE_CREATED_AT.timestamp())}, + ) @pytest.mark.asyncio diff --git a/tests/unit_tests/api/service/test_space_service.py b/tests/unit_tests/api/service/test_space_service.py index d92265743..f16b609a9 100644 --- a/tests/unit_tests/api/service/test_space_service.py +++ b/tests/unit_tests/api/service/test_space_service.py @@ -25,6 +25,7 @@ import time from langbot.pkg.api.http.service.space import SpaceService from langbot.pkg.entity.persistence.user import User +from langbot.pkg.utils import constants pytestmark = pytest.mark.asyncio @@ -573,10 +574,20 @@ class TestSpaceServiceExchangeOAuthCode: mock_session_obj.post.return_value.__aexit__ = AsyncMock(return_value=None) # Execute - result = await service.exchange_oauth_code('auth_code') + result = await service.exchange_oauth_code( + 'auth_code', + ['workspace-1'], + {'workspace-1': 1_700_000_000}, + ) # Verify assert result['access_token'] == 'new_access_token' + assert mock_session_obj.post.call_args.kwargs['json'] == { + 'code': 'auth_code', + 'instance_id': constants.instance_id, + 'workspace_uuids': ['workspace-1'], + 'workspace_created_ats': {'workspace-1': 1_700_000_000}, + } async def test_exchange_oauth_code_api_error(self): """Raises ValueError on API error.""" diff --git a/tests/unit_tests/telemetry/test_heartbeat.py b/tests/unit_tests/telemetry/test_heartbeat.py index 3d7e60620..86b70094b 100644 --- a/tests/unit_tests/telemetry/test_heartbeat.py +++ b/tests/unit_tests/telemetry/test_heartbeat.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +from datetime import datetime, timezone from types import SimpleNamespace import pytest @@ -14,6 +15,12 @@ def get_heartbeat_module(): return import_module('langbot.pkg.telemetry.heartbeat') +def test_workspace_created_timestamp_treats_naive_database_values_as_utc(): + heartbeat = get_heartbeat_module() + created_at = datetime(2026, 8, 4, 0, 0, 0) + assert heartbeat._workspace_created_timestamp(created_at) == 1785801600 + + def make_app(): ap = Mock() ap.instance_config = Mock() @@ -57,15 +64,17 @@ def make_app(): class TestBuildHeartbeatPayload: @pytest.mark.asyncio - async def test_payload_shape(self): + async def test_payload_shape(self, monkeypatch): heartbeat = get_heartbeat_module() + monkeypatch.setattr(heartbeat.constants, 'instance_id', 'instance-test') ap = make_app() payload = await heartbeat.build_heartbeat_payload(ap, workspace_uuid='workspace-a') assert payload['event_type'] == 'instance_heartbeat' assert payload['query_id'] == '' assert payload['workspace_uuid'] == 'workspace-a' - assert 'instance_id' not in payload + assert payload['instance_id'] + assert payload['workspace_create_ts'] == 0 assert 'instance_create_ts' in payload assert 'timestamp' in payload f = payload['features'] @@ -100,8 +109,9 @@ class TestBuildHeartbeatPayload: assert payload['features']['pipeline_count'] == -1 @pytest.mark.asyncio - async def test_cloud_counts_loaded_registries_without_tenant_sql(self): + async def test_cloud_counts_loaded_registries_without_tenant_sql(self, monkeypatch): heartbeat = get_heartbeat_module() + monkeypatch.setattr(heartbeat.constants, 'instance_id', 'instance-test') ap = make_app() ap.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime') ap.persistence_mgr.execute_async = AsyncMock( @@ -140,7 +150,11 @@ class TestBuildHeartbeatPayload: ap.workspace_service.list_active_execution_bindings = AsyncMock( return_value=[ SimpleNamespace(workspace_uuid='workspace-a', placement_generation=7), - SimpleNamespace(workspace_uuid='workspace-b', placement_generation=9), + SimpleNamespace( + workspace_uuid='workspace-b', + placement_generation=9, + workspace_created_at=datetime(2026, 8, 4, tzinfo=timezone.utc), + ), ], ) ap.platform_mgr._bots_by_key[('instance-a', 'workspace-b', 'bot-b')] = SimpleNamespace( @@ -150,7 +164,9 @@ class TestBuildHeartbeatPayload: payloads = await heartbeat.build_heartbeat_payloads(ap) assert [payload['workspace_uuid'] for payload in payloads] == ['workspace-a', 'workspace-b'] - assert all('instance_id' not in payload for payload in payloads) + assert all(payload['instance_id'] for payload in payloads) + assert payloads[0]['workspace_create_ts'] == 0 + assert payloads[1]['workspace_create_ts'] == 1785801600 by_workspace = {payload['workspace_uuid']: payload['features'] for payload in payloads} assert by_workspace['workspace-a']['pipeline_count'] == 2 assert by_workspace['workspace-a']['mcp_server_count'] == 3 diff --git a/tests/unit_tests/telemetry/test_telemetry.py b/tests/unit_tests/telemetry/test_telemetry.py index 2242a1aa4..291264f61 100644 --- a/tests/unit_tests/telemetry/test_telemetry.py +++ b/tests/unit_tests/telemetry/test_telemetry.py @@ -596,6 +596,36 @@ class TestTelemetryManagedRuntimeAuthentication: assert captured['headers'] == {'X-LangBot-Telemetry-Token': 'managed-runtime-secret'} +class TestAuthenticatedWorkspaceReporter: + @pytest.mark.asyncio + async def test_workspace_owner_access_token_is_sent_as_bearer(self): + telemetry = get_telemetry_module() + mock_app = Mock() + mock_app.logger = Mock() + mock_app.user_service = Mock() + mock_app.user_service.get_workspace_owner = AsyncMock( + return_value=Mock(user='owner@example.com', space_access_token='expired-token') + ) + mock_app.space_service = Mock() + mock_app.space_service.get_valid_access_token = AsyncMock(return_value='refreshed-workspace-owner-token') + manager = telemetry.TelemetryManager(mock_app) + manager.telemetry_config = {'url': 'https://example.com'} + + response = Mock(status_code=200, text='') + response.json = Mock(return_value={'code': 0}) + mock_client = Mock() + mock_client.post = Mock(return_value=response) + + with patch.object(httpx, 'AsyncClient', return_value=mock_client): + await manager.send({'query_id': 'q-1', 'workspace_uuid': 'workspace-1'}) + + mock_app.user_service.get_workspace_owner.assert_awaited_once_with('workspace-1') + mock_app.space_service.get_valid_access_token.assert_awaited_once_with('owner@example.com') + assert mock_client.post.call_args.kwargs['headers'] == { + 'Authorization': 'Bearer refreshed-workspace-owner-token' + } + + class TestStartSendTask: """Tests for start_send_task() method.""" diff --git a/tests/unit_tests/telemetry/test_workspace_identity.py b/tests/unit_tests/telemetry/test_workspace_identity.py index efbd68a11..24f263d90 100644 --- a/tests/unit_tests/telemetry/test_workspace_identity.py +++ b/tests/unit_tests/telemetry/test_workspace_identity.py @@ -7,25 +7,28 @@ from types import SimpleNamespace def test_standard_oss_instance_id_aligns_to_embedded_uuid(): from langbot.pkg.workspace.identity import workspace_uuid_from_instance_id - instance_uuid = "a711d9e4-0953-443f-a0e9-7dd50193a79f" + instance_uuid = 'a711d9e4-0953-443f-a0e9-7dd50193a79f' assert workspace_uuid_from_instance_id(instance_uuid) == instance_uuid - assert workspace_uuid_from_instance_id(f"instance_{instance_uuid}") == instance_uuid + assert workspace_uuid_from_instance_id(f'instance_{instance_uuid}') == instance_uuid def test_custom_legacy_instance_id_maps_to_stable_valid_uuid(): from langbot.pkg.workspace.identity import workspace_uuid_from_instance_id - first = workspace_uuid_from_instance_id("instance_migration_test") - second = workspace_uuid_from_instance_id("instance_migration_test") + first = workspace_uuid_from_instance_id('instance_migration_test') + second = workspace_uuid_from_instance_id('instance_migration_test') assert first == second assert str(uuid.UUID(first)) == first -def test_query_telemetry_identity_uses_execution_workspace_only(): +def test_query_telemetry_identity_reports_instance_and_workspace(): from langbot.pkg.telemetry.identity import workspace_identity - identity = workspace_identity(SimpleNamespace(workspace_uuid="workspace-a", instance_uuid="instance-a")) + identity = workspace_identity(SimpleNamespace(workspace_uuid='workspace-a', instance_uuid='instance-a')) - assert identity == {"workspace_uuid": "workspace-a"} + assert identity == { + 'instance_id': 'instance-a', + 'workspace_uuid': 'workspace-a', + }