Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 23 additions & 14 deletions autogpt_platform/backend/backend/copilot/sdk/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -2330,13 +2330,19 @@ def _on_stderr(line: str) -> None:
# --session-id here. CLI >=2.1.97 rejects the combination of
# --session-id + --resume unless --fork-session is also given.
sdk_options_kwargs["resume"] = resume_file
elif not has_history:
# T1 only: write CLI native session to a predictable path so
# upload_cli_session() can find it after the turn completes.
# On T2+ without --resume the T1 session file already exists at
# that path; passing --session-id again would fail with
# "Session ID already in use". The upload guard also skips T2+
# no-resume turns, so --session-id provides no benefit there.
else:
# Set session_id whenever NOT resuming so the CLI writes the
# native session file to a predictable path for
# upload_cli_session() after the turn. This covers:
Comment thread
majdyz marked this conversation as resolved.
Comment thread
majdyz marked this conversation as resolved.
# • T1 fresh: no prior history, first SDK turn.
# • Mode-switch T1: has_history=True (prior baseline turns in
# DB) but no CLI session file was ever uploaded — the CLI has
# never been invoked with this session_id before.
# • T2+ without --resume (restore failed): no session file was
# restored to local storage (restore_cli_session returned
# False), so no conflict with an existing file.
# When --resume is active the session_id is already implied by
# the resume file; passing it again would be rejected by the CLI.
sdk_options_kwargs["session_id"] = session_id
# Optional explicit Claude Code CLI binary path (decouples the
# bundled SDK version from the CLI version we run — needed because
Expand Down Expand Up @@ -2516,16 +2522,19 @@ def _on_stderr(line: str) -> None:
if ctx.use_resume and ctx.resume_file:
sdk_options_kwargs_retry["resume"] = ctx.resume_file
sdk_options_kwargs_retry.pop("session_id", None)
elif not has_history:
# T1 retry: keep session_id so the CLI writes to the
# predictable path for upload_cli_session().
elif "session_id" in sdk_options_kwargs:
# Initial invocation used session_id (T1 or mode-switch
# T1): keep it so the CLI writes the session file to the
# predictable path for upload_cli_session(). Storage is
# ephemeral per invocation, so no "Session ID already in
# use" conflict occurs — no prior file was restored.
sdk_options_kwargs_retry.pop("resume", None)
sdk_options_kwargs_retry["session_id"] = session_id
else:
# T2+ retry without --resume: do not pass --session-id.
# The T1 session file already exists at that path; re-using
# the same ID would fail with "Session ID already in use".
# The upload guard skips T2+ no-resume turns anyway.
# T2+ retry without --resume: initial invocation used
# --resume, which restored the T1 session file to local
# storage. Re-using session_id without --resume would
# fail with "Session ID already in use".
sdk_options_kwargs_retry.pop("resume", None)
sdk_options_kwargs_retry.pop("session_id", None)
state.options = ClaudeAgentOptions(**sdk_options_kwargs_retry) # type: ignore[arg-type] # dynamic kwargs
Expand Down
256 changes: 256 additions & 0 deletions autogpt_platform/backend/backend/copilot/sdk/service_helpers_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,9 @@
_is_prompt_too_long,
_is_tool_only_message,
_iter_sdk_messages,
_normalize_model_name,
_reduce_context,
_TokenUsage,
)

# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -331,3 +333,257 @@ def test_single_tool_block_is_parallel(self):
msg = MagicMock(spec=AssistantMessage)
msg.content = [self._make_tool_block()]
assert _is_tool_only_message(msg) is True


# ---------------------------------------------------------------------------
# _normalize_model_name — used by per-request model override
# ---------------------------------------------------------------------------


class TestNormalizeModelName:
"""Unit tests for the model-name normalisation helper.

The per-request model toggle calls _normalize_model_name with either
``"anthropic/claude-opus-4-6"`` (for 'advanced') or ``config.model`` (for
'standard'). These tests verify the OpenRouter/provider-prefix stripping
that keeps the value compatible with the Claude CLI.
"""

def test_strips_anthropic_prefix(self):
assert _normalize_model_name("anthropic/claude-opus-4-6") == "claude-opus-4-6"

def test_strips_openai_prefix(self):
assert _normalize_model_name("openai/gpt-4o") == "gpt-4o"

def test_strips_google_prefix(self):
assert _normalize_model_name("google/gemini-2.5-flash") == "gemini-2.5-flash"

def test_already_normalized_unchanged(self):
assert (
_normalize_model_name("claude-sonnet-4-20250514")
== "claude-sonnet-4-20250514"
)

def test_empty_string_unchanged(self):
assert _normalize_model_name("") == ""

def test_opus_model_roundtrip(self):
"""The exact string used for the 'opus' toggle strips correctly."""
assert _normalize_model_name("anthropic/claude-opus-4-6") == "claude-opus-4-6"

def test_sonnet_openrouter_model(self):
"""Sonnet model as stored in config (OpenRouter-prefixed) strips cleanly."""
assert _normalize_model_name("anthropic/claude-sonnet-4") == "claude-sonnet-4"


# ---------------------------------------------------------------------------
# _TokenUsage — null-safe accumulation (OpenRouter initial-stream-event bug)
# ---------------------------------------------------------------------------


class TestTokenUsageNullSafety:
"""Verify that ResultMessage.usage dicts with null-valued cache fields
(as emitted by OpenRouter for the initial streaming event before real
token counts are available) do not crash the accumulator.

Before the fix, dict.get("cache_read_input_tokens", 0) returned None
when the key existed with a null value, causing 'int += None' TypeError.
"""

def _apply_usage(self, usage: dict, acc: _TokenUsage) -> None:
"""Mirror the production accumulation in sdk/service.py."""
acc.prompt_tokens += usage.get("input_tokens") or 0
acc.cache_read_tokens += usage.get("cache_read_input_tokens") or 0
acc.cache_creation_tokens += usage.get("cache_creation_input_tokens") or 0
acc.completion_tokens += usage.get("output_tokens") or 0
Comment thread
coderabbitai[bot] marked this conversation as resolved.

def test_null_cache_tokens_do_not_crash(self):
"""OpenRouter initial event: cache keys present with null value."""
usage = {
"input_tokens": 0,
"output_tokens": 0,
"cache_read_input_tokens": None,
"cache_creation_input_tokens": None,
}
acc = _TokenUsage()
self._apply_usage(usage, acc) # must not raise TypeError
assert acc.prompt_tokens == 0
assert acc.cache_read_tokens == 0
assert acc.cache_creation_tokens == 0
assert acc.completion_tokens == 0

def test_real_cache_tokens_are_accumulated(self):
"""OpenRouter final event: real cache token counts are captured."""
usage = {
"input_tokens": 10,
"output_tokens": 349,
"cache_read_input_tokens": 16600,
"cache_creation_input_tokens": 512,
}
acc = _TokenUsage()
self._apply_usage(usage, acc)
assert acc.prompt_tokens == 10
assert acc.cache_read_tokens == 16600
assert acc.cache_creation_tokens == 512
assert acc.completion_tokens == 349

def test_absent_cache_keys_default_to_zero(self):
"""Minimal usage dict without cache keys defaults correctly."""
usage = {"input_tokens": 5, "output_tokens": 20}
acc = _TokenUsage()
self._apply_usage(usage, acc)
assert acc.prompt_tokens == 5
assert acc.cache_read_tokens == 0
assert acc.cache_creation_tokens == 0
assert acc.completion_tokens == 20

def test_multi_turn_accumulation(self):
"""Null event followed by real event: only real tokens counted."""
null_event = {
"input_tokens": 0,
"output_tokens": 0,
"cache_read_input_tokens": None,
"cache_creation_input_tokens": None,
}
real_event = {
"input_tokens": 10,
"output_tokens": 349,
"cache_read_input_tokens": 16600,
"cache_creation_input_tokens": 512,
}
acc = _TokenUsage()
self._apply_usage(null_event, acc)
self._apply_usage(real_event, acc)
assert acc.prompt_tokens == 10
assert acc.cache_read_tokens == 16600
assert acc.cache_creation_tokens == 512
assert acc.completion_tokens == 349


# ---------------------------------------------------------------------------
# session_id / resume selection logic
# ---------------------------------------------------------------------------


def _build_sdk_options(
use_resume: bool,
resume_file: str | None,
session_id: str,
) -> dict:
"""Mirror the session_id/resume selection in stream_chat_completion_sdk.

This helper encodes the exact branching so the unit tests stay in sync
with the production code without needing to invoke the full generator.
"""
kwargs: dict = {}
if use_resume and resume_file:
kwargs["resume"] = resume_file
else:
kwargs["session_id"] = session_id
return kwargs


def _build_retry_sdk_options(
initial_kwargs: dict,
ctx_use_resume: bool,
ctx_resume_file: str | None,
session_id: str,
) -> dict:
"""Mirror the retry branch in stream_chat_completion_sdk."""
retry: dict = dict(initial_kwargs)
if ctx_use_resume and ctx_resume_file:
retry["resume"] = ctx_resume_file
retry.pop("session_id", None)
elif "session_id" in initial_kwargs:
retry.pop("resume", None)
retry["session_id"] = session_id
else:
retry.pop("resume", None)
retry.pop("session_id", None)
return retry


class TestSdkSessionIdSelection:
"""Verify that session_id is set for all non-resume turns.

Regression test for the mode-switch T1 bug: when a user switches from
baseline mode (fast) to SDK mode (extended_thinking) mid-session, the
first SDK turn has has_history=True but no CLI session file. The old
code gated session_id on ``not has_history``, so mode-switch T1 never
got a session_id — the CLI used a random ID that couldn't be found on
the next turn, causing --resume to fail for the whole session.
"""

SESSION_ID = "sess-abc123"

def test_t1_fresh_sets_session_id(self):
"""T1 of a fresh session always gets session_id."""
opts = _build_sdk_options(
use_resume=False,
resume_file=None,
session_id=self.SESSION_ID,
)
assert opts.get("session_id") == self.SESSION_ID
assert "resume" not in opts

def test_mode_switch_t1_sets_session_id(self):
"""Mode-switch T1 (has_history=True, no CLI session) gets session_id.

Before the fix, the ``elif not has_history`` guard prevented this
case from setting session_id, causing all subsequent turns to run
without --resume.
"""
# Mode-switch T1: use_resume=False (no prior CLI session) and
# has_history=True (prior baseline turns in DB). The old code
# (``elif not has_history``) silently skipped this case.
opts = _build_sdk_options(
use_resume=False,
resume_file=None,
session_id=self.SESSION_ID,
)
assert opts.get("session_id") == self.SESSION_ID
assert "resume" not in opts

def test_t2_with_resume_uses_resume(self):
"""T2+ with a restored CLI session uses --resume, not session_id."""
opts = _build_sdk_options(
use_resume=True,
resume_file=self.SESSION_ID,
session_id=self.SESSION_ID,
)
assert opts.get("resume") == self.SESSION_ID
assert "session_id" not in opts

def test_t2_without_resume_sets_session_id(self):
"""T2+ when restore failed still gets session_id (no prior file on disk)."""
opts = _build_sdk_options(
use_resume=False,
resume_file=None,
session_id=self.SESSION_ID,
)
assert opts.get("session_id") == self.SESSION_ID
assert "resume" not in opts

def test_retry_keeps_session_id_for_t1(self):
"""Retry for T1 (or mode-switch T1) preserves session_id."""
initial = _build_sdk_options(False, None, self.SESSION_ID)
retry = _build_retry_sdk_options(initial, False, None, self.SESSION_ID)
assert retry.get("session_id") == self.SESSION_ID
assert "resume" not in retry

def test_retry_removes_session_id_for_t2_plus(self):
"""Retry for T2+ (initial used --resume) removes session_id to avoid conflict."""
initial = _build_sdk_options(True, self.SESSION_ID, self.SESSION_ID)
# T2+ retry where context reduction dropped --resume
retry = _build_retry_sdk_options(initial, False, None, self.SESSION_ID)
assert "session_id" not in retry
assert "resume" not in retry

def test_retry_t2_with_resume_sets_resume(self):
"""Retry that still uses --resume keeps --resume and drops session_id."""
initial = _build_sdk_options(True, self.SESSION_ID, self.SESSION_ID)
retry = _build_retry_sdk_options(
initial, True, self.SESSION_ID, self.SESSION_ID
)
assert retry.get("resume") == self.SESSION_ID
assert "session_id" not in retry
Loading