diff --git a/pyk/src/pyk/cterm/symbolic.py b/pyk/src/pyk/cterm/symbolic.py index f7c20cca794..64bcfec9771 100644 --- a/pyk/src/pyk/cterm/symbolic.py +++ b/pyk/src/pyk/cterm/symbolic.py @@ -314,6 +314,16 @@ def add_module(self, module: KFlatModule, name_as_id: bool = False) -> str: _kore_module = kflatmodule_to_kore(self._definition, module) return self._kore_client.add_module(_kore_module, name_as_id=name_as_id) + def set_client_label(self, label: str) -> None: + """Stamp `label` on every subsequent JSON-RPC request issued through this client. + + Called automatically by APRProver/ImpliesProver at the start of each + proof so that booster's per-line `{request: ...}` context self-identifies + which claim drove the work. Safe to call between proofs to re-tag the + same client for a different claim. + """ + self._kore_client.set_client_label(label) + def _smt_solver_error(self, err: SmtSolverError) -> CTermSMTError: kast = self.kore_to_kast(err.pattern) pretty_pattern = PrettyPrinter(self._definition).print(kast) @@ -326,6 +336,7 @@ def cterm_symbolic( definition_dir: Path, *, id: str | None = None, + client_label: str | None = None, port: int | None = None, kore_rpc_command: str | Iterable[str] | None = None, llvm_definition_dir: Path | None = None, @@ -365,7 +376,13 @@ def cterm_symbolic( simplify_each=simplify_each, no_post_exec_simplify=no_post_exec_simplify, ) as server: - with KoreClient('localhost', server.port, bug_report=bug_report, bug_report_id=id) as client: + with KoreClient( + 'localhost', + server.port, + bug_report=bug_report, + bug_report_id=id, + client_label=client_label, + ) as client: yield CTermSymbolic( client, definition, @@ -376,7 +393,9 @@ def cterm_symbolic( else: if port is None: raise ValueError('Missing port with start_server=False') - with KoreClient('localhost', port, bug_report=bug_report, bug_report_id=id) as client: + with KoreClient( + 'localhost', port, bug_report=bug_report, bug_report_id=id, client_label=client_label + ) as client: yield CTermSymbolic( client, definition, log_succ_rewrites=log_succ_rewrites, log_fail_rewrites=log_fail_rewrites ) diff --git a/pyk/src/pyk/kore/rpc.py b/pyk/src/pyk/kore/rpc.py index dff11c23b49..4d6966d46ac 100644 --- a/pyk/src/pyk/kore/rpc.py +++ b/pyk/src/pyk/kore/rpc.py @@ -193,6 +193,7 @@ def __init__( timeout: int | None = None, bug_report: BugReport | None = None, bug_report_id: str | None = None, + client_label: str | None = None, ): client_cache = {} self._clients = {} @@ -202,6 +203,7 @@ def __init__( timeout=timeout, bug_report=bug_report, bug_report_id=bug_report_id, + client_label=client_label, transport=default_transport, ) client_cache[(default_host, default_port)] = self._default_client @@ -212,7 +214,13 @@ def __init__( else: new_id = None if bug_report_id is None else bug_report_id + '_' + str(transport) new_client = JsonRpcClient( - host, port, timeout=timeout, bug_report=bug_report, bug_report_id=new_id, transport=transport + host, + port, + timeout=timeout, + bug_report=bug_report, + bug_report_id=new_id, + client_label=client_label, + transport=transport, ) self._update_clients(method, new_client) client_cache[(host, port)] = new_client @@ -237,6 +245,23 @@ def close(self) -> None: for client in clients: client.close() + def _unique_clients(self) -> list[JsonRpcClient]: + # Dispatch lists may share clients across methods (one JsonRpcClient per + # distinct (host, port)); collect each client only once by identity. + result = [self._default_client] + seen = {id(self._default_client)} + for clients in self._clients.values(): + for client in clients: + if id(client) not in seen: + seen.add(id(client)) + result.append(client) + return result + + def set_client_label(self, label: str) -> None: + """Push a new client_label onto every distinct underlying JsonRpcClient.""" + for client in self._unique_clients(): + client.set_client_label(label) + def request(self, method: str, **params: Any) -> dict[str, Any]: if method in self._clients: for client in self._clients[method]: @@ -253,6 +278,7 @@ class JsonRpcClient(ContextManager['JsonRpcClient']): _transport: Transport _req_id: int + _client_label: str _bug_report: BugReport | None _bug_report_id: str | None @@ -265,12 +291,17 @@ def __init__( timeout: int | None = None, bug_report: BugReport | None = None, bug_report_id: str | None = None, + client_label: str | None = None, transport: TransportType = TransportType.SINGLE_SOCKET, ): self._transport = self._create_transport(transport, host=host, port=port, timeout=timeout) self._req_id = 1 self._bug_report_id = bug_report_id self._bug_report = bug_report + # Stamped on every outgoing request-id as `{client_label}-NNN` so that + # booster's `{request: ...}` context lines self-identify the caller. + # Defaults to str(id(self)) so non-adopters keep the prior shape byte-for-byte. + self._client_label = client_label if client_label is not None else str(id(self)) @staticmethod def _create_transport(transport: TransportType, *, host: str, port: int, timeout: int | None) -> Transport: @@ -292,7 +323,7 @@ def close(self) -> None: self._transport.close() def request(self, method: str, **params: Any) -> dict[str, Any]: - req_id = f'{id(self)}-{self._req_id:03}' + req_id = f'{self._client_label}-{self._req_id:03}' self._req_id += 1 payload = { @@ -334,6 +365,14 @@ def _check(response: Mapping[str, Any]) -> None: assert response['error']['code'] not in {-32700, -32600}, 'Malformed JSON-RPC request' raise JsonRpcError(**response['error']) + def set_client_label(self, label: str) -> None: + """Set the prefix stamped on every subsequent request id. + + Persists until the next `set_client_label` (or close). Used by + APRProver/ImpliesProver to re-tag a shared client per claim. + """ + self._client_label = label + class KoreClientError(Exception, ABC): def __init__(self, message: str): @@ -895,6 +934,7 @@ def __init__( timeout: int | None = None, bug_report: BugReport | None = None, bug_report_id: str | None = None, + client_label: str | None = None, transport: TransportType = TransportType.SINGLE_SOCKET, dispatch: dict[str, list[tuple[str, int, TransportType]]] | None = None, ): @@ -908,9 +948,18 @@ def __init__( timeout=timeout, bug_report=bug_report, bug_report_id=bug_report_id, + client_label=client_label, dispatch=dispatch, ) + def set_client_label(self, label: str) -> None: + """Set the label stamped on every subsequent JSON-RPC request id. + + Persists until the next `set_client_label` (or close). Booster's + per-line `{request: }` context then self-identifies the caller. + """ + self._client.set_client_label(label) + def __enter__(self) -> KoreClient: return self diff --git a/pyk/src/pyk/proof/implies.py b/pyk/src/pyk/proof/implies.py index e5ca7d041b2..00608db5a50 100644 --- a/pyk/src/pyk/proof/implies.py +++ b/pyk/src/pyk/proof/implies.py @@ -466,7 +466,10 @@ def step_proof(self, step: ImpliesProofStep) -> list[ImpliesProofResult]: ] def init_proof(self, proof: ImpliesProof) -> None: - pass + # Stamp proof.id on every subsequent kore-RPC request so booster's + # per-line `{request: ...}` context self-identifies the claim driving + # the work — no PID join through pyk.log needed downstream. + self.kcfg_explore.cterm_symbolic.set_client_label(proof.id) def failure_info(self, proof: ImpliesProof) -> FailureInfo: # TODO add implementation diff --git a/pyk/src/pyk/proof/reachability.py b/pyk/src/pyk/proof/reachability.py index 928fe5c25a1..faa7d8f7c3a 100644 --- a/pyk/src/pyk/proof/reachability.py +++ b/pyk/src/pyk/proof/reachability.py @@ -756,6 +756,10 @@ def close(self) -> None: self.kcfg_explore.cterm_symbolic._kore_client.close() def init_proof(self, proof: APRProof) -> None: + # Stamp proof.id on every subsequent kore-RPC request so booster's + # per-line `{request: ...}` context self-identifies the claim driving + # the work — no PID join through pyk.log needed downstream. + self.kcfg_explore.cterm_symbolic.set_client_label(proof.id) main_module_name = self.main_module_name if self.extra_module: main_module_name = self.kcfg_explore.cterm_symbolic.add_module(self.extra_module, name_as_id=True) diff --git a/pyk/src/tests/unit/kore/test_client.py b/pyk/src/tests/unit/kore/test_client.py index 8842da2550d..a81f76aedef 100644 --- a/pyk/src/tests/unit/kore/test_client.py +++ b/pyk/src/tests/unit/kore/test_client.py @@ -90,7 +90,13 @@ def rpc_client(mock: Mock) -> MockClient: def kore_client(mock: Mock, mock_class: Mock) -> Iterator[KoreClient]: # noqa: N803 client = KoreClient('localhost', 3000) mock_class.assert_called_with( - 'localhost', 3000, timeout=None, bug_report=None, bug_report_id=None, transport=TransportType.SINGLE_SOCKET + 'localhost', + 3000, + timeout=None, + bug_report=None, + bug_report_id=None, + client_label=None, + transport=TransportType.SINGLE_SOCKET, ) assert client._client._default_client == mock yield client diff --git a/pyk/src/tests/unit/kore/test_client_label.py b/pyk/src/tests/unit/kore/test_client_label.py new file mode 100644 index 00000000000..7151470d94d --- /dev/null +++ b/pyk/src/tests/unit/kore/test_client_label.py @@ -0,0 +1,118 @@ +"""Unit tests for the caller-supplied `client_label` on JsonRpcClient / KoreClient. + +The label is stamped on every outgoing request-id (`{label}-NNN`) so booster's +per-line `{request: ...}` context self-identifies the caller (typically the +claim being discharged). `set_client_label(label)` mutates the prefix for all +subsequent requests; in normal pyk use, APRProver/ImpliesProver invoke this +automatically per proof so consumers never touch the API directly. +""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING +from unittest.mock import patch + +import pytest + +from pyk.kore.rpc import JsonRpcClient, KoreClient, SingleSocketTransport + +if TYPE_CHECKING: + from collections.abc import Iterator + from unittest.mock import Mock + + +@pytest.fixture +def mock_class() -> Iterator[Mock]: + patcher = patch('pyk.kore.rpc.SingleSocketTransport', spec=True) + yield patcher.start() + patcher.stop() + + +@pytest.fixture +def mock(mock_class: Mock) -> Mock: + mock = mock_class.return_value + assert isinstance(mock, SingleSocketTransport) + return mock # type: ignore + + +def _wire_capture(mock: Mock) -> list[dict]: + """Capture every outgoing payload and echo a success response.""" + captured: list[dict] = [] + + def respond(req: str, req_id: str, method_name: str) -> str: + payload = json.loads(req) + captured.append(payload) + return json.dumps({'jsonrpc': '2.0', 'id': payload['id'], 'result': {}}) + + mock.request.side_effect = respond + return captured + + +def test_default_label_uses_object_id(mock: Mock) -> None: + """Without a client_label, the request-id prefix is str(id(self)) — byte-stable with the legacy path.""" + captured = _wire_capture(mock) + client = JsonRpcClient('localhost', 3000) + expected_prefix = str(id(client)) + + client.request('execute') + client.request('simplify') + + assert [p['id'] for p in captured] == [f'{expected_prefix}-001', f'{expected_prefix}-002'] + + +def test_construction_label_is_used_as_prefix(mock: Mock) -> None: + captured = _wire_capture(mock) + client = JsonRpcClient('localhost', 3000, client_label='LEMMAS-SPEC.range-31') + + client.request('execute') + client.request('simplify') + + assert [p['id'] for p in captured] == ['LEMMAS-SPEC.range-31-001', 'LEMMAS-SPEC.range-31-002'] + + +def test_set_client_label_swaps_prefix_for_subsequent_requests(mock: Mock) -> None: + captured = _wire_capture(mock) + client = JsonRpcClient('localhost', 3000, client_label='claim-A') + + client.request('execute') + client.set_client_label('claim-B') + client.request('simplify') + client.request('implies') + + assert [p['id'] for p in captured] == ['claim-A-001', 'claim-B-002', 'claim-B-003'] + + +def test_set_client_label_persists_no_restoration(mock: Mock) -> None: + """The setter is permanent — there is no enclosing scope or restoration semantics.""" + captured = _wire_capture(mock) + client = JsonRpcClient('localhost', 3000, client_label='outer') + + client.set_client_label('inner') + client.request('execute') + client.request('simplify') + + assert [p['id'] for p in captured] == ['inner-001', 'inner-002'] + + +def test_kore_client_forwards_client_label(mock: Mock) -> None: + """KoreClient(client_label=...) plumbs through to the underlying JsonRpcClient.""" + captured = _wire_capture(mock) + kore_client = KoreClient('localhost', 3000, client_label='claim-X') + + assert kore_client._client._default_client._client_label == 'claim-X' + + kore_client._client._default_client.request('execute') + assert captured[-1]['id'] == 'claim-X-001' + + +def test_kore_client_set_client_label_forwards_to_underlying_clients(mock: Mock) -> None: + captured = _wire_capture(mock) + kore_client = KoreClient('localhost', 3000, client_label='construction-default') + + kore_client.set_client_label('claim-A') + kore_client._client._default_client.request('execute') + kore_client.set_client_label('claim-B') + kore_client._client._default_client.request('simplify') + + assert [p['id'] for p in captured] == ['claim-A-001', 'claim-B-002'] diff --git a/pyk/src/tests/unit/test_cterm_symbolic_client_label.py b/pyk/src/tests/unit/test_cterm_symbolic_client_label.py new file mode 100644 index 00000000000..804c8a9bc71 --- /dev/null +++ b/pyk/src/tests/unit/test_cterm_symbolic_client_label.py @@ -0,0 +1,93 @@ +"""Unit tests for the set_client_label plumbing on CTermSymbolic. + +CTermSymbolic.set_client_label(label) is the setter that APRProver/ImpliesProver +call automatically per proof so booster's per-line `{request: ...}` context +self-identifies the claim driving the work. Consumers normally do not call +this directly. +""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, cast +from unittest.mock import MagicMock, patch + +import pytest + +from pyk.cterm.symbolic import CTermSymbolic +from pyk.kore.rpc import KoreClient, SingleSocketTransport + +if TYPE_CHECKING: + from collections.abc import Iterator + from unittest.mock import Mock + + from pyk.kast.outer import KDefinition + + +@pytest.fixture +def mock_class() -> Iterator[Mock]: + patcher = patch('pyk.kore.rpc.SingleSocketTransport', spec=True) + yield patcher.start() + patcher.stop() + + +@pytest.fixture +def mock(mock_class: Mock) -> Mock: + m = mock_class.return_value + assert isinstance(m, SingleSocketTransport) + return m # type: ignore + + +def _wire_capture(mock: Mock) -> list[dict]: + captured: list[dict] = [] + + def respond(req: str, req_id: str, method_name: str) -> str: + payload = json.loads(req) + captured.append(payload) + return json.dumps({'jsonrpc': '2.0', 'id': payload['id'], 'result': {}}) + + mock.request.side_effect = respond + return captured + + +def _make_cterm_symbolic(client_label: str | None) -> tuple[CTermSymbolic, KoreClient]: + client = KoreClient('localhost', 3000, client_label=client_label) + # The definition is held by CTermSymbolic but never invoked by set_client_label, + # so a MagicMock satisfies the runtime requirement. + cterm = CTermSymbolic(client, cast('KDefinition', MagicMock())) + return cterm, client + + +def test_cterm_symbolic_set_client_label_swaps_prefix(mock: Mock) -> None: + captured = _wire_capture(mock) + cterm, client = _make_cterm_symbolic(client_label='construction-default') + + client._client._default_client.request('execute') + cterm.set_client_label('claim-A') + client._client._default_client.request('simplify') + cterm.set_client_label('claim-B') + client._client._default_client.request('implies') + + assert [p['id'] for p in captured] == [ + 'construction-default-001', + 'claim-A-002', + 'claim-B-003', + ] + + +def test_cterm_symbolic_serves_multiple_claims_in_sequence(mock: Mock) -> None: + """One CTermSymbolic discharging two claims in turn — prover stamps each label once.""" + captured = _wire_capture(mock) + cterm, client = _make_cterm_symbolic(client_label='unused-default') + + cterm.set_client_label('claim-A') + client._client._default_client.request('execute') + client._client._default_client.request('simplify') + cterm.set_client_label('claim-B') + client._client._default_client.request('execute') + + assert [p['id'] for p in captured] == [ + 'claim-A-001', + 'claim-A-002', + 'claim-B-003', + ] diff --git a/pyk/src/tests/unit/test_prover_client_label.py b/pyk/src/tests/unit/test_prover_client_label.py new file mode 100644 index 00000000000..ee20f6a3d4e --- /dev/null +++ b/pyk/src/tests/unit/test_prover_client_label.py @@ -0,0 +1,69 @@ +"""Auto-stamping of the kore-RPC `client_label` by APRProver / ImpliesProver. + +Both provers call `cterm_symbolic.set_client_label(proof.id)` from `init_proof` +so booster's per-line `{request: ...}` context self-identifies the claim +without the consumer touching the API. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from unittest.mock import MagicMock + +import pytest + +from pyk.kast.prelude.kbool import BOOL +from pyk.kast.prelude.kint import intToken +from pyk.kcfg.kcfg import KCFG +from pyk.proof.implies import EqualityProof, ImpliesProver +from pyk.proof.reachability import APRProof, APRProver + +from .test_kcfg import node, node_dicts + +if TYPE_CHECKING: + from pathlib import Path + + from pytest import TempPathFactory + + +@pytest.fixture(scope='function') +def proof_dir(tmp_path_factory: TempPathFactory) -> Path: + return tmp_path_factory.mktemp('proofs') + + +def test_apr_prover_init_proof_stamps_client_label(proof_dir: Path) -> None: + """APRProver.init_proof(proof) calls cterm_symbolic.set_client_label(proof.id) at the top.""" + kcfg_explore = MagicMock() + # APRProver.__init__ reads kcfg_explore.cterm_symbolic._definition.main_module_name; pin it. + kcfg_explore.cterm_symbolic._definition.main_module_name = 'TEST' + # init_proof iterates [proof.init, proof.target] and calls is_terminal on each; + # return False to avoid touching proof.add_terminal. + kcfg_explore.kcfg_semantics.is_terminal.return_value = False + + prover = APRProver(kcfg_explore=kcfg_explore) + proof = APRProof( + id='apr_proof_1', + kcfg=KCFG.from_dict({'nodes': node_dicts(1)}), + terminal=[], + init=node(1).id, + target=node(1).id, + logs={}, + proof_dir=proof_dir, + ) + + prover.init_proof(proof) + + kcfg_explore.cterm_symbolic.set_client_label.assert_called_once_with('apr_proof_1') + + +def test_implies_prover_init_proof_stamps_client_label(proof_dir: Path) -> None: + """ImpliesProver.init_proof(proof) calls cterm_symbolic.set_client_label(proof.id).""" + kcfg_explore = MagicMock() + proof = EqualityProof( + id='equality_proof_1', lhs_body=intToken(1), rhs_body=intToken(1), sort=BOOL, proof_dir=proof_dir + ) + prover = ImpliesProver(proof, kcfg_explore=kcfg_explore) + + prover.init_proof(proof) + + kcfg_explore.cterm_symbolic.set_client_label.assert_called_once_with('equality_proof_1')