Skip to content

Add session usage ledger and refine step-loop plan handling - #154

Merged
JohnRichard4096 merged 7 commits into
mainfrom
feat-usage-fix
Aug 15, 2026
Merged

Add session usage ledger and refine step-loop plan handling#154
JohnRichard4096 merged 7 commits into
mainfrom
feat-usage-fix

Conversation

@JohnRichard4096

@JohnRichard4096 JohnRichard4096 commented Aug 15, 2026

Copy link
Copy Markdown
Member

Summary by Sourcery

Introduce a run-scoped usage ledger, strengthen the step-driven ReAct workflow (plan visibility, autonomous revision, token budgeting), harden MCP client concurrency, and update docs/demos to reflect explicit step-loop opt-in and adapter/provider behavior.

New Features:

  • Add session-scoped usage ledger and proxy to track per-run model usage independently of final completion responses.
  • Expose the plan-revision built-in tool only when the native step loop is active and inject current plan status into the context so the model can revise plans mid-run.
  • Add deterministic tool failure hints after hard ERROR results to guide the model toward update_step-based plan revision instead of infinite retries.
  • Provide new demos illustrating autonomous plan revision in the step loop and generic OpenAI-compatible endpoint configuration.

Bug Fixes:

  • Ensure reasoning_content from thinking-mode providers is preserved and round-tripped on tool error paths to avoid HTTP 400 "must be passed back" failures.
  • Fix MCP client connection handling so concurrent simple_call invocations share a single connection without premature close and respect connection_ttl semantics.
  • Correct between-step token budgeting to derive prompt windows from recorded usage rather than ad-hoc accumulation, avoiding mis-triggered compression and budget resets.

Enhancements:

  • Refine stall detection, step lifecycle, and history compression to operate per iteration within the loop and to keep tool-call/result pairing intact while resetting baselines safely.
  • Streamline usage accounting by routing auxiliary LLM calls (decomposition, step summaries, memory abstraction, reasoning tools) through the usage proxy instead of manual aggregation.
  • Improve update_step behavior by echoing the revised plan and completed steps back to the model immediately after a plan change.
  • Relax and clarify the chat workflows so simple chat is the default and the step-driven ReAct loop is explicitly enabled via workflow selection.
  • Strengthen agent strategy initialization to ensure step tools are injected idempotently and peer input drainage and token windows are refreshed at step boundaries.

Documentation:

  • Revise troubleshooting guides (EN/ZH) to document reasoning_content requirements, stall detection placement, token burn controls, adapter/protocol usage, and common async/testing pitfalls.
  • Update workflow engine and builtins documentation (EN/ZH) to distinguish simple chat from step-driven pipelines and to describe the available pre-composed workflows and their selection.
  • Clarify adapter docs (EN/ZH) around protocol vs provider, showing how OpenAI-compatible and Anthropic endpoints are configured and how create_agent relates to ModelPreset and AgentRuntime.
  • Extend ChatObject, agent-strategy, streaming, event-hooks, and getting-started docs to explain explicit step-loop opt-in, step metadata emission, and usage of workflow=_step_workflow_rendered.
  • Document MCP server URL formats and transports (EN/ZH), including streamable+http(s) syntax, SSE shorthands, and connection TTL behavior.

Tests:

  • Add tests for reasoning_content preservation on tool error append, ensuring assistant messages in error paths carry provider reasoning when present.
  • Extend step loop tests to validate history compression behavior, token budget survival across baseline resets, tool pairing integrity, and use of the usage ledger proxy in threshold checks.
  • Add MCP client tests for simple_call success, error propagation, concurrent calls sharing connections, TTL -1 resident behavior, and safe close while calls are in flight.
  • Introduce usage ledger tests to verify registry lifecycle, snapshotting, prompt_since windows, and libchat gateway recording for both tools_caller and call_completion.
  • Update coverage-gap tests to ensure memory limiter usage is recorded via the usage ledger rather than direct extra_usage aggregation.

Chores:

  • Bump project version from 0.13.1 to 0.13.2 in pyproject.toml and update lockfile.
  • Normalize comments and docstrings from unicode arrows to ASCII arrows for consistency in tests and source comments.

@sourcery-ai sourcery-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry @JohnRichard4096, your pull request is larger than the review limit of 150000 diff characters

@JohnRichard4096

Copy link
Copy Markdown
Member Author

@sourcery-ai title

@sourcery-ai sourcery-ai Bot changed the title Feat usage fix Add session usage ledger and refine step-loop plan handling Aug 15, 2026
@sourcery-ai

sourcery-ai Bot commented Aug 15, 2026

Copy link
Copy Markdown
Contributor

Reviewer's Guide

This PR introduces a run-scoped usage ledger wired into libchat and strategies, refactors token accounting and step-loop behaviors to use the ledger instead of ad-hoc accumulation, hardens MCP client connection management under concurrency, and extends the ReAct step-loop strategy with plan-status injection, tool-failure guidance, and correct reasoning_content round-tripping, along with updated demos, docs, and tests.

Sequence diagram for tool_failure_hint_and_plan_revision_in_step_loop

sequenceDiagram
    participant Strategy as ReActAgentStrategy
    participant RunState as AgentRunState
    participant LLM as Model

    Strategy->>RunState: _init_run_state()
    Strategy->>LLM: single_execute (tools_caller)
    LLM-->>Strategy: ToolCall db_query
    Strategy->>Strategy: _exec_one(tc)
    Strategy-->>Strategy: func_response="ERROR: ..."
    Strategy->>Strategy: _append_tool_result_to_context(tool_call, func_response)
    Strategy->>Strategy: _maybe_inject_tool_failure_hint(tool_call, func_response)
    Strategy->>RunState: tool_error_hints += 1
    Strategy->>LLM: Message role=user (failure_hint)

    LLM-->>Strategy: ToolCall update_step
    Strategy->>Strategy: _handle_update_step(args)
    Strategy->>RunState: plan_revision += 1
    Strategy-->>LLM: ToolResult "<STEP_PLAN_UPDATED> revised plan: ..."

    Strategy->>RunState: leave_step()
    Strategy->>RunState: begin_step("execute")
    Strategy->>RunState: step_started_ts = time.time()
    Strategy->>Strategy: _inject_plan_status()
    Strategy-->>LLM: Message role=user ("[Plan status]\n- node_id [state]: ...")
Loading

File-Level Changes

Change Details Files
Introduce a session-scoped usage ledger and wire it through libchat, ChatObject, strategies, memory limiter, and tests to track process usage separately from final completion usage.
  • Add UsageEntry, UsageLedger, SessionUsageProxy, UsageRegistry, and UsageSnapshot for per-stream accounting.
  • Update StrategyContext/StrategyBase to expose a usage proxy instead of resp_extra_usage and remove gather_usage-based accumulation.
  • Wire ChatObject._entry to register/unregister usage ledgers, keep a UsageSnapshot, and pass the proxy into strategy context and memory limiter.
  • Modify libchat.call_completion and tools_caller to accept an optional usage proxy and record provider-reported usage via the gateway.
  • Adjust MemoryLimiter to accept a usage proxy, record provider usage when present, and fall back to local estimation when absent.
  • Add tests covering ledger behavior, registry lifecycle, libchat gateway recording, and integration with the workflow node _limiting_memory.
src/amrita_core/usage.py
src/amrita_core/agent/strategy.py
src/amrita_core/agent/context.py
src/amrita_core/chatmanager/chat_object.py
src/amrita_core/chatmanager/memory_limiter.py
src/amrita_core/libchat.py
src/amrita_core/contexts.py
tests/test_usage.py
tests/test_coverage_gaps.py
Refine step-loop token budgeting and history compression to use the ledger window, reset baselines correctly, and keep tool-call/result pairing intact while refreshing step timestamps.
  • Extend TokenBudget with a refresh_window method that reads prompt tokens since a given timestamp from the usage proxy.
  • Add step_started_ts and tool_error_hints to AgentRunState and set/reset them in begin_step.
  • Update _compress_history_between_steps to refresh the token window from the ledger, reset tokens and step_started_ts on early-return paths, and pass usage into call_completion.
  • Ensure iter_cond refreshes the token window via strategy.usage and respects exhausted/budget conditions.
  • Adjust tests to bind a UsageRegistry ledger, record UniResponseUsage per step, and verify compression behavior, baseline resets, and budget survival using prompt_since rather than direct token updates.
src/amrita_core/builtins/agent/state.py
src/amrita_core/builtins/agent/react_comm.py
src/amrita_core/components/react.py
tests/test_step_loop.py
Extend the ReActAgentStrategy step-loop behavior with plan-status injection at step boundaries, deterministic tool-failure hints to encourage update_step, and correct reasoning_content handling on error paths and builtin tools.
  • Add _step_tools_injected and _last_plan_snapshot fields to the base strategy and implement _ensure_step_tools to expose UPDATE_STEP_TOOL once when the native step loop activates.
  • Implement _inject_plan_status in ReActAgentStrategy to append a plaintext plan status snapshot plus guidance note at step intro only when the snapshot changes.
  • Call _ensure_step_tools and rs.tokens.refresh_window in intro_step, then inject plan status at every step intro in plan mode.
  • Add _maybe_inject_tool_failure_hint to append user guidance after hard ERROR tool results, using tool_error_hints to change messaging between first and subsequent failures.
  • Update _append_tool_result_to_context to invoke _maybe_inject_tool_failure_hint, ensuring hints come after the tool result to keep pairing intact.
  • Enhance update_step execution to echo the revised plan IDs and completion list in the ToolResult, so the model can confirm revisions immediately.
  • Extend _handle_error_append signatures to accept the original provider response and propagate response_msg.reasoning_content onto fabricated assistant messages, and thread response_msg through _execute_tool_loop error branches.
  • Pass usage proxies into tools_caller and single_execute calls so auxiliary tool rounds are accounted in the ledger.
src/amrita_core/builtins/agent/react_base.py
src/amrita_core/builtins/agent/react_comm.py
src/amrita_core/builtins/agent/react_hyb.py
Fix MCPClient connection lifecycle under concurrency and TTL, ensuring shared connections are not torn down while calls are in flight and close behavior respects active_calls and connection_ttl semantics.
  • Track _active_calls on MCPClient, incrementing at simple_call entry and decrementing in finally.
  • Wrap call_tool in async with on the fastmcp Client to use its reentrant context manager for shared connections.
  • Change simple_call finally logic to schedule close via a TTL-aware close() only when _active_calls drops to zero and connection_ttl != -1.
  • Guard _close so it becomes a no-op while _active_calls > 0 and only swaps out and __aexit__s the client when no calls are active.
  • Adjust tests to patch close instead of _close, ensure aexit does not suppress exceptions, and add new tests for concurrent simple_call behavior, TTL=-1 residency, and _close no-op semantics with in-flight calls.
src/amrita_core/tools/mcp.py
tests/test_mcp.py
Clarify workflows (simple chat vs step-driven ReAct), adapter/protocol usage, MCP server URL formats, and troubleshooting guidance in English and Chinese docs, along with demos updated to use generic OpenAI-compatible configuration.
  • Update workflow-engine and builtins docs to distinguish the default simple-chat pipeline from opt-in step-driven workflows, describe *_ONLY and CHATOBJECT_STEP_REACT variants, and document explicit workflow selection via workflow=_step_workflow_rendered or SIMPLE_STEP_REACT.
  • Clarify adapter semantics: OpenAIAdapter serves any OpenAI-compatible endpoint via base_url/model; AnthropicAdapter needs protocol="anthropic" via ModelPreset passed to AgentRuntime; create_agent has no protocol parameter and always uses the default adapter.
  • Document MCP server script formats including streamable+http(s)://, sse+http(s)://, sse:// shorthand, stdio://["cmd",...] and note that streamable (not stream) is the correct extra; expand usage to cover concurrency-safe connection pooling and TTL behavior.
  • Revise troubleshooting to cover stall detection placement, reasoning_content round-tripping rules per provider, request-id headers, token budget knobs, plan revision behavior, and async/testing pitfalls.
  • Adjust tutorials (chat-object, streaming, event-hooks, index) and concepts to emphasize that step metadata and update_step exist only when the step-loop workflow is active, and that simple chat is the default.
  • Update demos (step_loop_demo, hybrid_demo, peer_push_demo, new step_update_demo) to parameterize BASE_URL/MODEL via environment variables and describe OpenAI-compatible endpoints rather than DeepSeek-specific behavior.
  • Increment project version from 0.13.1 to 0.13.2 and update API reference docs for ModelPreset.protocol default and create_agent behavior.
docs/docs/guide/advanced/workflow-engine.md
docs/docs/zh/guide/advanced/workflow-engine.md
docs/docs/guide/builtins.md
docs/docs/zh/guide/builtins.md
docs/docs/guide/extensions-integration/adapters.md
docs/docs/zh/guide/extensions-integration/adapters.md
docs/docs/guide/extensions-integration/mcp-server.md
docs/docs/zh/guide/extensions-integration/mcp-server.md
docs/docs/guide/concepts/chat-object.md
docs/docs/zh/guide/concepts/chat-object.md
docs/docs/guide/concepts/agent-strategy.md
docs/docs/zh/guide/concepts/agent-strategy.md
docs/docs/guide/concepts/index.md
docs/docs/zh/guide/concepts/index.md
docs/docs/guide/tutorials/chat-object.md
docs/docs/zh/guide/tutorials/chat-object.md
docs/docs/guide/tutorials/event-hooks.md
docs/docs/zh/guide/tutorials/event-hooks.md
docs/docs/guide/tutorials/streaming.md
docs/docs/zh/guide/tutorials/streaming.md
docs/docs/guide/tutorials/index.md
docs/docs/zh/guide/tutorials/index.md
docs/docs/guide/getting-started/basic-example.md
docs/docs/zh/guide/getting-started/basic-example.md
docs/docs/guide/getting-started/minimal-example.md
docs/docs/zh/guide/getting-started/minimal-example.md
docs/docs/guide/api-reference/index.md
docs/docs/zh/guide/api-reference/index.md
docs/docs/guide/api-reference/classes/ModelPreset.md
docs/docs/zh/guide/api-reference/classes/ModelPreset.md
docs/docs/guide/advanced/step-loop.md
docs/docs/zh/guide/advanced/step-loop.md
docs/docs/guide/advanced/suspend.md
docs/docs/zh/guide/advanced/suspend.md
docs/docs/guide/agent-engineering/troubleshooting.md
docs/docs/zh/guide/agent-engineering/troubleshooting.md
docs/docs/guide/extensions-integration/mcp-server.md
docs/docs/zh/guide/extensions-integration/mcp-server.md
demo/step_loop_demo.py
demo/hybrid_demo.py
demo/peer_push_demo.py
demo/step_update_demo.py
pyproject.toml
Minor text/style fixes and type updates across constants, parser, metadata types, and tests to standardize arrow notation and integrate Function type usage.
  • Replace Unicode arrows with ASCII '->' in comments, docstrings, and user-facing text for consistency.
  • Update AgentStepDecomposeMetadata.descriptions docstring to describe a node-id-to-description map using '->'.
  • Adjust builtins consts and consts summaries to use '->' in reasoning-phase descriptions and formatting rules.
  • Change parser transport registry comments to use '->' notation and clarify shorthand behavior.
  • Update tests to import and use Function for ToolCall.function where appropriate instead of raw dicts with pyright ignores.
  • Add minor doc clarifications around HybridReActAgentStrategy reasoning_content behavior and DeepSeek/Anthropic requirements.
tests/test_step_loop.py
tests/test_parser.py
src/amrita_core/builtins/consts.py
src/amrita_core/builtins/types.py
src/amrita_core/consts.py
src/amrita_core/tools/_parser.py
docs/docs/zh/guide/api-reference/classes/HybridReActAgentStrategy.md

Tips and commands

Interacting with Sourcery

  • Trigger a new review: Comment @sourcery-ai review on the pull request.
  • Continue discussions: Reply directly to Sourcery's review comments.
  • Generate a GitHub issue from a review comment: Ask Sourcery to create an
    issue from a review comment by replying to it. You can also reply to a
    review comment with @sourcery-ai issue to create an issue from it.
  • Generate a pull request title: Write @sourcery-ai anywhere in the pull
    request title to generate a title at any time. You can also comment
    @sourcery-ai title on the pull request to (re-)generate the title at any time.
  • Generate a pull request summary: Write @sourcery-ai summary anywhere in
    the pull request body to generate a PR summary at any time exactly where you
    want it. You can also comment @sourcery-ai summary on the pull request to
    (re-)generate the summary at any time.
  • Generate reviewer's guide: Comment @sourcery-ai guide on the pull
    request to (re-)generate the reviewer's guide at any time.
  • Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
    pull request to resolve all Sourcery comments. Useful if you've already
    addressed all the comments and don't want to see them anymore.
  • Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
    request to dismiss all existing Sourcery reviews. Especially useful if you
    want to start fresh with a new review - don't forget to comment
    @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

  • Enable or disable review features such as the Sourcery-generated pull request
    summary, the reviewer's guide, and others.
  • Change the review language.
  • Add, remove or edit custom review instructions.
  • Adjust other review settings.

Getting Help

@JohnRichard4096
JohnRichard4096 merged commit 63b893d into main Aug 15, 2026
9 checks passed
@JohnRichard4096
JohnRichard4096 deleted the feat-usage-fix branch August 15, 2026 06:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant