agents icon indicating copy to clipboard operation
agents copied to clipboard

Make client-initiated direct state change opt-in

Open dhl opened this issue 9 months ago • 5 comments

Currently (as at e2d7b5c1cf08927aa0549445c15e30bb14d2dc9f), Agent and by extension AIChatAgent have the convenient but insecure default of allowing WebSocket clients to send 𝚌𝚏_𝚊𝚐𝚎𝚗𝚝_𝚜𝚝𝚊𝚝𝚎 messages to mutate internal agent state.

https://github.com/cloudflare/agents/blob/e2d7b5c1cf08927aa0549445c15e30bb14d2dc9f/packages/agents/src/index.ts#L341-L344

This feature makes it possible for client to call agent.setState to update state from client side with ease:

import { useState } from "react";
import { useAgent } from "agents/react";

function StateInterface() {
  const [state, setState] = useState({ counter: 0 });

  const agent = useAgent({
    agent: "thinking-agent",
    onStateUpdate: (newState) => setState(newState),
  });

  const increment = () => {
    agent.setState({ counter: state.counter + 1 });
  };

  return (
    <div>
      <div>Count: {state.counter}</div>
      <button onClick={increment}>Increment</button>
    </div>
  );
}

While this is great for demos and MVPs, developers unaware of the consequences of this convenience can leave their agents vulnerable to unexpected state changes or other bypasses.

To protect agents and developers, this client-initiated direct state change should become opt-in. The Synchronizing State section of the Agents API documentation should also warn developers of the potential danger of this convenience.

dhl avatar May 06 '25 09:05 dhl

It's been added specifically to synchronise state with connected clients, so I'm not sure we'd add an option to make this opt-in. If you need private state, it's proibably best to use the sql layer directly, and/or use the kv store on this.ctx.storage. did you accidentally leak state using this?

threepointone avatar May 06 '25 10:05 threepointone

Hi @threepointone, thanks for looking into this.

The main issue here is not in exposing state, but allowing uncontrolled client-side write access to the agent state.

With the current implementation, anything running client-side has full and direct write access to the two rows cf_state_row_id and cf_state_was_changed in the table cf_agents_state, as well as replacing the in memory state:

https://github.com/cloudflare/agents/blob/e2d7b5c1cf08927aa0549445c15e30bb14d2dc9f/packages/agents/src/index.ts#L423-L432

This means that if a developer were to use the state feature, they will not have a chance to validate state updates coming from the client side, and the code running Worker/Agent side cannot rely on the integrity of the state.

Even if a developer were to avoid this feature altogether and write directly to the embedded D1 database of the agent, this feature still leaves the cf_state_was_changed and cf_state_was_changed rows fully writable, effectively functioning as a scratchpad.

With the current implementation, it seems like the only way to opt out of this feature is to override the onMessage method directly in the constructor:

// Constructor of an Agent subclass
constructor(ctx: AgentContext, env: Env) {
    super(ctx, env);

    const inheritedOnMessage = this.onMessage;
    this.onMessage = async (connection: Connection, message: WSMessage) => {
      try {
        if (typeof message === "string") {
          const parsedMsg = JSON.parse(message);

          // Ignore `cf_agent_state` messages to prevent state replacement
          if ("type" in parsedMsg && parsedMsg.type === "cf_agent_state") {
            console.warn("Client-initiated state replacement is not allowed. Ignoring message.");
            return;
          }
        }

        // If message is not a string or not a 'cf_agent_state' message,
        // proceed with the inherited onMessage handler.
        return inheritedOnMessage(connection, message);
      } catch (e) {
        // If JSON.parse fails or any other error occurs,
        // falls back to process with the inherited handler.
        console.error("Error processing message in constructor override:", e);
        return inheritedOnMessage(connection, message);
      }
    };
  }

Given how easily this feature could be abused, making it an opt-in feature would be great. Otherwise, providing an easier way to opt-out and having the consequence of the state sync documented would really benefit new developers working with this SDK for the first time.

dhl avatar May 06 '25 16:05 dhl

You have a point. Let me think about it some more and I'll implement something.

threepointone avatar May 06 '25 16:05 threepointone

It also seems inherently racy (read-modify-write from client with stale data), and doesn't validate the input which can cause the agent to throw.

Frankly, this feature will never work right and should be fully removed, in favor of explicit RPC that can do it right. The only thing it's used for currently is having the local UX seem slightly faster (at the cost of hiding errors?). IMHO that sort of optimistic behavior should be part of the UI layer, like displaying counter + (delta_in_flight ?? 0) or as a button that animates to indicate "sending".

https://github.com/cloudflare/agents/issues/174#issuecomment-2806998927

tv42 avatar Jul 05 '25 00:07 tv42

just let us have private state...

https://github.com/cloudflare/agents/issues/312

jlia0 avatar Sep 15 '25 11:09 jlia0