Skip to content

Commit b8c6195

Browse files
alb-rlalb-rl
committed
Fix merge conflicts
1 parent 0ddad80 commit b8c6195

File tree

2 files changed

+236
-187
lines changed

2 files changed

+236
-187
lines changed

‎src/runloop_api_client/resources/devboxes/executions.py‎

Lines changed: 188 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from __future__ import annotations
44

5-
from typing import Optional
5+
from typing import Optional, cast
66

77
import httpx
88

@@ -16,8 +16,8 @@
1616
async_to_raw_response_wrapper,
1717
async_to_streamed_response_wrapper,
1818
)
19-
from ..._constants import DEFAULT_TIMEOUT
20-
from ..._streaming import Stream, AsyncStream
19+
from ..._constants import DEFAULT_TIMEOUT, RAW_RESPONSE_HEADER
20+
from ..._streaming import Stream, AsyncStream, ReconnectingStream, AsyncReconnectingStream
2121
from ..._exceptions import APIStatusError, APITimeoutError
2222
from ...lib.polling import PollingConfig, poll_until
2323
from ..._base_client import make_request_options
@@ -361,20 +361,53 @@ def stream_stderr_updates(
361361
raise ValueError(f"Expected a non-empty value for `devbox_id` but received {devbox_id!r}")
362362
if not execution_id:
363363
raise ValueError(f"Expected a non-empty value for `execution_id` but received {execution_id!r}")
364-
return self._get(
365-
f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_stderr_updates",
366-
options=make_request_options(
367-
extra_headers=extra_headers,
368-
extra_query=extra_query,
369-
extra_body=extra_body,
370-
timeout=timeout,
371-
query=maybe_transform(
372-
{"offset": offset}, execution_stream_stderr_updates_params.ExecutionStreamStderrUpdatesParams
364+
if extra_headers and extra_headers.get(RAW_RESPONSE_HEADER):
365+
return self._get(
366+
f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_stderr_updates",
367+
options=make_request_options(
368+
extra_headers=extra_headers,
369+
extra_query=extra_query,
370+
extra_body=extra_body,
371+
timeout=timeout,
372+
query=maybe_transform(
373+
{"offset": offset}, execution_stream_stderr_updates_params.ExecutionStreamStderrUpdatesParams
374+
),
373375
),
374-
),
375-
cast_to=ExecutionUpdateChunk,
376-
stream=True,
377-
stream_cls=Stream[ExecutionUpdateChunk],
376+
cast_to=DevboxAsyncExecutionDetailView,
377+
stream=True,
378+
stream_cls=Stream[ExecutionUpdateChunk],
379+
)
380+
381+
def create_stream(last_offset: str | None) -> Stream[ExecutionUpdateChunk]:
382+
new_offset = last_offset if last_offset is not None else (None if isinstance(offset, NotGiven) else offset)
383+
return self._get(
384+
f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_stderr_updates",
385+
options=make_request_options(
386+
extra_headers=extra_headers,
387+
extra_query=extra_query,
388+
extra_body=extra_body,
389+
timeout=timeout,
390+
query=maybe_transform(
391+
{"offset": new_offset},
392+
execution_stream_stderr_updates_params.ExecutionStreamStderrUpdatesParams,
393+
),
394+
),
395+
cast_to=DevboxAsyncExecutionDetailView,
396+
stream=True,
397+
stream_cls=Stream[ExecutionUpdateChunk],
398+
)
399+
400+
initial_stream = create_stream(None)
401+
402+
def get_offset(item: ExecutionUpdateChunk) -> str | None:
403+
value = getattr(item, "offset", None)
404+
if value is None:
405+
return None
406+
return str(value)
407+
408+
return cast(
409+
Stream[ExecutionUpdateChunk],
410+
ReconnectingStream(current_stream=initial_stream, stream_creator=create_stream, get_offset=get_offset),
378411
)
379412

380413
def stream_stdout_updates(
@@ -408,20 +441,53 @@ def stream_stdout_updates(
408441
raise ValueError(f"Expected a non-empty value for `devbox_id` but received {devbox_id!r}")
409442
if not execution_id:
410443
raise ValueError(f"Expected a non-empty value for `execution_id` but received {execution_id!r}")
411-
return self._get(
412-
f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_stdout_updates",
413-
options=make_request_options(
414-
extra_headers=extra_headers,
415-
extra_query=extra_query,
416-
extra_body=extra_body,
417-
timeout=timeout,
418-
query=maybe_transform(
419-
{"offset": offset}, execution_stream_stdout_updates_params.ExecutionStreamStdoutUpdatesParams
444+
if extra_headers and extra_headers.get(RAW_RESPONSE_HEADER):
445+
return self._get(
446+
f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_stdout_updates",
447+
options=make_request_options(
448+
extra_headers=extra_headers,
449+
extra_query=extra_query,
450+
extra_body=extra_body,
451+
timeout=timeout,
452+
query=maybe_transform(
453+
{"offset": offset}, execution_stream_stdout_updates_params.ExecutionStreamStdoutUpdatesParams
454+
),
420455
),
421-
),
422-
cast_to=ExecutionUpdateChunk,
423-
stream=True,
424-
stream_cls=Stream[ExecutionUpdateChunk],
456+
cast_to=DevboxAsyncExecutionDetailView,
457+
stream=True,
458+
stream_cls=Stream[ExecutionUpdateChunk],
459+
)
460+
461+
def create_stream(last_offset: str | None) -> Stream[ExecutionUpdateChunk]:
462+
new_offset = last_offset if last_offset is not None else (None if isinstance(offset, NotGiven) else offset)
463+
return self._get(
464+
f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_stdout_updates",
465+
options=make_request_options(
466+
extra_headers=extra_headers,
467+
extra_query=extra_query,
468+
extra_body=extra_body,
469+
timeout=timeout,
470+
query=maybe_transform(
471+
{"offset": new_offset},
472+
execution_stream_stdout_updates_params.ExecutionStreamStdoutUpdatesParams,
473+
),
474+
),
475+
cast_to=DevboxAsyncExecutionDetailView,
476+
stream=True,
477+
stream_cls=Stream[ExecutionUpdateChunk],
478+
)
479+
480+
initial_stream = create_stream(None)
481+
482+
def get_offset(item: ExecutionUpdateChunk) -> str | None:
483+
value = getattr(item, "offset", None)
484+
if value is None:
485+
return None
486+
return str(value)
487+
488+
return cast(
489+
Stream[ExecutionUpdateChunk],
490+
ReconnectingStream(current_stream=initial_stream, stream_creator=create_stream, get_offset=get_offset),
425491
)
426492

427493

@@ -739,20 +805,53 @@ async def stream_stderr_updates(
739805
raise ValueError(f"Expected a non-empty value for `devbox_id` but received {devbox_id!r}")
740806
if not execution_id:
741807
raise ValueError(f"Expected a non-empty value for `execution_id` but received {execution_id!r}")
742-
return await self._get(
743-
f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_stderr_updates",
744-
options=make_request_options(
745-
extra_headers=extra_headers,
746-
extra_query=extra_query,
747-
extra_body=extra_body,
748-
timeout=timeout,
749-
query=await async_maybe_transform(
750-
{"offset": offset}, execution_stream_stderr_updates_params.ExecutionStreamStderrUpdatesParams
808+
if extra_headers and extra_headers.get(RAW_RESPONSE_HEADER):
809+
return await self._get(
810+
f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_stderr_updates",
811+
options=make_request_options(
812+
extra_headers=extra_headers,
813+
extra_query=extra_query,
814+
extra_body=extra_body,
815+
timeout=timeout,
816+
query=await async_maybe_transform(
817+
{"offset": offset}, execution_stream_stderr_updates_params.ExecutionStreamStderrUpdatesParams
818+
),
751819
),
752-
),
753-
cast_to=ExecutionUpdateChunk,
754-
stream=True,
755-
stream_cls=AsyncStream[ExecutionUpdateChunk],
820+
cast_to=DevboxAsyncExecutionDetailView,
821+
stream=True,
822+
stream_cls=AsyncStream[ExecutionUpdateChunk],
823+
)
824+
825+
async def create_stream(last_offset: str | None) -> AsyncStream[ExecutionUpdateChunk]:
826+
new_offset = last_offset if last_offset is not None else (None if isinstance(offset, NotGiven) else offset)
827+
return await self._get(
828+
f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_stderr_updates",
829+
options=make_request_options(
830+
extra_headers=extra_headers,
831+
extra_query=extra_query,
832+
extra_body=extra_body,
833+
timeout=timeout,
834+
query=await async_maybe_transform(
835+
{"offset": new_offset},
836+
execution_stream_stderr_updates_params.ExecutionStreamStderrUpdatesParams,
837+
),
838+
),
839+
cast_to=DevboxAsyncExecutionDetailView,
840+
stream=True,
841+
stream_cls=AsyncStream[ExecutionUpdateChunk],
842+
)
843+
844+
initial_stream = await create_stream(None)
845+
846+
def get_offset(item: ExecutionUpdateChunk) -> str | None:
847+
value = getattr(item, "offset", None)
848+
if value is None:
849+
return None
850+
return str(value)
851+
852+
return cast(
853+
AsyncStream[ExecutionUpdateChunk],
854+
AsyncReconnectingStream(current_stream=initial_stream, stream_creator=create_stream, get_offset=get_offset),
756855
)
757856

758857
async def stream_stdout_updates(
@@ -786,20 +885,54 @@ async def stream_stdout_updates(
786885
raise ValueError(f"Expected a non-empty value for `devbox_id` but received {devbox_id!r}")
787886
if not execution_id:
788887
raise ValueError(f"Expected a non-empty value for `execution_id` but received {execution_id!r}")
789-
return await self._get(
790-
f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_stdout_updates",
791-
options=make_request_options(
792-
extra_headers=extra_headers,
793-
extra_query=extra_query,
794-
extra_body=extra_body,
795-
timeout=timeout,
796-
query=await async_maybe_transform(
797-
{"offset": offset}, execution_stream_stdout_updates_params.ExecutionStreamStdoutUpdatesParams
888+
# If caller requested a raw or streaming response wrapper, return the underlying stream as-is
889+
if extra_headers and extra_headers.get(RAW_RESPONSE_HEADER):
890+
return await self._get(
891+
f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_stdout_updates",
892+
options=make_request_options(
893+
extra_headers=extra_headers,
894+
extra_query=extra_query,
895+
extra_body=extra_body,
896+
timeout=timeout,
897+
query=await async_maybe_transform(
898+
{"offset": offset}, execution_stream_stdout_updates_params.ExecutionStreamStdoutUpdatesParams
899+
),
798900
),
799-
),
800-
cast_to=ExecutionUpdateChunk,
801-
stream=True,
802-
stream_cls=AsyncStream[ExecutionUpdateChunk],
901+
cast_to=DevboxAsyncExecutionDetailView,
902+
stream=True,
903+
stream_cls=AsyncStream[ExecutionUpdateChunk],
904+
)
905+
906+
async def create_stream(last_offset: str | None) -> AsyncStream[ExecutionUpdateChunk]:
907+
new_offset = last_offset if last_offset is not None else (None if isinstance(offset, NotGiven) else offset)
908+
return await self._get(
909+
f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_stdout_updates",
910+
options=make_request_options(
911+
extra_headers=extra_headers,
912+
extra_query=extra_query,
913+
extra_body=extra_body,
914+
timeout=timeout,
915+
query=await async_maybe_transform(
916+
{"offset": new_offset},
917+
execution_stream_stdout_updates_params.ExecutionStreamStdoutUpdatesParams,
918+
),
919+
),
920+
cast_to=DevboxAsyncExecutionDetailView,
921+
stream=True,
922+
stream_cls=AsyncStream[ExecutionUpdateChunk],
923+
)
924+
925+
initial_stream = await create_stream(None)
926+
927+
def get_offset(item: ExecutionUpdateChunk) -> str | None:
928+
value = getattr(item, "offset", None)
929+
if value is None:
930+
return None
931+
return str(value)
932+
933+
return cast(
934+
AsyncStream[ExecutionUpdateChunk],
935+
AsyncReconnectingStream(current_stream=initial_stream, stream_creator=create_stream, get_offset=get_offset),
803936
)
804937

805938

0 commit comments

Comments
 (0)