Skip to content

feat(adk): Revamp run_claude_agent_activity to use more streaming#309

Merged
declan-scale merged 15 commits intonextfrom
declan-scale/update-claude-utils
Apr 7, 2026
Merged

feat(adk): Revamp run_claude_agent_activity to use more streaming#309
declan-scale merged 15 commits intonextfrom
declan-scale/update-claude-utils

Conversation

@declan-scale
Copy link
Copy Markdown
Contributor

@declan-scale declan-scale commented Apr 7, 2026

Summary

Revamps run_claude_agent_activity to stream all content block types directly from the Claude SDK message stream, replacing the previous ClaudeMessageHandler indirection.

Streaming changes (activities.py):

  • Inline message handling: processes AssistantMessage content blocks (TextBlock, ThinkingBlock, ToolUseBlock) in iteration order with correct timestamps
  • TextBlock → streamed as text deltas via a persistent streaming context
  • ThinkingBlock → streamed as ReasoningContent (with summary extraction)
  • ToolUseBlock → streamed as ToolRequestContent with subagent span tracing
  • Tool results come from PostToolUse/PostToolUseFailure hooks (not from ToolResultBlock), giving richer content

Hooks refactor (hooks/hooks.py):

  • Extracted inline hook closures from activities.py into TemporalStreamingHooks class
  • auto_allow_hook (PreToolUse): auto-allows all tool permissions
  • post_tool_use_hook (PostToolUse): streams tool results to UI, closes subagent spans
  • post_tool_use_failure_hook (PostToolUseFailure): streams tool errors to UI, closes subagent spans
  • Uses proper Claude SDK types (HookInput, HookContext, SyncHookJSONOutput, etc.)
  • create_streaming_hooks() accepts shared subagent_spans dict so hooks and message-level streaming share state
  • User-provided hooks from claude_options are merged with the activity hooks

Tests (test_claude_agents_hooks.py):

  • 24 tests covering all 3 hook types + the factory function
  • Tests for: init, skip conditions, streaming context creation, subagent span lifecycle, error resilience

Test plan

  • rye run pytest tests/lib/test_claude_agents_hooks.py — 24 pass
  • rye run pytest tests/lib/test_claude_agents_activities.py -k "not TestRunClaudeAgentActivity" — 13 pass
  • Manual test: deploy an agent and verify tool calls, reasoning, and text stream correctly in the UI
  • Manual test: verify tool failures show "Error: ..." in the UI

🤖 Generated with Claude Code

Greptile Summary

This PR replaces the ClaudeMessageHandler indirection in run_claude_agent_activity with direct inline streaming of all AssistantMessage content blocks (TextBlock → text deltas, ThinkingBlock → ReasoningContent, ToolUseBlock → ToolRequestContent), and extracts the hook closures into a typed TemporalStreamingHooks class backed by proper Claude SDK types (HookInput, HookContext, HookJSONOutput). The subagent span lifecycle is now correctly split — spans are opened in stream_tool_request and closed in PostToolUse/PostToolUseFailure hooks or in the exception handler — resolving the span-leak concerns raised in previous review rounds.

Confidence Score: 5/5

Safe to merge; the refactor is well-structured with solid test coverage and previous P1 span-leak concerns are fully resolved

No P0 or P1 issues remain. The only finding is a minor P2 style inconsistency (duplicate ToolRequestContent construction in stream_tool_request vs. the shared-variable pattern already used in hooks). 24 new hook tests and updated activity tests provide good regression coverage.

No files require special attention

Important Files Changed

Filename Overview
src/agentex/lib/core/temporal/plugins/claude_agents/activities.py Core activity refactored to process all AssistantMessage content blocks inline with proper text-stream open/close lifecycle; subagent span cleanup correctly added to exception path
src/agentex/lib/core/temporal/plugins/claude_agents/hooks/hooks.py Hooks class refactored to use typed SDK HookInput/HookContext/HookJSONOutput; PreToolUse auto-allows, PostToolUse/PostToolUseFailure close subagent spans and stream results to UI
tests/lib/test_claude_agents_hooks.py New 24-test file covering all three hook types and the factory function with good error-resilience and shared subagent_spans lifecycle coverage
tests/lib/test_claude_agents_activities.py Updated to remove ClaudeMessageHandler mocks; new TestActivityMessageHandling class tests TextBlock, ToolUseBlock, and result message parsing

Sequence Diagram

sequenceDiagram
    participant W as Temporal Workflow
    participant A as run_claude_agent_activity
    participant SDK as ClaudeSDKClient
    participant H as TemporalStreamingHooks
    participant UI as AgentEx UI

    W->>A: invoke(prompt, options)
    A->>SDK: query(prompt)
    loop receive_response()
        SDK-->>A: AssistantMessage
        alt TextBlock
            A->>UI: stream_update(TextDelta)
        else ThinkingBlock
            A->>UI: stream_update(ReasoningContent)
        else ToolUseBlock
            A->>UI: stream_update(ToolRequestContent)
            A->>A: open subagent span (if Task tool)
        end
        SDK-->>H: PostToolUse hook
        H->>H: pop + close subagent span (if Task)
        H->>UI: stream_update(ToolResponseContent)
        SDK-->>H: PostToolUseFailure hook
        H->>H: pop + close subagent span (if Task)
        H->>UI: stream_update(ToolResponseContent Error)
        SDK-->>A: ResultMessage
        A->>A: record usage/cost/session_id
    end
    A->>UI: close_text_stream()
    A-->>W: {messages, session_id, usage, cost_usd}
Loading

Comments Outside Diff (1)

  1. tests/lib/test_claude_agents_activities.py, line 231-446 (link)

    P1 Stale @patch decorators break TestRunClaudeAgentActivity

    All three test methods patch agentex.lib.core.temporal.plugins.claude_agents.activities.ClaudeMessageHandler, but ClaudeMessageHandler is no longer imported in activities.py after this PR. By default, unittest.mock.patch raises AttributeError when the target attribute doesn't exist on the module. The PR test plan explicitly works around this with -k "not TestRunClaudeAgentActivity", confirming these tests are currently broken and provide no coverage for the new streaming logic.

    Prompt To Fix With AI
    This is a comment left during a code review.
    Path: tests/lib/test_claude_agents_activities.py
    Line: 231-446
    
    Comment:
    **Stale `@patch` decorators break `TestRunClaudeAgentActivity`**
    
    All three test methods patch `agentex.lib.core.temporal.plugins.claude_agents.activities.ClaudeMessageHandler`, but `ClaudeMessageHandler` is no longer imported in `activities.py` after this PR. By default, `unittest.mock.patch` raises `AttributeError` when the target attribute doesn't exist on the module. The PR test plan explicitly works around this with `-k "not TestRunClaudeAgentActivity"`, confirming these tests are currently broken and provide no coverage for the new streaming logic.
    
    How can I resolve this? If you propose a fix, please make it concise.
Prompt To Fix All With AI
This is a comment left during a code review.
Path: src/agentex/lib/core/temporal/plugins/claude_agents/activities.py
Line: 293-313

Comment:
**Duplicate `ToolRequestContent` construction**

Two identical `ToolRequestContent` objects are built for the same call — once for `initial_content` and once inside `StreamTaskMessageFull.content`. The hooks refactor already avoids this with a shared `response_content` variable; consistent use of that pattern here would be cleaner.

```suggestion
        req_content = ToolRequestContent(
            author="agent",
            name=block.name,
            arguments=block.input,
            tool_call_id=block.id,
        )
        try:
            async with adk.streaming.streaming_task_message_context(
                task_id=task_id,
                initial_content=req_content,
            ) as ctx:
                await ctx.stream_update(
                    StreamTaskMessageFull(
                        parent_task_message=ctx.task_message,
                        content=req_content,
                        type="full",
                    )
                )
```

How can I resolve this? If you propose a fix, please make it concise.

Reviews (6): Last reviewed commit: "Add task_id gaurd when None" | Re-trigger Greptile

@declan-scale declan-scale changed the base branch from main to next April 7, 2026 13:30
Comment thread src/agentex/lib/core/temporal/plugins/claude_agents/activities.py
@declan-scale declan-scale force-pushed the declan-scale/update-claude-utils branch from 0205e2e to 4c265b2 Compare April 7, 2026 14:12
@declan-scale declan-scale force-pushed the declan-scale/update-claude-utils branch from 4c265b2 to 91f1f6e Compare April 7, 2026 14:12
Comment thread src/agentex/lib/core/temporal/plugins/claude_agents/activities.py
@stainless-app stainless-app Bot force-pushed the next branch 2 times, most recently from 11ed880 to 464cd2d Compare April 7, 2026 18:48
@danielmillerp danielmillerp self-requested a review April 7, 2026 19:36
@declan-scale declan-scale merged commit 25069d3 into next Apr 7, 2026
9 checks passed
@declan-scale declan-scale deleted the declan-scale/update-claude-utils branch April 7, 2026 19:57
@stainless-app stainless-app Bot mentioned this pull request Apr 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants