"""Regression tests for P1 PII masking before engine requests.""" from __future__ import annotations import json import unittest from typing import Any from unittest.mock import patch from .contracts.engine_gateway import EngineGatewaySseLineDecoder from .engine_client import EngineClient, GenerateResponse from .services import guardrail, orchestrator, persona, state_machine RAW_PHONE = "010-1234-5678" RAW_EMAIL = "test@example.com" RAW_RRN = "990101-1234567" RAW_TEXT = f"My phone is {RAW_PHONE}, email {RAW_EMAIL}, and RRN {RAW_RRN}." RAW_VALUES = (RAW_PHONE, RAW_EMAIL, RAW_RRN) MASK_VALUES = ("[PHONE]", "[EMAIL]", "[RRN]") RAW_KO_NAME = "김서연" RAW_KO_ORG = "한신대학교" RAW_KO_DEPT = "상담심리학과" RAW_KO_TEXT = ( f"내담자 {RAW_KO_NAME}은 {RAW_KO_ORG} {RAW_KO_DEPT} 학생이고 " "연락은 하지 말아 주세요." ) RAW_KO_VALUES = (RAW_KO_NAME, RAW_KO_ORG, RAW_KO_DEPT) MASK_KO_VALUES = ("[NAME]", "[ORG]") def _json_blob(value: Any) -> str: return json.dumps(value, ensure_ascii=False, sort_keys=True, default=str) def _message_blob(messages: object) -> str: return "\n".join(message.content for message in messages) # type: ignore[attr-defined] def _initial_state() -> state_machine.SessionState: card = persona.P1 return state_machine.init_state( params=card.openness_params(), ) def _prepare_context() -> orchestrator.TurnContext: return orchestrator.prepare_turn( session_id="masking-session", case_id="masking-case", card=persona.P1, state=_initial_state(), learner_text=RAW_TEXT, memory=orchestrator.TurnMemory( recent_turns=[ { "speaker": "counselor", "text": "Previous learner contact was already masked: [PHONE] [EMAIL] [RRN].", } ], ), ) def _assert_no_raw_pii(test: unittest.TestCase, value: object) -> None: blob = _json_blob(value) for raw in RAW_VALUES: test.assertNotIn(raw, blob) def _assert_masked_pii_present(test: unittest.TestCase, value: object) -> None: blob = _json_blob(value) for masked in MASK_VALUES: test.assertIn(masked, blob) def _assert_no_raw_ko_pii(test: unittest.TestCase, value: object) -> None: blob = _json_blob(value) for raw in RAW_KO_VALUES: test.assertNotIn(raw, blob) def _assert_masked_ko_pii_present(test: unittest.TestCase, value: object) -> None: blob = _json_blob(value) for masked in MASK_KO_VALUES: test.assertIn(masked, blob) class CaptureGenerateEngine: def __init__(self, text: str = "Masked engine reply.") -> None: self.request = None self.payload: dict[str, Any] | None = None self._payload_builder = EngineClient(base_url="http://engine.test") self.text = text self.requests: list[Any] = [] async def generate(self, req): self.request = req self.requests.append(req) self.payload = self._payload_builder._payload(req) return GenerateResponse( text=self.text, model="fake-model", provider="fake-provider", tokens_in=3, tokens_out=4, cost_usd=0.0, ) class CaptureStreamEngine: engine_mode = "fake-provider" default_model = "fake-model" def __init__(self, chunks: list[str] | None = None) -> None: self.request = None self.payload: dict[str, Any] | None = None self._payload_builder = EngineClient(base_url="http://engine.test") self.chunks = chunks or ["Masked stream reply."] async def stream(self, req): self.request = req self.payload = self._payload_builder._payload(req) for chunk in self.chunks: yield "event: token" yield "data: " + json.dumps({"text": chunk}, ensure_ascii=False) yield "event: done" yield ( 'data: {"provider":"fake-provider","model":"fake-model",' '"tokens_in":5,"tokens_out":6,"cost_usd":0.0}' ) async def stream_packets(self, req): decoder = EngineGatewaySseLineDecoder() async for raw in self.stream(req): packet = decoder.feed_line(raw) if packet is not None: yield packet class OrchestratorMaskingGateTest(unittest.IsolatedAsyncioTestCase): def setUp(self) -> None: self.presidio_patch = patch.object( guardrail, "_try_load_presidio", return_value=(None, None), ) self.presidio_patch.start() guardrail.set_ko_pii_recognizer(None) self.addCleanup(self.presidio_patch.stop) self.addCleanup(guardrail.set_ko_pii_recognizer, None) def test_prepare_turn_keeps_raw_text_but_builds_masked_engine_messages(self) -> None: ctx = _prepare_context() self.assertEqual(ctx.learner_text_raw, RAW_TEXT) for raw in RAW_VALUES: self.assertIn(raw, ctx.learner_text_raw) self.assertNotIn(raw, ctx.learner_text_masked) self.assertNotIn(raw, _message_blob(ctx.messages)) for masked in MASK_VALUES: self.assertIn(masked, ctx.learner_text_masked) self.assertIn(masked, _message_blob(ctx.messages)) def test_prepare_turn_masks_raw_pii_from_context_inputs(self) -> None: ctx = orchestrator.prepare_turn( session_id="masking-session", case_id="masking-case", card=persona.P1, state=_initial_state(), learner_text="Current text has no identifiers.", memory=orchestrator.TurnMemory( recall_summary=f"Recall mentioned {RAW_PHONE}.", pinned_facts=[f"Pinned email {RAW_EMAIL}."], recent_turns=[ {"speaker": "counselor", "text": f"Previous raw RRN {RAW_RRN}."}, ], ), ) blob = _message_blob(ctx.messages) for raw in RAW_VALUES: self.assertNotIn(raw, blob) for masked in MASK_VALUES: self.assertIn(masked, blob) def test_mask_pii_masks_korean_name_and_institution_context(self) -> None: masked = guardrail.mask_pii( f"이름: {RAW_KO_NAME}, 소속은 {RAW_KO_ORG} {RAW_KO_DEPT}입니다." ) self.assertFalse(masked.used_presidio) self.assertIn("NAME", masked.entities) self.assertIn("ORG", masked.entities) for raw in RAW_KO_VALUES: self.assertNotIn(raw, masked.text_masked) self.assertIn("[NAME]", masked.text_masked) self.assertGreaterEqual(masked.text_masked.count("[ORG]"), 2) def test_mask_pii_does_not_mask_common_korean_context_words_as_names(self) -> None: masked = guardrail.mask_pii("학교 가는 게 힘들고 엄마랑 친구 이야기를 하면 불안해요.") self.assertEqual(masked.text_masked, "학교 가는 게 힘들고 엄마랑 친구 이야기를 하면 불안해요.") self.assertNotIn("NAME", masked.entities) self.assertNotIn("ORG", masked.entities) def test_mask_pii_does_not_mask_career_topic_as_name(self) -> None: masked = guardrail.mask_pii("아직 정해진 건 없어요. 진로는 그대로고 엄마한테 말도 못 했어요.") self.assertEqual( masked.text_masked, "아직 정해진 건 없어요. 진로는 그대로고 엄마한테 말도 못 했어요.", ) self.assertNotIn("NAME", masked.entities) def test_prepare_turn_masks_korean_pii_from_engine_messages(self) -> None: ctx = orchestrator.prepare_turn( session_id="masking-session", case_id="masking-case", card=persona.P1, state=_initial_state(), learner_text=RAW_KO_TEXT, memory=orchestrator.TurnMemory( recall_summary=f"지난 회기 요약에 {RAW_KO_NAME}과 {RAW_KO_ORG}가 남아 있었다.", pinned_facts=[f"소속 {RAW_KO_DEPT}"], recent_turns=[ {"speaker": "counselor", "text": f"{RAW_KO_NAME} 씨가 상담실에 왔다."}, ], ), ) blob = _message_blob(ctx.messages) _assert_no_raw_ko_pii(self, blob) _assert_masked_ko_pii_present(self, blob) for raw in RAW_KO_VALUES: self.assertIn(raw, ctx.learner_text_raw) self.assertNotIn(raw, ctx.learner_text_masked) def test_prepare_turn_applies_optional_ko_recognizer_before_regex_fallback(self) -> None: class FakeKoRecognizer: def analyze(self, text: str): spans = [] for entity_type, value in ( ("NAME", "보라별"), ("ORG", "미래학교상담연구랩"), ): start = text.index(value) spans.append(guardrail.PiiEntitySpan(entity_type, start, start + len(value))) return spans guardrail.set_ko_pii_recognizer(FakeKoRecognizer()) raw_text = "별명은 보라별이고 기관은 미래학교상담연구랩입니다. 전화는 010-1234-5678입니다." ctx = orchestrator.prepare_turn( session_id="masking-session", case_id="masking-case", card=persona.P1, state=_initial_state(), learner_text=raw_text, ) blob = _message_blob(ctx.messages) self.assertIn("보라별", ctx.learner_text_raw) self.assertIn("미래학교상담연구랩", ctx.learner_text_raw) for raw in ("보라별", "미래학교상담연구랩", "010-1234-5678"): self.assertNotIn(raw, ctx.learner_text_masked) self.assertNotIn(raw, blob) for masked in ("[NAME]", "[ORG]", "[PHONE]"): self.assertIn(masked, ctx.learner_text_masked) self.assertIn(masked, blob) def test_prepare_turn_threads_theory_mode_into_engine_messages(self) -> None: ctx = orchestrator.prepare_turn( session_id="theory-session", case_id="theory-case", card=persona.P2, state=_initial_state(), learner_text="그냥 아무것도 하기 싫어요.", theory_mode="cbt", ) blob = _message_blob(ctx.messages) self.assertIn("[L3-T 이론모드: CBT]", blob) self.assertIn("자동적 사고", blob) self.assertIn("행동활성화", blob) def test_prepare_turn_maps_recent_turns_from_client_ai_perspective(self) -> None: ctx = orchestrator.prepare_turn( session_id="history-session", case_id="history-case", card=persona.P1, state=_initial_state(), learner_text="그 말을 듣고 어떤 생각이 들었나요?", memory=orchestrator.TurnMemory( recent_turns=[ {"speaker": "counselor", "text": "왜 상담에 오게 됐나요?"}, {"speaker": "client", "text": "엄마가 가보라고 해서요."}, ], ), ) recent = ctx.messages[-3:-1] self.assertEqual([message.role for message in recent], ["user", "assistant"]) self.assertEqual(recent[0].content, "왜 상담에 오게 됐나요?") self.assertEqual(recent[1].content, "엄마가 가보라고 해서요.") async def test_run_turn_generate_sends_only_masked_engine_payload(self) -> None: ctx = _prepare_context() engine = CaptureGenerateEngine() audit_payloads: list[dict[str, Any]] = [] async def audit_hook(payload: dict[str, Any]) -> None: audit_payloads.append(payload) await orchestrator.run_turn_generate( ctx, engine, # type: ignore[arg-type] audit_hook=audit_hook, ) self.assertIsNotNone(engine.request) self.assertIsNotNone(engine.payload) _assert_no_raw_pii(self, engine.request.messages) _assert_no_raw_pii(self, engine.payload) _assert_masked_pii_present(self, engine.request.messages) _assert_masked_pii_present(self, engine.payload) self.assertEqual(ctx.learner_text_raw, RAW_TEXT) self.assertEqual(len(audit_payloads), 1) self.assertEqual(audit_payloads[0]["provider"], "fake-provider") self.assertEqual(audit_payloads[0]["model"], "fake-model") _assert_no_raw_pii(self, audit_payloads) for key in ("messages", "prompt", "text"): self.assertNotIn(key, audit_payloads[0]) async def test_run_turn_stream_sends_only_masked_engine_payload(self) -> None: ctx = _prepare_context() engine = CaptureStreamEngine() audit_payloads: list[dict[str, Any]] = [] async def audit_hook(payload: dict[str, Any]) -> None: audit_payloads.append(payload) events = [ event async for event in orchestrator.run_turn_stream( ctx, engine, # type: ignore[arg-type] audit_hook=audit_hook, ) ] self.assertEqual(events[-1].event, "done") self.assertEqual( "".join(str(event.data.get("text", "")) for event in events if event.event == "token"), "Masked stream reply.", ) self.assertIsNotNone(engine.request) self.assertIsNotNone(engine.payload) _assert_no_raw_pii(self, engine.request.messages) _assert_no_raw_pii(self, engine.payload) _assert_masked_pii_present(self, engine.request.messages) _assert_masked_pii_present(self, engine.payload) self.assertEqual(ctx.learner_text_raw, RAW_TEXT) self.assertEqual(len(audit_payloads), 1) self.assertEqual(audit_payloads[0]["tokens_in"], 5) self.assertEqual(audit_payloads[0]["tokens_out"], 6) _assert_no_raw_pii(self, audit_payloads) for key in ("messages", "prompt", "text"): self.assertNotIn(key, audit_payloads[0]) async def test_run_turn_generate_sends_only_masked_korean_pii(self) -> None: ctx = orchestrator.prepare_turn( session_id="masking-session", case_id="masking-case", card=persona.P1, state=_initial_state(), learner_text=RAW_KO_TEXT, ) engine = CaptureGenerateEngine() audit_payloads: list[dict[str, Any]] = [] async def audit_hook(payload: dict[str, Any]) -> None: audit_payloads.append(payload) await orchestrator.run_turn_generate( ctx, engine, # type: ignore[arg-type] audit_hook=audit_hook, ) self.assertIsNotNone(engine.request) self.assertIsNotNone(engine.payload) _assert_no_raw_ko_pii(self, engine.request.messages) _assert_no_raw_ko_pii(self, engine.payload) _assert_masked_ko_pii_present(self, engine.request.messages) _assert_masked_ko_pii_present(self, engine.payload) _assert_no_raw_ko_pii(self, audit_payloads) async def test_run_turn_generate_humanizes_masked_placeholder_reply(self) -> None: ctx = orchestrator.prepare_turn( session_id="masking-session", case_id="masking-case", card=persona.P1, state=_initial_state(), learner_text="진로 이야기를 이어가고 싶어요.", ) engine = CaptureGenerateEngine( text="아직 정해진 건 하나도 없어요. [NAME]는 그대로고, 엄마한테 뭐라고 말할지도 모르겠고요." ) result = await orchestrator.run_turn_generate( ctx, engine, # type: ignore[arg-type] ) self.assertNotIn("[NAME]", result.client_reply or "") self.assertIn("그 이름은 그대로고", result.client_reply or "") async def test_run_turn_generate_retries_once_after_role_meta_reply(self) -> None: class SequenceEngine(CaptureGenerateEngine): def __init__(self) -> None: super().__init__("") self.responses = [ "내담자 역할로 응답하겠습니다. 엄마가 가보라고 해서요.", "엄마가 그냥 가보라고 해서 왔어요.", ] async def generate(self, req): self.request = req self.requests.append(req) self.payload = self._payload_builder._payload(req) text = self.responses.pop(0) return GenerateResponse( text=text, model="fake-model", provider="fake-provider", tokens_in=3, tokens_out=4, cost_usd=0.0, ) ctx = orchestrator.prepare_turn( session_id="quality-session", case_id="quality-case", card=persona.P1, state=_initial_state(), learner_text="어머니가 오라고 하셨군요. 지금은 어떤 마음인가요?", ) engine = SequenceEngine() result = await orchestrator.run_turn_generate(ctx, engine) # type: ignore[arg-type] self.assertEqual(len(engine.requests), 2) self.assertEqual(result.client_reply, "엄마가 그냥 가보라고 해서 왔어요.") self.assertFalse(result.safety_flagged) async def test_run_turn_generate_returns_retryable_error_without_saving_bad_fallback(self) -> None: class AlwaysBadEngine(CaptureGenerateEngine): def __init__(self) -> None: super().__init__("") async def generate(self, req): self.request = req self.requests.append(req) self.payload = self._payload_builder._payload(req) return GenerateResponse( text="AI로서 내담자 역할로 응답하겠습니다.", model="fake-model", provider="fake-provider", tokens_in=3, tokens_out=4, cost_usd=0.0, ) ctx = orchestrator.prepare_turn( session_id="quality-session", case_id="quality-case", card=persona.P1, state=_initial_state(), learner_text="지금 이 자리에서 가장 말하기 어려운 게 뭔가요?", ) engine = AlwaysBadEngine() result = await orchestrator.run_turn_generate(ctx, engine) # type: ignore[arg-type] self.assertEqual(len(engine.requests), 2) self.assertIsNone(result.client_reply) self.assertTrue(result.safety_flagged) self.assertEqual(getattr(result, "output_error", None), "client_reply_quality_retryable") async def test_run_turn_stream_humanizes_split_masked_placeholder_reply(self) -> None: ctx = orchestrator.prepare_turn( session_id="masking-session", case_id="masking-case", card=persona.P1, state=_initial_state(), learner_text="진로 이야기를 이어가고 싶어요.", ) engine = CaptureStreamEngine( chunks=["아직 정해진 건 하나도 없어요. ", "[NA", "ME]는", " 그대로고요."] ) events = [ event async for event in orchestrator.run_turn_stream( ctx, engine, # type: ignore[arg-type] ) ] streamed = "".join( str(event.data.get("text", "")) for event in events if event.event == "token" ) self.assertNotIn("[NAME]", streamed) self.assertIn("그 이름은 그대로고요.", streamed) self.assertEqual(events[-1].event, "done") async def test_run_turn_stream_buffers_role_meta_reply_without_token_leak(self) -> None: ctx = orchestrator.prepare_turn( session_id="quality-stream-session", case_id="quality-stream-case", card=persona.P1, state=_initial_state(), learner_text="어머니가 오라고 하셨군요. 지금은 어떤 마음인가요?", ) engine = CaptureStreamEngine( chunks=[ "내담자 ", "역할로 응답하겠습니다. 엄마가 가보라고 해서요.", ] ) events = [ event async for event in orchestrator.run_turn_stream( ctx, engine, # type: ignore[arg-type] ) ] self.assertEqual([event.event for event in events], ["safety", "done"]) self.assertEqual( "".join(str(event.data.get("text", "")) for event in events if event.event == "token"), "", ) self.assertEqual(events[0].data["reason"], "client_reply_quality_retryable") self.assertTrue(events[-1].data["safety_flagged"]) self.assertEqual(events[-1].data["output_error"], "client_reply_quality_retryable") if __name__ == "__main__": unittest.main()