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

Đang mở
#2,246 1 bình luận 0 reaction 0 người được giao Xem trên GitHub
Ngôn ngữ chính
Python
Star
15.9k
Fork
1.4k
Merge trung bình
1 ngày 17 giờ
Pull request đã merge (30 ngày)
163

Mô tả

### 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.

Hướng dẫn đóng góp

Mở hướng dẫn đóng góp

Đánh giá

Issue này chưa được đánh giá.

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.