alishahryar1--free-claude-code
f0b31065ee
## Problem FCC hoisted inline Anthropic `system` messages into the top-level system prompt during request validation. Mid-conversation system messages are position-sensitive, so this applied later instructions retroactively, changed the existing prompt/cache prefix, and prevented provider conversion from seeing the original transcript. ## Changes - Preserve inline `system` messages, content, metadata, and ordering in Messages and token-count requests while keeping the top-level system prompt distinct. - Convert text-only inline system messages to OpenAI Chat `system` messages at the same transcript position; reject unrepresentable inline blocks before streaming instead of silently dropping them. - Remove the lossy normalization path and its unused role enum, and document protocol-model versus target-conversion ownership in `ARCHITECTURE.md`. - Cover API routing, model serialization, cache-prefix stability, text blocks, tool-result ordering, invalid content, and token counting; bump the package to `4.6.2`. - Verify all five local CI checks (2,287 tests) and the ordered transcript against NVIDIA NIM, OpenRouter, Gemini, DeepSeek, Mistral, and Hugging Face. <!-- greptile_comment --> <details open><summary><h3>Greptile Summary</h3></summary> This PR preserves inline Anthropic system messages through provider conversion. The main changes are: - Keeps top-level and inline system content separate and ordered. - Converts text-only inline system messages without moving them. - Rejects system blocks that OpenAI Chat cannot represent safely. - Updates request detection to ignore system context when counting user turns. - Adds serialization, routing, token-counting, and conversion coverage. - Updates the package version and architecture documentation. </details> <h3>Confidence Score: 5/5</h3> This looks safe to merge. The leading-system detection path ignores system entries when counting user turns. Inline system content remains ordered for provider conversion. Unsupported system blocks fail explicitly instead of being dropped. No blocking issues were found in the changed code. No files require attention. <details><summary><h3><a href="https://www.greptile.com/trex"><img alt="T-Rex" src="https://greptile-static-assets.s3.amazonaws.com/trex/trex_green.svg" height="20" align="absmiddle"></a> T-Rex Logs</h3></summary> **What T-Rex did** - Validated that the transcript roles now follow the order user, assistant, system, user and that the top-level prompt remains separate. - Verified that inline system content is no longer counted in message tokens and that cache\_control metadata survives parsing. - Confirmed that the converted OpenAI transcript preserves position and cache prefix. - Observed that a system message following a tool result is converted as assistant, tool, system. - Ran the focused pytest and confirmed 209 passed in 2.82s with exit code 0. <a href="https://app.greptile.com/trex/runs/14526513/artifacts"><picture><source media="(prefers-color-scheme: dark)" srcset="https://greptile-static-assets.s3.amazonaws.com/badges/ViewAllArtifactsDark.svg?v=4"><source media="(prefers-color-scheme: light)" srcset="https://greptile-static-assets.s3.amazonaws.com/badges/ViewAllArtifacts.svg?v=4"><img alt="View all artifacts" src="https://greptile-static-assets.s3.amazonaws.com/badges/ViewAllArtifacts.svg?v=4"></picture></a> <sub><a href="https://www.greptile.com/trex"><img alt="T-Rex" src="https://greptile-static-assets.s3.amazonaws.com/trex/trex_green.svg" height="14" align="absmiddle"></a> Ran code and verified through T-Rex</sub> </details> <details open><summary><h3>Important Files Changed</h3></summary> | Filename | Overview | |----------|----------| | src/free_claude_code/core/anthropic/models.py | Preserves system-role messages in the original transcript instead of hoisting them into the top-level prompt. | | src/free_claude_code/core/anthropic/conversion.py | Converts ordered text-only system messages and rejects unsupported system content before streaming. | | src/free_claude_code/api/detection.py | Builds a read-only semantic view of system context and conversational user turns for local request detection. | </details> <sub>Reviews (2): Last reviewed commit: ["Restore optimizations with inline system..."](https://github.com/alishahryar1/free-claude-code/commit/6605ede7f604053381552106489dbd16bcd37987) | [Re-trigger Greptile](https://app.greptile.com/api/retrigger?id=44423885)</sub> <!-- /greptile_comment -->
657 行
22 KiB
Python
657 行
22 KiB
Python
"""Tests for API request detection and token counting helpers."""
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from free_claude_code.api.command_utils import extract_command_prefix
|
|
from free_claude_code.api.detection import (
|
|
is_prefix_detection_request,
|
|
is_quota_check_request,
|
|
is_title_generation_request,
|
|
)
|
|
from free_claude_code.core.anthropic import get_token_count
|
|
from free_claude_code.core.anthropic.models import (
|
|
Message,
|
|
MessagesRequest,
|
|
SystemContent,
|
|
)
|
|
|
|
|
|
class TestQuotaCheckRequest:
|
|
"""Tests for is_quota_check_request function."""
|
|
|
|
def test_quota_check_simple_string(self):
|
|
"""Test quota check with simple string content."""
|
|
msg = MagicMock(spec=Message)
|
|
msg.role = "user"
|
|
msg.content = "Check my quota"
|
|
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.max_tokens = 1
|
|
req.messages = [msg]
|
|
|
|
assert is_quota_check_request(req) is True
|
|
|
|
def test_quota_check_case_insensitive(self):
|
|
"""Test quota check is case insensitive."""
|
|
msg = MagicMock(spec=Message)
|
|
msg.role = "user"
|
|
msg.content = "Check my QUOTA"
|
|
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.max_tokens = 1
|
|
req.messages = [msg]
|
|
|
|
assert is_quota_check_request(req) is True
|
|
|
|
def test_quota_check_list_content(self):
|
|
"""Test quota check with list content blocks."""
|
|
block = MagicMock()
|
|
block.text = "Check my quota"
|
|
|
|
msg = MagicMock(spec=Message)
|
|
msg.role = "user"
|
|
msg.content = [block]
|
|
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.max_tokens = 1
|
|
req.messages = [msg]
|
|
|
|
assert is_quota_check_request(req) is True
|
|
|
|
def test_not_quota_check_wrong_max_tokens(self):
|
|
"""Test not quota check when max_tokens != 1."""
|
|
msg = MagicMock(spec=Message)
|
|
msg.role = "user"
|
|
msg.content = "Check my quota"
|
|
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.max_tokens = 100
|
|
req.messages = [msg]
|
|
|
|
assert is_quota_check_request(req) is False
|
|
|
|
def test_not_quota_check_multiple_messages(self):
|
|
"""Test not quota check when multiple messages."""
|
|
msg1 = MagicMock(spec=Message)
|
|
msg1.role = "user"
|
|
msg1.content = "Check my quota"
|
|
|
|
msg2 = MagicMock(spec=Message)
|
|
msg2.role = "assistant"
|
|
msg2.content = "Hello"
|
|
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.max_tokens = 1
|
|
req.messages = [msg1, msg2]
|
|
|
|
assert is_quota_check_request(req) is False
|
|
|
|
def test_not_quota_check_wrong_role(self):
|
|
"""Test not quota check when role is not user."""
|
|
msg = MagicMock(spec=Message)
|
|
msg.role = "assistant"
|
|
msg.content = "Check my quota"
|
|
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.max_tokens = 1
|
|
req.messages = [msg]
|
|
|
|
assert is_quota_check_request(req) is False
|
|
|
|
def test_not_quota_check_no_quota_keyword(self):
|
|
"""Test not quota check when content doesn't contain quota."""
|
|
msg = MagicMock(spec=Message)
|
|
msg.role = "user"
|
|
msg.content = "Hello world"
|
|
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.max_tokens = 1
|
|
req.messages = [msg]
|
|
|
|
assert is_quota_check_request(req) is False
|
|
|
|
|
|
class TestTitleGenerationRequest:
|
|
"""Tests for is_title_generation_request function."""
|
|
|
|
def _title_gen_system(self) -> list[MagicMock]:
|
|
block = MagicMock()
|
|
block.text = (
|
|
"Generate a concise, sentence-case title (3-7 words) that captures the "
|
|
"main topic or goal of this coding session. Return JSON with a single "
|
|
'"title" field.'
|
|
)
|
|
return [block]
|
|
|
|
def test_title_generation_detected_via_system(self):
|
|
"""Title gen detected by session title system prompt (sentence-case / JSON)."""
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.system = self._title_gen_system()
|
|
req.tools = None
|
|
req.messages = []
|
|
|
|
assert is_title_generation_request(req) is True
|
|
|
|
def test_title_generation_not_detected_with_tools(self):
|
|
"""Not detected when tools are present (main conversation, not title gen)."""
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.system = self._title_gen_system()
|
|
req.tools = [MagicMock()]
|
|
req.messages = []
|
|
|
|
assert is_title_generation_request(req) is False
|
|
|
|
def test_title_generation_not_detected_no_system(self):
|
|
"""Not detected when system is absent."""
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.system = None
|
|
req.tools = None
|
|
req.messages = []
|
|
|
|
assert is_title_generation_request(req) is False
|
|
|
|
def test_title_generation_not_detected_unrelated_system(self):
|
|
"""Not detected when system prompt has no topic/title keywords."""
|
|
block = MagicMock()
|
|
block.text = "You are a helpful assistant."
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.system = [block]
|
|
req.tools = None
|
|
req.messages = []
|
|
|
|
assert is_title_generation_request(req) is False
|
|
|
|
def test_title_generation_return_json_coding_session_branch(self):
|
|
"""JSON title field + session wording matches without sentence-case phrase."""
|
|
block = MagicMock()
|
|
block.text = 'Return JSON with a single "title" field for this coding session.'
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.system = [block]
|
|
req.tools = None
|
|
req.messages = []
|
|
|
|
assert is_title_generation_request(req) is True
|
|
|
|
|
|
class TestExtractCommandPrefix:
|
|
"""Tests for extract_command_prefix function."""
|
|
|
|
def test_simple_command(self):
|
|
"""Test extraction of simple command."""
|
|
assert extract_command_prefix("git status") == "git status"
|
|
assert extract_command_prefix("ls -la") == "ls"
|
|
|
|
def test_two_word_commands(self):
|
|
"""Test extraction of two-word commands."""
|
|
assert extract_command_prefix("git commit -m 'message'") == "git commit"
|
|
assert extract_command_prefix("npm install package") == "npm install"
|
|
assert extract_command_prefix("docker run image") == "docker run"
|
|
assert extract_command_prefix("kubectl get pods") == "kubectl get"
|
|
|
|
def test_two_word_command_with_options(self):
|
|
"""Test two-word command with options only returns first word."""
|
|
assert extract_command_prefix("git -v") == "git"
|
|
assert extract_command_prefix("npm --version") == "npm"
|
|
|
|
def test_with_env_vars(self):
|
|
"""Test command with environment variables."""
|
|
assert extract_command_prefix("DEBUG=1 python script.py") == "DEBUG=1 python"
|
|
assert (
|
|
extract_command_prefix("API_KEY=secret node app.js")
|
|
== "API_KEY=secret node"
|
|
)
|
|
|
|
def test_single_word_commands(self):
|
|
"""Test single word commands."""
|
|
assert extract_command_prefix("ls") == "ls"
|
|
assert extract_command_prefix("python") == "python"
|
|
assert extract_command_prefix("make") == "make"
|
|
|
|
def test_command_injection_detected(self):
|
|
"""Test detection of command injection attempts."""
|
|
assert extract_command_prefix("`whoami`") == "command_injection_detected"
|
|
assert extract_command_prefix("$(whoami)") == "command_injection_detected"
|
|
assert (
|
|
extract_command_prefix("echo $(cat /etc/passwd)")
|
|
== "command_injection_detected"
|
|
)
|
|
|
|
def test_empty_command(self):
|
|
"""Test handling of empty commands."""
|
|
assert extract_command_prefix("") == "none"
|
|
assert extract_command_prefix(" ") == "none"
|
|
|
|
def test_complex_git_command(self):
|
|
"""Test complex git command extraction."""
|
|
assert extract_command_prefix("git log --oneline --graph") == "git log"
|
|
assert (
|
|
extract_command_prefix("git checkout -b feature-branch") == "git checkout"
|
|
)
|
|
|
|
def test_cargo_command(self):
|
|
"""Test cargo command extraction."""
|
|
assert extract_command_prefix("cargo build") == "cargo build"
|
|
assert extract_command_prefix("cargo test") == "cargo test"
|
|
assert extract_command_prefix("cargo --version") == "cargo"
|
|
|
|
|
|
class TestPrefixDetectionRequest:
|
|
"""Tests for is_prefix_detection_request function."""
|
|
|
|
def test_prefix_detection_with_policy_spec(self):
|
|
"""Test prefix detection with policy spec and command."""
|
|
msg = MagicMock(spec=Message)
|
|
msg.role = "user"
|
|
msg.content = "<policy_spec>policy</policy_spec> Command: git status"
|
|
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.messages = [msg]
|
|
|
|
is_prefix, command = is_prefix_detection_request(req)
|
|
assert is_prefix is True
|
|
assert command == "git status"
|
|
|
|
def test_prefix_detection_case_sensitive(self):
|
|
"""Test prefix detection is case sensitive for Command:."""
|
|
msg = MagicMock(spec=Message)
|
|
msg.role = "user"
|
|
msg.content = "<policy_spec>policy</policy_spec> command: git status"
|
|
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.messages = [msg]
|
|
|
|
is_prefix, command = is_prefix_detection_request(req)
|
|
assert is_prefix is False
|
|
assert command == ""
|
|
|
|
def test_not_prefix_detection_no_policy_spec(self):
|
|
"""Test not prefix detection without policy_spec."""
|
|
msg = MagicMock(spec=Message)
|
|
msg.role = "user"
|
|
msg.content = "Command: git status"
|
|
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.messages = [msg]
|
|
|
|
is_prefix, command = is_prefix_detection_request(req)
|
|
assert is_prefix is False
|
|
assert command == ""
|
|
|
|
def test_not_prefix_detection_multiple_messages(self):
|
|
"""Test not prefix detection with multiple messages."""
|
|
msg1 = MagicMock(spec=Message)
|
|
msg1.role = "user"
|
|
msg1.content = "<policy_spec>policy</policy_spec> Command: git status"
|
|
|
|
msg2 = MagicMock(spec=Message)
|
|
msg2.role = "assistant"
|
|
msg2.content = "OK"
|
|
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.messages = [msg1, msg2]
|
|
|
|
is_prefix, command = is_prefix_detection_request(req)
|
|
assert is_prefix is False
|
|
assert command == ""
|
|
|
|
def test_not_prefix_detection_wrong_role(self):
|
|
"""Test not prefix detection when message is not from user."""
|
|
msg = MagicMock(spec=Message)
|
|
msg.role = "assistant"
|
|
msg.content = "<policy_spec>policy</policy_spec> Command: git status"
|
|
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.messages = [msg]
|
|
|
|
is_prefix, command = is_prefix_detection_request(req)
|
|
assert is_prefix is False
|
|
assert command == ""
|
|
|
|
def test_prefix_detection_list_content(self):
|
|
"""Test prefix detection with list content blocks."""
|
|
block = MagicMock()
|
|
block.text = "<policy_spec>policy</policy_spec> Command: ls -la"
|
|
|
|
msg = MagicMock(spec=Message)
|
|
msg.role = "user"
|
|
msg.content = [block]
|
|
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.messages = [msg]
|
|
|
|
is_prefix, command = is_prefix_detection_request(req)
|
|
assert is_prefix is True
|
|
assert command == "ls -la"
|
|
|
|
|
|
class TestGetTokenCount:
|
|
"""Tests for get_token_count function."""
|
|
|
|
def test_empty_messages(self):
|
|
"""Test token count with empty messages."""
|
|
count = get_token_count([])
|
|
assert count >= 1 # Returns max(1, tokens)
|
|
|
|
def test_simple_message(self):
|
|
"""Test token count with simple text message."""
|
|
msg = MagicMock()
|
|
msg.content = "Hello world"
|
|
|
|
count = get_token_count([msg])
|
|
assert count > 0
|
|
# "Hello world" is ~2-3 tokens plus overhead
|
|
assert count >= 3
|
|
|
|
def test_special_token_text_is_counted_as_plain_text(self):
|
|
"""Tiktoken special-token strings should not break token estimates."""
|
|
msg = MagicMock()
|
|
msg.content = "<|endoftext|>"
|
|
|
|
count = get_token_count([msg], system="<|endoftext|>")
|
|
assert count > 0
|
|
|
|
def test_message_with_system_prompt(self):
|
|
"""Test token count includes system prompt."""
|
|
msg = MagicMock()
|
|
msg.content = "Hello"
|
|
|
|
count = get_token_count([msg], system="You are a helpful assistant")
|
|
assert count > 0
|
|
|
|
def test_message_with_list_content(self):
|
|
"""Test token count with list content blocks."""
|
|
text_block = MagicMock()
|
|
text_block.type = "text"
|
|
text_block.text = "Hello world"
|
|
|
|
msg = MagicMock()
|
|
msg.content = [text_block]
|
|
|
|
count = get_token_count([msg])
|
|
assert count > 0
|
|
|
|
def test_message_with_thinking_block(self):
|
|
"""Test token count includes thinking blocks."""
|
|
thinking_block = MagicMock()
|
|
thinking_block.type = "thinking"
|
|
thinking_block.thinking = "Let me think about this..."
|
|
|
|
msg = MagicMock()
|
|
msg.content = [thinking_block]
|
|
|
|
count = get_token_count([msg])
|
|
assert count > 0
|
|
|
|
def test_message_with_tool_use(self):
|
|
"""Test token count includes tool use blocks."""
|
|
tool_block = MagicMock()
|
|
tool_block.type = "tool_use"
|
|
tool_block.name = "search"
|
|
tool_block.input = {"query": "test"}
|
|
|
|
msg = MagicMock()
|
|
msg.content = [tool_block]
|
|
|
|
count = get_token_count([msg])
|
|
assert count > 0
|
|
|
|
def test_message_with_tool_result(self):
|
|
"""Test token count includes tool result blocks."""
|
|
result_block = MagicMock()
|
|
result_block.type = "tool_result"
|
|
result_block.content = "Search results here"
|
|
|
|
msg = MagicMock()
|
|
msg.content = [result_block]
|
|
|
|
count = get_token_count([msg])
|
|
assert count > 0
|
|
|
|
def test_message_with_tools(self):
|
|
"""Test token count includes tool definitions."""
|
|
msg = MagicMock()
|
|
msg.content = "Use the search tool"
|
|
|
|
tool = MagicMock()
|
|
tool.name = "search"
|
|
tool.description = "Search for information"
|
|
tool.input_schema = {"type": "object", "properties": {}}
|
|
|
|
count = get_token_count([msg], tools=[tool])
|
|
assert count > 0
|
|
|
|
def test_system_as_list(self):
|
|
"""Test token count with system as list of blocks."""
|
|
msg = MagicMock()
|
|
msg.content = "Hello"
|
|
|
|
block = MagicMock()
|
|
block.text = "System prompt"
|
|
|
|
count = get_token_count([msg], system=[block])
|
|
assert count > 0
|
|
|
|
def test_tool_result_with_dict_content(self):
|
|
"""Test token count with tool result containing dict content."""
|
|
result_block = MagicMock()
|
|
result_block.type = "tool_result"
|
|
result_block.content = {"result": "data"}
|
|
|
|
msg = MagicMock()
|
|
msg.content = [result_block]
|
|
|
|
count = get_token_count([msg])
|
|
assert count > 0
|
|
|
|
def test_multiple_messages_overhead(self):
|
|
"""Test that multiple messages include overhead."""
|
|
msg1 = MagicMock()
|
|
msg1.content = "Hi"
|
|
|
|
msg2 = MagicMock()
|
|
msg2.content = "Hello"
|
|
|
|
count_single = get_token_count([msg1])
|
|
count_double = get_token_count([msg1, msg2])
|
|
|
|
# Double message should have more tokens (including overhead)
|
|
assert count_double > count_single
|
|
|
|
def test_inline_system_message_contributes_to_token_count(self):
|
|
"""Mid-conversation system content remains part of the counted transcript."""
|
|
user_message = Message(role="user", content="Hello")
|
|
inline_system = Message(role="system", content="New instructions")
|
|
|
|
without_inline_system = get_token_count([user_message])
|
|
with_inline_system = get_token_count([user_message, inline_system])
|
|
|
|
assert with_inline_system > without_inline_system
|
|
|
|
def test_per_message_overhead_four_tokens(self):
|
|
"""Per-message overhead is 4 tokens (was 3)."""
|
|
msg = MagicMock()
|
|
msg.content = "x" # Minimal content
|
|
count = get_token_count([msg])
|
|
# 1 msg * 4 overhead + content tokens
|
|
assert count >= 5
|
|
|
|
def test_system_overhead_added(self):
|
|
"""System prompt adds ~4 tokens overhead."""
|
|
msg = MagicMock()
|
|
msg.content = "Hi"
|
|
count_no_sys = get_token_count([msg])
|
|
count_with_sys = get_token_count([msg], system="You are helpful")
|
|
assert count_with_sys >= count_no_sys + 4
|
|
|
|
def test_system_as_typed_content_blocks(self):
|
|
"""Typed system content blocks are counted."""
|
|
msg = MagicMock()
|
|
msg.content = "Hi"
|
|
count_no_sys = get_token_count([msg])
|
|
system_blocks = [
|
|
SystemContent(type="text", text="System prompt from typed block")
|
|
]
|
|
count_with_system = get_token_count([msg], system=system_blocks)
|
|
assert count_with_system > count_no_sys
|
|
|
|
def test_tool_use_includes_id(self):
|
|
"""Tool use blocks count id field."""
|
|
tool_block = MagicMock()
|
|
tool_block.type = "tool_use"
|
|
tool_block.name = "search"
|
|
tool_block.input = {"q": "test"}
|
|
tool_block.id = "call_abc123"
|
|
msg = MagicMock()
|
|
msg.content = [tool_block]
|
|
count = get_token_count([msg])
|
|
assert count > 0
|
|
|
|
def test_tool_result_includes_tool_use_id(self):
|
|
"""Tool result blocks count tool_use_id field."""
|
|
result_block = MagicMock()
|
|
result_block.type = "tool_result"
|
|
result_block.content = "ok"
|
|
result_block.tool_use_id = "call_xyz"
|
|
msg = MagicMock()
|
|
msg.content = [result_block]
|
|
count = get_token_count([msg])
|
|
assert count > 0
|
|
|
|
def test_unrecognized_block_type_fallback(self):
|
|
"""Unrecognized block types are tokenized via json.dumps fallback."""
|
|
unknown_block = {"type": "custom", "spec": "data"}
|
|
msg = MagicMock()
|
|
msg.content = [unknown_block]
|
|
count = get_token_count([msg])
|
|
assert count > 0
|
|
|
|
def test_message_with_image_block(self):
|
|
"""Test token count includes image blocks."""
|
|
image_block = MagicMock()
|
|
image_block.type = "image"
|
|
image_block.source = {
|
|
"type": "base64",
|
|
"media_type": "image/png",
|
|
"data": "x" * 3000,
|
|
}
|
|
msg = MagicMock()
|
|
msg.content = [image_block]
|
|
count = get_token_count([msg])
|
|
assert count >= 85
|
|
|
|
def test_image_block_with_dict_source(self):
|
|
"""Image block with dict-style source is counted."""
|
|
image_block = {"type": "image", "source": {"data": "a" * 10000}}
|
|
msg = MagicMock()
|
|
msg.content = [image_block]
|
|
count = get_token_count([msg])
|
|
assert count >= 85
|
|
|
|
def test_known_payload_estimate_range(self):
|
|
"""Known payload produces estimate within expected range (validation harness)."""
|
|
import tiktoken
|
|
|
|
enc = tiktoken.get_encoding("cl100k_base")
|
|
system_text = "You are a helpful assistant."
|
|
user_text = "Hello, how are you?"
|
|
sys_tokens = len(enc.encode(system_text))
|
|
user_tokens = len(enc.encode(user_text))
|
|
# Min: content tokens + system overhead (4) + per-msg overhead (4)
|
|
expected_min = sys_tokens + user_tokens + 4 + 4
|
|
msg = MagicMock()
|
|
msg.content = user_text
|
|
count = get_token_count([msg], system=system_text)
|
|
assert count >= expected_min, f"count={count} < expected_min={expected_min}"
|
|
|
|
|
|
# --- Parametrized Edge Case Tests ---
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"command,expected",
|
|
[
|
|
("git status", "git status"),
|
|
("ls -la", "ls"),
|
|
("git commit -m 'msg'", "git commit"),
|
|
("npm install pkg", "npm install"),
|
|
("ls", "ls"),
|
|
("python", "python"),
|
|
("", "none"),
|
|
(" ", "none"),
|
|
("`whoami`", "command_injection_detected"),
|
|
("$(whoami)", "command_injection_detected"),
|
|
("echo $(cat /etc/passwd)", "command_injection_detected"),
|
|
("git -v", "git"),
|
|
("DEBUG=1 python script.py", "DEBUG=1 python"),
|
|
("cargo build", "cargo build"),
|
|
("cargo --version", "cargo"),
|
|
],
|
|
ids=[
|
|
"git_status",
|
|
"ls_with_flag",
|
|
"git_commit",
|
|
"npm_install",
|
|
"bare_ls",
|
|
"bare_python",
|
|
"empty",
|
|
"whitespace",
|
|
"injection_backtick",
|
|
"injection_dollar",
|
|
"injection_echo",
|
|
"git_flag",
|
|
"env_var",
|
|
"cargo_build",
|
|
"cargo_flag",
|
|
],
|
|
)
|
|
def test_extract_command_prefix_parametrized(command, expected):
|
|
"""Parametrized command prefix extraction."""
|
|
assert extract_command_prefix(command) == expected
|
|
|
|
|
|
def test_extract_command_prefix_unterminated_quote():
|
|
"""Unterminated quote falls back to simple split (shlex.split ValueError)."""
|
|
result = extract_command_prefix("git commit -m 'unterminated")
|
|
# Should fall back to command.split()[0] = "git"
|
|
assert result == "git"
|
|
|
|
|
|
def test_extract_command_prefix_pipe():
|
|
"""Piped commands - shlex handles pipe character."""
|
|
result = extract_command_prefix("cat file.txt | grep pattern")
|
|
assert result in ("cat", "cat file.txt")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"content,max_tokens,role,expected",
|
|
[
|
|
("Check my quota", 1, "user", True),
|
|
("Check my QUOTA", 1, "user", True),
|
|
("Hello world", 1, "user", False),
|
|
("Check my quota", 100, "user", False),
|
|
("Check my quota", 1, "assistant", False),
|
|
],
|
|
ids=["basic", "case_insensitive", "no_keyword", "wrong_max_tokens", "wrong_role"],
|
|
)
|
|
def test_quota_check_parametrized(content, max_tokens, role, expected):
|
|
"""Parametrized quota check request detection."""
|
|
msg = MagicMock(spec=Message)
|
|
msg.role = role
|
|
msg.content = content
|
|
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.max_tokens = max_tokens
|
|
req.messages = [msg]
|
|
|
|
assert is_quota_check_request(req) is expected
|
|
|
|
|
|
def test_quota_check_empty_messages():
|
|
"""Quota check with empty message list should not crash."""
|
|
req = MagicMock(spec=MessagesRequest)
|
|
req.max_tokens = 1
|
|
req.messages = []
|
|
assert is_quota_check_request(req) is False
|