feat: report independent instance and workspace identities (#2394)

* feat: report independent instance and workspace identities

* test: include workspace in OAuth callback fixture

* ci: pin production cloud adapter to Space release

* fix: preserve authenticated Workspace telemetry attribution

* ci: pin production cloud adapter to final Space release

---------

Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
This commit is contained in:
Hyu
2026-08-04 17:27:23 +08:00
committed by GitHub
parent 7820949d3a
commit c08bfc8ced
14 changed files with 196 additions and 28 deletions
+13 -1
View File
@@ -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
@@ -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."""
+21 -5
View File
@@ -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
@@ -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."""
@@ -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',
}