ag-ui-protocol / ag-ui-protocol/ag-ui

[Bug]: ag-ui-claude-sdk only injects shared state into the model's context on a thread's first turn

未關閉
#2,246 1 則留言 0 個 reaction 已指派 0 人 在 GitHub 檢視
主要語言
Python
星號
15.9k
分支
1.4k
平均合併
1 天 17 小時
30 天內合併 PR
163

描述

### Pre-flight Checklist

- [x] I have searched [existing issues](https://github.com/ag-ui-protocol/ag-ui/issues) and this hasn't been reported yet.
- [x] I am using the **latest** version of AG-UI.

---

### Describe the Bug

Once a thread's session worker is created, the model never sees an updated `RunAgentInput.state` again — only the very first turn's state ever reaches it. `STATE_SNAPSHOT` is still echoed to the client correctly on every turn, so the client-side UI stays perfectly in sync, but the model itself answers as if shared state were still whatever it was on turn 1 (frozen), even when the client sends a materially different `state` on turn 2+.

This makes `RunAgentInput.state` effectively write-once per thread from the model's point of view, which breaks any multi-turn shared-state flow (human-in-the-loop approval, a collaboratively-edited list, etc.) the moment the conversation goes past its first message.

---

### Steps to Reproduce

Confirmed against the unmodified adapter — no app code involved — using a bare `ClaudeAgentAdapter` behind FastAPI (`system_prompt="You are a helpful assistant."`, no tools, `setting_sources=[]`), on both the PyPI release and current `main`.

1. Start a thread. Turn 1: POST `/agent` with `state: {"notes": []}` and a plain greeting message.
2. Turn 2, **same `thread_id`**: POST `/agent` again with `state: {"notes": ["hello world"]}` and a message asking the model to report exactly what's in `notes`.
3. Compare the `STATE_SNAPSHOT` event's `snapshot.notes` against what the model's text response claims is in `notes`.

---

### Expected Behavior

The model's answer on turn 2 should reflect the `state` sent on turn 2 (`["hello world"]`), the same way `STATE_SNAPSHOT` already does.

---

### Environment

```text
ag-ui-claude-sdk: 0.1.5 (PyPI)
also reproduced on main @ ab9ae4594147ac5ff4d8bf72264dee37c3e0cdf8

claude-agent-sdk: 0.2.128 (latest)
anthropic: 0.120.0
Python: 3.12.12
```

---

### Logs

Turn 1 (`state: {"notes": []}`):

```text
STATE_SNAPSHOT snapshot={"notes": []}
```

Turn 2, same thread (`state: {"notes": ["hello world"]}`, user asks "What is currently in the notes array in your shared state? Quote it exactly."):

```text
STATE_SNAPSHOT snapshot={"notes": ["hello world"]} <-- correct, echoed to client as always

TEXT_MESSAGE_CONTENT delta="The notes array in my shared state currently contains:\n\n```\n[]\n```\n\nIt's empty."
```

Identical result on both the PyPI release and current `main` — same two turns, same wrong answer.

---

### Root Cause

`ag_ui_claude_sdk/adapter.py`, in `run()`:

```python
if entry is None:
options = self.build_options(input_data, thread_id=thread_id) # <-- only place this is ever called
worker = SessionWorker(thread_id, options)
await worker.start()
...
else:
entry["active_runs"] = entry.get("active_runs", 0) + 1
...
worker = entry["worker"] # <-- existing options/system_prompt reused as-is
```

`build_options()` is what calls `build_state_context_addendum()` (`utils.py`) and appends the "## Current Shared State" block to `system_prompt` — but it only runs when a thread's `SessionWorker` is first created. Every later turn reuses that same worker, so its `system_prompt` — and the state snapshot baked into it — is frozen at whatever it was on turn 1.

Both branches converge right before the actual query:

```python
prompt, _ = process_messages(input_data)
message_stream = worker.query(prompt, session_id=thread_id)
```

`prompt` is the only thing that's genuinely fresh every turn. And there's no way to route around the frozen `system_prompt` from here either: `claude_agent_sdk.ClaudeSDKClient.query(prompt, session_id)` only accepts a prompt string — there is no supported way to refresh `system_prompt` (or any other session option) for an already-connected client. This isn't a bug in `claude_agent_sdk` itself — a long-lived CLI session is a reasonable design for conversational continuity — but it does mean `system_prompt` is structurally the wrong place to put anything that needs to change turn-to-turn.

---

### Suggested Fix / Discussion

`integrations/aws-strands` has the exact same architecture (one persistent agent instance cached per `thread_id`, `system_prompt` fixed at creation), and already solves this: `StrandsAgentConfig.state_context_builder` (`config.py:49`, `Callable[[RunAgentInput, str], str]`) is applied to the outgoing user-message text on **every** `run()` call regardless of whether the cached agent is being created or reused (`agent.py:839-851`, `957-978`), wrapped in a `try/except` that falls back to the original text on failure. Because it only touches the local message text used for the actual model call — never `input_data.messages` or anything echoed back to the AG-UI client — there's no snapshot/history to reconcile afterward either.

Notably, `ag_ui_claude_sdk/utils.py`'s own `process_messages()` docstring already says *"Similar to AWS Strands pattern: validates full message history..."* — this integration modeled part of itself on aws-strands already, just not this part.

Concretely, three small, additive changes to `adapter.py` (no protocol/wire changes):

**1. Add an opt-in config option**, mirroring aws-strands's exact type signature for consistency across the repo:

```diff
def __init__(
self,
name: str,
options: Union["ClaudeAgentOptions", dict, None] = None,
description: str = "",
max_workers: int = 1000,
worker_ttl_seconds: float = 1800,
query_timeout_seconds: Optional[float] = 300,
+ state_context_builder: Optional[Callable[["RunAgentInput", str], str]] = None,
):
...
+ self._state_context_builder = state_context_builder
```

**2. Apply it right where `prompt` is built**, after the new-worker/reused-worker branches have already converged — so it fires on every turn, not just the first:

```diff
prompt, _ = process_messages(input_data)
+if self._state_context_builder:
+ try:
+ prompt = self._state_context_builder(input_data, prompt)
+ except Exception as e:
+ logger.warning(f"state_context_builder failed: {e}", exc_info=True)
message_stream = worker.query(prompt, session_id=thread_id)
```

**3. Leave the existing `build_state_context_addendum`/system-prompt path untouched** — it's still correct for a thread's first turn, and this is purely additive, so there's no breaking change for anyone not using `state_context_builder`. Docs/docstring should note that the system-prompt-based state addendum is first-turn-only, and point at `state_context_builder` for anything that needs to stay fresh across a resumed session.

Happy to PR this if I can get assigned — it's small and localized to `adapter.py`, same shape as the aws-strands precedent.

貢獻指南

開啟貢獻指南

評估

這個 Issue 還沒有評估資料。

把新 issue 寄到你的電子郵件信箱

精選適合新手參與的 GitHub issue 摘要。