feat(streaming): stream tool call argument deltas in TemporalStreamingModel

[vkalmathscale]

Summary

TemporalStreamingModel already streams text deltas and reasoning summary deltas to Redis via StreamingTaskMessageContext, but ResponseFunctionCallArgumentsDeltaEvent was being silently buffered into function_calls_in_progress[...]['arguments'] with no per-delta publish. Consumers only saw the completed tool call surface later (after the activity returned, via downstream hooks if any).

For write-heavy tools — write_file, apply_patch, anything that puts a 2–20KB string into a single argument — the model spends multiple seconds generating the argument body, and the UI sees nothing until the entire activity finishes. The result is a frozen UI followed by an abrupt jump when the activity returns.

This PR threads tool-call argument deltas through the same streaming machinery used for text and reasoning, riding on the CoalescingBuffer + StreamingMode infrastructure added in #333. The buffer's merge helpers already key on tool_call_id for ToolRequestDelta, so coalescing, mode dispatch, and opt-out are inherited from existing infra.

Design

TemporalStreamingModel now opens a streaming_task_message_context per function call (keyed off the call's output_index), with initial_content=ToolRequestContent(...) and the model's configured streaming_mode. Three event handlers participate:

Event	Behavior added
ResponseOutputItemAddedEvent (type=function_call)	Open the per-call streaming context and stash it on function_calls_in_progress[output_index]['context'].
ResponseFunctionCallArgumentsDeltaEvent	Emit StreamTaskMessageDelta(delta=ToolRequestDelta(arguments_delta=..., tool_call_id=..., name=...)) into the per-call context. The coalescing buffer merges consecutive deltas with the same tool_call_id.
ResponseOutputItemDoneEvent (type=function_call)	Parse the accumulated args (with a graceful empty-dict fallback on JSONDecodeError), emit a final StreamTaskMessageFull(content=ToolRequestContent(...)), and close the context.

End-of-loop cleanup defensively closes any function-call contexts that didn't see a Done event (truncated stream or mid-stream exception).

ModelResponse output is unchanged: output_items still receives the same complete ResponseFunctionToolCall. Activity determinism is unaffected — streaming is a side effect.

What this does NOT change

• Text and reasoning streaming paths are untouched.
• StreamingMode is already the on/off knob. No new flag. streaming_mode="off" suppresses tool-arg deltas the same way it suppresses text deltas. "per_token" publishes immediately; "coalesced" (default) batches at 50ms / 128 chars.
• TemporalStreamingHooks.on_tool_start is unchanged. It still fires after the activity returns and still emits a ToolRequestContent Full message via the stream_lifecycle_content activity. See Caveats.

Caveats

1. Overlap with TemporalStreamingHooks.on_tool_start. Users who pass TemporalStreamingHooks to Runner.run will now see two persisted task_messages per tool call: one created by the model (delta stream + final Full) and one created by the hook (Full only). Both land on the same Redis topic task:{task_id} with different parent_task_message.ids, so a default UI will render two cards for the same logical tool call.

   This needs a follow-up to decide which path owns the canonical ToolRequest emission. Options for review discussion:

   • Silence on_tool_start's Full emit when the model is also emitting (auto-detect via a workflow-instance flag, mirroring how _task_id / _trace_id are threaded today).
   • Remove on_tool_start's Full emit entirely in a follow-up major bump (the model becomes the single source of truth for ToolRequest events).

   Until that follow-up, users who want streamed tool args without duplicate emits should subclass TemporalStreamingHooks and override on_tool_start to a no-op.

2. Coalescing windows still apply. With the default 50ms / 128-char window, tool args render in ~50ms-granularity chunks rather than per-token. This is the same tradeoff already made for text streaming in perf(streaming): coalesce per-token publishes to Redis (50ms / 128-char window) #333, and the right default for write-heavy tools (UX value is "watch the artifact appear", not "see each token").

3. Malformed argument JSON. If the model produces invalid JSON for the args (truncated stream, hallucinated structure), the path logs a WARNING and emits arguments={} in the final ToolRequestContent. The raw delta stream is preserved on the consumer side regardless — only the structured final view falls back.

Test plan

• Two new unit tests in test_streaming_model.py::TestStreamingModelFunctionCallArgsStreaming:
  • Happy path: well-formed args produce one streaming context opened with ToolRequestContent, one StreamTaskMessageDelta(ToolRequestDelta) per ArgumentsDelta event preserving the delta text, and one final StreamTaskMessageFull(ToolRequestContent) with parsed args.
  • Malformed args: emits arguments={} in the final Full and logs a WARNING.
• Full test_streaming_model.py suite passes (42/42).
• ruff check clean on both modified files.
• Manual smoke: deploy to a dev environment with an agent that calls a write-heavy tool, confirm UI sees tool args streaming in coalesced batches.
• Manual smoke: streaming_mode="off" suppresses tool-arg deltas (only the final persisted message exists on close).

cc reviewers familiar with #333's CoalescingBuffer design.

Greptile Summary

This PR extends TemporalStreamingModel to publish tool-call argument deltas to Redis in real time, eliminating the UI freeze while the model generates write-heavy tool arguments like write_file. It opens a per-function-call streaming_task_message_context on ResponseOutputItemAddedEvent, emits a ToolRequestDelta per ResponseFunctionCallArgumentsDeltaEvent, and closes with a final ToolRequestContent Full message on ResponseOutputItemDoneEvent, riding the same CoalescingBuffer + StreamingMode infrastructure introduced in #333.

• Model file: Three updated event handlers (Added/Delta/Done for function calls), a per-call context slot in function_calls_in_progress, and a defensive post-loop cleanup. ModelResponse output and activity determinism are unchanged.
• Test file: Two new tests covering the happy-path delta/Full sequence and the malformed-JSON fallback; both use a real TaskMessage to pass Pydantic validation.
• Known caveat (acknowledged in PR): Users who also pass TemporalStreamingHooks will see duplicate ToolRequest task messages per tool call until a follow-up resolves ownership."

Confidence Score: 5/5

Safe to merge; the new streaming path is a pure side effect that does not alter ModelResponse output or activity determinism.

The change threads tool-call argument deltas through the existing streaming infrastructure without touching the text or reasoning paths. The double-close concern from the prior review thread is correctly addressed via finally: call_data['context'] = None. The orphan-cleanup block only runs on normal stream exhaustion, consistent with the pre-existing pattern for streaming_context and reasoning_context. The duplicate-emit overlap with TemporalStreamingHooks.on_tool_start is explicitly documented and deferred.

No files require special attention beyond the acknowledged TemporalStreamingHooks overlap documented in the PR description.

Important Files Changed

Filename	Overview
src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py	Adds per-function-call streaming contexts so tool argument deltas are published to Redis in real time; follows existing text/reasoning streaming pattern correctly.
src/agentex/lib/core/temporal/plugins/openai_agents/tests/test_streaming_model.py	Adds two focused tests covering the happy-path delta/full sequence and the malformed-JSON fallback; mirrors production event ordering and uses a real TaskMessage for Pydantic validation.

Sequence Diagram

sequenceDiagram
    participant OAI as OpenAI Stream
    participant TSM as TemporalStreamingModel
    participant CTX as per-call StreamingContext
    participant Redis as Redis

    OAI->>TSM: ResponseOutputItemAddedEvent (function_call)
    TSM->>CTX: "open context with ToolRequestContent(arguments={})"
    CTX->>Redis: persist initial ToolRequestContent

    loop For each argument chunk
        OAI->>TSM: ResponseFunctionCallArgumentsDeltaEvent
        TSM->>CTX: stream_update StreamTaskMessageDelta(ToolRequestDelta)
        CTX->>Redis: publish delta coalesced at 50ms/128ch
    end

    OAI->>TSM: ResponseFunctionCallArgumentsDoneEvent
    Note over TSM: overwrite call_data arguments with authoritative string

    OAI->>TSM: ResponseOutputItemDoneEvent (function_call)
    TSM->>CTX: stream_update StreamTaskMessageFull(ToolRequestContent parsed args)
    CTX->>Redis: persist final ToolRequestContent
    TSM->>CTX: close()

    OAI->>TSM: ResponseCompletedEvent
    Note over TSM: output_items = response.output for ModelResponse

    Note over TSM: post-loop cleanup closes contexts missing Done event

Loading

Prompt To Fix All With AI

Fix the following 2 code review issues. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 2
src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py:978-987
The orphan-cleanup loop is placed after `async for event in stream:` but before the `except` block, so it only runs when the iterator exhausts normally (e.g. truncated stream where `ResponseOutputItemDoneEvent` was never sent). If the loop body throws any exception — network error, `Temporal.CancelledError`, a bad delta value — the outer `except Exception` at line 1141 re-raises immediately, skipping this block and leaking every open function-call context. The comment's "error mid-call" language implies it covers exceptions too, which it does not.

```suggestion
                # Defensive: close any function call contexts that didn't see a
                # ResponseOutputItemDoneEvent (truncated stream).
                # NOTE: this block only runs when the stream exhausts normally.
                # Exceptions thrown mid-loop bypass this.
                for call_data in function_calls_in_progress.values():
                    call_context = call_data.get('context')
                    if call_context is not None:
                        try:
                            await call_context.close()
                        except Exception as e:
                            logger.warning(f"Failed to close orphaned tool request context: {e}")
                        call_data['context'] = None
```

### Issue 2 of 2
src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py:741-747
**`arguments` accumulation is redundant for the final parse** — `call_data['arguments']` is seeded from `item.arguments`, incremented per delta, and then fully overwritten by `ResponseFunctionCallArgumentsDoneEvent` with the SDK's authoritative string. The `json.loads()` in `ResponseOutputItemDoneEvent` therefore only ever sees the `Done` value. A clarifying comment would prevent future readers from assuming the concatenation is load-bearing.

_{Reviews (3): Last reviewed commit: "Merge remote-tracking branch 'origin/nex..." | Re-trigger Greptile}