flaxo/wynba

public

AI coding agent!!

Initial commit: coding-agent4915718a month ago443 lines · 14.0 KB

1	use std::path::PathBuf;
2
3	use anyhow::Result
4	use rusqlite::{params,
5	use serde_json::Value
6
7	pub type DbConn
8
9	#[allow(dead_code)]
10	#[derive(Clone
11	pub struct Session
12	pub id: i64
13	pub uuid: String
14	pub name: String
15	pub provider: String
16	pub model: String
17	pub created_at: i64
18	pub updated_at: i64
19	}
20
21	#[allow(dead_code)]
22	pub struct StoredMessage
23	pub id: i64
24	pub session_id: i64
25	pub role: String
26	pub content: Option
27	pub tool_call_id: Option
28	pub name: Option
29	pub tool_calls: Option
30	pub images: Option
31	pub created_at: i64
32	}
33
34	pub fn default_db_path()
35	if let Ok
36	&& !p.trim
37	return PathBuf::from
38	}
39	let home =
40	PathBuf::from(
41	}
42
43	pub fn open(path
44	if let Some
45	std::fs::create_dir_all
46	}
47	let conn =
48	conn.execute_batch(
49	"PRAGMA journal_mode=WAL;
50	PRAGMA foreign_keys=ON;
51	CREATE TABLE IF NOT EXISTS sessions (
52	id INTEGER PRIMARY KEY AUTOINCREMENT,
53	name TEXT NOT NULL DEFAULT '',
54	provider TEXT NOT NULL,
55	model TEXT NOT NULL,
56	uuid TEXT,
57	created_at INTEGER NOT NULL,
58	updated_at INTEGER NOT NULL
59	);
60	CREATE TABLE IF NOT EXISTS messages (
61	id INTEGER PRIMARY KEY AUTOINCREMENT,
62	session_id INTEGER NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
63	role TEXT NOT NULL,
64	content TEXT,
65	tool_call_id TEXT,
66	name TEXT,
67	tool_calls TEXT,
68	created_at INTEGER NOT NULL
69	);
70	CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id, id);
71	CREATE TABLE IF NOT EXISTS provider_configs (
72	provider TEXT PRIMARY KEY,
73	config TEXT NOT NULL,
74	updated_at INTEGER NOT NULL
75	);
76	CREATE TABLE IF NOT EXISTS file_backups (
77	id INTEGER PRIMARY KEY AUTOINCREMENT,
78	session_id INTEGER NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
79	path TEXT NOT NULL,
80	kind TEXT NOT NULL,
81	backup_path TEXT,
82	msg_before INTEGER NOT NULL,
83	created_at INTEGER NOT NULL
84	);
85	")?;
86	migrate_sessions(&conn)?
87	migrate_messages_images(&conn)?
88	Ok(conn)
89	}
90
91	fn migrate_messages_images(conn:
92	let has =
93	let mut stmt
94	let mut rows
95	let mut found
96	while let Some
97	if row.get::
98	found = true;
99	break;
100	}
101	}
102	found
103	};
104	if !has {
105	conn.execute("ALTER TABLE messages ADD COLUMN images TEXT
106	}
107	Ok(())
108	}
109
110	fn migrate_sessions(conn:
111	let has_uuid =
112	let mut stmt
113	let mut rows
114	let mut found
115	while let Some
116	if row.get::
117	found = true;
118	break;
119	}
120	}
121	found
122	};
123	if !has_uuid {
124	conn.execute("ALTER TABLE sessions ADD COLUMN uuid TEXT
125	}
126	let tx =
127	{
128	let mut stmt
129	let ids: Vec
130	.query_map([], \|r
131	.collect::<rusqlite::
132	for id in ids
133	tx.execute(
134	"UPDATE sessions SET uuid=?2 WHERE id=?1",
135	params![id,
136	)?;
137	}
138	}
139	tx.commit()?;
140	conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_sessions_uuid ON sessions(uuid)
141	Ok(())
142	}
143
144	pub fn create_session(conn
145	let now =
146	let name =
147	let uuid =
148	conn.execute(
149	"INSERT INTO sessions (name, provider, model, uuid, created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5, ?5)",
150	params![name,
151	)?;
152	Ok((conn.
153	}
154
155	fn map_session(r:
156	Ok(Session
157	id: r.get
158	uuid: r.get
159	name: r.get
160	provider: r.get
161	model: r.get
162	created_at: r.get
163	updated_at: r.get
164	})
165	}
166
167	pub fn get_session(conn
168	let mut stmt
169	let mut rows
170	if let Some
171	Ok(Some
172	} else {
173	Ok(None
174	}
175	}
176
177	pub fn get_session_by_uuid(conn
178	let mut stmt
179	let mut rows
180	if let Some
181	Ok(Some
182	} else {
183	Ok(None
184	}
185	}
186
187	pub fn get_session_by_ref(conn
188	if let Some
189	return Ok(Some
190	}
191	if let Ok
192	return get_session(conn,
193	}
194	Ok(None
195	}
196
197	pub fn list_sessions(conn
198	let mut stmt
199	let rows =
200	let mut out
201	for row in rows
202	out.push(row
203	}
204	Ok(out)
205	}
206
207	pub fn delete_session(conn
208	let n =
209	Ok(n)
210	}
211
212	pub fn update_session_model(conn
213	let now =
214	conn.execute(
215	"UPDATE sessions SET provider=?2, model=?3, updated_at=?4 WHERE id=?1",
216	params![id,
217	)?;
218	Ok(())
219	}
220
221	pub fn touch_session(conn
222	let now =
223	conn.execute("UPDATE sessions SET updated_at=?2 WHERE id=?1
224	Ok(())
225	}
226
227	pub fn set_session_name(conn
228	conn.execute("UPDATE sessions SET name=?2 WHERE id=?1
229	Ok(())
230	}
231
232	#[allow(clippy::
233	pub fn add_message(
234	conn: &Connection
235	session_id: i64
236	role: &str
237	content: Option
238	tool_call_id: Option
239	name: Option
240	tool_calls: Option
241	images: Option
242	) -> Result<i64
243	let now =
244	conn.execute(
245	"INSERT INTO messages (session_id, role, content, tool_call_id, name, tool_calls, images, created_at) VALUES (?1,?2,?3,?4,?5,?6,?7,?8)",
246	params![session_id,
247	)?;
248	Ok(conn.
249	}
250
251	pub fn get_messages(conn
252	let mut stmt
253	"SELECT id, session_id, role, content, tool_call_id, name, tool_calls, images, created_at FROM messages WHERE session_id=?1 ORDER BY id",
254	)?;
255	let rows =
256	Ok(StoredMessage
257	id: r.get
258	session_id: r.get
259	role: r.get
260	content: r.get
261	tool_call_id: r.get
262	name: r.get
263	tool_calls: r.get
264	images: r.get
265	created_at: r.get
266	})
267	})?;
268	let mut out
269	for row in rows
270	out.push(row
271	}
272	Ok(out)
273	}
274
275	pub fn current_message_id(conn
276	let id: Option
277	.query_row(
278	"SELECT MAX(id) FROM messages WHERE session_id=?1",
279	params![session_id],
280	\|r\| r.
281	)
282	.optional()?;
283	Ok(id.
284	}
285
286	pub fn revert_messages(conn
287	let all =
288	let deleted: Vec
289	let tx =
290	tx.execute(
291	"DELETE FROM messages WHERE session_id=?1 AND id>=?2",
292	params![session_id,
293	)?;
294	tx.commit()?;
295	Ok(deleted)
296	}
297
298	pub fn copy_messages_up_to(
299	conn: &Connection
300	src_session: i64
301	dst_session: i64
302	upto_msg_id: i64
303	) -> Result<()
304	let msgs =
305	let tx =
306	for m in msgs
307	if m.id >
308	continue;
309	}
310	tx.execute(
311	"INSERT INTO messages (session_id, role, content, tool_call_id, name, tool_calls, images, created_at) VALUES (?1,?2,?3,?4,?5,?6,?7,?8)",
312	params![
313	dst_session,
314	m.role,
315	m.content,
316	m.tool_call_id,
317	m.name,
318	m.tool_calls,
319	m.images,
320	m.created_at
321	],
322	)?;
323	}
324	tx.commit()?;
325	Ok(())
326	}
327
328	pub fn record_file_backup(
329	conn: &Connection
330	session_id: i64
331	path: &str
332	kind: &str
333	backup_path: Option
334	) -> Result<()
335	let now =
336	let msg_before =
337	conn.execute(
338	"INSERT INTO file_backups (session_id, path, kind, backup_path, msg_before, created_at) VALUES (?1,?2,?3,?4,?5,?6)",
339	params![session_id,
340	)?;
341	Ok(())
342	}
343
344	pub fn restore_file_backups(conn
345	let mut stmt
346	"SELECT id, path, kind, backup_path FROM file_backups WHERE session_id=?1 AND msg_before>=?2 ORDER BY id DESC",
347	)?;
348	let rows =
349	Ok((
350	r.get::<_
351	r.get::<_
352	r.get::<_
353	))
354	})?;
355	let mut affected
356	for row in rows
357	affected.push(row
358	}
359	let mut count
360	for (path, kind
361	match kind.as_str()
362	"created" => {
363	if std::fs
364	count += 1;
365	}
366	}
367	"before" \| "deleted"
368	if let Some
369	&& let Ok
370	&& let Some
371	{
372	let _ =
373	if std::fs
374	count += 1;
375	}
376	}
377	}
378	_ => {}
379	}
380	}
381	conn.execute(
382	"DELETE FROM file_backups WHERE session_id=?1 AND msg_before>=?2",
383	params![session_id,
384	)?;
385	for (_, _
386	if let Some
387	let _ =
388	}
389	}
390	Ok(count)
391	}
392
393	pub fn backup_dir(session_id
394	let home =
395	std::path::PathBuf
396	.join(".coding-agent")
397	.join("backups")
398	.join(session_id.to_string
399	}
400
401	use rusqlite::OptionalExtension
402
403	pub fn set_provider_config(conn
404	let now =
405	let s =
406	conn.execute(
407	"INSERT INTO provider_configs (provider, config, updated_at) VALUES (?1, ?2, ?3)
408	ON CONFLICT(provider) DO UPDATE SET config=excluded.config, updated_at=excluded.updated_at",
409	params![provider,
410	)?;
411	Ok(())
412	}
413
414	pub fn get_provider_config(conn
415	let mut stmt
416	let mut rows
417	if let Some
418	Ok(serde_json::
419	} else {
420	Ok(None
421	}
422	}
423
424	pub fn delete_provider_config(conn
425	Ok(conn.
426	}
427
428	pub fn list_provider_configs(conn
429	let mut stmt
430	let rows =
431	let provider: String
432	let config: String
433	Ok((provider,
434	})?;
435	let mut out
436	for row in rows
437	let (provider,
438	if let Ok
439	out.push((provider
440	}
441	}
442	Ok(out)
443	}