diff --git a/.gitignore b/.gitignore index 80c1ceff7..a9d436958 100644 --- a/.gitignore +++ b/.gitignore @@ -181,3 +181,7 @@ node_modules/ # PyDoc generated docs .generated_docs/ + +# Local-only scratch and credentials (never committed). Anchored, so it cannot also +# swallow a nested directory that happens to share the name. +/.local/ diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a0046d0c..4f01fa289 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,14 @@ ## Unreleased +### Bug Fixes + +- **Only flush traces that no other task is still building**: Concurrent tasks share one `GalileoLogger`, and therefore one internal trace list, whenever they run on the same thread against the same project and log stream. A flush treated that shared list as if it belonged to the caller, which produced three silent failures: a trace appended while a flush was in flight was discarded; two overlapping flushes each sent the whole list, so a trace went out once per in-flight flush; and a flush sent traces a sibling task had started and not concluded, so they reached Galileo with no output, no duration and no status code. + + A flush now hands the batch off rather than clearing the list afterwards — so a trace appended mid-flush is picked up by the next flush instead of being lost — and it sends only the caller's own trace plus traces that no live context is still building. Traces held back leave with their owner's flush, or once that owner finishes. The batch is restored if the send fails or is cancelled, and finished traces still leave together in one request, so this costs no extra requests. + + Most likely to be seen with `@log` on an async function plus `asyncio.gather`, since a decorated coroutine necessarily holds its trace open across every await inside it. The LangChain callbacks are affected by the dropped and duplicated sends, but not by the premature one: they start, log and conclude a trace with no await in between, so no sibling can observe a half-built trace on that path. ([#634](https://github.com/rungalileo/galileo-python/pull/634)) + ### Features - **New `generated_output` field**: Add `generated_output` field to `DatasetRecord` for storing model-generated outputs separately from ground truth. This allows you to track both the expected output (ground truth) and the actual model output in the same dataset record. In the UI, this field is displayed as "Generated Output". diff --git a/README.md b/README.md index f9ce03b93..191c755cc 100644 --- a/README.md +++ b/README.md @@ -233,6 +233,10 @@ call_openai() galileo_context.flush() ``` +**What a flush uploads.** A flush uploads the trace the calling code is building, plus any trace that no live task or thread is still building. If a concurrent task has started a trace and not finished it, that trace stays queued and is uploaded by that task's own flush instead, so it is never sent without its output and spans. + +**A flush is not a receipt for your trace.** `galileo_context.flush()` returns nothing and swallows upload errors, so a normal return is not a confirmation that a given trace was sent. `logger.flush()` returns the batch that was uploaded, which is still not the fate of your own trace: it returns an empty list in three different situations — nothing was pending, a concurrent flush had already carried your trace, or an upload error was swallowed. An empty list therefore does not mean your trace was not sent. + Using the Langchain callback handler: ```python diff --git a/src/galileo/decorator.py b/src/galileo/decorator.py index a0a91c814..f374e7601 100644 --- a/src/galileo/decorator.py +++ b/src/galileo/decorator.py @@ -633,6 +633,10 @@ def _safe_prepare_call( bool True if preparation succeeded, False if it failed """ + # Captured before the attempt, not read in the handler below: an enclosing decorated call's + # span means the trace being built is that call's, not this one's, and the answer must not + # depend on how far `_prepare_call` got before raising. + outermost = not _get_or_init_list(_span_stack_context) try: self._prepare_call(span_type, span_params, dataset_record) return True @@ -641,8 +645,30 @@ def _safe_prepare_call( _logger.error("Galileo logging initialization failed: %s", e, exc_info=True) else: _logger.warning("Galileo logging initialization failed, continuing without logging: %s", e) + # Preparation may already have started a trace, and returning False means the caller + # skips _finalize_call - so nothing else would ever report this context as done with + # it, leaving it held back from every flush. This context is not building it. + # + # Only when this call was the outermost one, though: an enclosing call is still building + # the trace, and releasing it on that call's behalf would expose it to a sibling's flush + # mid-build. That call reports its own hand-off when it returns. + if outermost: + self._release_trace_being_built() return False + def _release_trace_being_built(self) -> None: + """Report that this context is no longer building the trace it had started, if any.""" + try: + logger = self.get_logger_instance() + trace_id = logger._current_trace_id() + if trace_id is not None: + logger._mark_trace_finished(trace_id) + except Exception as e: + # Debug, not a warning: the most likely raiser is `get_logger_instance()` above, and the + # caller that could not build a logger has already reported that failure. There is no + # trace to release in that case, so a second message per decorated call is pure noise. + _logger.debug("Could not release the trace being built: %s", e) + def _prepare_call( self, span_type: SPAN_TYPE | None, span_params: dict[str, Any], dataset_record: DatasetRecord | None ) -> None: @@ -698,6 +724,19 @@ def _prepare_call( span = client_instance.add_workflow_span(input=input_, name=name, created_at=created_at) _get_or_init_list(_span_stack_context).append(span) + # This context owns the trace until the outermost decorated call returns, which is what + # _finalize_call reports. Concurrent tasks share one logger and one trace list, so a + # sibling's flush must not carry the trace away while spans are still being added to it. + # Re-marked on every entry because a reused trace was released by the previous call. + # + # Deliberately the last statement: the caller skips _finalize_call when _prepare_call + # raises, so marking any earlier would strand the trace - held back from every flush with + # nothing left to release it. Nothing above yields, so the trace cannot be observed by + # another task before this runs. + trace_being_built = _trace_context.get() + if trace_being_built is not None: + client_instance._mark_trace_unfinished(trace_being_built.id) + def _get_input_from_func_args( self, *, is_method: bool = False, func_args: tuple = (), func_kwargs: dict | None = None ) -> Any: @@ -748,9 +787,23 @@ def _finalize_call( ------- The original result, possibly wrapped if it's a generator """ - if inspect.isgenerator(result): - return self._wrap_sync_generator_result(span_type, span_params, result) - if inspect.isasyncgen(result): + if inspect.isgenerator(result) or inspect.isasyncgen(result): + # The wrappers below report the hand-off from `_handle_call_result` when the generator + # is exhausted, but only `_async_log` keeps the wrapper it is handed: `_sync_log` + # discards it and returns the raw generator, and both generator kinds route through + # `_sync_log` (`asyncio.iscoroutinefunction` is False for an async generator function). + # So on that path nothing would ever release the trace. Reported here instead, where + # both paths pass. The stack still holds this call's own span - only + # `_handle_call_result` pops it - so the outermost call is the one that leaves it alone. + # Only a workflow, agent or untyped call pushed a span in `_prepare_call`: a + # non-concludable span type pushed nothing, so for it a stack of one holds an enclosing + # call's span, and reporting the hand-off would release a trace that call is still + # building. + pushed_own_span = not span_type or is_concludable_span_type(span_type) + if len(_get_or_init_list(_span_stack_context)) <= (1 if pushed_own_span else 0): + self._release_trace_being_built() + if inspect.isgenerator(result): + return self._wrap_sync_generator_result(span_type, span_params, result) return self._wrap_async_generator_result(span_type, span_params, result) return self._handle_call_result(span_type, span_params, result) @@ -905,6 +958,14 @@ def _handle_call_result(self, span_type: SPAN_TYPE | None, span_params: dict[str except Exception as e: _logger.error(f"Failed to create trace for span '{span_name}' (type: {span_type}): {e}", exc_info=True) + # The outermost decorated call has returned, so this context has stopped adding to the + # trace and any flush may now send it. The trace itself is deliberately left open for a + # later decorated call in this context to reuse, so concluding is not the signal here. + if not _get_or_init_list(_span_stack_context): + trace = _trace_context.get() + if trace is not None: + logger._mark_trace_finished(trace.id) + return result def _wrap_sync_generator_result( @@ -1092,7 +1153,14 @@ def flush( on_error: Callable[[Exception], None] | None = None, ) -> None: """ - Upload all captured traces under a project and log stream context to Galileo. + Upload traces captured under a project and log stream context to Galileo. + + Uploads the trace the calling code is building, plus every trace that no live task or thread is + still building. A trace another context is part-way through building stays queued and leaves with + that context's own flush instead, so it is not sent without its output, spans and duration. + + Nothing is returned and upload errors are swallowed (see ``on_error``), so a normal return is not + a confirmation that a given trace was uploaded. If no project or log stream is provided, then the currently initialized context is used. @@ -1147,9 +1215,11 @@ def _on_flush_error(exc: Exception) -> None: def flush_all(self) -> None: """ - Upload all captured traces under all contexts to Galileo. + Upload traces captured under all contexts to Galileo. - This method flushes all traces regardless of project or log stream. + This method flushes every cached logger regardless of project or log stream. Each one uploads the + same set as `flush()` does: the trace the calling code is building, plus every trace that no live + task or thread is still building. """ GalileoLoggerSingleton().flush_all() _span_stack_context.set([]) diff --git a/src/galileo/logger/logger.py b/src/galileo/logger/logger.py index 8b38d0142..9986339b3 100644 --- a/src/galileo/logger/logger.py +++ b/src/galileo/logger/logger.py @@ -6,8 +6,10 @@ import json import logging import os +import threading import time import uuid +import weakref from collections.abc import Callable from datetime import datetime from typing import TYPE_CHECKING, Any, Union @@ -203,6 +205,13 @@ class GalileoLogger(TracesLogger): _traces_client: Union["Traces", "IngestTraces"] | None = None _task_handler: ThreadPoolTaskHandler _trace_completion_submitted: bool + # Traces a context is still building, each mapped to a weak reference to the asyncio task or + # thread building it. `self.traces` is shared across concurrent tasks + # (GalileoLoggerSingleton._get_key keys on thread+project+log_stream), so a flush must not ship + # a sibling's unfinished trace - but it must ship one whose owner has gone away, or nothing + # ever will: the claim in `_flush_batch` comes from the caller's own parent chain, and an + # abandoned trace's chain died with its owner. + _traces_being_built: dict[uuid.UUID, weakref.ref[Any] | None] def __init__( self, @@ -261,6 +270,7 @@ def __init__( mode = _get_mode_or_default(mode) self.mode: LoggerModeType = mode self._task_counter = 0 + self._traces_being_built = {} self._ingestion_hook = ingestion_hook if self._ingestion_hook and self.mode == "distributed": @@ -560,6 +570,7 @@ def add_trace( ) trace._parent = None self.traces.append(trace) + self._mark_trace_unfinished(trace.id) self._set_current_parent(trace) return trace @@ -927,6 +938,97 @@ def _update_step_streaming(self, step: StepWithChildSpans, is_complete: bool = F def previous_parent(self) -> StepWithChildSpans | None: return self._parent_stack[-2] if len(self._parent_stack) > 1 else None + def _current_trace_id(self) -> uuid.UUID | None: + """Id of the trace this context is currently building, if any. + + The parent chain lives in a ContextVar, so this only answers for the calling + context — never for a concurrent task sharing the same logger. + """ + parent_stack = self._parent_stack + return parent_stack[0].id if parent_stack else None + + @staticmethod + def _current_owner_ref() -> weakref.ref[Any] | None: + """Weak reference to the task or thread building a trace, for later liveness checks.""" + try: + owner: Any = asyncio.current_task() + except RuntimeError: + # No running event loop: the caller is plain synchronous code. + owner = None + try: + return weakref.ref(owner if owner is not None else threading.current_thread()) + except TypeError: + # Not weak-referenceable. Treated as alive, i.e. exactly the behaviour of a marker + # with no liveness information at all. + return None + + def _mark_trace_unfinished(self, trace_id: uuid.UUID) -> None: + """Record that a context is building this trace, so other contexts' flushes skip it. + + The guarantee for anything that marks a trace: a flush from another context skips it while + the marking task is alive, or while the marking thread is alive and is not itself the thread + asking for the flush. A marker that never releases it therefore costs a trace that ships + unconcluded once its owner finishes - not one that is never sent at all. + """ + self._traces_being_built[trace_id] = self._current_owner_ref() + + def _mark_trace_finished(self, trace_id: uuid.UUID) -> None: + """Record that nobody is building this trace any more, so any flush may send it. + + `conclude()` implies this, but a caller can be done with a trace without concluding it: + `@log` leaves the trace open so a later decorated call in the same context can reuse it, + and reports the hand-off through here instead. + """ + self._traces_being_built.pop(trace_id, None) + + def _still_being_built(self, trace_id: uuid.UUID, flushing_thread: threading.Thread) -> bool: + """Whether a live context is still adding to this trace. + + A trace whose owner has finished is being built by nobody: no context will conclude it or + add to it, and none can claim it either, since the claim comes from the caller's own parent + chain. Holding it back would mean never sending it. + + Parameters + ---------- + trace_id : uuid.UUID + Trace to test. + flushing_thread : threading.Thread + Thread the flush was requested from, captured before dispatch. A synchronous owner that + is this thread is done with the trace: a thread cannot be flushing and mid-build at once. + """ + if trace_id not in self._traces_being_built: + return False + owner_ref = self._traces_being_built[trace_id] + if owner_ref is None: + return True + owner = owner_ref() + if owner is None: + # The owner was collected, so it cannot still be running. + return False + if isinstance(owner, asyncio.Task): + return not owner.done() + if isinstance(owner, threading.Thread): + # `is_alive()` alone would hold a trace for the life of a pool worker, which outlives the + # job that started the trace - and the next job on that worker cannot claim it, because + # pools that copy the context per job (anyio, so FastAPI's sync endpoints) hand it an + # empty parent chain. Comparing identity, not name: anyio names every worker alike. + return owner is not flushing_thread and owner.is_alive() + return True + + def reset_parent_tracking(self) -> None: + """Drop this context's parent chain, abandoning the trace it was building. + + Also reports that nobody is building that trace any more, so a later flush may send it + rather than holding it back. This reaches the caller's own trace only - the id comes from + the caller's parent chain, which is empty in any other context - so it is not a way to + release a trace abandoned by a different task or thread. + """ + trace_id = self._current_trace_id() + if trace_id is not None: + # Abandoned, not still running: a later flush must not hold it back forever. + self._mark_trace_finished(trace_id) + super().reset_parent_tracking() + @nop_sync @warn_catch_exception(exceptions=(Exception,)) def has_active_trace(self) -> bool: @@ -1978,6 +2080,9 @@ def _conclude( # Navigate up to parent via _parent pointer finished_step = current_parent + if finished_step._parent is None: + # Root of the chain: this context's trace is finished and may now be flushed. + self._mark_trace_finished(finished_step.id) self._set_current_parent(current_parent._parent) return (finished_step, self.current_parent()) @@ -2045,7 +2150,11 @@ def conclude( @nop_sync def flush(self, on_error: Callable[[Exception], None] | None = None) -> list[LoggedTrace]: """ - Upload all traces to Galileo. + Upload traces to Galileo. + + Sends this caller's own trace, plus every trace no live context is still building. A trace + another context is part-way through building leaves with that context's flush instead, so + that it is not sent without its output, spans and duration. Parameters ---------- @@ -2060,11 +2169,16 @@ def flush(self, on_error: Callable[[Exception], None] | None = None) -> list[Log list[LoggedTrace] The list of uploaded traces. """ + # Resolved before dispatch so the claim reflects the caller, not whatever the + # EventLoopThreadPool thread happens to see in its copy of the context. The thread is + # captured for the same reason: from here `_flush_batch` runs on a pool thread. + claimed_trace_id = self._current_trace_id() + flushing_thread = threading.current_thread() try: try: if self.mode == "distributed": return async_run(self._flush_distributed()) - return async_run(self._flush_batch()) + return async_run(self._flush_batch(claimed_trace_id, flushing_thread)) finally: # Reset parent tracking in the main thread (async_run uses thread pool). # Using finally ensures cleanup even if ingestion fails. @@ -2088,17 +2202,22 @@ def flush(self, on_error: Callable[[Exception], None] | None = None) -> list[Log @async_warn_catch_exception(exceptions=(Exception,)) async def async_flush(self) -> list[LoggedTrace]: """ - Async upload all traces to Galileo. + Async upload traces to Galileo. + + Sends this caller's own trace, plus every trace no live context is still building — see + `flush()`. Returns ------- list[LoggedTrace] The list of uploaded traces. """ + claimed_trace_id = self._current_trace_id() + flushing_thread = threading.current_thread() try: if self.mode == "distributed": return await self._flush_distributed() - return await self._flush_batch() + return await self._flush_batch(claimed_trace_id, flushing_thread) finally: # Reset parent tracking. Using finally ensures cleanup even if ingestion fails. self._set_current_parent(None) @@ -2174,14 +2293,17 @@ def _wait_for_pending_span_ingests(self, timeout_seconds: int) -> None: def _auto_conclude_trace(self) -> None: """Helper to auto-conclude any unconcluded trace/spans before flushing. - Note: We assume at most one active trace at a time. add_trace() enforces this - by raising an error if current_parent() is not None. + Only concludes the chain belonging to the calling context: `self.traces` can hold + concurrent tasks' traces, and concluding one of those on its owner's behalf would + stamp it with an output derived from someone else's spans. """ if not self.traces: return - # Use the last trace in self.traces (should be the only active trace) - trace = self.traces[-1] + parent_stack = self._parent_stack + # Root of *this* context's chain. The fallback only matters in distributed mode, + # where start_trace() resets self.traces to the single trace being built. + trace = parent_stack[0] if parent_stack else self.traces[-1] # Don't auto-conclude stub traces - they're owned by the upstream service # Downstream services that receive distributed tracing headers create stubs @@ -2190,7 +2312,7 @@ def _auto_conclude_trace(self) -> None: return # If there are unconcluded items in the stack, conclude them - if self._parent_stack: + if parent_stack: self._logger.info("Concluding unconcluded spans before flush...") # Get output from last child span if trace has no explicit output output, redacted_output = GalileoLogger._get_last_output(trace) @@ -2228,25 +2350,80 @@ async def _flush_distributed(self) -> list[LoggedTrace]: self._logger.info("All distributed tracing requests are complete.") self.traces = [] + # Only batch mode reads this; clearing it alongside the traces keeps a long-lived + # distributed logger from accumulating entries for traces it has already sent. + self._traces_being_built.clear() self._set_current_parent(None) return [] - async def _flush_batch(self) -> list[LoggedTrace]: - """Flush in batch mode: conclude unconcluded traces and send all traces to backend.""" + async def _flush_batch( + self, claimed_trace_id: uuid.UUID | None, flushing_thread: threading.Thread + ) -> list[LoggedTrace]: + """Flush in batch mode: conclude unconcluded traces and send all traces to backend. + + Parameters + ---------- + claimed_trace_id : Optional[uuid.UUID] + Trace the flushing context was building, from `_current_trace_id()`. It ships + even if unconcluded — the caller asked for it — while traces belonging to other + contexts are held back until they finish. + flushing_thread : threading.Thread + Thread the flush was requested from, captured before dispatch. Both must be resolved + in the caller's own thread and context; neither is readable from here, because the + synchronous `flush()` runs this coroutine on an `EventLoopThreadPool` thread. + """ if not self.traces: self._logger.info("No traces to flush.") return [] + # Snapshot before anything can clear a mark: `_auto_conclude_trace()` releases the caller's + # own trace, and the partition below pops every mark it detaches. A failed send puts the + # batch back, so it has to put back what the batch was marked with too - otherwise a trace + # its owner is still building returns unmarked and the next flush from another context + # carries it away half-built. + marks_before_send = dict(self._traces_being_built) + self._auto_conclude_trace() + # Detach the batch up-front. `self.traces` can be shared by concurrent tasks (see + # GalileoLoggerSingleton._get_key, which keys on thread+project+log_stream), so a + # trace appended while the request below is in flight must land in the *new* list + # and be picked up by the next flush. Clearing after the await instead would + # discard it silently: it was never in the frozen payload, and the next flush would + # find an empty list and no-op while still reporting success. + # + # A trace another context started but has not concluded stays behind rather than + # riding along: `ingest_traces` serialises the payload synchronously, so sending it + # now would put it on the wire without its output, spans or duration, and — since + # the batch is detached — its owner's own flush would find nothing left to send. + # Only while that context is still alive, though: a trace whose owner has finished + # without concluding it can never be claimed by anyone, so holding it back would mean + # never sending it at all. + # This loop must stay free of awaits so no *task* can append between the partition and the + # rebinding. It does not exclude a second OS thread: `_get_key` keys on thread name, so anyio + # workers share one logger and one list, and that window is not closed here. + ready: list[LoggedTrace] = [] + still_running: list[LoggedTrace] = [] + for trace in self.traces: + if trace.id != claimed_trace_id and self._still_being_built(trace.id, flushing_thread): + still_running.append(trace) + else: + ready.append(trace) + logged_traces, self.traces = ready, still_running + for trace in logged_traces: + self._traces_being_built.pop(trace.id, None) + + if not logged_traces: + self._logger.info("No traces ready to flush.") + return [] + if self.local_metrics: self._logger.info("Computing metrics for local scorers...") # TODO: parallelize, possibly with asyncio to_thread/gather - for trace in self.traces: + for trace in logged_traces: populate_local_metrics(trace, self.local_metrics) - logged_traces = self.traces trace_count = len(logged_traces) self._logger.info(f"Flushing {trace_count} {'trace' if trace_count == 1 else 'traces'}...") @@ -2257,6 +2434,32 @@ async def _flush_batch(self) -> list[LoggedTrace]: experiment_id=self.experiment_id, ) + try: + await self._send_ingest_request(traces_ingest_request) + except BaseException: + # The batch was detached before the send, so put it back for the next flush. + # Without this, detaching up-front would turn a failed send into data loss, + # whereas previously the traces stayed in place and were retried. + # BaseException, not Exception: `asyncio.CancelledError` does not derive from + # Exception, so cancelling an in-flight `async_flush()` would otherwise drop the + # detached batch outright. + self.traces = logged_traces + self.traces + for trace in logged_traces: + if trace.id in marks_before_send: + # setdefault, not assignment: a context that took the mark again while the send + # was in flight owns the trace now, and its reference must not be overwritten by + # the stale one. Membership test, not truthiness: `None` is a legitimate value, + # meaning an owner that cannot be weak-referenced and so counts as alive. + self._traces_being_built.setdefault(trace.id, marks_before_send[trace.id]) + raise + + self._logger.info(f"Successfully flushed {trace_count} {'trace' if trace_count == 1 else 'traces'}.") + + self._set_current_parent(None) # Reset parent tracking + return logged_traces + + async def _send_ingest_request(self, traces_ingest_request: TracesIngestRequest) -> None: + """Hand a built ingest request to the ingestion hook, or to the traces client.""" if self._ingestion_hook: if inspect.iscoroutinefunction(self._ingestion_hook): await self._ingestion_hook(traces_ingest_request) @@ -2276,12 +2479,6 @@ async def _flush_batch(self) -> list[LoggedTrace]: else: await self._traces_client.ingest_traces(traces_ingest_request) - self._logger.info(f"Successfully flushed {trace_count} {'trace' if trace_count == 1 else 'traces'}.") - - self.traces = [] - self._set_current_parent(None) # Reset parent tracking - return logged_traces - @nop_sync @warn_catch_exception(exceptions=(Exception,)) def terminate(self) -> None: @@ -2310,10 +2507,14 @@ def terminate(self) -> None: self._auto_conclude_trace() self._wait_for_all_tasks_sync(timeout_seconds=terminate_timeout_seconds) self.traces = [] + self._traces_being_built.clear() self._set_current_parent(None) else: # Batch mode: try flush() but don't fail if async_run has issues during shutdown try: + # Lifecycle end: nothing can still be running, so let the final flush + # ship traces that were never concluded instead of stranding them. + self._traces_being_built.clear() self.flush() except RuntimeError as e: # Event loop might be closed during shutdown, log warning but don't crash diff --git a/src/galileo/openai/__init__.py b/src/galileo/openai/__init__.py index 0f94ea4eb..9f6194262 100644 --- a/src/galileo/openai/__init__.py +++ b/src/galileo/openai/__init__.py @@ -278,6 +278,13 @@ def _wrap( raise exc_info return openai_response except Exception as ex: + if should_complete_trace: + # This call started the trace and owns concluding it, but is unwinding before it could. + # A flush skips traces another context is still building, so without reporting the + # hand-off nothing would send this one until the process exits. + trace_id = galileo_logger._current_trace_id() + if trace_id is not None: + galileo_logger._mark_trace_finished(trace_id) _logger.error(f"Error while processing OpenAI request: {ex}") raise RuntimeError("Failed to process the OpenAI Request") from ex diff --git a/src/galileo/utils/singleton.py b/src/galileo/utils/singleton.py index a8338aaca..9eb2f43fc 100644 --- a/src/galileo/utils/singleton.py +++ b/src/galileo/utils/singleton.py @@ -239,6 +239,10 @@ def flush( and cleared. Otherwise, only the specific logger corresponding to the provided key (project, log_stream) is flushed and removed. + Each logger uploads the trace the calling code is building, plus every trace that no live task or + thread is still building; a trace another context is part-way through building stays queued for + that context's own flush. + Parameters ---------- project (Optional[str], optional) @@ -261,7 +265,12 @@ def flush( self._galileo_loggers[key].flush() def flush_all(self) -> None: - """Flush (upload and clear) all GalileoLogger instances.""" + """Flush (upload and clear) all GalileoLogger instances. + + Each logger uploads the trace the calling code is building, plus every trace that no live task or + thread is still building; a trace another context is part-way through building stays queued for + that context's own flush. + """ with self._lock: # Terminate and clear all logger instances. for logger in self._galileo_loggers.values(): diff --git a/tests/test_decorator.py b/tests/test_decorator.py index 96394664e..64d1d018b 100644 --- a/tests/test_decorator.py +++ b/tests/test_decorator.py @@ -1,3 +1,4 @@ +import asyncio from typing import NoReturn from unittest.mock import Mock, patch from uuid import UUID @@ -1670,3 +1671,596 @@ def llm_call(query: str) -> str: # Then: debug is called, not warning mock_logger.debug.assert_called_once() mock_logger.warning.assert_not_called() + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_concurrent_decorated_coroutines_do_not_flush_each_others_traces( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, reset_context +) -> None: + """A flush must not carry a concurrent task's still-running decorated trace. + + ``@log`` on an async function necessarily holds its trace open across every await inside the + function, and concurrent tasks under one ``galileo_context`` share a single ``GalileoLogger`` + and therefore a single trace list. Flushing from one task while another is mid-await used to + send the sibling's trace, which is serialised before the request is awaited and so goes out + with no output - and, because the batch is detached from the list, the sibling's own flush + then finds nothing left to send. + + This is the documented usage shape rather than an exotic one: plain ``@log`` plus + ``asyncio.gather``, and the blocking ``flush()``. No ``async_flush()`` is involved. + """ + # Given: two decorated coroutines under one context, one held mid-await + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + london_started = asyncio.Event() + london_may_finish = asyncio.Event() + ingest_payloads: list[list[tuple[str, str | None]]] = [] + + def record_payload(request) -> dict: + # Snapshotted synchronously: the real client serialises before its first await, so a + # trace sent mid-flight cannot be repaired by a later conclude. + ingest_payloads.append([(trace.name, trace.output) for trace in request.traces]) + return {} + + mock_traces_client_instance.ingest_traces.side_effect = record_payload + + @log + async def london_forecast() -> str: + london_started.set() + await london_may_finish.wait() + return "rainy in London" + + @log + async def new_york_forecast() -> str: + return "sunny in New York" + + async def london_request() -> None: + await london_forecast() + galileo_context.get_logger_instance(project="project-concurrent", log_stream="stream-concurrent").flush() + + with galileo_context(project="project-concurrent", log_stream="stream-concurrent"): + logger = galileo_context.get_logger_instance(project="project-concurrent", log_stream="stream-concurrent") + london = asyncio.create_task(london_request()) + await asyncio.wait_for(london_started.wait(), timeout=5) + + # When: the other coroutine finishes and flushes while London is still awaiting + await new_york_forecast() + logger.flush() + + # Then: only the finished trace was sent, and London waits for its own flush + assert ingest_payloads == [[("new_york_forecast", "sunny in New York")]] + assert [trace.name for trace in logger.traces] == ["london_forecast"] + + # When: London finishes and flushes from its own task + london_may_finish.set() + await london + + # Then: London was sent exactly once, carrying its own output + assert ingest_payloads == [[("new_york_forecast", "sunny in New York")], [("london_forecast", "rainy in London")]] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_context_exit_flush_sends_traces_from_finished_tasks( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, reset_context +) -> None: + """Traces from tasks that finished without flushing must still leave at context exit. + + ``@log`` leaves its trace open on purpose so a later decorated call in the same context can + reuse it, so "not concluded" cannot mean "still being built" - if it did, tasks that finish + without flushing would have their traces held back from the exit flush and stranded until the + process ends. The decorator therefore reports the hand-off when its outermost call returns. + + Nothing flushes per task here: ``galileo_context.__exit__`` is the only flush. + """ + # Given: three decorated coroutines that run concurrently and never flush themselves + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + ingest_payloads: list[list[str]] = [] + + def record_payload(request) -> dict: + ingest_payloads.append([trace.name for trace in request.traces]) + return {} + + mock_traces_client_instance.ingest_traces.side_effect = record_payload + + @log + async def forecast(city: str) -> str: + await asyncio.sleep(0) # yield, so the three tasks genuinely interleave + return f"{city}: done" + + # When: they all finish inside the context, and only the exit flush runs + with galileo_context(project="project-exit-flush", log_stream="stream-exit-flush"): + logger = galileo_context.get_logger_instance(project="project-exit-flush", log_stream="stream-exit-flush") + await asyncio.gather(*(forecast(city) for city in ("New York", "London", "Tokyo"))) + assert len(logger.traces) == 3 + assert ingest_payloads == [] + + # Then: the exit flush carried all three rather than stranding them + assert ingest_payloads == [["forecast", "forecast", "forecast"]] + assert logger.traces == [] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_reused_trace_is_protected_while_a_second_decorated_call_builds_it( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, reset_context +) -> None: + """A trace reused by a second decorated call is off limits again while that call runs. + + A context that finishes one decorated call releases its trace so any flush can send it, but + ``_prepare_call`` reuses that same still-open trace for the next decorated call in the context. + The release therefore has to be re-taken on entry, or the second call's spans are exposed to a + concurrent task's flush - the original defect, reached through the reuse path. + """ + # Given: a task that completes one decorated call, then starts a second on the same trace + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + second_call_started = asyncio.Event() + second_call_may_finish = asyncio.Event() + ingest_payloads: list[list[str]] = [] + + def record_payload(request) -> dict: + ingest_payloads.append([trace.name for trace in request.traces]) + return {} + + mock_traces_client_instance.ingest_traces.side_effect = record_payload + + @log + async def first_step() -> str: + return "first step done" + + @log + async def second_step() -> str: + second_call_started.set() + await second_call_may_finish.wait() + return "second step done" + + @log + async def new_york_forecast() -> str: + return "sunny in New York" + + async def two_step_request() -> None: + await first_step() + await second_step() + + with galileo_context(project="project-trace-reuse", log_stream="stream-trace-reuse"): + logger = galileo_context.get_logger_instance(project="project-trace-reuse", log_stream="stream-trace-reuse") + request = asyncio.create_task(two_step_request()) + await asyncio.wait_for(second_call_started.wait(), timeout=5) + + # When: another task flushes while the second call is still building the reused trace + await new_york_forecast() + logger.flush() + + # Then: the reused trace stayed behind + assert ingest_payloads == [["new_york_forecast"]] + assert [trace.name for trace in logger.traces] == ["first_step"] + + second_call_may_finish.set() + await request + + # Then: it left at context exit, once, with both steps' spans on it + assert ingest_payloads == [["new_york_forecast"], ["first_step"]] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_trace_is_not_stranded_when_span_setup_fails( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, reset_context +) -> None: + """A trace is never marked as being built unless something will release it. + + The wrappers skip ``_finalize_call`` when ``_prepare_call`` raises, and ``_finalize_call`` is + the only thing that reports the decorator's hand-off. Claiming the trace before the span setup + that might raise would therefore strand it: held back from every flush, with nothing left to + release it, until the process exits. + + The failing call runs in its own task and the flush comes from outside it, because a flush in + the owning context claims its own trace and would mask the leak. The task is also held alive + across the flush: a finished owner is treated as no longer building the trace, so letting the + task complete would mask the leak a second way and leave this test passing with the release + removed. + """ + # Given: a decorated call in another task whose span setup fails after the trace exists + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + ingest_payloads: list[list[str]] = [] + + def record_payload(request) -> dict: + ingest_payloads.append([trace.name for trace in request.traces]) + return {} + + mock_traces_client_instance.ingest_traces.side_effect = record_payload + + @log + async def forecast() -> str: + return "sunny in New York" + + with galileo_context(project="project-span-setup-fails", log_stream="stream-span-setup-fails"): + logger = galileo_context.get_logger_instance( + project="project-span-setup-fails", log_stream="stream-span-setup-fails" + ) + + request_failed = asyncio.Event() + task_may_finish = asyncio.Event() + + async def failing_request() -> None: + with patch.object(type(logger), "add_workflow_span", side_effect=RuntimeError("span setup exploded")): + await forecast() + request_failed.set() + await task_may_finish.wait() + + # When: the span setup inside _prepare_call raises, so _finalize_call is skipped + owner = asyncio.create_task(failing_request()) + await asyncio.wait_for(request_failed.wait(), timeout=5) + assert len(logger.traces) == 1 + + # Then: a flush from outside that task carries the trace even though its owner is still + # alive, so only the release in _safe_prepare_call can have let it go + logger.flush() + payloads_after_foreign_flush = list(ingest_payloads) + + task_may_finish.set() + await owner + + assert payloads_after_foreign_flush == [["forecast"]] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_decorated_generator_releases_its_trace_when_the_call_returns( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, reset_context +) -> None: + """A decorated generator must report its hand-off like any other decorated call. + + ``_finalize_call`` wraps a generator result and returns the wrapper, but ``_sync_log`` discards + that return value and hands back the unwrapped generator, so the wrapper is never iterated and + ``_handle_call_result`` - the only place the decorator reports it has stopped building the trace + - never runs. The trace is then held back from every flush. + + The owning task is still alive when the foreign flush happens, so nothing but this hand-off can + release the trace: an owner-liveness check cannot rescue it, which is what makes this the shape + that pins the release rather than the backstop. + """ + # Given: a decorated generator consumed inside a task that then stays alive + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + ingest_payloads: list[list[str]] = [] + + def record_payload(request) -> dict: + ingest_payloads.append([trace.name for trace in request.traces]) + return {} + + mock_traces_client_instance.ingest_traces.side_effect = record_payload + + @log + def stream_forecast(): + yield "sunny " + yield "in New York" + + generator_consumed = asyncio.Event() + task_may_finish = asyncio.Event() + + with galileo_context(project="project-decorated-generator", log_stream="stream-decorated-generator"): + logger = galileo_context.get_logger_instance( + project="project-decorated-generator", log_stream="stream-decorated-generator" + ) + + async def consuming_task() -> None: + assert list(stream_forecast()) == ["sunny ", "in New York"] + generator_consumed.set() + await task_may_finish.wait() + + consumer = asyncio.create_task(consuming_task()) + await asyncio.wait_for(generator_consumed.wait(), timeout=5) + assert len(logger.traces) == 1 + + # When: a context that does not own the trace flushes while the owning task is still alive + logger.flush() + payloads_after_foreign_flush = list(ingest_payloads) + + task_may_finish.set() + await consumer + + # Then: the flush carried the trace instead of holding it back for an owner that was done + assert payloads_after_foreign_flush == [["stream_forecast"]] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_decorated_async_generator_releases_its_trace_when_the_call_returns( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, reset_context +) -> None: + """An async generator reaches the same unreleased path as a sync one. + + ``asyncio.iscoroutinefunction`` is False for an async generator function, so ``@log`` routes it + through ``_sync_log`` too, and its wrapper is discarded the same way. + """ + # Given: a decorated async generator consumed inside a task that then stays alive + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + ingest_payloads: list[list[str]] = [] + + def record_payload(request) -> dict: + ingest_payloads.append([trace.name for trace in request.traces]) + return {} + + mock_traces_client_instance.ingest_traces.side_effect = record_payload + + @log + async def stream_forecast_async(): + yield "rainy " + yield "in London" + + generator_consumed = asyncio.Event() + task_may_finish = asyncio.Event() + + with galileo_context(project="project-decorated-async-generator", log_stream="stream-decorated-async-generator"): + logger = galileo_context.get_logger_instance( + project="project-decorated-async-generator", log_stream="stream-decorated-async-generator" + ) + + async def consuming_task() -> None: + assert [item async for item in stream_forecast_async()] == ["rainy ", "in London"] + generator_consumed.set() + await task_may_finish.wait() + + consumer = asyncio.create_task(consuming_task()) + await asyncio.wait_for(generator_consumed.wait(), timeout=5) + assert len(logger.traces) == 1 + + # When: a context that does not own the trace flushes while the owning task is still alive + logger.flush() + payloads_after_foreign_flush = list(ingest_payloads) + + task_may_finish.set() + await consumer + + # Then: the flush carried the trace instead of holding it back for an owner that was done + assert payloads_after_foreign_flush == [["stream_forecast_async"]] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_nested_decorated_generator_does_not_release_the_outer_trace( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, reset_context +) -> None: + """A nested decorated generator must not report the hand-off on its caller's behalf. + + ``_prepare_call`` pushes a span only for a workflow, agent or untyped call, so a generator + decorated with a non-concludable span type pushes nothing. Reporting the hand-off whenever the + span stack holds at most one entry therefore reported it for the *enclosing* call's span, and a + sibling's flush carried the outer trace away while its body was still running - without its + output, and detached from the list, so the outer call's own flush had nothing left to send. + + The outer call is parked on an await when the foreign flush happens, so its task is alive and an + owner-liveness check cannot rescue the trace: only the guard on the stack length can. + """ + # Given: an outer decorated coroutine that consumes a nested decorated generator, then parks + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + ingest_payloads: list[list[tuple[str, str | None]]] = [] + + def record_payload(request) -> dict: + ingest_payloads.append([(trace.name, trace.output) for trace in request.traces]) + return {} + + mock_traces_client_instance.ingest_traces.side_effect = record_payload + + @log(span_type="llm") + def stream_tokens(): + yield "sunny " + yield "in New York" + + outer_parked = asyncio.Event() + outer_may_finish = asyncio.Event() + + @log + async def outer_forecast() -> str: + assert list(stream_tokens()) == ["sunny ", "in New York"] + outer_parked.set() + await outer_may_finish.wait() + return "sunny in New York" + + with galileo_context(project="project-nested-generator", log_stream="stream-nested-generator"): + logger = galileo_context.get_logger_instance( + project="project-nested-generator", log_stream="stream-nested-generator" + ) + + async def outer_request() -> None: + await outer_forecast() + logger.flush() + + owner = asyncio.create_task(outer_request()) + await asyncio.wait_for(outer_parked.wait(), timeout=5) + + # When: a context that does not own the trace flushes while the outer call is mid-await + logger.flush() + + # Then: the outer trace stayed behind for the task that is still building it + assert ingest_payloads == [] + assert [trace.name for trace in logger.traces] == ["outer_forecast"] + + outer_may_finish.set() + await owner + + # Then: it was sent exactly once, by its owner, carrying the output the early release destroyed + assert ingest_payloads == [[("outer_forecast", "sunny in New York")]] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_nested_decorated_async_generator_does_not_release_the_outer_trace( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, reset_context +) -> None: + """An async generator nested in a decorated call must not release its caller's trace either. + + Both generator kinds route through ``_sync_log``, and neither pushes a span when decorated with a + non-concludable span type, so the same off-by-one reached them both. + """ + # Given: an outer decorated coroutine that consumes a nested decorated async generator + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + ingest_payloads: list[list[tuple[str, str | None]]] = [] + + def record_payload(request) -> dict: + ingest_payloads.append([(trace.name, trace.output) for trace in request.traces]) + return {} + + mock_traces_client_instance.ingest_traces.side_effect = record_payload + + @log(span_type="tool") + async def stream_tokens_async(): + yield "rainy " + yield "in London" + + outer_parked = asyncio.Event() + outer_may_finish = asyncio.Event() + + @log + async def outer_forecast_async() -> str: + assert [token async for token in stream_tokens_async()] == ["rainy ", "in London"] + outer_parked.set() + await outer_may_finish.wait() + return "rainy in London" + + with galileo_context(project="project-nested-async-generator", log_stream="stream-nested-async-generator"): + logger = galileo_context.get_logger_instance( + project="project-nested-async-generator", log_stream="stream-nested-async-generator" + ) + + async def outer_request() -> None: + await outer_forecast_async() + logger.flush() + + owner = asyncio.create_task(outer_request()) + await asyncio.wait_for(outer_parked.wait(), timeout=5) + + # When: a context that does not own the trace flushes while the outer call is mid-await + logger.flush() + + # Then: the outer trace stayed behind for the task that is still building it + assert ingest_payloads == [] + assert [trace.name for trace in logger.traces] == ["outer_forecast_async"] + + outer_may_finish.set() + await owner + + # Then: it was sent exactly once, by its owner, carrying its output + assert ingest_payloads == [[("outer_forecast_async", "rainy in London")]] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_nested_span_setup_failure_does_not_release_the_outer_trace( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, reset_context +) -> None: + """A nested decorated call whose setup fails must not release its caller's trace. + + ``_safe_prepare_call`` reports the hand-off when ``_prepare_call`` raises, because the wrappers + then skip ``_finalize_call`` and nothing else would report it. But the trace it reports on is the + root of the caller's parent chain, which a nested call shares with every enclosing call - so + reporting it unconditionally handed the outer call's trace to any concurrent flush while the + outer call was still running. + + The failure mode here is synthetic: every logger call ``_prepare_call`` makes swallows its own + exceptions, so in production the escaping raiser is the bare ``traces[-1]`` index. Patching the + span setup is a stand-in that reaches the same release site. + """ + # Given: an outer decorated coroutine whose nested decorated call fails its span setup + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + ingest_payloads: list[list[tuple[str, str | None]]] = [] + + def record_payload(request) -> dict: + ingest_payloads.append([(trace.name, trace.output) for trace in request.traces]) + return {} + + mock_traces_client_instance.ingest_traces.side_effect = record_payload + + @log + def nested_step() -> str: + return "nested step ran" + + outer_parked = asyncio.Event() + outer_may_finish = asyncio.Event() + + @log + async def outer_forecast() -> str: + nested_step() + outer_parked.set() + await outer_may_finish.wait() + return "sunny in New York" + + with galileo_context(project="project-nested-setup-fails", log_stream="stream-nested-setup-fails"): + logger = galileo_context.get_logger_instance( + project="project-nested-setup-fails", log_stream="stream-nested-setup-fails" + ) + + real_add_workflow_span = type(logger).add_workflow_span + span_setups = {"count": 0} + + def only_the_nested_call_fails(self, *args, **kwargs): + span_setups["count"] += 1 + if span_setups["count"] == 2: + raise RuntimeError("nested span setup exploded") + return real_add_workflow_span(self, *args, **kwargs) + + async def outer_request() -> None: + await outer_forecast() + logger.flush() + + with patch.object(type(logger), "add_workflow_span", only_the_nested_call_fails): + owner = asyncio.create_task(outer_request()) + await asyncio.wait_for(outer_parked.wait(), timeout=5) + + # When: a context that does not own the trace flushes while the outer call is mid-await + logger.flush() + + # Then: the outer trace stayed behind for the task that is still building it + assert ingest_payloads == [] + assert [trace.name for trace in logger.traces] == ["outer_forecast"] + + outer_may_finish.set() + await owner + + # Then: it was sent exactly once, by its owner, carrying the output the early release destroyed + assert ingest_payloads == [[("outer_forecast", "sunny in New York")]] diff --git a/tests/test_logger_batch.py b/tests/test_logger_batch.py index bf13a7832..12a5e51bd 100644 --- a/tests/test_logger_batch.py +++ b/tests/test_logger_batch.py @@ -1,6 +1,9 @@ +import asyncio +import contextvars import datetime import json import logging +import threading import uuid from collections import deque from unittest.mock import AsyncMock, Mock, patch @@ -2477,3 +2480,940 @@ def test_ingest_traces_reuses_existing_client( # Then: no additional Traces client was created (reuses the existing one) assert mock_traces_cls.call_count == call_count_before + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_trace_added_during_ingest_is_not_dropped( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """A trace appended while a flush is in flight must still reach the backend. + + Concurrent tasks sharing one ``GalileoLogger`` also share its trace list. Because the + ingest payload is frozen before the network await, a trace appended during that await is + not part of it, so the flush must hand its batch off rather than clear the list + afterwards - a blind clear discards that trace, and the next flush then finds an empty + list, sends nothing, and still reports success. + + Only the network egress (``ingest_traces``) is mocked; the shared list, the payload built + before the await, the batch hand-off, and the empty-list early return are all real. The + mock's only job is to hold the await open deterministically, because in production that + window is ordinary network latency and timing-dependent tests are unreliable. + + ``_flush_batch()`` is awaited directly rather than via the public sync ``flush()``, + which blocks its OS thread and so cannot interleave. + """ + # Given: a logger whose in-flight ingest is held open until a second trace is appended + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + + ingest_in_flight = asyncio.Event() + second_trace_appended = asyncio.Event() + ingested_trace_names: list[str] = [] + + async def hold_the_window_open(request: TracesIngestRequest) -> dict: + ingested_trace_names.extend(trace.name for trace in request.traces) + if not ingest_in_flight.is_set(): + ingest_in_flight.set() + await second_trace_appended.wait() + return {} + + mock_traces_client_instance.ingest_traces = AsyncMock(side_effect=hold_the_window_open) + + logger.start_trace(input="weather in New York?", name="New York") + logger.conclude(output="sunny in New York") + + # When: a second trace is appended while the first flush sits in its ingest await, + # then each trace is flushed + first_flush = asyncio.create_task(logger._flush_batch(None, threading.current_thread())) + await ingest_in_flight.wait() + + logger.start_trace(input="weather in London?", name="London") + logger.conclude(output="rainy in London") + second_trace_appended.set() + + await first_flush + await logger._flush_batch(None, threading.current_thread()) + + # Then: both traces were sent, and neither was silently discarded + assert sorted(set(ingested_trace_names)) == ["London", "New York"] + assert logger.traces == [] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_failed_ingest_retains_traces_for_next_flush( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """A failed send must leave its traces queued rather than dropping them. + + ``_flush_batch`` detaches the batch from ``self.traces`` before sending, so it has to put + the batch back if the send raises. Otherwise avoiding a silent drop under concurrency + would introduce a silent drop on every ingest error. + """ + # Given: a logger whose first ingest attempt fails and whose second succeeds + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + logger.start_trace(input="input", name="test-trace") + logger.conclude(output="output") + + mock_traces_client_instance.ingest_traces = AsyncMock(side_effect=ConnectionError("backend unavailable")) + + # When: the flush fails + with pytest.raises(ConnectionError): + await logger._flush_batch(None, threading.current_thread()) + + # Then: the trace is still queued, and a later successful flush sends it + assert [trace.name for trace in logger.traces] == ["test-trace"] + + mock_traces_client_instance.ingest_traces = AsyncMock(return_value={}) + await logger._flush_batch(None, threading.current_thread()) + + mock_traces_client_instance.ingest_traces.assert_called_once() + payload: TracesIngestRequest = mock_traces_client_instance.ingest_traces.call_args.args[0] + assert [trace.name for trace in payload.traces] == ["test-trace"] + assert logger.traces == [] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_cancelled_flush_retains_traces_for_next_flush( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """A cancelled flush must leave its traces queued rather than dropping them. + + ``_flush_batch`` detaches its batch before sending, so it has to put the batch back whenever + the send does not complete. ``asyncio.CancelledError`` does not derive from ``Exception``, so a + restore guarded by ``except Exception`` would let a cancelled flush lose the batch outright - + worse than clearing after the await, where cancellation simply skipped the clear. + + Cancellation is driven through the public ``async_flush()``, the path a caller reaches via + ``asyncio.wait_for`` or task-group teardown. Neither of its decorators intercepts + ``CancelledError``: ``async_warn_catch_exception(exceptions=(Exception,))`` does not match it, + and ``nop_async`` has no handler at all. The sync ``flush()`` runs on a pool thread and so + cannot be cancelled from here. + """ + # Given: a logger whose ingest parks indefinitely once it is in flight + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + + ingest_in_flight = asyncio.Event() + never_released = asyncio.Event() + + async def park_in_flight(request: TracesIngestRequest) -> dict: + ingest_in_flight.set() + await never_released.wait() + return {} + + mock_traces_client_instance.ingest_traces = AsyncMock(side_effect=park_in_flight) + + logger.start_trace(input="input", name="test-trace") + logger.conclude(output="output") + + # When: the flush is cancelled while sitting in its ingest await + flush = asyncio.create_task(logger.async_flush()) + await asyncio.wait_for(ingest_in_flight.wait(), timeout=5) + flush.cancel() + + with pytest.raises(asyncio.CancelledError): + await flush + + # Then: the trace is still queued, and a later successful flush sends it + assert [trace.name for trace in logger.traces] == ["test-trace"] + + mock_traces_client_instance.ingest_traces = AsyncMock(return_value={}) + await logger._flush_batch(None, threading.current_thread()) + + mock_traces_client_instance.ingest_traces.assert_called_once() + payload: TracesIngestRequest = mock_traces_client_instance.ingest_traces.call_args.args[0] + assert [trace.name for trace in payload.traces] == ["test-trace"] + assert logger.traces == [] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_concurrent_flushes_do_not_duplicate_traces( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """Two flushes in flight at once must not both send the same trace. + + Because ``_flush_batch`` detaches its batch before sending, a flush that starts while another + is still awaiting picks up only the traces appended since. Leaving the list in place until + after the await instead hands the second flush a batch that still holds the first flush's + traces, sending them twice - traffic the backend hides by collapsing repeated trace ids. + + Both flushes are held inside their ingest await simultaneously, so the overlap is real rather + than sequential, and payloads are asserted by multiplicity: a duplicate send shows up as a + repeated name rather than being folded away. + """ + # Given: a logger whose ingests all park until released, recording each payload as it arrives + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + + release_ingest = asyncio.Event() + first_ingest_entered = asyncio.Event() + second_ingest_entered = asyncio.Event() + ingest_payloads: list[list[str]] = [] + + async def park_until_released(request: TracesIngestRequest) -> dict: + ingest_payloads.append([trace.name for trace in request.traces]) + if len(ingest_payloads) == 1: + first_ingest_entered.set() + else: + second_ingest_entered.set() + await release_ingest.wait() + return {} + + mock_traces_client_instance.ingest_traces = AsyncMock(side_effect=park_until_released) + + logger.start_trace(input="weather in New York?", name="New York") + logger.conclude(output="sunny in New York") + + # When: a second flush reaches its ingest while the first is still in flight, then both finish + first_flush = asyncio.create_task(logger._flush_batch(None, threading.current_thread())) + await asyncio.wait_for(first_ingest_entered.wait(), timeout=5) + + logger.start_trace(input="weather in London?", name="London") + logger.conclude(output="rainy in London") + + second_flush = asyncio.create_task(logger._flush_batch(None, threading.current_thread())) + await asyncio.wait_for(second_ingest_entered.wait(), timeout=5) + + release_ingest.set() + await asyncio.gather(first_flush, second_flush) + + # Then: each trace was sent exactly once, carried by exactly one of the two flushes + assert sorted(name for payload in ingest_payloads for name in payload) == ["London", "New York"] + assert [len(payload) for payload in ingest_payloads] == [1, 1] + assert mock_traces_client_instance.ingest_traces.call_count == 2 + assert logger.traces == [] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_flush_does_not_ship_another_tasks_unconcluded_trace( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """A flush must leave behind traces another task is still building. + + Concurrent tasks sharing one ``GalileoLogger`` share its trace list, and a flush has no + reason to send a trace that its owner has not finished: ``ingest_traces`` serialises the + payload synchronously before its network await, so a trace picked up mid-flight goes on + the wire without its output, duration, status code, or any span added after that moment. + A later ``conclude()`` mutates a model that has already been serialised and so cannot + repair it. + + Each task is an ``asyncio.Task`` so it gets its own copy of the context, which is what + makes the two parent chains independent while the trace list stays shared - the real + shape of the defect. Ordering is driven by events rather than sleeps. + """ + # Given: two tasks sharing a logger, one holding a trace open while the other flushes + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + + london_trace_started = asyncio.Event() + london_may_conclude = asyncio.Event() + ingest_payloads: list[list[str]] = [] + + async def record_payload(request: TracesIngestRequest) -> dict: + ingest_payloads.append([trace.name for trace in request.traces]) + return {} + + mock_traces_client_instance.ingest_traces = AsyncMock(side_effect=record_payload) + + async def london_task() -> None: + logger.start_trace(input="weather in London?", name="London") + london_trace_started.set() + await london_may_conclude.wait() + logger.conclude(output="rainy in London") + await logger.async_flush() + + london = asyncio.create_task(london_task()) + await asyncio.wait_for(london_trace_started.wait(), timeout=5) + + # When: New York finishes and flushes while London is still mid-trace + logger.start_trace(input="weather in New York?", name="New York") + logger.conclude(output="sunny in New York") + await logger.async_flush() + + # Then: only New York was sent, and London is still queued for its own flush + assert ingest_payloads == [["New York"]] + assert [trace.name for trace in logger.traces] == ["London"] + + london_may_conclude.set() + await london + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_trace_held_open_during_a_sibling_flush_is_sent_once_with_its_output( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """The customer-visible symptom: every trace reaches the backend complete, exactly once. + + A sibling flush that carries an unconcluded trace both corrupts it - ``output=None`` on the + wire - and consumes it, because the batch is detached from the shared list. Its owner's own + flush then finds nothing left to send and reports success, so nothing anywhere signals that + the trace was shipped half-built. The backend rejects the repeat as a duplicate id, so there + is no second chance either. + + Asserts on the serialised outputs, not just the names, because a trace can arrive under the + right name and still be missing everything that makes it useful. + """ + # Given: two tasks sharing a logger, each concluding its own trace and flushing + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + + london_trace_started = asyncio.Event() + new_york_flush_done = asyncio.Event() + ingested: list[tuple[str, str | None]] = [] + + async def record_payload(request: TracesIngestRequest) -> dict: + # Snapshot synchronously: the real client calls model_dump() before its first await, + # so reading these later would observe mutations that never reached the wire. + ingested.extend((trace.name, trace.output) for trace in request.traces) + return {} + + mock_traces_client_instance.ingest_traces = AsyncMock(side_effect=record_payload) + + async def london_task() -> None: + logger.start_trace(input="weather in London?", name="London") + london_trace_started.set() + await new_york_flush_done.wait() + logger.conclude(output="rainy in London") + await logger.async_flush() + + london = asyncio.create_task(london_task()) + await asyncio.wait_for(london_trace_started.wait(), timeout=5) + + # When: New York's whole trace happens inside London's, and both flush + logger.start_trace(input="weather in New York?", name="New York") + logger.conclude(output="sunny in New York") + await logger.async_flush() + new_york_flush_done.set() + await london + + # Then: each trace was ingested exactly once, carrying its own output + assert sorted(ingested) == [("London", "rainy in London"), ("New York", "sunny in New York")] + assert logger.traces == [] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_auto_conclude_does_not_borrow_another_tasks_output( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """Auto-concluding a trace must derive its output from the caller's own spans. + + ``_flush_batch`` concludes whatever the flushing context left open, and a trace with no + explicit output inherits its last child span's. Selecting that child from the last entry of + the shared trace list picks a concurrent task's trace whenever one arrived more recently, so + the flushing task's trace is stamped with an answer computed for somebody else's question - + a wrong output rather than a missing one, and invisible in the console. + + The flushing task never calls ``conclude()``, so the auto-conclude path inside the flush is + what runs. London starts second purely to occupy the last slot in the list, and both tasks + are created before either starts a trace: a task created while the creating context already + holds an open trace inherits it through the copied context, and ``add_trace()`` then refuses + to start a second one. + """ + # Given: two tasks sharing a logger, with the sibling's trace last in the shared list + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + mock_traces_client_instance.ingest_traces = AsyncMock(return_value={}) + + new_york_trace_started = asyncio.Event() + london_trace_started = asyncio.Event() + new_york_flush_done = asyncio.Event() + new_york_trace: list[LoggedTrace] = [] + + async def new_york_task() -> None: + new_york_trace.append(logger.start_trace(input="weather in New York?", name="New York")) + logger.add_llm_span(input="weather in New York?", output="sunny in New York", model="gpt-4o") + new_york_trace_started.set() + await london_trace_started.wait() + await logger.async_flush() + new_york_flush_done.set() + + async def london_task() -> None: + await new_york_trace_started.wait() + logger.start_trace(input="weather in London?", name="London") + logger.add_llm_span(input="weather in London?", output="rainy in London", model="gpt-4o") + london_trace_started.set() + await new_york_flush_done.wait() + + # When: New York flushes without concluding, so the flush auto-concludes on its behalf + await asyncio.wait_for(asyncio.gather(new_york_task(), london_task()), timeout=5) + + # Then: New York's output came from New York's own span + assert new_york_trace[0].output is not None + assert "sunny in New York" in new_york_trace[0].output + assert "rainy in London" not in new_york_trace[0].output + + +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +def test_flush_sends_the_callers_own_unconcluded_trace( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """Flushing your own open trace still sends it. + + Holding back traces that are still being built must not turn into "unconcluded means never + sent": a caller that flushes while its own trace is open has asked for that trace, and the + flush concludes it on the way out. Only *other* contexts' unfinished traces wait. + """ + # Given: a logger with an open, never-concluded trace + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + logger.start_trace(input="weather in New York?", name="New York") + logger.add_llm_span(input="weather in New York?", output="sunny in New York", model="gpt-4o") + + # When: the same context flushes without concluding + logger.flush() + + # Then: the trace was sent, carrying the output inherited from its span + mock_traces_client_instance.ingest_traces.assert_called_once() + payload: TracesIngestRequest = mock_traces_client_instance.ingest_traces.call_args.args[0] + assert [trace.name for trace in payload.traces] == ["New York"] + assert payload.traces[0].output is not None + assert "sunny in New York" in payload.traces[0].output + assert logger.traces == [] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_terminate_sends_a_trace_abandoned_by_a_finished_task( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """A trace whose owner never concluded it must still be sent at lifecycle end. + + Holding back traces that another context is still building leaves nobody to flush one whose + owner has gone away without concluding it - the context that could claim it no longer exists. + ``terminate()`` is the last chance: nothing can still be running, so its final flush sends + everything rather than stranding a trace in memory until the process dies. + """ + # Given: a task that starts a trace, never concludes it, and finishes + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + + async def abandoning_task() -> None: + logger.start_trace(input="weather in London?", name="London") + logger.add_llm_span(input="weather in London?", output="rainy in London", model="gpt-4o") + + await asyncio.create_task(abandoning_task()) + assert [trace.name for trace in logger.traces] == ["London"] + + # When: the logger is terminated from a context that never owned that trace + logger.terminate() + + # Then: the abandoned trace was sent rather than stranded + mock_traces_client_instance.ingest_traces.assert_called_once() + payload: TracesIngestRequest = mock_traces_client_instance.ingest_traces.call_args.args[0] + assert [trace.name for trace in payload.traces] == ["London"] + assert logger.traces == [] + + +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +def test_reset_parent_tracking_does_not_strand_its_trace( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """Abandoning a parent chain releases its trace for the next flush. + + ``reset_parent_tracking()`` drops the chain without concluding it, so the trace it was + building has no owner left to claim it. Treating that as "still being built" would hold it + back from every subsequent flush until the process exits. ``galileo_context`` setup calls this + on every entry, so the case is routine rather than exotic. + + The abandoning thread is kept alive and is not the thread that flushes, because a live owner + that *is* the flushing thread is already treated as finished - so doing both on one thread + would leave this test passing with the release removed. + """ + # Given: a live thread that abandoned its open trace by resetting parent tracking + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + trace_abandoned = threading.Event() + owner_may_finish = threading.Event() + + def abandon_a_trace() -> None: + logger.start_trace(input="weather in New York?", name="New York") + logger.reset_parent_tracking() + assert logger.current_parent() is None + trace_abandoned.set() + owner_may_finish.wait(timeout=5) + + owner = threading.Thread(target=abandon_a_trace, name="trace-owner") + owner.start() + assert trace_abandoned.wait(timeout=5) + + # When: a different thread flushes while the abandoning thread is still running + logger.flush() + + # Then: the abandoned trace was sent, not held back + mock_traces_client_instance.ingest_traces.assert_called_once() + payload: TracesIngestRequest = mock_traces_client_instance.ingest_traces.call_args.args[0] + assert [trace.name for trace in payload.traces] == ["New York"] + assert logger.traces == [] + + owner_may_finish.set() + owner.join(timeout=5) + + +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +def test_flush_still_coalesces_concluded_traces_into_one_request( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """Batch mode still batches. + + Deciding per-trace whether to send must not degrade into one request per trace: accumulating + many traces and flushing once is the pattern batch mode exists to serve, and splitting it + would multiply requests for every caller in exchange for a concurrency fix they may not need. + """ + # Given: three traces accumulated and concluded before any flush + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + for city in ("New York", "London", "Tokyo"): + logger.start_trace(input=f"weather in {city}?", name=city) + logger.conclude(output=f"forecast for {city}") + + # When: the logger flushes once + logger.flush() + + # Then: all three left in a single request + mock_traces_client_instance.ingest_traces.assert_called_once() + payload: TracesIngestRequest = mock_traces_client_instance.ingest_traces.call_args.args[0] + assert [trace.name for trace in payload.traces] == ["New York", "London", "Tokyo"] + assert logger.traces == [] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_failed_flush_keeps_holding_back_another_tasks_trace( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """A failed send must restore its own batch without adopting a sibling's open trace. + + The restore path puts the detached batch back into the shared list, which is also where a + held-back trace is waiting. If the restore blurred the two, a retry would sweep up the + sibling's unfinished trace - reintroducing the defect on exactly the path added to guard + against data loss. + """ + # Given: a sibling holding a trace open, a concluded trace of our own, and a failing ingest + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + + london_trace_started = asyncio.Event() + london_may_conclude = asyncio.Event() + ingest_payloads: list[list[str]] = [] + + async def london_task() -> None: + logger.start_trace(input="weather in London?", name="London") + london_trace_started.set() + await london_may_conclude.wait() + logger.conclude(output="rainy in London") + await logger.async_flush() + + async def record_payload(request: TracesIngestRequest) -> dict: + ingest_payloads.append([trace.name for trace in request.traces]) + return {} + + london = asyncio.create_task(london_task()) + await asyncio.wait_for(london_trace_started.wait(), timeout=5) + + logger.start_trace(input="weather in New York?", name="New York") + logger.conclude(output="sunny in New York") + + mock_traces_client_instance.ingest_traces = AsyncMock(side_effect=ConnectionError("backend unavailable")) + + # When: our flush fails and is retried after the backend recovers + await logger.async_flush() + assert sorted(trace.name for trace in logger.traces) == ["London", "New York"] + + mock_traces_client_instance.ingest_traces = AsyncMock(side_effect=record_payload) + await logger.async_flush() + + # Then: the retry sent only our trace, and London still waits for its own flush + assert ingest_payloads == [["New York"]] + assert [trace.name for trace in logger.traces] == ["London"] + + london_may_conclude.set() + await london + + # Then: London's own flush sends it, complete + assert ingest_payloads == [["New York"], ["London"]] + assert logger.traces == [] + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_failed_flush_restores_the_ownership_mark_with_the_batch( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """A failed send must restore what its batch was marked with, not only the batch. + + A flush takes its own trace whether or not it is finished with it, so a trace still being built + can sit inside a batch whose send fails. Restoring the list without the marks would hand that + trace to the next flush from any other context - the premature send this guard exists to + prevent, reintroduced on the path that exists to prevent data loss. + + The other direction, a restore that adopts a trace it never detached, is pinned by + `test_failed_flush_keeps_holding_back_another_tasks_trace`; neither test covers both. + """ + # Given: an owner part-way through a trace, and an ingest that fails + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + + owner_flush_failed = asyncio.Event() + owner_may_finish = asyncio.Event() + ingest_payloads: list[list[str]] = [] + + async def record_payload(request: TracesIngestRequest) -> dict: + ingest_payloads.append([trace.name for trace in request.traces]) + return {} + + async def owner_task() -> None: + logger.start_trace(input="weather in London?", name="London") + logger.add_llm_span(input="weather in London?", output="checking London", model="gpt-4o") + # When: the owner's own flush detaches its unconcluded trace and the send fails + await logger.async_flush() + owner_flush_failed.set() + # Kept alive deliberately: a trace whose owner has finished is nobody's to build any more, + # so the flush below would be let through on liveness alone and pin nothing. + await owner_may_finish.wait() + + mock_traces_client_instance.ingest_traces = AsyncMock(side_effect=ConnectionError("backend unavailable")) + owner = asyncio.create_task(owner_task()) + await asyncio.wait_for(owner_flush_failed.wait(), timeout=5) + + # Then: the trace is back, and back with its owner still recorded + assert [trace.name for trace in logger.traces] == ["London"] + assert logger.traces[0].id in logger._traces_being_built + + # When: another context flushes while that owner is still running + mock_traces_client_instance.ingest_traces = AsyncMock(side_effect=record_payload) + await logger.async_flush() + + # Then: nothing was sent - London waits for its owner rather than leaving half-built + assert ingest_payloads == [] + assert [trace.name for trace in logger.traces] == ["London"] + + # Then: the restored mark holds the trace for its owner without stranding it + owner_may_finish.set() + await owner + await logger.async_flush() + assert ingest_payloads == [["London"]] + assert logger.traces == [] + + +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +def test_flush_sends_the_callers_own_trace_even_when_concluding_it_fails( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """The flushing caller's trace is sent even if the flush cannot conclude it. + + Deciding what to send from "has this been concluded" alone would strand any trace whose + auto-conclude did not take effect - and that path swallows its own exceptions, so the failure + is invisible. A real case exists today: a child span whose output type cannot be assigned to + its parent raises inside ``_auto_conclude_trace()``, which ``warn_catch_exception`` eats. The + caller therefore claims its own trace explicitly, independently of whether concluding worked. + + ``_auto_conclude_trace`` is stubbed to a no-op rather than reproducing that type mismatch, so + the guarantee is pinned on its own terms and survives whatever happens to the mismatch bug. + """ + # Given: a logger with an open trace whose auto-conclude will silently do nothing + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + logger.start_trace(input="weather in New York?", name="New York") + logger.add_llm_span(input="weather in New York?", output="sunny in New York", model="gpt-4o") + + # When: the owning context flushes + with patch.object(GalileoLogger, "_auto_conclude_trace", lambda self: None): + logger.flush() + + # Then: the trace was still sent, and its id was released rather than tracked forever + mock_traces_client_instance.ingest_traces.assert_called_once() + payload: TracesIngestRequest = mock_traces_client_instance.ingest_traces.call_args.args[0] + assert [trace.name for trace in payload.traces] == ["New York"] + assert logger.traces == [] + assert logger._traces_being_built == {} + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_flush_sends_a_trace_whose_owning_task_has_finished( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """A trace is only held back while the context building it is still alive. + + Holding a trace back until its owner finishes is what stops a sibling's flush shipping it + half-built. But an owner can go away without concluding - user code raising between + ``start_trace()`` and ``conclude()``, or ``_conclude()`` swallowing a coercion error - and then + no context can claim it: the claim is resolved from the caller's parent chain, and the owner's + chain died with its task. Without a liveness check the trace is skipped by every flush and + accumulates in memory, which is a worse outcome than the half-built send this rule prevents. + + The abandoning call runs in its own task and the flush comes from outside it, because a flush in + the owning context claims its own trace and would mask the leak. + """ + # Given: a task that starts a trace, adds a span, and finishes without concluding + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + + ingest_payloads: list[list[str]] = [] + + async def record_payload(request: TracesIngestRequest) -> dict: + ingest_payloads.append([trace.name for trace in request.traces]) + return {} + + mock_traces_client_instance.ingest_traces = AsyncMock(side_effect=record_payload) + + async def abandoning_task() -> None: + logger.start_trace(input="weather in London?", name="London") + logger.add_llm_span(input="weather in London?", output="rainy in London", model="gpt-4o") + + await asyncio.create_task(abandoning_task()) + assert [trace.name for trace in logger.traces] == ["London"] + + # When: a context that never owned that trace flushes, after its owner has finished + await logger.async_flush() + + # Then: the trace was sent rather than held back until the process exits + assert ingest_payloads == [["London"]] + assert logger.traces == [] + assert logger._traces_being_built == {} + + +@pytest.mark.asyncio +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +async def test_reset_parent_tracking_in_another_context_does_not_strand_a_trace( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """Abandoning a trace from a context that does not own it must not strand it either. + + ``reset_parent_tracking()`` releases the trace it abandons, but only the caller's own: it + resolves the id from the caller's parent chain, which is empty in any other context. So the + call a reader would reach for to clean up after a finished task does nothing for that task's + trace, and only the owner's liveness distinguishes "abandoned" from "still being built". + """ + # Given: a task that starts a trace, adds a span, and finishes without concluding + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + + ingest_payloads: list[list[str]] = [] + + async def record_payload(request: TracesIngestRequest) -> dict: + ingest_payloads.append([trace.name for trace in request.traces]) + return {} + + mock_traces_client_instance.ingest_traces = AsyncMock(side_effect=record_payload) + + async def abandoning_task() -> None: + logger.start_trace(input="weather in London?", name="London") + logger.add_llm_span(input="weather in London?", output="rainy in London", model="gpt-4o") + + await asyncio.create_task(abandoning_task()) + + # When: another context resets its own parent tracking and then flushes + logger.reset_parent_tracking() + await logger.async_flush() + + # Then: the abandoned trace was sent + assert ingest_payloads == [["London"]] + assert logger.traces == [] + + +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +def test_flush_sends_a_trace_abandoned_on_a_still_alive_pool_thread( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """A worker thread outliving the job that abandoned a trace must not hold it forever. + + For a synchronous caller the owner recorded is the thread, and a pool worker stays alive long + after the job that started the trace returned. The next job on that worker cannot claim the trace + either: pools that copy the context per job - anyio, and therefore FastAPI's synchronous + endpoints - hand it an empty parent chain. Liveness alone would hold the trace for the life of + the worker, where before this rule existed it was sent on the first flush. + + Each job runs inside its own ``copy_context()``, because that isolation is the property that + matters rather than the pool library: ``ThreadPoolExecutor`` does not copy, so its second job + still sees the first job's parent chain and claims the trace. The flush goes through the public + ``flush()`` so the thread is captured where a real caller captures it - reading it inside + ``_flush_batch`` would see an ``EventLoopThreadPool`` thread instead. + """ + # Given: a job on a worker thread that starts a trace and returns without concluding it + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + + ingest_payloads: list[list[str]] = [] + + async def record_payload(request: TracesIngestRequest) -> dict: + ingest_payloads.append([trace.name for trace in request.traces]) + return {} + + mock_traces_client_instance.ingest_traces = AsyncMock(side_effect=record_payload) + + def abandoning_job() -> None: + logger.start_trace(input="weather in London?", name="London") + logger.add_llm_span(input="weather in London?", output="rainy in London", model="gpt-4o") + + queued_between_jobs: list[list[str]] = [] + + def worker() -> None: + contextvars.copy_context().run(abandoning_job) + queued_between_jobs.append([trace.name for trace in logger.traces]) + contextvars.copy_context().run(logger.flush) + + # When: a later job on the same still-alive worker flushes, in a context of its own + worker_thread = threading.Thread(target=worker, name="AnyIO worker thread") + worker_thread.start() + worker_thread.join(timeout=30) + + # Then: the abandoned trace was sent, not held until the worker dies or the process exits + assert not worker_thread.is_alive() + assert queued_between_jobs == [["London"]] + assert ingest_payloads == [["London"]] + assert logger.traces == [] + assert logger._traces_being_built == {} + + +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +def test_flush_holds_back_a_trace_another_live_thread_is_building( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + """Releasing a trace whose owner is the flushing thread must not release another thread's. + + This is the other half of that rule, and it matters most where it is least visible: anyio names + every worker ``"AnyIO worker thread"``, so they resolve to one ``GalileoLogger`` and share one + trace list. Comparing thread *names* rather than identity would therefore let one worker's flush + ship another worker's half-built trace - the failure the hold-back exists to prevent - so the + flushing thread here is deliberately given the builder's name. + """ + # Given: a worker thread that is still building a trace, parked while it holds it open + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + logger = GalileoLogger(project="my_project", log_stream="my_log_stream") + + ingest_payloads: list[list[str]] = [] + + async def record_payload(request: TracesIngestRequest) -> dict: + ingest_payloads.append([trace.name for trace in request.traces]) + return {} + + mock_traces_client_instance.ingest_traces = AsyncMock(side_effect=record_payload) + + trace_started = threading.Event() + let_builder_finish = threading.Event() + + def building_job() -> None: + logger.start_trace(input="weather in London?", name="London") + logger.add_llm_span(input="weather in London?", output="rainy in London", model="gpt-4o") + trace_started.set() + let_builder_finish.wait(timeout=30) + + builder = threading.Thread(target=building_job, name="AnyIO worker thread") + builder.start() + try: + assert trace_started.wait(timeout=30) + + # When: a different, still-live thread carrying the same name flushes + flusher = threading.Thread(target=logger.flush, name="AnyIO worker thread") + flusher.start() + flusher.join(timeout=30) + assert not flusher.is_alive() + finally: + # Always release the builder: an unreleased event would park an xdist worker until the + # suite-wide `--timeout` fires. + let_builder_finish.set() + builder.join(timeout=30) + + # Then: the trace stayed behind for the thread that is still building it + assert ingest_payloads == [] + assert [trace.name for trace in logger.traces] == ["London"] + assert len(logger._traces_being_built) == 1 diff --git a/tests/test_openai.py b/tests/test_openai.py index 9b5c524a3..0aec454c9 100644 --- a/tests/test_openai.py +++ b/tests/test_openai.py @@ -1,3 +1,4 @@ +import asyncio from unittest.mock import AsyncMock, Mock, patch import httpx @@ -663,3 +664,59 @@ def test_responses_api_streaming( assert payload.traces[0].spans[0].input == [Message(content="Say hello", role=MessageRole.user)] assert payload.traces[0].spans[0].output == Message(content="This is a test response", role=MessageRole.assistant) + + +@patch("openai.resources.chat.Completions.create", side_effect=openai.APIConnectionError(request=Request("GET", "url"))) +@patch("galileo.logger.logger.LogStreams") +@patch("galileo.logger.logger.Projects") +@patch("galileo.logger.logger.Traces") +@pytest.mark.asyncio +async def test_trace_is_not_stranded_when_the_openai_call_raises_a_non_status_error( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, openai_create +) -> None: + """A failed OpenAI call must not leave its trace unflushable. + + The wrapper starts the trace itself when there is no active one, and only ``APIStatusError`` is + routed into the response processing that concludes it - a connection error or timeout unwinds + straight to the wrapper's error path, leaving the trace unconcluded. Since a flush skips traces + another context is still building, the wrapper has to report that it is no longer building this + one, or nothing would send it until the process exits. + + The call runs in its own task and the flush comes from outside it, because a flush in the owning + context claims its own trace and would mask the leak. The task is also held alive across the + flush: a finished owner is treated as no longer building the trace, so letting the task complete + would mask the leak a second way and leave this test passing with the release removed. + """ + # Given: an OpenAI call that fails with an error the wrapper does not route into processing + mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + + galileo_context.reset() + OpenAIGalileo().register_tracing() + + request_failed = asyncio.Event() + task_may_finish = asyncio.Event() + + async def failing_request() -> None: + with pytest.raises(RuntimeError): + openai.chat.completions.create( + messages=[{"role": "user", "content": "Say this is a test"}], model="gpt-3.5-turbo" + ) + request_failed.set() + await task_may_finish.wait() + + # When: the failure happens in another task that is still alive, and the flush comes from this one + owner = asyncio.create_task(failing_request()) + await asyncio.wait_for(request_failed.wait(), timeout=5) + galileo_context.flush() + ingest_calls_after_foreign_flush = mock_traces_client_instance.ingest_traces.call_args_list[:] + + task_may_finish.set() + await owner + + # Then: the trace was still sent rather than held back until process exit + openai_create.assert_called_once() + assert len(ingest_calls_after_foreign_flush) == 1 + payload = ingest_calls_after_foreign_flush[0][0][0] + assert len(payload.traces) == 1