fix(agent): stop duplicate streaming with server-held keys

apps/agent/src/model.py

@@ -6,6 +6,7 @@
 from langchain.agents.middleware import AgentMiddleware
 from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
 from langchain_core.messages import AIMessageChunk
+from langchain_core.runnables import RunnableBinding
 from pydantic import Field
 from langchain_anthropic import ChatAnthropic
 from langchain_core.language_models.chat_models import BaseChatModel
@@ -146,31 +147,48 @@ def _target(self, kwargs):
         tool_options = options.pop("_credential_tool_options", {})
         if tools is not None:
             model = model.bind_tools(tools, **tool_options)
+        # Call the provider's private methods so only this wrapper's run reports
+        # events. A public nested call inherits the graph's callbacks and streams
+        # every token and tool-call delta to AG-UI a second time.
+        if isinstance(model, RunnableBinding):
+            options = {**model.kwargs, **options}
+            model = model.bound
         return model, options
 
     def _generate(self, messages, stop=None, run_manager=None, **kwargs):
         with provider_error_boundary():
             model, options = self._target(kwargs)
-            message = model.invoke(messages, stop=stop, **options)
-            return ChatResult(generations=[ChatGeneration(message=message)])
+            return model._generate(messages, stop=stop, **options)
 
     async def _agenerate(self, messages, stop=None, run_manager=None, **kwargs):
         with provider_error_boundary():
             model, options = self._target(kwargs)
-            message = await model.ainvoke(messages, stop=stop, **options)
-            return ChatResult(generations=[ChatGeneration(message=message)])
+            return await model._agenerate(messages, stop=stop, **options)
 
     def _stream(self, messages, stop=None, run_manager=None, **kwargs):
         with provider_error_boundary():
             model, options = self._target(kwargs)
-            for chunk in model.stream(messages, stop=stop, **options):
-                yield ChatGenerationChunk(message=as_chunk(chunk))
+            if not streams(model):
+                yield as_generation_chunk(model._generate(messages, stop=stop, **options))
+                return
+            yield from model._stream(messages, stop=stop, **options)
 
     async def _astream(self, messages, stop=None, run_manager=None, **kwargs):
         with provider_error_boundary():
             model, options = self._target(kwargs)
-            async for chunk in model.astream(messages, stop=stop, **options):
-                yield ChatGenerationChunk(message=as_chunk(chunk))
+            if not streams(model):
+                yield as_generation_chunk(await model._agenerate(messages, stop=stop, **options))
+                return
+            async for chunk in model._astream(messages, stop=stop, **options):
+                yield chunk
+
+
+def streams(model):
+    return type(model)._stream is not BaseChatModel._stream or type(model)._astream is not BaseChatModel._astream
+
+
+def as_generation_chunk(result):
+    return ChatGenerationChunk(message=as_chunk(result.generations[0].message))
 
 
 def as_chunk(message):

apps/agent/tests/test_model.py

@@ -84,3 +84,32 @@ def test_missing_provider_key_has_actionable_error(monkeypatch, model, key, valu
         monkeypatch.setenv(key, value)
     with pytest.raises(ValueError, match=key):
         build_model()
+
+
+def test_credential_scoped_model_streams_each_chunk_once():
+    import asyncio
+    from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
+    from langchain_core.messages import AIMessage
+    from langchain_core.runnables import RunnableLambda
+    from src.model import CredentialScopedModel
+
+    class ToolFake(GenericFakeChatModel):
+        def bind_tools(self, tools, **kwargs):
+            return self.bind(tools=tools)
+
+    inner = ToolFake(messages=iter([AIMessage(content="a b c")]))
+    # Bound tools wrap the provider in a RunnableBinding, as in the real agent.
+    wrapped = CredentialScopedModel(fallback=inner).bind_tools([{"name": "t"}])
+
+    # Like a LangGraph node: the parent run's callbacks are inherited by calls inside it.
+    async def node(_):
+        return await wrapped.ainvoke("hi")
+
+    async def stream_events():
+        return [
+            e async for e in RunnableLambda(node).astream_events("x", version="v2")
+            if e["event"] == "on_chat_model_stream"
+        ]
+
+    tokens = [e["data"]["chunk"].content for e in asyncio.run(stream_events())]
+    assert "".join(tokens) == "a b c"