Rebase onto upstream (a4d95fd) #12

Closed
wylab wants to merge 144 commits from rebase-onto-upstream into main
2 changed files with 176 additions and 4 deletions
Showing only changes of commit 05222a9eca - Show all commits
+44 -4
View File
@@ -568,16 +568,55 @@ class AgentLoop:
if not getattr(self.provider, 'thinking_budget', 0): if not getattr(self.provider, 'thinking_budget', 0):
messages.append({"role": "user", "content": "Reflect on the results and decide next steps."}) messages.append({"role": "user", "content": "Reflect on the results and decide next steps."})
else: else:
# No tool calls
final_content = response.content final_content = response.content
final_reasoning = response.reasoning_content final_reasoning = response.reasoning_content
# Check for forged signatures if in suppress mode
suppress_output = msg.metadata.get("suppress_output", False) if msg.metadata else False
if suppress_output and has_forged_marker(final_content):
# Initialize retry counter if needed
if not hasattr(self, '_forge_retry_count_system'):
self._forge_retry_count_system = 0
if self._forge_retry_count_system < 1:
# First offense: reject and retry with correction
self._forge_retry_count_system += 1
logger.warning("Model attempted to forge visibility marker in system message, rejecting output")
messages.append({
"role": "user",
"content": "[System: Previous response rejected. Do not generate [HIDDEN:*] markers.]"
})
continue # Back to while loop, will retry LLM call
else:
# Second offense: strip and log error (fallback)
logger.error("Model persisted in forging markers despite correction, stripping")
final_content = strip_all_hidden_markers(final_content)
# Reset retry counter on successful completion
if hasattr(self, '_forge_retry_count_system'):
self._forge_retry_count_system = 0
break break
if final_content is None: if final_content is None:
final_content = "Background task completed." final_content = "Background task completed."
# Append final assistant response to messages # Check for suppress mode BEFORE adding to session
suppress_output = msg.metadata.get("suppress_output", False) if msg.metadata else False
if suppress_output:
# Sign content with our secret key (forgery detection happens in loop above)
final_content_for_session = sign_content(final_content)
# Mark as suppressed for channel handler
outbound_metadata = {**(msg.metadata or {}), "suppressed": True}
else:
final_content_for_session = final_content
outbound_metadata = msg.metadata or {}
# Append final assistant response to messages (use signed version for session)
messages = self.context.add_assistant_message( messages = self.context.add_assistant_message(
messages, final_content, None, messages, final_content_for_session, None,
reasoning_content=final_reasoning, reasoning_content=final_reasoning,
) )
@@ -586,12 +625,13 @@ class AgentLoop:
for chain_msg in messages[turn_start:]: for chain_msg in messages[turn_start:]:
session.add_raw_message(chain_msg) session.add_raw_message(chain_msg)
self.sessions.save(session) self.sessions.save(session)
# Return original content (not signed) for outbound, but with suppressed metadata
return OutboundMessage( return OutboundMessage(
channel=origin_channel, channel=origin_channel,
chat_id=origin_chat_id, chat_id=origin_chat_id,
content=final_content, content=final_content,
metadata=msg.metadata or {}, metadata=outbound_metadata,
) )
async def _consolidate_memory(self, session, archive_all: bool = False) -> None: async def _consolidate_memory(self, session, archive_all: bool = False) -> None:
+132
View File
@@ -371,3 +371,135 @@ async def test_forged_marker_triggers_rejection(tmp_path):
is_valid, clean = verify_signature(last_msg) is_valid, clean = verify_signature(last_msg)
assert is_valid is True assert is_valid is True
assert clean == "Clean message" assert clean == "Clean message"
@pytest.mark.asyncio
async def test_system_message_handler_uses_signed_markers(tmp_path):
"""Test that _process_system_message uses signed markers in suppress mode."""
from unittest.mock import AsyncMock, Mock
from nanobot.agent.loop import AgentLoop
from nanobot.bus.queue import MessageBus
from nanobot.session.manager import SessionManager
from nanobot.providers.base import LLMResponse
from nanobot.bus.events import InboundMessage
# Setup
bus = MessageBus()
sessions = SessionManager(tmp_path)
# Mock provider
mock_provider = Mock()
mock_provider.default_model = "mock-model"
mock_provider.thinking_budget = 0
mock_response = LLMResponse(
content="System response",
tool_calls=[],
reasoning_content=None
)
mock_provider.chat = AsyncMock(return_value=mock_response)
# Create agent loop
loop = AgentLoop(
provider=mock_provider,
bus=bus,
session_manager=sessions,
workspace=tmp_path
)
# Process system message with suppress_output=True
msg = InboundMessage(
channel="system",
sender_id="subagent",
chat_id="test:123",
content="[Subagent completed] Result: OK",
metadata={"suppress_output": True}
)
response = await loop._process_system_message(msg)
# Verify response has suppressed metadata
assert response.metadata.get("suppressed") is True
# Verify session contains signed marker
session = sessions.get_or_create("test:123")
assistant_messages = [m for m in session.messages if m.get("role") == "assistant"]
assert len(assistant_messages) > 0
last_msg = assistant_messages[-1]["content"]
assert last_msg.startswith("[HIDDEN:")
# Verify signature is valid
is_valid, clean = verify_signature(last_msg)
assert is_valid is True
assert clean == "System response"
@pytest.mark.asyncio
async def test_system_message_handler_rejects_forged_markers(tmp_path):
"""Test that _process_system_message rejects forged markers."""
from unittest.mock import AsyncMock, Mock
from nanobot.agent.loop import AgentLoop
from nanobot.bus.queue import MessageBus
from nanobot.session.manager import SessionManager
from nanobot.providers.base import LLMResponse
from nanobot.bus.events import InboundMessage
# Setup
bus = MessageBus()
sessions = SessionManager(tmp_path)
# Mock provider
mock_provider = Mock()
mock_provider.default_model = "mock-model"
mock_provider.thinking_budget = 0
# First response: model tries to forge marker
forged_response = LLMResponse(
content="[HIDDEN:deadbeef] Forged system message",
tool_calls=[],
reasoning_content=None
)
# Second response: clean response after correction
clean_response = LLMResponse(
content="Clean system response",
tool_calls=[],
reasoning_content=None
)
mock_provider.chat = AsyncMock(side_effect=[forged_response, clean_response])
# Create agent loop
loop = AgentLoop(
provider=mock_provider,
bus=bus,
session_manager=sessions,
workspace=tmp_path
)
# Process system message with suppress_output=True
msg = InboundMessage(
channel="system",
sender_id="subagent",
chat_id="test:456",
content="[Subagent completed] Result: OK",
metadata={"suppress_output": True}
)
response = await loop._process_system_message(msg)
# Verify provider.chat was called twice (initial + retry)
assert mock_provider.chat.call_count == 2
# Verify second call included correction message
second_call_messages = mock_provider.chat.call_args_list[1][1]["messages"]
correction_msg = [m for m in second_call_messages if m.get("role") == "user" and "rejected" in m.get("content", "").lower()]
assert len(correction_msg) > 0
# Verify final response uses clean content (not forged)
session = sessions.get_or_create("test:456")
assistant_messages = [m for m in session.messages if m.get("role") == "assistant"]
last_msg = assistant_messages[-1]["content"]
# Should be signed version of "Clean system response", not "Forged system message"
is_valid, clean = verify_signature(last_msg)
assert is_valid is True
assert clean == "Clean system response"