Merge branch 'strict-workspaces'

daemon/ipc.rs

@@ -7,7 +7,7 @@ use std::{
 use crate::socket::{create_niri_socket, tell};
 use anyhow::{Error, Result, anyhow};
 use niri_ipc::{Event, Request};
-use niri_tag::{TagCmd, TagEvent};
+use niri_tag::{TagCmd, TagEvent, TagState};
 use nix::unistd::geteuid;
 
 use smol::{
@@ -74,10 +74,14 @@ enum EventsReceivable {
 }
 
 #[allow(unreachable_code)]
-pub async fn event_provider(rx: channel::Receiver<TagEvent>) -> Result<()> {
+pub async fn event_provider(
+    rx: channel::Receiver<TagEvent>,
+    fullstate_tx: channel::Sender<channel::Sender<HashMap<u8, TagState>>>,
+) -> Result<()> {
     tracing::debug!("creating event provider");
     let listen = create_provider_socket("event provider", "niri-tag-events").await?;
     let mut sockets = BTreeMap::new();
+    tracing::debug!("beginning event provider loop");
     loop {
         use EventsReceivable::*;
         let recvd: EventsReceivable = future::or(
@@ -98,8 +102,22 @@ pub async fn event_provider(rx: channel::Receiver<TagEvent>) -> Result<()> {
             BadSockets(Vec<String>),
         }
         let res = match recvd {
-            Conn(addr, socket) => {
-                sockets.insert(addr, socket);
+            Conn(addr, mut socket) => {
+                tracing::debug!("received a new event provider connection");
+                sockets.insert(addr, socket.clone());
+                let (t, r) = smol::channel::bounded(1);
+                tracing::debug!("sending fullstate request");
+                fullstate_tx.send(t).await?;
+                match r.recv().await {
+                    Ok(fullstate) => {
+                        tracing::debug!("received fullstate, sending");
+                        let data = serde_json::to_string(&TagEvent::TagFullState(fullstate))?;
+                        if let Err(e) = socket.write_all(&[data.as_bytes(), b"\n"].concat()).await {
+                            tracing::error!("Failed to send fullstate to socket: {}", e);
+                        }
+                    }
+                    Err(e) => tracing::error!("Failed to receive fullstate: {}", e),
+                }
                 Res::Ok
             }
             Event(e) => {
@@ -113,7 +131,7 @@ pub async fn event_provider(rx: channel::Receiver<TagEvent>) -> Result<()> {
                     .into_iter()
                     .fold(Vec::new(), |mut acc, (addr, res)| {
                         if let Err(e) = res {
-                            tracing::warn!("error on event provider socket {}: {}", addr, e);
+                            tracing::warn!("error from event provider client {}: {}", addr, e);
                             acc.push(addr.to_owned());
                         }
                         acc
@@ -131,7 +149,6 @@ pub async fn event_provider(rx: channel::Receiver<TagEvent>) -> Result<()> {
             });
         }
     }
-    tracing::debug!("beginning ipc provider loop");
 
     Ok(())
 }

daemon/main.rs

@@ -21,13 +21,14 @@ fn main() -> Result<()> {
     smol::spawn(ipc::ipc_provider(ipc_tx)).detach();
     // spawn socket listener for events
     let (event_tx, event_rx) = smol::channel::unbounded();
-    smol::spawn(ipc::event_provider(event_rx)).detach();
+    let (fullstate_tx, fullstate_rx) = smol::channel::unbounded();
+    smol::spawn(ipc::event_provider(event_rx, fullstate_tx)).detach();
     // begin managing niri tags
     smol::block_on(async {
         let niri_tag = manager::NiriTag::new(event_tx)
             .await
             .context("Initialising niri tag manager")
             .unwrap();
-        niri_tag.manage_tags(niri_rx, ipc_rx).await
+        niri_tag.manage_tags(niri_rx, ipc_rx, fullstate_rx).await
     })
 }

daemon/manager.rs

@@ -4,7 +4,7 @@ use niri_ipc::{
     Action, Event, Reply, Request, Response, Window, Workspace, WorkspaceReferenceArg,
     state::{EventStreamState, EventStreamStatePart},
 };
-use niri_tag::{TagCmd, TagEvent};
+use niri_tag::{TagCmd, TagEvent, TagState};
 use smol::{
     channel::{self, Sender},
     future,
@@ -16,6 +16,7 @@ use std::collections::HashMap;
 pub struct NiriTag {
     tags: HashMap<u8, bool>,
     windows: HashMap<u64, u8>,
+    active_ws: Vec<u64>,
     state: EventStreamState,
     socket: BufReader<UnixStream>,
     ev_tx: channel::Sender<TagEvent>,
@@ -33,6 +34,7 @@ impl NiriTag {
             windows: HashMap::new(),
             state: EventStreamState::default(),
             socket: create_niri_socket().await?,
+            active_ws: Vec::new(),
             ev_tx,
         })
     }
@@ -234,6 +236,25 @@ impl NiriTag {
                 let _ = self.state.apply(ev.clone());
                 return self.handle_event(ev).await;
             }
+            Receivable::FullState(tx) => {
+                tracing::debug!("received request for full state");
+                let fullstate: HashMap<u8, TagState> = self
+                    .tags
+                    .iter()
+                    .map(|(&t, &enabled)| {
+                        (
+                            t,
+                            TagState {
+                                enabled,
+                                occupied: self.windows.values().filter(|w_t| **w_t == t).count()
+                                    > 0,
+                                urgent: false, // urgency is TODO
+                            },
+                        )
+                    })
+                    .collect();
+                return tx.send(fullstate).await.map_err(|e| anyhow!(e));
+            }
             Receivable::TagCmd(cmd) => match cmd {
                 TagCmd::AddTagToWin(t) => {
                     let win = self.get_focused_window().await?;
@@ -308,8 +329,45 @@ impl NiriTag {
                 self.windows.remove(&id);
                 Ok(())
             }
-            // WorkspaceActivated { .. } => (),
-            // WorkspacesChanged { .. } => (),
+            WorkspaceActivated { id, .. } => {
+                if !self.active_ws.contains(&id) {
+                    let q = query(&mut self.socket, Request::Workspaces).await?;
+                    let wsid = if let Reply::Ok(Response::Workspaces(workspaces)) = q {
+                        let new_ws = workspaces
+                            .iter()
+                            .find(|ws| ws.id == id)
+                            .expect("Activated workspace not found in workspace query");
+                        workspaces
+                            .iter()
+                            .find(|ws| {
+                                ws.output == new_ws.output
+                                    && ws.id != new_ws.id
+                                    && self.active_ws.contains(&ws.id)
+                            })
+                            .expect("Could not find a valid niri-tag workspace to return to")
+                            .id
+                    } else {
+                        return Err(anyhow!("Invalid response to workspace query"));
+                    };
+                    tell(
+                        &mut self.socket,
+                        Request::Action(Action::FocusWorkspace {
+                            reference: WorkspaceReferenceArg::Id(wsid),
+                        }),
+                    )
+                    .await
+                } else {
+                    Ok(())
+                }
+            }
+            WorkspacesChanged { workspaces } => {
+                self.active_ws = workspaces
+                    .into_iter()
+                    .filter(|ws| ws.is_active)
+                    .map(|ws| ws.id)
+                    .collect();
+                Ok(())
+            }
             // WorkspaceUrgencyChanged { .. } => (),
             WindowsChanged { windows } => {
                 for w in windows {
@@ -331,16 +389,20 @@ impl NiriTag {
         mut self,
         ev_rx: channel::Receiver<Event>,
         tag_rx: channel::Receiver<TagCmd>,
+        fullstate_rx: channel::Receiver<channel::Sender<HashMap<u8, TagState>>>,
     ) -> Result<()> {
         // base tag is always visible
         self.tags.insert(0, true);
 
         loop {
-            let recvd: Receivable =
-                future::or(async { ev_rx.recv().await.map(Receivable::Event) }, async {
-                    tag_rx.recv().await.map(Receivable::TagCmd)
-                })
-                .await?;
+            let recvd: Receivable = future::or(
+                async { ev_rx.recv().await.map(Receivable::Event) },
+                future::or(
+                    async { tag_rx.recv().await.map(Receivable::TagCmd) },
+                    async { fullstate_rx.recv().await.map(Receivable::FullState) },
+                ),
+            )
+            .await?;
             tracing::debug!("manager received {:?}", recvd);
 
             let res = self.handle_recvd(recvd).await;
@@ -358,4 +420,5 @@ impl NiriTag {
 enum Receivable {
     Event(Event),
     TagCmd(TagCmd),
+    FullState(channel::Sender<HashMap<u8, TagState>>),
 }

lib/main.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use serde::{Deserialize, Serialize};
 
 #[derive(Serialize, Deserialize, Debug)]
@@ -19,4 +21,12 @@ pub enum TagEvent {
     TagUrgent(u8),
     TagEnabled(u8),
     TagDisabled(u8),
+    TagFullState(HashMap<u8, TagState>),
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct TagState {
+    pub enabled: bool,
+    pub occupied: bool,
+    pub urgent: bool,
 }