mirror of
https://github.com/tinyhumansai/openhuman.git
synced 2026-07-27 21:08:00 +00:00
feat(memory): singleton ingestion + status RPC + UI pill (#1126)
This commit is contained in:
Steven Enamakel
@@ -0,0 +1,80 @@
|
||||
import { act, renderHook, waitFor } from '@testing-library/react';
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { useMemoryIngestionStatus } from '../useMemoryIngestionStatus';
|
||||
|
||||
const mockCallCoreRpc = vi.fn();
|
||||
|
||||
vi.mock('../../services/coreRpcClient', () => ({
|
||||
callCoreRpc: (args: unknown) => mockCallCoreRpc(args),
|
||||
}));
|
||||
|
||||
describe('useMemoryIngestionStatus', () => {
|
||||
beforeEach(() => {
|
||||
mockCallCoreRpc.mockReset();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it('maps the snake_case RPC envelope into camelCase status', async () => {
|
||||
mockCallCoreRpc.mockResolvedValue({
|
||||
running: true,
|
||||
current_document_id: 'doc-1',
|
||||
current_title: 'Notes',
|
||||
current_namespace: 'global',
|
||||
queue_depth: 2,
|
||||
last_completed_at: 1700000000000,
|
||||
last_document_id: 'doc-0',
|
||||
last_success: true,
|
||||
});
|
||||
|
||||
const { result } = renderHook(() => useMemoryIngestionStatus());
|
||||
|
||||
await waitFor(() => expect(result.current.loading).toBe(false));
|
||||
|
||||
expect(result.current.status).toEqual({
|
||||
running: true,
|
||||
currentDocumentId: 'doc-1',
|
||||
currentTitle: 'Notes',
|
||||
currentNamespace: 'global',
|
||||
queueDepth: 2,
|
||||
lastCompletedAt: 1700000000000,
|
||||
lastDocumentId: 'doc-0',
|
||||
lastSuccess: true,
|
||||
});
|
||||
expect(result.current.error).toBeNull();
|
||||
expect(mockCallCoreRpc).toHaveBeenCalledWith({ method: 'openhuman.memory_ingestion_status' });
|
||||
});
|
||||
|
||||
it('reports an error when the RPC fails and keeps idle defaults', async () => {
|
||||
mockCallCoreRpc.mockRejectedValue(new Error('boom'));
|
||||
|
||||
const { result } = renderHook(() => useMemoryIngestionStatus());
|
||||
|
||||
await waitFor(() => expect(result.current.loading).toBe(false));
|
||||
|
||||
expect(result.current.status.running).toBe(false);
|
||||
expect(result.current.status.queueDepth).toBe(0);
|
||||
expect(result.current.error).toBe('boom');
|
||||
});
|
||||
|
||||
it('refresh() re-issues the RPC and updates status', async () => {
|
||||
mockCallCoreRpc
|
||||
.mockResolvedValueOnce({ running: false, queue_depth: 0 })
|
||||
.mockResolvedValueOnce({ running: true, queue_depth: 1, current_document_id: 'doc-2' });
|
||||
|
||||
const { result } = renderHook(() => useMemoryIngestionStatus());
|
||||
await waitFor(() => expect(result.current.loading).toBe(false));
|
||||
expect(result.current.status.running).toBe(false);
|
||||
|
||||
await act(async () => {
|
||||
await result.current.refresh();
|
||||
});
|
||||
|
||||
expect(result.current.status.running).toBe(true);
|
||||
expect(result.current.status.queueDepth).toBe(1);
|
||||
expect(result.current.status.currentDocumentId).toBe('doc-2');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,97 @@
|
||||
import { useCallback, useEffect, useRef, useState } from 'react';
|
||||
|
||||
import { callCoreRpc } from '../services/coreRpcClient';
|
||||
|
||||
export interface MemoryIngestionStatus {
|
||||
running: boolean;
|
||||
currentDocumentId?: string;
|
||||
currentTitle?: string;
|
||||
currentNamespace?: string;
|
||||
queueDepth: number;
|
||||
lastCompletedAt?: number;
|
||||
lastDocumentId?: string;
|
||||
lastSuccess?: boolean;
|
||||
}
|
||||
|
||||
interface IngestionStatusEnvelope {
|
||||
running: boolean;
|
||||
current_document_id?: string;
|
||||
current_title?: string;
|
||||
current_namespace?: string;
|
||||
queue_depth: number;
|
||||
last_completed_at?: number;
|
||||
last_document_id?: string;
|
||||
last_success?: boolean;
|
||||
}
|
||||
|
||||
const DEFAULT_POLL_MS = 4000;
|
||||
const FAST_POLL_MS = 1500;
|
||||
|
||||
const EMPTY_STATUS: MemoryIngestionStatus = { running: false, queueDepth: 0 };
|
||||
|
||||
/**
|
||||
* Polls `openhuman.memory_ingestion_status`. Polls faster while a job is
|
||||
* running or queued so the UI reacts quickly when ingestion finishes;
|
||||
* relaxes to a slower cadence at idle.
|
||||
*/
|
||||
export function useMemoryIngestionStatus(): {
|
||||
status: MemoryIngestionStatus;
|
||||
loading: boolean;
|
||||
error: string | null;
|
||||
refresh: () => void;
|
||||
} {
|
||||
const [status, setStatus] = useState<MemoryIngestionStatus>(EMPTY_STATUS);
|
||||
const [loading, setLoading] = useState(true);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
const cancelledRef = useRef(false);
|
||||
|
||||
const fetchOnce = useCallback(async () => {
|
||||
try {
|
||||
const env = await callCoreRpc<IngestionStatusEnvelope>({
|
||||
method: 'openhuman.memory_ingestion_status',
|
||||
});
|
||||
if (cancelledRef.current) return;
|
||||
setStatus({
|
||||
running: env.running,
|
||||
currentDocumentId: env.current_document_id,
|
||||
currentTitle: env.current_title,
|
||||
currentNamespace: env.current_namespace,
|
||||
queueDepth: env.queue_depth ?? 0,
|
||||
lastCompletedAt: env.last_completed_at,
|
||||
lastDocumentId: env.last_document_id,
|
||||
lastSuccess: env.last_success,
|
||||
});
|
||||
setError(null);
|
||||
} catch (err) {
|
||||
if (cancelledRef.current) return;
|
||||
setError(err instanceof Error ? err.message : String(err));
|
||||
} finally {
|
||||
if (!cancelledRef.current) setLoading(false);
|
||||
}
|
||||
}, []);
|
||||
|
||||
const statusRef = useRef(status);
|
||||
statusRef.current = status;
|
||||
|
||||
useEffect(() => {
|
||||
cancelledRef.current = false;
|
||||
let timer: ReturnType<typeof setTimeout> | null = null;
|
||||
|
||||
const tick = async () => {
|
||||
await fetchOnce();
|
||||
if (cancelledRef.current) return;
|
||||
const live = statusRef.current;
|
||||
const delay = live.running || live.queueDepth > 0 ? FAST_POLL_MS : DEFAULT_POLL_MS;
|
||||
timer = setTimeout(tick, delay);
|
||||
};
|
||||
|
||||
void tick();
|
||||
|
||||
return () => {
|
||||
cancelledRef.current = true;
|
||||
if (timer) clearTimeout(timer);
|
||||
};
|
||||
}, [fetchOnce]);
|
||||
|
||||
return { status, loading, error, refresh: fetchOnce };
|
||||
}
|
||||
@@ -17,6 +17,7 @@ import {
|
||||
useIntelligenceSocketManager,
|
||||
} from '../hooks/useIntelligenceSocket';
|
||||
import { useIntelligenceStats } from '../hooks/useIntelligenceStats';
|
||||
import { useMemoryIngestionStatus } from '../hooks/useMemoryIngestionStatus';
|
||||
import { useScreenIntelligenceItems } from '../hooks/useScreenIntelligenceItems';
|
||||
import { useSubconscious } from '../hooks/useSubconscious';
|
||||
import type {
|
||||
@@ -31,6 +32,7 @@ type IntelligenceTab = 'memory' | 'subconscious' | 'dreams';
|
||||
|
||||
export default function Intelligence() {
|
||||
const { aiStatus } = useIntelligenceStats();
|
||||
const { status: ingestionStatus } = useMemoryIngestionStatus();
|
||||
|
||||
const [activeTab, setActiveTab] = useState<IntelligenceTab>('memory');
|
||||
const [sourceFilter, setSourceFilter] = useState<ActionableItemSource | 'all'>('all');
|
||||
@@ -312,6 +314,24 @@ export default function Intelligence() {
|
||||
<span className="text-xs text-stone-400">{systemStatusLabel}</span>
|
||||
</div>
|
||||
)}
|
||||
{activeTab === 'memory' &&
|
||||
(ingestionStatus.running || ingestionStatus.queueDepth > 0) && (
|
||||
<div
|
||||
className="flex items-center gap-1.5 px-2 py-0.5 rounded-full border border-amber-200 bg-amber-50 text-amber-700"
|
||||
title={
|
||||
ingestionStatus.running
|
||||
? ingestionStatus.currentTitle
|
||||
? `Ingesting: ${ingestionStatus.currentTitle}`
|
||||
: 'Memory ingestion running'
|
||||
: 'Memory ingestion queued'
|
||||
}>
|
||||
<div className="w-1.5 h-1.5 rounded-full bg-amber-500 animate-pulse" />
|
||||
<span className="text-[11px] font-medium">
|
||||
{ingestionStatus.running ? 'Ingesting' : 'Queued'}
|
||||
{ingestionStatus.queueDepth > 0 && ` · ${ingestionStatus.queueDepth}`}
|
||||
</span>
|
||||
</div>
|
||||
)}
|
||||
{activeTab === 'memory' && (
|
||||
<button
|
||||
onClick={usingMemoryData ? refreshConscious : handleAnalyzeNow}
|
||||
|
||||
@@ -72,6 +72,23 @@ pub enum DomainEvent {
|
||||
/// this variant is a hook for future ingestion subscribers to react to pull
|
||||
/// requests. See `src/openhuman/memory/ops.rs` for the RPC handlers.
|
||||
MemorySyncRequested { channel_id: Option<String> },
|
||||
/// A memory ingestion job started running on the local extraction LLM.
|
||||
/// Ingestion is singleton — this fires once, then a matching
|
||||
/// [`Self::MemoryIngestionCompleted`] follows when the job finishes.
|
||||
MemoryIngestionStarted {
|
||||
document_id: String,
|
||||
title: String,
|
||||
namespace: String,
|
||||
queue_depth: usize,
|
||||
},
|
||||
/// A memory ingestion job finished (successfully or with an error).
|
||||
MemoryIngestionCompleted {
|
||||
document_id: String,
|
||||
namespace: String,
|
||||
success: bool,
|
||||
elapsed_ms: u64,
|
||||
queue_depth: usize,
|
||||
},
|
||||
|
||||
// ── Channels ────────────────────────────────────────────────────────
|
||||
/// An inbound channel message from the transport layer, ready for processing.
|
||||
@@ -366,7 +383,9 @@ impl DomainEvent {
|
||||
|
||||
Self::MemoryStored { .. }
|
||||
| Self::MemoryRecalled { .. }
|
||||
| Self::MemorySyncRequested { .. } => "memory",
|
||||
| Self::MemorySyncRequested { .. }
|
||||
| Self::MemoryIngestionStarted { .. }
|
||||
| Self::MemoryIngestionCompleted { .. } => "memory",
|
||||
|
||||
Self::ChannelInboundMessage { .. }
|
||||
| Self::ChannelMessageReceived { .. }
|
||||
|
||||
@@ -8,10 +8,13 @@
|
||||
//! from the actual extraction process.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::core::event_bus::{publish_global, DomainEvent};
|
||||
use crate::openhuman::memory::ingestion::MemoryIngestionConfig;
|
||||
use crate::openhuman::memory::ingestion_state::IngestionState;
|
||||
use crate::openhuman::memory::store::{NamespaceDocumentInput, UnifiedMemory};
|
||||
|
||||
/// A job submitted to the ingestion worker.
|
||||
@@ -37,6 +40,8 @@ pub struct IngestionJob {
|
||||
pub struct IngestionQueue {
|
||||
/// Sender half of the job queue channel.
|
||||
tx: mpsc::UnboundedSender<IngestionJob>,
|
||||
/// Shared state — singleton lock, queue depth, status snapshot.
|
||||
state: IngestionState,
|
||||
}
|
||||
|
||||
impl IngestionQueue {
|
||||
@@ -52,9 +57,12 @@ impl IngestionQueue {
|
||||
/// worker has shut down (e.g., during application termination) and the
|
||||
/// job was dropped.
|
||||
pub fn submit(&self, job: IngestionJob) -> bool {
|
||||
self.state.enqueue();
|
||||
match self.tx.send(job) {
|
||||
Ok(()) => true,
|
||||
Err(e) => {
|
||||
// Worker is gone — undo the enqueue bump so depth stays accurate.
|
||||
self.state.dequeue();
|
||||
log::warn!(
|
||||
"[memory:ingestion_queue] failed to enqueue job (worker gone?): {}",
|
||||
e.0.document.title,
|
||||
@@ -63,6 +71,13 @@ impl IngestionQueue {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a clone of the shared ingestion state. Use this to drive the
|
||||
/// status RPC or to share the singleton lock with synchronous ingest
|
||||
/// paths that bypass the queue.
|
||||
pub fn state(&self) -> IngestionState {
|
||||
self.state.clone()
|
||||
}
|
||||
}
|
||||
|
||||
/// Start the background ingestion worker.
|
||||
@@ -77,14 +92,23 @@ impl IngestionQueue {
|
||||
/// any number of producers. The worker runs on a dedicated tokio task,
|
||||
/// processing jobs sequentially so ingestion work stays serialized.
|
||||
pub fn start_worker(memory: Arc<UnifiedMemory>) -> IngestionQueue {
|
||||
// Create an unbounded channel for the ingestion jobs.
|
||||
let state = IngestionState::new();
|
||||
start_worker_with_state(memory, state)
|
||||
}
|
||||
|
||||
/// Start a worker bound to a caller-supplied [`IngestionState`]. Useful when
|
||||
/// the synchronous ingest path needs to share the same singleton lock and
|
||||
/// snapshot as the queue worker.
|
||||
pub fn start_worker_with_state(
|
||||
memory: Arc<UnifiedMemory>,
|
||||
state: IngestionState,
|
||||
) -> IngestionQueue {
|
||||
let (tx, rx) = mpsc::unbounded_channel::<IngestionJob>();
|
||||
|
||||
// Spawn the worker loop as a background task.
|
||||
tokio::spawn(ingestion_worker(memory, rx));
|
||||
tokio::spawn(ingestion_worker(memory, rx, state.clone()));
|
||||
|
||||
log::info!("[memory:ingestion_queue] background worker started");
|
||||
IngestionQueue { tx }
|
||||
IngestionQueue { tx, state }
|
||||
}
|
||||
|
||||
/// The main worker loop for background document ingestion.
|
||||
@@ -99,6 +123,7 @@ pub fn start_worker(memory: Arc<UnifiedMemory>) -> IngestionQueue {
|
||||
async fn ingestion_worker(
|
||||
memory: Arc<UnifiedMemory>,
|
||||
mut rx: mpsc::UnboundedReceiver<IngestionJob>,
|
||||
state: IngestionState,
|
||||
) {
|
||||
log::debug!("[memory:ingestion_queue] worker loop entered");
|
||||
|
||||
@@ -113,8 +138,24 @@ async fn ingestion_worker(
|
||||
doc_id={document_id}, title={title}",
|
||||
);
|
||||
|
||||
// Perform the graph extraction. This is the most resource-intensive step.
|
||||
match memory
|
||||
// Acquire the singleton lock so only one ingestion runs at a time
|
||||
// (covers both queue worker and synchronous callers sharing this
|
||||
// state). Decrement the pending-queue counter only after we hold the
|
||||
// lock — while we're blocked waiting on it the job is still queued.
|
||||
let _guard = state.acquire().await;
|
||||
state.dequeue();
|
||||
|
||||
let queue_depth = state.snapshot().queue_depth;
|
||||
state.mark_running(&document_id, &title, &namespace);
|
||||
publish_global(DomainEvent::MemoryIngestionStarted {
|
||||
document_id: document_id.clone(),
|
||||
title: title.clone(),
|
||||
namespace: namespace.clone(),
|
||||
queue_depth,
|
||||
});
|
||||
|
||||
let started = Instant::now();
|
||||
let success = match memory
|
||||
.extract_graph(&document_id, &job.document, &job.config)
|
||||
.await
|
||||
{
|
||||
@@ -127,14 +168,27 @@ async fn ingestion_worker(
|
||||
result.relation_count,
|
||||
result.chunk_count,
|
||||
);
|
||||
true
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!(
|
||||
"[memory:ingestion_queue] extraction failed namespace={namespace} \
|
||||
doc_id={document_id} title={title}: {e}",
|
||||
);
|
||||
false
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let elapsed_ms = started.elapsed().as_millis() as u64;
|
||||
let completed_at_ms = chrono::Utc::now().timestamp_millis();
|
||||
state.mark_completed(&document_id, success, completed_at_ms);
|
||||
publish_global(DomainEvent::MemoryIngestionCompleted {
|
||||
document_id,
|
||||
namespace,
|
||||
success,
|
||||
elapsed_ms,
|
||||
queue_depth: state.snapshot().queue_depth,
|
||||
});
|
||||
}
|
||||
|
||||
log::info!("[memory:ingestion_queue] worker shut down (channel closed)");
|
||||
|
||||
@@ -0,0 +1,182 @@
|
||||
//! Shared state + singleton lock for memory ingestion.
|
||||
//!
|
||||
//! Memory ingestion runs the local extraction LLM and must not run more than
|
||||
//! once concurrently — otherwise multiple jobs contend for the same local AI
|
||||
//! and either thrash or fail. [`IngestionState`] enforces the singleton via
|
||||
//! [`tokio::sync::Mutex`] and exposes a snapshot suitable for the
|
||||
//! `openhuman.memory_ingestion_status` RPC.
|
||||
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
use parking_lot::RwLock;
|
||||
use serde::Serialize;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
/// Snapshot of ingestion state, surfaced over RPC.
|
||||
#[derive(Debug, Clone, Default, Serialize)]
|
||||
pub struct IngestionStatusSnapshot {
|
||||
/// Whether an ingestion job is currently running.
|
||||
pub running: bool,
|
||||
/// Document id of the in-flight job, if any.
|
||||
pub current_document_id: Option<String>,
|
||||
/// Document title of the in-flight job, if any (best-effort).
|
||||
pub current_title: Option<String>,
|
||||
/// Namespace of the in-flight job, if any.
|
||||
pub current_namespace: Option<String>,
|
||||
/// Number of jobs waiting in the queue (not counting the running one).
|
||||
pub queue_depth: usize,
|
||||
/// Unix-ms timestamp of when the most recent job completed.
|
||||
pub last_completed_at: Option<i64>,
|
||||
/// Document id of the most recent completed job.
|
||||
pub last_document_id: Option<String>,
|
||||
/// Whether the most recent job succeeded.
|
||||
pub last_success: Option<bool>,
|
||||
}
|
||||
|
||||
/// Shared ingestion state + singleton lock. Cheap to clone.
|
||||
#[derive(Clone)]
|
||||
pub struct IngestionState {
|
||||
inner: Arc<IngestionStateInner>,
|
||||
}
|
||||
|
||||
struct IngestionStateInner {
|
||||
/// Singleton lock — held while a job is running.
|
||||
run_lock: Mutex<()>,
|
||||
/// Queue depth — bumped on submit, decremented when the worker pulls a job.
|
||||
queue_depth: AtomicUsize,
|
||||
/// Snapshot for status RPC.
|
||||
snapshot: RwLock<IngestionStatusSnapshot>,
|
||||
}
|
||||
|
||||
impl Default for IngestionState {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl IngestionState {
|
||||
/// Create a fresh state with empty snapshot and zero queue depth.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
inner: Arc::new(IngestionStateInner {
|
||||
run_lock: Mutex::new(()),
|
||||
queue_depth: AtomicUsize::new(0),
|
||||
snapshot: RwLock::new(IngestionStatusSnapshot::default()),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Bump the pending-queue depth (call on `submit`).
|
||||
pub fn enqueue(&self) {
|
||||
self.inner.queue_depth.fetch_add(1, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
/// Decrement pending-queue depth (call when the worker has pulled a job
|
||||
/// off the channel and is about to acquire the run lock).
|
||||
pub fn dequeue(&self) {
|
||||
self.inner.queue_depth.fetch_sub(1, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
/// Acquire the singleton run lock. Holders run ingestion serialised; any
|
||||
/// other caller blocks until the holder drops the guard.
|
||||
pub async fn acquire(&self) -> tokio::sync::MutexGuard<'_, ()> {
|
||||
self.inner.run_lock.lock().await
|
||||
}
|
||||
|
||||
/// Mark a job as in-flight in the snapshot. Caller must already hold
|
||||
/// [`Self::acquire`].
|
||||
pub fn mark_running(&self, document_id: &str, title: &str, namespace: &str) {
|
||||
let mut snap = self.inner.snapshot.write();
|
||||
snap.running = true;
|
||||
snap.current_document_id = Some(document_id.to_string());
|
||||
snap.current_title = Some(title.to_string());
|
||||
snap.current_namespace = Some(namespace.to_string());
|
||||
}
|
||||
|
||||
/// Mark the in-flight job as finished.
|
||||
pub fn mark_completed(&self, document_id: &str, success: bool, completed_at_ms: i64) {
|
||||
let mut snap = self.inner.snapshot.write();
|
||||
snap.running = false;
|
||||
snap.current_document_id = None;
|
||||
snap.current_title = None;
|
||||
snap.current_namespace = None;
|
||||
snap.last_completed_at = Some(completed_at_ms);
|
||||
snap.last_document_id = Some(document_id.to_string());
|
||||
snap.last_success = Some(success);
|
||||
}
|
||||
|
||||
/// Returns a clone of the current snapshot. Includes live queue depth.
|
||||
pub fn snapshot(&self) -> IngestionStatusSnapshot {
|
||||
let mut snap = self.inner.snapshot.read().clone();
|
||||
snap.queue_depth = self.inner.queue_depth.load(Ordering::SeqCst);
|
||||
snap
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::Arc;
|
||||
use tokio::time::{sleep, Duration};
|
||||
|
||||
#[tokio::test]
|
||||
async fn singleton_serialises_concurrent_acquires() {
|
||||
let state = IngestionState::new();
|
||||
let counter = Arc::new(parking_lot::Mutex::new(0u32));
|
||||
let max_concurrent = Arc::new(parking_lot::Mutex::new(0u32));
|
||||
|
||||
let mut handles = Vec::new();
|
||||
for _ in 0..4 {
|
||||
let state = state.clone();
|
||||
let counter = Arc::clone(&counter);
|
||||
let max_concurrent = Arc::clone(&max_concurrent);
|
||||
handles.push(tokio::spawn(async move {
|
||||
let _g = state.acquire().await;
|
||||
let now = {
|
||||
let mut c = counter.lock();
|
||||
*c += 1;
|
||||
*c
|
||||
};
|
||||
{
|
||||
let mut m = max_concurrent.lock();
|
||||
if now > *m {
|
||||
*m = now;
|
||||
}
|
||||
}
|
||||
sleep(Duration::from_millis(20)).await;
|
||||
*counter.lock() -= 1;
|
||||
}));
|
||||
}
|
||||
|
||||
for h in handles {
|
||||
h.await.unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(*max_concurrent.lock(), 1, "ingestion must be singleton");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_reports_running_and_queue_depth() {
|
||||
let state = IngestionState::new();
|
||||
state.enqueue();
|
||||
state.enqueue();
|
||||
let snap = state.snapshot();
|
||||
assert_eq!(snap.queue_depth, 2);
|
||||
assert!(!snap.running);
|
||||
|
||||
state.dequeue();
|
||||
state.mark_running("doc-1", "title", "ns");
|
||||
let snap = state.snapshot();
|
||||
assert_eq!(snap.queue_depth, 1);
|
||||
assert!(snap.running);
|
||||
assert_eq!(snap.current_document_id.as_deref(), Some("doc-1"));
|
||||
|
||||
state.mark_completed("doc-1", true, 12345);
|
||||
let snap = state.snapshot();
|
||||
assert!(!snap.running);
|
||||
assert_eq!(snap.last_document_id.as_deref(), Some("doc-1"));
|
||||
assert_eq!(snap.last_success, Some(true));
|
||||
assert_eq!(snap.last_completed_at, Some(12345));
|
||||
}
|
||||
}
|
||||
@@ -11,6 +11,7 @@ pub mod embeddings;
|
||||
pub mod global;
|
||||
pub mod ingestion;
|
||||
pub mod ingestion_queue;
|
||||
pub mod ingestion_state;
|
||||
pub mod ops;
|
||||
pub mod rpc_models;
|
||||
pub mod schemas;
|
||||
@@ -24,6 +25,7 @@ pub use ingestion::{
|
||||
MemoryIngestionRequest, MemoryIngestionResult, DEFAULT_MEMORY_EXTRACTION_MODEL,
|
||||
};
|
||||
pub use ingestion_queue::{IngestionJob, IngestionQueue};
|
||||
pub use ingestion_state::{IngestionState, IngestionStatusSnapshot};
|
||||
pub use ops as rpc;
|
||||
pub use ops::*;
|
||||
pub use rpc_models::*;
|
||||
|
||||
@@ -902,6 +902,28 @@ pub struct SyncAllResult {
|
||||
pub requested: bool,
|
||||
}
|
||||
|
||||
/// Result returned by `memory_ingestion_status`. Mirrors
|
||||
/// [`crate::openhuman::memory::IngestionStatusSnapshot`] but is the public RPC
|
||||
/// shape — the indirection keeps internal renames from breaking the wire
|
||||
/// contract.
|
||||
#[derive(Debug, Clone, Default, serde::Serialize)]
|
||||
pub struct IngestionStatusResult {
|
||||
pub running: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub current_document_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub current_title: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub current_namespace: Option<String>,
|
||||
pub queue_depth: usize,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub last_completed_at: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub last_document_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub last_success: Option<bool>,
|
||||
}
|
||||
|
||||
/// Per-namespace outcome for `memory_learn_all`.
|
||||
#[derive(Debug, serde::Serialize)]
|
||||
pub struct NamespaceLearnResult {
|
||||
@@ -972,6 +994,30 @@ pub async fn memory_sync_all() -> Result<RpcOutcome<SyncAllResult>, String> {
|
||||
Ok(RpcOutcome::new(SyncAllResult { requested: true }, vec![]))
|
||||
}
|
||||
|
||||
/// Returns the current memory-ingestion status: whether a job is running, the
|
||||
/// in-flight document, queue depth, and the most recent completion. Read-only,
|
||||
/// safe to poll.
|
||||
pub async fn memory_ingestion_status() -> Result<RpcOutcome<IngestionStatusResult>, String> {
|
||||
let snapshot = match crate::openhuman::memory::global::client_if_ready() {
|
||||
Some(c) => c.ingestion_state().snapshot(),
|
||||
// Memory not yet initialised — report idle, no in-flight job.
|
||||
None => Default::default(),
|
||||
};
|
||||
Ok(RpcOutcome::new(
|
||||
IngestionStatusResult {
|
||||
running: snapshot.running,
|
||||
current_document_id: snapshot.current_document_id,
|
||||
current_title: snapshot.current_title,
|
||||
current_namespace: snapshot.current_namespace,
|
||||
queue_depth: snapshot.queue_depth,
|
||||
last_completed_at: snapshot.last_completed_at,
|
||||
last_document_id: snapshot.last_document_id,
|
||||
last_success: snapshot.last_success,
|
||||
},
|
||||
vec![],
|
||||
))
|
||||
}
|
||||
|
||||
/// Run the tree summarizer over all (or a constrained set of) namespaces.
|
||||
///
|
||||
/// Enumerates namespaces via `namespace_list`, then for each runs
|
||||
|
||||
@@ -55,6 +55,7 @@ pub fn all_controller_schemas() -> Vec<ControllerSchema> {
|
||||
schemas("sync_channel"),
|
||||
schemas("sync_all"),
|
||||
schemas("learn_all"),
|
||||
schemas("ingestion_status"),
|
||||
]
|
||||
}
|
||||
|
||||
@@ -169,6 +170,10 @@ pub fn all_registered_controllers() -> Vec<RegisteredController> {
|
||||
schema: schemas("learn_all"),
|
||||
handler: handle_learn_all,
|
||||
},
|
||||
RegisteredController {
|
||||
schema: schemas("ingestion_status"),
|
||||
handler: handle_ingestion_status,
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
@@ -993,6 +998,63 @@ pub fn schemas(function: &str) -> ControllerSchema {
|
||||
],
|
||||
},
|
||||
|
||||
"ingestion_status" => ControllerSchema {
|
||||
namespace: "memory",
|
||||
function: "ingestion_status",
|
||||
description: "Returns the current memory-ingestion status (whether a job is running, the in-flight document, queue depth, and most recent completion). Safe to poll.",
|
||||
inputs: vec![],
|
||||
outputs: vec![
|
||||
FieldSchema {
|
||||
name: "running",
|
||||
ty: TypeSchema::Bool,
|
||||
comment: "True while an ingestion job is running on the local extraction LLM.",
|
||||
required: true,
|
||||
},
|
||||
FieldSchema {
|
||||
name: "current_document_id",
|
||||
ty: TypeSchema::Option(Box::new(TypeSchema::String)),
|
||||
comment: "Document id of the in-flight job.",
|
||||
required: false,
|
||||
},
|
||||
FieldSchema {
|
||||
name: "current_title",
|
||||
ty: TypeSchema::Option(Box::new(TypeSchema::String)),
|
||||
comment: "Title of the in-flight document.",
|
||||
required: false,
|
||||
},
|
||||
FieldSchema {
|
||||
name: "current_namespace",
|
||||
ty: TypeSchema::Option(Box::new(TypeSchema::String)),
|
||||
comment: "Namespace of the in-flight document.",
|
||||
required: false,
|
||||
},
|
||||
FieldSchema {
|
||||
name: "queue_depth",
|
||||
ty: TypeSchema::U64,
|
||||
comment: "Number of jobs waiting behind the current one.",
|
||||
required: true,
|
||||
},
|
||||
FieldSchema {
|
||||
name: "last_completed_at",
|
||||
ty: TypeSchema::Option(Box::new(TypeSchema::I64)),
|
||||
comment: "Unix-ms timestamp of the most recent completion.",
|
||||
required: false,
|
||||
},
|
||||
FieldSchema {
|
||||
name: "last_document_id",
|
||||
ty: TypeSchema::Option(Box::new(TypeSchema::String)),
|
||||
comment: "Document id of the most recent completed job.",
|
||||
required: false,
|
||||
},
|
||||
FieldSchema {
|
||||
name: "last_success",
|
||||
ty: TypeSchema::Option(Box::new(TypeSchema::Bool)),
|
||||
comment: "Whether the most recent job succeeded.",
|
||||
required: false,
|
||||
},
|
||||
],
|
||||
},
|
||||
|
||||
// ----- fallback -----
|
||||
_other => ControllerSchema {
|
||||
namespace: "memory",
|
||||
@@ -1209,6 +1271,10 @@ fn handle_learn_all(params: Map<String, Value>) -> ControllerFuture {
|
||||
})
|
||||
}
|
||||
|
||||
fn handle_ingestion_status(_params: Map<String, Value>) -> ControllerFuture {
|
||||
Box::pin(async move { to_json(rpc::memory_ingestion_status().await?) })
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -29,6 +29,7 @@ const ALL_FUNCTIONS: &[&str] = &[
|
||||
"sync_channel",
|
||||
"sync_all",
|
||||
"learn_all",
|
||||
"ingestion_status",
|
||||
];
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -16,6 +16,7 @@ use crate::openhuman::memory::ingestion::{
|
||||
MemoryIngestionConfig, MemoryIngestionRequest, MemoryIngestionResult,
|
||||
};
|
||||
use crate::openhuman::memory::ingestion_queue::{self, IngestionJob, IngestionQueue};
|
||||
use crate::openhuman::memory::ingestion_state::IngestionState;
|
||||
use crate::openhuman::memory::store::types::{
|
||||
NamespaceDocumentInput, NamespaceMemoryHit, NamespaceRetrievalContext,
|
||||
};
|
||||
@@ -93,7 +94,10 @@ impl MemoryClient {
|
||||
let inner = Arc::new(memory);
|
||||
|
||||
// Start the background worker for document ingestion and graph extraction.
|
||||
let ingestion_queue = ingestion_queue::start_worker(Arc::clone(&inner));
|
||||
// The worker shares its IngestionState with the synchronous ingest path
|
||||
// below so all ingestion is singleton-serialised.
|
||||
let ingestion_queue =
|
||||
ingestion_queue::start_worker_with_state(Arc::clone(&inner), IngestionState::new());
|
||||
|
||||
Ok(Self {
|
||||
inner,
|
||||
@@ -141,11 +145,62 @@ impl MemoryClient {
|
||||
/// Perform a full ingestion (chunking, embedding, extraction) synchronously.
|
||||
///
|
||||
/// Unlike `put_doc`, this waits for the entire process to complete.
|
||||
/// Serialised against the background worker via the shared
|
||||
/// [`IngestionState`] singleton lock — only one ingestion runs at a time.
|
||||
pub async fn ingest_doc(
|
||||
&self,
|
||||
request: MemoryIngestionRequest,
|
||||
) -> Result<MemoryIngestionResult, String> {
|
||||
self.inner.ingest_document(request).await
|
||||
let state = self.ingestion_queue.state();
|
||||
let _guard = state.acquire().await;
|
||||
|
||||
let title = request.document.title.clone();
|
||||
let namespace = request.document.namespace.clone();
|
||||
// Synthetic id until upsert assigns one — purely for the snapshot.
|
||||
let placeholder_id = format!("sync:{title}");
|
||||
|
||||
let queue_depth = state.snapshot().queue_depth;
|
||||
state.mark_running(&placeholder_id, &title, &namespace);
|
||||
crate::core::event_bus::publish_global(
|
||||
crate::core::event_bus::DomainEvent::MemoryIngestionStarted {
|
||||
document_id: placeholder_id.clone(),
|
||||
title,
|
||||
namespace: namespace.clone(),
|
||||
queue_depth,
|
||||
},
|
||||
);
|
||||
|
||||
let started = std::time::Instant::now();
|
||||
let outcome = self.inner.ingest_document(request).await;
|
||||
let elapsed_ms = started.elapsed().as_millis() as u64;
|
||||
let success = outcome.is_ok();
|
||||
|
||||
// Use the same placeholder id as the matching MemoryIngestionStarted
|
||||
// event so subscribers can correlate start/complete pairs. The real
|
||||
// upstream-assigned document id is available on `Ok(outcome)` for
|
||||
// callers that need it.
|
||||
state.mark_completed(
|
||||
&placeholder_id,
|
||||
success,
|
||||
chrono::Utc::now().timestamp_millis(),
|
||||
);
|
||||
crate::core::event_bus::publish_global(
|
||||
crate::core::event_bus::DomainEvent::MemoryIngestionCompleted {
|
||||
document_id: placeholder_id,
|
||||
namespace,
|
||||
success,
|
||||
elapsed_ms,
|
||||
queue_depth: state.snapshot().queue_depth,
|
||||
},
|
||||
);
|
||||
|
||||
outcome
|
||||
}
|
||||
|
||||
/// Returns the shared ingestion state — singleton lock + status snapshot.
|
||||
/// Used by the `openhuman.memory_ingestion_status` RPC handler.
|
||||
pub fn ingestion_state(&self) -> IngestionState {
|
||||
self.ingestion_queue.state()
|
||||
}
|
||||
|
||||
/// Specialized method for syncing skill data into memory.
|
||||
|
||||
@@ -1016,6 +1016,26 @@ async fn json_rpc_memory_sync_and_learn() {
|
||||
"non-existent namespace must be filtered out"
|
||||
);
|
||||
|
||||
// ── memory_ingestion_status: idle on a fresh store ──────────────────────
|
||||
let ing_status = post_json_rpc(
|
||||
&rpc_base,
|
||||
7006,
|
||||
"openhuman.memory_ingestion_status",
|
||||
json!({}),
|
||||
)
|
||||
.await;
|
||||
let ing_result = assert_no_jsonrpc_error(&ing_status, "memory_ingestion_status");
|
||||
assert_eq!(
|
||||
ing_result.get("running"),
|
||||
Some(&json!(false)),
|
||||
"ingestion must be idle on a fresh store, got: {ing_result}"
|
||||
);
|
||||
assert_eq!(
|
||||
ing_result.get("queue_depth").and_then(Value::as_u64),
|
||||
Some(0),
|
||||
"queue_depth must be 0 on a fresh store"
|
||||
);
|
||||
|
||||
mock_join.abort();
|
||||
rpc_join.abort();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user