Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: preserve Responses reasoning state and compact context (#5172)

Browse Source
This commit is contained in:
chengyongru
2026-07-30 22:39:43 +08:00
committed by GitHub
parent 511c764f45
commit 6a1a45d07a
37 changed files with 4778 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/context.py
+32 -9
View File
@@ -225,9 +225,6 @@ class ContextBuilder:
if current_role == "user"
else []
)
user_content = self.build_user_content(current_message, image_paths=media)
blocks = list(runtime_context_blocks or ()) if current_role == "user" else []
merged, runtime_context_meta = append_runtime_context(user_content, blocks)
messages: list[dict[str, Any]] = [
{
"role": "system",
@@ -243,21 +240,47 @@ class ContextBuilder:
},
*history,
]
current = self.build_current_message(
current_message,
media=media,
current_role=current_role,
runtime_context_blocks=runtime_context_blocks,
)
if messages[-1].get("role") == current_role:
last = dict(messages[-1])
last["content"] = self._merge_message_content(last.get("content"), merged)
if current_role == "user" and runtime_context_meta is not None:
last["content"] = self._merge_message_content(
last.get("content"),
current.get("content"),
)
current_meta = current.get("_meta")
if current_role == "user" and isinstance(current_meta, dict):
internal_meta = dict(last.get("_meta") or {})
internal_meta[RUNTIME_CONTEXT_MESSAGE_META] = runtime_context_meta
internal_meta.update(cast(dict[str, Any], current_meta))
last["_meta"] = internal_meta
messages[-1] = last
return messages
current: dict[str, Any] = {"role": current_role, "content": merged}
if current_role == "user" and runtime_context_meta is not None:
current["_meta"] = {RUNTIME_CONTEXT_MESSAGE_META: runtime_context_meta}
messages.append(current)
return messages
def build_current_message(
self,
current_message: str,
*,
media: list[str] | None = None,
current_role: str = "user",
runtime_context_blocks: Sequence[RuntimeContextBlock] | None = None,
) -> dict[str, Any]:
"""Build only the fresh turn message without merging it into history."""
content = self.build_user_content(current_message, image_paths=media)
blocks = list(runtime_context_blocks or ()) if current_role == "user" else []
merged, runtime_context_meta = append_runtime_context(content, blocks)
current: dict[str, Any] = {"role": current_role, "content": merged}
if current_role == "user" and runtime_context_meta is not None:
current["_meta"] = {
RUNTIME_CONTEXT_MESSAGE_META: runtime_context_meta,
}
return current
def build_user_content(
self,
text: str,
nanobot/agent/loop.py
+117 -5
View File
@@ -49,7 +49,7 @@ from nanobot.bus.queue import MessageBus
from nanobot.bus.runtime_events import RuntimeEventBus
from nanobot.command import CommandContext, CommandRouter, register_builtin_commands
from nanobot.config.schema import AgentDefaults, ModelPresetConfig
from nanobot.providers.base import LLMProvider
from nanobot.providers.base import LLMProvider, ProviderConversationState
from nanobot.providers.factory import ProviderSnapshot
from nanobot.runtime_context import (
RUNTIME_CONTEXT_HISTORY_META,
@@ -106,6 +106,7 @@ if TYPE_CHECKING:
from nanobot.triggers.local_store import LocalTriggerStore
_T = TypeVar("_T")
_SUBAGENT_PROVIDER_TASK_META = "subagent_provider_task_id"
class TurnKind(Enum):
@@ -126,6 +127,7 @@ class TurnContext:
history: list[dict[str, Any]] = field(default_factory=list)
initial_messages: list[dict[str, Any]] = field(default_factory=list)
provider_state: ProviderConversationState | None = field(default=None, repr=False)
request_context: RequestContext | None = None
runtime_context_blocks: list[RuntimeContextBlock] = field(default_factory=list)
attributes: dict[str, Any] = field(default_factory=dict)
@@ -243,6 +245,8 @@ class AgentLoop:
_RUNTIME_CHECKPOINT_KEY = "runtime_checkpoint"
_PENDING_USER_TURN_KEY = "pending_user_turn"
_PROVIDER_STATE_CHECKPOINT_VERSION_KEY = "provider_state_checkpoint_version"
_PROVIDER_STATE_CHECKPOINT_VERSION = "v1"
def __init__(
self,
@@ -857,6 +861,7 @@ class AgentLoop:
turn_scopes: list[AbstractContextManager[Any]] | None = None,
tools: ToolRegistry | None = None,
request_context: RequestContext | None = None,
provider_state: ProviderConversationState | None = None,
) -> tuple[str | None, list[str], list[dict[str, Any]], str, bool]:
"""Run the agent iteration loop.
@@ -872,7 +877,18 @@ class AgentLoop:
async def _checkpoint(payload: dict[str, Any]) -> None:
if session is None:
return
self._set_runtime_checkpoint(session, payload)
public_payload = dict(payload)
private_state = public_payload.pop("provider_state", None)
public_payload.pop(self._PROVIDER_STATE_CHECKPOINT_VERSION_KEY, None)
if "provider_state" in payload and (
private_state is None
or isinstance(private_state, ProviderConversationState)
):
session.provider_state = private_state
public_payload[self._PROVIDER_STATE_CHECKPOINT_VERSION_KEY] = (
self._PROVIDER_STATE_CHECKPOINT_VERSION
)
self._set_runtime_checkpoint(session, public_payload)
async def _drain_pending(*, limit: int = _MAX_INJECTIONS_PER_TURN) -> list[dict[str, Any]]:
"""Drain follow-up messages from the pending queue.
@@ -1070,6 +1086,7 @@ class AgentLoop:
session_metadata=session_metadata,
message_metadata=metadata,
),
provider_state=provider_state,
))
finally:
turn_scope_stack.close()
@@ -1077,6 +1094,8 @@ class AgentLoop:
reset_request_context(request_token)
reset_file_states(file_state_token)
self._last_usage = result.usage
if session is not None and not ephemeral:
session.provider_state = result.provider_state
if result.stop_reason == "max_iterations":
logger.warning("Max iterations ({}) reached", self.max_iterations)
should_stream = turn_continuation.should_stream_budget_response(
@@ -1660,14 +1679,24 @@ class AgentLoop:
"extend_to_user": is_subagent,
}
ctx.history = session.get_history(**_hist_kwargs)
stored_state = session.provider_state
subagent_followup_persisted = False
if is_subagent:
# Keep the durable internal delivery as an assistant record, but
# present this completion to the model as fresh follow-up input.
# Providers without assistant-prefill support drop trailing
# assistant messages, so using the persisted record as the current
# prompt would hide an independently dispatched subagent result.
if self._persist_subagent_followup(session, ctx.msg):
subagent_followup_persisted = self._persist_subagent_followup(
session,
ctx.msg,
)
if subagent_followup_persisted:
logger.debug("Subagent result persisted for session {}", ctx.session_key)
# Establish a durable, replay-safe baseline before any fallible
# provider compatibility or prompt assembly work. A compatible
# staged state replaces this in a second atomic save below.
session.provider_state = None
self.sessions.save(session)
ctx.input_persisted_early = True
ctx.delivery.record_runtime(runtime)
@@ -1675,13 +1704,65 @@ class AgentLoop:
ctx.request_context = self._request_context_for_turn(ctx)
if ctx.kind is TurnKind.USER:
ctx.runtime_context_blocks = await self._resolve_runtime_context_for_turn(ctx)
ctx.initial_messages = self._build_initial_messages(ctx)
staged_provider_state = False
if stored_state is not None and runtime.provider.can_resume_conversation_state(
stored_state,
runtime.model,
):
current_provider_message = self.context.build_current_message(
ctx.msg.content,
media=ctx.msg.media if ctx.kind is TurnKind.USER and ctx.msg.media else None,
runtime_context_blocks=ctx.runtime_context_blocks,
)
task_id = ctx.msg.metadata.get("subagent_task_id") if is_subagent else None
already_staged = False
if isinstance(task_id, str) and task_id:
internal_meta = current_provider_message.get("_meta")
current_provider_message["_meta"] = {
**(
cast(dict[str, Any], internal_meta)
if isinstance(internal_meta, dict)
else {}
),
_SUBAGENT_PROVIDER_TASK_META: task_id,
}
already_staged = any(
isinstance(message.get("_meta"), dict)
and cast(dict[str, Any], message["_meta"]).get(
_SUBAGENT_PROVIDER_TASK_META
)
== task_id
for message in stored_state.pending_messages
)
ctx.provider_state = (
stored_state
if already_staged
else stored_state.with_pending_messages([
*stored_state.pending_messages,
current_provider_message,
])
)
if (
not ctx.ephemeral
and (ctx.kind is TurnKind.USER or subagent_followup_persisted)
):
session.provider_state = ctx.provider_state
staged_provider_state = True
elif stored_state is not None:
session.provider_state = None
if ctx.kind is TurnKind.USER:
ctx.input_persisted_early = self._persist_user_message_early(
ctx.msg,
session,
runtime_context_blocks=ctx.runtime_context_blocks,
)
if staged_provider_state and not ctx.input_persisted_early:
session.provider_state = stored_state
elif subagent_followup_persisted and staged_provider_state:
# Upgrade the replay-safe baseline to the resumable state before
# prompt assembly and the first model checkpoint.
self.sessions.save(session)
ctx.initial_messages = self._build_initial_messages(ctx)
if ctx.on_progress is None:
ctx.on_progress = ctx.delivery.progress_callback()
@@ -1715,6 +1796,7 @@ class AgentLoop:
turn_scopes=ctx.turn_scopes,
tools=ctx.tools,
request_context=ctx.request_context,
provider_state=ctx.provider_state,
)
final_content, _, all_msgs, stop_reason, had_injections = result
ctx.final_content = final_content
@@ -2052,7 +2134,36 @@ class AgentLoop:
):
overlap = size
break
session.messages.extend(restored_messages[overlap:])
appended_messages = restored_messages[overlap:]
session.messages.extend(appended_messages)
assistant_message_data = (
cast(dict[str, Any], assistant_message)
if isinstance(assistant_message, dict)
else None
)
provider_state_is_synchronized = (
checkpoint_data.get(self._PROVIDER_STATE_CHECKPOINT_VERSION_KEY)
== self._PROVIDER_STATE_CHECKPOINT_VERSION
)
phase = checkpoint_data.get("phase")
exact_final_response = (
phase == "final_response"
and assistant_message_data is not None
and assistant_message_data.get("role") == "assistant"
and not bool(checkpoint_data.get("completed_tool_results"))
and not bool(checkpoint_data.get("pending_tool_calls"))
)
exact_completed_tools = (
phase == "tools_completed"
and assistant_message_data is not None
and assistant_message_data.get("role") == "assistant"
and not bool(checkpoint_data.get("pending_tool_calls"))
)
if not (
provider_state_is_synchronized
and (exact_final_response or exact_completed_tools)
):
session.provider_state = None
self._clear_pending_user_turn(session)
self._clear_runtime_checkpoint(session)
@@ -2073,6 +2184,7 @@ class AgentLoop:
"timestamp": datetime.now().isoformat(),
}
)
session.provider_state = None
session.updated_at = datetime.now()
self._clear_pending_user_turn(session)
nanobot/agent/memory.py
+3
View File
@@ -931,6 +931,7 @@ class Consolidator:
session_key=session.key,
)
session.last_consolidated = end_idx
session.provider_state = None
self.sessions.save(session)
return summary
@@ -1136,6 +1137,7 @@ class Consolidator:
if summary:
last_summary = summary
session.last_consolidated = end_idx
session.provider_state = None
self.sessions.save(session)
if not summary:
# LLM is degraded — stop hammering it this call;
@@ -1205,6 +1207,7 @@ class Consolidator:
# Preserve history and advance only the replay boundary.
session.last_consolidated = len(session.messages) - len(visible_suffix)
session.provider_state = None
self.sessions.save(session)
logger.info(
nanobot/agent/runner.py
+166 -28
View File
@@ -19,7 +19,17 @@ from nanobot.agent.context_governance import (
)
from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext
from nanobot.agent.tools.registry import ToolRegistry, is_tool_error_result
from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest
from nanobot.providers.base import (
LLMProvider,
LLMResponse,
ProviderCallContext,
ProviderConversationState,
ToolCallRequest,
)
from nanobot.providers.conversation_state import (
ProviderConversationStateController,
allows_conversation_message_merge,
)
from nanobot.runtime_context import (
RUNTIME_CONTEXT_MESSAGE_META,
detach_runtime_context,
@@ -104,6 +114,7 @@ class AgentRunSpec:
goal_active_predicate: Callable[[], bool] | None = None
goal_continue_message: GoalContinueMessage | None = None
finalize_on_max_iterations: bool = True
provider_state: ProviderConversationState | None = None
@dataclass(slots=True)
@@ -120,6 +131,7 @@ class AgentRunResult:
had_injections: bool = False
# Terminal tail to emit when the preceding final-content prefix was already streamed.
pending_stream_content: str | None = None
provider_state: ProviderConversationState | None = field(default=None, repr=False)
class AgentRunner:
@@ -161,6 +173,7 @@ class AgentRunner:
and messages[-1].get("role") == "user"
and not is_hidden_history_message(injection)
and not is_hidden_history_message(messages[-1])
and allows_conversation_message_merge(messages[-1])
):
merged = dict(messages[-1])
left_meta = merged.get("_meta")
@@ -231,6 +244,7 @@ class AgentRunner:
assistant_message: dict[str, Any] | None,
injection_cycles: int,
*,
conversation_state: ProviderConversationStateController | None = None,
phase: str = "after error",
iteration: int | None = None,
allow_goal_continue: bool = False,
@@ -258,16 +272,21 @@ class AgentRunner:
if assistant_message is not None:
messages.append(assistant_message)
if iteration is not None:
checkpoint: dict[str, Any] = {
"phase": "final_response",
"iteration": iteration,
"model": spec.runtime.model,
"assistant_message": assistant_message,
"completed_tool_results": [],
"pending_tool_calls": [],
}
if conversation_state is not None:
checkpoint["provider_state"] = conversation_state.checkpoint(
messages
)
await self._emit_checkpoint(
spec,
{
"phase": "final_response",
"iteration": iteration,
"model": spec.runtime.model,
"assistant_message": assistant_message,
"completed_tool_results": [],
"pending_tool_calls": [],
},
checkpoint,
)
self._append_injected_messages(messages, injections)
if real_injection:
@@ -420,6 +439,12 @@ class AgentRunner:
injection_cycles = 0
compacted_tool_call_ids: set[str] = set()
pending_stream_content: str | None = None
conversation_state = ProviderConversationStateController(
provider=spec.runtime.provider,
model=spec.runtime.model,
messages=messages,
state=spec.provider_state,
)
governance_config = ContextGovernanceConfig(
provider=spec.runtime.provider,
model=spec.runtime.model,
@@ -450,7 +475,20 @@ class AgentRunner:
session_key=spec.session_key,
)
await hook.before_iteration(context)
response = await self._request_model(spec, messages_for_model, hook, context)
provider_context = conversation_state.prepare_request(
messages,
context_window_tokens=spec.runtime.context_window_tokens,
model_messages=messages_for_model,
)
response = await self._request_model(
spec,
messages_for_model,
hook,
context,
conversation_state=conversation_state,
provider_context=provider_context,
)
conversation_state.observe_response(response, messages)
context.response = response
context.tool_calls = list(response.tool_calls)
@@ -480,6 +518,10 @@ class AgentRunner:
reasoning_content=response.reasoning_content,
thinking_blocks=response.thinking_blocks,
)
assistant_message = conversation_state.project_response_message(
assistant_message,
response,
)
messages.append(assistant_message)
await self._emit_checkpoint(
spec,
@@ -544,6 +586,15 @@ class AgentRunner:
length_recovery_parts.clear()
continue
break
checkpoint_model_messages = (
self.context_governor.prepare_for_model(
governance_config,
messages,
compacted_tool_call_ids,
)
if response.provider_state is not None
else None
)
await self._emit_checkpoint(
spec,
{
@@ -553,6 +604,10 @@ class AgentRunner:
"assistant_message": assistant_message,
"completed_tool_results": completed_tool_results,
"pending_tool_calls": [],
"provider_state": conversation_state.checkpoint(
messages,
model_messages=checkpoint_model_messages,
),
},
)
empty_content_retries = 0
@@ -575,7 +630,11 @@ class AgentRunner:
)
clean = hook.finalize_content(context, response.content)
if response.finish_reason not in ("error", "length") and is_blank_text(clean):
if (
response.finish_reason
not in {"error", "length", "refusal", "content_filter"}
and is_blank_text(clean)
):
empty_content_retries += 1
if empty_content_retries < _MAX_EMPTY_RETRIES:
logger.warning(
@@ -598,7 +657,12 @@ class AgentRunner:
if hook.wants_streaming():
await hook.on_stream_end(context, resuming=False)
retry_messages = self._finalization_retry_messages(messages_for_model)
response = await self._request_finalization_retry(spec, messages_for_model)
response = await self._request_finalization_retry(
spec,
messages_for_model,
transcript=messages,
conversation_state=conversation_state,
)
retry_usage = self._usage_or_estimate(spec, retry_messages, response)
self._accumulate_usage(usage, retry_usage)
raw_usage = self._merge_usage(raw_usage, retry_usage)
@@ -623,10 +687,13 @@ class AgentRunner:
if hook.wants_streaming():
context.stream_continues_current_message = True
await hook.on_stream_end(context, resuming=True)
messages.append(build_assistant_message(
clean,
reasoning_content=response.reasoning_content,
thinking_blocks=response.thinking_blocks,
messages.append(conversation_state.project_response_message(
build_assistant_message(
clean,
reasoning_content=response.reasoning_content,
thinking_blocks=response.thinking_blocks,
),
response,
))
messages.append(build_length_recovery_message(clean or ""))
await hook.after_iteration(context)
@@ -656,15 +723,22 @@ class AgentRunner:
reasoning_content=response.reasoning_content,
thinking_blocks=response.thinking_blocks,
)
assistant_message = conversation_state.project_response_message(
assistant_message,
response,
)
# Check for mid-turn injections BEFORE signaling stream end.
# If injections are found we keep the stream alive (resuming=True)
# so streaming channels don't prematurely finalize the card.
should_continue, injection_cycles = await self._try_drain_injections(
spec, messages, assistant_message, injection_cycles,
conversation_state=conversation_state,
phase="after final response",
iteration=iteration,
allow_goal_continue=True,
allow_goal_continue=(
response.finish_reason not in {"refusal", "content_filter"}
),
)
if should_continue:
had_injections = True
@@ -717,11 +791,17 @@ class AgentRunner:
continue
break
messages.append(assistant_message or build_assistant_message(
clean,
reasoning_content=response.reasoning_content,
thinking_blocks=response.thinking_blocks,
))
messages.append(
assistant_message
or conversation_state.project_response_message(
build_assistant_message(
clean,
reasoning_content=response.reasoning_content,
thinking_blocks=response.thinking_blocks,
),
response,
)
)
await self._emit_checkpoint(
spec,
{
@@ -731,6 +811,7 @@ class AgentRunner:
"assistant_message": messages[-1],
"completed_tool_results": [],
"pending_tool_calls": [],
"provider_state": conversation_state.checkpoint(messages),
},
)
if length_recovery_parts:
@@ -764,6 +845,7 @@ class AgentRunner:
hook,
messages,
usage,
conversation_state,
)
if terminal_content is None:
terminal_content = self._max_iterations_fallback(spec)
@@ -787,6 +869,7 @@ class AgentRunner:
tool_events=tool_events,
had_injections=had_injections,
pending_stream_content=pending_stream_content,
provider_state=conversation_state.finish(messages),
)
def _build_request_kwargs(
@@ -817,6 +900,8 @@ class AgentRunner:
context: AgentHookContext,
*,
malformed_retry: bool = False,
conversation_state: ProviderConversationStateController,
provider_context: ProviderCallContext | None = None,
) -> LLMResponse:
timeout_s: float | None = spec.llm_timeout_s
if timeout_s is None:
@@ -886,6 +971,7 @@ class AgentRunner:
coro = spec.runtime.provider.chat_stream_with_retry(
**kwargs,
provider_context=provider_context,
on_content_delta=_stream,
on_thinking_delta=_thinking,
on_tool_call_delta=_provider_tool_event,
@@ -920,11 +1006,15 @@ class AgentRunner:
coro = spec.runtime.provider.chat_stream_with_retry(
**kwargs,
provider_context=provider_context,
on_content_delta=_stream_progress,
on_tool_call_delta=_provider_tool_event,
)
else:
coro = spec.runtime.provider.chat_with_retry(**kwargs)
coro = spec.runtime.provider.chat_with_retry(
**kwargs,
provider_context=provider_context,
)
# Streaming requests also have provider-level idle timeouts
# (NANOBOT_STREAM_IDLE_TIMEOUT_S), but a stream that keeps producing
@@ -986,6 +1076,10 @@ class AgentRunner:
return await self._request_model(
spec, retry_messages, hook, context,
malformed_retry=True,
conversation_state=conversation_state,
provider_context=conversation_state.independent_request_context(
context_window_tokens=spec.runtime.context_window_tokens,
),
)
if (
all_dropped
@@ -998,7 +1092,13 @@ class AgentRunner:
fallback_messages = self._malformed_tool_call_retry_messages(
messages, response.content,
)
return await self._request_no_tools(spec, fallback_messages)
return await self._request_no_tools(
spec,
fallback_messages,
provider_context=conversation_state.independent_request_context(
context_window_tokens=spec.runtime.context_window_tokens,
),
)
return response
@staticmethod
@@ -1031,6 +1131,10 @@ class AgentRunner:
original_finish_reason,
)
response.tool_calls = valid
# The opaque candidate still contains every raw function_call item.
# Advancing it after dropping even one call would replay an unmatched
# call without a corresponding tool output on the next request.
response.provider_state = None
if not valid:
response.finish_reason = "stop"
return (dropped, not valid, original_finish_reason)
@@ -1060,9 +1164,27 @@ class AgentRunner:
self,
spec: AgentRunSpec,
messages: list[dict[str, Any]],
*,
transcript: list[dict[str, Any]],
conversation_state: ProviderConversationStateController,
) -> LLMResponse:
retry_messages = self._finalization_retry_messages(messages)
return await self._request_no_tools(spec, retry_messages)
provider_context = conversation_state.prepare_request(
transcript,
context_window_tokens=spec.runtime.context_window_tokens,
supplemental_messages=[retry_messages[-1]],
)
response = await self._request_no_tools(
spec,
retry_messages,
provider_context=provider_context,
)
conversation_state.observe_response(
response,
transcript,
adopt_candidate_state=False,
)
return response
@staticmethod
def _finalization_retry_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
@@ -1076,10 +1198,17 @@ class AgentRunner:
hook: AgentHook,
messages: list[dict[str, Any]],
usage: dict[str, int],
conversation_state: ProviderConversationStateController,
) -> str | None:
retry_messages = self._budget_exhausted_finalization_messages(messages)
try:
response = await self._request_no_tools(spec, retry_messages)
response = await self._request_no_tools(
spec,
retry_messages,
provider_context=conversation_state.independent_request_context(
context_window_tokens=spec.runtime.context_window_tokens,
),
)
except Exception:
logger.exception(
"Budget-exhausted finalization failed for {}; using fallback",
@@ -1115,9 +1244,18 @@ class AgentRunner:
self,
spec: AgentRunSpec,
messages: list[dict[str, Any]],
*,
provider_context: ProviderCallContext | None = None,
) -> LLMResponse:
kwargs = self._build_request_kwargs(spec, messages, tools=None)
return await spec.runtime.provider.chat_with_retry(**kwargs)
kwargs = self._build_request_kwargs(
spec,
messages,
tools=None,
)
return await spec.runtime.provider.chat_with_retry(
**kwargs,
provider_context=provider_context,
)
@staticmethod
def _budget_exhausted_finalization_messages(