From 97f7c52c1ef7298e628a85dee41d088230c0cab2 Mon Sep 17 00:00:00 2001 From: Travis James Date: Sat, 3 Oct 2026 22:31:09 -0500 Subject: [PATCH] fix(worker): attribute learning to project/team/role and commit the prompt snapshot LearningJob gains optional projectId/teamId/roleId. enqueue_memory now derives surreal-memory scope keys from them (user_id = project id; agent_id = /, else @project; shared scope = @global) instead of always sending agent_id null, which left every record unattributed. Legacy add_task_step operations get non-null keys too. Already-hashed payloads are left untouched so stored payload hashes and receipts stay valid. process_job commits the prompt snapshot after upsert so session records become visible to pk context. Co-Authored-By: Claude Opus 5.5 --- pk-learning-worker/src/main.rs | 80 ++++++- pk-learning-worker/tests/attribution.rs | 269 ++++++++++++++++++++++++ 2 files changed, 342 insertions(+), 7 deletions(-) create mode 100644 pk-learning-worker/tests/attribution.rs diff --git a/pk-learning-worker/src/main.rs b/pk-learning-worker/src/main.rs index 00f9fd6..801855b 100644 --- a/pk-learning-worker/src/main.rs +++ b/pk-learning-worker/src/main.rs @@ -100,6 +100,16 @@ struct LearningJob { scope: LearningScope, #[serde(default)] attempt: u32, + /// Resolved by the hook's single project-id resolver; falls back to + /// `project_scope(project_root)` when absent (jobs queued by older hooks). + #[serde(default)] + project_id: Option, + /// Agent-team id of the authoring agent, when it resolved to a team role. + #[serde(default)] + team_id: Option, + /// Role id within `team_id`. + #[serde(default)] + role_id: Option, } #[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)] @@ -556,6 +566,17 @@ async fn process_job(root: &Path, pending_path: &Path) -> Result<()> { ) .await?; } + // Publish the store's current entries as the committed prompt snapshot that + // `pk context` reads. Without this, session records were written to disk + // but never became recallable. + pk_store::commit_prompt_snapshot( + &target_kb, + match job.scope { + LearningScope::Project => "project", + LearningScope::Shared => "shared", + }, + store.snapshot().await?, + )?; append_learning_log(&job, &packet)?; enqueue_memory(root, &job, &packet)?; durable_rename(&processing, &completed)?; @@ -724,14 +745,38 @@ fn append_learning_log(job: &LearningJob, packet: &str) -> Result<()> { Ok(()) } +/// The surreal-memory scope keys for a job (design: one `agent_id` per +/// visibility level, never null). A role-attributed project lesson is private +/// to `/`; an unattributed one is project-visible; shared-scope +/// learning is global. +fn memory_identity(job: &LearningJob) -> (String, String) { + let present = |value: &Option| { + value + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_owned) + }; + match job.scope { + LearningScope::Shared => ("@global".to_owned(), "@global".to_owned()), + LearningScope::Project => { + let user_id = + present(&job.project_id).unwrap_or_else(|| project_scope(&job.project_root)); + let agent_id = match (present(&job.team_id), present(&job.role_id)) { + (Some(team), Some(role)) => format!("{team}/{role}"), + _ => "@project".to_owned(), + }; + (user_id, agent_id) + } + } +} + fn enqueue_memory(root: &Path, job: &LearningJob, packet: &str) -> Result<()> { + let (user_id, agent_id) = memory_identity(job); let arguments = json!({ "content": packet, - "user_id": match job.scope { - LearningScope::Project => project_scope(&job.project_root), - LearningScope::Shared => "global".to_owned(), - }, - "agent_id": null, + "user_id": user_id, + "agent_id": agent_id, "session_id": job.session_id, "categories": ["karpathy", "session-learning"] }); @@ -1295,6 +1340,9 @@ fn normalize_operation(mut operation: MemoryOperation) -> Result Result { match method { + // Already-normalized payloads pass through unchanged: their stored + // payload hash (and any server receipt) is bound to these exact bytes. + // New operations get non-null scope keys at the source (memory bridge). "add_memory" | "create_task_stream" => Ok(arguments.clone()), "add_task_step" if arguments.get("stream_name").is_some() => Ok(arguments.clone()), "add_task_step" => { @@ -1312,8 +1360,8 @@ fn normalize_payload(method: &str, arguments: &Value) -> Result { "name": description, "description": description, "idempotency_key": description, - "agent_id": null, - "user_id": null + "agent_id": scope_key(arguments, "agent_id", "@project"), + "user_id": scope_key(arguments, "user_id", &fallback_project_id()) })) } "complete_step" if arguments.get("idempotency_key").is_some() => Ok(arguments.clone()), @@ -1530,6 +1578,24 @@ fn print_status(root: &Path, stale_after: Duration, json_output: bool) -> Result Ok(()) } +fn scope_key(arguments: &Value, key: &str, fallback: &str) -> Value { + match arguments.get(key).and_then(Value::as_str).map(str::trim) { + Some(value) if !value.is_empty() => Value::String(value.to_owned()), + _ => Value::String(fallback.to_owned()), + } +} + +/// Project scope for a legacy bridge operation that did not name one: the +/// explicit environment id, else an explicit `project:unknown` sentinel that the +/// surreal-memory re-key operation can repair (never a null key). +fn fallback_project_id() -> String { + std::env::var("PROMETHEUS_PROJECT_ID") + .ok() + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) + .unwrap_or_else(|| "project:unknown".to_owned()) +} + fn truncate_chars(value: &str, limit: usize) -> String { value.chars().take(limit).collect() } diff --git a/pk-learning-worker/tests/attribution.rs b/pk-learning-worker/tests/attribution.rs new file mode 100644 index 0000000..fbca86b --- /dev/null +++ b/pk-learning-worker/tests/attribution.rs @@ -0,0 +1,269 @@ +//! Learning jobs carry the authoring agent's identity, and the worker must turn +//! it into surreal-memory scope keys (never a null `agent_id`), and must commit +//! the prompt snapshot so `pk context` can recall the session record. + +use serde_json::{json, Value}; +use std::{ + fs, + io::{ErrorKind, Read, Write}, + net::TcpListener, + path::{Path, PathBuf}, + process::Command, + thread, + time::{Duration, Instant}, +}; + +fn transcript(dir: &Path, text: &str) -> PathBuf { + let path = dir.join(format!("transcript-{}.jsonl", text.len())); + fs::write( + &path, + serde_json::to_string(&json!({ + "type": "assistant", + "message": {"role": "assistant", "content": [{"type": "text", "text": text}]} + })) + .unwrap() + + "\n", + ) + .unwrap(); + path +} + +fn queue_job(queue: &Path, event_id: &str, project: &Path, transcript: &Path, extra: Value) { + let mut job = json!({ + "schemaVersion": 2, + "eventId": event_id, + "eventType": "subagent_stop", + "harness": "fixture", + "sessionId": "fixture-session", + "projectRoot": project, + "transcriptPath": transcript, + "capturedAt": "2026-10-04T00:00:00Z", + "payloadDigest": "fixture", + "attempt": 0 + }); + for (key, value) in extra.as_object().unwrap() { + job[key] = value.clone(); + } + fs::write( + queue.join("pending").join(format!("{event_id}.json")), + serde_json::to_vec_pretty(&job).unwrap(), + ) + .unwrap(); +} + +/// Every queued memory operation, in any state directory, keyed by operation id. +fn memory_operations(queue: &Path) -> Vec { + let mut operations = Vec::new(); + for state in fs::read_dir(queue.join("memory")).unwrap().flatten() { + if !state.path().is_dir() { + continue; + } + for file in fs::read_dir(state.path()).unwrap().flatten() { + if file.path().extension().and_then(|value| value.to_str()) == Some("json") { + operations.push(serde_json::from_slice(&fs::read(file.path()).unwrap()).unwrap()); + } + } + } + operations +} + +#[test] +fn jobs_become_attributed_memory_operations_and_recallable_snapshot_entries() { + let fixture = tempfile::tempdir().unwrap(); + let home = fixture.path().join("home"); + let project = fixture.path().join("project"); + let queue = home.join(".prometheus/learning-queue"); + fs::create_dir_all(queue.join("pending")).unwrap(); + fs::create_dir_all(project.join(".git")).unwrap(); + + let attributed = "a".repeat(64); + let unattributed = "b".repeat(64); + queue_job( + &queue, + &attributed, + &project, + &transcript( + fixture.path(), + "LESSON: anchor matchers to the plugin prefix.", + ), + json!({"projectId": "project:fixture", "teamId": "tlm-fixture", "roleId": "api-dev"}), + ); + queue_job( + &queue, + &unattributed, + &project, + &transcript(fixture.path(), "Main-thread turn with no team role."), + json!({}), + ); + + let output = Command::new(env!("CARGO_BIN_EXE_prometheus-learning-worker")) + .env("HOME", &home) + .env("RUST_LOG", "error") + .env_remove("PROMETHEUS_LEARNING_QUEUE") + .args(["--memory-url", "http://127.0.0.1:1", "run-once"]) + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + + let operations = memory_operations(&queue); + let by_id = |id: &str| { + operations + .iter() + .find(|operation| operation["operationId"] == id) + .unwrap_or_else(|| panic!("no memory operation {id}: {operations:#?}")) + }; + let attributed_args = &by_id(&attributed)["arguments"]; + assert_eq!(attributed_args["user_id"], "project:fixture"); + assert_eq!(attributed_args["agent_id"], "tlm-fixture/api-dev"); + + let unattributed_args = &by_id(&unattributed)["arguments"]; + assert_eq!(unattributed_args["agent_id"], "@project"); + let user_id = unattributed_args["user_id"].as_str().unwrap(); + assert!(user_id.starts_with("project:"), "{unattributed_args:#}"); + + for operation in &operations { + for key in ["user_id", "agent_id"] { + assert!( + operation["arguments"][key].is_string(), + "operation with a null {key}: {operation:#}" + ); + } + } + + // The committed prompt snapshot is what `pk context` reads. + let snapshot = + pk_store::read_prompt_snapshot(&project.join(".prometheus/knowledge"), "project") + .expect("worker must commit a project prompt snapshot"); + let ids: Vec<&str> = snapshot + .entries + .iter() + .map(|entry| entry.id.as_str()) + .collect(); + for event_id in [&attributed, &unattributed] { + let expected = format!("karpathy-session-{}", &event_id[..16]); + assert!( + ids.contains(&expected.as_str()), + "{expected} missing from snapshot {ids:?}" + ); + } +} + +#[test] +fn legacy_task_step_operations_are_submitted_with_non_null_scope_keys() { + let fixture = tempfile::tempdir().unwrap(); + let home = fixture.path().join("home"); + let queue = home.join(".prometheus/learning-queue"); + fs::create_dir_all(queue.join("memory/pending")).unwrap(); + fs::write( + queue.join("memory/pending/legacy-step.json"), + serde_json::to_vec_pretty(&json!({ + "schemaVersion": 1, + "operationId": "legacy-step", + "method": "add_task_step", + "arguments": {"stream": "legacy:fixture:phase", "description": "change-001"}, + "dependencies": [], + "state": "pending", + "queuedAt": "2026-10-04T00:00:00Z", + "lastError": null, + "receipt": null + })) + .unwrap(), + ) + .unwrap(); + + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + listener.set_nonblocking(true).unwrap(); + let capture = thread::spawn(move || { + let deadline = Instant::now() + Duration::from_secs(60); + let mut bodies = Vec::new(); + while Instant::now() < deadline { + match listener.accept() { + Ok((mut stream, _)) => { + stream.set_nonblocking(false).unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(5))) + .unwrap(); + let mut request = Vec::new(); + let mut buffer = [0_u8; 8192]; + loop { + match stream.read(&mut buffer) { + Ok(0) => break, + Ok(count) => { + request.extend_from_slice(&buffer[..count]); + let text = String::from_utf8_lossy(&request); + if let Some(split) = text.find("\r\n\r\n") { + let length = text[..split] + .lines() + .find_map(|line| { + line.to_ascii_lowercase() + .strip_prefix("content-length:") + .map(|value| { + value.trim().parse::().unwrap_or(0) + }) + }) + .unwrap_or(0); + if request.len() >= split + 4 + length { + break; + } + } + } + Err(_) => break, + } + } + let text = String::from_utf8_lossy(&request).into_owned(); + // Receipt lookup first: an unknown operation is 404, which + // makes the worker submit it. Refuse the submission itself + // so nothing is marked delivered. + let reply: &[u8] = if text.starts_with("GET /ready ") { + b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: 32\r\nconnection: close\r\n\r\n{\"capabilities\":{\"ledger\":true}}" + } else if text.starts_with("GET ") { + b"HTTP/1.1 404 Not Found\r\ncontent-type: application/json\r\ncontent-length: 2\r\nconnection: close\r\n\r\n{}" + } else { + b"HTTP/1.1 503 Service Unavailable\r\ncontent-length: 0\r\nconnection: close\r\n\r\n" + }; + let _ = stream.write_all(reply); + let done = text.contains("stream_name"); + bodies.push(text); + if done { + break; + } + } + Err(error) if error.kind() == ErrorKind::WouldBlock => { + thread::sleep(Duration::from_millis(20)) + } + Err(error) => panic!("accept failed: {error}"), + } + } + bodies + }); + + let output = Command::new(env!("CARGO_BIN_EXE_prometheus-learning-worker")) + .env("HOME", &home) + .env("RUST_LOG", "error") + .env("PROMETHEUS_PROJECT_ID", "project:legacy-fixture") + .env_remove("PROMETHEUS_LEARNING_QUEUE") + .args(["--memory-url", &format!("http://{address}"), "run-once"]) + .output() + .unwrap(); + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + + let bodies = capture.join().unwrap(); + let submitted = bodies + .iter() + .find(|body| body.contains("stream_name")) + .unwrap_or_else(|| panic!("no task-step submission captured: {bodies:#?}")); + let json_start = submitted.find("\r\n\r\n").unwrap() + 4; + let body: Value = serde_json::from_str(&submitted[json_start..]).unwrap(); + let arguments = &body["payload"]; + assert_eq!(arguments["agent_id"], "@project", "{body:#}"); + assert_eq!(arguments["user_id"], "project:legacy-fixture", "{body:#}"); +}