Skip to content

fix(compact): preserve active skill recovery state #1744

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Komzpa wants to merge 1 commit into from
Closed

fix(compact): preserve active skill recovery state #1744

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
17 changes: 17 additions & 0 deletions app/modules/proxy/_service/http_bridge/helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -627,6 +627,23 @@ def _has_http_bridge_response_output_marker(item: JsonValue) -> bool:
return status in {"completed", "in_progress"}


def _http_bridge_pending_response_events_seen(pending_states: Sequence[Any]) -> int:
return max(
(
max(
int(getattr(state, "response_event_count", 0)),
int(
getattr(state, "response_id", None) is not None
or getattr(state, "latency_response_created_ms", None) is not None
or bool(getattr(state, "downstream_visible", False))
),
)
for state in pending_states
),
default=0,
)


def _http_bridge_input_item_type(item: JsonValue) -> str | None:
if not isinstance(item, dict):
return None
Expand Down
105 changes: 79 additions & 26 deletions app/modules/proxy/_service/http_bridge/streaming.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1511,18 +1511,29 @@ def classify_durable_full_resend(
if durable_lookup is not None and not _http_bridge_models_compatible(durable_lookup.model, payload.model)
else None
)
durable_model_transition_uses_fresh_replay = (
durable_model_transition_lookup is not None
and not forwarded_request
and rewritten_file_account_id is None
and durable_full_resend_fresh_payload is not None
and durable_full_resend_has_safe_fresh_context
and durable_full_resend_is_account_neutral is True
)
durable_model_transition_requires_owner = durable_model_transition_lookup is not None and (
payload.previous_response_id is not None
or bridge_session_key.strength == "hard"
or (
bridge_session_key.affinity_kind == "prompt_cache"
and _http_bridge_request_stage(
headers=headers,
payload=payload,
durable_lookup=durable_model_transition_lookup,
not durable_model_transition_uses_fresh_replay
and (
payload.previous_response_id is not None
or bridge_session_key.strength == "hard"
or (
bridge_session_key.affinity_kind == "prompt_cache"
and _http_bridge_request_stage(
headers=headers,
payload=payload,
durable_lookup=durable_model_transition_lookup,
)
== "follow_up"
and durable_model_transition_lookup.latest_turn_state is not None
)
== "follow_up"
and durable_model_transition_lookup.latest_turn_state is not None
)
)
if durable_model_transition_lookup is not None:
Expand All @@ -1536,7 +1547,36 @@ def classify_durable_full_resend(
model_class=_extract_model_class(payload.model) if payload.model else None,
owner_check_applied=durable_model_transition_requires_owner,
)
if is_http_bridge_account_neutral_replay(
if durable_model_transition_uses_fresh_replay:
replay_kind, replay_key = make_http_bridge_account_neutral_replay_key(uuid4().hex)
bridge_session_key = _HTTPBridgeSessionKey(
replay_kind,
replay_key,
bridge_session_key.api_key_id,
strength="soft",
)
affinity = _AffinityPolicy()
incoming_turn_state_header = None
incoming_session_header = None
session_header_fallback_key = None
effective_payload = durable_full_resend_fresh_payload
untrimmed_effective_payload = durable_full_resend_fresh_payload
force_local_recovery_creation = True
preferred_account_has_continuity_provenance = False
_log_http_bridge_event(
"model_transition_fresh_resend",
bridge_session_key,
account_id=durable_model_transition_lookup.account_id,
model=payload.model,
detail=(
"outcome=account_neutral_full_resend_without_owner,"
f"previous_model={durable_model_transition_lookup.model}"
),
cache_key_family=bridge_session_key.affinity_kind,
model_class=_extract_model_class(payload.model) if payload.model else None,
owner_check_applied=False,
)
elif is_http_bridge_account_neutral_replay(
kind=durable_model_transition_lookup.canonical_kind,
key=durable_model_transition_lookup.canonical_key,
):
Expand Down Expand Up @@ -1719,6 +1759,7 @@ def classify_durable_full_resend(
durable_lookup.account_id
if (
durable_lookup is not None
and not durable_model_transition_uses_fresh_replay
and (
request_state.previous_response_id is not None
or bridge_session_key.strength == "hard"
Expand All @@ -1735,6 +1776,7 @@ def classify_durable_full_resend(
request_state.preferred_account_id is None
and durable_model_transition_lookup is not None
and durable_model_transition_requires_owner
and not durable_model_transition_uses_fresh_replay
):
request_state.preferred_account_id = durable_model_transition_lookup.account_id
local_previous_response_owner: str | None = None
Expand Down Expand Up @@ -1856,6 +1898,7 @@ def classify_durable_full_resend(

def durable_full_resend_allows_account_neutral_replay() -> bool:
nonlocal durable_full_resend_fresh_payload
nonlocal durable_full_resend_has_safe_fresh_context
nonlocal durable_full_resend_is_account_neutral
nonlocal durable_full_resend_retains_prior_output

Expand All @@ -1881,7 +1924,17 @@ def durable_full_resend_allows_account_neutral_replay() -> bool:
stored_count=eligibility_projection.stored_prefix_count,
canonical_lite_developer_index=eligibility_projection.canonical_lite_developer_index,
)
if not durable_full_resend_retains_prior_output:
durable_full_resend_has_safe_fresh_context = durable_full_resend_retains_prior_output or (
durable_lookup is not None
and durable_lookup.latest_pending_tool_calls is not None
and responses_input_suffix_matches_pending_tool_calls(
eligibility_projection.input_items,
stored_count=eligibility_projection.stored_prefix_count,
pending_tool_calls=durable_lookup.latest_pending_tool_calls,
canonical_lite_developer_index=eligibility_projection.canonical_lite_developer_index,
)
)
if not durable_full_resend_has_safe_fresh_context:
return False
replay_projection = project_responses_input_for_account_neutral_fresh_replay(
cast(list[JsonValue], payload.input),
Expand All @@ -1892,7 +1945,7 @@ def durable_full_resend_allows_account_neutral_replay() -> bool:
durable_full_resend_fresh_payload = _http_bridge_payload_without_previous_response_id(
payload
).model_copy(update={"input": replay_projection.input_items})
if not durable_full_resend_retains_prior_output:
if not durable_full_resend_has_safe_fresh_context:
return False
if durable_full_resend_is_account_neutral is None:
durable_full_resend_is_account_neutral = _http_bridge_payload_is_account_neutral_fresh_replay(
Expand Down Expand Up @@ -2813,16 +2866,16 @@ def switch_to_account_neutral_replay() -> None:
previous_request_state.proxy_injected_anchor_had_full_resend_payload
)
request_state.fresh_upstream_request_text = fresh_upstream_request_text
# The trim branch only fires when the untrimmed payload
# is a true full resend whose prefix exactly matches the
# already-stored context, so the unanchored request text
# is a safe fresh-turn replay target regardless of
# whether the anchor came from the durable or
# session-level injection path. Injection-only re-prepares
# keep the replay-safety decision made when the anchor was
# injected.
# The trim branch proves the upstream submission can omit the
# stored prefix, but it does not by itself prove that dropping
# the injected anchor is safe. Keep the original anchor site's
# decision unless this was a durable full-resend proof with a
# verified safe fresh suffix. Session-level anchors may still be
# compacted follow-ups whose prior context only exists behind
# previous_response_id.
request_state.fresh_upstream_request_is_retry_safe = (
(durable_full_resend_anchor_count is None or durable_full_resend_has_safe_fresh_context)
previous_request_state.fresh_upstream_request_is_retry_safe
or (durable_full_resend_anchor_count is not None and durable_full_resend_has_safe_fresh_context)
if store_context_trim_applied
else previous_request_state.fresh_upstream_request_is_retry_safe
)
Expand Down Expand Up @@ -2937,6 +2990,7 @@ async def rollback_pre_dispatch_recovery_claim() -> None:
owner_check_applied=True,
)
replacement_preferred_account_id = request_state.preferred_account_id
replacement_excluded_account_ids = set(request_state.excluded_account_ids)
if request_state.previous_response_id is not None and replacement_preferred_account_id is None:
replacement_preferred_account_id = session.account.id
elif replacement_preferred_account_id is None:
Expand All @@ -2949,9 +3003,8 @@ async def rollback_pre_dispatch_recovery_claim() -> None:
# impossible (fallback_on_preferred_account_unavailable is
# False for exactly this pinned case below) and would
# keep poisoning every later recovery call on this
# request, since excluded_account_ids persists on
# request_state.
request_state.excluded_account_ids.add(session.account.id)
# request.
replacement_excluded_account_ids.add(session.account.id)
while True:
try:
replacement_session = await self._get_or_create_http_bridge_session(
Expand Down Expand Up @@ -2984,7 +3037,7 @@ async def rollback_pre_dispatch_recovery_claim() -> None:
request_usage_budget=request_state.request_usage_budget,
request_deadline=request_deadline,
session_header_fallback_key=session_header_fallback_key,
exclude_account_ids=request_state.excluded_account_ids or None,
exclude_account_ids=replacement_excluded_account_ids or None,
deferred_account_backoff_lifecycle=request_state.deferred_account_backoff_lifecycle,
defer_account_health_writes=request_state.api_key_reservation is not None,
)
Expand Down
90 changes: 66 additions & 24 deletions app/modules/proxy/api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6150,7 +6150,7 @@ async def _compact_responses(
await reservation_cleanup.release(action="compact response")
result_payload = result.model_dump(mode="json", exclude_none=True)
if codex_session_affinity:
result_payload = _normalize_codex_remote_compaction_v2_result(result, result_payload)
result_payload = _normalize_codex_remote_compaction_v2_result(result, result_payload, payload)
return JSONResponse(
content=result_payload,
headers=rate_limit_headers,
Expand All @@ -6160,6 +6160,7 @@ async def _compact_responses(
def _normalize_codex_remote_compaction_v2_result(
payload: CompactResponsePayload,
result_payload: dict[str, JsonValue],
compact_request: ResponsesCompactRequest,
) -> dict[str, JsonValue]:
compaction_item = _compact_response_output_item(payload)
if compaction_item is None:
Expand Down Expand Up @@ -6211,6 +6212,41 @@ def _normalize_compaction_output_item(item: Mapping[str, JsonValue]) -> dict[str
return normalized


def _compaction_output_item_id(value: JsonValue) -> str | None:
if not isinstance(value, str):
return None
item_id = value.strip()
if not item_id.startswith("cmp"):
return None
return item_id


def _compaction_item_texts(value: JsonValue) -> list[str]:
if isinstance(value, str):
return [value]
if not is_json_mapping(value):
return []
content = value.get("content")
if isinstance(content, str):
return [content]
if is_json_mapping(content):
content_items: list[JsonValue] = [content]
elif is_json_list(content):
content_items = content
else:
return []
texts: list[str] = []
for part in content_items:
if isinstance(part, str):
texts.append(part)
continue
if is_json_mapping(part):
text = part.get("text")
if isinstance(text, str):
texts.append(text)
return texts


def _json_mapping_from_model_or_mapping(value: object) -> Mapping[str, JsonValue] | None:
if is_json_mapping(value):
return value
Expand All @@ -6236,13 +6272,15 @@ async def _synthetic_compaction_response_stream(
response_id: str,
usage: object | None,
) -> AsyncIterator[str]:
item = dict(compact_item)
item.setdefault("status", "completed")
compact_output: dict[str, JsonValue] = dict(compact_item)
compact_output.setdefault("status", "completed")
output_items: list[dict[str, JsonValue]] = [compact_output]
completed_response: dict[str, JsonValue] = {
"id": response_id,
"object": "response",
"status": "completed",
"output": [item],
# list is invariant, so the JSON-valid element type needs a cast here.
"output": cast(JsonValue, output_items),
}
usage_mapping = _json_mapping_from_model_or_mapping(usage)
if usage_mapping is not None:
Expand All @@ -6259,29 +6297,33 @@ async def _synthetic_compaction_response_stream(
},
}
)
yield format_sse_event(
{
"type": "response.output_item.added",
"sequence_number": 1,
"output_index": 0,
"item": {
**item,
"status": "in_progress",
},
}
)
yield format_sse_event(
{
"type": "response.output_item.done",
"sequence_number": 2,
"output_index": 0,
"item": item,
}
)
sequence_number = 1
for output_index, item in enumerate(output_items):
yield format_sse_event(
{
"type": "response.output_item.added",
"sequence_number": sequence_number,
"output_index": output_index,
"item": {
**item,
"status": "in_progress",
},
}
)
sequence_number += 1
yield format_sse_event(
{
"type": "response.output_item.done",
"sequence_number": sequence_number,
"output_index": output_index,
"item": item,
}
)
sequence_number += 1
yield format_sse_event(
{
"type": "response.completed",
"sequence_number": 3,
"sequence_number": sequence_number,
"response": completed_response,
}
)
Expand Down
Loading
Loading