@@ -123,96 +123,97 @@ async def _handle_run_action_with_partial_results(
123123
124124 # From here to the final result, this connection is the origin of the
125125 # run: an ER that asks a question mid-run is asking the client that
126- # started it and nobody else (ADR-0082 rule 1). Set on the streamed
127- # paths, which are the ones that still hold their caller's connection;
128- # the dispatch below binds it to the run id it mints, which is what the
129- # ER names when it asks.
130- with elicitation_bridge .originating_client (writer ):
131- stream = await partial_results_service .run_action_with_partial_results (
132- action_name = action_name ,
133- project_path = project_path ,
134- params = params .get ("params" , {}),
135- partial_result_token = token ,
136- run_trigger = trigger ,
137- dev_env = dev_env ,
138- ws_context = ws_context ,
139- result_formats = result_formats ,
140- progress_token = progress_token ,
141- selected_interpreters = selected_interpreters ,
142- )
143-
144- # Opt-in (collect-style callers like MCP): accumulate the `json` format of
145- # each partial per project so it can be type-safely merged into the response.
146- merge_results_enabled = options .get ("mergeResults" , False )
147- json_by_project : dict [str , list [dict ]] = {}
148-
149- async def _forward_partials () -> int :
150- count = 0
151- async for value in stream :
152- count += 1
153- if merge_results_enabled and isinstance (value , dict ):
154- project_str = value .get ("project" , "" )
155- result_by_format = value .get ("resultByFormat" ) or {}
156- json_by_project .setdefault (project_str , []).append (
157- result_by_format .get ("json" )
158- )
159- logger .trace (
160- f"run+partialResults: sending partial #{ count } for token={ token } , keys={ list (value .keys ()) if isinstance (value , dict ) else type (value )} "
161- )
162- _notify_client (
163- writer ,
164- "actions/partialResult" ,
165- {"token" : token , "value" : value },
126+ # started it and nobody else (ADR-0082 rule 1). Constructed on the
127+ # streamed paths, which are the ones that still hold their caller's
128+ # connection; the dispatch below binds it to the run id it mints,
129+ # which is what the ER names when it asks.
130+ origin = elicitation_bridge .RunDispatchOrigin (connection = writer )
131+ stream = await partial_results_service .run_action_with_partial_results (
132+ action_name = action_name ,
133+ project_path = project_path ,
134+ params = params .get ("params" , {}),
135+ partial_result_token = token ,
136+ run_trigger = trigger ,
137+ dev_env = dev_env ,
138+ ws_context = ws_context ,
139+ result_formats = result_formats ,
140+ progress_token = progress_token ,
141+ selected_interpreters = selected_interpreters ,
142+ origin = origin ,
143+ )
144+
145+ # Opt-in (collect-style callers like MCP): accumulate the `json` format of
146+ # each partial per project so it can be type-safely merged into the response.
147+ merge_results_enabled = options .get ("mergeResults" , False )
148+ json_by_project : dict [str , list [dict ]] = {}
149+
150+ async def _forward_partials () -> int :
151+ count = 0
152+ async for value in stream :
153+ count += 1
154+ if merge_results_enabled and isinstance (value , dict ):
155+ project_str = value .get ("project" , "" )
156+ result_by_format = value .get ("resultByFormat" ) or {}
157+ json_by_project .setdefault (project_str , []).append (
158+ result_by_format .get ("json" )
166159 )
167- await writer .drain ()
168- return count
160+ logger .trace (
161+ f"run+partialResults: sending partial #{ count } for token={ token } , keys={ list (value .keys ()) if isinstance (value , dict ) else type (value )} "
162+ )
163+ _notify_client (
164+ writer ,
165+ "actions/partialResult" ,
166+ {"token" : token , "value" : value },
167+ )
168+ await writer .drain ()
169+ return count
169170
170- async def _forward_progress () -> None :
171- if stream .progress_stream is None or progress_token is None :
172- return
173- async for value in stream .progress_stream :
174- logger .trace (
175- f"run+partialResults: sending progress type={ value .get ('type' )} for token={ progress_token } "
176- )
177- _notify_client (
178- writer ,
179- "actions/progress" ,
180- {"token" : progress_token , "value" : value },
181- )
182- await writer .drain ()
171+ async def _forward_progress () -> None :
172+ if stream .progress_stream is None or progress_token is None :
173+ return
174+ async for value in stream .progress_stream :
175+ logger .trace (
176+ f"run+partialResults: sending progress type={ value .get ('type' )} for token={ progress_token } "
177+ )
178+ _notify_client (
179+ writer ,
180+ "actions/progress" ,
181+ {"token" : progress_token , "value" : value },
182+ )
183+ await writer .drain ()
183184
184- partial_count = 0
185- async with asyncio .TaskGroup () as forward_tg :
186- partials_task = forward_tg .create_task (_forward_partials ())
187- forward_tg .create_task (_forward_progress ())
188- partial_count = partials_task .result ()
189-
190- final = await stream .final_result ()
191-
192- if merge_results_enabled and json_by_project :
193- return_code = final .get ("returnCode" , 0 ) if isinstance (final , dict ) else 0
194- results : dict [str , dict ] = {}
195- for project_str , payloads in json_by_project .items ():
196- merged_json = await merge_partial_results_for_action (
197- project_path = pathlib .Path (project_str ),
198- action_name = action_name ,
199- json_payloads = payloads ,
200- ws_context = ws_context ,
201- )
202- if merged_json is not None :
203- results [project_str ] = {
204- action_source : {
205- "resultByFormat" : {"json" : merged_json },
206- "returnCode" : return_code ,
207- }
185+ partial_count = 0
186+ async with asyncio .TaskGroup () as forward_tg :
187+ partials_task = forward_tg .create_task (_forward_partials ())
188+ forward_tg .create_task (_forward_progress ())
189+ partial_count = partials_task .result ()
190+
191+ final = await stream .final_result ()
192+
193+ if merge_results_enabled and json_by_project :
194+ return_code = final .get ("returnCode" , 0 ) if isinstance (final , dict ) else 0
195+ results : dict [str , dict ] = {}
196+ for project_str , payloads in json_by_project .items ():
197+ merged_json = await merge_partial_results_for_action (
198+ project_path = pathlib .Path (project_str ),
199+ action_name = action_name ,
200+ json_payloads = payloads ,
201+ ws_context = ws_context ,
202+ )
203+ if merged_json is not None :
204+ results [project_str ] = {
205+ action_source : {
206+ "resultByFormat" : {"json" : merged_json },
207+ "returnCode" : return_code ,
208208 }
209- if results :
210- final = {** final , "results" : results }
209+ }
210+ if results :
211+ final = {** final , "results" : results }
211212
212- logger .trace (
213- f"run+partialResults: done, sent { partial_count } partials, final keys={ list (final .keys ()) if isinstance (final , dict ) else type (final )} "
214- )
215- return final
213+ logger .trace (
214+ f"run+partialResults: done, sent { partial_count } partials, final keys={ list (final .keys ()) if isinstance (final , dict ) else type (final )} "
215+ )
216+ return final
216217
217218
218219async def _handle_run_action_with_partial_results_task (
@@ -301,13 +302,11 @@ async def _handle_run_batch_with_partial_results(
301302
302303 params = params or {}
303304 # The connection that asked for the batch is the origin of every run in it,
304- # including the ones the per-project tasks below dispatch: a task copies the
305- # context it is created in, so setting this before they exist is what makes
306- # it reach them (ADR-0082 rule 1).
307- with (
308- telemetry .attach_incoming_traceparent (params ),
309- elicitation_bridge .originating_client (writer ),
310- ):
305+ # including the ones the per-project tasks below dispatch: `origin` is
306+ # captured in `_stream_action`'s closure, so a task created from it carries
307+ # the same descriptor to every dispatch it makes (ADR-0082 rule 1).
308+ origin = elicitation_bridge .RunDispatchOrigin (connection = writer )
309+ with telemetry .attach_incoming_traceparent (params ):
311310 parsed = _parse_run_batch_params (params )
312311 token = params ["partialResultToken" ]
313312
@@ -438,6 +437,7 @@ async def _on_partial(
438437 merge_results = parsed .merge_results ,
439438 on_partial = _on_partial ,
440439 selected_interpreters = selected_interpreters ,
440+ origin = origin ,
441441 )
442442 if parsed .merge_results :
443443 merged_results .setdefault (str (project_path ), {})[action_source ] = {
@@ -461,6 +461,7 @@ async def _on_partial(
461461 ws_context = ws_context ,
462462 initialize_all_handlers = True ,
463463 result_formats = parsed .result_formats ,
464+ origin = origin ,
464465 ) as ctx :
465466 async for value in ctx :
466467 partial_count += 1
@@ -727,6 +728,10 @@ async def _forward_progress() -> None:
727728 result_formats = parsed .result_formats ,
728729 initialize_all_handlers = True ,
729730 progress_token = progress_token ,
731+ # Progress is forwarded to this connection for the whole run, so
732+ # it is just as much the run's origin as on the partial-results
733+ # paths: an ER that elicits mid-run has a client to ask.
734+ origin = elicitation_bridge .RunDispatchOrigin (connection = writer ),
730735 )
731736 return {
732737 "resultByFormat" : result .result_by_format ,
@@ -906,6 +911,9 @@ async def _forward_to_client() -> None:
906911 result_formats = parsed .result_formats ,
907912 payload_overrides_by_project = parsed .params_by_project or None ,
908913 progress_token_by_project = progress_token_by_project ,
914+ # Aggregated progress goes to this connection until the batch
915+ # ends, so it is the origin of every run in the batch.
916+ origin = elicitation_bridge .RunDispatchOrigin (connection = writer ),
909917 )
910918 finally :
911919 # Cancel get_progress tasks first, then drain slot lists through aggregator,
0 commit comments